[None][refactor] Add derived properties for the thop.attention call site - #14279
Merged
Conversation
Collaborator
Author
|
/bot run |
Contributor
📝 WalkthroughWalkthroughThis PR refactors the TensorRT-LLM attention backend to use structured parameters and derive kernel-ready values from metadata. It introduces ChangesAttention Backend Refactor: Structured Parameters and Sync Testing
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/unittest/_torch/attention_backend/test_attention_op_sync.py (2)
305-305: 💤 Low valueMinor: File should end with a newline.
Most style guides and editors expect files to end with a trailing newline character.
🤖 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` at line 305, Add a trailing newline character at the end of the file tests/unittest/_torch/attention_backend/test_attention_op_sync.py so the file ends with a newline (ensure the final byte is '\n'); open the file, place the cursor at EOF and save to include the newline, then run linters/tests to confirm the style warning is resolved.
282-298: 💤 Low valueMaintenance note:
expected_localsrequires manual updates when_runchanges.This hardcoded list couples the test to
_run's internal structure. The comment on lines 295-297 already notes thatsink_token_lengthwill be removed with a future PR. Consider adding a brief inline comment at the top of this set documenting that updates are expected when_runlocals change, to help future maintainers.This is acceptable as-is since the sync test will fail fast if the lists diverge.
🤖 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 282 - 298, The hardcoded set expected_locals in the test_attention_op_sync.py couples the test to the internal locals of the _run function; update the test by adding a brief inline comment immediately above the expected_locals definition (the set named expected_locals) that notes this list must be manually updated whenever _run's locals change (and mention the future removal of sink_token_length), so future maintainers understand the coupling and know to keep the set in sync with _run.
🤖 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.
Nitpick comments:
In `@tests/unittest/_torch/attention_backend/test_attention_op_sync.py`:
- Line 305: Add a trailing newline character at the end of the file
tests/unittest/_torch/attention_backend/test_attention_op_sync.py so the file
ends with a newline (ensure the final byte is '\n'); open the file, place the
cursor at EOF and save to include the newline, then run linters/tests to confirm
the style warning is resolved.
- Around line 282-298: The hardcoded set expected_locals in the
test_attention_op_sync.py couples the test to the internal locals of the _run
function; update the test by adding a brief inline comment immediately above the
expected_locals definition (the set named expected_locals) that notes this list
must be manually updated whenever _run's locals change (and mention the future
removal of sink_token_length), so future maintainers understand the coupling and
know to keep the set in sync with _run.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3f9b377c-6744-4992-ba75-f13f66bce952
📒 Files selected for processing (3)
tensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/trtllm.pytests/unittest/_torch/attention_backend/test_attention_op_sync.py
Collaborator
|
PR_Github #49007 [ run ] triggered by Bot. Commit: |
PerkzZheng
approved these changes
May 19, 2026
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 19, 2026
Address review feedback on the thop.attention derived-properties refactor: - Move ``_THOP_EXCLUDED_FIELDS`` / ``_THOP_LITERAL_NONE`` to ``trtllm.py``; they describe the TRTLLM backend's call site, not the shared interface. - Drop ``effective_flash_mla_tile_scheduler_metadata`` and ``effective_flash_mla_num_splits``: the raw fields are already gated by ``enable_flash_mla`` and default to ``None``, so the "view" properties re-stated a guarantee the producer already enforces. - Move ``spec_decoding_tensor_params`` from a ``TrtllmAttention`` method to a ``TrtllmAttentionMetadata`` @Property; everything the list reads lives on metadata, including ``is_sm_version_trtllm_gen_kernel``. - Extend the static sync test to walk @Property bodies on ``AttentionForwardArgs`` so that a field is recognized as consumed iff it's read either directly at the call site or transitively via a @Property that is itself accessed there. This lets the excluded-field set shrink to the three fields the TRTLLM backend genuinely ignores (``topk_indices`` / ``is_generation`` / ``attention_mask_data``). Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #49083 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49007 [ run ] completed with state |
Collaborator
|
PR_Github #49083 [ run ] completed with state
|
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 19, 2026
Address the eight-item review feedback on PR NVIDIA#14279. After this change, every kwarg at the ``thop.attention`` call site sources from ``self`` / ``metadata`` / ``forward_args`` (incl. ``forward_args.sparse``) or is an allowlisted literal — no per-call locals remain. Schema changes (``AttentionForwardArgs``): - Drop ``out_scale_sf`` and ``effective_out_scale`` @Property; the caller in ``modules/attention.py`` now picks ``input_scale`` vs ``inv_input_scale`` based on ``output_sf`` and assigns ``out_scale`` directly. (Item 3) - Rename ``kv_scales_sf`` → ``kv_scale_quant_orig`` and ``kv_scales_sf_inv`` → ``kv_scale_orig_quant`` to mirror the C++ kwarg names. (Item 8) - Add ``is_fused_qkv: bool = False`` and ``update_kv_cache: bool = True``; ``TrtllmAttention.forward`` populates them from ``k`` and ``metadata.is_cross``. (Item 7) - Add ``sparse: AttentionSparseArgs = field(default_factory=...)`` so the sparse sub-bag lives inside ``forward_args`` instead of being threaded separately. (Item 1) ``TrtllmAttention`` cleanups: - Remove ``self.kv_scale_orig_quant`` / ``self.kv_scale_quant_orig`` / ``self.kv_cache_scaling_factor`` — they were stuck at ``torch.ones(1)``, a no-op. The four MLA preprocessing ops now pass ``None`` for the scales, which the C++ side accepts when KV cache quant is inactive. (Item 8) - Add ``self.local_layer_idx`` field, written at the top of ``_run``; rename the C++ ``thop.attention`` kwarg ``layer_idx`` → ``local_layer_idx`` to disambiguate from the module-level ``self.layer_idx``. (Item 6) - Mutate ``forward_args.attention_window_size`` in-place when it's ``None`` so the call site reads it as ``forward_args.<attr>``. (Item 5) - Drop ``sparse_args`` from the ``_run`` signature; access via ``forward_args.sparse.<field>`` instead. (Item 1) - Regroup the call site: ``spec_decoding_tensor_params`` joins the metadata block; ``sparse_mla_topk_lens`` / ``compressed_kv_cache_pool_ptr`` retain literal ``None`` with a comment marking them as DeepSeek V4 placeholders. (Items 2, 4) Static sync test: - Extend ``_classify_kwargs`` to walk arbitrary-depth attribute chains (``forward_args.sparse.X``) and resolve intermediate steps via ``__dataclass_fields__`` types. - Drop ``sparse_args`` from the source-class set; sparse fields are now validated through the ``forward_args.sparse`` chain. - Trim ``expected_locals`` to just ``q``/``k``/``v`` and the literal ``sink_token_length=0`` — all other per-call locals are gone. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #49185 [ run ] triggered by Bot. Commit: |
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 19, 2026
…ssert Follow-up to PR NVIDIA#14279 review: - Strip comments that explained where a value used to live or where it is set/read elsewhere; the surrounding code is self-evident now. Comments retained where they call out a current invariant or non-obvious branch (e.g. the SM90 paged-context override, the ``SkipSoftmax`` carve-out). - Move the ``is_fused_qkv`` shape assert into ``TrtllmAttention.forward`` next to where ``forward_args.is_fused_qkv`` is computed from ``k`` / ``metadata.is_cross`` — validating the caller belongs at the point where the value is derived, not in ``_run`` where it's already trusted. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #49194 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49185 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast --add-multi-gpu-test |
Collaborator
|
PR_Github #49221 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49194 [ run ] completed with state |
Collaborator
|
PR_Github #49221 [ run ] completed with state |
yihwang-nv
approved these changes
May 21, 2026
yuxianq
force-pushed
the
derived-attn-props
branch
from
May 21, 2026 17:11
f1b159b to
5c0153a
Compare
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 21, 2026
Address review feedback on the thop.attention derived-properties refactor: - Move ``_THOP_EXCLUDED_FIELDS`` / ``_THOP_LITERAL_NONE`` to ``trtllm.py``; they describe the TRTLLM backend's call site, not the shared interface. - Drop ``effective_flash_mla_tile_scheduler_metadata`` and ``effective_flash_mla_num_splits``: the raw fields are already gated by ``enable_flash_mla`` and default to ``None``, so the "view" properties re-stated a guarantee the producer already enforces. - Move ``spec_decoding_tensor_params`` from a ``TrtllmAttention`` method to a ``TrtllmAttentionMetadata`` @Property; everything the list reads lives on metadata, including ``is_sm_version_trtllm_gen_kernel``. - Extend the static sync test to walk @Property bodies on ``AttentionForwardArgs`` so that a field is recognized as consumed iff it's read either directly at the call site or transitively via a @Property that is itself accessed there. This lets the excluded-field set shrink to the three fields the TRTLLM backend genuinely ignores (``topk_indices`` / ``is_generation`` / ``attention_mask_data``). Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 21, 2026
Address the eight-item review feedback on PR NVIDIA#14279. After this change, every kwarg at the ``thop.attention`` call site sources from ``self`` / ``metadata`` / ``forward_args`` (incl. ``forward_args.sparse``) or is an allowlisted literal — no per-call locals remain. Schema changes (``AttentionForwardArgs``): - Drop ``out_scale_sf`` and ``effective_out_scale`` @Property; the caller in ``modules/attention.py`` now picks ``input_scale`` vs ``inv_input_scale`` based on ``output_sf`` and assigns ``out_scale`` directly. (Item 3) - Rename ``kv_scales_sf`` → ``kv_scale_quant_orig`` and ``kv_scales_sf_inv`` → ``kv_scale_orig_quant`` to mirror the C++ kwarg names. (Item 8) - Add ``is_fused_qkv: bool = False`` and ``update_kv_cache: bool = True``; ``TrtllmAttention.forward`` populates them from ``k`` and ``metadata.is_cross``. (Item 7) - Add ``sparse: AttentionSparseArgs = field(default_factory=...)`` so the sparse sub-bag lives inside ``forward_args`` instead of being threaded separately. (Item 1) ``TrtllmAttention`` cleanups: - Remove ``self.kv_scale_orig_quant`` / ``self.kv_scale_quant_orig`` / ``self.kv_cache_scaling_factor`` — they were stuck at ``torch.ones(1)``, a no-op. The four MLA preprocessing ops now pass ``None`` for the scales, which the C++ side accepts when KV cache quant is inactive. (Item 8) - Add ``self.local_layer_idx`` field, written at the top of ``_run``; rename the C++ ``thop.attention`` kwarg ``layer_idx`` → ``local_layer_idx`` to disambiguate from the module-level ``self.layer_idx``. (Item 6) - Mutate ``forward_args.attention_window_size`` in-place when it's ``None`` so the call site reads it as ``forward_args.<attr>``. (Item 5) - Drop ``sparse_args`` from the ``_run`` signature; access via ``forward_args.sparse.<field>`` instead. (Item 1) - Regroup the call site: ``spec_decoding_tensor_params`` joins the metadata block; ``sparse_mla_topk_lens`` / ``compressed_kv_cache_pool_ptr`` retain literal ``None`` with a comment marking them as DeepSeek V4 placeholders. (Items 2, 4) Static sync test: - Extend ``_classify_kwargs`` to walk arbitrary-depth attribute chains (``forward_args.sparse.X``) and resolve intermediate steps via ``__dataclass_fields__`` types. - Drop ``sparse_args`` from the source-class set; sparse fields are now validated through the ``forward_args.sparse`` chain. - Trim ``expected_locals`` to just ``q``/``k``/``v`` and the literal ``sink_token_length=0`` — all other per-call locals are gone. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 21, 2026
…ssert Follow-up to PR NVIDIA#14279 review: - Strip comments that explained where a value used to live or where it is set/read elsewhere; the surrounding code is self-evident now. Comments retained where they call out a current invariant or non-obvious branch (e.g. the SM90 paged-context override, the ``SkipSoftmax`` carve-out). - Move the ``is_fused_qkv`` shape assert into ``TrtllmAttention.forward`` next to where ``forward_args.is_fused_qkv`` is computed from ``k`` / ``metadata.is_cross`` — validating the caller belongs at the point where the value is derived, not in ``_run`` where it's already trusted. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
|
PR_Github #50427 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50353 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50453 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50427 [ run ] completed with state |
This was referenced May 27, 2026
After commit 574c518 renamed AttentionForwardArgs.kv_scales_sf and kv_scales_sf_inv to kv_scale_quant_orig and kv_scale_orig_quant, two references in FlashInferTrtllmGenAttention._get_kv_scale_params were left pointing at the old (now-nonexistent) field names. Because AttentionForwardArgs uses slots=True, hitting these references at runtime raises AttributeError. The path is gated by TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=1, so default-config tests do not exercise it; the failure surfaces only when the trtllm-gen kernel is enabled (e.g. test_attention_mla with the env override). Static-only fix: pure rename to the current field names. No GPU validation possible without rebuilding the .so on a Blackwell node, but the change cannot regress anything else — these are the only references to the old names in tensorrt_llm/. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
CI build #39880 surfaced NVFP4 accuracy regressions on DGX_B200: - TestLlama3_1_8B.test_nvfp4: 5.517 vs threshold 23.479 - TestLlama3_1_8B.test_nvfp4_streaming: same - TestQwen3_30B_A3B.test_nvfp4[latency_moe_*]: same root cause Root cause: refactor commit 574c518 consolidated main's out_scale/out_scale_sf pair into a single forward_args.out_scale field, picking at the modules/attention.py call site based on ``output_sf is not None``. But ``output_sf`` at that point is the ``_attn_impl`` function argument, which forward_impl() never populates (the buffer is created later inside the backend's forward by ``create_output``). So the predicate was always False and the NVFP4 path silently picked ``inv_input_scale`` instead of ``input_scale`` -> kernel-side rescale wrong -> ~5x accuracy drop. Main avoided this by passing both scales through to trtllm.py and deciding inside ``_run`` after ``output_sf`` is allocated. Fix: * interface.py: re-introduce ``out_scale_sf`` field on AttentionForwardArgs. * modules/attention.py: populate both ``out_scale`` (= ``inv_input_scale``) and ``out_scale_sf`` (= ``input_scale``) whenever quantized output is enabled, matching main. * trtllm.py: in ``_run``, after the fp8-KV fallback, promote ``out_scale_sf`` -> ``out_scale`` when ``forward_args.output_sf`` is allocated. The static sync test's ``_THOP_EXCLUDED_FIELDS`` gains ``out_scale_sf`` since the kernel binding still takes only a single ``out_scale``. Validation on B200 (umbriel-b200-081, sm_100): * TestLlama3_1_8B::test_nvfp4 PASSED (MMLU 62.622 vs threshold 61.343) * TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_4,64] PASSED * test_attention_op_sync.py (5 tests) PASSED Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
After the AttentionForwardArgs field rename in 574c518, the trtllm-gen backend's _get_kv_scale_params kept a two-tier lookup (per-call forward_args.kv_scale_* OR layer-level attention_layer.kv_scale_*). The layer-level fallback is dead for the PyTorch backend: * The PyTorch model loader never writes to TrtllmAttention. kv_cache_scaling_factor — it stays at the torch.ones(1) initial value, so attention_layer.kv_scale_quant_orig is always 1.0. * The attention-kernel side (mlaPreprocessOp / MMHA / unfused / sm_120 XQA) now defaults to 1.0 when the scale pointer is null (commits 319ec1c, 5040cc7), reproducing the exact same effective value. Sourcing only from forward_args also removes the implicit dependency on attention_layer's kv_scale_* attributes, allowing the layer-level mirrors to be dropped from TrtllmAttention.__init__ in a follow-up. Validation on B200 (umb-b200-236): * TestLlama3_1_8B::test_nvfp4 with TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=1 PASSED (exercises the simplified _get_kv_scale_params). Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
yuxianq
force-pushed
the
derived-attn-props
branch
from
May 27, 2026 06:48
2aea9fc to
5e5334f
Compare
Collaborator
|
PR_Github #50489 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50453 [ run ] completed with state |
Add ``test_excluded_fields_match_real_fields`` to the AST sync test. It walks ``AttentionForwardArgs`` (recursively into nested dataclass sub-bags) and asserts every entry in ``_THOP_EXCLUDED_FIELDS`` names a real field. A dead entry would silently allow a newly-added real field to slip past ``test_every_forward_args_field_is_consumed``. While at it, drop the dead ``is_generation`` entry — it was added in 574c518 but never matched a field on the renamed ``AttentionForwardArgs`` (the matching field was removed earlier in the same refactor iteration). The companion sanity check on ``_THOP_LITERALS`` already exists in ``test_literal_kwargs_match_allowlist`` (both directions: unknowns and stales), so this commit covers only the gap on the exclude set. Validation: ran the full sync test locally (no GPU); all 6 tests pass, including the new check. Temporarily re-adding ``is_generation`` makes the new check fail with a clear message, as expected. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast --add-multi-gpu-test |
Collaborator
|
PR_Github #50497 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50489 [ run ] completed with state |
L0 build #40005 surfaced 5 illegal-memory-access failures on
RTXPro6000D and 1 CUBLAS_STATUS_EXECUTION_FAILED on DGX_H100:
- TestLagunaXS.{test_fp8,test_nvfp4} (RTXPro6000D, sm_120)
- test_ptp_quickstart_advanced[{Llama3.1-8B-NVFP4,Qwen3-30B-A3B_{fp8,nvfp4}}_hf]
(RTXPro6000D, sm_120)
- TestNemotronV3Nano.test_fp8 (DGX_H100, sm_90)
Cross-PR scan: 0 other PRs hit these tests; PR NVIDIA#14279 alone since
commit 5e5334f (which dropped the older ecf3fa5 "kv_scale
fallback" commit during the rebase that simplified trtllm_gen.py).
Root cause: dropping ecf3fa5 removed two things that turn out to be
load-bearing for fp8 KV cache models on the default-path (non-trtllm-gen)
attention kernels:
1. Layer-level mirrors ``self.kv_cache_scaling_factor`` /
``self.kv_scale_quant_orig`` / ``self.kv_scale_orig_quant`` on
``TrtllmAttention.__init__``.
2. The ``_run``-side fallback that defaults
``forward_args.kv_scale_*`` to those mirrors when the caller
leaves them ``None``.
I earlier (correctly) noted the PyTorch backend never writes real
values to ``kv_cache_scaling_factor`` and concluded the mirrors were
1.0-no-ops the kernel could re-derive from a nullptr default. That
conclusion was incomplete: several XQA kernels still deref
``kvCacheScale[0]`` whenever ``isKVCacheQuantized`` is true and do
NOT check for nullptr —
* cpp/kernels/xqa/mha.cu (lines 1744, 2495) - sm_100/sm_120
* cpp/kernels/xqa/mha_sm90.cu (lines 885, 1127, 1530, 1620) - sm_90
The earlier null-safe deref fixes (commits 319ec1c, 5040cc7)
only covered the MMHA template, the unfused KV-cache write path, and
``mla_sm120.cu`` — the non-MLA XQA kernels were missed. Passing a
nullptr to them on fp8-KV models reads invalid memory at the kernel
entry → illegal memory access at executor init → cascades to
CUBLAS_STATUS_EXECUTION_FAILED on subsequent calls.
Fix: restore the layer-level mirrors and the ``_run`` fallback. This
is verbatim equivalent to the trtllm.py state at the previously-passing
build NVIDIA#2656 (commit 2aea9fc) modulo this commit message's comment
edits. ``modules/attention.py`` continues to populate the per-call
fields only for fp4 KV cache; fp8-KV models fall through to the
layer-level mirrors and the kernel receives a valid 1.0 pointer.
Local validation: my single-arch / multi-arch sm_120 builds cannot
reproduce the CI's complete FMHA cubin set (e4m3 + sm_120 + head_dim 128
+ ReturnSoftmaxStats variants are not generated by setup.py /
fmha_cubin.cpp in the local build) — TestLagunaXS::test_fp8 fails at
executor init with "FMHA kernels are not found" before reaching the
attention runtime path. This is a build-environment gap, not a code
regression. The fix is byte-equivalent (modulo comments) to commit
2aea9fc which DID pass this test on RTXPro6000D in CI build NVIDIA#2656.
A follow-up should add null-safe ``kvCacheScale[0]`` guards in
``mha.cu`` and ``mha_sm90.cu`` so the layer-level mirrors are no
longer load-bearing, mirroring the pattern already in
``mla_sm120.cu``. Tracked separately.
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50539 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50497 [ run ] completed with state |
Collaborator
|
PR_Github #50539 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast --add-multi-gpu-test |
Collaborator
|
PR_Github #50640 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50640 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast --add-multi-gpu-test |
Collaborator
|
PR_Github #50712 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50712 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Tests
Description
Convert the inline
thop.attention(...)call inTrtllmAttention._runfrom a ~75-arg positional call into an explicit-kwarg call where every kwarg is sourced assource.attributefrom one of the rich objects (self/metadata/forward_args/sparse_args). Section comments at the call site mirror the rich-object boundaries.This is a foundation PR — additive and behavior-preserving — extracted from a larger ongoing refactor of the attention dispatch path. It does not extract a
_call_thop_attentionwrapper method (that comes in a follow-up); the explicit-kwarg call remains inline in_runwith section-comment organization.Three pieces
1.
interface.py— properties + dataclass + constantsAttentionForwardArgs:mask_type@property— replaces the legacyTrtllmAttention._get_mask_type.effective_out_scale@property— replaces the inlineout_scale_sf if output_sf is not None else out_scaleternary.mrope_rotary_cos_sin/mrope_position_deltas@property— pull from themrope_configdict.enable_attn_nvfp4_outputfield (no readers).AttentionSparseArgsdataclass for the sparse-attention args bag (sparse_kv_indices,sparse_kv_offsets,sparse_attn_indices,sparse_attn_offsets,sparse_attn_indices_block_size)._THOP_EXCLUDED_FIELDSand_THOP_LITERAL_NONEfor the sync test.2.
trtllm.py— metadata field, ~13 @Property accessors, call-site rewriteTrtllmAttentionMetadata: adduse_paged_context_fmhafield (init fromruntime_featuresin__post_init__;forward()applies per-call SM/sparse/MLA overrides). Addnum_sparse_topk: int = 0baseline. Add 6 derived metadata@propertyaccessors (effective_workspace,helix_tensor_params,spec_decoding_bool_params,spec_decoding_position_offsets_for_cpp,effective_flash_mla_tile_scheduler_metadata,effective_flash_mla_num_splits,max_context_length).TrtllmAttention: add 5rotary_embedding_*@propertyaccessors backed byself.rope_params. Add 2skip_softmax_threshold_scale_factor_{prefill,decode}@propertyaccessors backed byself.sparse_attention_config. Add aspec_decoding_tensor_params(metadata)method (not@propertybecause the SM check lives on the backend)._get_mask_typehelper; removeoutput/output_sffrom_run's signature (now sourced fromforward_args.outputafterforward()writes it back); remove all inline pre-computations absorbed by the new properties._runsignature to takesparse_args: AttentionSparseArgsinstead of the 7 loose sparse params.thop.attention(...)call as explicit kwargs with 7 section comments (Inputs / Per-step batch state / AttentionForwardArgs / TrtllmAttention / Sparse-specific / Per-call locals / Literals).3.
test_attention_op_sync.py— AST-only sync test (new, GPU-free)Parses the AST of
TrtllmAttention._run, locates thethop.attention(...)call, and enforces five invariants:nb::arg(...)fromcpp/tensorrt_llm/nanobind/thop/bindings.cppbecauseinspect.signature(thop.attention)returns(*args, **kwargs)for nanobind-bound functions.source.attrkwarg resolves toattron exactly one of the four source classes.kwarg=Noneis only legal if the kwarg is in_THOP_LITERAL_NONE.AttentionForwardArgsfield is consumed (directly, via a@property, or in_THOP_EXCLUDED_FIELDS).AttentionSparseArgsfield is consumed.The test is the load-bearing piece — it makes the derived-properties refactor self-enforcing: if a property is added but never consumed, the test fails; if a kwarg drifts vs the C++ binding, the test fails; if a dataclass field is silently dropped, the test fails.
Test Coverage
Validated on B200 (
umbriel-b200-049, incremental build):tests/unittest/_torch/attention_backend/test_attention_op_sync.pytests/unittest/api_stabilitytests/unittest/_torch/attention/test_attention.pythop.attentionpath)tests/unittest/_torch/attention/test_attention_mla.pyTotal: 138 passed, 0 failed.
Rebase notes
drop-sink-token-lengthlands first: drop thesink_token_length=0kwarg from the call site and the"sink_token_length"entry fromexpected_localsin the sync test.clean-attn-fwd-argslands first: rebase is mechanical; anyAttentionForwardArgsfield renames cascade to the@propertybodies and_THOP_EXCLUDED_FIELDS.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.