From fb045754dc58e6c35abee66da27e9cbb6cecd15d Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 19 Aug 2026 06:23:36 -0700 Subject: [PATCH 1/3] [https://nvbugs/6571418][fix] Restore DeepSeek-V4-Pro GSM8K accuracy Preserve the legacy BF16 precision boundary in the fused QNorm path before FP8 conversion and update the fused references to guard it. Rely on the existing statistical accuracy gate instead of comparing against the raw reference twice, then remove the B200 waiver. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../kernels/deepseekV4QNormKernel.cu | 24 +++++++++++++++---- .../defs/accuracy/test_llm_api_pytorch.py | 5 +--- .../custom_ops/test_deepseek_v4_q_norm.py | 12 ++++++---- .../modeling/test_modeling_deepseekv4.py | 15 +++++++----- 4 files changed, 37 insertions(+), 19 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu index 47d7640e7e62..4f5c932a42bf 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu @@ -88,6 +88,14 @@ struct Vec2Traits<__nv_bfloat16> } }; +template +__device__ __forceinline__ float2 normalizeAndRoundToInput(float2 value, float normScale) +{ + float2 const normalized{value.x * normScale, value.y * normScale}; + auto const rounded = Vec2Traits::fromFloat2(normalized); + return Vec2Traits::toFloat2(rounded); +} + __device__ __forceinline__ float warpReduceSum(float value) { for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) @@ -274,7 +282,11 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 sumSquares = warpReduceSum(sumSquares); float const normScale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); - float const fp8Scale = normScale * quantScale; + + // Preserve the legacy two-kernel precision contract: RMSNorm stores T + // (bf16 or half), then the quantization/RoPE kernel reloads T before the + // FP8 conversion. Skipping this round changes enough FP8 codes to regress + // DeepSeek-V4 model accuracy. // Position depends on the token, which every lane of the warp shares, so this is // warp-uniform and hoisted out of the store loop. @@ -335,7 +347,8 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 for (int j = 0; j < kPairsPerVec; ++j) { float2 const v = Vec2Traits::toFloat2(srcPairs[j]); - out.pairs[j] = __nv_fp8x2_e4m3(float2{v.x * fp8Scale, v.y * fp8Scale}); + float2 const normalized = normalizeAndRoundToInput(v, normScale); + out.pairs[j] = __nv_fp8x2_e4m3(float2{normalized.x * quantScale, normalized.y * quantScale}); } *reinterpret_cast::Type*>(nopeOut + vecIdx * kEltsPerVec) = out.packed; @@ -349,7 +362,7 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 for (int j = 0; j < kPairsPerVec; ++j) { float2 const v = Vec2Traits::toFloat2(srcPairs[j]); - float2 const normalized{v.x * normScale, v.y * normScale}; + float2 const normalized = normalizeAndRoundToInput(v, normScale); float2 const coef = cos_sin_cache[static_cast(kRopeDim) * positionId + ropePairBase + j]; float2 const rotated{ coef.x * normalized.x - coef.y * normalized.y, coef.x * normalized.y + coef.y * normalized.x}; @@ -377,7 +390,8 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 for (int i = 0; i < kPairsPerLane - 1; ++i) { int const pairIdx = i * kWarpSize + laneId; - float2 const scaled{values[i].x * fp8Scale, values[i].y * fp8Scale}; + float2 const normalized = normalizeAndRoundToInput(values[i], normScale); + float2 const scaled{normalized.x * quantScale, normalized.y * quantScale}; nopeOutPair[pairIdx] = __nv_fp8x2_e4m3(scaled); } @@ -386,7 +400,7 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 constexpr int i = kPairsPerLane - 1; int const pairIdx = i * kWarpSize + laneId; // in [kNopePairs, kPairsPerRow) int const ropePairIdx = pairIdx - kNopePairs; // in [0, kRopeDim/2) - float2 const normalized{values[i].x * normScale, values[i].y * normScale}; + float2 const normalized = normalizeAndRoundToInput(values[i], normScale); if constexpr (kFuseRope) { float2 const coef = cos_sin_cache[static_cast(kRopeDim) * positionId + ropePairIdx]; diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 456f6895c7a9..43967fb760e3 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4208,11 +4208,8 @@ def test_gsm8k_full_accuracy(self): spec_dec_algo=llm.args.speculative_config.decoding_type) assert acc_params.num_samples == GSM8K.NUM_SAMPLES with mock.patch.dict(os.environ, {"INTEGRATION_TEST": "0"}): - score = task.evaluate( + task.evaluate( llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - assert score >= acc_params.ref_accuracy, ( - f"GSM8K accuracy {score:.3f} is below recorded reference " - f"{acc_params.ref_accuracy:.3f}") @pytest.mark.timeout(14400) diff --git a/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py b/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py index 0309b9dcf6cd..79917c10ed7a 100644 --- a/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py +++ b/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py @@ -101,15 +101,19 @@ def _reference_q_norm_fused_fp8( eps: float, quant_scale_qkv: float, ): - """Reference: per-row RMSNorm in fp32; split last column-axis into nope/rope; - nope path multiplied by quant_scale_qkv then cast to fp8_e4m3; rope path cast - back to input dtype. + """Reference: per-row RMSNorm in fp32, rounded back to the input dtype; + split the last column-axis into nope/rope; multiply the nope path by + quant_scale_qkv and cast it to fp8_e4m3. + + The input-dtype round matches the legacy two-kernel path's norm store and + reload. It is part of the model's numerical contract: skipping it changes + enough FP8 codes to regress DeepSeek-V4 GSM8K accuracy. """ num_tokens = q.shape[0] rope_dim = head_dim - nope_dim q_view = q.view(num_tokens * num_heads, head_dim).float() inv_rms = torch.rsqrt(q_view.pow(2).mean(dim=-1, keepdim=True) + eps) - normalized = q_view * inv_rms + normalized = (q_view * inv_rms).to(q.dtype).float() nope_fp32 = normalized[:, :nope_dim] * quant_scale_qkv rope_fp32 = normalized[:, nope_dim:] diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 3abeb9471426..7544060e8814 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -1328,11 +1328,15 @@ def _make_cos_sin(device: torch.device) -> torch.Tensor: def _reference( q: torch.Tensor, cos_sin: torch.Tensor, positions: torch.Tensor, num_heads: int ) -> torch.Tensor: - """RMS-norm over the whole 512-wide head, rotate the tail, scale for FP8.""" + """RMS-norm, round to BF16, rotate the tail, and scale for FP8. + + The BF16 round matches the unfused norm-store/RoPE-load path and is required + before quantization to preserve DeepSeek-V4 model accuracy. + """ num_tokens = q.shape[0] x = q.view(num_tokens, num_heads, HEAD_DIM).float() inv_rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) - normed = x * inv_rms + normed = (x * inv_rms).to(q.dtype).float() nope = normed[..., :NOPE_DIM] * QUANT_SCALE @@ -1382,10 +1386,9 @@ def _assert_matches(quant_q, q_pe, reference, num_heads): A relative tolerance on the dequantized values has to be at least one e4m3 step (~13%) to absorb rounding, and that is wide enough to swallow real bugs -- normalizing over 448 dims instead of 512 is only a 6.9% shift. So quantize - the reference the same way and require the codes to agree. The kernel folds - inv_rms and the quant scale into a single multiply where the reference uses - two, so a few values sit on the other side of a rounding boundary; those get - a small budget, capped at one FP8 step each. + the reference the same way and require the codes to agree. CUDA and PyTorch + can still choose adjacent FP8 values at exact rounding midpoints, so those + get a small budget capped at one FP8 step each. """ num_tokens = quant_q.shape[0] got = quant_q.view(num_tokens, num_heads, HEAD_DIM).float() From dfcd2881e0a6dd7414143b99f5963aad97ea7892 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:49:04 -0700 Subject: [PATCH 2/3] [https://nvbugs/6571418][fix] Preserve post-RoPE precision boundary Round fused RoPE output through the input dtype before FP8 conversion to match the legacy RoPE store/reload semantics. Tighten the reference and cover BF16/FP16 context and generation paths. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../kernels/deepseekV4QNormKernel.cu | 25 ++++++++++++------- .../modeling/test_modeling_deepseekv4.py | 23 ++++++++++------- 2 files changed, 30 insertions(+), 18 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu index 4f5c932a42bf..3fc590c9af2f 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu @@ -89,13 +89,18 @@ struct Vec2Traits<__nv_bfloat16> }; template -__device__ __forceinline__ float2 normalizeAndRoundToInput(float2 value, float normScale) +__device__ __forceinline__ float2 roundToInput(float2 value) { - float2 const normalized{value.x * normScale, value.y * normScale}; - auto const rounded = Vec2Traits::fromFloat2(normalized); + auto const rounded = Vec2Traits::fromFloat2(value); return Vec2Traits::toFloat2(rounded); } +template +__device__ __forceinline__ float2 normalizeAndRoundToInput(float2 value, float normScale) +{ + return roundToInput(float2{value.x * normScale, value.y * normScale}); +} + __device__ __forceinline__ float warpReduceSum(float value) { for (int mask = kWarpSize / 2; mask > 0; mask >>= 1) @@ -283,10 +288,9 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 sumSquares = warpReduceSum(sumSquares); float const normScale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); - // Preserve the legacy two-kernel precision contract: RMSNorm stores T - // (bf16 or half), then the quantization/RoPE kernel reloads T before the - // FP8 conversion. Skipping this round changes enough FP8 codes to regress - // DeepSeek-V4 model accuracy. + // Preserve the legacy multi-kernel precision contract: RMSNorm stores T + // (bf16 or half) before RoPE/quantization, and RoPE stores T again before + // the standalone FP8 conversion. Skipping either round changes FP8 codes. // Position depends on the token, which every lane of the warp shares, so this is // warp-uniform and hoisted out of the store loop. @@ -366,7 +370,9 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 float2 const coef = cos_sin_cache[static_cast(kRopeDim) * positionId + ropePairBase + j]; float2 const rotated{ coef.x * normalized.x - coef.y * normalized.y, coef.x * normalized.y + coef.y * normalized.x}; - out.pairs[j] = __nv_fp8x2_e4m3(float2{rotated.x * quantScale, rotated.y * quantScale}); + float2 const roundedRotated = roundToInput(rotated); + out.pairs[j] + = __nv_fp8x2_e4m3(float2{roundedRotated.x * quantScale, roundedRotated.y * quantScale}); } *reinterpret_cast::Type*>(nopeOut + vecIdx * kEltsPerVec) = out.packed; @@ -406,8 +412,9 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 float2 const coef = cos_sin_cache[static_cast(kRopeDim) * positionId + ropePairIdx]; float2 const rotated{ coef.x * normalized.x - coef.y * normalized.y, coef.x * normalized.y + coef.y * normalized.x}; + float2 const roundedRotated = roundToInput(rotated); nopeOutPair[kNopePairs + ropePairIdx] - = __nv_fp8x2_e4m3(float2{rotated.x * quantScale, rotated.y * quantScale}); + = __nv_fp8x2_e4m3(float2{roundedRotated.x * quantScale, roundedRotated.y * quantScale}); } else { diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 7544060e8814..95dcba834853 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -1328,10 +1328,10 @@ def _make_cos_sin(device: torch.device) -> torch.Tensor: def _reference( q: torch.Tensor, cos_sin: torch.Tensor, positions: torch.Tensor, num_heads: int ) -> torch.Tensor: - """RMS-norm, round to BF16, rotate the tail, and scale for FP8. + """RMS-norm, round to the input dtype, rotate and round again, then scale for FP8. - The BF16 round matches the unfused norm-store/RoPE-load path and is required - before quantization to preserve DeepSeek-V4 model accuracy. + The two input-dtype rounds match the unfused norm-store/RoPE-load and + RoPE-store/quantization-load boundaries. """ num_tokens = q.shape[0] x = q.view(num_tokens, num_heads, HEAD_DIM).float() @@ -1349,6 +1349,7 @@ def _reference( rotated = torch.empty_like(rope) rotated[..., 0::2] = cos * even - sin * odd rotated[..., 1::2] = cos * odd + sin * even + rotated = rotated.to(q.dtype).float() return torch.cat([nope, rotated * QUANT_SCALE], dim=-1) @@ -1395,8 +1396,8 @@ def _assert_matches(quant_q, q_pe, reference, num_heads): expected = reference.to(torch.float8_e4m3fn).float() differing = got != expected - frac = differing.float().mean().item() - assert frac < 0.01, f"{frac:.4%} of FP8 codes differ from the reference" + num_differing = int(differing.sum().item()) + assert num_differing <= 16, f"{num_differing} FP8 codes differ from the reference" if differing.any(): scale = torch.maximum(got.abs(), expected.abs()).clamp_min(1e-6) @@ -1414,7 +1415,8 @@ def _assert_matches(quant_q, q_pe, reference, num_heads): [(4, 2), (6, 3)], ids=["heads4_seqlen2_pow2", "heads6_seqlen3_divide"], ) -def test_fused_rope_generation_positions(num_heads: int, seq_len: int) -> None: +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_fused_rope_generation_positions(num_heads: int, seq_len: int, dtype: torch.dtype) -> None: """Uniform query length: position = cache_len[batch] - seq_len + local_token. The parameters flip both power-of-two shortcuts the kernel takes (`row / @@ -1426,7 +1428,7 @@ def test_fused_rope_generation_positions(num_heads: int, seq_len: int) -> None: num_seqs = 3 num_tokens = num_seqs * seq_len - q = torch.randn((num_tokens, num_heads * HEAD_DIM), dtype=torch.bfloat16, device=device) + q = torch.randn((num_tokens, num_heads * HEAD_DIM), dtype=dtype, device=device) cos_sin = _make_cos_sin(device) cache_seq_lens = torch.tensor([16, 40, 71], dtype=torch.int32, device=device) @@ -1444,7 +1446,10 @@ def test_fused_rope_generation_positions(num_heads: int, seq_len: int) -> None: [(4, 0), (6, 5)], ids=["heads4_fresh_prefill", "heads6_chunked_prefill"], ) -def test_fused_rope_context_positions(num_heads: int, cached_offset: int) -> None: +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_fused_rope_context_positions( + num_heads: int, cached_offset: int, dtype: torch.dtype +) -> None: """Ragged: position = local_token + (cache_len[seq] - current_seq_len). `cached_offset > 0` is the chunked-prefill / block-reuse case, where part of @@ -1463,7 +1468,7 @@ def test_fused_rope_context_positions(num_heads: int, cached_offset: int) -> Non cache_seq_lens = torch.tensor( [s + cached_offset for s in seq_lens], dtype=torch.int32, device=device ) - q = torch.randn((num_tokens, num_heads * HEAD_DIM), dtype=torch.bfloat16, device=device) + q = torch.randn((num_tokens, num_heads * HEAD_DIM), dtype=dtype, device=device) cos_sin = _make_cos_sin(device) quant_q, q_pe = _run_op(q, num_heads, cos_sin, cache_seq_lens, 0, cu_q_seqlens) From 18e84c0900711f288b38521b8e99be2cb4b79e94 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Sat, 22 Aug 2026 09:13:35 -0700 Subject: [PATCH 3/3] [https://nvbugs/6571418][test] Preserve DeepSeek-V4-Pro accuracy gate Keep the full accuracy test identical to main so the kernel fix does not lower its existing ref_accuracy requirement. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 43967fb760e3..456f6895c7a9 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4208,8 +4208,11 @@ def test_gsm8k_full_accuracy(self): spec_dec_algo=llm.args.speculative_config.decoding_type) assert acc_params.num_samples == GSM8K.NUM_SAMPLES with mock.patch.dict(os.environ, {"INTEGRATION_TEST": "0"}): - task.evaluate( + score = task.evaluate( llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} is below recorded reference " + f"{acc_params.ref_accuracy:.3f}") @pytest.mark.timeout(14400)