From f3225bceb3670d4d701d19882d020ba57e5fd1b0 Mon Sep 17 00:00:00 2001 From: Kaixi Date: Mon, 17 Aug 2026 07:06:41 +0200 Subject: [PATCH] [DSV4] Emit TMA-aligned UE8M0 scales for FP8 einsum (#34277) --- .../fp8_wo_a_group_major_quant.cuh | 61 +++++++++++++++---- .../kernels/ops/attention/dsv4/fp8_wo_a.py | 30 ++++++--- .../kernels/ops/attention/test_fp8_wo_a.py | 58 +++++++++++++----- 3 files changed, 113 insertions(+), 36 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh index de1c781de..152fe8ed0 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh @@ -2,9 +2,10 @@ // // This is intentionally narrower than the generic per_token_group_quant_8bit_v2 // kernel: input is a [T, G, D] view with contiguous hidden groups, output_q is -// contiguous [T, G, D], group_size is fixed to 128, scales are fp32 UE8M0 -// power-of-two values, and output_s is a logical [T, G, D/128] view backed by -// group-major [G, T, D/128] storage. +// contiguous [T, G, D], group_size is fixed to 128, and output_s stores four +// packed UE8M0 exponent bytes per int32. Its physical layout is +// [G, ceil((D/128)/4), align_up(T, 4)]; the Python wrapper returns the logical +// [T, G, ceil((D/128)/4)] view consumed natively by DeepGEMM. // // The generic kernel cannot read the strided DSV4 view while producing // contiguous [T, G, D] codes and group-major scales without an extra full-tensor @@ -34,6 +35,8 @@ constexpr uint32_t THREADS_PER_GROUP = 8; constexpr uint32_t SUBWARPS_PER_BLOCK = 16; constexpr uint32_t INPUT_VEC_NUM_BYTES = 32; constexpr uint32_t INPUT_INT4_SIZE = INPUT_VEC_NUM_BYTES / sizeof(int4); +constexpr int UE8M0_SCALES_PER_PACK = 4; +static_assert(UE8M0_SCALES_PER_PACK == sizeof(int32_t)); template SGL_DEVICE float GroupReduceMax(float val) { @@ -51,11 +54,12 @@ template __global__ void fp8_wo_a_group_major_quant_ue8m0_kernel( const T* __restrict__ input, fp8_e4m3_t* __restrict__ output_q, - float* __restrict__ output_s, + int32_t* __restrict__ output_s, int64_t total_scale_groups, - int64_t num_tokens, int hidden_dim_groups, + int packed_hidden_dim_groups, int num_outer_groups, + int64_t aligned_num_tokens, int64_t input_stride_t) { device::PDLWaitPrimary(); @@ -95,7 +99,6 @@ __global__ void fp8_wo_a_group_major_quant_ue8m0_kernel( constexpr float kFp8MaxInv = 1.0f / kFP8E4M3Max; const int32_t scale_ue8m0 = cast_to_ue8m0(local_absmax * kFp8MaxInv); const float y_scale = inv_scale_ue8m0(scale_ue8m0); - const float y_scale_inv = __uint_as_float(static_cast(scale_ue8m0) << 23); int4 output_buf; auto* output_buf_ptr = reinterpret_cast(&output_buf); @@ -108,7 +111,23 @@ __global__ void fp8_wo_a_group_major_quant_ue8m0_kernel( *reinterpret_cast(output_q + output_group_start_offset + lane_id * INPUT_VEC_SIZE) = output_buf; if (lane_id == 0) { - output_s[(outer_idx * num_tokens + token_idx) * hidden_dim_groups + hidden_group] = y_scale_inv; + const int hidden_pack_idx = hidden_group / UE8M0_SCALES_PER_PACK; + const int pack_byte_idx = hidden_group % UE8M0_SCALES_PER_PACK; + const int64_t scale_word_offset = + (static_cast(outer_idx) * packed_hidden_dim_groups + hidden_pack_idx) * aligned_num_tokens + + token_idx; + auto* scale_output = + reinterpret_cast(output_s) + scale_word_offset * UE8M0_SCALES_PER_PACK + pack_byte_idx; + *scale_output = static_cast(scale_ue8m0); + + // DeepGEMM consumes complete int32 packs. Zero bytes without a matching + // hidden group so allocator garbage cannot become an activation scale. + if (hidden_group == hidden_dim_groups - 1) { +#pragma unroll + for (int byte_idx = pack_byte_idx + 1; byte_idx < UE8M0_SCALES_PER_PACK; ++byte_idx) { + scale_output[byte_idx - pack_byte_idx] = 0; + } + } } } @@ -125,24 +144,39 @@ struct FP8WoAGroupMajorQuantUE8M0Kernel { auto TSize = SymbolicSize{"num_tokens"}; auto GSize = SymbolicSize{"num_outer_groups"}; auto DSize = SymbolicSize{"hidden_dim"}; - auto SSize = SymbolicSize{"hidden_dim_groups"}; + auto PSize = SymbolicSize{"packed_hidden_dim_groups"}; + auto ASize = SymbolicSize{"aligned_num_tokens"}; TensorMatcher({TSize, GSize, DSize}).with_strides({-1, DSize, 1}).with_dtype().with_device(device).verify(input); TensorMatcher({TSize, GSize, DSize}).with_dtype().with_device(device).verify(output_q); - TensorMatcher({GSize, TSize, SSize}).with_dtype().with_device(device).verify(output_s); + TensorMatcher({GSize, PSize, ASize}).with_dtype().with_device(device).verify(output_s); const auto num_tokens = TSize.unwrap(); const auto num_outer_groups = GSize.unwrap(); const auto hidden_dim = DSize.unwrap(); - const auto hidden_dim_groups = SSize.unwrap(); + const auto packed_hidden_dim_groups = PSize.unwrap(); + const auto aligned_num_tokens = ASize.unwrap(); const auto input_stride_t = input.stride(0); constexpr int64_t kInputAlignElements = sizeof(int4) / sizeof(T); RuntimeCheck(hidden_dim % GROUP_SIZE == 0, "hidden_dim must be divisible by 128"); - RuntimeCheck(hidden_dim_groups == hidden_dim / GROUP_SIZE, "output_s hidden dim mismatch"); + const auto hidden_dim_groups = hidden_dim / GROUP_SIZE; + RuntimeCheck( + packed_hidden_dim_groups == (hidden_dim_groups + UE8M0_SCALES_PER_PACK - 1) / UE8M0_SCALES_PER_PACK, + "packed output_s hidden dim mismatch"); + RuntimeCheck( + aligned_num_tokens == (num_tokens + UE8M0_SCALES_PER_PACK - 1) / UE8M0_SCALES_PER_PACK * UE8M0_SCALES_PER_PACK, + "output_s token dim must be TMA-aligned"); + RuntimeCheck( + output_s.stride(0) == packed_hidden_dim_groups * aligned_num_tokens && + output_s.stride(1) == aligned_num_tokens && output_s.stride(2) == 1, + "output_s must use contiguous [group, packed_hidden, aligned_token] storage"); RuntimeCheck( reinterpret_cast(input.data_ptr()) % sizeof(int4) == 0, "input base pointer must be 16-byte aligned"); + RuntimeCheck( + reinterpret_cast(output_s.data_ptr()) % sizeof(int4) == 0, + "output_s base pointer must be 16-byte aligned"); RuntimeCheck( num_tokens <= 1 || input_stride_t % kInputAlignElements == 0, "input token stride must preserve 16-byte vector-load alignment"); @@ -157,11 +191,12 @@ struct FP8WoAGroupMajorQuantUE8M0Kernel { fp8_wo_a_group_major_quant_ue8m0_kernel, static_cast(input.data_ptr()), static_cast(output_q.data_ptr()), - static_cast(output_s.data_ptr()), + static_cast(output_s.data_ptr()), total_scale_groups, - static_cast(num_tokens), static_cast(hidden_dim_groups), + static_cast(packed_hidden_dim_groups), static_cast(num_outer_groups), + static_cast(aligned_num_tokens), static_cast(input_stride_t)); } }; diff --git a/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py b/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py index df4709fee..39cb66e85 100644 --- a/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py +++ b/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: from tvm_ffi.module import Module _GROUP_SIZE = 128 +_UE8M0_SCALES_PER_PACK = 4 @cache_once @@ -71,23 +72,34 @@ def sglang_per_token_group_quant_fp8_dsv4_wo_a( """Quantize DSV4 wo_a activations for DeepGEMM fp8_einsum. The input is a [T, G, D] bf16/fp16 tensor whose hidden dimension is - contiguous. The output codes are contiguous [T, G, D] fp8 values. The scale - tensor is returned as logical [T, G, D/128] fp32 UE8M0 values backed by - contiguous [G, T, D/128] storage, so each group/head [T, S] panel is - contiguous for the DeepGEMM recipe=(1, 1, 128) consumer. Group size is fixed - to 128 and the absmax floor is fixed to 1e-10. + contiguous. The output codes are contiguous [T, G, D] fp8 values. Four + UE8M0 scale exponent bytes are packed into each int32. The returned logical + scale tensor has shape [T, G, ceil((D/128)/4)] and DeepGEMM's native + TMA-aligned strides, so fp8_einsum can consume it without a scale-layout + conversion kernel. Group size is fixed to 128 and the absmax floor is fixed + to 1e-10. """ num_tokens, num_groups, hidden = x.shape hidden_groups = hidden // _GROUP_SIZE + packed_hidden_groups = ( + hidden_groups + _UE8M0_SCALES_PER_PACK - 1 + ) // _UE8M0_SCALES_PER_PACK + aligned_num_tokens = ( + (num_tokens + _UE8M0_SCALES_PER_PACK - 1) + // _UE8M0_SCALES_PER_PACK + * _UE8M0_SCALES_PER_PACK + ) x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn) x_s_storage = torch.empty( - (num_groups, num_tokens, hidden_groups), + (num_groups, packed_hidden_groups, aligned_num_tokens), device=x.device, - dtype=torch.float32, + dtype=torch.int32, ) if x.numel() > 0: fp8_wo_a_group_major_quant_ue8m0(x, x_q, x_s_storage) - # DeepGEMM fp8_einsum consumes each group/head [T, S] scale panel contiguously. - return x_q, x_s_storage.transpose(0, 1) + # DeepGEMM permutes this to [G, T, packed_hidden_groups], where tokens are + # contiguous and the packed-hidden stride is aligned_num_tokens. + x_s = x_s_storage.transpose(-1, -2)[:, :num_tokens, :].transpose(0, 1) + return x_q, x_s diff --git a/test/registered/kernels/ops/attention/test_fp8_wo_a.py b/test/registered/kernels/ops/attention/test_fp8_wo_a.py index 230833b84..456f88877 100644 --- a/test/registered/kernels/ops/attention/test_fp8_wo_a.py +++ b/test/registered/kernels/ops/attention/test_fp8_wo_a.py @@ -1,8 +1,8 @@ """DeepSeek-V4 wo_a FP8 activation quant for DeepGEMM fp8_einsum. -Covers the dedicated DSV4 wo_a quant helper: bit-exact FP8/scales against the -ordinary flat UE8M0 quant values, group-major scale storage, large / DSV4-shaped -token axes, and the DeepGEMM fp8_einsum consumer contract. +Covers the dedicated DSV4 wo_a quant helper: bit-exact FP8/packed scales against +the ordinary flat UE8M0 quant values, DeepGEMM's native TMA-aligned scale +layout, large / DSV4-shaped token axes, and the fp8_einsum consumer contract. """ import unittest @@ -58,26 +58,41 @@ class TestDeepSeekV4FP8WoA(CustomTestCase): torch.randn(T, G + 1, D, device=device, dtype=torch.float32) * 0.25 ).to(dtype) o = storage[:, 1:, :] - self.assertFalse(o.is_contiguous()) + if T > 1: + self.assertFalse(o.is_contiguous()) self.assertEqual(o.stride(-1), 1) return o def _assert_matches_flat_reference(self, o, o_fp8, o_s): T, G, D = o.shape q_ref, s_ref = self._flat_reference(o) + packed_ref = ( + self.deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor( + s_ref.transpose(0, 1).contiguous() + ).transpose(0, 1) + ) torch.cuda.synchronize() + hidden_groups = D // _GROUP_SIZE + packed_hidden_groups = (hidden_groups + 3) // 4 + aligned_num_tokens = (T + 3) // 4 * 4 self.assertEqual(o_fp8.shape, (T, G, D)) self.assertEqual(o_fp8.dtype, fp8_dtype) - self.assertEqual(o_s.shape, (T, G, D // _GROUP_SIZE)) - self.assertEqual(o_s.dtype, torch.float32) - self.assertEqual(o_s.stride(), (D // _GROUP_SIZE, T * (D // _GROUP_SIZE), 1)) - self.assertTrue(o_s[:, 0, :].is_contiguous()) + self.assertEqual(o_s.shape, (T, G, packed_hidden_groups)) + self.assertEqual(o_s.dtype, torch.int32) + self.assertEqual( + o_s.stride(), + ( + 1, + packed_hidden_groups * aligned_num_tokens, + aligned_num_tokens, + ), + ) self.assertTrue( torch.equal(o_fp8.view(torch.int8), q_ref.view(torch.int8)), "fp8 codes differ", ) - self.assertTrue(torch.equal(o_s, s_ref), "scales differ") + self.assertTrue(torch.equal(o_s, packed_ref), "packed scales differ") def test_dsv4_wo_a_quant_matches_flat_reference(self): torch.manual_seed(1) @@ -85,6 +100,7 @@ class TestDeepSeekV4FP8WoA(CustomTestCase): device = torch.device("cuda") for dtype, T, G, D in [ + (torch.bfloat16, 1, 2, 256), (torch.bfloat16, 9, 5, 384), (torch.float16, 7, 3, 512), ]: @@ -105,9 +121,8 @@ class TestDeepSeekV4FP8WoA(CustomTestCase): self.assertEqual(o_fp8.shape, o.shape) self.assertEqual(o_fp8.dtype, fp8_dtype) - self.assertEqual(o_s.shape, (0, 3, 2)) - self.assertEqual(o_s.dtype, torch.float32) - self.assertEqual(o_s.stride(), (2, 2, 1)) + self.assertEqual(o_s.shape, (0, 3, 1)) + self.assertEqual(o_s.dtype, torch.int32) def test_dsv4_wo_a_quant_large_token_dimension(self): torch.manual_seed(2) @@ -152,7 +167,7 @@ class TestDeepSeekV4FP8WoA(CustomTestCase): finally: fp8_wo_a_module._jit_module = original_jit_module - def test_fp8_wo_a_einsum_uses_group_major_activation_scales(self): + def test_fp8_wo_a_einsum_uses_tma_aligned_activation_scales(self): torch.manual_seed(0) torch.cuda.manual_seed_all(0) @@ -183,6 +198,8 @@ class TestDeepSeekV4FP8WoA(CustomTestCase): weight_s = transform_scale_ue8m0(weight_s_raw, mn=R) q_dsv4, s_dsv4 = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) + q_ref, s_ref = self._flat_reference(o) + s_ref_group_major = s_ref.transpose(0, 1).contiguous().transpose(0, 1) out = torch.empty(T, G, R, device=device, dtype=torch.bfloat16) self.deep_gemm.fp8_einsum( "bhr,hdr->bhd", @@ -191,11 +208,24 @@ class TestDeepSeekV4FP8WoA(CustomTestCase): out, recipe=(1, 1, _GROUP_SIZE), ) + out_ref = torch.empty_like(out) + self.deep_gemm.fp8_einsum( + "bhr,hdr->bhd", + (q_ref, s_ref_group_major), + (weight_fp8, weight_s), + out_ref, + recipe=(1, 1, _GROUP_SIZE), + ) torch.cuda.synchronize() + self.assertTrue( + torch.equal(out, out_ref), + "einsum outputs differ between packed and fp32 scales", + ) + o_dequant = q_dsv4.float().view( T, G, D // _GROUP_SIZE, _GROUP_SIZE - ) * s_dsv4.unsqueeze(-1) + ) * s_ref.unsqueeze(-1) weight_dequant = block_quant_dequant( weight_fp8, weight_s_raw,