refactor(pytorch): derive CUDA step metadata from selected operators

[grimoire]

Motivation

CUDA step metadata was previously prepared mainly from model configuration flags in update_step_context. This can diverge from the operators actually selected by backend dispatch, especially when:

• different models require different metadata;
• one model contains attention operators with different configurations;
• FA3, FlashMLA, or other kernels own specialized metadata;
• CUDA graph capture requires persistent or lazily initialized metadata buffers.

Changes

• Collect metadata providers from the CUDA operators instantiated during model construction.
• Resolve a model-owned CudaStepMetaPlan without traversing the completed model.
• Build shared sequence metadata once per step.
• Move specialized metadata construction to operator-owned builders for:
  • default Triton attention;
  • FA3;
  • FlashMLA;
  • gated delta rule.
• Deduplicate operators with the same metadata configuration while assigning different groups to different configurations.
• Add typed graph-buffer and operator-metadata contracts to attention builders.
• Add a pre-capture lifecycle hook for lazily initialized FlashMLA metadata.
• Preserve typed FA3 metadata when refreshing legacy fields before graph capture.
• Keep the existing model-config-driven path as a compatibility fallback for unsupported implementations.
• Add a default no-op model build context to the base backend, leaving non-CUDA backends unchanged.

Compatibility

• Existing single-attention models continue to receive the legacy metadata fields.
• Models with multiple attention configurations use group-specific metadata.
• Unknown or unsupported implementations fall back to the previous preparation path.
• CAMB, MACA, and other non-CUDA backends do not participate in CUDA metadata collection.
• CUDA graph state is rebuilt together with the model during level-2 sleep/wakeup.

[Copilot]

Pull request overview

This PR refactors CUDA step-metadata preparation in the PyTorch backend so metadata is derived from the CUDA operators actually instantiated during model construction, enabling grouped (per-configuration) metadata and better alignment with CUDA graph capture requirements while retaining a legacy, model-config-driven fallback path.

Changes:

• Introduces a model-owned CudaStepMetaPlan (resolved during model build) and runner-owned graph buffers to build/fill grouped attention metadata and per-step updaters.
• Moves FlashMLA and FA3 metadata construction into operator-owned builder/update helpers and routes CUDAGraph capture/warmup through the resolved plan when supported.
• Adds a backend model_build_context scope used during model construction to collect implementations and attach the resolved plan (no-op for non-CUDA backends).

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
lmdeploy/pytorch/models/utils/cudagraph.py	Threads an optional CudaStepMetaPlan through cudagraph buffer allocation/fill, with compatibility fallbacks for legacy paths.
lmdeploy/pytorch/models/patch.py	Wraps model build in backend model_build_context and passes an explicit StepContextManager into model construction.
lmdeploy/pytorch/model_inputs.py	Adds storage on StepContextManager for backend-resolved step-metadata plans.
lmdeploy/pytorch/backends/cuda/step_metadata.py	Adds the core step-metadata planning, grouped builder contracts, and build-time collection mechanism.
lmdeploy/pytorch/backends/cuda/op_backend.py	Switches step-context preparation to use CudaSequenceMetadata + implementation-derived plans, retaining a legacy fallback.
lmdeploy/pytorch/backends/cuda/graph_runner.py	Wires the resolved step-metadata plan into cudagraph capture and pre-capture lifecycle hooks.
lmdeploy/pytorch/backends/cuda/gated_delta_rule.py	Refactors gated-delta prefill prep into a reusable helper and integrates it as a step-updater provider.
lmdeploy/pytorch/backends/cuda/attention/mla.py	Adds FlashMLA typed metadata + builder contract, including cudagraph buffer handling and pre-capture behavior.
lmdeploy/pytorch/backends/cuda/attention/fa3.py	Adds FA3 typed metadata + builder contract and consolidates scheduler-metadata construction/update helpers.
lmdeploy/pytorch/backends/cuda/attention/default.py	Adds shared sequence-to-attention metadata projection and default attention builder/registration hooks.
lmdeploy/pytorch/backends/base.py	Adds a default no-op model_build_context for backend-specific model-build scoping.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.