From 17dbd429306d2fe9cc6be04e57aab314bd697f5e Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 19 Aug 2026 10:20:05 -0500 Subject: [PATCH] [ROCm] Add UE8M0 scale packing for Triton silu_mul_quant (#37835) Signed-off-by: Andreas Karatzas --- .../moe/test_silu_mul_fp8_quant_deep_gemm.py | 6 +-- .../experts/batched_deep_gemm_moe.py | 51 ++++++++++++++++--- 2 files changed, 44 insertions(+), 13 deletions(-) diff --git a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py index e9a246f51b..85b992ad65 100644 --- a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py +++ b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py @@ -239,12 +239,8 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt scale_fmts = [ DeepGemmQuantScaleFMT.FLOAT32, DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0, + DeepGemmQuantScaleFMT.UE8M0, ] - # UE8M0 (int32 packed) scales require the C++ kernel which is - # not available on ROCm (#ifndef USE_ROCM). - # https://github.com/ROCm/aiter/issues/2420 - if current_platform.is_cuda(): - scale_fmts.append(DeepGemmQuantScaleFMT.UE8M0) # Run the SiLU V2 kernel for scale_fmt in scale_fmts: diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index af9d170fcb..c578fb641b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -220,6 +220,22 @@ def persistent_masked_m_silu_mul_quant( # Triton fallback for ROCm and XPU -- the C++ kernel is guarded by # #ifndef USE_ROCM in activation_kernels.cu. # https://github.com/ROCm/aiter/issues/2420 + # For UE8M0 (int32 packed scales), compute with float32 scales + # first, then pack afterwards. + is_packed_ue8m0 = quant_scale_fmt == DeepGemmQuantScaleFMT.UE8M0 + if is_packed_ue8m0: + f32_shape, f32_strides, _ = scales_shape_stride_dtype( + E, T, G, DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 + ) + y_s_f32 = torch.empty_strided( + f32_shape, + f32_strides, + dtype=torch.float32, + device=y.device, + ) + else: + y_s_f32 = y_s + stride_cnt_e = tokens_per_expert.stride()[0] # Static grid over experts and H-groups. @@ -231,14 +247,16 @@ def persistent_masked_m_silu_mul_quant( fp8_min, fp8_max = get_fp8_min_max() eps: float = 1e-10 - assert y_s.dtype == torch.float32, ( - "_silu_mul_fp8_quant_deep_gemm Triton fallback does not " - f"support {y_s.dtype} scales. Only torch.float32 supported." - ) + if not is_packed_ue8m0: + assert y_s.dtype == torch.float32, ( + "_silu_mul_fp8_quant_deep_gemm Triton fallback does not " + f"support {y_s.dtype} scales. Only torch.float32 supported." + ) + f32_strides = y_s_f32.stride() _silu_mul_fp8_quant_deep_gemm[grid]( y, y_q, - y_s, + y_s_f32, tokens_per_expert, H, group_size, @@ -248,9 +266,9 @@ def persistent_masked_m_silu_mul_quant( stride_yq_e, stride_yq_t, stride_yq_h, - ys_strides[0], - ys_strides[1], - ys_strides[2], + f32_strides[0], + f32_strides[1], + f32_strides[2], stride_cnt_e, eps, fp8_min, @@ -261,6 +279,23 @@ def persistent_masked_m_silu_mul_quant( num_warps=1, ) + if is_packed_ue8m0: + # Pack float32 scales into int32 UE8M0 format: + # extract exponent bits (bits 30:23) from float32. + E_dim, T_dim, G_dim = y_s_f32.shape + y_s_cont = y_s_f32.contiguous() + i32_pad = round_up(G_dim, 4) - G_dim + y_s_u8 = (y_s_cont.view(torch.int32) >> 23).to(torch.uint8) + if i32_pad > 0: + y_s_u8 = torch.nn.functional.pad(y_s_u8, (0, i32_pad)) + # y_s has shape (E, T, G//4) with stride (T*G//4, 1, T) + packed = y_s_u8.view(torch.int32) + # Copy with matching strides + for e_idx in range(E_dim): + nt = tokens_per_expert[e_idx].item() + if nt > 0: + y_s[e_idx, :nt].copy_(packed[e_idx, :nt]) + return y_q, y_s