[Fix] Two root causes of the H100 deepep TBO CI break: scale-tensor use-after-free + missing non-finite quant sanitization (#32188)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh

@@ -47,9 +47,15 @@ template <>
 struct WeightTrait<fp8_e4m3_t> {
   using packed2_t = fp8x2_e4m3_t;
   static constexpr float kMaxValue = DTypeTrait<fp8_e4m3_t>::kFloatMax;
-  // SATFINITE conversion saturates to +-448, no need to clip
+  // SATFINITE saturates +-inf / out-of-range values, but converts NaN to an
+  // fp8 NaN code. IEEE fminf/fmaxf return the non-NaN operand, so clamping
+  // first quantizes non-finite inputs to +-448 -- matching the v1/v2/Triton
+  // kernels. CUDA-graph capture warmup runs the model on whatever the
+  // (reused, uninitialized) buffers contain, and relies on this: an fp8 NaN
+  // code would poison the downstream GEMM and trip the sampler NaN check.
+  // For finite inputs the clamp is bit-identical to bare SATFINITE.
   SGL_DEVICE static packed2_t quant(const float2 v) {
-    return packed2_t{v};
+    return packed2_t{float2{fminf(fmaxf(v.x, -kMaxValue), kMaxValue), fminf(fmaxf(v.y, -kMaxValue), kMaxValue)}};
   }
 };
 
@@ -284,9 +290,14 @@ struct QuantTrait {
       scale_inv = static_cast<uint8_t>(exp);
       const float quant_scale = inv_scale_ue8m0(exp);
       const auto scale2 = cast<T2>(float2{quant_scale, quant_scale});
+      // Finite scaled values already lie in +-448 (2^exp >= amax/448), so the
+      // clamp only sanitizes non-finite inputs (see WeightTrait<fp8_e4m3_t>);
+      // __hmin2/__hmax2 return the non-NaN operand.
+      const auto lo2 = cast<T2>(float2{-kMaxValue, -kMaxValue});
+      const auto hi2 = cast<T2>(float2{kMaxValue, kMaxValue});
 #pragma unroll
       for (uint32_t i = 0; i < kVecSize / 2; ++i) {
-        out[i] = static_cast<Q2>(__hmul2(in[i], scale2));
+        out[i] = static_cast<Q2>(__hmin2(__hmax2(__hmul2(in[i], scale2), lo2), hi2));
       }
     } else {
       // fp32 scale: multiply in fp32 (hmul2 brings too much precision loss)

python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py

@@ -18,7 +18,29 @@ logger = logging.getLogger(__name__)
 
 if ENABLE_JIT_DEEPGEMM:
     import deep_gemm
-    from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor  # noqa: F401
+    from deep_gemm.utils.layout import (
+        get_mn_major_tma_aligned_tensor as _get_mn_major_tma_aligned_tensor,
+    )
+
+    def get_mn_major_tma_aligned_tensor(sf: torch.Tensor) -> torch.Tensor:
+        """Transform ``sf`` into an MN-major, TMA-aligned layout for DeepGEMM.
+
+        When ``sf`` is already in that layout, sgl-deep-gemm's fast path
+        (<= 0.1.4.post1) returns a NON-OWNING ``torch::from_blob`` alias of
+        ``sf`` across the TVM-FFI boundary. Callers rebind the result over
+        their only reference (``x = get_mn_major_tma_aligned_tensor(x)``),
+        which frees the storage while the GEMM still reads through the alias
+        -- a use-after-free that surfaces as NaN logits or "pointer resides
+        on host memory" during CUDA graph capture once the allocator reuses
+        the block. Hand back ``sf`` itself in that case so ownership is
+        preserved.
+        """
+        out = _get_mn_major_tma_aligned_tensor(sf)
+        if out.data_ptr() == sf.data_ptr():
+            assert out.shape == sf.shape and out.stride() == sf.stride()
+            return sf
+        return out
+
 
 _SANITY_CHECK = envs.SGLANG_DEEPGEMM_SANITY_CHECK.get()
 

test/registered/kernels/ops/quantization/test_per_token_group_quant.py

@@ -431,6 +431,82 @@ def test_masked_fused():
         assert torch.all(x_q[e, m:].view(torch.int8) == 0), "padding touched"
 
 
+@pytest.mark.parametrize("poison", [float("nan"), float("inf"), -float("inf")])
+@pytest.mark.parametrize("scale_ue8m0", [False, True])
+@pytest.mark.parametrize("masked", [False, True])
+def test_non_finite_inputs_are_sanitized(poison, scale_ue8m0, masked):
+    """CUDA-graph capture warmup runs the model on reused, uninitialized
+    buffers, so quant inputs can contain NaN/Inf bit patterns. The v1/v2/Triton
+    kernels clamp before converting (IEEE fminf/fmaxf drop the NaN operand),
+    quantizing non-finite values to +-fp8_max; emitting fp8 NaN codes instead
+    poisons the downstream GEMM and trips the sampler NaN check
+    (TestTBOWithTPAttn H100 CI). Pin the sanitizing behavior."""
+    torch.manual_seed(0)
+    if masked:
+        x = torch.randn(4, 32, 512, device="cuda", dtype=torch.bfloat16)
+        x[1, 3, 100] = poison
+        x[2, 0, 300] = poison
+        masked_m = torch.tensor([32, 16, 4, 0], device="cuda", dtype=torch.int32)
+    else:
+        x = torch.randn(16, 512, device="cuda", dtype=torch.bfloat16)
+        x[3, 100] = poison
+        x[7, 500] = poison
+        masked_m = None
+
+    x_q, x_s = per_token_group_quant(
+        x,
+        group_size=G,
+        scale_ue8m0=scale_ue8m0,
+        masked_m=masked_m,
+        column_major_scales=scale_ue8m0,
+    )
+    torch.cuda.synchronize()
+    if masked:
+        rows = [x_q[e, :m] for e, m in enumerate(masked_m.tolist()) if m > 0]
+        written = torch.cat([r.reshape(-1, x_q.shape[-1]) for r in rows])
+    else:
+        written = x_q
+    assert not torch.isnan(written.float()).any(), "quant emitted fp8 NaN codes"
+
+
+def test_mn_major_tma_aligned_transform_keeps_ownership():
+    """The masked fused quant emits scales already MN-major/TMA-aligned, which
+    makes deep_gemm's get_mn_major_tma_aligned_tensor hit its no-op fast path.
+    In sgl-deep-gemm <= 0.1.4.post1 that path returns a non-owning
+    ``torch::from_blob`` alias across TVM-FFI; the production caller rebinds
+    the result over its only reference, so an alias frees the scale storage
+    before the down GEMM reads it (NaN logits / host-pointer crash under CUDA
+    graph capture). The wrapper must hand back the input tensor itself."""
+    from sglang.srt.layers import deep_gemm_wrapper
+
+    if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
+        pytest.skip("deep_gemm unavailable")
+
+    E, N, num_groups = 18, 128, 20  # N 4-aligned -> already-TMA-aligned layout
+    s = torch.empty((E, num_groups, N), device="cuda", dtype=torch.float32)
+    s = s.transpose(-1, -2)
+    s.copy_(torch.rand(E, N, num_groups, device="cuda") + 1.0)
+    expected = s.clone()
+
+    out = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(s)
+    assert out is s, "already-aligned input must be returned as-is (owning)"
+
+    # The production pattern: rebind + allocator churn between quant and GEMM.
+    del s
+    reuse = torch.full((E, num_groups, N), float("nan"), device="cuda")
+    torch.cuda.synchronize()
+    assert not torch.isnan(out).any(), "scale storage was freed and reused"
+    torch.testing.assert_close(out, expected)
+    del reuse
+
+    # Row-major input still takes the real transform into an owning buffer.
+    row_major = expected.contiguous()
+    out2 = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(row_major)
+    assert out2.data_ptr() != row_major.data_ptr()
+    assert out2.stride(-2) == 1
+    torch.testing.assert_close(out2, row_major)
+
+
 @pytest.mark.parametrize("out_dtype,column_major_scales,scale_ue8m0", AUTO_ALLOC_CASES)
 def test_auto_allocation(out_dtype, column_major_scales, scale_ue8m0):
     """Omitting output_q/output_s allocates them per out_dtype / major mode /