diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu index 47d7640e7e62..3fc590c9af2f 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu @@ -88,6 +88,19 @@ struct Vec2Traits<__nv_bfloat16> } }; +template +__device__ __forceinline__ float2 roundToInput(float2 value) +{ + 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) @@ -274,7 +287,10 @@ __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 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::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,11 +366,13 @@ __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}; - 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; @@ -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(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(values[i], normScale); if constexpr (kFuseRope) { 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/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..95dcba834853 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 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)