[None][chore] Drop sink_token_length from PyTorch attention surface - #14275
Merged
Conversation
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (20)
💤 Files with no reviewable changes (9)
📝 WalkthroughWalkthroughThis PR removes the ChangesStreamingLLM sink_token_length Deprecation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
|
PR_Github #48962 [ run ] triggered by Bot. Commit: |
1 task
jieli-matrix
approved these changes
May 19, 2026
Collaborator
|
PR_Github #48962 [ run ] completed with state
|
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 19, 2026
…_kv_cache_for_mla caller PR NVIDIA#14275 dropped ``sink_token_length`` from the four MLA helper ops in commit b9dc10d, but the AutoDeploy chunked-prefill caller for ``load_chunked_kv_cache_for_mla`` (trtllm_mla.py:1271) was missed and still passed a literal ``0, # sink_token_length`` between ``attention_window_size`` and ``beam_width``. The schema now declares 18 parameters and the call passed 19, breaking ``test_trtllm_mla_out_buffer_padding_cached_kv[cuda-dtype0-4]`` on H100 CI with: RuntimeError: trtllm::load_chunked_kv_cache_for_mla() expected at most 18 argument(s) but received 19 argument(s). Drop the stale literal so the positional list matches the schema. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #49225 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49225 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49287 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49287 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
Author
|
/bot run --disable-fail-fast |
Continue the cleanup one layer deeper into the C++ thop helper.
`buildPagedKvCacheBuffers` used to take a `sink_token_length`
parameter and forward it into `buildKvCacheBuffers<KVBlockArray>`.
All thop callers either already pass 0 (MLA helpers, dsv3 RoPE) or
no longer have a meaningful non-zero value to pass (trtllm-gen QKV
preprocess — its own `sink_token_length` is still consumed by
`p.sinkTokenLength` for the kernel preprocessing struct, but the
KV cache buffer it constructs no longer cares).
Changes:
- cpp/tensorrt_llm/thop/attentionOp.h: drop `sink_token_length`
from `buildPagedKvCacheBuffers` declaration.
- cpp/tensorrt_llm/thop/attentionOp.cpp: drop from definition.
The internal `buildKvCacheBuffers<KVBlockArray>` call still
requires the param (deeper C++ plumbing untouched), so pass a
literal 0 there.
- 5 callers updated to stop forwarding the value:
mlaPreprocessOp.cpp (3 callsites in load_paged_kv_cache_for_mla,
load_chunked_kv_cache_for_mla, mla_rope_append_paged_kv_assign_q),
dsv3RopeOp.cpp (mla_rope_generation), and
trtllmGenQKVProcessOp.cpp.
Validation on B200 (umbriel-b200-013, incremental build):
- tests/unittest/_torch/attention/test_attention.py: 17 passed
- tests/unittest/_torch/attention/test_attention_mla.py: 80 passed
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
`KVCacheManager.get_max_atten_window_upper_bound` was the last
internal Python helper that took `sink_token_len`. Its sole caller
(`_validate_and_adjust_attention_windows`) has been passing 0
since the previous commits, and the function's `sink_bubble_len`
math collapses to a no-op when `sink_token_len == 0`:
sink_tokens_in_last_block = 0 % tokens_per_block # = 0
sink_bubble_len = 0
max_atten_window_upper_bound = max_token_num - 0 # = max_token_num
Drop the parameter and the now-dead bubble math. Sample evaluations
confirm identical results for the two interesting regimes:
- beam=1, 100 blocks * 32 tokens, max_seq=2048 -> 3200 (= 100*32)
- beam=2, same, max_seq>upper -> 1568 (= 50*32 - 32)
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Continue the cleanup into the trtllm-gen path. Every Python caller
in this pipeline already passed 0 — the `is_supported()` rejection
branch was provably dead since nothing produced a non-zero value.
Changes:
- cpp/tensorrt_llm/kernels/gptKernels.h: default-init
`BuildDecoderInfoParams::sinkTokenLength` to 0 so callers that
previously had to assign it explicitly can omit the write.
`QKVPreprocessingParams::sink_token_len` already defaults to {0}.
- cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp:
- `build_decoder_info`: drop `sink_token_length` param +
schema entry + the explicit `p.sinkTokenLength = ...` write.
- `qkv_preprocessing` / `kv_cache_postprocessing`: drop both
`sink_token_len` (kernel-prep) and `sink_token_length`
(cyclic-KV) params + schema entries + the
`p.sink_token_len = ...` write.
- tensorrt_llm/_torch/attention_backend/trtllm_gen.py: drop the
`sink_token_length` field from `TrtllmGenAttentionParams`, the
dead rejection branch in `is_supported`, the docstring lines,
and all 6 sites that forwarded `params.sink_token_length` (or
literal 0) to the C++ ops.
- tensorrt_llm/_torch/attention_backend/trtllm.py: drop the
positional `0` (sink_token_length) from the `trtllm_gen_attention(...)`
call now that the wrapper no longer takes the parameter.
- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py: drop
`sink_token_length` from the `build_decoder_info` register_fake
stub to match the updated schema.
Validation on B200 (umbriel-b200-013):
- tests/unittest/_torch/attention/test_attention.py: 17 passed
- tests/unittest/_torch/attention/test_attention_mla.py: 80 passed
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
…_kv_cache_for_mla caller PR NVIDIA#14275 dropped ``sink_token_length`` from the four MLA helper ops in commit b9dc10d, but the AutoDeploy chunked-prefill caller for ``load_chunked_kv_cache_for_mla`` (trtllm_mla.py:1271) was missed and still passed a literal ``0, # sink_token_length`` between ``attention_window_size`` and ``beam_width``. The schema now declares 18 parameters and the call passed 19, breaking ``test_trtllm_mla_out_buffer_padding_cached_kv[cuda-dtype0-4]`` on H100 CI with: RuntimeError: trtllm::load_chunked_kv_cache_for_mla() expected at most 18 argument(s) but received 19 argument(s). Drop the stale literal so the positional list matches the schema. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
…op.attention call The ``thop.attention()`` call inside ``_handle_prefill_thop_cached_kv``'s chunked-prefill loop (trtllm_mla.py:1304) was missing the ``sparse_mla_topk_lens`` positional argument between ``sparse_attn_indices_block_size`` (schema pos 68) and ``skip_softmax_threshold_scale_factor_prefill`` (schema pos 71). The placeholder ``0, # sparse_mla_topk`` was intended to be ``num_sparse_topk`` (schema pos 69) but its label was abbreviated and the following ``sparse_mla_topk_lens`` (Tensor? = None, schema pos 70) was omitted entirely. Without it, every subsequent positional arg was off-by-one and the binding-arity validator rejected the call on H100: TypeError: attention(): incompatible function arguments. ... Invoked with types: ... int, int, NoneType, NoneType, ... (one Tensor?/None short of the 88-arg schema). Add the missing ``None, # sparse_mla_topk_lens`` line and rename the sibling comment to ``num_sparse_topk`` for consistency with the three other ``thop.attention`` call sites in this file. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
…tion in cached_kv path A second stale ``0, # sink_token_length`` placeholder remained in the final-causal ``thop.attention()`` call inside ``_handle_prefill_thop_cached_kv`` (trtllm_mla.py:1422), between ``attention_window_size`` and ``beam_width``. The prior ``ff551611cc`` only added the missing ``sparse_mla_topk_lens`` arg to the chunked-loop call; the final-causal call still had one extra positional ``0``, so it still passed 82 positional + 2 kwargs to a schema expecting 88. Drop the leftover line and add the missing position-name comments to match the three other ``thop.attention`` call sites in this file. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
…cy D200 This branch modifies ``tensorrt_llm/llmapi/llm_args.py`` (deprecating ``sink_token_length``), which triggers the baseline-gated ``ruff-legacy`` pre-commit on that file. The file already had 24 D200 violations on ``main`` while the baseline allows 23, so the lint reported one regression (line numbers shifted but the count didn't change). Collapse the existing multi-line ``CudaGraphConfig`` docstring (D200: a one-line docstring should fit on one line) to bring the count back to 23 and unblock CI without modifying any of the other 23 pre-existing violations. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
yuxianq
force-pushed
the
drop-sink-token-length
branch
from
May 22, 2026 07:38
234e1b0 to
b4face3
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49890 [ run ] triggered by Bot. Commit: |
Post-rebase clang-format cleanup on the call sites where I removed ``sink_token_length`` from ``buildPagedKvCacheBuffers(...)`` — the argument-list reflow left some long lines that exceeded the project's column limit. Apply the formatter to bring them in line with the project style. No functional change; matches what ``pre-commit run clang-format`` produces locally. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49897 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49890 [ run ] completed with state |
Collaborator
|
PR_Github #49897 [ run ] completed with state
|
…and header Build #39475 failed Build-x86_64 NVIDIA#2806 and Build-SBSA NVIDIA#2786 with a linker error: undefined references to ``trtllmGenContextPreprocess``, ``trtllmGenContextPostprocess``, and ``trtllmGenGenerationPreprocess`` — all three with one extra trailing ``long`` parameter compared to the defined symbols in ``libth_common.so``. Root cause: ``trtllmGenQKVProcessOp.cpp`` (the definitions) had ``sink_token_length`` removed in a prior commit, but two surfaces still declared / passed it: 1. ``cpp/tensorrt_llm/thop/trtllmGenFusedOps.h`` (declarations consumed by ``bindings.cpp``) — still listed ``int64_t sink_token_length`` in all three signatures, so the compiler generated calls with the wrong arity that the linker could not resolve. 2. ``cpp/tensorrt_llm/nanobind/thop/bindings.cpp`` — the two lambda wrappers (``trtllmGenContextPreprocessBinding`` and ``trtllmGenGenerationPreprocessBinding``) still took ``sink_token_length`` as a parameter and forwarded it to the underlying call. The three ``m.def(...)`` registrations also still listed ``nb::arg("sink_token_length")``. Fix: drop ``sink_token_length`` from all six remaining call-sites — three in the header and three in the nanobind bindings. The Python callers in ``trtllm_gen.py`` already don't pass it. Verified: pre-commit (clang-format / ruff / etc.) all pass; no residual ``sink_token_length`` reference in the trtllm-gen surface. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49914 [ run ] triggered by Bot. Commit: |
Superjomn
approved these changes
May 22, 2026
Superjomn
left a comment
Superjomn
left a comment
Collaborator
There was a problem hiding this comment.
LGTM on the llmapi changes.
Collaborator
|
PR_Github #49914 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49969 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49969 [ run ] completed with state |
MrGeva
approved these changes
May 24, 2026
yizhang-nv
approved these changes
May 25, 2026
yuxianq
added a commit
to yuxianq/TensorRT-LLM
that referenced
this pull request
May 25, 2026
Step 2 of the squash + 2-step rebase workflow. Step 1 used ``git rebase -X ours`` so the tree silently took main's lines wherever it could; this commit re-applies the parts of the derived-attn-props refactor that step 1 dropped on the floor. Two displaced regions: 1. ``tensorrt_llm/_torch/attention_backend/trtllm.py`` — the ``thop.attention(...)`` call inside ``TrtllmAttention._run`` was left as a syntactically-valid hybrid: main's positional args followed by my kwarg form, with the positional args referencing locals (output, workspace, layer_idx, ...) that the refactor moved into ``forward_args.X`` properties. Restore the kwarg-only form, drop the dead prelude that built the now-unused locals, and remove ``sink_token_length`` from both the call site and ``_THOP_LITERALS`` (main dropped it from the C++ signature in NVIDIA#14275). 2. ``cpp/tensorrt_llm/nanobind/thop/bindings.cpp`` — the ``nb::arg("layer_idx")`` for ``thop.attention`` reverted to main's spelling even though the C++ function signature in attentionOp.h/.cpp already renamed the parameter to ``local_layer_idx``. Restore the ``nb::arg("local_layer_idx")`` so the binding name matches the underlying parameter and the Python kwarg ``local_layer_idx=self.local_layer_idx`` resolves. Verification: - pre-commit clean on all 10 touched files (autoflake + clang-format ran) - ``test_attention_op_sync.py``: 5/5 pass (call-site kwargs match the ``nb::arg`` list and ``_THOP_LITERALS`` allowlist) - All 6 touched Python modules import cleanly Reviewability: ``git diff rebase-step1..HEAD`` shows every line of manual judgement in this rebase, isolated from main's churn. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
KleinBlueC
pushed a commit
to KleinBlueC/TensorRT-LLM
that referenced
this pull request
May 26, 2026
…VIDIA#14275) Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
bmarimuthu-nv
pushed a commit
to nv-auto-deploy/TensorRT-LLM
that referenced
this pull request
May 28, 2026
…VIDIA#14275) Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
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
9 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
Deprecated Features
sink_token_lengthparameter in KvCacheConfig is now deprecated and ignored on the PyTorch backend.Removed Features
Documentation
Tests
Description
StreamingLLM has not been functional on the PyTorch backend since the non-cyclic KV-cache + paged-context FMHA migration (
commands/build.py:317-318already asserted this for the legacy TensorRT path). On the PyTorch path,sink_token_lengthwas silently ignored by the attention kernel while the KV-cache manager still reserved sink blocks — a footgun that produced silent miscomputation if a user ever setKvCacheConfig.sink_token_length=N.This PR cleans up
sink_token_lengthfrom the entire PyTorch attention surface, leaving the deeper C++ plumbing (AttentionOp::EnqueueParams, XQA, MMHA,KVBlockArray, executor_KvCacheConfig, decoderDecoderState) intact. 5 atomic commits, peelable layer by layer:b9dc10d43f— Drop fromthop.attentionC++ signature + nanobind binding +RunnerBase; drop from 4 MLA helper ops (load_paged_kv_cache_for_mla,load_chunked_kv_cache_for_mla,mla_rope_append_paged_kv_assign_q,mla_rope_generation) — schema strings + signatures + bodies +register_fake; clean Python callers intrtllm.py,sparse/dsa.py, AutoDeploy custom ops,resource_manager.py, andllm_args.py; delete two already-skipped StreamingLLM tests; updatekvcache.md+attention.md.0e20321523— Drop redundant explicit= 0;writes (EnqueueParams::sink_token_lengthalready defaults to 0 incommon/attentionOp.h); trim awkward implementation-detail prose from the deprecation notes.fbe45530d2— Dropsink_token_lengthfrombuildPagedKvCacheBuffersdeclaration + definition + 5 callers (mlaPreprocessOp x3, dsv3RopeOp x1, trtllmGenQKVProcessOp x1); hardcode 0 inside the body where it hands off to the deeperbuildKvCacheBuffers<KVBlockArray>template.2cb6cd7b40— Dropsink_token_lenfromKVCacheManager.get_max_atten_window_upper_boundand its single internal caller. Math collapses cleanly when sink_token_len is always 0; sample values verified (3200, 1568) identical to prior behavior.f5a72127a9— Drop sink_token_length from the trtllm-gen QKV pipeline: default-initBuildDecoderInfoParams::sinkTokenLength{0}ingptKernels.h; drop the param + schema + writes fromtrtllmGenQKVProcessOp.cpp(build_decoder_info, qkv_preprocessing, kv_cache_postprocessing); remove the dataclass field, deadis_supportedrejection, and 6 forwarding sites intrtllm_gen.py; drop the matching positional0intrtllm.pyandregister_fakestub field incpp_custom_ops.py.KvCacheConfig.sink_token_lengthis retained as a deprecated Pydantic shim (required bymirror_pybind_fields(_KvCacheConfig)since the executor C++ struct still owns the field) but no longer passed through_to_pybind— any user-supplied non-None value is silently dropped, closing the prior footgun. ADeprecationWarningfires when the field is accessed.Per
docs/source/features/kvcache.md, this field was already pre-announced as "deprecated and will be removed in a future release" — this PR cashes in that deprecation on the PyTorch path.Test Coverage
Validated on B200 (
umbriel-b200-013) after each rebuild:tests/unittest/_torch/attention/test_attention.pythop.attentionpath)tests/unittest/_torch/attention/test_attention_mla.pytests/unittest/api_stabilitytests/unittest/llmapi/test_llm_args.py::test_KvCacheConfig_declarationLocal Pydantic sanity:
KvCacheConfig(sink_token_length=4)syntactically accepted but emitsDeprecationWarningcfg._to_pybind().sink_token_length is None— silently dropped before reaching C++ executor_validate_and_adjust_attention_windowsandget_max_atten_window_upper_boundsignatures no longer takesink_token_len(gth)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.