| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``. |
| 4 | |
| 5 | Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16); |
| 6 | activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling |
| 7 | matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple |
| 8 | of that is dequantized to BF16 once at load and served by a plain linear |
| 9 | instead, since ``can_implement`` does not filter on K. |
| 10 | """ |
| 11 | |
| 12 | import torch |
| 13 | from torch.nn.parameter import Parameter |
| 14 | |
| 15 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 16 | MXFP8_BLOCK_SIZE, |
| 17 | MXFP8_SCALE_DTYPE, |
| 18 | dequant_mxfp8_to_bf16, |
| 19 | mxfp8_e4m3_quantize, |
| 20 | ) |
| 21 | 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 |
| 29 | def _mxfp8_linear_kernel( |
| 30 | x_ptr, |
| 31 | xs_ptr, |
| 32 | w_ptr, |
| 33 | ws_ptr, |
| 34 | out_ptr, |
| 35 | M, |
| 36 | N, |
| 37 | K, |
| 38 | stride_xm, |
| 39 | stride_xk, |
| 40 | stride_xsm, |
| 41 | stride_xsk, |
| 42 | stride_wn, |
| 43 | stride_wk, |
| 44 | stride_wsn, |
| 45 | stride_wsk, |
| 46 | stride_om, |
| 47 | stride_on, |
| 48 | BLOCK_M: tl.constexpr, |
| 49 | BLOCK_N: tl.constexpr, |
| 50 | BLOCK_K: tl.constexpr, |
| 51 | ): |
| 52 | pid_m = tl.program_id(0) |
| 53 | pid_n = tl.program_id(1) |
| 54 | offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) |
| 55 | offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) |
| 56 | offs_k = tl.arange(0, BLOCK_K) |
| 57 | offs_sk = tl.arange(0, BLOCK_K // 32) |
| 58 | m_mask = offs_m < M |
| 59 | n_mask = offs_n < N |
| 60 | |
| 61 | x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk |
| 62 | xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk |
| 63 | w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk |
| 64 | ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk |
| 65 | |
| 66 | acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) |
| 67 | for _ in range(0, tl.cdiv(K, BLOCK_K)): |
| 68 | x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0) |
| 69 | w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0) |
| 70 | xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0) |
| 71 | ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0) |
| 72 | acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3") |
| 73 | x_ptrs += BLOCK_K * stride_xk |
| 74 | w_ptrs += BLOCK_K * stride_wk |
| 75 | xs_ptrs += (BLOCK_K // 32) * stride_xsk |
| 76 | ws_ptrs += (BLOCK_K // 32) * stride_wsk |
| 77 | |
| 78 | o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on |
| 79 | tl.store( |
| 80 | o_ptrs, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :] |
| 81 | ) |
| 82 | |
| 83 | |
| 84 | def _mxfp8_dot_scaled_linear( |
| 85 | x: torch.Tensor, # [M, K] bf16/fp16 |
| 86 | w: torch.Tensor, # [N, K] fp8 e4m3 |
| 87 | w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0) |
| 88 | ) -> torch.Tensor: |
| 89 | M, K = x.shape |
| 90 | N = w.shape[0] |
| 91 | x_q, x_scale = mxfp8_e4m3_quantize(x) |
| 92 | out = torch.empty((M, N), dtype=x.dtype, device=x.device) |
| 93 | BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K) |
| 94 | grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) |
| 95 | _mxfp8_linear_kernel[grid]( |
| 96 | x_q, |
| 97 | x_scale, |
| 98 | w, |
| 99 | w_scale, |
| 100 | out, |
| 101 | M, |
| 102 | N, |
| 103 | K, |
| 104 | x_q.stride(0), |
| 105 | x_q.stride(1), |
| 106 | x_scale.stride(0), |
| 107 | x_scale.stride(1), |
| 108 | w.stride(0), |
| 109 | w.stride(1), |
| 110 | w_scale.stride(0), |
| 111 | w_scale.stride(1), |
| 112 | out.stride(0), |
| 113 | out.stride(1), |
| 114 | BLOCK_M=BLOCK_M, |
| 115 | BLOCK_N=BLOCK_N, |
| 116 | BLOCK_K=BLOCK_K, |
| 117 | num_warps=num_warps, |
| 118 | num_stages=num_stages, |
| 119 | ) |
| 120 | return out |
| 121 | |
| 122 | |
| 123 | # gfx950's async-copy LDS (default in the triton 3.8 this branch pins) leaves no |
| 124 | # room for num_stages=3 at BLOCK_K=256; use 2 there. |
| 125 | _BK256_STAGES = 3 |
| 126 | if current_platform.is_rocm(): |
| 127 | from vllm.platforms.rocm import on_gfx950 |
| 128 | |
| 129 | if on_gfx950(): |
| 130 | _BK256_STAGES = 2 |
| 131 | |
| 132 | |
| 133 | def _select_cfg(M, N, K): |
| 134 | """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950. |
| 135 | |
| 136 | The M-bucketed, shape-adaptive tile selection here is the speedup over the |
| 137 | upstream 2-bucket launcher. Tiles are pipelined (num_stages>=2, larger BLOCK_K) |
| 138 | and occupancy- and shape-aware: keyed on the LOCAL (M, N, K), so it adapts to the |
| 139 | TP-sharded shapes (e.g. MiniMax-M3 TP=4 vs TP=8, where local N and K differ) — |
| 140 | large-K prefill uses BLOCK_K=256; short-K (K=768) widens N. BLOCK_K must divide K |
| 141 | (the K-loop is unmasked), so every BLOCK_K below is guarded to be K-divisible |
| 142 | (served K: 384/768/1024/2048/6144). |
| 143 | """ |
| 144 | if M <= 64: |
| 145 | # decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The |
| 146 | # lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in |
| 147 | # parallel) + LARGE BLOCK_K (fewer K-iters, bigger coalesced weight loads). |
| 148 | # Tuned by CUDA-graph replay latency. Optimal at both TP=4 and TP=8. |
| 149 | if K % 1024 == 0: # K=2048, 6144 -> graph-best 16x16x1024 (all M) |
| 150 | return 16, 16, 1024, 2, 2 |
| 151 | if K % 512 == 0: |
| 152 | return 16, 16, 512, 2, 3 |
| 153 | if K % 256 == 0: # K=768 (shared_down) -> graph-best 16x32x256 |
| 154 | return 16, 32, 256, 4, 3 |
| 155 | return 16, 32, 128, 4, 3 |
| 156 | # mid-M (65..256) on SMALL local-N: still occupancy-bound (a 64x64 tile makes too |
| 157 | # few N-tiles), so the narrow-BLOCK_N decode-style tile fills the CUs better. |
| 158 | # N<=1536 covers the real fused-qkv local N at TP=8: q heads shard but the GQA KV |
| 159 | # (4) + sparse-indexer (4) heads are < TP=8, so vLLM replicates them to 1/rank -> |
| 160 | # N = 1024 + 4*128 = 1536 (not 2560/2=1280). For the wider 1280<N<=1536 band the |
| 161 | # narrow tile only wins up to M=128 (at M=256 the 64x64 tile is better), so cap it |
| 162 | # there; N<=1280 keeps the narrow tile through M=256. TP=4 qkv N=2560 is unchanged. |
| 163 | if (M <= 256 and N <= 1280) or (M <= 128 and N <= 1536): |
| 164 | if K % 1024 == 0: |
| 165 | return 16, 16, 1024, 2, 2 |
| 166 | if K % 512 == 0: |
| 167 | return 16, 16, 512, 2, 3 |
| 168 | if K % 256 == 0: |
| 169 | return 16, 16, 256, 2, 3 |
| 170 | return 16, 16, 128, 2, 3 |
| 171 | # right-sized launch grid (host-side), used to gate the tall 256-BLOCK_M tile. |
| 172 | occ = triton.cdiv(M, 256) * triton.cdiv(N, 128) |
| 173 | if K <= 1024: # short-K (shared_down K=384/768; TP=8 o_proj K=1024) |
| 174 | if M <= 256: |
| 175 | return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2) |
| 176 | # large prefill. (The former 256x128x256 tall tile was faster only on triton |
| 177 | # 3.6; on triton 3.7 its large BLOCK_M register/LDS footprint spills or hits |
| 178 | # "out of resources", so use 128x128x256 -- within the known-good footprint.) |
| 179 | if M >= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256: |
| 180 | return 128, 128, 256, 8, _BK256_STAGES |
| 181 | return 128, 128, 128, 8, 3 |
| 182 | # large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked): |
| 183 | # served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since |
| 184 | # the entry requires K%128==0) for any other K to stay correct. |
| 185 | if M <= 256: # conc~128 decode + small prefill chunk: occupancy tile |
| 186 | if K % 512 == 0: |
| 187 | return 64, 64, 512, 8, 2 |
| 188 | return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2) |
| 189 | if M <= 1024: # medium prefill chunk: BN=64 keeps small-N occupied |
| 190 | return (128, 64, 256, 8, 3) if K % 256 == 0 else (128, 64, 128, 8, 3) |
| 191 | # large prefill (M > 1024). The previously graph-tuned tall 256x128x256 and deep |
| 192 | # 128x128x512 tiles won only on triton 3.6; on triton 3.7 their larger BLOCK_M / |
| 193 | # BLOCK_K register+LDS footprint spills (or hits "out of resources" on stricter |
| 194 | # ROCm/triton builds). The 128x128x256 tile is equal-or-faster on triton 3.7 (the |
| 195 | # M=4096,N=2048,K=6144 shape: ~104us vs the 256x128x256 tile's ~184us), within ~5% |
| 196 | # on 3.6, and inside the footprint of the tiles used elsewhere in this selector. |
| 197 | # Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K |
| 198 | # / very-large-M shapes. |
| 199 | if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))): |
| 200 | return 128, 128, 256, 8, _BK256_STAGES |
| 201 | # small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the |
| 202 | # N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there). |
| 203 | if N <= 1024 and K % 256 == 0: |
| 204 | return 128, 64, 256, 8, 3 |
| 205 | return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3) |
| 206 | |
| 207 | |
| 208 | _DOT_SCALED_K_ALIGN = 128 |
| 209 | |
| 210 | |
| 211 | def _as_block32_scale(weight_scale: torch.Tensor) -> torch.Tensor | None: |
| 212 | """The [N / 32, K / 32] block scales, if every 32 rows share their scales.""" |
| 213 | N = weight_scale.shape[0] |
| 214 | if N % BLOCK_ROWS != 0: |
| 215 | return None |
| 216 | blocks = weight_scale.view(N // BLOCK_ROWS, BLOCK_ROWS, -1) |
| 217 | if not torch.equal(blocks, blocks[:, :1].expand_as(blocks)): |
| 218 | return None |
| 219 | return blocks[:, 0].contiguous() |
| 220 | |
| 221 | |
| 222 | class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): |
| 223 | """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``.""" |
| 224 | |
| 225 | supports_pre_processed_weights = True |
| 226 | |
| 227 | @classmethod |
| 228 | def is_supported( |
| 229 | cls, compute_capability: int | None = None |
| 230 | ) -> tuple[bool, str | None]: |
| 231 | if not current_platform.is_rocm(): |
| 232 | return False, "not ROCm" |
| 233 | # supports_mx() == gfx95x (CDNA4 native microscaling hardware). On other |
| 234 | # archs dot_scaled would upcast to BF16, so the kernel selector falls |
| 235 | # through to the BF16 emulation (hipBLASLt) path instead. |
| 236 | if not current_platform.supports_mx(): |
| 237 | return False, "native MX requires CDNA4 (gfx95x)" |
| 238 | return True, None |
| 239 | |
| 240 | @classmethod |
| 241 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 242 | return True, None |
| 243 | |
| 244 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 245 | weight = layer.weight.data # [N, K] fp8 |
| 246 | N, K = weight.shape |
| 247 | scale_k = K // MXFP8_BLOCK_SIZE |
| 248 | weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() |
| 249 | block_scale = _as_block32_scale(weight_scale) |
| 250 | if block_scale is not None: |
| 251 | # Checkpoints with 32x32 scale blocks (DeepSeek V4/V4.1) get them |
| 252 | # back from the loader expanded per row; keep the blocks and use |
| 253 | # the block-scaled GEMM, which also handles any K % 32 == 0. |
| 254 | weight_scale = block_scale |
| 255 | elif K % _DOT_SCALED_K_ALIGN != 0: |
| 256 | weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale) |
| 257 | layer.weight = Parameter(weight.contiguous(), requires_grad=False) |
| 258 | layer.weight_scale = Parameter(weight_scale, requires_grad=False) |
| 259 | |
| 260 | def apply_weights( |
| 261 | self, |
| 262 | layer: torch.nn.Module, |
| 263 | x: torch.Tensor, |
| 264 | bias: torch.Tensor | None = None, |
| 265 | ) -> torch.Tensor: |
| 266 | if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE: |
| 267 | raise ValueError( |
| 268 | f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got " |
| 269 | f"{layer.weight_scale.dtype}." |
| 270 | ) |
| 271 | out_shape = (*x.shape[:-1], layer.weight.shape[0]) |
| 272 | x2d = x.reshape(-1, x.shape[-1]) |
| 273 | if layer.weight_scale.shape[0] != layer.weight.shape[0]: |
| 274 | # One scale row per 32 weight rows (see process_weights_after_loading). |
| 275 | x_q, x_scale = mxfp8_e4m3_quantize(x2d) |
| 276 | out = rocm_mxfp8_block32_gemm( |
| 277 | x_q, x_scale, layer.weight, layer.weight_scale, x.dtype |
| 278 | ) |
| 279 | elif layer.weight.element_size() >= 2: |
| 280 | out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype)) |
| 281 | elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0: |
| 282 | out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale) |
| 283 | else: |
| 284 | # process_weights_after_loading dequantizes these weights once, so |
| 285 | # this runs only if it did not; dot_scaled would be invalid here. |
| 286 | w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale) |
| 287 | out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype) |
| 288 | out = out.reshape(out_shape) |
| 289 | if bias is not None: |
| 290 | out = out + bias |
| 291 | return out |