@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 737 lines 29.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Portions adapted from ROCm/aiter (https://github.com/ROCm/aiter/pull/5750):
5 # SPDX-License-Identifier: MIT
6 # Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
7 #
8 # Permission is hereby granted, free of charge, to any person obtaining a copy
9 # of this software and associated documentation files (the "Software"), to deal
10 # in the Software without restriction, including without limitation the rights
11 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12 # copies of the Software, and to permit persons to whom the Software is
13 # furnished to do so, subject to the following conditions:
14 #
15 # The above copyright notice and this permission notice shall be included in
16 # all copies or substantial portions of the Software.
17 #
18 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
24 # SOFTWARE.
25 """MXFP8 GEMM on 32x32 block-scaled weights for gfx950 (``tl.dot_scaled``).
26
27 ``y = x @ w.T`` with an MXFP8 activation (e4m3 values, one E8M0 scale per
28 row and 32 K, ``[M, K / 32]``) and an e4m3 weight whose E8M0 scales stay in
29 the checkpoint's 32x32 blocks, ``[N / 32, K / 32]``, instead of being
30 expanded to every row. Each weight scale byte is then read once per 32 output
31 rows, and small-M shapes can use the packed kernel below.
32
33 Two kernels, picked per shape from a table tuned on MI355X:
34
35 * a tiled kernel for larger M;
36 * a packed kernel for small M, which fills the MFMA's rows with K panels
37 instead of tokens and keeps only the matching-panel (block-diagonal)
38 products, so a handful of tokens still streams the weight at full width.
39
40 Either can split K. The partials are normally summed in the same launch by
41 the last program of each output tile to finish, in split order, so the result
42 does not depend on scheduling. Otherwise a second launch reduces them.
43
44 The packed kernel and the in-launch split-K reduction on one XCD
45 (``_split_tile``, ``_sum_splits``, ``_split_counters``, ``_k_partition``) are
46 adapted from the group32 GEMM in ROCm/aiter#5750.
47 """
48
49 from typing import NamedTuple
50
51 import torch
52
53 from vllm.triton_utils import tl, triton
54 from vllm.utils.torch_utils import direct_register_custom_op
55
56 BLOCK_ROWS = 32
57
58 # gfx950 hands out workgroups to its XCDs round-robin by program id, and 8 is
59 # a multiple of the XCD count in every partition mode. _split_tile hardcodes
60 # the same stride.
61 _XCD_STRIDE = 8
62 # One arrival counter per output tile of an in-launch split-K GEMM.
63 _SPLIT_COUNTERS = 1 << 16
64
65
66 @triton.jit
67 def _split_tile(N: tl.constexpr, BLOCK_N: tl.constexpr, SPLITS: tl.constexpr):
68 """(pid_m, pid_n, tile, split) of a 1D grid that keeps a tile's splits on
69 one XCD, so they meet in that XCD's L2."""
70 pid = tl.program_id(0)
71 tile = pid // (8 * SPLITS) * 8 + pid % 8
72 grid_n: tl.constexpr = (N + BLOCK_N - 1) // BLOCK_N
73 return tile // grid_n, tile % grid_n, tile, pid // 8 % SPLITS
74
75
76 @triton.jit
77 def _sum_splits(
78 acc, out_ptrs, out_mask, slot_ptr, count_ptr, split, SPLITS: tl.constexpr
79 ):
80 """Publish this split's partial; the tile's last arrival sums all of them
81 in split order and re-arms the counter."""
82 BM: tl.constexpr = acc.shape[0]
83 BN: tl.constexpr = acc.shape[1]
84 local = tl.arange(0, BM)[:, None] * BN + tl.arange(0, BN)[None, :]
85 tl.store(slot_ptr + split * BM * BN + local, acc)
86 # The partial must be in L2 before the arrival is counted. The splits
87 # share an XCD, so no device-scope release (an L2 writeback) is needed.
88 tl.inline_asm_elementwise(
89 "s_waitcnt vmcnt(0)", "=v,v", [split], dtype=tl.int32, is_pure=False, pack=1
90 )
91 tl.debug_barrier()
92 if tl.atomic_add(count_ptr, 1, sem="acq_rel", scope="cta") == SPLITS - 1:
93 total = tl.zeros((BM, BN), dtype=tl.float32)
94 for s in tl.static_range(SPLITS):
95 total += tl.load(slot_ptr + s * BM * BN + local, cache_modifier=".cv")
96 tl.store(out_ptrs, total.to(out_ptrs.dtype.element_ty), mask=out_mask)
97 tl.store(count_ptr, 0)
98
99
100 @triton.jit(do_not_specialize=["M"])
101 def _block32_tiled_kernel(
102 x_ptr,
103 xs_ptr,
104 w_ptr,
105 ws_ptr,
106 out_ptr,
107 slot_ptr,
108 count_ptr,
109 M,
110 N: tl.constexpr,
111 K: tl.constexpr,
112 stride_om,
113 BLOCK_M: tl.constexpr,
114 BLOCK_N: tl.constexpr,
115 BLOCK_K: tl.constexpr,
116 K_PER_SPLIT: tl.constexpr,
117 N_FIRST: tl.constexpr,
118 EVEN_K: tl.constexpr,
119 FUSED_SPLITS: tl.constexpr,
120 ):
121 if FUSED_SPLITS > 1:
122 pid_m, pid_n, tile, pid_k = _split_tile(N, BLOCK_N, FUSED_SPLITS)
123 if pid_m * BLOCK_M >= M:
124 return
125 else:
126 # N_FIRST launches output-column tiles fastest, so consecutive
127 # programs share an activation tile.
128 if N_FIRST:
129 pid_n, pid_m = tl.program_id(0), tl.program_id(1)
130 else:
131 pid_m, pid_n = tl.program_id(0), tl.program_id(1)
132 pid_k = tl.program_id(2)
133 offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
134 offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
135 m_mask = offs_m < M
136 n_mask = offs_n < N
137 k0 = pid_k * K_PER_SPLIT
138 offs_k = k0 + tl.arange(0, BLOCK_K)
139 offs_sk = k0 // 32 + tl.arange(0, BLOCK_K // 32)
140 x_ptrs = x_ptr + offs_m[:, None] * K + offs_k[None, :]
141 xs_ptrs = xs_ptr + offs_m[:, None] * (K // 32) + offs_sk[None, :]
142 w_ptrs = w_ptr + offs_n[:, None] * K + offs_k[None, :]
143 ws_ptrs = ws_ptr + (offs_n[:, None] // 32) * (K // 32) + offs_sk[None, :]
144 acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
145 for kk in range(0, K_PER_SPLIT, BLOCK_K):
146 if EVEN_K:
147 x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
148 xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
149 w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
150 ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
151 else:
152 k_ok = (offs_k + kk) < K
153 s_ok = (offs_sk + kk // 32) < K // 32
154 x = tl.load(x_ptrs, mask=m_mask[:, None] & k_ok[None, :], other=0.0)
155 xs = tl.load(xs_ptrs, mask=m_mask[:, None] & s_ok[None, :], other=0)
156 w = tl.load(w_ptrs, mask=n_mask[:, None] & k_ok[None, :], other=0.0)
157 ws = tl.load(ws_ptrs, mask=n_mask[:, None] & s_ok[None, :], other=0)
158 acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
159 x_ptrs += BLOCK_K
160 w_ptrs += BLOCK_K
161 xs_ptrs += BLOCK_K // 32
162 ws_ptrs += BLOCK_K // 32
163 out_mask = m_mask[:, None] & n_mask[None, :]
164 if FUSED_SPLITS > 1:
165 _sum_splits(
166 acc,
167 out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
168 out_mask,
169 slot_ptr + tile * (FUSED_SPLITS * BLOCK_M * BLOCK_N),
170 count_ptr + tile,
171 pid_k,
172 FUSED_SPLITS,
173 )
174 else:
175 out_ptrs = (
176 out_ptr + pid_k * M * N + offs_m[:, None] * stride_om + offs_n[None, :]
177 )
178 tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=out_mask)
179
180
181 @triton.jit(do_not_specialize=["M"])
182 def _block32_packed_kernel(
183 x_ptr,
184 xs_ptr,
185 w_ptr,
186 ws_ptr,
187 out_ptr,
188 slot_ptr,
189 count_ptr,
190 M,
191 N: tl.constexpr,
192 K: tl.constexpr,
193 stride_om,
194 BLOCK_M: tl.constexpr,
195 BLOCK_N: tl.constexpr,
196 BLOCK_K: tl.constexpr,
197 K_PACK: tl.constexpr,
198 K_PER_SPLIT: tl.constexpr,
199 EVEN_K: tl.constexpr,
200 FUSED_SPLITS: tl.constexpr,
201 ):
202 if FUSED_SPLITS > 1:
203 pid_m, pid_n, tile, pid_k = _split_tile(N, BLOCK_N, FUSED_SPLITS)
204 if pid_m * BLOCK_M >= M:
205 return
206 k_start = pid_k * K_PER_SPLIT
207 k_stop = tl.minimum(k_start + K_PER_SPLIT, K)
208 else:
209 pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2)
210 if K_PER_SPLIT >= K:
211 k_start = 0
212 k_stop = K
213 else:
214 k_start = pid_k * K_PER_SPLIT
215 k_stop = tl.minimum(k_start + K_PER_SPLIT, K)
216 # Tile row r is token r // K_PACK reading K panel r % K_PACK (and likewise
217 # for weight rows), so one dot covers K_PACK consecutive BLOCK_K slices.
218 rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M * K_PACK) // K_PACK
219 x_panel = tl.arange(0, BLOCK_M * K_PACK) % K_PACK
220 cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N * K_PACK) // K_PACK
221 w_panel = tl.arange(0, BLOCK_N * K_PACK) % K_PACK
222 ks = tl.arange(0, BLOCK_K)
223 gs = tl.arange(0, BLOCK_K // 32)
224 row_ok = rows < M
225 col_ok = cols < N
226 acc = tl.zeros((BLOCK_M * K_PACK, BLOCK_N * K_PACK), dtype=tl.float32)
227 for base in range(k_start, k_stop, BLOCK_K * K_PACK):
228 xk = base + x_panel[:, None] * BLOCK_K + ks[None, :]
229 wk = base + w_panel[:, None] * BLOCK_K + ks[None, :]
230 xg = base // 32 + x_panel[:, None] * (BLOCK_K // 32) + gs[None, :]
231 wg = base // 32 + w_panel[:, None] * (BLOCK_K // 32) + gs[None, :]
232 x_mask = row_ok[:, None]
233 w_mask = col_ok[:, None]
234 xs_mask = row_ok[:, None]
235 ws_mask = col_ok[:, None]
236 if not EVEN_K:
237 x_mask = x_mask & (xk < K)
238 w_mask = w_mask & (wk < K)
239 xs_mask = xs_mask & (xg < K // 32)
240 ws_mask = ws_mask & (wg < K // 32)
241 x = tl.load(x_ptr + rows[:, None] * K + xk, mask=x_mask, other=0.0)
242 w = tl.load(w_ptr + cols[:, None] * K + wk, mask=w_mask, other=0.0)
243 xs = tl.load(xs_ptr + rows[:, None] * (K // 32) + xg, mask=xs_mask, other=0)
244 ws = tl.load(
245 ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=0
246 )
247 acc = tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3", acc=acc)
248 # Keep only products of matching K panels: the block diagonal.
249 blocks = acc.reshape(BLOCK_M, K_PACK, BLOCK_N, K_PACK).trans(0, 2, 1, 3)
250 pair = tl.arange(0, K_PACK)
251 diagonal = tl.where(
252 pair[None, None, :, None] == pair[None, None, None, :], blocks, 0.0
253 )
254 y = tl.sum(tl.sum(diagonal, 3), 2)
255 out_rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
256 out_cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
257 out_ptrs = out_ptr + out_rows[:, None] * stride_om + out_cols[None, :]
258 out_mask = (out_rows[:, None] < M) & (out_cols[None, :] < N)
259 if FUSED_SPLITS > 1:
260 _sum_splits(
261 y,
262 out_ptrs,
263 out_mask,
264 slot_ptr + tile * (FUSED_SPLITS * BLOCK_M * BLOCK_N),
265 count_ptr + tile,
266 pid_k,
267 FUSED_SPLITS,
268 )
269 else:
270 tl.store(
271 out_ptrs + pid_k * M * N, y.to(out_ptr.dtype.element_ty), mask=out_mask
272 )
273
274
275 @triton.jit(do_not_specialize=["M"])
276 def _block32_splitk_reduce_kernel(
277 partial_ptr,
278 out_ptr,
279 M,
280 N: tl.constexpr,
281 stride_om,
282 SPLIT_K: tl.constexpr,
283 BLOCK_M: tl.constexpr,
284 BLOCK_N: tl.constexpr,
285 ):
286 offs_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
287 offs_n = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
288 mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
289 acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
290 for split in tl.static_range(SPLIT_K):
291 acc += tl.load(
292 partial_ptr + split * M * N + offs_m[:, None] * N + offs_n[None, :],
293 mask=mask,
294 other=0.0,
295 )
296 tl.store(
297 out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
298 acc.to(out_ptr.dtype.element_ty),
299 mask=mask,
300 )
301
302
303 class _Config(NamedTuple):
304 packed: bool
305 block_m: int
306 block_n: int
307 block_k: int
308 # K panels per MFMA tile (packed) or K partitions reduced by a second
309 # launch (tiled).
310 k_split: int
311 n_first: bool
312 num_warps: int
313 num_stages: int
314 matrix_instr_nonkdim: int
315 waves_per_eu: int
316 # K partitions summed within the launch.
317 fused_splits: int = 1
318
319
320 def _tiled(
321 bm, bn, bk, split_k, n_first, warps, stages, nonkdim, waves, fused=1
322 ) -> _Config:
323 return _Config(
324 False, bm, bn, bk, split_k, bool(n_first), warps, stages, nonkdim, waves, fused
325 )
326
327
328 def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Config:
329 return _Config(
330 True, bm, bn, bk, k_pack, False, warps, stages, nonkdim, waves, fused
331 )
332
333
334 # (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8 linears at
335 # TP2/TP4, tuned on MI355X under HIP graphs with weights cold in the LLC. A
336 # tier serves every M up to its bound; larger M reuse the last tier.
337 _TUNED: dict[tuple[int, int], list[tuple[int, _Config]]] = {
338 # shared expert gate_up, TP4
339 (1152, 5120): [
340 (1, _packed(8, 16, 1024, 2, 2, 2, 16, 1, 3)),
341 (2, _packed(4, 32, 256, 4, 2, 2, 16, 1, 5)),
342 (4, _packed(4, 32, 256, 4, 2, 2, 16, 0, 5)),
343 (8, _packed(8, 16, 1024, 2, 2, 2, 16, 0, 3)),
344 (16, _tiled(16, 32, 512, 1, True, 2, 3, 16, 0, 5)),
345 (24, _tiled(32, 32, 512, 1, True, 4, 3, 16, 1, 5)),
346 (32, _packed(32, 32, 1024, 1, 4, 2, 16, 0, 5)),
347 (48, _tiled(64, 32, 512, 1, True, 4, 3, 16, 1, 5)),
348 (64, _tiled(64, 32, 512, 1, True, 8, 3, 16, 1, 5)),
349 (96, _tiled(32, 16, 1024, 1, True, 2, 3, 16, 1)),
350 (128, _tiled(64, 32, 512, 1, True, 8, 2, 16, 0, 3)),
351 (160, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0, 2)),
352 (192, _tiled(32, 32, 1024, 1, True, 4, 2, 16, 0)),
353 (256, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 3)),
354 (320, _tiled(64, 32, 256, 1, True, 8, 3, 16, 0, 2)),
355 (384, _tiled(32, 32, 512, 1, True, 4, 2, 16, 0)),
356 (512, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
357 (768, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
358 (1024, _tiled(128, 128, 256, 2, False, 4, 2, 16, 1)),
359 (1536, _tiled(128, 128, 256, 2, False, 4, 2, 16, 0)),
360 (2048, _tiled(128, 128, 256, 1, False, 4, 2, 32, 0)),
361 (3072, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
362 (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
363 (8192, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
364 ],
365 # fused wq_a|wkv (replicated)
366 (1792, 5120): [
367 (1, _packed(8, 16, 1024, 2, 2, 3, 16, 0)),
368 (8, _packed(8, 16, 1024, 2, 2, 3, 16, 1)),
369 (16, _tiled(16, 64, 1024, 1, True, 4, 2, 16, 0, 5)),
370 (32, _packed(16, 16, 1024, 1, 2, 3, 16, 1)),
371 (48, _tiled(64, 32, 512, 1, True, 8, 2, 16, 0, 4)),
372 (64, _tiled(64, 32, 512, 1, True, 8, 3, 16, 1, 4)),
373 (96, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0, 2)),
374 (128, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 4)),
375 (192, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 3)),
376 (256, _tiled(64, 64, 512, 1, True, 4, 2, 16, 1, 2)),
377 (384, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
378 (512, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
379 (768, _tiled(128, 128, 256, 1, True, 4, 2, 16, 1, 2)),
380 (1024, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
381 (2048, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
382 (3072, _tiled(128, 128, 256, 1, False, 8, 2, 32, 1)),
383 (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
384 (8192, _tiled(128, 128, 128, 1, False, 8, 2, 16, 1)),
385 ],
386 # shared expert gate_up, TP2
387 (2304, 5120): [
388 (1, _packed(8, 16, 1024, 2, 2, 3, 16, 1)),
389 (8, _packed(8, 16, 1024, 2, 2, 3, 16, 0)),
390 (16, _tiled(16, 64, 1024, 1, True, 4, 3, 16, 1, 5)),
391 (24, _packed(32, 32, 512, 1, 4, 2, 16, 0, 5)),
392 (32, _packed(32, 32, 1024, 1, 4, 2, 16, 1, 3)),
393 (48, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0, 3)),
394 (64, _tiled(64, 64, 512, 1, True, 8, 2, 16, 0, 5)),
395 (128, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0, 3)),
396 (160, _tiled(64, 64, 512, 1, True, 8, 2, 16, 1, 2)),
397 (192, _tiled(64, 64, 512, 1, True, 4, 2, 16, 0, 2)),
398 (256, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
399 (320, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
400 (384, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
401 (512, _tiled(128, 128, 256, 1, True, 8, 2, 16, 0, 3)),
402 (768, _tiled(128, 128, 256, 1, True, 4, 2, 16, 1, 2)),
403 (1536, _tiled(128, 128, 256, 1, False, 4, 2, 32, 0)),
404 (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
405 (3072, _tiled(128, 128, 128, 1, True, 8, 2, 16, 0)),
406 (4096, _tiled(128, 128, 256, 1, False, 4, 2, 32, 1)),
407 (8192, _tiled(128, 128, 128, 1, False, 8, 2, 32, 0)),
408 ],
409 # indexer wq_b
410 (4096, 1280): [
411 (1, _packed(4, 16, 256, 4, 4, 2, 16, 1)),
412 (2, _packed(4, 16, 256, 4, 4, 2, 16, 0)),
413 (4, _packed(4, 16, 256, 4, 4, 2, 16, 1)),
414 (8, _packed(8, 16, 256, 4, 4, 3, 16, 1)),
415 (16, _packed(16, 16, 512, 2, 2, 3, 16, 1)),
416 (32, _packed(32, 16, 1024, 1, 2, 3, 16, 1)),
417 (64, _tiled(32, 32, 512, 1, True, 4, 3, 16, 1)),
418 (128, _tiled(64, 32, 512, 1, True, 8, 3, 16, 0)),
419 (192, _tiled(64, 32, 256, 1, True, 4, 3, 16, 0)),
420 (256, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
421 (384, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
422 (512, _tiled(128, 64, 256, 1, True, 8, 3, 32, 0)),
423 (1024, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
424 (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
425 (3072, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
426 (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
427 (8192, _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)),
428 ],
429 # shared expert down, TP4
430 (5120, 576): [
431 (8, _packed(8, 32, 256, 2, 2, 2, 16, 1)),
432 (16, _packed(16, 32, 512, 1, 2, 2, 16, 0)),
433 (32, _packed(32, 32, 256, 1, 8, 3, 16, 1)),
434 (64, _tiled(32, 32, 128, 1, True, 4, 3, 16, 0)),
435 (96, _tiled(32, 64, 256, 1, True, 4, 3, 16, 0)),
436 (128, _tiled(32, 32, 256, 1, True, 2, 2, 16, 0)),
437 (192, _tiled(64, 64, 256, 1, True, 4, 3, 32, 0)),
438 (256, _tiled(64, 64, 256, 1, True, 2, 3, 32, 1)),
439 (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
440 (384, _tiled(128, 64, 256, 1, True, 8, 3, 32, 1)),
441 (512, _tiled(64, 64, 128, 1, True, 4, 2, 32, 0)),
442 (768, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
443 (1024, _tiled(128, 64, 128, 1, True, 4, 2, 32, 0)),
444 (3072, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
445 (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
446 (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
447 ],
448 # shared expert down, TP2
449 (5120, 1152): [
450 (1, _packed(4, 32, 128, 4, 2, 3, 16, 1)),
451 (8, _packed(8, 32, 256, 2, 4, 3, 16, 0)),
452 (32, _packed(32, 32, 512, 1, 2, 3, 16, 1)),
453 (128, _tiled(32, 32, 128, 1, True, 4, 3, 16, 0)),
454 (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
455 (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
456 (384, _tiled(128, 64, 256, 1, True, 8, 3, 32, 1)),
457 (512, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
458 (768, _tiled(128, 128, 256, 1, True, 4, 2, 32, 0)),
459 (1024, _tiled(128, 64, 128, 1, True, 8, 2, 32, 0)),
460 (1536, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
461 (2048, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
462 (3072, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)),
463 (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
464 (8192, _tiled(128, 64, 128, 1, True, 8, 2, 32, 1)),
465 ],
466 # wo_b, TP4
467 (5120, 2048): [
468 (1, _packed(4, 32, 256, 4, 2, 2, 16, 0)),
469 (4, _packed(4, 32, 256, 4, 2, 2, 16, 1)),
470 (8, _packed(8, 32, 256, 4, 2, 3, 16, 0)),
471 (16, _tiled(16, 32, 1024, 1, True, 2, 3, 16, 0)),
472 (32, _tiled(32, 32, 512, 1, False, 4, 3, 16, 1)),
473 (64, _tiled(32, 32, 512, 1, True, 4, 2, 16, 0)),
474 (96, _tiled(32, 64, 512, 1, True, 8, 2, 16, 0)),
475 (128, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
476 (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
477 (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
478 (384, _tiled(128, 64, 256, 1, True, 8, 3, 32, 0)),
479 (768, _tiled(128, 128, 256, 1, True, 4, 2, 32, 0)),
480 (1024, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
481 (1536, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
482 (2048, _tiled(128, 128, 256, 1, True, 8, 2, 32, 1)),
483 (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
484 (8192, _tiled(128, 128, 128, 1, True, 2, 2, 32, 1)),
485 ],
486 # wo_b, TP2
487 (5120, 4096): [
488 (4, _packed(4, 32, 512, 4, 2, 2, 16, 0)),
489 (8, _packed(8, 16, 128, 4, 2, 2, 16, 0, 2)),
490 (16, _tiled(16, 32, 256, 1, True, 2, 2, 16, 0, 4)),
491 (24, _tiled(32, 32, 256, 1, True, 4, 3, 16, 0, 4)),
492 (32, _tiled(32, 64, 512, 1, True, 8, 3, 16, 0, 3)),
493 (48, _tiled(64, 64, 512, 1, True, 8, 2, 16, 1, 3)),
494 (64, _tiled(64, 64, 512, 1, True, 8, 2, 16, 0, 3)),
495 (128, _tiled(64, 64, 256, 1, True, 4, 2, 16, 0, 3)),
496 (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 0)),
497 (256, _tiled(128, 128, 256, 1, True, 4, 2, 16, 1, 2)),
498 (320, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
499 (384, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0, 2)),
500 (768, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
501 (1024, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
502 (1536, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
503 (2048, _tiled(128, 128, 256, 1, True, 8, 2, 32, 1)),
504 (3072, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
505 (4096, _tiled(128, 128, 128, 1, True, 8, 2, 16, 1)),
506 (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
507 ],
508 # wq_b, TP4
509 (8192, 1280): [
510 (1, _packed(8, 32, 256, 2, 4, 3, 16, 1)),
511 (8, _packed(8, 32, 256, 2, 4, 3, 16, 0)),
512 (16, _packed(16, 32, 1024, 1, 2, 3, 16, 1)),
513 (32, _packed(32, 32, 512, 1, 2, 3, 16, 0)),
514 (48, _tiled(64, 32, 512, 1, True, 8, 3, 16, 1)),
515 (64, _packed(32, 64, 256, 1, 4, 3, 16, 1)),
516 (128, _tiled(64, 64, 512, 1, True, 4, 2, 32, 1)),
517 (192, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0)),
518 (256, _tiled(128, 64, 256, 1, True, 8, 3, 32, 1)),
519 (384, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
520 (512, _tiled(128, 128, 256, 1, True, 4, 2, 32, 0)),
521 (1024, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
522 (1536, _tiled(128, 128, 128, 1, True, 8, 2, 16, 0)),
523 (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
524 (8192, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)),
525 ],
526 # wq_b, TP2
527 (16384, 1280): [
528 (2, _packed(4, 32, 128, 4, 2, 3, 16, 1)),
529 (4, _packed(4, 32, 128, 4, 4, 2, 16, 0)),
530 (8, _packed(8, 32, 256, 2, 4, 2, 16, 0)),
531 (32, _packed(32, 64, 256, 1, 4, 3, 16, 1)),
532 (64, _tiled(64, 64, 256, 1, False, 4, 3, 32, 1)),
533 (128, _tiled(128, 64, 256, 1, True, 8, 3, 32, 0)),
534 (256, _tiled(128, 128, 256, 1, True, 4, 2, 32, 1)),
535 (384, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
536 (512, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)),
537 (768, _tiled(128, 128, 128, 1, True, 8, 2, 16, 1)),
538 (1024, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)),
539 (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)),
540 (8192, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)),
541 ],
542 }
543
544
545 def _default_config(M: int, N: int, K: int) -> _Config:
546 """Untuned shapes: the tiers most tuned shapes settle on."""
547 if M <= 8:
548 return _packed(8, 16, 256, 2, 2, 3, 16, 0)
549 if M <= 32:
550 return _packed(32, 32, 256, 1, 2, 3, 16, 0)
551 if M <= 256:
552 return _tiled(64, 64, 256, 1, True, 4, 2, 16, 0)
553 return _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)
554
555
556 def _config(M: int, N: int, K: int) -> _Config:
557 tiers = _TUNED.get((N, K))
558 if tiers is None:
559 return _default_config(M, N, K)
560 for max_m, cfg in tiers:
561 if max_m >= M:
562 return cfg
563 return tiers[-1][1]
564
565
566 _split_counters_by_stream: dict[tuple[int, int], torch.Tensor] = {}
567
568
569 def _split_counters(device: torch.device) -> torch.Tensor | None:
570 """Zeroed arrival counters, one set per stream so concurrent launches never
571 share a tile's counter. Each launch leaves the counters it used at zero.
572
573 None when a stream first needs them inside a CUDA graph capture: that
574 allocation would come from the graph's pool, where it can take the address
575 of an intermediate that every replay overwrites. vLLM warms up on its
576 capture streams first, so this only affects captures without a warmup.
577 """
578 key = (device.index, torch.accelerator.current_stream(device).stream_id)
579 counters = _split_counters_by_stream.get(key)
580 if counters is None:
581 if torch.cuda.is_current_stream_capturing():
582 return None
583 counters = torch.zeros(_SPLIT_COUNTERS, dtype=torch.int32, device=device)
584 _split_counters_by_stream[key] = counters
585 return counters
586
587
588 def _k_partition(K: int, splits: int, step: int) -> tuple[int, int]:
589 """K per split, rounded up to whole steps, and the number of non-empty
590 splits."""
591 size = triton.cdiv(triton.cdiv(K, splits), step) * step
592 return size, triton.cdiv(K, size)
593
594
595 def _launch(
596 x: torch.Tensor,
597 x_scale: torch.Tensor,
598 weight: torch.Tensor,
599 weight_scale: torch.Tensor,
600 out: torch.Tensor,
601 cfg: _Config,
602 ) -> None:
603 M, K = x.shape
604 N = weight.shape[0]
605 opts = dict(
606 num_warps=cfg.num_warps,
607 num_stages=cfg.num_stages,
608 matrix_instr_nonkdim=cfg.matrix_instr_nonkdim,
609 waves_per_eu=cfg.waves_per_eu,
610 )
611 grid_m, grid_n = triton.cdiv(M, cfg.block_m), triton.cdiv(N, cfg.block_n)
612 tiles = grid_m * grid_n
613 step = cfg.block_k * cfg.k_split if cfg.packed else cfg.block_k
614 splits = cfg.fused_splits if cfg.fused_splits > 1 or cfg.packed else cfg.k_split
615 k_per_split, splits = _k_partition(K, splits, step)
616 counters = None
617 if cfg.fused_splits > 1 and splits > 1 and tiles <= _SPLIT_COUNTERS:
618 counters = _split_counters(out.device)
619 if counters is not None:
620 fused = splits
621 target = out
622 slots = torch.empty(
623 tiles * fused * cfg.block_m * cfg.block_n,
624 dtype=torch.float32,
625 device=out.device,
626 )
627 grid: tuple[int, ...] = (triton.cdiv(tiles, _XCD_STRIDE) * _XCD_STRIDE * fused,)
628 else:
629 # Without counters, split K through a second launch.
630 fused = 1
631 target = (
632 out
633 if splits == 1
634 else torch.empty((splits, M, N), dtype=torch.float32, device=x.device)
635 )
636 slots = counters = target # Unused by the kernels.
637 if cfg.packed or not cfg.n_first:
638 grid = (grid_m, grid_n, splits)
639 else:
640 grid = (grid_n, grid_m, splits)
641 args = (x, x_scale, weight, weight_scale, target, slots, counters, M, N, K)
642 stride = out.stride(0) if splits == 1 or fused > 1 else N
643 if cfg.packed:
644 _block32_packed_kernel[grid](
645 *args,
646 stride,
647 BLOCK_M=cfg.block_m,
648 BLOCK_N=cfg.block_n,
649 BLOCK_K=cfg.block_k,
650 K_PACK=cfg.k_split,
651 K_PER_SPLIT=k_per_split,
652 EVEN_K=K % step == 0,
653 FUSED_SPLITS=fused,
654 **opts,
655 )
656 else:
657 _block32_tiled_kernel[grid](
658 *args,
659 stride,
660 BLOCK_M=cfg.block_m,
661 BLOCK_N=cfg.block_n,
662 BLOCK_K=cfg.block_k,
663 K_PER_SPLIT=k_per_split,
664 N_FIRST=cfg.n_first,
665 EVEN_K=K % k_per_split == 0,
666 FUSED_SPLITS=fused,
667 **opts,
668 )
669 if fused == 1 and splits > 1:
670 _block32_splitk_reduce_kernel[(triton.cdiv(M, 32), triton.cdiv(N, 64))](
671 target,
672 out,
673 M,
674 N,
675 out.stride(0),
676 SPLIT_K=splits,
677 BLOCK_M=32,
678 BLOCK_N=64,
679 )
680
681
682 def _rocm_mxfp8_block32_gemm_impl(
683 x: torch.Tensor,
684 x_scale: torch.Tensor,
685 weight: torch.Tensor,
686 weight_scale: torch.Tensor,
687 out_dtype: torch.dtype,
688 ) -> torch.Tensor:
689 M, K = x.shape
690 out = torch.empty((M, weight.shape[0]), dtype=out_dtype, device=x.device)
691 if M > 0:
692 _launch(x, x_scale, weight, weight_scale, out, _config(M, weight.shape[0], K))
693 return out
694
695
696 def _rocm_mxfp8_block32_gemm_fake(
697 x: torch.Tensor,
698 x_scale: torch.Tensor,
699 weight: torch.Tensor,
700 weight_scale: torch.Tensor,
701 out_dtype: torch.dtype,
702 ) -> torch.Tensor:
703 return x.new_empty((x.shape[0], weight.shape[0]), dtype=out_dtype)
704
705
706 # An opaque op: the kernel and tile choice branch on the token count, which a
707 # compiled caller must not freeze at its trace-time value.
708 direct_register_custom_op(
709 op_name="rocm_mxfp8_block32_gemm",
710 op_func=_rocm_mxfp8_block32_gemm_impl,
711 fake_impl=_rocm_mxfp8_block32_gemm_fake,
712 )
713
714
715 def rocm_mxfp8_block32_gemm(
716 x: torch.Tensor,
717 x_scale: torch.Tensor,
718 weight: torch.Tensor,
719 weight_scale: torch.Tensor,
720 out_dtype: torch.dtype,
721 ) -> torch.Tensor:
722 """``x @ weight.T`` for MXFP8 ``x`` and a 32x32 block-scaled MXFP8 weight.
723
724 Args:
725 x: [M, K] e4m3 activation, contiguous.
726 x_scale: [M, K / 32] E8M0 (uint8) activation scales.
727 weight: [N, K] e4m3 weight, contiguous.
728 weight_scale: [ceil(N / 32), K / 32] E8M0 (uint8) weight block scales.
729 out_dtype: Output dtype.
730
731 Returns:
732 The [M, N] product.
733
734 """
735 return torch.ops.vllm.rocm_mxfp8_block32_gemm(
736 x, x_scale, weight, weight_scale, out_dtype
737 )