[yihwang-nv]

@coderabbitai summary

Description

AttentionMetadata is no longer extended by subclassing. It becomes a small core (batch layout, sequence/KV lengths, paged KV, workspace, CUDA-graph buffers) plus pluggable MetadataExtension objects. Every feature, sparse algorithm and model that used to add fields or override methods on TrtllmAttentionMetadata (or subclass it) now attaches an extension, reached with metadata.ext(T) (raises if absent) or metadata.find_ext(T) (returns None).

The goal is to let modeling_v2 targets reuse the runtime's metadata instead of reimplementing it, mirroring the C++ V2 refactor on the Python side. No C++ file is touched.

Framework (tensorrt_llm/_torch/attention/backends/extensions.py)

• MetadataExtension: no-op lifecycle hooks (allocate, prepare_scope, pre_derive_lengths, prepare, post_derive, begin_forward, on_kv_lens_changed, on_spec_dec_param, on_spec_dec_restored, draft_swap/draft_restore, adp_token_counts, workspace_bytes, kv_cache_pools, ...), requires dependencies and graph_temporary_attrs.
• ExtensionRegistry: dependency-ordered attach, lookup by exact type then isinstance, rebinding on metadata copies (CUDA graph, cross-attention).
• Core prepare() is a template: scopes → pre_derive_lengths → _prepare_core() → extension prepare → post_derive.
• AttentionBackend.metadata_extensions + Backend.create_metadata(...); sparse backends declare their algorithm extension instead of a metadata subclass.

Extensions

Kind	Extensions
Feature (backends/features/)	HelixMetadata, DraftKVMetadata, FlashMlaMetadata, SpecDecodingMetadata, RecurrentMetadata, Fp4MlaState
Sparse algorithm (replaces deleted subclasses)	RocketMetadata, QSAMetadata, DSAMetadata, DeepseekV4Metadata, CSA2Metadata, MiniMaxM3TritonMetadata, MiniMaxM3MsaMetadata; shared SparseAlgorithmExtension / MTPIndexShareExtension

Cross-attention, beam/cache indirection, MLA cached-context state, mixed-precision KV staging and executor-owned batch fields stay in the core.

Consumers updated: FMHA libraries, MLA/attention modules, FP4 MLA, DSA/DSv4/CSA2/QSA/MiniMax modules, models (qwen3_next, qwen4_exp, nemotron_h, glm5_next, minimaxm3, deepseekv4, step3p7), mamba/KDA mixers and cache managers, speculative workers, executor (metadata construction, decoder inputs/runner, CUDA-graph runner), modeling_v2 targets, and ATTENTION_DEVELOPER_GUIDE.md.

Behaviour changes to review

1. Every sparse backend on the TRTLLM core now receives cache_indirection (matters only with beam search).
2. FlashMLA's draft swap now also runs for DSA/DeepSeek-V4 (their old prepare_for_draft_forward overrides skipped super()).
3. In a spec-dec step, extensions get on_kv_lens_changed(True) only; the old QSA path also ran a second plain on_update_kv_lens.
4. CSA2 ADP input stage uses the core num_tokens (was padded_num_tokens or num_tokens).
5. CSA2: request metadata without a KV manager reports tokens_per_block None (was 128); tiles report 128 via kv_cache_pools.
6. The strict CUDA-graph buffer check also snapshots extension tensors (keys "<ExtClass>.<attr>").
7. Fp4MlaState prepare runs after the core prepare and rebinds prompt_lens_*_runtime itself.
8. The legacy sparse_attention_config= / enable_flash_mla= / enable_ragged_verification= metadata constructor kwargs are removed; build metadata with Backend.create_metadata(...).

Commits

The change was ported onto a newer main. The first commit is the original change cherry-picked verbatim and contains raw conflict markers (it does not build on its own); the second commit is the conflict resolution plus porting of code main added since; the third fixes two test doubles. Please review the combined diff; the PR is intended to be squash-merged.

Test Coverage

Single-GPU slice on B200 (SM100), file-sharded: all of tests/unittest/_torch/attention except multi_gpu, plus executor/engine/test_metadata.py, executor/test_cuda_graph_capture_replay.py, executor/kv_cache/test_mamba_cache_manager.py, modeling/test_qsa_runtime_wiring.py, modeling/test_glm5_next_contracts.py, modules/mamba/test_gdn_kernel_optimizations.py, modules/kimi_kda/, models/test_minimax_m3.py, modeling_v2/models/qwen4_exp/, speculative/hw_agnostic/, visual_gen/test_trtllm_attention_metadata.py.

• ~9,550 passed in the first pass. The failures were: 4 tests using SimpleNamespace metadata doubles without ext() (fixed in the third commit); trtllm-gen FMHA JIT failures caused by a local environment NVRTC mismatch (pip nvidia-cuda-nvrtc 13.0 loaded against CUDA 13.4 headers); and the qwen3_8_flash_next_nvfp4__sm_103__tp1 row-tile test, which requires SM103.
• Every file with a failure re-run with the matching NVRTC: 1,534 passed, 0 failed.
• tests/unittest/_torch/attention/test_attention_op_sync.py statically checks the thop.attention call against the extension fields.

Not yet run: multi-GPU Helix (tests/unittest/_torch/attention/multi_gpu), SM103-only tests.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

🤖 Generated with Claude Code