[ROCm][Perf] MXFP8 GEMM on native 32x32 block scales for gfx950 (#58510)
Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Fangzhou Ai committed
Sep 24, 2026 at 22:26 UTC
3059155b47cee5269eaf446e24c9d0f8f908c7c9
5 files changed
+899
-7
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+15
-2
@@ -1204,8 +1204,14 @@ def test_sparse_attn_decode_mxfp8_output() -> None:
1204
(1100, 8),
1205
],
1206
)
1207
+@pytest.mark.parametrize("block_scales", [False, True])
1208
@torch.inference_mode()
1208
-def test_rocm_mxfp8_wo_a_bmm(num_tokens: int, n_groups: int) -> None:
1209
+def test_rocm_mxfp8_wo_a_bmm(
1210
+ num_tokens: int, n_groups: int, block_scales: bool
1211
+) -> None:
1212
+ from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
1213
+ _as_block32_scale,
1214
+ )
1215
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1216
_mxfp8_e4m3_quantize_torch,
1217
)
@@ -1220,6 +1226,11 @@ def test_rocm_mxfp8_wo_a_bmm(num_tokens: int, n_groups: int) -> None:
1226
w, w_scale = _mxfp8_e4m3_quantize_torch(
1227
torch.randn(n_groups * o_lora_rank, group_dim, device=device)
1228
)
1229
+ if block_scales:
1230
+ # A 32x32 checkpoint loads as per-row scales that repeat every 32 rows;
1231
+ # RocmDotScaledMxfp8LinearKernel compacts them before wo_a sees them.
1232
+ w_scale = _as_block32_scale(w_scale[::32].repeat_interleave(32, dim=0))
1233
+ assert w_scale is not None
1234
wo_a = SimpleNamespace(weight=w, weight_scale=w_scale)
1235
1236
out = rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank)
@@ -1227,7 +1238,9 @@ def test_rocm_mxfp8_wo_a_bmm(num_tokens: int, n_groups: int) -> None:
1238
expected = torch.einsum(
1239
"tgd,grd->tgr",
1240
_mxfp8_dequant(a, a_scale).view(num_tokens, n_groups, group_dim),
1230
- _mxfp8_dequant(w, w_scale).view(n_groups, o_lora_rank, group_dim),
1241
+ _mxfp8_dequant(
1242
+ w, w_scale.repeat_interleave(w.shape[0] // w_scale.shape[0], 0)
1243
+ ).view(n_groups, o_lora_rank, group_dim),
1244
)
1245
torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2)
1246
tests/kernels/quantization/test_rocm_mxfp8_linear.py
+141
@@ -117,6 +117,147 @@ def test_rocm_mxfp8_linear_handles_extreme_magnitudes():
117
assert (out[0] == 0).all(), "an all-zero activation row must give a zero row"
118
119
120
+# (name, N, K) of the DeepSeek-V4.1-Flash MXFP8 linears at TP=2 and TP=4. The
121
+# checkpoint stores their scales in 32x32 blocks; shared_down at TP=4 has
122
+# K % 128 != 0, which only the block-scaled GEMM runs natively.
123
+V41_BLOCK32_SHAPES = [
124
+ ("fused_wqa_wkv", 1280 + 512, 5120),
125
+ ("wq_b_tp2", 64 * 512 // 2, 1280),
126
+ ("indexer_wq_b", 32 * 128, 1280),
127
+ ("wo_b_tp2", 5120, 8192 // 2),
128
+ ("shared_gate_up_tp2", 2 * 2304 // 2, 5120),
129
+ ("shared_down_tp2", 5120, 2304 // 2),
130
+ ("shared_down_tp4", 5120, 2304 // 4),
131
+]
132
+
133
+
134
+def _make_block32_layer(n: int, k: int, device: str) -> torch.nn.Module:
135
+ """An MXFP8 layer with 32x32 block scales, expanded per row as loaded."""
136
+ blocks = (torch.randn(n, k, device=device) * 0.05).view(n // 32, 32, k // 32, 32)
137
+ amax = blocks.abs().amax(dim=(1, 3)).clamp(min=torch.finfo(torch.float32).tiny)
138
+ bits = (torch.ceil(torch.log2(amax / 448.0)) + 127.0).clamp(0, 254)
139
+ weight = (blocks * torch.exp2(127.0 - bits)[:, None, :, None]).view(n, k)
140
+ layer = torch.nn.Module()
141
+ layer.weight = torch.nn.Parameter(
142
+ weight.to(torch.float8_e4m3fn), requires_grad=False
143
+ )
144
+ layer.weight_scale = torch.nn.Parameter(
145
+ bits.to(torch.uint8).repeat_interleave(32, dim=0), requires_grad=False
146
+ )
147
+ return layer
148
+
149
+
150
+@pytest.mark.parametrize("name,n,k", V41_BLOCK32_SHAPES, ids=lambda v: str(v))
151
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 33, 192, 1024])
152
+@torch.inference_mode()
153
+def test_rocm_mxfp8_linear_block32_matches_dequant_reference(name, n, k, num_tokens):
154
+ """32x32 block scales stay compact and the block-scaled GEMM is exact.
155
+
156
+ The token counts cover the packed, tiled and split-K configurations.
157
+ """
158
+ from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
159
+ Mxfp8LinearLayerConfig,
160
+ )
161
+ from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
162
+ RocmDotScaledMxfp8LinearKernel,
163
+ )
164
+
165
+ torch.manual_seed(num_tokens)
166
+ device = "cuda"
167
+ layer = _make_block32_layer(n, k, device)
168
+ row_scale = layer.weight_scale.data.clone()
169
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
170
+ kernel.process_weights_after_loading(layer)
171
+ assert layer.weight.dtype == torch.float8_e4m3fn
172
+ assert layer.weight_scale.shape == (n // 32, k // 32)
173
+
174
+ x = torch.randn(num_tokens, k, device=device, dtype=torch.bfloat16) * 0.5
175
+ out = kernel.apply_weights(layer, x)
176
+
177
+ x_q, x_scale = mxfp8_e4m3_quantize(x)
178
+ expected = (
179
+ dequant_mxfp8_to_bf16(x_q, x_scale).float()
180
+ @ dequant_mxfp8_to_bf16(layer.weight, row_scale).float().T
181
+ )
182
+ assert out.shape == (num_tokens, n)
183
+ rel = (out.float() - expected).norm() / expected.norm()
184
+ assert rel < 5e-3, f"{name}: relative error {rel:.4f}"
185
+
186
+
187
+@torch.inference_mode()
188
+def test_rocm_mxfp8_linear_keeps_per_row_scales():
189
+ """Scales that differ within a 32-row block are not compacted."""
190
+ from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
191
+ Mxfp8LinearLayerConfig,
192
+ )
193
+ from vllm.model_executor.kernels.linear.mxfp8.rocm_native import (
194
+ RocmDotScaledMxfp8LinearKernel,
195
+ )
196
+
197
+ layer = _make_block32_layer(64, 256, "cuda")
198
+ layer.weight_scale.data[1, 0] += 1
199
+ kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(bmm_batch_size=None))
200
+ kernel.process_weights_after_loading(layer)
201
+ assert layer.weight_scale.shape == (64, 256 // 32)
202
+
203
+
204
+def _split_k_case(packed: bool):
205
+ """A 3-way in-launch split-K on shared_down TP4, with K and M tails."""
206
+ from vllm.model_executor.kernels.linear.mxfp8 import rocm_block32_gemm as g
207
+
208
+ n, k, m = 5120, 576, 19
209
+ layer = _make_block32_layer(n, k, "cuda")
210
+ x_q, x_scale = mxfp8_e4m3_quantize(
211
+ torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
212
+ )
213
+ expected = (
214
+ dequant_mxfp8_to_bf16(x_q, x_scale).float()
215
+ @ dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale).float().T
216
+ )
217
+ if packed:
218
+ cfg = g._packed(16, 32, 128, 1, 2, 2, 16, 0, 3)
219
+ else:
220
+ cfg = g._tiled(32, 64, 128, 1, True, 4, 2, 16, 0, 3)
221
+ args = (x_q, x_scale, layer.weight, layer.weight_scale.data[::32].contiguous())
222
+ return g, cfg, args, expected
223
+
224
+
225
+def _rel_err(out: torch.Tensor, expected: torch.Tensor) -> float:
226
+ return ((out.float() - expected).norm() / expected.norm()).item()
227
+
228
+
229
+@pytest.mark.parametrize("packed", [True, False])
230
+@torch.inference_mode()
231
+def test_rocm_mxfp8_block32_split_k_in_launch(packed):
232
+ """Splits summed by the tile's last program are exact and repeatable."""
233
+ torch.manual_seed(0)
234
+ g, cfg, args, expected = _split_k_case(packed)
235
+ out = torch.empty_like(expected, dtype=torch.bfloat16)
236
+ g._launch(*args, out, cfg)
237
+ assert _rel_err(out, expected) < 5e-3
238
+ first = out.clone()
239
+ for _ in range(20):
240
+ g._launch(*args, out, cfg)
241
+ assert torch.equal(out, first)
242
+
243
+
244
+@pytest.mark.parametrize("packed", [True, False])
245
+@torch.inference_mode()
246
+def test_rocm_mxfp8_block32_split_k_first_use_in_capture(packed):
247
+ """A stream whose first split-K launch is being captured has no counters
248
+ yet; it must fall back to an unfused split instead of failing."""
249
+ torch.manual_seed(0)
250
+ g, cfg, args, expected = _split_k_case(packed)
251
+ out = torch.zeros_like(expected, dtype=torch.bfloat16)
252
+ stream = torch.cuda.Stream()
253
+ graph = torch.cuda.CUDAGraph()
254
+ with torch.cuda.graph(graph, stream=stream):
255
+ g._launch(*args, out, cfg)
256
+ graph.replay()
257
+ torch.accelerator.synchronize()
258
+ assert _rel_err(out, expected) < 5e-3
259
+
260
+
261
@pytest.mark.parametrize("num_rows", [1, 64, 65])
262
@torch.inference_mode()
263
def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows):
vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
new
+711
@@ -0,0 +1,711 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MXFP8 GEMM on 32x32 block-scaled weights for gfx950 (``tl.dot_scaled``).
4
+
5
+``y = x @ w.T`` with an MXFP8 activation (e4m3 values, one E8M0 scale per
6
+row and 32 K, ``[M, K / 32]``) and an e4m3 weight whose E8M0 scales stay in
7
+the checkpoint's 32x32 blocks, ``[N / 32, K / 32]``, instead of being
8
+expanded to every row. Each weight scale byte is then read once per 32 output
9
+rows, and small-M shapes can use the packed kernel below.
10
+
11
+Two kernels, picked per shape from a table tuned on MI355X:
12
+
13
+* a tiled kernel for larger M;
14
+* a packed kernel for small M, which fills the MFMA's rows with K panels
15
+ instead of tokens and keeps only the matching-panel (block-diagonal)
16
+ products, so a handful of tokens still streams the weight at full width.
17
+
18
+Either can split K. The partials are normally summed in the same launch by
19
+the last program of each output tile to finish, in split order, so the result
20
+does not depend on scheduling. Otherwise a second launch reduces them.
21
+"""
22
+
23
+from typing import NamedTuple
24
+
25
+import torch
26
+
27
+from vllm.triton_utils import tl, triton
28
+from vllm.utils.torch_utils import direct_register_custom_op
29
+
30
+BLOCK_ROWS = 32
31
+
32
+# gfx950 hands out workgroups to its XCDs round-robin by program id, and 8 is
33
+# a multiple of the XCD count in every partition mode. _split_tile hardcodes
34
+# the same stride.
35
+_XCD_STRIDE = 8
36
+# One arrival counter per output tile of an in-launch split-K GEMM.
37
+_SPLIT_COUNTERS = 1 << 16
38
+
39
+
40
+@triton.jit
41
+def _split_tile(N: tl.constexpr, BLOCK_N: tl.constexpr, SPLITS: tl.constexpr):
42
+ """(pid_m, pid_n, tile, split) of a 1D grid that keeps a tile's splits on
43
+ one XCD, so they meet in that XCD's L2."""
44
+ pid = tl.program_id(0)
45
+ tile = pid // (8 * SPLITS) * 8 + pid % 8
46
+ grid_n: tl.constexpr = (N + BLOCK_N - 1) // BLOCK_N
47
+ return tile // grid_n, tile % grid_n, tile, pid // 8 % SPLITS
48
+
49
+
50
+@triton.jit
51
+def _sum_splits(
52
+ acc, out_ptrs, out_mask, slot_ptr, count_ptr, split, SPLITS: tl.constexpr
53
+):
54
+ """Publish this split's partial; the tile's last arrival sums all of them
55
+ in split order and re-arms the counter."""
56
+ BM: tl.constexpr = acc.shape[0]
57
+ BN: tl.constexpr = acc.shape[1]
58
+ local = tl.arange(0, BM)[:, None] * BN + tl.arange(0, BN)[None, :]
59
+ tl.store(slot_ptr + split * BM * BN + local, acc)
60
+ # The partial must be in L2 before the arrival is counted. The splits
61
+ # share an XCD, so no device-scope release (an L2 writeback) is needed.
62
+ tl.inline_asm_elementwise(
63
+ "s_waitcnt vmcnt(0)", "=v,v", [split], dtype=tl.int32, is_pure=False, pack=1
64
+ )
65
+ tl.debug_barrier()
66
+ if tl.atomic_add(count_ptr, 1, sem="acq_rel", scope="cta") == SPLITS - 1:
67
+ total = tl.zeros((BM, BN), dtype=tl.float32)
68
+ for s in tl.static_range(SPLITS):
69
+ total += tl.load(slot_ptr + s * BM * BN + local, cache_modifier=".cv")
70
+ tl.store(out_ptrs, total.to(out_ptrs.dtype.element_ty), mask=out_mask)
71
+ tl.store(count_ptr, 0)
72
+
73
+
74
+@triton.jit(do_not_specialize=["M"])
75
+def _block32_tiled_kernel(
76
+ x_ptr,
77
+ xs_ptr,
78
+ w_ptr,
79
+ ws_ptr,
80
+ out_ptr,
81
+ slot_ptr,
82
+ count_ptr,
83
+ M,
84
+ N: tl.constexpr,
85
+ K: tl.constexpr,
86
+ stride_om,
87
+ BLOCK_M: tl.constexpr,
88
+ BLOCK_N: tl.constexpr,
89
+ BLOCK_K: tl.constexpr,
90
+ K_PER_SPLIT: tl.constexpr,
91
+ N_FIRST: tl.constexpr,
92
+ EVEN_K: tl.constexpr,
93
+ FUSED_SPLITS: tl.constexpr,
94
+):
95
+ if FUSED_SPLITS > 1:
96
+ pid_m, pid_n, tile, pid_k = _split_tile(N, BLOCK_N, FUSED_SPLITS)
97
+ if pid_m * BLOCK_M >= M:
98
+ return
99
+ else:
100
+ # N_FIRST launches output-column tiles fastest, so consecutive
101
+ # programs share an activation tile.
102
+ if N_FIRST:
103
+ pid_n, pid_m = tl.program_id(0), tl.program_id(1)
104
+ else:
105
+ pid_m, pid_n = tl.program_id(0), tl.program_id(1)
106
+ pid_k = tl.program_id(2)
107
+ offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
108
+ offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
109
+ m_mask = offs_m < M
110
+ n_mask = offs_n < N
111
+ k0 = pid_k * K_PER_SPLIT
112
+ offs_k = k0 + tl.arange(0, BLOCK_K)
113
+ offs_sk = k0 // 32 + tl.arange(0, BLOCK_K // 32)
114
+ x_ptrs = x_ptr + offs_m[:, None] * K + offs_k[None, :]
115
+ xs_ptrs = xs_ptr + offs_m[:, None] * (K // 32) + offs_sk[None, :]
116
+ w_ptrs = w_ptr + offs_n[:, None] * K + offs_k[None, :]
117
+ ws_ptrs = ws_ptr + (offs_n[:, None] // 32) * (K // 32) + offs_sk[None, :]
118
+ acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
119
+ for kk in range(0, K_PER_SPLIT, BLOCK_K):
120
+ if EVEN_K:
121
+ x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
122
+ xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=127)
123
+ w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
124
+ ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=127)
125
+ else:
126
+ k_ok = (offs_k + kk) < K
127
+ s_ok = (offs_sk + kk // 32) < K // 32
128
+ x = tl.load(x_ptrs, mask=m_mask[:, None] & k_ok[None, :], other=0.0)
129
+ xs = tl.load(xs_ptrs, mask=m_mask[:, None] & s_ok[None, :], other=127)
130
+ w = tl.load(w_ptrs, mask=n_mask[:, None] & k_ok[None, :], other=0.0)
131
+ ws = tl.load(ws_ptrs, mask=n_mask[:, None] & s_ok[None, :], other=127)
132
+ acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
133
+ x_ptrs += BLOCK_K
134
+ w_ptrs += BLOCK_K
135
+ xs_ptrs += BLOCK_K // 32
136
+ ws_ptrs += BLOCK_K // 32
137
+ out_mask = m_mask[:, None] & n_mask[None, :]
138
+ if FUSED_SPLITS > 1:
139
+ _sum_splits(
140
+ acc,
141
+ out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
142
+ out_mask,
143
+ slot_ptr + tile * (FUSED_SPLITS * BLOCK_M * BLOCK_N),
144
+ count_ptr + tile,
145
+ pid_k,
146
+ FUSED_SPLITS,
147
+ )
148
+ else:
149
+ out_ptrs = (
150
+ out_ptr + pid_k * M * N + offs_m[:, None] * stride_om + offs_n[None, :]
151
+ )
152
+ tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=out_mask)
153
+
154
+
155
+@triton.jit(do_not_specialize=["M"])
156
+def _block32_packed_kernel(
157
+ x_ptr,
158
+ xs_ptr,
159
+ w_ptr,
160
+ ws_ptr,
161
+ out_ptr,
162
+ slot_ptr,
163
+ count_ptr,
164
+ M,
165
+ N: tl.constexpr,
166
+ K: tl.constexpr,
167
+ stride_om,
168
+ BLOCK_M: tl.constexpr,
169
+ BLOCK_N: tl.constexpr,
170
+ BLOCK_K: tl.constexpr,
171
+ K_PACK: tl.constexpr,
172
+ K_PER_SPLIT: tl.constexpr,
173
+ EVEN_K: tl.constexpr,
174
+ FUSED_SPLITS: tl.constexpr,
175
+):
176
+ if FUSED_SPLITS > 1:
177
+ pid_m, pid_n, tile, pid_k = _split_tile(N, BLOCK_N, FUSED_SPLITS)
178
+ if pid_m * BLOCK_M >= M:
179
+ return
180
+ k_start = pid_k * K_PER_SPLIT
181
+ k_stop = tl.minimum(k_start + K_PER_SPLIT, K)
182
+ else:
183
+ pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
184
+ if K_PER_SPLIT >= K:
185
+ k_start = 0
186
+ k_stop = K
187
+ else:
188
+ k_start = pid_k * K_PER_SPLIT
189
+ k_stop = tl.minimum(k_start + K_PER_SPLIT, K)
190
+ # Tile row r is token r // K_PACK reading K panel r % K_PACK (and likewise
191
+ # for weight rows), so one dot covers K_PACK consecutive BLOCK_K slices.
192
+ rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M * K_PACK) // K_PACK
193
+ x_panel = tl.arange(0, BLOCK_M * K_PACK) % K_PACK
194
+ cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N * K_PACK) // K_PACK
195
+ w_panel = tl.arange(0, BLOCK_N * K_PACK) % K_PACK
196
+ ks = tl.arange(0, BLOCK_K)
197
+ gs = tl.arange(0, BLOCK_K // 32)
198
+ row_ok = rows < M
199
+ col_ok = cols < N
200
+ acc = tl.zeros((BLOCK_M * K_PACK, BLOCK_N * K_PACK), dtype=tl.float32)
201
+ for base in range(k_start, k_stop, BLOCK_K * K_PACK):
202
+ xk = base + x_panel[:, None] * BLOCK_K + ks[None, :]
203
+ wk = base + w_panel[:, None] * BLOCK_K + ks[None, :]
204
+ xg = base // 32 + x_panel[:, None] * (BLOCK_K // 32) + gs[None, :]
205
+ wg = base // 32 + w_panel[:, None] * (BLOCK_K // 32) + gs[None, :]
206
+ x_mask = row_ok[:, None]
207
+ w_mask = col_ok[:, None]
208
+ xs_mask = row_ok[:, None]
209
+ ws_mask = col_ok[:, None]
210
+ if not EVEN_K:
211
+ x_mask = x_mask & (xk < K)
212
+ w_mask = w_mask & (wk < K)
213
+ xs_mask = xs_mask & (xg < K // 32)
214
+ ws_mask = ws_mask & (wg < K // 32)
215
+ x = tl.load(x_ptr + rows[:, None] * K + xk, mask=x_mask, other=0.0)
216
+ w = tl.load(w_ptr + cols[:, None] * K + wk, mask=w_mask, other=0.0)
217
+ xs = tl.load(xs_ptr + rows[:, None] * (K // 32) + xg, mask=xs_mask, other=127)
218
+ ws = tl.load(
219
+ ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=127
220
+ )
221
+ acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
222
+ # Keep only products of matching K panels: the block diagonal.
223
+ blocks = acc.reshape(BLOCK_M, K_PACK, BLOCK_N, K_PACK).trans(0, 2, 1, 3)
224
+ pair = tl.arange(0, K_PACK)
225
+ diagonal = tl.where(
226
+ pair[None, None, :, None] == pair[None, None, None, :], blocks, 0.0
227
+ )
228
+ y = tl.sum(tl.sum(diagonal, 3), 2)
229
+ out_rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
230
+ out_cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
231
+ out_ptrs = out_ptr + out_rows[:, None] * stride_om + out_cols[None, :]
232
+ out_mask = (out_rows[:, None] < M) & (out_cols[None, :] < N)
233
+ if FUSED_SPLITS > 1:
234
+ _sum_splits(
235
+ y,
236
+ out_ptrs,
237
+ out_mask,
238
+ slot_ptr + tile * (FUSED_SPLITS * BLOCK_M * BLOCK_N),
239
+ count_ptr + tile,
240
+ pid_k,
241
+ FUSED_SPLITS,
242
+ )
243
+ else:
244
+ tl.store(
245
+ out_ptrs + pid_k * M * N, y.to(out_ptr.dtype.element_ty), mask=out_mask
246
+ )
247
+
248
+
249
+@triton.jit(do_not_specialize=["M"])
250
+def _block32_splitk_reduce_kernel(
251
+ partial_ptr,
252
+ out_ptr,
253
+ M,
254
+ N: tl.constexpr,
255
+ stride_om,
256
+ SPLIT_K: tl.constexpr,
257
+ BLOCK_M: tl.constexpr,
258
+ BLOCK_N: tl.constexpr,
259
+):
260
+ offs_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
261
+ offs_n = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
262
+ mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
263
+ acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
264
+ for split in tl.static_range(SPLIT_K):
265
+ acc += tl.load(
266
+ partial_ptr + split * M * N + offs_m[:, None] * N + offs_n[None, :],
267
+ mask=mask,
268
+ other=0.0,
269
+ )
270
+ tl.store(
271
+ out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
272
+ acc.to(out_ptr.dtype.element_ty),
273
+ mask=mask,
274
+ )
275
+
276
+
277
+class _Config(NamedTuple):
278
+ packed: bool
279
+ block_m: int
280
+ block_n: int
281
+ block_k: int
282
+ # K panels per MFMA tile (packed) or K partitions reduced by a second
283
+ # launch (tiled).
284
+ k_split: int
285
+ n_first: bool
286
+ num_warps: int
287
+ num_stages: int
288
+ matrix_instr_nonkdim: int
289
+ waves_per_eu: int
290
+ # K partitions summed within the launch.
291
+ fused_splits: int = 1
292
+
293
+
294
+def _tiled(
295
+ bm, bn, bk, split_k, n_first, warps, stages, nonkdim, waves, fused=1
296
+) -> _Config:
297
+ return _Config(
298
+ False, bm, bn, bk, split_k, bool(n_first), warps, stages, nonkdim, waves, fused
299
+ )
300
+
301
+
302
+def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Config:
303
+ return _Config(
304
+ True, bm, bn, bk, k_pack, False, warps, stages, nonkdim, waves, fused
305
+ )
306
+
307
+
308
+# (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8 linears at
309
+# TP2/TP4, tuned on MI355X under HIP graphs with weights cold in the LLC. A
310
+# tier serves every M up to its bound; larger M reuse the last tier.
311
+_TUNED: dict[tuple[int, int], list[tuple[int, _Config]]] = {
312
+ # shared expert gate_up, TP4
313
+ (1152, 5120): [
314
+ (1, _packed(8, 16, 1024, 2, 2, 2, 16, 1, 3)),
315
+ (2, _packed(4, 32, 256, 4, 2, 2, 16, 1, 5)),
316
+ (4, _packed(4, 32, 256, 4, 2, 2, 16, 0, 5)),
317
+ (8, _packed(8, 16, 1024, 2, 2, 2, 16, 0, 3)),
318
+ (16, _tiled(16, 32, 512, 1, True, 2, 3, 16, 0, 5)),
319
+ (24, _tiled(32, 32, 512, 1, True, 4, 3, 16, 1, 5)),
320
+ (32, _packed(32, 32, 1024, 1, 4, 2, 16, 0, 5)),
321
+ (48, _tiled(64, 32, 512, 1, True, 4, 3, 16, 1, 5)),
322
+ (64, _tiled(64, 32, 512, 1, True, 8, 3, 16, 1, 5)),
323
+ (96, _tiled(32, 16, 1024, 1, True, 2, 3, 16, 1)),
324
+ (128, _tiled(64, 32, 512, 1, True, 8, 2, 16, 0, 3)),
325
+ (160, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0, 2)),
326
+ (192, _tiled(32, 32, 1024, 1, True, 4, 2, 16, 0)),
327
+ (256, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 3)),
328
+ (320, _tiled(64, 32, 256, 1, True, 8, 3, 16, 0, 2)),
329
+ (384, _tiled(32, 32, 512, 1, True, 4, 2, 16, 0)),
330
+ (512, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
331
+ (768, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
332
+ (1024, _tiled(128, 128, 256, 2, False, 4, 2, 16, 1)),
333
+ (1536, _tiled(128, 128, 256, 2, False, 4, 2, 16, 0)),
334
+ (2048, _tiled(128, 128, 256, 1, False, 4, 2, 32, 0)),
335
+ (3072, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
336
+ (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
337
+ (8192, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
338
+ ],
339
+ # fused wq_a|wkv (replicated)
340
+ (1792, 5120): [
341
+ (1, _packed(8, 16, 1024, 2, 2, 3, 16, 0)),
342
+ (8, _packed(8, 16, 1024, 2, 2, 3, 16, 1)),
343
+ (16, _tiled(16, 64, 1024, 1, True, 4, 2, 16, 0, 5)),
344
+ (32, _packed(16, 16, 1024, 1, 2, 3, 16, 1)),
345
+ (48, _tiled(64, 32, 512, 1, True, 8, 2, 16, 0, 4)),
346
+ (64, _tiled(64, 32, 512, 1, True, 8, 3, 16, 1, 4)),
347
+ (96, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0, 2)),
348
+ (128, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 4)),
349
+ (192, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 3)),
350
+ (256, _tiled(64, 64, 512, 1, True, 4, 2, 16, 1, 2)),
351
+ (384, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
352
+ (512, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
353
+ (768, _tiled(128, 128, 256, 1, True, 4, 2, 16, 1, 2)),
354
+ (1024, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
355
+ (2048, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
356
+ (3072, _tiled(128, 128, 256, 1, False, 8, 2, 32, 1)),
357
+ (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
358
+ (8192, _tiled(128, 128, 128, 1, False, 8, 2, 16, 1)),
359
+ ],
360
+ # shared expert gate_up, TP2
361
+ (2304, 5120): [
362
+ (1, _packed(8, 16, 1024, 2, 2, 3, 16, 1)),
363
+ (8, _packed(8, 16, 1024, 2, 2, 3, 16, 0)),
364
+ (16, _tiled(16, 64, 1024, 1, True, 4, 3, 16, 1, 5)),
365
+ (24, _packed(32, 32, 512, 1, 4, 2, 16, 0, 5)),
366
+ (32, _packed(32, 32, 1024, 1, 4, 2, 16, 1, 3)),
367
+ (48, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0, 3)),
368
+ (64, _tiled(64, 64, 512, 1, True, 8, 2, 16, 0, 5)),
369
+ (128, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 3)),
370
+ (160, _tiled(64, 64, 512, 1, True, 8, 2, 16, 1, 2)),
371
+ (192, _tiled(64, 64, 512, 1, True, 4, 2, 16, 0, 2)),
372
+ (256, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
373
+ (320, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
374
+ (384, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
375
+ (512, _tiled(128, 128, 256, 1, True, 8, 2, 16, 0, 3)),
376
+ (768, _tiled(128, 128, 256, 1, True, 4, 2, 16, 1, 2)),
377
+ (1536, _tiled(128, 128, 256, 1, False, 4, 2, 32, 0)),
378
+ (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
379
+ (3072, _tiled(128, 128, 128, 1, True, 8, 2, 16, 0)),
380
+ (4096, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
381
+ (8192, _tiled(128, 128, 128, 1, False, 8, 2, 32, 0)),
382
+ ],
383
+ # indexer wq_b
384
+ (4096, 1280): [
385
+ (1, _packed(4, 16, 256, 4, 4, 2, 16, 1)),
386
+ (2, _packed(4, 16, 256, 4, 4, 2, 16, 0)),
387
+ (4, _packed(4, 16, 256, 4, 4, 2, 16, 1)),
388
+ (8, _packed(8, 16, 256, 4, 4, 3, 16, 1)),
389
+ (16, _packed(16, 16, 512, 2, 2, 3, 16, 1)),
390
+ (32, _packed(32, 16, 1024, 1, 2, 3, 16, 1)),
391
+ (64, _tiled(32, 32, 512, 1, True, 4, 3, 16, 1)),
392
+ (128, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0)),
393
+ (192, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0)),
394
+ (256, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
395
+ (384, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
396
+ (512, _tiled(128, 64, 256, 1, True, 8, 3, 32, 0)),
397
+ (1024, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
398
+ (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
399
+ (3072, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
400
+ (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
401
+ (8192, _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)),
402
+ ],
403
+ # shared expert down, TP4
404
+ (5120, 576): [
405
+ (8, _packed(8, 32, 256, 2, 2, 2, 16, 1)),
406
+ (16, _packed(16, 32, 512, 1, 2, 2, 16, 0)),
407
+ (32, _packed(32, 32, 256, 1, 8, 3, 16, 1)),
408
+ (64, _tiled(32, 32, 128, 1, True, 4, 3, 16, 0)),
409
+ (96, _tiled(32, 64, 256, 1, True, 4, 3, 16, 0)),
410
+ (128, _tiled(32, 32, 256, 1, True, 2, 2, 16, 0)),
411
+ (192, _tiled(64, 64, 256, 1, True, 4, 3, 32, 0)),
412
+ (256, _tiled(64, 64, 256, 1, True, 2, 3, 32, 1)),
413
+ (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
414
+ (384, _tiled(128, 64, 256, 1, True, 8, 3, 32, 1)),
415
+ (512, _tiled(64, 64, 128, 1, True, 4, 2, 32, 0)),
416
+ (768, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
417
+ (1024, _tiled(128, 64, 128, 1, True, 4, 2, 32, 0)),
418
+ (3072, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
419
+ (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
420
+ (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
421
+ ],
422
+ # shared expert down, TP2
423
+ (5120, 1152): [
424
+ (1, _packed(4, 32, 128, 4, 2, 3, 16, 1)),
425
+ (8, _packed(8, 32, 256, 2, 4, 3, 16, 0)),
426
+ (32, _packed(32, 32, 512, 1, 2, 3, 16, 1)),
427
+ (128, _tiled(32, 32, 128, 1, True, 4, 3, 16, 0)),
428
+ (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
429
+ (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
430
+ (384, _tiled(128, 64, 256, 1, True, 8, 3, 32, 1)),
431
+ (512, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
432
+ (768, _tiled(128, 128, 256, 1, True, 4, 2, 32, 0)),
433
+ (1024, _tiled(128, 64, 128, 1, True, 8, 2, 32, 0)),
434
+ (1536, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
435
+ (2048, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
436
+ (3072, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)),
437
+ (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
438
+ (8192, _tiled(128, 64, 128, 1, True, 8, 2, 32, 1)),
439
+ ],
440
+ # wo_b, TP4
441
+ (5120, 2048): [
442
+ (1, _packed(4, 32, 256, 4, 2, 2, 16, 0)),
443
+ (4, _packed(4, 32, 256, 4, 2, 2, 16, 1)),
444
+ (8, _packed(8, 32, 256, 4, 2, 3, 16, 0)),
445
+ (16, _tiled(16, 32, 1024, 1, True, 2, 3, 16, 0)),
446
+ (32, _tiled(32, 32, 512, 1, False, 4, 3, 16, 1)),
447
+ (64, _tiled(32, 32, 512, 1, True, 4, 2, 16, 0)),
448
+ (96, _tiled(32, 64, 512, 1, True, 8, 2, 16, 0)),
449
+ (128, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
450
+ (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
451
+ (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
452
+ (384, _tiled(128, 64, 256, 1, True, 8, 3, 32, 0)),
453
+ (768, _tiled(128, 128, 256, 1, True, 4, 2, 32, 0)),
454
+ (1024, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
455
+ (1536, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
456
+ (2048, _tiled(128, 128, 256, 1, True, 8, 2, 32, 1)),
457
+ (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
458
+ (8192, _tiled(128, 128, 128, 1, True, 2, 2, 32, 1)),
459
+ ],
460
+ # wo_b, TP2
461
+ (5120, 4096): [
462
+ (4, _packed(4, 32, 512, 4, 2, 2, 16, 0)),
463
+ (8, _packed(8, 16, 128, 4, 2, 2, 16, 0, 2)),
464
+ (16, _tiled(16, 32, 256, 1, True, 2, 2, 16, 0, 4)),
465
+ (24, _tiled(32, 32, 256, 1, True, 4, 3, 16, 0, 4)),
466
+ (32, _tiled(32, 64, 512, 1, True, 8, 3, 16, 0, 3)),
467
+ (48, _tiled(64, 64, 512, 1, True, 8, 2, 16, 1, 3)),
468
+ (64, _tiled(64, 64, 512, 1, True, 8, 2, 16, 0, 3)),
469
+ (128, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
470
+ (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
471
+ (256, _tiled(128, 128, 256, 1, True, 4, 2, 16, 1, 2)),
472
+ (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
473
+ (384, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0, 2)),
474
+ (768, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
475
+ (1024, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
476
+ (1536, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
477
+ (2048, _tiled(128, 128, 256, 1, True, 8, 2, 32, 1)),
478
+ (3072, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
479
+ (4096, _tiled(128, 128, 128, 1, True, 8, 2, 16, 1)),
480
+ (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
481
+ ],
482
+ # wq_b, TP4
483
+ (8192, 1280): [
484
+ (1, _packed(8, 32, 256, 2, 4, 3, 16, 1)),
485
+ (8, _packed(8, 32, 256, 2, 4, 3, 16, 0)),
486
+ (16, _packed(16, 32, 1024, 1, 2, 3, 16, 1)),
487
+ (32, _packed(32, 32, 512, 1, 2, 3, 16, 0)),
488
+ (48, _tiled(64, 32, 512, 1, True, 8, 3, 16, 1)),
489
+ (64, _packed(32, 64, 256, 1, 4, 3, 16, 1)),
490
+ (128, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
491
+ (192, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
492
+ (256, _tiled(128, 64, 256, 1, True, 8, 3, 32, 1)),
493
+ (384, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
494
+ (512, _tiled(128, 128, 256, 1, True, 4, 2, 32, 0)),
495
+ (1024, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
496
+ (1536, _tiled(128, 128, 128, 1, True, 8, 2, 16, 0)),
497
+ (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
498
+ (8192, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)),
499
+ ],
500
+ # wq_b, TP2
501
+ (16384, 1280): [
502
+ (2, _packed(4, 32, 128, 4, 2, 3, 16, 1)),
503
+ (4, _packed(4, 32, 128, 4, 4, 2, 16, 0)),
504
+ (8, _packed(8, 32, 256, 2, 4, 2, 16, 0)),
505
+ (32, _packed(32, 64, 256, 1, 4, 3, 16, 1)),
506
+ (64, _tiled(64, 64, 256, 1, False, 4, 3, 32, 1)),
507
+ (128, _tiled(128, 64, 256, 1, True, 8, 3, 32, 0)),
508
+ (256, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
509
+ (384, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
510
+ (512, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
511
+ (768, _tiled(128, 128, 128, 1, True, 8, 2, 16, 1)),
512
+ (1024, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
513
+ (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
514
+ (8192, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)),
515
+ ],
516
+}
517
+
518
+
519
+def _default_config(M: int, N: int, K: int) -> _Config:
520
+ """Untuned shapes: the tiers most tuned shapes settle on."""
521
+ if M <= 8:
522
+ return _packed(8, 16, 256, 2, 2, 3, 16, 0)
523
+ if M <= 32:
524
+ return _packed(32, 32, 256, 1, 2, 3, 16, 0)
525
+ if M <= 256:
526
+ return _tiled(64, 64, 256, 1, True, 4, 2, 16, 0)
527
+ return _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)
528
+
529
+
530
+def _config(M: int, N: int, K: int) -> _Config:
531
+ tiers = _TUNED.get((N, K))
532
+ if tiers is None:
533
+ return _default_config(M, N, K)
534
+ for max_m, cfg in tiers:
535
+ if max_m >= M:
536
+ return cfg
537
+ return tiers[-1][1]
538
+
539
+
540
+_split_counters_by_stream: dict[tuple[int, int], torch.Tensor] = {}
541
+
542
+
543
+def _split_counters(device: torch.device) -> torch.Tensor | None:
544
+ """Zeroed arrival counters, one set per stream so concurrent launches never
545
+ share a tile's counter. Each launch leaves the counters it used at zero.
546
+
547
+ None when a stream first needs them inside a CUDA graph capture: that
548
+ allocation would come from the graph's pool, where it can take the address
549
+ of an intermediate that every replay overwrites. vLLM warms up on its
550
+ capture streams first, so this only affects captures without a warmup.
551
+ """
552
+ key = (device.index, torch.accelerator.current_stream(device).stream_id)
553
+ counters = _split_counters_by_stream.get(key)
554
+ if counters is None:
555
+ if torch.cuda.is_current_stream_capturing():
556
+ return None
557
+ counters = torch.zeros(_SPLIT_COUNTERS, dtype=torch.int32, device=device)
558
+ _split_counters_by_stream[key] = counters
559
+ return counters
560
+
561
+
562
+def _k_partition(K: int, splits: int, step: int) -> tuple[int, int]:
563
+ """K per split, rounded up to whole steps, and the number of non-empty
564
+ splits."""
565
+ size = triton.cdiv(triton.cdiv(K, splits), step) * step
566
+ return size, triton.cdiv(K, size)
567
+
568
+
569
+def _launch(
570
+ x: torch.Tensor,
571
+ x_scale: torch.Tensor,
572
+ weight: torch.Tensor,
573
+ weight_scale: torch.Tensor,
574
+ out: torch.Tensor,
575
+ cfg: _Config,
576
+) -> None:
577
+ M, K = x.shape
578
+ N = weight.shape[0]
579
+ opts = dict(
580
+ num_warps=cfg.num_warps,
581
+ num_stages=cfg.num_stages,
582
+ matrix_instr_nonkdim=cfg.matrix_instr_nonkdim,
583
+ waves_per_eu=cfg.waves_per_eu,
584
+ )
585
+ grid_m, grid_n = triton.cdiv(M, cfg.block_m), triton.cdiv(N, cfg.block_n)
586
+ tiles = grid_m * grid_n
587
+ step = cfg.block_k * cfg.k_split if cfg.packed else cfg.block_k
588
+ splits = cfg.fused_splits if cfg.fused_splits > 1 or cfg.packed else cfg.k_split
589
+ k_per_split, splits = _k_partition(K, splits, step)
590
+ counters = None
591
+ if cfg.fused_splits > 1 and splits > 1 and tiles <= _SPLIT_COUNTERS:
592
+ counters = _split_counters(out.device)
593
+ if counters is not None:
594
+ fused = splits
595
+ target = out
596
+ slots = torch.empty(
597
+ tiles * fused * cfg.block_m * cfg.block_n,
598
+ dtype=torch.float32,
599
+ device=out.device,
600
+ )
601
+ grid: tuple[int, ...] = (triton.cdiv(tiles, _XCD_STRIDE) * _XCD_STRIDE * fused,)
602
+ else:
603
+ # Without counters, split K through a second launch.
604
+ fused = 1
605
+ target = (
606
+ out
607
+ if splits == 1
608
+ else torch.empty((splits, M, N), dtype=torch.float32, device=x.device)
609
+ )
610
+ slots = counters = target # Unused by the kernels.
611
+ if cfg.packed or not cfg.n_first:
612
+ grid = (grid_m, grid_n, splits)
613
+ else:
614
+ grid = (grid_n, grid_m, splits)
615
+ args = (x, x_scale, weight, weight_scale, target, slots, counters, M, N, K)
616
+ stride = out.stride(0) if splits == 1 or fused > 1 else N
617
+ if cfg.packed:
618
+ _block32_packed_kernel[grid](
619
+ *args,
620
+ stride,
621
+ BLOCK_M=cfg.block_m,
622
+ BLOCK_N=cfg.block_n,
623
+ BLOCK_K=cfg.block_k,
624
+ K_PACK=cfg.k_split,
625
+ K_PER_SPLIT=k_per_split,
626
+ EVEN_K=K % step == 0,
627
+ FUSED_SPLITS=fused,
628
+ **opts,
629
+ )
630
+ else:
631
+ _block32_tiled_kernel[grid](
632
+ *args,
633
+ stride,
634
+ BLOCK_M=cfg.block_m,
635
+ BLOCK_N=cfg.block_n,
636
+ BLOCK_K=cfg.block_k,
637
+ K_PER_SPLIT=k_per_split,
638
+ N_FIRST=cfg.n_first,
639
+ EVEN_K=K % k_per_split == 0,
640
+ FUSED_SPLITS=fused,
641
+ **opts,
642
+ )
643
+ if fused == 1 and splits > 1:
644
+ _block32_splitk_reduce_kernel[(triton.cdiv(M, 32), triton.cdiv(N, 64))](
645
+ target,
646
+ out,
647
+ M,
648
+ N,
649
+ out.stride(0),
650
+ SPLIT_K=splits,
651
+ BLOCK_M=32,
652
+ BLOCK_N=64,
653
+ )
654
+
655
+
656
+def _rocm_mxfp8_block32_gemm_impl(
657
+ x: torch.Tensor,
658
+ x_scale: torch.Tensor,
659
+ weight: torch.Tensor,
660
+ weight_scale: torch.Tensor,
661
+ out_dtype: torch.dtype,
662
+) -> torch.Tensor:
663
+ M, K = x.shape
664
+ out = torch.empty((M, weight.shape[0]), dtype=out_dtype, device=x.device)
665
+ if M > 0:
666
+ _launch(x, x_scale, weight, weight_scale, out, _config(M, weight.shape[0], K))
667
+ return out
668
+
669
+
670
+def _rocm_mxfp8_block32_gemm_fake(
671
+ x: torch.Tensor,
672
+ x_scale: torch.Tensor,
673
+ weight: torch.Tensor,
674
+ weight_scale: torch.Tensor,
675
+ out_dtype: torch.dtype,
676
+) -> torch.Tensor:
677
+ return x.new_empty((x.shape[0], weight.shape[0]), dtype=out_dtype)
678
+
679
+
680
+# An opaque op: the kernel and tile choice branch on the token count, which a
681
+# compiled caller must not freeze at its trace-time value.
682
+direct_register_custom_op(
683
+ op_name="rocm_mxfp8_block32_gemm",
684
+ op_func=_rocm_mxfp8_block32_gemm_impl,
685
+ fake_impl=_rocm_mxfp8_block32_gemm_fake,
686
+)
687
+
688
+
689
+def rocm_mxfp8_block32_gemm(
690
+ x: torch.Tensor,
691
+ x_scale: torch.Tensor,
692
+ weight: torch.Tensor,
693
+ weight_scale: torch.Tensor,
694
+ out_dtype: torch.dtype,
695
+) -> torch.Tensor:
696
+ """``x @ weight.T`` for MXFP8 ``x`` and a 32x32 block-scaled MXFP8 weight.
697
+
698
+ Args:
699
+ x: [M, K] e4m3 activation, contiguous.
700
+ x_scale: [M, K / 32] E8M0 (uint8) activation scales.
701
+ weight: [N, K] e4m3 weight, contiguous.
702
+ weight_scale: [ceil(N / 32), K / 32] E8M0 (uint8) weight block scales.
703
+ out_dtype: Output dtype.
704
+
705
+ Returns:
706
+ The [M, N] product.
707
+
708
+ """
709
+ return torch.ops.vllm.rocm_mxfp8_block32_gemm(
710
+ x, x_scale, weight, weight_scale, out_dtype
711
+ )
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+26
-2
@@ -22,6 +22,7 @@ from vllm.platforms import current_platform
22
from vllm.triton_utils import tl, triton
23
24
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
25
+from .rocm_block32_gemm import BLOCK_ROWS, rocm_mxfp8_block32_gemm
26
27
28
@triton.jit
@@ -209,6 +210,17 @@ def _select_cfg(M, N, K):
210
_DOT_SCALED_K_ALIGN = 128
211
212
213
+def _as_block32_scale(weight_scale: torch.Tensor) -> torch.Tensor | None:
214
+ """The [N / 32, K / 32] block scales, if every 32 rows share their scales."""
215
+ N = weight_scale.shape[0]
216
+ if N % BLOCK_ROWS != 0:
217
+ return None
218
+ blocks = weight_scale.view(N // BLOCK_ROWS, BLOCK_ROWS, -1)
219
+ if not torch.equal(blocks, blocks[:, :1].expand_as(blocks)):
220
+ return None
221
+ return blocks[:, 0].contiguous()
222
+
223
+
224
class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
225
"""Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
226
@@ -236,7 +248,13 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
248
N, K = weight.shape
249
scale_k = K // MXFP8_BLOCK_SIZE
250
weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
239
- if K % _DOT_SCALED_K_ALIGN != 0:
251
+ block_scale = _as_block32_scale(weight_scale)
252
+ if block_scale is not None:
253
+ # Checkpoints with 32x32 scale blocks (DeepSeek V4/V4.1) get them
254
+ # back from the loader expanded per row; keep the blocks and use
255
+ # the block-scaled GEMM, which also handles any K % 32 == 0.
256
+ weight_scale = block_scale
257
+ elif K % _DOT_SCALED_K_ALIGN != 0:
258
weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
259
layer.weight = Parameter(weight.contiguous(), requires_grad=False)
260
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
@@ -254,7 +272,13 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
272
)
273
out_shape = (*x.shape[:-1], layer.weight.shape[0])
274
x2d = x.reshape(-1, x.shape[-1])
257
- if layer.weight.element_size() >= 2:
275
+ if layer.weight_scale.shape[0] != layer.weight.shape[0]:
276
+ # One scale row per 32 weight rows (see process_weights_after_loading).
277
+ x_q, x_scale = mxfp8_e4m3_quantize(x2d)
278
+ out = rocm_mxfp8_block32_gemm(
279
+ x_q, x_scale, layer.weight, layer.weight_scale, x.dtype
280
+ )
281
+ elif layer.weight.element_size() >= 2:
282
out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype))
283
elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0:
284
out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+6
-3
@@ -1811,7 +1811,7 @@ def _mxfp8_wo_a_bmm_kernel(
1811
a_ptr, # [T, G * K] e4m3
1812
as_ptr, # [T, G * K // 32] E8M0
1813
w_ptr, # [G * R, K] e4m3
1814
- ws_ptr, # [G * R, K // 32] E8M0
1814
+ ws_ptr, # [G * R // SCALE_ROWS, K // 32] E8M0
1815
out_ptr, # [T, G * R]
1816
num_tokens,
1817
stride_at,
@@ -1821,6 +1821,7 @@ def _mxfp8_wo_a_bmm_kernel(
1821
stride_out,
1822
R: tl.constexpr,
1823
K: tl.constexpr,
1824
+ SCALE_ROWS: tl.constexpr,
1825
BLOCK_M: tl.constexpr,
1826
BLOCK_N: tl.constexpr,
1827
BLOCK_K: tl.constexpr,
@@ -1841,7 +1842,7 @@ def _mxfp8_wo_a_bmm_kernel(
1842
as_ptr + offs_m[:, None] * stride_ast + (group * (K // 32) + offs_sk)[None, :]
1843
)
1844
w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :]
1844
- ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :]
1845
+ ws_ptrs = ws_ptr + (offs_n // SCALE_ROWS)[:, None] * stride_wsn + offs_sk[None, :]
1846
1847
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
1848
for _ in range(K // BLOCK_K):
@@ -1926,6 +1927,7 @@ def _rocm_mxfp8_wo_a_bmm_impl(
1927
out.stride(0),
1928
R=o_lora_rank,
1929
K=group_dim,
1930
+ SCALE_ROWS=weight.shape[0] // weight_scale.shape[0],
1931
BLOCK_M=block_m,
1932
BLOCK_N=block_n,
1933
BLOCK_K=block_k,
@@ -1967,7 +1969,8 @@ def rocm_mxfp8_wo_a_bmm(
1969
``a`` is the [T, G * K] e4m3 attention output and ``a_scale`` its
1970
[T, G * K // 32] E8M0 scales, as the sparse decode reduce writes them.
1971
The weight is the checkpoint's MXFP8 ``wo_a`` as loaded, [G * R, K] with
1970
- [G * R, K // 32] scales, so there is no dequantized copy to keep.
1972
+ either [G * R // 32, K // 32] block scales or [G * R, K // 32] per-row
1973
+ scales, so there is no dequantized copy to keep.
1974
Returns [T, G * R].
1975
"""
1976
return torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(