From 75febe4bb811b7f485ce832c72aa6ec98abc8788 Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Thu, 21 May 2026 02:33:38 -0700 Subject: [PATCH 1/2] [nvbugs/6186880][fix] Handle non-quantized MoE in DeepEP post-quant dispatch and ensure FP4 DSL kv stride Two issues caused TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_*] to fail on B300: 1. DeepEP post-quant dispatch was entered whenever the comm strategy declared nvfp4 support, but the MoE backend may still return hidden_states_sf=None when the local module is excluded from quantization (e.g. an MTP layer's MoE in a partially-quantized checkpoint). Passing (hidden_states, None) to deep_ep_buffer.dispatch causes recv_x to come back as a single tensor, triggering 'too many values to unpack (expected 2)'. Route to the pre-quant (plain) dispatch path when hidden_states_sf is None. 2. The FP4 paged-MQA-logits CuTe DSL kernel binding asserts kv_flat.stride(0) == phys_block_kv * (D//2 + 4). When kv_fused is a non-contiguous view of a larger buffer (FP4 indexer-K cache slice), reshape preserves the larger row stride and the TVM-FFI bound check rejects the call. Force kv_flat to be contiguous so the row stride always matches the kernel's compile-time fake stride. Removes the corresponding waivers now that the test passes. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 10 ++++++++-- .../_torch/modules/fused_moe/communication/deep_ep.py | 7 ++++++- tests/integration/test_lists/waives.txt | 4 ---- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 013a30523e78..9108ee2adb5c 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -6981,8 +6981,14 @@ def forward( else: w_2d = weights.reshape(B, N).t() - # Flatten fused KV to [num_phys_blocks, block_bytes] - kv_flat = kv_fused.reshape(num_phys_blocks, -1) + # Flatten fused KV to [num_phys_blocks, block_bytes]. + # The compiled DSL kernel's TVM-FFI binding asserts + # `kv_flat.stride(0) == phys_block_kv * (D//2 + 4)`. Force a + # contiguous tensor so the row stride matches even when + # `kv_fused` is a non-contiguous view of a larger buffer + # (e.g. an FP4 indexer-K cache slice with row stride > row + # width). See nvbugs/6186880. + kv_flat = kv_fused.reshape(num_phys_blocks, -1).contiguous() # Allocate output with alignment padding SPLIT_KV = compute_block_kv * 2 # NUM_MATH_WG = 2 diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py index d3e1c68edd98..b7d962806ef4 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep.py @@ -190,7 +190,12 @@ def dispatch( if token_final_scales is not None and token_final_scales.dtype != torch.float32: token_final_scales = token_final_scales.to(torch.float32) - if not self.supports_post_quant_dispatch(): + # Even when the comm strategy declares post-quant dispatch support, the + # caller may still pass `hidden_states_sf=None` if the local MoE module + # is excluded from quantization (e.g. an MTP layer with a non-quantized + # MoE in a partially-quantized checkpoint). In that case the data is + # actually unquantized, so dispatch the plain hidden_states. + if not self.supports_post_quant_dispatch() or hidden_states_sf is None: # Pre-quant dispatch (unquantized data) ( hidden_states, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6f895057854d..2961b116b2e4 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -32,10 +32,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[disable_s accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[latency_default] SKIP (https://nvbugs/6185196) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline] SKIP (https://nvbugs/6185196) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_mtp1] SKIP (https://nvbugs/6185196) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp0] SKIP (https://nvbugs/6186880) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp1] SKIP (https://nvbugs/6186880) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp2] SKIP (https://nvbugs/6186880) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[fp4_indexer_dsl_mtp3] SKIP (https://nvbugs/6186880) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[baseline] SKIP (https://nvbugs/6185196) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] SKIP (https://nvbugs/5989920) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6084720) From 7822d38c26b61586e541e57fd6f04a3e0e65705b Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Thu, 21 May 2026 04:13:52 -0700 Subject: [PATCH 2/2] [nvbugs/6186880][fix] Drop kv_flat.contiguous() workaround per reviewer PR #14133 updated the FP4 paged-MQA-logits CuTe DSL op to accept non-contiguous kv tensors, so the .contiguous() forcing in this custom op is no longer needed and was triggering an extra copy kernel. Revert to the original reshape-only path. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 9108ee2adb5c..013a30523e78 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -6981,14 +6981,8 @@ def forward( else: w_2d = weights.reshape(B, N).t() - # Flatten fused KV to [num_phys_blocks, block_bytes]. - # The compiled DSL kernel's TVM-FFI binding asserts - # `kv_flat.stride(0) == phys_block_kv * (D//2 + 4)`. Force a - # contiguous tensor so the row stride matches even when - # `kv_fused` is a non-contiguous view of a larger buffer - # (e.g. an FP4 indexer-K cache slice with row stride > row - # width). See nvbugs/6186880. - kv_flat = kv_fused.reshape(num_phys_blocks, -1).contiguous() + # Flatten fused KV to [num_phys_blocks, block_bytes] + kv_flat = kv_fused.reshape(num_phys_blocks, -1) # Allocate output with alignment padding SPLIT_KV = compute_block_kv * 2 # NUM_MATH_WG = 2