[https://nvbugs/6571418][fix] Restore DeepSeek-V4-Pro GSM8K accuracy

cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu

@@ -88,6 +88,19 @@ struct Vec2Traits<__nv_bfloat16>
     }
 };
 
+template <typename T>
+__device__ __forceinline__ float2 roundToInput(float2 value)
+{
+    auto const rounded = Vec2Traits<T>::fromFloat2(value);
+    return Vec2Traits<T>::toFloat2(rounded);
+}
+
+template <typename T>
+__device__ __forceinline__ float2 normalizeAndRoundToInput(float2 value, float normScale)
+{
+    return roundToInput<T>(float2{value.x * normScale, value.y * normScale});
+}
+
 __device__ __forceinline__ float warpReduceSum(float value)
 {
     for (int mask = kWarpSize / 2; mask > 0; mask >>= 1)
@@ -274,7 +287,10 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8
 
     sumSquares = warpReduceSum(sumSquares);
     float const normScale = rsqrtf(sumSquares / static_cast<float>(kHeadDim) + eps);
-    float const fp8Scale = normScale * quantScale;
+
+    // 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.
@@ -335,7 +351,8 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8
                 for (int j = 0; j < kPairsPerVec; ++j)
                 {
                     float2 const v = Vec2Traits<T>::toFloat2(srcPairs[j]);
-                    out.pairs[j] = __nv_fp8x2_e4m3(float2{v.x * fp8Scale, v.y * fp8Scale});
+                    float2 const normalized = normalizeAndRoundToInput<T>(v, normScale);
+                    out.pairs[j] = __nv_fp8x2_e4m3(float2{normalized.x * quantScale, normalized.y * quantScale});
                 }
                 *reinterpret_cast<typename Fp8VecStore<kEltsPerVec>::Type*>(nopeOut + vecIdx * kEltsPerVec)
                     = out.packed;
@@ -349,11 +366,13 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8
                 for (int j = 0; j < kPairsPerVec; ++j)
                 {
                     float2 const v = Vec2Traits<T>::toFloat2(srcPairs[j]);
-                    float2 const normalized{v.x * normScale, v.y * normScale};
+                    float2 const normalized = normalizeAndRoundToInput<T>(v, normScale);
                     float2 const coef = cos_sin_cache[static_cast<int64_t>(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<T>(rotated);
+                    out.pairs[j]
+                        = __nv_fp8x2_e4m3(float2{roundedRotated.x * quantScale, roundedRotated.y * quantScale});
                 }
                 *reinterpret_cast<typename Fp8VecStore<kEltsPerVec>::Type*>(nopeOut + vecIdx * kEltsPerVec)
                     = out.packed;
@@ -377,7 +396,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<T>(values[i], normScale);
+            float2 const scaled{normalized.x * quantScale, normalized.y * quantScale};
             nopeOutPair[pairIdx] = __nv_fp8x2_e4m3(scaled);
         }
 
@@ -386,14 +406,15 @@ __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<T>(values[i], normScale);
             if constexpr (kFuseRope)
             {
                 float2 const coef = cos_sin_cache[static_cast<int64_t>(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<T>(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
             {

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:]

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 the input dtype, rotate and round again, then scale for FP8.
+
+    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()
     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
 
@@ -1345,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)
 
@@ -1382,18 +1387,17 @@ 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()
     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)
@@ -1411,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 /
@@ -1423,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)
 
@@ -1441,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
@@ -1460,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)