[https://nvbugs/6695563][fix] Densify the bmm LHS with a.contiguous() gated on get_sm_version() in (120… - #18678
Open
trtllm-agent wants to merge 1 commit into
Open
Conversation
… chunked-prefill path TestDeepSeekV3Lite::test_guided_decoding_4gpus[*-mtp_nextn=2] died on RTX PRO 6000 Blackwell (SM120) with an illegal memory access surfacing inside MLA.forward_context_with_chunked_prefill. A CUDA core dump identified the faulting kernel as nvjet_sm120_tst_mma_32x16x128_8_16x8x128_tmaAB_alignCD4_bz_TNNN with CUDBG_EXCEPTION_WARP_MMU_FAULT -- a cuBLAS GEMM, not an attention kernel. MLA._bmm_bf16_out passes its LHS as a head-major transpose of a [tokens, heads, dim] buffer, so the batch stride is the token count rather than the tile extent. On SM120/121 cuBLAS selects a TMA-based nvjet kernel for that layout, which cuTensorMapEncodeTiled cannot describe; the kernel faults with an MMU page fault that surfaces asynchronously at the next synchronize -- here inside the chunked-prefill loop, which is why the reported stack pointed at kv_b_proj rather than at the bmm. This is the same root cause as https://nvbugs/6700707, which reached it via generation CUDA-graph capture. Densify the LHS on those architectures so a non-TMA kernel is selected, mirroring the existing SM120 accommodation in fp8_block_scaling_bmm_out. Also keep the host_total_kv_lens save as a value copy: host_total_kv_lens[0] returns a 0-d view aliasing the element the chunk loop overwrites in place, so the post-loop restore wrote the last chunk's value back into itself and left the context total-KV length under-reported. Unwaive the two guided_decoding_4gpus mtp_nextn=2 cases the fault was masking. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
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 (2)
💤 Files with no reviewable changes (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughThe change fixes KV-length restoration during chunked prefill, makes BF16 BMM inputs contiguous on SM120/121, and removes two DeepSeekV3Lite guided-decoding test waivers. ChangesMLA attention validation
Estimated code review effort: 1 (Trivial) | ~5 minutes Merge Risk: ⚪ Minimal · up to The change restores chunked-prefill KV-length state correctly and applies the SM120/121 BF16 BMM layout workaround, with no remaining merge-blocking risk identified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description explains the root cause, fix, affected GPUs, test command, regression checks, and bug link. It does not use the template's exact Description, Test Coverage, and PR Checklist headings, and it does not include the formal checklist, but the required technical context is mostly complete.
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
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
2 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
a.contiguous()gated on get_sm_version() in (120, 121) so a non-TMA kernel is selected, mirroring the existing SM120 branch in fp8_block_scaling_bmm_out; kept the host_total_kv_lens.item()value-copy and removed the 2 waivers.pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_guided_decoding_4gpus[xgrammar-mtp_nextn=2] -vTest plan
Links
Reproduction comparison
Signature: tensorrt_llm.executor.utils.RequestError: Event loop terminated with error: Fatal error: CUDA error: an illegal memory access was encountered
Dev Engineer Review
MLA._bmm_bf16_outnow callsa.contiguous()on SM120/121 GPUs. This avoids the incompatible TMA-based cuBLAS kernel layout.host_total_kv_lensas a value copy. This prevents aliasing with values modified by the chunk loop.QA Engineer Review
tests/integration/test_lists/waives.txt.llguidanceandxgrammar.