diff --git a/python/sglang/kernels/aot/csrc/elementwise/dsv4_norm_rope.cu b/python/sglang/kernels/aot/csrc/elementwise/dsv4_norm_rope.cu index 138daf094..a69802c4c 100644 --- a/python/sglang/kernels/aot/csrc/elementwise/dsv4_norm_rope.cu +++ b/python/sglang/kernels/aot/csrc/elementwise/dsv4_norm_rope.cu @@ -25,6 +25,7 @@ limitations under the License. #else #include #include +#include #include #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) { diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_cvt.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_cvt.cuh new file mode 100644 index 000000000..3bdd9c310 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_cvt.cuh @@ -0,0 +1,47 @@ +#include +#include + +#include + +#include + +#include + +#include +#include + +// 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(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(); + + TensorMatcher({N}).with_strides({1}).with_dtype().with_device(device_).verify(src); + TensorMatcher({N}).with_strides({1}).with_dtype().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(dst.data_ptr()), static_cast(src.data_ptr()), num_pairs); +} + +} // namespace sglang diff --git a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp8_utils.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp8_utils.cuh index 2d1972b2a..7b600211c 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp8_utils.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/fp8_utils.cuh @@ -7,6 +7,13 @@ #include #ifndef USE_ROCM #include +#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 +#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(x8) | (static_cast(y8) << 8); } +#endif // HIP_FP8_TYPE_OCP && !HIP_FP8_TYPE_FNUZ #endif } // namespace deepseek_v4::fp8 diff --git a/python/sglang/kernels/ops/attention/dsv4/fp8_cvt.py b/python/sglang/kernels/ops/attention/dsv4/fp8_cvt.py new file mode 100644 index 000000000..9ff05c89d --- /dev/null +++ b/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 diff --git a/test/registered/e2e/dsv4/test_dsv4_fp8_cast.py b/test/registered/e2e/dsv4/test_dsv4_fp8_cast.py new file mode 100644 index 000000000..42c5095e3 --- /dev/null +++ b/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() diff --git a/test/registered/kernels/ops/attention/test_dsv4_indexer_quant.py b/test/registered/kernels/ops/attention/test_dsv4_indexer_quant.py index f7ba90917..fe2d562b9 100644 --- a/test/registered/kernels/ops/attention/test_dsv4_indexer_quant.py +++ b/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