[None][fix] Fix fused mHC output reuse and extend compressor next_n#16221
Open
heyuhhh wants to merge 3 commits into
Open
[None][fix] Fix fused mHC output reuse and extend compressor next_n#16221heyuhhh wants to merge 3 commits into
heyuhhh wants to merge 3 commits into
Conversation
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>
heyuhhh
requested review from
QiJune,
longlee0622,
mingyangHao and
pengbowang-nv
Contributor
📝 WalkthroughWalkthroughChangesCompressed KV execution
Fused-HC scratch reuse
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CompressedKvMetadata
participant GenerationPositionIds
participant CuNewCompKv
CompressedKvMetadata->>GenerationPositionIds: pass cu_new_comp_kv
GenerationPositionIds->>CuNewCompKv: search compact output ownership
CuNewCompKv-->>GenerationPositionIds: return request indices and offsets
GenerationPositionIds-->>CompressedKvMetadata: return compressed position IDs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 (4)
tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py (2)
1542-1542: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd type annotations to the new test.
Annotate both integer parameters and the
Nonereturn type.As per coding guidelines, Python functions must always be annotated.
🤖 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/sparse/deepseek_v4/test_compressor_kernel.py` at line 1542, Add type annotations to test_decode_rejects_unsupported_next_n, declaring next_n and storage_next_n as integers and the function return type as None.Source: Coding guidelines
1410-1416: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover the HD128 single-warp configurations for
next_n > 4.The new 5–8 cases exercise HD512 only, while the dispatch matrix also adds CR=128/HD=128/NRW=1 configurations. Add at least one
basic_hd128_next8case to catch head-dimension-specific dispatch or indexing regressions.As per path instructions, test changes need actionable coverage follow-up when a changed branch remains unexercised.
🤖 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/sparse/deepseek_v4/test_compressor_kernel.py` around lines 1410 - 1416, Add a single-warp HD128 parameterized test case with next_n greater than 4, preferably `basic_hd128_next8`, to the relevant test parameter list alongside the existing `basic_hd512_next5`–`next8` cases. Ensure it uses the HD128 configuration and matches the expected argument shape used by the test.Source: Path instructions
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py (1)
929-938: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the compiled helper’s return type.
_compute_gen_compressed_position_idsmutates buffers and returns nothing; add-> None.As per coding guidelines, Python functions must always be annotated.
🤖 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/sparse/deepseek_v4/deepseek_v4.py` around lines 929 - 938, The _compute_gen_compressed_position_ids function lacks a return annotation; add -> None to its signature, since it mutates buffers and returns no value.Source: Coding guidelines
tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py (1)
365-367: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType-annotate the new helper and tests.
Add parameter annotations,
-> list[int]for_active_compressed_position_ids, and-> Nonefor both test functions.As per coding guidelines, Python functions must always be annotated.
Also applies to: 426-428, 439-440
🤖 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/sparse/deepseek_v4/test_compressor_module.py` around lines 365 - 367, The new helper and tests lack required type annotations. Add appropriate parameter annotations to _active_compressed_position_ids and declare its return type as list[int]; annotate both test functions with their parameter types and -> None return types.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`:
- Around line 929-938: The _compute_gen_compressed_position_ids function lacks a
return annotation; add -> None to its signature, since it mutates buffers and
returns no value.
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py`:
- Line 1542: Add type annotations to test_decode_rejects_unsupported_next_n,
declaring next_n and storage_next_n as integers and the function return type as
None.
- Around line 1410-1416: Add a single-warp HD128 parameterized test case with
next_n greater than 4, preferably `basic_hd128_next8`, to the relevant test
parameter list alongside the existing `basic_hd512_next5`–`next8` cases. Ensure
it uses the HD128 configuration and matches the expected argument shape used by
the test.
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py`:
- Around line 365-367: The new helper and tests lack required type annotations.
Add appropriate parameter annotations to _active_compressed_position_ids and
declare its return type as list[int]; annotate both test functions with their
parameter types and -> None return types.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0cabe84d-2556-49e3-8c1e-1836df366c34
📒 Files selected for processing (8)
cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cucpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cucpp/tensorrt_llm/thop/compressorOp.cpptensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytensorrt_llm/_torch/modules/mhc/mhc_cuda.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pytests/unittest/_torch/modules/test_mhc.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58629 [ run ] triggered by Bot. Commit: |
liji-nv
reviewed
Jul 10, 2026
| comb_mix_cur = torch.empty((B, n2), dtype=torch.float32, device=device) | ||
| layer_input_cur = torch.empty((B, hidden_size), dtype=torch.bfloat16, device=device) | ||
| if ws_ks == 1: | ||
| y_acc_ws = torch.empty((B, shape_n), dtype=torch.float32, device=device) |
Collaborator
There was a problem hiding this comment.
The remaining size is no longer large, maybe we just remove the cache?
longlee0622
added a commit
to longlee0622/TensorRT-LLM
that referenced
this pull request
Jul 10, 2026
Signed-off-by: Jonas Li <6110159+longlee0622@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
3 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
Bug Fixes
Tests
Description
This PR fixes two common correctness issues that affect speculative decoding workloads:
next_n <= 4tonext_n <= 8, including explicit THOP/launcher validation and the required CUDA template instantiations.cu_new_comp_kvorder used by the compressor kernel. This prevents multi-request batches from assigning RoPE positions from the wrong request when requests produce different numbers of compressed tokens.The branch contains common fixes only and does not modify DSpark-specific code.
Test Coverage
Regression coverage added or retained for:
next_n=5..8, includingnext_n=8and invalid-range rejection.Validation performed on the final main-based branch:
git diff --check origin/main..HEADpassed.PR Checklist