| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DSV4.1 small-batch MXFP8 WO-A with inverse RoPE and quantization.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | from typing import Any |
| 7 | |
| 8 | import cutlass |
| 9 | import cutlass.cute as cute |
| 10 | import cutlass.utils as utils |
| 11 | import torch |
| 12 | from cuda.bindings.driver import CUstream |
| 13 | from cutlass import ( |
| 14 | BFloat16, |
| 15 | Float8E4M3FN, |
| 16 | Float32, |
| 17 | Int32, |
| 18 | Int64, |
| 19 | Uint8, |
| 20 | Uint32, |
| 21 | Uint64, |
| 22 | ) |
| 23 | from cutlass._mlir.dialects import llvm |
| 24 | from cutlass.cute.nvgpu import cpasync, tcgen05 |
| 25 | from cutlass.cute.runtime import make_fake_tensor |
| 26 | from cutlass.cutlass_dsl import dsl_user_op |
| 27 | |
| 28 | from vllm.cute_utils import _tcgen05, mbarrier, recast_val, simple_tma_copy |
| 29 | from vllm.model_executor.warmup.jit_warmup import WarmupIntRange, kernel_launcher |
| 30 | from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( |
| 31 | CuTeDSLLaunchSpec, |
| 32 | VllmCuTeDSLJitKernel, |
| 33 | ) |
| 34 | |
| 35 | # Each head is [nope | rope] as in fused_inv_rope_fp8_quant: only the last |
| 36 | # _ROPE_DIM columns carry RoPE. |
| 37 | _HEAD_DIM = 512 |
| 38 | _ROPE_DIM = 64 |
| 39 | |
| 40 | |
| 41 | @dsl_user_op |
| 42 | def _ue8m0_rp(value, *, loc=None, ip=None): |
| 43 | result = llvm.inline_asm( |
| 44 | Uint32.mlir_type, |
| 45 | [value.ir_value(loc=loc, ip=ip)], |
| 46 | "{ .reg .b16 sf; cvt.rp.satfinite.ue8m0x2.f32 sf, $1, $1; " |
| 47 | "cvt.u32.u16 $0, sf; and.b32 $0, $0, 255; }", |
| 48 | "=r,f", |
| 49 | has_side_effects=False, |
| 50 | loc=loc, |
| 51 | ip=ip, |
| 52 | ) |
| 53 | return Uint32(result) |
| 54 | |
| 55 | |
| 56 | @dsl_user_op |
| 57 | def _rope_fma(value, cos, other, sin, *, subtract=False, loc=None, ip=None): |
| 58 | """Preserve multiply/FMA rounding when RoPE loads are hoisted.""" |
| 59 | result = llvm.inline_asm( |
| 60 | Float32.mlir_type, |
| 61 | [v.ir_value(loc=loc, ip=ip) for v in (value, cos, other, sin)], |
| 62 | "{ .reg .f32 product; mul.rn.f32 product, $3, $4; " |
| 63 | + ( |
| 64 | "neg.f32 product, product; fma.rn.f32 $0, $1, $2, product; }" |
| 65 | if subtract |
| 66 | else "fma.rn.f32 $0, $1, $2, product; }" |
| 67 | ), |
| 68 | "=f,f,f,f,f", |
| 69 | has_side_effects=False, |
| 70 | loc=loc, |
| 71 | ip=ip, |
| 72 | ) |
| 73 | return Float32(result) |
| 74 | |
| 75 | |
| 76 | @cute.jit |
| 77 | def _scale(amax: Float32): |
| 78 | exponent = _ue8m0_rp(amax * Float32(1.0 / 448.0)) |
| 79 | inv = recast_val((Uint32(254) - exponent) << 23, Float32) |
| 80 | if exponent == 0: |
| 81 | inv = Float32(0) |
| 82 | return exponent, inv |
| 83 | |
| 84 | |
| 85 | class FusedWoAKernel(VllmCuTeDSLJitKernel["FusedWoAKernel.CompileKey"]): |
| 86 | @dataclass(frozen=True) |
| 87 | class CompileKey: |
| 88 | # T=1..96 uses at most 46 keys per (n_groups, heads_per_group, o_lora_rank). |
| 89 | tokens_per_tile: int |
| 90 | token_tiles: int |
| 91 | full_tiles: bool |
| 92 | n_groups: int |
| 93 | heads_per_group: int |
| 94 | o_lora_rank: int |
| 95 | |
| 96 | @staticmethod |
| 97 | def kernel(compile_key: CompileKey) -> Any: |
| 98 | tokens = compile_key.tokens_per_tile |
| 99 | token_tiles = compile_key.token_tiles |
| 100 | full_tiles = compile_key.full_tiles |
| 101 | acc_cols = 1 << (tokens - 1).bit_length() |
| 102 | tile_m = max(8, acc_cols) |
| 103 | tile_n = 128 |
| 104 | tile_k = 128 |
| 105 | # tcgen05 MXFP8 MMA consumes K = 32, one MX block, per instruction. |
| 106 | mma_k = 32 |
| 107 | # Each group's heads split K across one cluster, one head per CTA. |
| 108 | heads_per_group = compile_key.heads_per_group |
| 109 | tiles_per_group = compile_key.o_lora_rank // tile_n |
| 110 | n_tiles = compile_key.n_groups * tiles_per_group |
| 111 | # Each CTA keeps one whole head of K in SMEM, one TMA box per stage. |
| 112 | num_stages = _HEAD_DIM // tile_k |
| 113 | w_stage_bytes = tile_n * tile_k |
| 114 | nope_dim = _HEAD_DIM - _ROPE_DIM |
| 115 | threads_per_token = _HEAD_DIM // 4 |
| 116 | sf_col = max(16, tile_m) |
| 117 | tmem_cols = sf_col * 2 |
| 118 | threads = min(acc_cols * threads_per_token, 512) |
| 119 | tokens_per_iter = threads // threads_per_token |
| 120 | |
| 121 | @cute.kernel |
| 122 | def device_kernel( |
| 123 | x: cute.Tensor, |
| 124 | positions: cute.Tensor, |
| 125 | rope: cute.Tensor, |
| 126 | w: cpasync.TmaInfo, |
| 127 | ws: cpasync.TmaInfo, |
| 128 | q: cute.Tensor, |
| 129 | qs: cute.Tensor, |
| 130 | num_tokens: Int32, |
| 131 | ): |
| 132 | total_tokens = num_tokens |
| 133 | if cutlass.const_expr(full_tiles): |
| 134 | total_tokens = Int32(tokens * token_tiles) |
| 135 | tid, _, _ = cute.arch.thread_idx() |
| 136 | bid, token_tile, _ = cute.arch.block_idx() |
| 137 | token_start = Int32(0) |
| 138 | if cutlass.const_expr(token_tiles > 1): |
| 139 | token_start = token_tile * tokens |
| 140 | warp = cute.arch.make_warp_uniform(tid // 32) |
| 141 | lane = tid % 32 |
| 142 | split = bid % heads_per_group |
| 143 | tile = bid // heads_per_group |
| 144 | group = tile // tiles_per_group |
| 145 | smem = utils.SmemAllocator() |
| 146 | sW = smem.allocate_tensor( |
| 147 | Float8E4M3FN, |
| 148 | w.smem_layout.outer, |
| 149 | byte_alignment=128, |
| 150 | swizzle=w.smem_layout.inner, |
| 151 | ) |
| 152 | sX = smem.allocate_tensor( |
| 153 | Float8E4M3FN, |
| 154 | cute.make_layout( |
| 155 | (tile_m, tile_k, num_stages), stride=(tile_k, 1, tile_m * tile_k) |
| 156 | ), |
| 157 | byte_alignment=1024, |
| 158 | swizzle=cute.make_swizzle(3, 4, 3), |
| 159 | ) |
| 160 | # tcgen05.cp 32x128b.warpx4 order: row r at byte (r % 32) * 16 + |
| 161 | # (r // 32) * 4, one 4-byte cell of MX-block scales per row. |
| 162 | sW_SF = smem.allocate_tensor( |
| 163 | Int32, |
| 164 | cute.make_layout(((32, 4), num_stages), stride=((4, 1), tile_n)), |
| 165 | 128, |
| 166 | ) |
| 167 | sW_SF_raw = smem.allocate_tensor( |
| 168 | Int32, cute.make_layout((tile_n, num_stages)), 128 |
| 169 | ) |
| 170 | sX_SF = smem.allocate_tensor( |
| 171 | Uint8, |
| 172 | cute.make_layout( |
| 173 | ((32, 4), tile_k // mma_k, num_stages), |
| 174 | stride=((16, 4), 1, tile_n * 4), |
| 175 | ), |
| 176 | 128, |
| 177 | ) |
| 178 | copy_fp8x4 = cute.make_copy_atom( |
| 179 | cute.nvgpu.CopyUniversalOp(), Float8E4M3FN, num_bits_per_copy=32 |
| 180 | ) |
| 181 | partial = smem.allocate_tensor( |
| 182 | Float32, |
| 183 | cute.make_layout( |
| 184 | (tile_n, cute.ceil_div(tokens, heads_per_group), heads_per_group) |
| 185 | ), |
| 186 | 128, |
| 187 | ) |
| 188 | mbar_loaded = smem.allocate_array(Int64, num_stages) |
| 189 | mbar_ws_loaded = smem.allocate_array(Int64, 1) |
| 190 | mbar_done = smem.allocate_array(Int64, 1) |
| 191 | mbar_reduced = smem.allocate_array(Int64, 1) |
| 192 | taddr = smem.allocate(Int32, 4) |
| 193 | if tid == 0: |
| 194 | for stage in cutlass.range_constexpr(num_stages): |
| 195 | # Wait for weight TMA and this stage's quantization threads. |
| 196 | cute.arch.mbarrier_init(mbar_loaded + stage, 1 + threads // 4) |
| 197 | cute.arch.mbarrier_init(mbar_ws_loaded, 1) |
| 198 | cute.arch.mbarrier_init(mbar_done, 1) |
| 199 | cute.arch.mbarrier_init(mbar_reduced, 1) |
| 200 | cute.arch.mbarrier_init_fence() |
| 201 | # Local users of the barriers (warp 0's TMA below) need the init too. |
| 202 | cute.arch.sync_threads() |
| 203 | # Paired with cluster_wait() below: peers must see mbar_reduced |
| 204 | # initialized before their st.async; the split hides the barrier. |
| 205 | cute.arch.cluster_arrive_relaxed() |
| 206 | if warp == 0: |
| 207 | # Raw DeepGEMM MN-major scales; the MMA warp transposes them to |
| 208 | # the UTCCP order, which TMA cannot scatter at 4-byte granularity. |
| 209 | ws_tiles = cute.zipped_divide(ws.tma_tensor, (tile_n, num_stages)) |
| 210 | with cute.arch.elect_one(): |
| 211 | mbarrier.arrive_expect_tx(mbar_ws_loaded, tile_n * num_stages * 4) |
| 212 | simple_tma_copy( |
| 213 | ws.atom, |
| 214 | ws_tiles[None, (tile % tiles_per_group, split, group)], |
| 215 | sW_SF_raw, |
| 216 | mbar_ws_loaded, |
| 217 | ) |
| 218 | tiles = cute.zipped_divide(w.tma_tensor, (tile_n, tile_k)) |
| 219 | for stage in cutlass.range_constexpr(num_stages): |
| 220 | with cute.arch.elect_one(): |
| 221 | mbarrier.arrive_expect_tx(mbar_loaded + stage, w_stage_bytes) |
| 222 | simple_tma_copy( |
| 223 | w.atom, |
| 224 | tiles[None, (tile, split * num_stages + stage)], |
| 225 | sW[None, None, stage], |
| 226 | mbar_loaded + stage, |
| 227 | ) |
| 228 | # x and positions come from the PDL predecessor; everything above |
| 229 | # only touches weights. Padded sX rows only feed unread MMA columns. |
| 230 | cute.arch.griddepcontrol_wait() |
| 231 | # TMEM is allocated at runtime, and a PDL predecessor on this SM may |
| 232 | # hold it until it exits; allocating earlier can deadlock. |
| 233 | if warp == 0: |
| 234 | cute.arch.alloc_tmem(tmem_cols, taddr) |
| 235 | cute.arch.relinquish_tmem_alloc_permit() |
| 236 | copy_x = cute.make_copy_atom( |
| 237 | cute.nvgpu.CopyUniversalOp(), BFloat16, num_bits_per_copy=64 |
| 238 | ) |
| 239 | n_iters = cute.ceil_div(tokens, tokens_per_iter) |
| 240 | qid = (tid + 32) % threads_per_token |
| 241 | k = qid * 4 |
| 242 | # Tail rows repeat the last input; their output stores are masked. |
| 243 | # Issue every token's x load before any math so their latencies overlap. |
| 244 | x_bf16 = cute.make_rmem_tensor((4, n_iters), BFloat16) |
| 245 | for iteration in cutlass.range_constexpr(n_iters): |
| 246 | token = iteration * tokens_per_iter + tid // threads_per_token |
| 247 | if cutlass.const_expr(tokens % tokens_per_iter == 0) or token < tokens: |
| 248 | x_src = cute.local_tile( |
| 249 | x[ |
| 250 | cute.min(token_start + token, total_tokens - 1), |
| 251 | group * heads_per_group + split, |
| 252 | None, |
| 253 | ], |
| 254 | (4,), |
| 255 | (qid,), |
| 256 | ) |
| 257 | cute.copy(copy_x, x_src, x_bf16[None, iteration]) |
| 258 | # Overlap the RoPE loads across tokens before consuming any of them. |
| 259 | cos_reg = cute.make_rmem_tensor((2, n_iters), Float32) |
| 260 | sin_reg = cute.make_rmem_tensor((2, n_iters), Float32) |
| 261 | if cutlass.const_expr(n_iters > 1): |
| 262 | copy_rope = cute.make_copy_atom( |
| 263 | cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=32 |
| 264 | ) |
| 265 | if k >= nope_dim: |
| 266 | for iteration in cutlass.range_constexpr(n_iters): |
| 267 | token = iteration * tokens_per_iter + tid // threads_per_token |
| 268 | if ( |
| 269 | cutlass.const_expr(tokens % tokens_per_iter == 0) |
| 270 | or token < tokens |
| 271 | ): |
| 272 | src_token = cute.min(token_start + token, total_tokens - 1) |
| 273 | pos = positions[src_token] |
| 274 | freq = (k - nope_dim) // 2 |
| 275 | rope_row = rope[pos, None] |
| 276 | cos_src = cute.local_tile(rope_row, (2,), (freq // 2,)) |
| 277 | cute.copy(copy_rope, cos_src, cos_reg[None, iteration]) |
| 278 | sin_tile = (freq + _ROPE_DIM // 2) // 2 |
| 279 | sin_src = cute.local_tile(rope_row, (2,), (sin_tile,)) |
| 280 | cute.copy(copy_rope, sin_src, sin_reg[None, iteration]) |
| 281 | for iteration in cutlass.range_constexpr(n_iters): |
| 282 | token = iteration * tokens_per_iter + tid // threads_per_token |
| 283 | if cutlass.const_expr(tokens % tokens_per_iter == 0) or token < tokens: |
| 284 | x_f32 = cute.make_rmem_tensor((4,), Float32) |
| 285 | x_f32.store(x_bf16[None, iteration].load().to(Float32)) |
| 286 | if k >= nope_dim: |
| 287 | # Inverse RoPE on the interleaved pairs (k, k+1) and |
| 288 | # (k+2, k+3); each rope row is cos || sin. |
| 289 | for pair in cutlass.range_constexpr(2): |
| 290 | if cutlass.const_expr(n_iters > 1): |
| 291 | cos = cos_reg[pair, iteration] |
| 292 | sin = sin_reg[pair, iteration] |
| 293 | else: |
| 294 | pos = positions[ |
| 295 | cute.min(token_start + token, total_tokens - 1) |
| 296 | ] |
| 297 | freq = (k - nope_dim) // 2 + pair |
| 298 | cos = Float32(rope[pos, freq]) |
| 299 | sin = Float32(rope[pos, freq + _ROPE_DIM // 2]) |
| 300 | even, odd = x_f32[2 * pair], x_f32[2 * pair + 1] |
| 301 | x_f32[2 * pair] = _rope_fma(even, cos, odd, sin) |
| 302 | x_f32[2 * pair + 1] = _rope_fma( |
| 303 | odd, cos, even, sin, subtract=True |
| 304 | ) |
| 305 | amax = cute.arch.fmax( |
| 306 | cute.arch.fmax(cute.abs(x_f32[0]), cute.abs(x_f32[1])), |
| 307 | cute.arch.fmax(cute.abs(x_f32[2]), cute.abs(x_f32[3])), |
| 308 | ) |
| 309 | amax = cute.arch.warp_reduction_max(amax, threads_in_group=8) |
| 310 | exponent, inv = _scale(cute.arch.fmax(amax, Float32(1e-10))) |
| 311 | x_fp8 = cute.make_rmem_tensor((4,), Float8E4M3FN) |
| 312 | x_fp8.store((x_f32.load() * inv).to(Float8E4M3FN)) |
| 313 | cute.copy( |
| 314 | copy_fp8x4, |
| 315 | x_fp8, |
| 316 | cute.local_tile( |
| 317 | sX[token, None, k // tile_k], (4,), (qid % 32,) |
| 318 | ), |
| 319 | ) |
| 320 | if lane % 8 == 0: |
| 321 | sX_SF[token, k % tile_k // mma_k, k // tile_k] = Uint8(exponent) |
| 322 | cute.arch.fence_proxy("async.shared", space="cta") |
| 323 | mbarrier.arrive(mbar_loaded + qid // 32, order="release") |
| 324 | if warp == 0: |
| 325 | base = cute.make_tensor(taddr, cute.make_layout(1))[0] |
| 326 | # sW / sX are 128B-swizzled K-major; one stage's K fills the |
| 327 | # swizzle atom, so the leading byte offset is unused. |
| 328 | sdesc = _tcgen05.make_sdesc_128B_swizzle(LBO=0) |
| 329 | # Unswizzled 32 x 16 B scale tiles: SBO = one 8-row core matrix |
| 330 | # (16 B units at bit 32); bit 46 is the SM100 descriptor version. |
| 331 | sf_sbo = 8 * 16 |
| 332 | sfdesc = Uint64((sf_sbo >> 4 << 32) | (1 << 46)) |
| 333 | # M = tile_n weight rows (A), N = tile_m token rows (B). |
| 334 | idesc = _tcgen05.make_mxfp8_idesc(tile_n, tile_m) |
| 335 | # sW_SF's layout is the UTCCP order, so a plain copy transposes; |
| 336 | # fence the generic-proxy writes before tcgen05.cp reads them. |
| 337 | cute.arch.mbarrier_wait(mbar_ws_loaded, 0) |
| 338 | for stage in cutlass.range_constexpr(num_stages): |
| 339 | for i in cutlass.range_constexpr(4): |
| 340 | sW_SF[i * 32 + lane, stage] = sW_SF_raw[i * 32 + lane, stage] |
| 341 | cute.arch.fence_proxy("async.shared", space="cta") |
| 342 | cute.arch.sync_warp() |
| 343 | for stage in cutlass.range_constexpr(num_stages): |
| 344 | cute.arch.mbarrier_wait(mbar_loaded + stage, 0) |
| 345 | _tcgen05.fence_after_thread_sync() |
| 346 | if cutlass.const_expr(stage == num_stages - 1): |
| 347 | with cute.arch.elect_one(): |
| 348 | cute.arch.griddepcontrol_launch_dependents() |
| 349 | adesc = sdesc | (sW[None, None, stage].iterator.toint() >> 4) |
| 350 | bdesc = sdesc | (sX[None, None, stage].iterator.toint() >> 4) |
| 351 | _tcgen05.cp( |
| 352 | base + sf_col, |
| 353 | sfdesc | (sW_SF[None, stage].iterator.toint() >> 4), |
| 354 | "32x128b", |
| 355 | "warpx4", |
| 356 | ) |
| 357 | _tcgen05.cp( |
| 358 | base + sf_col + 4, |
| 359 | sfdesc | (sX_SF[None, None, stage].iterator.toint() >> 4), |
| 360 | "32x128b", |
| 361 | "warpx4", |
| 362 | ) |
| 363 | # Advance mma_k FP8 bytes (16 B units) and pick MX block kk of |
| 364 | # each 4-byte scale (B sf_id at bit 4, A sf_id at bit 29). |
| 365 | for kk in cutlass.range_constexpr(tile_k // mma_k): |
| 366 | _tcgen05.mma_mxfp8( |
| 367 | base, |
| 368 | adesc + kk * mma_k // 16, |
| 369 | bdesc + kk * mma_k // 16, |
| 370 | idesc + ((kk << 4) | (kk << 29)), |
| 371 | base + sf_col, |
| 372 | base + sf_col + 4, |
| 373 | stage > 0 or kk > 0, |
| 374 | ) |
| 375 | _tcgen05.commit(mbar_done) |
| 376 | # Only the MMA warp waits; barrier 2 releases the TMEM readers. |
| 377 | cute.arch.mbarrier_wait(mbar_done, 0) |
| 378 | if tid < tile_n: |
| 379 | cute.arch.barrier(barrier_id=2, number_of_threads=tile_n) |
| 380 | # Tokens t with t % heads_per_group == split arrive from every peer |
| 381 | # (this CTA included), tile_n Float32 rows each. |
| 382 | if tid == 0: |
| 383 | owned = cute.ceil_div(tokens - split, heads_per_group) |
| 384 | mbarrier.arrive_expect_tx( |
| 385 | mbar_reduced, owned * tile_n * heads_per_group * 4 |
| 386 | ) |
| 387 | # Pairs with cluster_arrive_relaxed() above: every peer's |
| 388 | # mbar_reduced is initialized before the st.async below. |
| 389 | cute.arch.cluster_wait() |
| 390 | # quack's process-wide const_expr-if rewrite can leave these unbound |
| 391 | # after the dynamic path above; bind them so the regions join. |
| 392 | amax, exponent, inv = Float32(0), Uint32(0), Float32(0) |
| 393 | if tid < tile_n: |
| 394 | base = cute.make_tensor(taddr, cute.make_layout(1))[0] |
| 395 | _tcgen05.fence_after_thread_sync() |
| 396 | acc = cute.make_rmem_tensor(acc_cols, Float32) |
| 397 | if cutlass.const_expr(tokens == 1): |
| 398 | acc[0] = _tcgen05.ld(warp * 32, base, "32x32b", 1) |
| 399 | else: |
| 400 | acc.store(_tcgen05.ld(warp * 32, base, "32x32b", acc_cols)) |
| 401 | _tcgen05.wait_ld() |
| 402 | for token in cutlass.range_constexpr(tokens): |
| 403 | ptr = cute.domain_offset( |
| 404 | (tid, token // heads_per_group, split), partial |
| 405 | ).iterator |
| 406 | cute.arch.store_async_dsmem( |
| 407 | ptr, |
| 408 | recast_val(acc[token], Int32), |
| 409 | mbar_reduced, |
| 410 | token % heads_per_group, |
| 411 | ) |
| 412 | if split < tokens: |
| 413 | # One warp waits for the peers' bytes; named barrier 1 (0 is |
| 414 | # sync_threads) releases the other tile_n epilogue threads. |
| 415 | if warp == 0: |
| 416 | cute.arch.mbarrier_wait(mbar_reduced, 0) |
| 417 | cute.arch.barrier(barrier_id=1, number_of_threads=tile_n) |
| 418 | for i in cutlass.range_constexpr( |
| 419 | cute.ceil_div(tokens, heads_per_group) |
| 420 | ): |
| 421 | token = split + i * heads_per_group |
| 422 | if cutlass.const_expr( |
| 423 | full_tiles |
| 424 | and ( |
| 425 | tokens <= heads_per_group |
| 426 | or tokens % heads_per_group == 0 |
| 427 | ) |
| 428 | ) or (token < tokens and token_start + token < total_tokens): |
| 429 | value = Float32(0) |
| 430 | for peer in cutlass.range_constexpr(heads_per_group): |
| 431 | value += partial[tid, i, peer] |
| 432 | value = Float32(BFloat16(value)) |
| 433 | amax = cute.arch.warp_redux_sync(value, "max", abs=True) |
| 434 | exponent, inv = _scale(amax) |
| 435 | q[token_start + token, tile * tile_n + tid] = Float8E4M3FN( |
| 436 | value * inv |
| 437 | ) |
| 438 | if lane == 0: |
| 439 | row = token_start + token |
| 440 | offset = row % 32 * 16 + row // 32 * 4 + warp |
| 441 | qs[tile * 512 + offset] = Uint8(exponent) |
| 442 | # Every scale tile owns its padding, with no second memset kernel. |
| 443 | for i in cutlass.range_constexpr(4): |
| 444 | offset = i * tile_n + tid |
| 445 | row = offset // 16 + (offset % 16 // 4) * 32 |
| 446 | if token_tile == 0 and split == 0 and row >= total_tokens: |
| 447 | qs[tile * 512 + offset] = Uint8(0) |
| 448 | # Each receiver drained all writes into its own inbox. Only local |
| 449 | # TMEM readers must finish before deallocation; no peer reads SMEM. |
| 450 | if tid < tile_n: |
| 451 | cute.arch.barrier(barrier_id=2, number_of_threads=tile_n) |
| 452 | if warp == 0: |
| 453 | cute.arch.dealloc_tmem( |
| 454 | cute.make_ptr( |
| 455 | Float32, |
| 456 | cute.make_tensor(taddr, cute.make_layout(1))[0], |
| 457 | cute.AddressSpace.tmem, |
| 458 | ), |
| 459 | tmem_cols, |
| 460 | ) |
| 461 | |
| 462 | @cute.jit |
| 463 | def host_entrypoint( |
| 464 | x: cute.Tensor, |
| 465 | positions: cute.Tensor, |
| 466 | rope: cute.Tensor, |
| 467 | w: cute.Tensor, |
| 468 | ws: cute.Tensor, |
| 469 | q: cute.Tensor, |
| 470 | qs: cute.Tensor, |
| 471 | num_tokens: Int32, |
| 472 | stream: CUstream, |
| 473 | ): |
| 474 | layout = cute.make_composed_layout( |
| 475 | cute.make_swizzle(3, 4, 3), |
| 476 | 0, |
| 477 | cute.make_layout( |
| 478 | (tile_n, tile_k, num_stages), stride=(tile_k, 1, w_stage_bytes) |
| 479 | ), |
| 480 | ) |
| 481 | tma = cpasync.make_tiled_tma_atom( |
| 482 | cpasync.CopyBulkTensorTileG2SOp(cta_group=tcgen05.CtaGroup.ONE), |
| 483 | w, |
| 484 | layout, |
| 485 | (tile_n, tile_k), |
| 486 | ) |
| 487 | ws_tma = cpasync.make_tiled_tma_atom( |
| 488 | cpasync.CopyBulkTensorTileG2SOp(cta_group=tcgen05.CtaGroup.ONE), |
| 489 | cute.make_tensor(ws.iterator, cute.select(ws.layout, mode=[1, 2, 0])), |
| 490 | cute.make_layout((tile_n, num_stages)), |
| 491 | (tile_n, num_stages), |
| 492 | ) |
| 493 | device_kernel(x, positions, rope, tma, ws_tma, q, qs, num_tokens).launch( |
| 494 | grid=(n_tiles * heads_per_group, token_tiles, 1), |
| 495 | block=(threads, 1, 1), |
| 496 | cluster=(heads_per_group, 1, 1), |
| 497 | stream=stream, |
| 498 | use_pdl=True, |
| 499 | ) |
| 500 | |
| 501 | return host_entrypoint |
| 502 | |
| 503 | def dispatch( # type: ignore[override] |
| 504 | self, *, tokens: int, n_groups: int, heads_per_group: int, o_lora_rank: int |
| 505 | ) -> CompileKey: |
| 506 | token_tiles = (tokens + 31) // 32 |
| 507 | # Host integer arithmetic: balance tiles in multiples of four tokens. |
| 508 | tile_tokens = ( |
| 509 | tokens |
| 510 | if token_tiles == 1 |
| 511 | else (tokens + 4 * token_tiles - 1) // (4 * token_tiles) * 4 |
| 512 | ) |
| 513 | return self.CompileKey( |
| 514 | tokens_per_tile=tile_tokens, |
| 515 | token_tiles=token_tiles, |
| 516 | full_tiles=tokens % tile_tokens == 0, |
| 517 | n_groups=n_groups, |
| 518 | heads_per_group=heads_per_group, |
| 519 | o_lora_rank=o_lora_rank, |
| 520 | ) |
| 521 | |
| 522 | def get_warmup_keys( |
| 523 | self, *, max_tokens: int, n_groups: int, heads_per_group: int, o_lora_rank: int |
| 524 | ) -> list[CompileKey]: |
| 525 | # Trace all accepted token counts and deduplicate shared tile shapes. |
| 526 | return self._trace_dispatch(self.dispatch)( |
| 527 | tokens=WarmupIntRange(1, max_tokens + 1), |
| 528 | n_groups=n_groups, |
| 529 | heads_per_group=heads_per_group, |
| 530 | o_lora_rank=o_lora_rank, |
| 531 | ) |
| 532 | |
| 533 | def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: |
| 534 | tokens = ( |
| 535 | compile_key.tokens_per_tile * compile_key.token_tiles |
| 536 | if compile_key.full_tiles |
| 537 | else cute.sym_int() |
| 538 | ) |
| 539 | groups = compile_key.n_groups |
| 540 | rank = compile_key.o_lora_rank |
| 541 | n = groups * rank |
| 542 | k = compile_key.heads_per_group * _HEAD_DIM |
| 543 | return ( |
| 544 | make_fake_tensor( |
| 545 | BFloat16, |
| 546 | (tokens, groups * compile_key.heads_per_group, _HEAD_DIM), |
| 547 | (cute.sym_int(divisibility=8), _HEAD_DIM, 1), |
| 548 | assumed_align=16, |
| 549 | ), |
| 550 | make_fake_tensor(Int64, (tokens,), (1,)), |
| 551 | make_fake_tensor(Float32, (cute.sym_int(), _ROPE_DIM), (_ROPE_DIM, 1)), |
| 552 | make_fake_tensor(Float8E4M3FN, (n, k), (k, 1), assumed_align=16), |
| 553 | make_fake_tensor( |
| 554 | Int32, |
| 555 | (groups, rank, k // 128), |
| 556 | (rank * k // 128, 1, rank), |
| 557 | assumed_align=16, |
| 558 | ), |
| 559 | make_fake_tensor(Float8E4M3FN, (tokens, n), (n, 1)), |
| 560 | make_fake_tensor(Uint8, (128 * n // 32,), (1,)), |
| 561 | Int32(0), |
| 562 | ) |
| 563 | |
| 564 | @kernel_launcher |
| 565 | def __call__( |
| 566 | self, |
| 567 | *, |
| 568 | x: torch.Tensor, |
| 569 | positions: torch.Tensor, |
| 570 | rope: torch.Tensor, |
| 571 | weight: torch.Tensor, |
| 572 | weight_scale: torch.Tensor, |
| 573 | ) -> CuTeDSLLaunchSpec["FusedWoAKernel.CompileKey"]: |
| 574 | """Return FP8 WO-A output and FlashInfer F8_128x4 scale bytes.""" |
| 575 | tokens = x.shape[0] |
| 576 | groups, rank = weight_scale.shape[:2] |
| 577 | n = groups * rank |
| 578 | q = torch.empty((tokens, n), device=x.device, dtype=torch.float8_e4m3fn) |
| 579 | scales = torch.empty(128 * n // 32, device=x.device, dtype=torch.uint8) |
| 580 | launch_args = ( |
| 581 | x, |
| 582 | positions, |
| 583 | rope, |
| 584 | weight.view(n, -1), |
| 585 | weight_scale, |
| 586 | q, |
| 587 | scales, |
| 588 | tokens, |
| 589 | ) |
| 590 | compile_key = self.dispatch( |
| 591 | tokens=tokens, |
| 592 | n_groups=groups, |
| 593 | heads_per_group=x.shape[1] // groups, |
| 594 | o_lora_rank=rank, |
| 595 | ) |
| 596 | return compile_key, launch_args, (q, scales) |
| 597 | |
| 598 | |
| 599 | _FUSED_WO_A_KERNEL = FusedWoAKernel() |