[None][feat] Default on FlashInferTrtllmGenAttention

[yihwang-nv]

Description

Built on top of #14279 ([None][refactor] Add derived properties for the thop.attention call site). That PR refactors the thop.attention(...) call site to use explicit kwargs sourced from rich objects; this PR stacks the following on top of it:

Default on FlashInferTrtllmGenAttention

Switches TrtllmGenAttention from an opt-in (env-gated) to the default decode backend on Blackwell+, replacing the previous default with FlashInferTrtllmGenAttention.

[yihwang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #50471 [ run ] triggered by Bot. Commit: 48a828e Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR hardens attention kernel and backend implementations against null KV cache scale pointers, reorganizes MLA parameter passing, adds batch-aware workspace sizing, and refactors the TRT-LLM-Gen backend with a unified parameter container and consolidated support checking.

Changes

Attention system null-safety and backend reorganization

Layer / File(s)	Summary
Kernel-side null-safe KV cache scales cpp/kernels/xqa/mla_sm120.cu, cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h, cpp/tensorrt_llm/kernels/flashMLA/flash_fwd_mla_kernel.h, cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h	CUDA kernels now safely check for null KV cache scale pointers in both producer and consumer paths, defaulting to 1.0f when pointers are absent rather than dereferencing unconditionally.
Layer index parameter rename and binding updates cpp/tensorrt_llm/thop/attentionOp.h, cpp/tensorrt_llm/thop/attentionOp.cpp, cpp/tensorrt_llm/nanobind/thop/bindings.cpp	The attention() C++ function parameter layer_idx is renamed to local_layer_idx throughout the signature, implementation, logging, and nanobind Python binding registration.
Optional helper and MLA KV scale parameter cleanup cpp/tensorrt_llm/nanobind/thop/bindings.cpp, cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp	Introduces generic optionalToObject<T> template for nanobind std::optional conversion; removes redundant KV scale parameters from loadPagedKVCacheForMLA, loadChunkedKVCacheForMLA, and MLARopeAppendPagedKVAssignQ, keeping only contextually appropriate scale tensors and updating Torch operator schemas.
Batch-size-aware FlashInfer counter workspace computation cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp	Replaces hardcoded 8 MiB counter workspace constant with computeFlashinferTrtllmGenCounterWorkspaceBytes(batchSize) helper; updates zeroFlashinferTrtllmGenCounterWorkspaceAsync() to accept computed workspace size and clears only the minimum needed for active batch in generation preprocessing.
Attention interface contract: sparse args and scale parameters tensorrt_llm/_torch/attention_backend/interface.py	Introduces AttentionSparseArgs dataclass for sparse-attention indices/offsets, adds kv_scale_orig_quant/kv_scale_quant_orig fields to AttentionForwardArgs, removes legacy scale-factor fields, adds control flags and mask_type property mapping masks to C++ integer types.
TrtllmAttentionMetadata properties and initialization tensorrt_llm/_torch/attention_backend/trtllm.py	Extends metadata with use_paged_context_fmha flag, num_sparse_topk field, and new properties for workspace selection, helix parameters, spec-decoding configuration, and RoPE/skip-softmax thresholds; adds cached local_layer_idx handling and RoPE configuration properties to TrtllmAttention.
TrtllmAttention kernel dispatch and execution refactoring tensorrt_llm/_torch/attention_backend/trtllm.py	Refactors _run() to select between TRT-LLM-Gen and regular thop.attention via backend support check, invokes thop.attention with keyword arguments from new metadata accessors; updates forward() to derive use_paged_context_fmha from SM/algorithm constraints, populates forward_args.sparse with AttentionSparseArgs, and consolidates MLA calls to pass None for unused KV scale parameters.
TRT-LLM-Gen backend major refactoring with FmhaParams tensorrt_llm/_torch/attention_backend/trtllm_gen.py	Replaces EnqueueParams with unified FmhaParams dataclass, removes external TrtllmGenSupportChecker, integrates support validation into FlashInferTrtllmGenAttention.__init__(), converts attention() method to forward() signature, reorganizes is_supported() with sparse gating and updated KV-cache checks, refactors context/generation paths to populate FmhaParams and invoke dispatch methods with new parameter layouts for C++ and FlashInfer.
MLA operator calls and Attention module quantization refactoring tensorrt_llm/_torch/attention_backend/sparse/dsa.py, tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py, tensorrt_llm/_torch/modules/attention.py	Updates MLA operator invocations to pass None for unused KV scale slots; refactors Attention._attn_impl to replace legacy scale-factor fields with kv_scale_orig_quant/kv_scale_quant_orig derived from qkv_proj when FP4 KV-cache is enabled.
Sync verification test for Python-C++ THOP binding consistency tests/unittest/_torch/attention_backend/test_attention_op_sync.py	Adds pytest module that parses attentionOp.h C++ header and TrtllmAttention._run() Python AST to enforce strict sync on thop.attention() call-site keyword arguments, including parameter names, type categories, literal values, and recursive consumption of all AttentionForwardArgs dataclass fields.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

• wenmingw
• yuxianq
• liji-nv
• yizhang-nv
• Shixiaowei02

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 44.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title is somewhat generic and lacks specificity about the main change being the default enablement of FlashInferTrtllmGenAttention, which is buried in the word 'Default'.	Consider a more descriptive title such as '[None][feat] Default on FlashInferTrtllmGenAttention for Blackwell+ decode' to better communicate the scope and impact of the change.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description check	✅ Passed	The PR description provides a clear summary of all changes, rationale, and test coverage, but lacks explicit sections matching the template structure (Summary, Description, Test Coverage, Checklist).
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.

_{✏️ 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 and usage tips.}

[coderabbitai]

Actionable comments posted: 11

🧹 Nitpick comments (1)

tests/unittest/_torch/attention_backend/test_attention_op_sync.py (1)

496-519: QA list updates look unnecessary here.

This is a narrow unittest-only static sync check under tests/unittest/, so I don't think it needs a tests/integration/test_lists/qa/ update.

🤖 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_backend/test_attention_op_sync.py` around
lines 496 - 519, The PR included unrelated updates to QA lists under
tests/integration/test_lists/qa/ that are unnecessary for this unittest-only
change; revert or remove those QA-list changes so the only changes are the unit
test edits in tests/unittest/_torch/attention_backend/test_attention_op_sync.py
(refer to the functions test_every_forward_args_field_is_consumed and
test_no_unexpected_other_kwargs), then run the test suite or git diff to confirm
no files under tests/integration/test_lists/qa/ are modified before committing.

🤖 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 `@cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp`:
- Around line 216-220: The chunked MLA path needs the FP8-only quantization
guard restored: in loadChunkedKVCacheForMLA() ensure you only take the
FP8-specific branch (the code that uses kv_scale_quant_orig_ptr and the FP8 read
path) when hasFp8KvCache() is true (or quant_mode indicates FP8); otherwise
route to the non-FP8/unquantized TCache instantiations so element types match.
Update the dispatch that currently only special-cases hasFp8KvCache() to
explicitly check quant_mode/hasFp8KvCache() before using kv_scale_quant_orig and
the FP8 reader, preventing non-FP8 quantized caches from being read with the
wrong element type.

In `@tensorrt_llm/_torch/attention_backend/interface.py`:
- Around line 756-765: The mask_type property currently raises for any
non-CAUSAL/FULL mask which breaks the custom-mask path; update mask_type to
detect CustomAttentionMask (e.g., isinstance(self.attention_mask,
CustomAttentionMask)) and return the corresponding integer
(int(AttentionMaskType.custom) or the enum value used for custom masks) while
keeping the existing returns for PredefinedAttentionMask.CAUSAL and .FULL, so
TrtllmAttention._run() and AttentionForwardArgs.attention_mask_data can use the
custom-mask path instead of failing early.

In `@tensorrt_llm/_torch/attention_backend/trtllm_gen.py`:
- Around line 506-510: The positive support result is being returned from the
cache before phase-specific validation, allowing a generation-only probe to
wrongly bypass later context checks; update the probe caching in the method that
uses self._support_result so that you do not short-circuit on a cached True
before running the [Context]/[Generation] validations — either move the
write/read of self._support_result to occur after all phase-specific checks, or
make the cache keyed by phase (e.g., store separate tokens like
self._support_result_context and self._support_result_generation or a dict keyed
by phase) so that a True for one phase does not satisfy a different phase;
ensure the code paths that call the context-specific checks reference the new
phase-aware cache and only return cached True when the stored entry matches the
current probe phase.
- Around line 743-750: Restore layer-level KV-scale fallback: the FmhaParams
construction currently copies kv_scale_orig_quant and kv_scale_quant_orig
directly from fwd (which can be None), so change the assignment in the
FmhaParams(...) call in trtllm_gen.py to use the layer/model defaults when the
forward-level values are None (e.g., set kv_scale_orig_quant =
fwd.kv_scale_orig_quant if not None else self._kv_scale_orig_quant and
kv_scale_quant_orig = fwd.kv_scale_quant_orig if not None else
self._kv_scale_quant_orig) so the C++ preprocess/postprocess path receives the
proper fallback.
- Around line 887-930: The call to thop.trtllm_gen_context_preprocess (and the
similar calls to trtllm_gen_context_postprocess and
trtllm_gen_generation_preprocess) incorrectly passes params.sink_token_length as
a positional argument; remove params.sink_token_length from those positional
argument lists so subsequent arguments (e.g., params.num_tokens,
params.batch_size, params.input_seq_length, params.max_past_kv_length, etc.)
align with the updated C++ signatures. Locate the calls by the function names
trtllm_gen_context_preprocess, trtllm_gen_context_postprocess, and
trtllm_gen_generation_preprocess and delete the params.sink_token_length token,
keeping the original order of the remaining parameters so they bind to the
correct runtime parameters.

In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 177-184: The method spec_decoding_position_offsets_for_cpp
currently reshapes spec_decoding_position_offsets directly and ignores the
precomputed view stored in spec_decoding_position_offsets_cpp; change it to
return self.spec_decoding_position_offsets_cpp (which update_spec_dec_param
already populates with the active cpp_query_len and correct 2D layout) when that
attribute is set, falling back to the existing logic only if it's None—this
prevents using _internal_buf_dim/max_num_requests reshaping that exposes
padded/stale offsets to the C++ kernel.
- Around line 1453-1462: The fallback for forward_args.kv_scale_quant_orig
should be derived as the reciprocal of whichever kv_scale_orig_quant value is
actually being used (which may have just been assigned to
forward_args.kv_scale_orig_quant), rather than reusing a possibly-stale
self.kv_scale_quant_orig; change the second branch so that if
forward_args.kv_scale_quant_orig is None you compute it from
forward_args.kv_scale_orig_quant (e.g. torch.reciprocal or 1.0 / tensor) after
the first fallback assignment, ensuring forward_args.kv_scale_quant_orig stays
consistent with the live kv_scale value (referencing
forward_args.kv_scale_orig_quant, forward_args.kv_scale_quant_orig,
self.kv_scale_orig_quant, and self.kv_scale_quant_orig).

In `@tensorrt_llm/_torch/modules/attention.py`:
- Around line 744-754: The compiled LoRA path still ends up allocating FP8/NVFP4
outputs because attn_custom_op_inplace() calls _attn_impl(..., has_lora=False)
and create_attn_outputs() isn't aware of LoRA; fix by threading the has_lora
flag through attn_custom_op_inplace() and create_attn_outputs() so both
allocation and the custom op respect LoRA: update forward_impl to pass the real
has_lora into attn_custom_op_inplace(), modify attn_custom_op_inplace() to
accept and forward has_lora into _attn_impl(...) and create_attn_outputs(), and
change create_attn_outputs() to skip creating quantized buffers (force BF16 /
set out_scale=None) when has_lora is True.

In `@tests/unittest/_torch/attention_backend/test_attention_op_sync.py`:
- Around line 61-63: The test currently only verifies keyword names against the
C++ declaration (_HEADER / attentionOp.h); update it to also parse and assert
the nanobind binding names used for thop.attention in
cpp/tensorrt_llm/nanobind/thop/bindings.cpp so the runtime kwarg contract
matches the header; specifically, read the binding parameter names (the "_a"
identifiers) from bindings.cpp and compare them to the header-derived list and
to the Python-callable names used by thop.attention (also add the same check for
the other block mentioned around lines 132-179) so any drift in bindings.cpp
will fail the test.
- Around line 446-460: Adjust _collect_chains so it only counts attribute reads
that actually flow into the thop.attention(...) call: locate Call AST nodes
where func is an Attribute with attr "attention" and base Name "thop" (i.e., the
thop.attention call sites inside TrtllmAttention._run), then when walking
Attribute nodes only add a chain if that Attribute node is a descendant of one
of those Call nodes (appears in the call's args or keywords) and its root Name
matches the provided root; update _collect_chains to first collect those
thop.attention Call nodes and filter Attribute occurrences by ancestry to those
calls so preprocessing/branch-only reads are not treated as consumed.

---

Nitpick comments:
In `@tests/unittest/_torch/attention_backend/test_attention_op_sync.py`:
- Around line 496-519: The PR included unrelated updates to QA lists under
tests/integration/test_lists/qa/ that are unnecessary for this unittest-only
change; revert or remove those QA-list changes so the only changes are the unit
test edits in tests/unittest/_torch/attention_backend/test_attention_op_sync.py
(refer to the functions test_every_forward_args_field_is_consumed and
test_no_unexpected_other_kwargs), then run the test suite or git diff to confirm
no files under tests/integration/test_lists/qa/ are modified before committing.

🪄 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: 3371b37f-9f40-45f8-9c6d-bb963ed18227

📥 Commits

Reviewing files that changed from the base of the PR and between da0359b and 48a828e.

📒 Files selected for processing (16)

• cpp/kernels/xqa/mla_sm120.cu
• cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h
• cpp/tensorrt_llm/kernels/flashMLA/flash_fwd_mla_kernel.h
• cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h
• cpp/tensorrt_llm/nanobind/thop/bindings.cpp
• cpp/tensorrt_llm/thop/attentionOp.cpp
• cpp/tensorrt_llm/thop/attentionOp.h
• cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp
• cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp
• tensorrt_llm/_torch/attention_backend/interface.py
• tensorrt_llm/_torch/attention_backend/sparse/dsa.py
• tensorrt_llm/_torch/attention_backend/trtllm.py
• tensorrt_llm/_torch/attention_backend/trtllm_gen.py
• tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py
• tensorrt_llm/_torch/modules/attention.py
• tests/unittest/_torch/attention_backend/test_attention_op_sync.py

💤 Files with no reviewable changes (1)

• tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py

[tensorrt-cicd]

PR_Github #50471 [ run ] completed with state SUCCESS. Commit: 48a828e
/LLM/main/L0_MergeRequest_PR pipeline #39982 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

[yihwang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #50826 [ run ] triggered by Bot. Commit: 04aba35 Link to invocation

[tensorrt-cicd]

PR_Github #50826 [ run ] completed with state FAILURE. Commit: 04aba35
/LLM/main/L0_MergeRequest_PR pipeline #40296 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

[yihwang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52720 [ run ] triggered by Bot. Commit: 48b9461 Link to invocation

[tensorrt-cicd]

PR_Github #52636 [ run ] completed with state ABORTED. Commit: c99928c

Link to invocation

[yihwang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52740 [ run ] triggered by Bot. Commit: 48b9461 Link to invocation

[tensorrt-cicd]

PR_Github #52720 [ run ] completed with state ABORTED. Commit: 48b9461

Link to invocation

[yihwang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52780 [ run ] triggered by Bot. Commit: 48b9461 Link to invocation

[tensorrt-cicd]

PR_Github #52740 [ run ] completed with state ABORTED. Commit: 48b9461

Link to invocation

[yihwang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52810 [ run ] triggered by Bot. Commit: 48b9461 Link to invocation

[tensorrt-cicd]

PR_Github #52780 [ run ] completed with state ABORTED. Commit: 48b9461

Link to invocation

[tensorrt-cicd]

PR_Github #52810 [ run ] completed with state SUCCESS. Commit: 48b9461
/LLM/main/L0_MergeRequest_PR pipeline #42062 completed with status: 'SUCCESS'

CI Report

Link to invocation