[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.
-// Supports both E4M3FN (MI350X, gfx950) and E4M3FNUZ (MI300X, gfx942).
+#ifdef SGL_ROCM_FP8_HW_CVT
+// 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