[ROCm] Add UE8M0 scale packing for Triton silu_mul_quant (#37835)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Aug 19, 2026 at 10:20 UTC
17dbd429306d2fe9cc6be04e57aab314bd697f5e
2 files changed
+44
-13
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+1
-5
@@ -239,12 +239,8 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt
239
scale_fmts = [
240
DeepGemmQuantScaleFMT.FLOAT32,
241
DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
242
+ DeepGemmQuantScaleFMT.UE8M0,
243
]
243
- # UE8M0 (int32 packed) scales require the C++ kernel which is
244
- # not available on ROCm (#ifndef USE_ROCM).
245
- # https://github.com/ROCm/aiter/issues/2420
246
- if current_platform.is_cuda():
247
- scale_fmts.append(DeepGemmQuantScaleFMT.UE8M0)
244
245
# Run the SiLU V2 kernel
246
for scale_fmt in scale_fmts:
vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
+43
-8
@@ -220,6 +220,22 @@ def persistent_masked_m_silu_mul_quant(
220
# Triton fallback for ROCm and XPU -- the C++ kernel is guarded by
221
# #ifndef USE_ROCM in activation_kernels.cu.
222
# https://github.com/ROCm/aiter/issues/2420
223
+ # For UE8M0 (int32 packed scales), compute with float32 scales
224
+ # first, then pack afterwards.
225
+ is_packed_ue8m0 = quant_scale_fmt == DeepGemmQuantScaleFMT.UE8M0
226
+ if is_packed_ue8m0:
227
+ f32_shape, f32_strides, _ = scales_shape_stride_dtype(
228
+ E, T, G, DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
229
+ )
230
+ y_s_f32 = torch.empty_strided(
231
+ f32_shape,
232
+ f32_strides,
233
+ dtype=torch.float32,
234
+ device=y.device,
235
+ )
236
+ else:
237
+ y_s_f32 = y_s
238
+
239
stride_cnt_e = tokens_per_expert.stride()[0]
240
241
# Static grid over experts and H-groups.
@@ -231,14 +247,16 @@ def persistent_masked_m_silu_mul_quant(
247
248
fp8_min, fp8_max = get_fp8_min_max()
249
eps: float = 1e-10
234
- assert y_s.dtype == torch.float32, (
235
- "_silu_mul_fp8_quant_deep_gemm Triton fallback does not "
236
- f"support {y_s.dtype} scales. Only torch.float32 supported."
237
- )
250
+ if not is_packed_ue8m0:
251
+ assert y_s.dtype == torch.float32, (
252
+ "_silu_mul_fp8_quant_deep_gemm Triton fallback does not "
253
+ f"support {y_s.dtype} scales. Only torch.float32 supported."
254
+ )
255
+ f32_strides = y_s_f32.stride()
256
_silu_mul_fp8_quant_deep_gemm[grid](
257
y,
258
y_q,
241
- y_s,
259
+ y_s_f32,
260
tokens_per_expert,
261
H,
262
group_size,
@@ -248,9 +266,9 @@ def persistent_masked_m_silu_mul_quant(
266
stride_yq_e,
267
stride_yq_t,
268
stride_yq_h,
251
- ys_strides[0],
252
- ys_strides[1],
253
- ys_strides[2],
269
+ f32_strides[0],
270
+ f32_strides[1],
271
+ f32_strides[2],
272
stride_cnt_e,
273
eps,
274
fp8_min,
@@ -261,6 +279,23 @@ def persistent_masked_m_silu_mul_quant(
279
num_warps=1,
280
)
281
282
+ if is_packed_ue8m0:
283
+ # Pack float32 scales into int32 UE8M0 format:
284
+ # extract exponent bits (bits 30:23) from float32.
285
+ E_dim, T_dim, G_dim = y_s_f32.shape
286
+ y_s_cont = y_s_f32.contiguous()
287
+ i32_pad = round_up(G_dim, 4) - G_dim
288
+ y_s_u8 = (y_s_cont.view(torch.int32) >> 23).to(torch.uint8)
289
+ if i32_pad > 0:
290
+ y_s_u8 = torch.nn.functional.pad(y_s_u8, (0, i32_pad))
291
+ # y_s has shape (E, T, G//4) with stride (T*G//4, 1, T)
292
+ packed = y_s_u8.view(torch.int32)
293
+ # Copy with matching strides
294
+ for e_idx in range(E_dim):
295
+ nt = tokens_per_expert[e_idx].item()
296
+ if nt > 0:
297
+ y_s[e_idx, :nt].copy_(packed[e_idx, :nt])
298
+
299
return y_q, y_s
300
301