[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()