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.

[ErenAta16]

Deriving the metadata from the operators that were actually selected, rather than re-deducing it from config flags, is clearly the right direction — the divergence cases in the motivation are real and the config-flag approach can't express "one model, two attention operators with different configurations". The CudaSequenceMetadata.from_step_context stacking so both cumulative lengths come out of one cumsum is a nice touch too.

One thing about the collection mechanism I'd want settled before it lands.

The registration silently no-ops outside the scope

_active_implementations: ContextVar[list[Any] | None] = ContextVar(
    'cuda_step_meta_implementations', default=None)

def register_step_metadata_impl(impl: Any) -> None:
    implementations = _active_implementations.get()
    if implementations is not None:
        implementations.append(impl)

If an operator is constructed outside collect_step_metadata(...), its registration is discarded with no error and no log line. The model then runs with a CudaStepMetaPlan that is missing that operator, and the symptom appears later — wrong metadata, or a missing buffer at capture time — rather than at the point the registration was lost.

That is the same failure this PR exists to remove. The old path could diverge because config flags disagreed with the selected operators; the new path can diverge because a registration didn't land. The second is quieter.

Two ways it happens, both measured:

construction outside the scope   -> collected: []      (silent)
construction in another thread   -> collected: []      (silent)
ThreadPoolExecutor               -> collected: []      (silent)
contextvars.copy_context() + run -> collected: ['x']   (works)

ContextVar values do not propagate into threading.Thread or a ThreadPoolExecutor worker — a new thread starts from an empty context and sees default=None. Whether lmdeploy ever builds CUDA operators off the constructing thread I can't tell from outside, but the engine is threaded enough that I would not want it to depend on nobody ever doing so, and lazy or deferred operator construction has the same effect without any threading at all.

The cheap version of the fix is to make the loss visible rather than to change the mechanism:

def register_step_metadata_impl(impl: Any) -> None:
    implementations = _active_implementations.get()
    if implementations is None:
        logger.debug('step metadata impl %r registered outside a build scope; ignored', impl)
        return
    implementations.append(impl)

A debug line turns "metadata is subtly wrong on this model" into one grep. If registration outside a scope is genuinely a programming error rather than a supported case, an assert would be better still — the current if ... is not None reads like it is tolerating something deliberate, and the docstring ("only inside the active CUDA build scope") suggests it is not.

A smaller one in the same area

collect_step_metadata assigns the plan inside the try, after yield:

try:
    yield
    ctx_mgr.backend_step_meta_plan = CudaStepMetaPlan.from_implementations(implementations)
finally:
    _active_implementations.reset(token)

so if model construction raises, backend_step_meta_plan keeps whatever it had before — possibly a plan from a previous model on a reused ctx_mgr. The finally correctly resets the ContextVar either way, so this is only about the plan attribute. Clearing it before the yield, or setting it to None in the failure path, would make a half-built model fail loudly instead of running against a stale plan.

Both notes are about the seam rather than the metadata work itself, which reads well — the comment explaining why FA3 exposes a different type from fill_cudagraph_buffers than it allocates is exactly the sort of thing that saves the next reader an hour.

[RunningLeon]

LGTM

[lvhan028]

May reslove the conflicts