[AMD] fix: use the hardware fp8 e4m3 convert on gfx950 (#37140)

Signed-off-by: amd-danli103 <danli103@amd.com>

python/sglang/kernels/aot/csrc/elementwise/dsv4_norm_rope.cu

@@ -25,6 +25,7 @@ limitations under the License.
 #else
 #include <hip/hip_bf16.h>
 #include <hip/hip_fp16.h>
+#include <hip/hip_fp8.h>
 #include <hip/hip_runtime.h>
 #endif
 
@@ -128,6 +129,17 @@ __device__ __forceinline__ fp8x2_e4m3_t pack_fp8(float x, float y) {
   y = fmaxf(fminf(y, kFP8Max), -kFP8Max);
   return __nv_fp8x2_e4m3(float2{x, y});
 }
+#elif HIP_FP8_TYPE_OCP && !HIP_FP8_TYPE_FNUZ
+// gfx950/gfx12xx write OCP e4m3 natively, so take v_cvt_pk_fp8_f32 -- RNE, both lanes in
+// one instruction. Not gfx942: hardware only converts to fnuz there, and this kernel writes
+// E4M3FN on every arch (the indexer caller allocates float8_e4m3fn), so gfx942 keeps the
+// software cast below. Testing FNUZ too because HIP sets both macros on the host pass and
+// on targets outside its list. Clip rather than ask for __HIP_SATFINITE: the x2 fast path
+// converts the value it was handed, not the clamped one (ROCm 7.2).
+__device__ __forceinline__ fp8x2_e4m3_t pack_fp8(float x, float y) {
+  const float2 v{fmaxf(fminf(x, kFP8Max), -kFP8Max), fmaxf(fminf(y, kFP8Max), -kFP8Max)};
+  return __hip_cvt_float2_to_fp8x2(v, __HIP_NOSAT, __HIP_E4M3);
+}
 #else
 // Software float -> FP8 E4M3 conversion for ROCm
 __device__ __forceinline__ uint8_t cvt_float_to_fp8_e4m3(float val) {

python/sglang/kernels/jit/csrc/deepseek_v4/fp8_cvt.cuh

@@ -0,0 +1,47 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+// Elementwise float -> fp8 e4m3 through the same `pack_fp8` every fp8 store in the
+// DSv4 tree goes through. It is here so that cast can be pinned against torch on its
+// own: a wrong rounding or saturation boundary in there does not show up as a failure
+// in the fused kernels -- it just looks like fp8 quantizing worse than it should.
+
+namespace sglang {
+
+constexpr size_t kCvtBlockSize = 256;
+
+__global__ void cvt_fp8_e4m3_kernel(uint8_t* dst, const float* src, size_t num_pairs) {
+  const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= num_pairs) return;
+  reinterpret_cast<fp8x2_e4m3_t*>(dst)[idx] = deepseek_v4::fp8::pack_fp8(src[2 * idx], src[2 * idx + 1]);
+}
+
+void cvt_fp8_e4m3(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
+  using namespace host;
+
+  auto N = SymbolicSize{"num_elements"};
+  auto device_ = SymbolicDevice{};
+  device_.set_options<kDLGPU>();
+
+  TensorMatcher({N}).with_strides({1}).with_dtype<float>().with_device(device_).verify(src);
+  TensorMatcher({N}).with_strides({1}).with_dtype<uint8_t>().with_device(device_).verify(dst);
+
+  const size_t num_elements = N.unwrap();
+  // pack_fp8 converts two values at a time
+  RuntimeCheck(num_elements > 0 && num_elements % 2 == 0, "num_elements must be even and non-zero, got ", num_elements);
+
+  const size_t num_pairs = num_elements / 2;
+  LaunchKernel(div_ceil(num_pairs, kCvtBlockSize), kCvtBlockSize, device_.unwrap())(
+      cvt_fp8_e4m3_kernel, static_cast<uint8_t*>(dst.data_ptr()), static_cast<const float*>(src.data_ptr()), num_pairs);
+}
+
+}  // namespace sglang

python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp8_utils.cuh

@@ -7,6 +7,13 @@
 #include <cstdint>
 #ifndef USE_ROCM
 #include <cuda_fp8.h>
+#elif defined(__gfx950__) || defined(__gfx1200__) || defined(__gfx1201__)
+// Only on the arches that take the hardware branch below. hip_fp8.h is what defines
+// HIP_FP8_TYPE_FNUZ, and nothing else in this include tree pulls it in, so gating on
+// those macros instead would also flip the software cast's arch constants on gfx942 --
+// it picks fn today because the macro is not visible there.
+#include <hip/hip_fp8.h>
+#define SGL_ROCM_FP8_HW_CVT 1
 #endif
 
 // Small helpers shared by the DeepSeek-V4 FP8/UE8M0 quantization kernels
@@ -45,8 +52,21 @@ SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) {
   return fp8x2_e4m3_t{fp32x2_t{fp8_e4m3_clip(x), fp8_e4m3_clip(y)}};
 }
 #else
-// Software float -> FP8 E4M3 conversion for ROCm/HIP.
+#ifdef SGL_ROCM_FP8_HW_CVT
-// Supports both E4M3FN (MI350X, gfx950) and E4M3FNUZ (MI300X, gfx942).
+// gfx950/gfx12xx do both lanes in one v_cvt_pk_fp8_f32 (RNE), and the flavour it produces
+// is the OCP one kFP8E4M3Max already assumes there. Clip first rather than passing
+// __HIP_SATFINITE -- the x2 fast path converts the value it was handed, not the clamped
+// one (ROCm 7.2).
+//
+// gfx942 keeps the software cast below, top-segment bug and all -- this instruction does
+// not produce the fnuz flavour that arch needs, so it takes a separate fix.
+SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) {
+  const fp32x2_t v{fp8_e4m3_clip(x), fp8_e4m3_clip(y)};
+  return __hip_cvt_float2_to_fp8x2(v, __HIP_NOSAT, __HIP_E4M3);
+}
+#else
+// Software float -> FP8 E4M3 conversion for the archs the branch above skips: gfx942,
+// plus any target with no native fp8 convert.
 SGL_DEVICE uint8_t cvt_float_to_fp8_e4m3(float val) {
   val = fp8_e4m3_clip(val);
   if (val == 0.0f) return 0;
@@ -117,6 +137,7 @@ SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) {
   uint8_t y8 = cvt_float_to_fp8_e4m3(y);
   return static_cast<uint16_t>(x8) | (static_cast<uint16_t>(y8) << 8);
 }
+#endif  // HIP_FP8_TYPE_OCP && !HIP_FP8_TYPE_FNUZ
 #endif
 
 }  // namespace deepseek_v4::fp8

python/sglang/kernels/ops/attention/dsv4/fp8_cvt.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import cache_once, load_jit
+
+from .utils import make_name
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_fp8_cvt_module() -> Module:
+    return load_jit(
+        make_name("fp8_cvt"),
+        cuda_files=["deepseek_v4/fp8_cvt.cuh"],
+        cuda_wrappers=[("cvt_fp8_e4m3", "cvt_fp8_e4m3")],
+    )
+
+
+def cvt_fp8_e4m3(src: torch.Tensor) -> torch.Tensor:
+    """Cast fp32 to fp8 e4m3 through the same ``pack_fp8`` the fp8 stores use.
+
+    Nothing in the serving path calls this -- it is here so the conversion can be
+    compared against torch on its own. Inside the fused kernels every value goes
+    through a quantization scale first, which turns a wrong conversion byte into
+    "fp8 is a bit lossy" rather than a failure.
+
+    Args:
+        src: contiguous 1D fp32 CUDA/HIP tensor of even length -- the conversion runs
+            two values at a time.
+
+    Returns:
+        uint8 tensor of the same length holding the raw e4m3 bytes.
+
+    Raises:
+        RuntimeError: if the length is zero or odd, or a tensor does not match.
+    """
+    dst = torch.empty_like(src, dtype=torch.uint8)
+    _jit_fp8_cvt_module().cvt_fp8_e4m3(dst, src)
+    return dst

test/registered/e2e/dsv4/test_dsv4_fp8_cast.py

@@ -0,0 +1,116 @@
+"""The DSv4 fp8 e4m3 conversion on its own, byte for byte against torch.
+
+Every fp8 store in the DSv4 tree quantizes through ``pack_fp8``, and none of the
+callers can see when it is wrong: the value has already been divided by a
+quantization scale, so a bad rounding or saturation boundary comes back as fp8 being
+lossier than it should be, not as a failure. ``pack_fp8`` on ROCm used to be a
+hand-written bit twiddle, and it had two: the whole top exponent segment saturated
+to the max normal, and the binade under the min subnormal flushed to zero instead of
+rounding up to it. Both are pinned below.
+
+gfx950 only. gfx942 still runs the software cast, bugs and all, because the
+instruction cannot produce the fnuz bytes that arch writes; that one needs its own fix.
+The reference is torch's own cast.
+"""
+
+import math
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4.fp8_cvt import cvt_fp8_e4m3
+from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype, fp8_max
+from sglang.srt.utils import is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+# the default amd runner is mi300, where pack_fp8 still takes the software path this
+# does not cover
+register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+DEVICE = torch.device("cuda")
+
+# start of the top exponent segment, i.e. the largest power of two the format holds
+TOP_BINADE = 2.0 ** math.floor(math.log2(fp8_max))
+
+
+def _representable():
+    vals = torch.arange(256, dtype=torch.uint8).view(fp8_dtype).float()
+    return vals[vals.isfinite()]
+
+
+def _domain():
+    """every representable value, every midpoint between two of them, every bf16"""
+    vals = _representable()
+    mids = ((vals[:, None] + vals[None, :]) / 2).flatten()
+    cases = torch.cat(
+        [
+            vals,
+            mids,
+            torch.linspace(-fp8_max, fp8_max, 100003),
+            torch.arange(1 << 16, dtype=torch.int32).view(torch.bfloat16).float(),
+        ]
+    )
+    # stay inside the range: past the max the two casts are allowed to disagree on
+    # whether to clamp or produce NaN, which is not what this is testing
+    cases = cases[cases.isfinite() & (cases.abs() <= fp8_max)].unique()
+    if cases.numel() % 2:
+        cases = cases[:-1]
+    return cases.contiguous()
+
+
+def _as_bytes(x):
+    # the conversion runs two values at a time, so the length has to stay even --
+    # e4m3fnuz has an odd number of representable values (only 0x80 is NaN)
+    if x.numel() % 2:
+        x = x[:-1]
+    x = x.contiguous().to(DEVICE)
+    return cvt_fp8_e4m3(x), x.to(fp8_dtype).view(torch.uint8)
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
+    "the gfx95 path of pack_fp8 is what this pins",
+)
+class TestDsv4Fp8Cast(CustomTestCase):
+    def test_matches_torch_over_the_whole_range(self):
+        cases = _domain()
+        got, want = _as_bytes(cases)
+        bad = got != want
+        if bool(bad.any()):
+            v = cases.to(DEVICE)[bad]
+            worst = v.abs().argmax()
+            self.fail(
+                f"{int(bad.sum())} of {cases.numel()} bytes differ, "
+                f"|v| in [{v.abs().min():.4e}, {v.abs().max():.4e}]; e.g. "
+                f"{v[worst]:.6g} -> {got[bad][worst].item():#04x} "
+                f"(torch {want[bad][worst].item():#04x})"
+            )
+
+    def test_top_exponent_segment_is_not_saturated(self):
+        # testing the exponent alone here used to write every value from TOP_BINADE up
+        # to the max out as the max normal
+        vals = _representable()
+        # pack_fp8 clips to fp8_max, so anything the format holds above it never comes
+        # out of the conversion
+        top = vals[(vals.abs() >= TOP_BINADE) & (vals.abs() <= fp8_max)]
+        self.assertGreater(top.numel(), 2)
+        got, want = _as_bytes(top)
+        self.assertTrue(torch.equal(got, want))
+        # and they really are distinct values, not all the same byte
+        self.assertGreater(int(got.unique().numel()), 2)
+
+    def test_binade_below_the_min_subnormal_rounds_up(self):
+        min_subnormal = _representable().abs()
+        min_subnormal = min_subnormal[min_subnormal > 0].min().item()
+        # (midpoint, min subnormal): rounds up. the midpoint itself is a tie and goes
+        # to even, i.e. to zero
+        band = torch.linspace(min_subnormal / 2, min_subnormal, 2049)[1:-1]
+        band = torch.cat([band, -band])
+        got, want = _as_bytes(band.contiguous())
+        self.assertTrue(torch.equal(got, want))
+        self.assertTrue(bool((got & 0x7F).all()))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_dsv4_indexer_quant.py

@@ -17,19 +17,22 @@ branches introduced by the scheduling optimization.
 from __future__ import annotations
 
 import pytest
+import sgl_kernel  # noqa: F401  the ROCm path dispatches to torch.ops.sgl_kernel
 import torch
 
 from sglang.kernels.ops.attention.dsv4 import (
     fused_q_indexer_rope_first_quant,
     fused_q_indexer_rope_hadamard_quant,
 )
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import is_gfx95_supported, is_hip
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 _is_hip = is_hip()
 
 register_cuda_ci(est_time=13, stage="base-b", runner_config="1-gpu-large")
-register_amd_ci(est_time=45, suite="jit-kernel-unit-test-amd")
+# the mi35x suite rather than the default AMD one: the ROCm case below is gfx95-only, and
+# everything else in here skips on HIP, so the mi300 registration only ever produced skips
+register_amd_ci(est_time=45, suite="stage-b-test-1-gpu-small-amd-mi35x")
 
 HEAD_DIM = 128
 ROPE_DIM = 64
@@ -45,10 +48,10 @@ N_HEADS = 64
 BATCHES = [1, 8, 64, 256, 512, 2048]
 
 
-def _skip_if_unavailable():
+def _skip_if_unavailable(hip_ok=False):
     if not torch.cuda.is_available():
         pytest.skip("CUDA required")
-    if _is_hip:
+    if _is_hip and not hip_ok:
         pytest.skip("Indexer fused Q kernel is CUDA-specific")
 
 
@@ -73,7 +76,14 @@ def _fp8_dequant_ok(q_fp8, ref, scale):
 @pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
 @pytest.mark.parametrize("batch", BATCHES)
 def test_v4_rope_hadamard_quant_matches_reference(batch, pos_dtype):
-    _skip_if_unavailable()
+    # runs on gfx95 too: elementwise.py routes this one to the AOT op there, and that op
+    # carries its own copy of the cast, so this is the only coverage it gets
+    _skip_if_unavailable(hip_ok=True)
+    if _is_hip:
+        if not is_gfx95_supported():
+            pytest.skip("gfx942 keeps the software cast in the AOT copy")
+        if pos_dtype is torch.int64:
+            pytest.skip("the ROCm AOT op takes int32 positions only")
     dev = "cuda"
     g = torch.Generator(device=dev).manual_seed(0)
     q = torch.randn(
@@ -157,6 +167,9 @@ def test_v32_rope_first_quant_matches_reference(batch):
 # Strided weight (the non-contiguous wk slice) matches contiguous (V4 path).
 # ----------------------------------------------------------------------------
 def test_v4_strided_weight_matches_contiguous():
+    # stays CUDA-only. the ROCm op reads the weight linearly, so a non-contiguous slice
+    # comes out wrong there -- unrelated to the cast, and latent, since the indexer hands
+    # it the contiguous weights_proj output
     _skip_if_unavailable()
     dev = "cuda"
     B = 512  # grid-stride regime