[ROCm] Bump AITER to 0.1.21.post1 (#52826)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: fai <fangzhouai@gmail.com>
Rohan Potdar committed
Sep 3, 2026 at 18:28 UTC
579aef4e8dfa2c9d498a702dd04c5ec5db333266
3 files changed
+128
-140
docker/Dockerfile.rocm_base
+1
-1
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.11.0"
9
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
10
ARG FA_BRANCH="0e60e394"
11
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12
-ARG AITER_BRANCH="v0.1.19"
12
+ARG AITER_BRANCH="v0.1.21.post1"
13
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14
ARG MORI_BRANCH="v1.1.0"
15
ARG MORI_REPO="https://github.com/ROCm/mori.git"
tests/kernels/quantization/test_rocm_mxfp4.py
+10
-1
@@ -50,7 +50,16 @@ SKINNY_GEMM_PASS_RATES = {
50
}
51
PRESHUFFLED_SHAPES = [
52
(64, 4096, 8192),
53
- (32, 8192, 8192),
53
+ # aiter 0.1.20 Triton preshuffled fp4 GEMM OOBs (GPU memory access fault ->
54
+ # SIGABRT, crashing the process) for (M=32, N=8192, K=8192) on gfx950, despite
55
+ # advertising the (N, K) as tuned. A hard GPU fault can't be xfail'd, so skip it
56
+ # until the aiter kernel is fixed (ROCm/aiter#4867).
57
+ pytest.param(
58
+ (32, 8192, 8192),
59
+ marks=pytest.mark.skip(
60
+ reason="aiter 0.1.20 preshuffled fp4 GEMM OOB (M32,N8192,K8192, gfx950)"
61
+ ),
62
+ ),
63
]
64
65
vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py
+117
-138
@@ -6,13 +6,11 @@ import functools
6
import json
7
import os
8
9
-import flydsl.compiler as flyc
10
-import flydsl.expr as fx
9
import torch
10
from aiter.fused_moe import moe_sorting as aiter_moe_sorting
13
-from aiter.ops.flydsl.kernels.moe_gemm_2stage import (
14
- compile_moe_gemm1,
15
- compile_moe_gemm2,
11
+from aiter.ops.flydsl.kernels.moe_2stage_a16wmix import (
12
+ flydsl_a16w4_gemm1,
13
+ flydsl_a16w4_gemm2,
14
)
15
16
from vllm.logger import init_logger
@@ -21,33 +19,58 @@ from vllm.utils.torch_utils import direct_register_custom_op
19
20
logger = init_logger(__name__)
21
24
-_FLYDSL_MOE_GEMM1_CACHE: dict = {}
25
-_FLYDSL_MOE_GEMM2_CACHE: dict = {}
26
-
22
+# Valid tiles for the aiter a16w-mix int4 kernels (aiter#4646): stage1 requires
23
+# tile_n in {64,128} and tile_k in {128,256} (k_wave=1); stage2 requires tile_n2==128
24
+# and tile_k2 in {128,256}. The pre-#4646 defaults (tile_k=512/64, tile_n2=256) are no
25
+# longer registered and make the port emit an empty-fragment kernel.
26
_FLYDSL_MOE_DEFAULT_CONFIG = {
28
- 1: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256},
29
- 2: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128},
30
- 4: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128},
31
- 8: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256},
32
- 16: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 256},
33
- 24: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
34
- 32: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
35
- 48: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
27
+ 1: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
28
+ 2: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
29
+ 4: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
30
+ 8: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
31
+ 16: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
32
+ 24: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
33
+ 32: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
34
+ 48: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
35
64: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
37
- 128: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
38
- 256: {"tile_m": 16, "tile_n": 128, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
39
- 512: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
40
- 1024: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
41
- 2048: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64},
42
- 4096: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
43
- 8192: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64},
36
+ 128: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
37
+ 256: {"tile_m": 16, "tile_n": 128, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
38
+ 512: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
39
+ 1024: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
40
+ 2048: {"tile_m": 64, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
41
+ 4096: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
42
+ 8192: {"tile_m": 64, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
43
}
44
46
-
47
-def _as_flydsl_ptr(tensor: torch.Tensor | None):
48
- """Adapt a torch tensor to the pointer ABI expected by AITER FlyDSL."""
49
- data_ptr = None if tensor is None else tensor.data_ptr()
50
- return flyc.from_c_void_p(fx.Uint8, data_ptr)
45
+# aiter#4646's a16w-mix int4 kernels only register a fixed tile set. Tiles from before
46
+# the bump (default config or tuned JSONs) carry tile_k=512/64 and tile_n2=256, which
47
+# emit an empty-fragment kernel (IndexError at emit / wrong output). Snap any resolved
48
+# tiles to the nearest registered value so a stale tuned config can't crash the kernel.
49
+_FLYDSL_INT4_TILE_M = (16, 32, 64, 128)
50
+_FLYDSL_INT4_TILE_N1 = (64, 128)
51
+_FLYDSL_INT4_TILE_K = (128, 256)
52
+
53
+
54
+def _snap_flydsl_int4_tiles(tile_m, tile_n, tile_k, tile_n2, tile_k2):
55
+ def _nearest(v, allowed):
56
+ return min(allowed, key=lambda a: (abs(a - v), a))
57
+
58
+ snapped = (
59
+ _nearest(tile_m, _FLYDSL_INT4_TILE_M),
60
+ _nearest(tile_n, _FLYDSL_INT4_TILE_N1),
61
+ _nearest(tile_k, _FLYDSL_INT4_TILE_K),
62
+ 128, # stage2 N tile is only registered at 128
63
+ _nearest(tile_k2, _FLYDSL_INT4_TILE_K),
64
+ )
65
+ if snapped != (tile_m, tile_n, tile_k, tile_n2, tile_k2):
66
+ logger.warning_once(
67
+ "FlyDSL int4 MoE tiles %s are not registered post-aiter#4646; "
68
+ "using %s. Retune the config for the a16w-mix kernel.",
69
+ (tile_m, tile_n, tile_k, tile_n2, tile_k2),
70
+ snapped,
71
+ scope="local",
72
+ )
73
+ return snapped
74
75
76
def moe_sorting(
@@ -185,6 +208,9 @@ def fused_flydsl_moe_impl(
208
tile_k = tuned_config["tile_k"]
209
tile_n2 = tuned_config["tile_n2"]
210
tile_k2 = tuned_config["tile_k2"]
211
+ tile_m, tile_n, tile_k, tile_n2, tile_k2 = _snap_flydsl_int4_tiles(
212
+ tile_m, tile_n, tile_k, tile_n2, tile_k2
213
+ )
214
215
routing = build_routing_buffers(
216
topk_ids=topk_ids,
@@ -202,123 +228,76 @@ def fused_flydsl_moe_impl(
228
blocks,
229
) = routing
230
205
- scale_x_1d = torch.empty((0,), device=device, dtype=torch.float32)
231
+ # aiter#4646 (v0.1.20) removed compile_moe_gemm{1,2} + the flyc.compile pointer ABI
232
+ # and replaced them with the a16w-mix wrappers, which compile+launch internally
233
+ # (functools-cached) and take torch tensors directly. w_dtype="int4" is the a16wi4
234
+ # (W4A16) path. The stage-1 intermediate now lives in sorted
235
+ # [sorted_size, inter_dim] layout, and topk weighting is applied in stage-2 via
236
+ # sorted_weights (the port has no stage-1 doweight).
237
sorted_weights_1d = sorted_weights.view(-1).contiguous()
207
- out_stage1 = torch.empty(
208
- (tokens, topk, inter_dim), device=device, dtype=out_torch_dtype
209
- )
238
+ empty_scale = torch.empty((0,), device=device, dtype=torch.uint8)
239
+ cumsum_tensor = num_valid_ids.to(torch.int32).contiguous()
240
+ m_indices = sorted_token_ids.to(torch.int32).contiguous()
241
211
- stream = torch.cuda.current_stream()
212
-
213
- key1 = (
214
- model_dim,
215
- inter_dim,
216
- num_experts,
217
- topk,
218
- in_dtype,
219
- out_dtype,
220
- group_size,
221
- tile_m,
222
- tile_n,
223
- tile_k,
224
- bool(doweight_stage1),
225
- False,
242
+ # stage1: bf16 A x int4 W1 + SiLU -> bf16 sorted intermediate.
243
+ inter_sorted = torch.empty(
244
+ (sorted_size, inter_dim), device=device, dtype=torch.bfloat16
245
)
227
-
228
- compiled_exe1 = _FLYDSL_MOE_GEMM1_CACHE.get(key1)
229
- args1 = (
230
- _as_flydsl_ptr(out_stage1),
231
- _as_flydsl_ptr(hidden_states),
232
- _as_flydsl_ptr(w1),
233
- _as_flydsl_ptr(scale_x_1d),
234
- _as_flydsl_ptr(w1_scale),
235
- _as_flydsl_ptr(sorted_token_ids),
236
- _as_flydsl_ptr(sorted_expert_ids),
237
- _as_flydsl_ptr(sorted_weights_1d),
238
- _as_flydsl_ptr(num_valid_ids),
239
- tokens,
240
- inter_dim,
241
- model_dim,
242
- int(blocks),
243
- stream,
246
+ w1_scale_u8 = (
247
+ w1_scale.view(torch.uint8).contiguous().view(-1)
248
+ if w1_scale is not None
249
+ else empty_scale
250
)
245
- if compiled_exe1 is None:
246
- exe1 = compile_moe_gemm1(
247
- model_dim=model_dim,
248
- inter_dim=inter_dim,
249
- experts=num_experts,
250
- topk=topk,
251
- in_dtype=in_dtype,
252
- out_dtype=out_dtype,
253
- group_size=group_size,
254
- tile_m=tile_m,
255
- tile_n=tile_n,
256
- tile_k=tile_k,
257
- doweight_stage1=bool(doweight_stage1),
258
- use_cshuffle_epilog=False,
259
- scale_is_bf16=scale_is_bf16,
260
- )
261
- compiled_exe1 = flyc.compile(exe1, *args1)
262
- _FLYDSL_MOE_GEMM1_CACHE[key1] = compiled_exe1
263
- else:
264
- compiled_exe1(*args1)
265
-
266
- a2_1d = out_stage1.view(-1).contiguous()
267
- a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32)
268
- out_stage2 = torch.empty((tokens, model_dim), device=device, dtype=out_torch_dtype)
269
- doweight_stage2 = not bool(doweight_stage1)
270
-
271
- key2 = (
272
- model_dim,
273
- inter_dim,
274
- num_experts,
275
- topk,
276
- in_dtype,
277
- out_dtype,
278
- group_size,
279
- tile_m,
280
- tile_n2,
281
- tile_k2,
282
- bool(doweight_stage2),
251
+ flydsl_a16w4_gemm1(
252
+ a_bf16=hidden_states.to(torch.bfloat16).contiguous(),
253
+ w1_u8=w1.view(torch.uint8).contiguous(),
254
+ w1_scale_u8=w1_scale_u8,
255
+ sorted_expert_ids=sorted_expert_ids,
256
+ cumsum_tensor=cumsum_tensor,
257
+ m_indices=m_indices,
258
+ inter_sorted_bf16=inter_sorted,
259
+ n_tokens=tokens,
260
+ NE=num_experts,
261
+ D_HIDDEN=model_dim,
262
+ D_INTER=inter_dim,
263
+ topk=topk,
264
+ tile_m=tile_m,
265
+ tile_n=tile_n,
266
+ tile_k=tile_k,
267
+ waves_per_eu=None,
268
+ act="silu",
269
+ w_dtype="int4",
270
+ w_layout="standard",
271
)
272
285
- compiled_exe2 = _FLYDSL_MOE_GEMM2_CACHE.get(key2)
286
- args2 = (
287
- _as_flydsl_ptr(out_stage2),
288
- _as_flydsl_ptr(a2_1d),
289
- _as_flydsl_ptr(w2),
290
- _as_flydsl_ptr(a2_scale_1d),
291
- _as_flydsl_ptr(w2_scale),
292
- _as_flydsl_ptr(sorted_token_ids),
293
- _as_flydsl_ptr(sorted_expert_ids),
294
- _as_flydsl_ptr(sorted_weights_1d),
295
- _as_flydsl_ptr(num_valid_ids),
296
- tokens,
297
- model_dim,
298
- inter_dim,
299
- int(blocks),
300
- stream,
273
+ # stage2: bf16 inter x int4 W2 -> weighted reduce into [tokens, model_dim].
274
+ out_stage2 = torch.zeros((tokens, model_dim), device=device, dtype=out_torch_dtype)
275
+ w2_scale_u8 = (
276
+ w2_scale.view(torch.uint8).contiguous().view(-1)
277
+ if w2_scale is not None
278
+ else empty_scale
279
+ )
280
+ flydsl_a16w4_gemm2(
281
+ inter_sorted_bf16=inter_sorted,
282
+ w2_u8=w2.view(torch.uint8).contiguous(),
283
+ w2_scale_u8=w2_scale_u8,
284
+ sorted_expert_ids=sorted_expert_ids,
285
+ cumsum_tensor=cumsum_tensor,
286
+ sorted_token_ids=sorted_token_ids,
287
+ sorted_weights=sorted_weights_1d,
288
+ flat_out=out_stage2.view(-1),
289
+ M_logical=tokens,
290
+ max_sorted=sorted_size,
291
+ NE=num_experts,
292
+ D_HIDDEN=model_dim,
293
+ D_INTER=inter_dim,
294
+ topk=topk,
295
+ tile_m=tile_m,
296
+ tile_n=tile_n2,
297
+ tile_k=tile_k2,
298
+ waves_per_eu=None,
299
+ w_dtype="int4",
300
)
302
- out_stage2.zero_()
303
- if compiled_exe2 is None:
304
- exe2 = compile_moe_gemm2(
305
- model_dim=model_dim,
306
- inter_dim=inter_dim,
307
- experts=num_experts,
308
- topk=topk,
309
- in_dtype=in_dtype,
310
- out_dtype=out_dtype,
311
- group_size=group_size,
312
- tile_m=tile_m,
313
- tile_n=tile_n2,
314
- tile_k=tile_k2,
315
- doweight_stage2=bool(doweight_stage2),
316
- scale_is_bf16=scale_is_bf16,
317
- )
318
- compiled_exe2 = flyc.compile(exe2, *args2)
319
- _FLYDSL_MOE_GEMM2_CACHE[key2] = compiled_exe2
320
- else:
321
- compiled_exe2(*args2)
301
return out_stage2
302
303