[Fix] [NVIDIA] modelopt: Keep SM12x FP8 batch-invariant in deterministic mode

python/sglang/kernels/ops/gemm/fp8_kernel.py

@@ -990,10 +990,10 @@ def scaled_mm_kernel(
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         masks_k = offsets_k < K
         masks_a = masks_am[:, None] & masks_k[None, :]
-        a = tl.load(a_ptrs, mask=masks_a)
+        a = tl.load(a_ptrs, mask=masks_a, other=0.0)
 
         masks_b = masks_k[:, None] & masks_bn[None, :]
-        b = tl.load(b_ptrs, mask=masks_b)
+        b = tl.load(b_ptrs, mask=masks_b, other=0.0)
 
         # Accumulate results.
         accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -2381,6 +2381,39 @@ def use_aiter_bpreshuffle_gemm(output_size: int) -> bool:
     return _use_aiter and output_size % 64 == 0
 
 
+def apply_fp8_linear_deterministic(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_scale: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Static FP8 linear with one reduction schedule across batch sizes."""
+    output_shape = [*input.shape[:-1], weight.shape[1]]
+    input_2d = input.reshape(-1, input.shape[-1]).contiguous()
+    qinput, x_scale = static_quant_fp8(input_2d, input_scale, repeat_scale=False)
+    # Keep K reduction and N tiling fixed so prefill and cache hits agree.
+    # Only the row tile varies, avoiding wasted decode work and small tiles
+    # on long prefills without changing each row's reduction.
+    m = input_2d.shape[0]
+    block_m = 16 if m <= 16 else 32 if m < 128 else 128
+    output = triton_scaled_mm(
+        qinput,
+        weight,
+        x_scale,
+        weight_scale,
+        input.dtype,
+        bias,
+        block_size_m=block_m,
+        block_size_n=64,
+        block_size_k=128,
+        use_heuristic=False,
+        num_warps=4,
+        num_stages=3,
+    )
+    return output.view(*output_shape)
+
+
 def apply_fp8_linear_bmm_flashinfer(
     input: torch.Tensor,
     weight: torch.Tensor,

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -42,6 +42,7 @@
 from sglang.srt.layers.quantization.fp8_utils import (
     apply_fp8_linear,
     apply_fp8_linear_bmm_flashinfer,
+    apply_fp8_linear_deterministic,
     can_auto_enable_marlin_fp8,
     cutlass_fp8_supported,
     flashinfer_per_tensor_fp8_supported,
@@ -65,7 +66,7 @@
 )
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.utils import alias_or_bind_derived_param, copy_or_rebind_param
-from sglang.srt.runtime_context import get_platform
+from sglang.srt.runtime_context import get_exec, get_platform
 from sglang.srt.utils.common import (
     get_device_capability,
     is_cuda,
@@ -657,6 +658,17 @@ def apply(
                 size_k=layer.input_size_per_partition,
                 bias=bias,
             )
+        if (
+            self.use_sm120_fp8
+            and get_exec().deterministic.enable_deterministic_inference
+        ):
+            return apply_fp8_linear_deterministic(
+                input=x,
+                weight=layer.weight,
+                weight_scale=layer.weight_scale,
+                input_scale=layer.input_scale,
+                bias=bias,
+            )
         if self.use_sm120_fp8:
             from sglang.kernels.ops.gemm import try_sm120_fp8_linear
 

test/registered/kernels/ops/gemm/test_sm120_fp8_linear.py

@@ -3,11 +3,14 @@
 from __future__ import annotations
 
 import sys
+from types import SimpleNamespace
 
 import pytest
 import torch
 
 from sglang.kernels.ops.gemm import try_sm120_fp8_linear
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
+from sglang.srt.layers.quantization import modelopt_quant
 from sglang.srt.layers.quantization.fp8_utils import apply_fp8_linear_bmm_flashinfer
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -92,6 +95,77 @@ def _run_sm120(args, *, m=None, vector_scales=False, bias=None):
     )
 
 
+@pytest.mark.parametrize("small_m,large_m", [(1, 64), (8, 128), (120, 248)])
+@pytest.mark.parametrize("n,k", [(4096, 4096), (16384, 5120), (5120, 6144), (96, 80)])
+def test_deterministic_modelopt_fp8_preserves_cached_suffix(
+    monkeypatch, small_m, large_m, n, k
+):
+    """Cold prefill and a shorter cache-hit batch must score identical rows."""
+    monkeypatch.setattr(
+        modelopt_quant,
+        "get_exec",
+        lambda: SimpleNamespace(
+            deterministic=SimpleNamespace(enable_deterministic_inference=True)
+        ),
+        raising=False,
+    )
+    method = object.__new__(modelopt_quant.ModelOptFp8LinearMethod)
+    method.use_marlin = False
+    method.use_sm120_fp8 = True
+    method.cutlass_fp8_supported = True
+    x, weight, weight_scale, input_scale = _make_inputs(large_m, n, k, seed=29)
+    layer = SimpleNamespace(
+        weight=weight,
+        weight_scale=weight_scale,
+        input_scale=input_scale,
+        sm120_fp8_alpha=input_scale * weight_scale,
+        use_flashinfer_bmm=True,
+    )
+    full = method.apply(layer, x)
+    suffix = method.apply(layer, x[-small_m:].contiguous())
+    torch.testing.assert_close(suffix, full[-small_m:], rtol=0, atol=0)
+    quantized, _ = static_quant_fp8(x, input_scale, repeat_scale=False)
+    reference = (quantized.float() @ weight.float()) * input_scale * weight_scale
+    relative_rmse = (
+        full.float() - reference
+    ).square().mean().sqrt() / reference.square().mean().sqrt().clamp_min(1e-8)
+    assert relative_rmse < 0.005
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_deterministic_fp8_graph_replay_with_bias_and_strided_input(dtype):
+    deterministic_linear = modelopt_quant.apply_fp8_linear_deterministic
+    x, weight, weight_scale, input_scale = _make_inputs(6, 96, 80, seed=17)
+    x = x.to(dtype).view(2, 3, 80).transpose(0, 1)
+    bias = torch.randn(96, dtype=dtype, device="cuda")
+    args = (x, weight, weight_scale, input_scale, bias)
+    deterministic_linear(*args)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = deterministic_linear(*args)
+
+    graph.replay()
+    before = actual.clone()
+    x.copy_(torch.randn_like(x).mul_(2))
+    graph.replay()
+    expected = deterministic_linear(*args)
+    assert actual.shape == (3, 2, 96)
+    assert actual.dtype == dtype
+    assert not torch.equal(actual, before)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    quantized, _ = static_quant_fp8(
+        x.reshape(-1, 80).contiguous(), input_scale, repeat_scale=False
+    )
+    reference = (
+        (quantized.float() @ weight.float()) * input_scale * weight_scale + bias
+    ).view_as(actual)
+    relative_rmse = (actual.float() - reference).square().mean().sqrt() / (
+        reference.square().mean().sqrt().clamp_min(1e-8)
+    )
+    assert relative_rmse < 0.005
+
+
 @pytest.mark.parametrize("seed", [0, 1, 7])
 @pytest.mark.parametrize("m,n,k", SUPPORTED_SHAPES)
 def test_supported_shapes_match_flashinfer(m: int, n: int, k: int, seed: int):