1- diff --git a/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu b/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu
2- index 4fbaf08e..5bad403d 100644
3- --- a/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu
4- +++ b/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu
5- @@ -35,7 +35,9 @@ __device__ __forceinline__ float silu(const float& val) {
6- float t = __tanhf(half);
7- return half * (1.0f + t);
8- #else
9- - return val / (1.0f + __expf(-val));
10- + // Match sgl_kernel.silu_and_mul on SM90. That kernel uses expf and
11- + // exposes only the final SiLU*up BF16 rounding boundary.
12- + return val / (1.0f + expf(-val));
13- #endif
14- }
15-
16- @@ -266,7 +268,8 @@ __global__ void per_token_group_quant_8bit_kernel(
17- // TODO can this be removed?
18- const int scale_expert_stride,
19- const int scale_hidden_stride,
20- - const int num_tokens_per_expert) {
21- + const int num_tokens_per_expert,
22- + const float max_8bit) {
23- using dst_dtype_info = DtypeInfo<DST_DTYPE>;
24- using scale_element_t = std::conditional_t<SCALE_UE8M0, uint8_t, float>;
25- static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
26- @@ -347,8 +350,11 @@ __global__ void per_token_group_quant_8bit_kernel(
27- for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) {
28- float val;
29- if constexpr (FUSE_SILU_AND_MUL) {
30- - // TODO maybe vectorize
31- - T val_lowprec = static_cast<T>(silu(static_cast<float>(input_primary_vec[j]))) * input_secondary_vec[j];
32- + // Keep the same single visible low-precision boundary as the
33- + // contiguous activation path. Rounding SiLU before multiplying
34- + // changes the second MoE GEMM's FP8 inputs.
35- + T val_lowprec = static_cast<T>(
36- + silu(static_cast<float>(input_primary_vec[j])) * static_cast<float>(input_secondary_vec[j]));
37- val = static_cast<float>(val_lowprec);
38- input_primary_vec[j] = val_lowprec;
39- } else {
40- @@ -362,8 +368,15 @@ __global__ void per_token_group_quant_8bit_kernel(
41- local_absmax = GroupReduceMax<THREADS_PER_SUBWARP>(local_absmax, lane_id);
42-
43- float y_scale, y_scale_inv;
44- - calculate_fp8_scales<SCALE_UE8M0, dst_dtype_info>(local_absmax, y_scale, y_scale_inv);
45- - float2 y_scale_repeated = {y_scale, y_scale};
46- + if constexpr (SCALE_UE8M0) {
47- + calculate_fp8_scales<true, dst_dtype_info>(local_absmax, y_scale, y_scale_inv);
48- + } else {
49- + // Preserve the scale and division sequence of the legacy unmasked
50- + // quantizer. Reciprocal reuse under -use_fast_math moves a few
51- + // values across an FP8 rounding boundary.
52- + y_scale_inv = __fdiv_rn(local_absmax, max_8bit);
53- + y_scale = 1.0f / y_scale_inv;
54- + }
55-
56- if (lane_id == 0) {
57- *scale_output = extract_required_scale_format<SCALE_UE8M0>(y_scale_inv);
58- @@ -373,19 +386,33 @@ __global__ void per_token_group_quant_8bit_kernel(
59- static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE * sizeof(DST_DTYPE));
60-
61- if constexpr (std::is_same_v<DST_DTYPE, c10::Float8_e4m3fn>) {
62- - const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf);
63- - static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE / 2 * sizeof(__nv_fp8x2_storage_t));
64- - static_assert(INPUT_PRIMARY_VEC_SIZE % 2 == 0);
65- + if constexpr (SCALE_UE8M0) {
66- + const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf);
67- + static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE / 2 * sizeof(__nv_fp8x2_storage_t));
68- + static_assert(INPUT_PRIMARY_VEC_SIZE % 2 == 0);
69-
70- #pragma unroll
71- - for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 2) {
72- - float2 inputx2 = {static_cast<float>(input_primary_vec[j]), static_cast<float>(input_primary_vec[j + 1])};
73- - float2 outputx2 = fmul2_rn(inputx2, y_scale_repeated);
74- + for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 2) {
75- + float2 inputx2 = {
76- + static_cast<float>(input_primary_vec[j]), static_cast<float>(input_primary_vec[j + 1])};
77- + float2 outputx2 = fmul2_rn(inputx2, {y_scale, y_scale});
78-
79- - outputx2.x = fminf(fmaxf(outputx2.x, dst_dtype_info::MIN), dst_dtype_info::MAX);
80- - outputx2.y = fminf(fmaxf(outputx2.y, dst_dtype_info::MIN), dst_dtype_info::MAX);
81- + outputx2.x = fminf(fmaxf(outputx2.x, dst_dtype_info::MIN), dst_dtype_info::MAX);
82- + outputx2.y = fminf(fmaxf(outputx2.y, dst_dtype_info::MIN), dst_dtype_info::MAX);
83-
84- - output_buf_ptr[j / 2] = __nv_cvt_float2_to_fp8x2(outputx2, __NV_SATFINITE, __NV_E4M3);
85- + output_buf_ptr[j / 2] = __nv_cvt_float2_to_fp8x2(outputx2, __NV_SATFINITE, __NV_E4M3);
86- + }
87- + } else {
88- + // Match the legacy SM90 quantizer's scalar conversion exactly;
89- + // packed conversion makes a different tie choice for rare values.
90- + const auto output_buf_ptr = reinterpret_cast<DST_DTYPE*>(&output_buf);
91- + #pragma unroll
92- + for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) {
93- + float val = static_cast<float>(input_primary_vec[j]);
94- + float q_val =
95- + fminf(fmaxf(__fdiv_rn(val, y_scale_inv), dst_dtype_info::MIN), dst_dtype_info::MAX);
96- + output_buf_ptr[j] = DST_DTYPE(q_val);
97- + }
98- }
99- } else {
100- const auto output_buf_ptr = reinterpret_cast<DST_DTYPE*>(&output_buf);
101- @@ -474,7 +501,8 @@ void sgl_per_token_group_quant_8bit_v2(
102- hidden_dim_num_groups, \
103- scale_expert_stride, \
104- scale_hidden_stride, \
105- - num_tokens_per_expert); \
106- + num_tokens_per_expert, \
107- + static_cast<float>(max_8bit)); \
108- } while (0)
109-
110- #define LAUNCH_KERNEL(GROUP_SIZE, T, DST_DTYPE) \
111- @@ -497,13 +525,51 @@ void sgl_per_token_group_quant_8bit_v2(
112- NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, true); \
113- } \
114- } else { \
115- - LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true); \
116- + if (masked_layout) { \
117- + LAUNCH_KERNEL_INNER( \
118- + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, false); \
119- + } else { \
120- + LAUNCH_KERNEL_INNER( \
121- + NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, false); \
122- + } \
123- } \
124- } else { \
125- - LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true); \
126- + if (fuse_silu_and_mul) { \
127- + if (masked_layout) { \
128- + LAUNCH_KERNEL_INNER( \
129- + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \
130- + } else { \
131- + LAUNCH_KERNEL_INNER( \
132- + NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \
133- + } \
134- + } else { \
135- + if (masked_layout) { \
136- + LAUNCH_KERNEL_INNER( \
137- + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, false); \
138- + } else { \
139- + LAUNCH_KERNEL_INNER( \
140- + NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, false); \
141- + } \
142- + } \
143- } \
144- } else { \
145- - LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false); \
146- + if (fuse_silu_and_mul) { \
147- + if (masked_layout) { \
148- + LAUNCH_KERNEL_INNER( \
149- + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \
150- + } else { \
151- + LAUNCH_KERNEL_INNER( \
152- + NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \
153- + } \
154- + } else { \
155- + if (masked_layout) { \
156- + LAUNCH_KERNEL_INNER( \
157- + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, false); \
158- + } else { \
159- + LAUNCH_KERNEL_INNER( \
160- + NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, false); \
161- + } \
162- + } \
163- } \
164- } while (0)
165-
1661diff --git a/python/sglang/kernels/aot/tests/test_per_token_group_quant_8bit.py b/python/sglang/kernels/aot/tests/test_per_token_group_quant_8bit.py
1672index 0432f4ec..e7db42a2 100644
1683--- a/python/sglang/kernels/aot/tests/test_per_token_group_quant_8bit.py
@@ -171,16 +6,16 @@ index 0432f4ec..e7db42a2 100644
1716
1727 import pytest
1738 import torch
174- + from sgl_kernel import sgl_per_token_group_quant_8bit, silu_and_mul
9+ + from sgl_kernel import silu_and_mul
17510 from sgl_kernel.test_utils import (
17611 assert_all_close_or_tiny_diff,
17712 create_per_token_group_quant_test_data,
178- @@ -19,6 +20,94 @@ from sglang.srt.utils import is_hip
13+ @@ -19,6 +20,91 @@ from sglang.srt.utils import is_hip
17914 _is_hip = is_hip()
18015 fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
18116
18217+
183- + @pytest.mark.parametrize("backend", ["jit", "jit_v2", "aot" ])
18+ + @pytest.mark.parametrize("backend", ["jit", "jit_v2"])
18419+ def test_masked_fused_silu_quant_matches_contiguous_sm90(backend):
18520+ if torch.cuda.get_device_capability() != (9, 0):
18621+ pytest.skip("This test covers the Hopper train-rollout alignment path")
@@ -223,12 +58,9 @@ index 0432f4ec..e7db42a2 100644
22358+ device="cuda",
22459+ dtype=torch.float32,
22560+ ).transpose(-1, -2)
226- + if backend == "jit_v2":
227- + from sglang.kernels.ops.quantization.per_token_group_quant_8bit_v2 import (
228- + per_token_group_quant_8bit_v2 as quantize,
229- + )
230- + else:
231- + quantize = sgl_per_token_group_quant_8bit
61+ + from sglang.kernels.ops.quantization.per_token_group_quant_8bit_v2 import (
62+ + per_token_group_quant_8bit_v2 as quantize,
63+ + )
23264+ quantize(
23365+ gate_up,
23466+ fused_q,
0 commit comments