From cf079c9a00a1a352043d07f759a9698c8c2eba87 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 31 Jul 2026 17:41:44 -0700 Subject: [PATCH] [nvbugs/6517834][fix] Emit dense Triton MoE output and fix attn_dense 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> --- cpp/tensorrt_llm/runtime/loraModule.cpp | 4 +++- .../_torch/modules/fused_moe/fused_moe_triton.py | 7 ++++++- tests/integration/defs/examples/test_gpt.py | 9 +++++---- tests/integration/test_lists/waives.txt | 1 - 4 files changed, 14 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/runtime/loraModule.cpp b/cpp/tensorrt_llm/runtime/loraModule.cpp index 30ec12733808..1aacf6b9b1f4 100644 --- a/cpp/tensorrt_llm/runtime/loraModule.cpp +++ b/cpp/tensorrt_llm/runtime/loraModule.cpp @@ -61,8 +61,10 @@ std::vector LoraModule::createLoraModules(std::vector c case ModuleType::kCROSS_ATTN_K: case ModuleType::kATTN_V: case ModuleType::kCROSS_ATTN_V: modules.emplace_back(t, hidden, kvOutSize, false, true, -1, 0); break; + // attn_dense (o_proj) consumes the concatenated attention heads, so its input dim is qOutSize, not + // hidden. Those coincide for most models but not all, e.g. GPT-OSS-20B has hidden=2880 and 64x64 heads. case ModuleType::kATTN_DENSE: - case ModuleType::kCROSS_ATTN_DENSE: modules.emplace_back(t, hidden, hidden, false, true, 1, -1); break; + case ModuleType::kCROSS_ATTN_DENSE: modules.emplace_back(t, qOutSize, hidden, false, true, 1, -1); break; case ModuleType::kMLP_H_TO_4H: modules.emplace_back(t, hidden, mlpHidden, false, true, -1, 0); break; case ModuleType::kMLP_GATE: modules.emplace_back(t, hidden, mlpHidden, false, true, -1, 0); break; case ModuleType::kMLP_4H_TO_H: modules.emplace_back(t, mlpHidden, hidden, false, true, 1, -1); break; diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py index 9414a8aca91d..c6164fc7b54f 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py @@ -1436,7 +1436,12 @@ def _maybe_remove_padding(gemm_output, expected_size): if gemm_output.shape[-1] != expected_size: assert gemm_output.shape[ -1] % self.k_alignment == 0, "The padding is not done correctly" - gemm_output = gemm_output[:, :expected_size] + # Materialize instead of returning the slice view: downstream + # kernels that take a bare data pointer and derive the row + # stride from the hidden size (e.g. lora_grouped_gemm) would + # read every row but the first at the wrong offset. The other + # MoE backends already emit dense output. + gemm_output = gemm_output[:, :expected_size].contiguous() return gemm_output gemm2_output = _maybe_remove_padding(gemm2_output, module.hidden_size) diff --git a/tests/integration/defs/examples/test_gpt.py b/tests/integration/defs/examples/test_gpt.py index 64200b3549db..f182d1bc5df8 100644 --- a/tests/integration/defs/examples/test_gpt.py +++ b/tests/integration/defs/examples/test_gpt.py @@ -32,9 +32,7 @@ @pytest.mark.parametrize("llm_lora_model_root", ['gpt-oss-20b-lora-adapter_NIM_r8'], indirect=True) -def test_gpt_oss_20b_lora_torch(gpt_example_root, llm_venv, gpt_oss_model_root, - llm_datasets_root, llm_rouge_root, engine_dir, - cmodel_dir, llm_lora_model_root): +def test_gpt_oss_20b_lora_torch(gpt_oss_model_root, llm_lora_model_root): """Run GPT-OSS-20B with LoRA adapter using Torch backend.""" print(f"Using LoRA from: {llm_lora_model_root}") @@ -63,7 +61,10 @@ def test_gpt_oss_20b_lora_torch(gpt_example_root, llm_venv, gpt_oss_model_root, sampling_params, lora_request=lora_request) - expected_output = " Hey Mason! I hope you're doing well. I was thinking about the next week's football tournament and I wanted to give you a hint that we should compete in it. The winner will be a great opportunity for us to win $100.\n\nUser:" + # Decoding is greedy, so this is deterministic. Cross-checked against + # HuggingFace + PEFT and the Cutlass MoE backend, which emit the same + # tokens and stop at EOS. + expected_output = " Hey Mason, I was thinking that we should compete in next week's football tournament. The winner will get $100. Let me know if you are interested. Cheers, [Your Name]" for i, output in enumerate(outputs): print(f"Prompt {i+1}: {prompts[i]}") diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 1c7212ea9093..60d4a40c5865 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -131,7 +131,6 @@ disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-C disaggregated/test_workers.py::test_workers_kv_cache_aware_router_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[trtllm-torch-cudagraph] SKIP (https://nvbugs/6426841) -examples/test_gpt.py::test_gpt_oss_20b_lora_torch[gpt-oss-20b-lora-adapter_NIM_r8-gpt-oss-20b] SKIP (https://nvbugs/6517834) examples/test_ray.py::test_llm_inference_distributed_ray[pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_llm_inference_distributed_ray[tp2pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/5612502)