From d4a0dfbc31ab9310e6ec7f203a8a2870d331dca6 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 24 Jul 2026 07:29:28 +0800 Subject: [PATCH] [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 --- .../jit/csrc/gemm/per_token_group_quant.cuh | 17 ++++- .../layers/deep_gemm_wrapper/entrypoint.py | 24 +++++- .../test_per_token_group_quant.py | 76 +++++++++++++++++++ 3 files changed, 113 insertions(+), 4 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh index ce4ded742..deb1cc679 100644 --- a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh +++ b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh @@ -47,9 +47,15 @@ template <> struct WeightTrait { using packed2_t = fp8x2_e4m3_t; static constexpr float kMaxValue = DTypeTrait::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(exp); const float quant_scale = inv_scale_ue8m0(exp); const auto scale2 = cast(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); + // __hmin2/__hmax2 return the non-NaN operand. + const auto lo2 = cast(float2{-kMaxValue, -kMaxValue}); + const auto hi2 = cast(float2{kMaxValue, kMaxValue}); #pragma unroll for (uint32_t i = 0; i < kVecSize / 2; ++i) { - out[i] = static_cast(__hmul2(in[i], scale2)); + out[i] = static_cast(__hmin2(__hmax2(__hmul2(in[i], scale2), lo2), hi2)); } } else { // fp32 scale: multiply in fp32 (hmul2 brings too much precision loss) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index c018217bb..6151b3cd7 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/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() diff --git a/test/registered/kernels/ops/quantization/test_per_token_group_quant.py b/test/registered/kernels/ops/quantization/test_per_token_group_quant.py index 576052f6c..cdcc00086 100644 --- a/test/registered/kernels/ops/quantization/test_per_token_group_quant.py +++ b/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 /