[https://nvbugs/6517834][fix] Fix attn_dense LoRA input dim and Triton MoE padded-view output - #17153
Open
trtllm-agent wants to merge 1 commit into
Open
[https://nvbugs/6517834][fix] Fix attn_dense LoRA input dim and Triton MoE padded-view output#17153trtllm-agent wants to merge 1 commit into
trtllm-agent wants to merge 1 commit into
Conversation
… LoRA in_dim The test failed on setup because it requested six fixtures its body never uses; gpt_example_root pip-installs examples/models/core/gpt/requirements.txt, whose directory no longer exists. Dropping them exposed two LoRA defects and a stale reference value. The node ID is unchanged: none of the six are parametrized, and pytest_generate_tests early-returns for every function but test_unittests_v2, so the two indirect params remain the only contributors. TritonFusedMoE was the only MoE backend returning a non-dense gemm2 output: it de-padded by slicing, yielding a view whose row stride is the N-aligned padded width. Cutlass passes unpadded_hidden_size into its kernel, and DeepGemm and Marlin write a fresh tensor, so all three already emit dense rows. Consumers that take a bare data pointer and derive the row stride from the hidden size -- the LoRA grouped GEMM among them, whose ABI carries no stride -- therefore read every row past the first at the wrong offset, silently. Instrumenting the run shows the padded view reaching the op in 46 of 72 calls, as the next layer's attention q/k/v hidden_states. Fixed at the producer so all consumers and both LoRA dispatch paths benefit. The branch only runs when padding was applied (hidden_size % 256 != 0), so aligned models are untouched. attn_dense (o_proj) consumes the concatenated attention heads, so its LoRA in_dim is numHeads * attnHeadSize, not hidden. These coincide for most models but not GPT-OSS-20B: hidden is 2880 while 64 heads of 64 give 4096, and the adapter's o_proj.lora_A is correspondingly (8, 4096) on disk. loraUtils only checks <=, so the undersized dim silently made weightsOut start inside lora_A. expected_output was recorded against a different base model -- the adapter declares base_model_name_or_path workspace-gpt-oss-20b/nemo2, not the HF checkpoint the test loads -- on a build carrying both defects above. It never reaches EOS, running to the 50-token cap and starting a fresh User: turn, and it contradicts its own prompt by saying the winner "will be a great opportunity" rather than will get $100. Replaced with the reference generation, which TensorRT-LLM, the CUTLASS MoE backend, and HuggingFace with PEFT all emit token-for-token, terminating at EOS after 43 tokens. Decoding is greedy, so this is deterministic rather than sampled. A sweep of 12 LoRA scales, 6 qkv/o splits, and 5 prompt formulations run entirely under HuggingFace plus PEFT caps at 0.58 against the old value, so no implementation could reach the 0.8 gate. The 0.8 threshold is unchanged and still rejects every regression it was written for: the padded-stride output scores 0.61, an unapplied adapter 0.02, a wrong LoRA scale 0.34, and the previous value 0.47. Reverting only the producer fix makes the test fail again, confirming it is load-bearing. Waiver removed. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
trtllm-agent
requested review from
BowenFu,
LarryXFly,
allisonlim-nv,
brnguyen2,
mikeiovine,
sunnyqgg,
xxi-nv,
yingguo-trt,
yiqingy0,
yuanjingx87 and
yufeiwu-nv
August 1, 2026 01:55
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 (4)
💤 Files with no reviewable changes (1)
WalkthroughChangesLoRA dimension and integration test updates
MXFP4 MoE output contiguity
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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
LoraModule::createLoraModulesbuilt theattn_dense/cross_attn_densemodules withhiddenas the input dimension, buto_projconsumes the concatenated attention heads, so its true input dim isnum_heads * head_dim— identical tohiddenfor most models but not for GPT-OSS-20B (hidden=2880 vs 64×64=4096), which sized the adapter buffers wrong. Independently,TritonMXFP4FusedMoEMethodreturned a non-contiguous slice view when stripping alignment padding, andlora_grouped_gemmtakes a bare data pointer and derives the row stride from the hidden size, so every row but the first was read at the wrong offset. The test itself could never pass regardless: its fixtures pip-installed a deletedrequirements.txt(setup error) and its golden string encoded the pre-fix buggy output.qOutSizefor the dense-projection input dim inloraModule.cpp, and.contiguous()on the de-padded Triton MoE output so it matches the dense output every other MoE backend already emits. On the test side, dropped the four unused fixtures that caused the setup error and refreshed the golden string to the output now produced, cross-validated as byte-identical against HuggingFace+PEFT and the Cutlass MoE backend; the waiver was then removed. Both product fixes were shown load-bearing by ablation, and the golden was chosen over relaxing the 0.8 ROUGE gate whose discriminating power was measured against four injected-defect variants (0.61/0.02/0.34/0.47 — all correctly rejected).Test plan
Links
Dev Engineer Review
attn_denseandcross_attn_densenow useqOutSize(num_heads * head_dim) as the LoRA input dimension.hiddento4096.lora_grouped_gemmconsumes them.test_gpt_oss_20b_lora_torchwas removed.QA Engineer Review
test_gpt_oss_20b_lora_torch.tests/integration/test_lists/waives.txt.0.8ROUGE threshold.