[TRTLLM-13580][test] Add model-derived PyTorch attention backend test suite

[yuxianq]

What

A unified, model-derived test suite for the PyTorch dense attention backends
(TrtllmAttention, FlashInferAttention, VanillaAttention). VanillaAttention
is the golden reference; TRTLLM and FlashInfer are validated against it through a
real KVCacheManager (standalone, no PyExecutor).

Coverage

• Model-derived sweep (model_configs.py): the distinct attention configs of
  the supported models, so each case maps to a real workload — GQA / MHA / MQA,
  sliding window, neox / gptj / mrope / yarn RoPE, head_dim 64 / 128 / 256, MLA,
  cross-attention, and bidirectional no-cache (DiT / encoder). Each config is
  crossed over batch phase {context, gen, mixed}, precision {bf16, fp8-KV} and
  KV-cache manager {v1, v2} at page_size=32 / HND layout; the non-default values
  (page_size=64, NHD layout, fp16) are exercised on a representative GQA/MHA/MQA
  subset to bound the matrix.
• MLA: absorbed-generation golden (MQA-on-latent) added to VanillaAttention,
  plus the up-projected context pass (asymmetric K/V). SDPA uses enable_gqa.
• CUDA-graph reuse: gen-only batches are also replayed through a captured
  graph and must match the eager golden.
• Cross-attention (incl. q == kv on TRTLLM via the standard metadata
  plumbing), FlashInfer kv_layout {NHD, HND}, and fp8 KV context.
• Per-backend unsupported (dtype / layout / feature) combinations are skipped via
  a capability matrix rather than failing, keeping the suite green while
  documenting the gaps (TRTLLM NHD layout, fp8 generation scale state,
  q != kv cross isolation).

Validation

311 cases, green on H200 (sm90) and B200 (sm100).

Notes

This branch deliberately omits the optional env-gated forward-capture / replay /
minimizer tooling; the model-derived synthetic sweep is the sole driver.

Summary by CodeRabbit

Release Notes

• New Features

  • Added support for MLA (absorbed latent attention) execution path in attention computation.
  • Optimized attention computation to leverage PyTorch's native GQA broadcasting.

• Tests

  • Expanded attention backend testing infrastructure with comprehensive validation harness and model configuration suite.

[coderabbitai]

📝 Walkthrough

Walkthrough

Extends VanillaAttention with absorbed MLA execution paths (_mla_forward_generation, _mla_forward_context, mla_rope_generation) and switches GQA handling from repeat_kv to scaled_dot_product_attention(..., enable_gqa=True). Simultaneously introduces a comprehensive unified attention backend test harness: BackendCase, model config catalog, capability matrix, KV cache fill utilities, and parameterized backend tests.

Changes

VanillaAttention MLA support and unified attention backend test harness

Layer / File(s)	Summary
VanillaAttention MLA implementation and GQA refactor tensorrt_llm/_torch/attention_backend/vanilla.py	Adds MLA parameter detection in __init__, support_mla(), mla_rope_generation(), _mla_forward_generation() (paged latent cache append + MQA SDPA), _mla_forward_context() (causal asymmetric head-dim SDPA), and MLA routing in forward(); replaces repeat_kv with enable_gqa=True in _single_request_attn_forward and _no_kv_cache_sdpa_fallback.
Test model configs catalog and backend capability matrix tests/unittest/_torch/attention/model_configs.py, tests/unittest/_torch/attention/capability_matrix.py	Introduces ModelAttnConfig frozen dataclass and four config groups (_STANDARD, _MLA, _CROSS, _NO_CACHE) exported as MODEL_CONFIGS; defines BACKEND_CAPS feature matrix, required_features(), and unsupported_reason() with SM-version gating and TRTLLM-specific skip conditions.
Shared KV cache fill and RoPE reference utilities tests/unittest/_torch/attention/kv_cache_utils.py	Adds fill_kv_cache_logical writing logical per-sequence KV tensors into paged cache with NHD/HND layout and fp8 zero-init; adds make_position_ids for packed varlen positions; adds lru_cache'd _rotary_embedding and apply_rope using the deterministic pure-PyTorch path.
BackendCase dataclass, input generators, and cache managers tests/unittest/_torch/attention/attention_test_harness.py	Introduces harness module with tolerance defaults, BackendCase kw-only dataclass with all config fields and derived properties, generate_inputs, _build_kv_cache_manager (v1/v2, NVFP4), generate_mla_gen_inputs, _build_mla_kv_cache_manager, _fill_mla_cache, generate_mla_context_inputs, and _tolerances.
Harness execution: _maybe_rope, _capture_replay, run_backend, run_case tests/unittest/_torch/attention/attention_test_harness.py	Adds _maybe_rope for fused/non-fused RoPE routing; _capture_replay for CUDA-graph decode capture-then-replay; run_backend dispatching MLA, standard, and cross-attention paths with cache fill, FlashInfer layout, and cu_seqlens; run_case as the golden-vs-backend verification loop with tolerance assertions and aggregated failure reporting.
Parameterized backend test suite and vanilla SDPA stub fix tests/unittest/_torch/attention/test_attention_backends.py, tests/unittest/_torch/attention/test_vanilla_attention.py	Builds MODEL_CASES by expanding MODEL_CONFIGS across precision/layout/page/v1-v2 dimensions; adds test_attention_backend_model and test_split_consistency; fixes fake_sdpa to accept **kwargs instead of a positional scale parameter.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description check	✅ Passed	The pull request description is comprehensive and well-structured, clearly explaining what, why, coverage scope, and validation results, matching the template structure with detailed explanation of objectives.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title directly matches the PR's main objective: introducing a unified, model-derived test suite for PyTorch attention backends.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 7

🧹 Nitpick comments (3)

tests/unittest/_torch/attention/capability_matrix.py (1)

58-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Align Vanilla MLA capability with the backend contract.

VanillaAttention.support_mla() now returns True, but BACKEND_CAPS["VANILLA"]["mla"] is still False. This drift can lead to incorrect skip decisions or contract-check noise when the matrix is reused.

Suggested patch

     "VANILLA": dict(
         paged=False,
         fp8_kv=True,
         fp4_kv=False,
         sliding_window=True,
         no_cache=True,
         sparse=False,
-        mla=False,
+        mla=True,
         cross_attn=True,
         kv_layouts=("NHD",),  # reads the NHD get_buffers view
     ),

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/attention/capability_matrix.py` around lines 58 - 67,
The BACKEND_CAPS["VANILLA"] capability matrix has mla set to False, but
VanillaAttention.support_mla() now returns True, causing a mismatch between the
capability declaration and the actual backend contract. Update the mla field in
the VANILLA dictionary from False to True to align the capability matrix with
the current backend implementation.

tensorrt_llm/_torch/attention_backend/vanilla.py (1)

620-633: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add negative-path MLA coverage for the new routing branch.

Coverage is sufficient for MLA happy paths, but insufficient for input-contract failures in this new branch. Please add focused tests in tests/unittest/_torch/attention/test_vanilla_attention.py for:

• AttentionInputType.mixed raising ValueError,
• context_only with missing k/v,
• generation_only with missing latent_cache.

As per path instructions, "tests/**: Act as a QA engineer reviewing test changes and coverage ... suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/attention_backend/vanilla.py` around lines 620 - 633, Add
three focused test cases to the test suite for the MLA routing logic in the
attention backend to cover negative paths: first, verify that passing
AttentionInputType.mixed to the forward path raises ValueError as expected;
second, test that AttentionInputType.context_only fails appropriately when k or
v parameters are None (triggering the assertion); and third, test that
AttentionInputType.generation_only fails when forward_args.latent_cache is None
(triggering the assertion). These tests should be added alongside existing
happy-path MLA tests to ensure the assertion guards and error handling in the
conditional branches checking forward_args.attention_input_type are properly
validated.

Source: Path instructions

tests/unittest/_torch/attention/test_attention_backends.py (1)

187-189: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Make “golden-only” cases explicit to keep backend coverage honest.

run_case() can succeed with only "VANILLA" when all non-golden backends are unsupported on the current platform. That should be surfaced as an explicit skip (or assertion) so these don’t look like full backend validation passes.

Suggested patch

 `@pytest.mark.parametrize`("name", list(MODEL_CASES), ids=lambda n: n)
 def test_attention_backend_model(name):
-    run_case(MODEL_CASES[name])
+    results = run_case(MODEL_CASES[name])
+    if set(results) == {"VANILLA"}:
+        pytest.skip("No non-golden backend supports this case on this platform")

Coverage assessment:

• tests/unittest/_torch/attention/test_attention_backends.py: sufficient for supported backend/case combinations.
• Follow-up needed in either this file or tests/unittest/_torch/attention/attention_test_harness.py to mark insufficient backend-executed coverage when cases are golden-only.

As per path instructions, tests/** reviews should assess coverage sufficiency and provide concrete file-level follow-up guidance.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/attention/test_attention_backends.py` around lines 187
- 189, The test_attention_backend_model function does not explicitly surface
when run_case() only executes with the "VANILLA" golden backend due to all
non-golden backends being unsupported on the current platform, which masks
incomplete backend coverage. Modify the test or the run_case() function to
detect this golden-only scenario and make it explicit by either skipping the
test with a clear message or raising an assertion that indicates insufficient
backend coverage was detected, ensuring that partial backend validation is not
reported as a full pass.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/attention/attention_test_harness.py`:
- Around line 773-775: The current loop skips backends when unsupported_reason
returns non-None but does not validate that at least one comparison backend was
executed. Track which backends are skipped during the loop over
BACKENDS_UNDER_TEST, and after the loop completes in run_case, check if all
backends were filtered out (meaning only VANILLA golden backend was tested). If
no comparison backends ran, call pytest.skip with a message indicating why the
test is being skipped, ensuring that test results are not marked as passing when
only the golden backend executed without any backend validation.
- Around line 604-607: The MLA batch handling logic routes all non-context-only
cases to _run_mla_gen_backend, which hardcodes num_contexts=0, meaning mixed MLA
batches (cases where is_mla is true but is_context_only is false) are not
exercising their context requests during testing. Either implement proper mixed
MLA batch support that doesn't ignore context phases, or add a guard to reject
or skip such cases if they're not supported in the current test configuration.
Reference the is_mla and is_context_only flags in the conditional logic where
_run_mla_gen_backend is called to implement the appropriate handling.
- Around line 325-340: The _fill_mla_cache function writes only the cached
latent prefix to the MLA cache pool but leaves unwritten portions of blocks
uninitialized with stale data (NaNs), making tests flaky. Zero the entire buffer
returned by mgr.get_buffers(layer_idx) at the beginning of the function (before
the loop that iterates through blocks_per_req) to ensure all unwritten portions
contain zeros instead of stale values that could be read by decode kernels.
- Around line 161-164: The zip() call in the token_nums method does not validate
that both input lists have equal length, which could silently truncate results
if there's a mismatch between num_cached_tokens and kv_new_lens. Add the
strict=True parameter to the zip() function call to enable strict length
checking, which will raise a ValueError at runtime if the two lists have
different lengths, helping catch invariant violations.

In `@tests/unittest/_torch/attention/kv_cache_utils.py`:
- Around line 52-57: The broad except Exception clause in the
kv_cache_manager.get_buffers() call block is silently swallowing unexpected
failures like bad layouts, manager bugs, or allocation failures, which masks
real issues and leaves the cache uninitialized. Remove the except Exception:
return block entirely and keep only the TypeError except clause to handle the
old get_buffers(layer_idx) signature compatibility. If there is a documented
specific exception type for unsupported cache cases that should be handled,
catch only that specific exception type instead of the broad Exception base
class.
- Around line 115-117: The zip() call in the loop iterating over seq_lens and
num_cached_tokens is missing the strict=True parameter, which allows silent
truncation if the two sequences have different lengths and causes incorrect RoPE
application to packed tokens. Add strict=True to the zip() function call to
enforce that both seq_lens and num_cached_tokens have equal length, raising a
ValueError if they don't match.

In `@tests/unittest/_torch/attention/test_attention_backends.py`:
- Around line 152-154: The zip call in the seq_lens_kv construction is missing
the strict=True parameter, which allows silent truncation if the two lists being
zipped have different lengths. This can cause invalid test data to be generated
without raising an error. Add strict=True to the zip function call that combines
_PHASES["ctx"]["seq_lens"] with the [2, 3, 6] list to ensure both lists have
matching lengths and fail fast if they do not.

---

Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/vanilla.py`:
- Around line 620-633: Add three focused test cases to the test suite for the
MLA routing logic in the attention backend to cover negative paths: first,
verify that passing AttentionInputType.mixed to the forward path raises
ValueError as expected; second, test that AttentionInputType.context_only fails
appropriately when k or v parameters are None (triggering the assertion); and
third, test that AttentionInputType.generation_only fails when
forward_args.latent_cache is None (triggering the assertion). These tests should
be added alongside existing happy-path MLA tests to ensure the assertion guards
and error handling in the conditional branches checking
forward_args.attention_input_type are properly validated.

In `@tests/unittest/_torch/attention/capability_matrix.py`:
- Around line 58-67: The BACKEND_CAPS["VANILLA"] capability matrix has mla set
to False, but VanillaAttention.support_mla() now returns True, causing a
mismatch between the capability declaration and the actual backend contract.
Update the mla field in the VANILLA dictionary from False to True to align the
capability matrix with the current backend implementation.

In `@tests/unittest/_torch/attention/test_attention_backends.py`:
- Around line 187-189: The test_attention_backend_model function does not
explicitly surface when run_case() only executes with the "VANILLA" golden
backend due to all non-golden backends being unsupported on the current
platform, which masks incomplete backend coverage. Modify the test or the
run_case() function to detect this golden-only scenario and make it explicit by
either skipping the test with a clear message or raising an assertion that
indicates insufficient backend coverage was detected, ensuring that partial
backend validation is not reported as a full pass.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 58c8a2b0-b99e-4fdb-a8e5-d62e474ab7f9

📥 Commits

Reviewing files that changed from the base of the PR and between 31d4301 and 91f4828.

📒 Files selected for processing (7)

• tensorrt_llm/_torch/attention_backend/vanilla.py
• tests/unittest/_torch/attention/attention_test_harness.py
• tests/unittest/_torch/attention/capability_matrix.py
• tests/unittest/_torch/attention/kv_cache_utils.py
• tests/unittest/_torch/attention/model_configs.py
• tests/unittest/_torch/attention/test_attention_backends.py
• tests/unittest/_torch/attention/test_vanilla_attention.py

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55200 [ run ] triggered by Bot. Commit: 91f4828 Link to invocation

[tensorrt-cicd]

PR_Github #55200 [ run ] completed with state SUCCESS. Commit: 91f4828
/LLM/main/L0_MergeRequest_PR pipeline #44162 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55403 [ run ] triggered by Bot. Commit: 91f4828 Link to invocation

[tensorrt-cicd]

PR_Github #55403 [ run ] completed with state FAILURE. Commit: 91f4828
/LLM/main/L0_MergeRequest_PR pipeline #44349 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55490 [ run ] triggered by Bot. Commit: 87b9820 Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55496 [ run ] triggered by Bot. Commit: 3a43fbe Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55490 [ run ] completed with state ABORTED. Commit: 87b9820

Link to invocation

[tensorrt-cicd]

PR_Github #55500 [ run ] triggered by Bot. Commit: 54fdad2 Link to invocation

[tensorrt-cicd]

PR_Github #55496 [ run ] completed with state ABORTED. Commit: 3a43fbe

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55504 [ run ] triggered by Bot. Commit: b8040f2 Link to invocation

[tensorrt-cicd]

PR_Github #55500 [ run ] completed with state ABORTED. Commit: 54fdad2

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55508 [ run ] triggered by Bot. Commit: ac0b235 Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56353 [ run ] triggered by Bot. Commit: f784379 Link to invocation

[tensorrt-cicd]

PR_Github #56353 [ run ] completed with state SUCCESS. Commit: f784379
/LLM/main/L0_MergeRequest_PR pipeline #45202 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56518 [ run ] triggered by Bot. Commit: bd73fa6 Link to invocation

[tensorrt-cicd]

PR_Github #56518 [ run ] completed with state SUCCESS. Commit: bd73fa6
/LLM/main/L0_MergeRequest_PR pipeline #45355 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56606 [ run ] triggered by Bot. Commit: bd73fa6 Link to invocation

[tensorrt-cicd]

PR_Github #56606 [ run ] completed with state SUCCESS. Commit: bd73fa6
/LLM/main/L0_MergeRequest_PR pipeline #45434 completed with status: 'SUCCESS'

CI Report

Link to invocation