| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from functools import cache |
| 4 | |
| 5 | import cutlass |
| 6 | import cutlass.cute as cute |
| 7 | import torch |
| 8 | from cuda.bindings.driver import CUstream |
| 9 | from cutlass import BFloat16, Float8E4M3FN, Float32, Int64, Uint8, Uint16, Uint32 |
| 10 | |
| 11 | from vllm.cute_utils import _TORCH_TO_CUTE_DTYPE, cvt |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | |
| 15 | def _make_fake_tensor(dtype, shape, divisibility): |
| 16 | stride = tuple( |
| 17 | 1 if i == len(shape) - 1 else cute.sym_int64(divisibility=divisibility) |
| 18 | for i in range(len(shape)) |
| 19 | ) |
| 20 | return cute.runtime.make_fake_tensor( |
| 21 | dtype, |
| 22 | shape, |
| 23 | stride, |
| 24 | assumed_align=divisibility * dtype.width // 8, |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | def is_fused_q_cutedsl_supported( |
| 29 | q_pe: torch.Tensor, |
| 30 | index_q: torch.Tensor | None, |
| 31 | ql_nope: torch.Tensor, |
| 32 | *, |
| 33 | has_indexer: bool, |
| 34 | quantize_mqa: bool, |
| 35 | ) -> bool: |
| 36 | if not ( |
| 37 | current_platform.has_device_capability(100) |
| 38 | and q_pe.dtype == torch.bfloat16 |
| 39 | and q_pe.shape[-1] == 64 |
| 40 | and ( |
| 41 | not quantize_mqa |
| 42 | or (ql_nope.dtype == torch.bfloat16 and ql_nope.shape[-1] == 512) |
| 43 | ) |
| 44 | # One warp per head group; a non-multiple would trip the kernel's own |
| 45 | # assert instead of falling back to Triton. |
| 46 | and q_pe.shape[1] % 4 == 0 |
| 47 | ): |
| 48 | return False |
| 49 | return not has_indexer or ( |
| 50 | index_q is not None |
| 51 | and index_q.dtype == torch.bfloat16 |
| 52 | and index_q.shape[1] % 16 == 0 |
| 53 | and index_q.shape[-1] == 128 |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | def fused_q_cutedsl( |
| 58 | positions: torch.Tensor, |
| 59 | q_pe: torch.Tensor, |
| 60 | rope_cache: torch.Tensor, |
| 61 | ql_nope: torch.Tensor, |
| 62 | q_scale: torch.Tensor, |
| 63 | mqa_output: torch.Tensor, |
| 64 | idx_q: torch.Tensor, |
| 65 | idx_rope_cache: torch.Tensor, |
| 66 | idx_weights: torch.Tensor, |
| 67 | idx_weights_softmax_scale: float, |
| 68 | idx_weights_head_scale: float, |
| 69 | idx_q_fp8: torch.Tensor, |
| 70 | idx_weights_out: torch.Tensor, |
| 71 | has_indexer: bool = True, |
| 72 | index_rope_interleave: bool = True, |
| 73 | quantize_mqa: bool = True, |
| 74 | ) -> None: |
| 75 | _, num_heads, rope_dim = q_pe.shape |
| 76 | nope_dim = ql_nope.shape[-1] if quantize_mqa else 0 |
| 77 | _, num_idx_heads, idx_dim = idx_q.shape |
| 78 | |
| 79 | if has_indexer: |
| 80 | idx_rope_type = _TORCH_TO_CUTE_DTYPE[idx_rope_cache.dtype] |
| 81 | idx_weights_type = _TORCH_TO_CUTE_DTYPE[idx_weights.dtype] |
| 82 | else: |
| 83 | idx_dim = num_idx_heads = 0 |
| 84 | idx_q = idx_rope_cache = idx_q_fp8 = None |
| 85 | idx_weights = idx_weights_out = None |
| 86 | idx_rope_type = idx_weights_type = None |
| 87 | |
| 88 | rope_type = _TORCH_TO_CUTE_DTYPE[rope_cache.dtype] |
| 89 | compiled = FusedQKernel.compile( |
| 90 | rope_dim, |
| 91 | nope_dim, |
| 92 | num_heads, |
| 93 | rope_type, |
| 94 | idx_dim, |
| 95 | num_idx_heads, |
| 96 | idx_rope_type, |
| 97 | idx_weights_type, |
| 98 | index_rope_interleave, |
| 99 | ) |
| 100 | compiled( |
| 101 | positions, |
| 102 | q_pe, |
| 103 | rope_cache, |
| 104 | ql_nope if quantize_mqa else None, |
| 105 | q_scale.view(1) if quantize_mqa else None, |
| 106 | mqa_output, |
| 107 | idx_q, |
| 108 | idx_rope_cache, |
| 109 | idx_weights, |
| 110 | idx_q_fp8, |
| 111 | idx_weights_out, |
| 112 | float(idx_weights_softmax_scale * idx_weights_head_scale), |
| 113 | ) |
| 114 | |
| 115 | |
| 116 | class FusedQKernel: |
| 117 | def __init__( |
| 118 | self, |
| 119 | rope_dim: int, |
| 120 | nope_dim: int, |
| 121 | num_heads: int, |
| 122 | idx_dim: int, |
| 123 | num_idx_heads: int, |
| 124 | index_rope_interleave: bool, |
| 125 | ) -> None: |
| 126 | assert rope_dim == 64 |
| 127 | assert nope_dim in (0, 512) |
| 128 | assert idx_dim in (128, 0) |
| 129 | |
| 130 | self.rope_dim = rope_dim |
| 131 | self.nope_dim = nope_dim |
| 132 | self.num_heads = num_heads |
| 133 | self.idx_dim = idx_dim |
| 134 | self.num_idx_heads = num_idx_heads |
| 135 | self.index_rope_interleave = index_rope_interleave |
| 136 | |
| 137 | # mqa: rope_dim=64, nope_dim=512, num_heads=64/TP |
| 138 | # indexer: rope_dim=64, nope_dim=64, num_heads=32 |
| 139 | |
| 140 | self.num_warps = 4 |
| 141 | assert num_heads % self.num_warps == 0 |
| 142 | assert num_idx_heads % (4 * self.num_warps) == 0 |
| 143 | self.num_ctas_per_tok = num_heads // self.num_warps |
| 144 | self.num_ctas_per_idx_tok = num_idx_heads // (4 * self.num_warps) |
| 145 | |
| 146 | @cute.jit |
| 147 | def __call__( |
| 148 | self, |
| 149 | positions: cute.Tensor, |
| 150 | q_pe: cute.Tensor, |
| 151 | rope_cache: cute.Tensor, |
| 152 | ql_nope: cute.Tensor, |
| 153 | q_scale: cute.Tensor, |
| 154 | mqa_output: cute.Tensor, |
| 155 | idx_q: cute.Tensor, |
| 156 | idx_rope_cache: cute.Tensor, |
| 157 | idx_weights: cute.Tensor, |
| 158 | idx_q_fp8: cute.Tensor, |
| 159 | idx_weights_out: cute.Tensor, |
| 160 | weight_scale: Float32, |
| 161 | stream: CUstream, |
| 162 | ): |
| 163 | num_tokens = positions.shape[0] |
| 164 | if cutlass.const_expr(self.idx_dim == 0): |
| 165 | grid = (num_tokens, self.num_ctas_per_tok, 1) |
| 166 | else: |
| 167 | num_mqa_ctas = num_tokens * self.num_ctas_per_tok |
| 168 | num_idx_ctas = num_tokens * self.num_ctas_per_idx_tok |
| 169 | grid = (num_mqa_ctas + num_idx_ctas, 1, 1) |
| 170 | |
| 171 | self.kernel( |
| 172 | positions, |
| 173 | q_pe, |
| 174 | rope_cache, |
| 175 | ql_nope, |
| 176 | q_scale, |
| 177 | mqa_output, |
| 178 | idx_q, |
| 179 | idx_rope_cache, |
| 180 | idx_weights, |
| 181 | idx_q_fp8, |
| 182 | idx_weights_out, |
| 183 | weight_scale, |
| 184 | ).launch( |
| 185 | grid=grid, |
| 186 | block=(self.num_warps * 32, 1, 1), |
| 187 | stream=stream, |
| 188 | use_pdl=True, |
| 189 | ) |
| 190 | |
| 191 | @cute.kernel |
| 192 | def kernel( |
| 193 | self, |
| 194 | positions: cute.Tensor, |
| 195 | q_pe: cute.Tensor, |
| 196 | q_pe_rope_cache: cute.Tensor, |
| 197 | ql_nope: cute.Tensor, |
| 198 | q_scale: cute.Tensor, |
| 199 | mqa_output: cute.Tensor, |
| 200 | idx_q: cute.Tensor, |
| 201 | idx_q_rope_cache: cute.Tensor, |
| 202 | idx_weights: cute.Tensor, |
| 203 | idx_q_fp8: cute.Tensor, |
| 204 | idx_weights_out: cute.Tensor, |
| 205 | weight_scale: Float32, |
| 206 | ): |
| 207 | if cutlass.const_expr(self.idx_dim == 0): |
| 208 | token_id, group_id, _ = cute.arch.block_idx() |
| 209 | self.mqa( |
| 210 | positions, |
| 211 | q_pe, |
| 212 | q_pe_rope_cache, |
| 213 | ql_nope, |
| 214 | q_scale, |
| 215 | mqa_output, |
| 216 | token_id, |
| 217 | group_id, |
| 218 | ) |
| 219 | else: |
| 220 | # CTA-specialization |
| 221 | bid, _, _ = cute.arch.block_idx() |
| 222 | num_mqa_ctas = positions.shape[0] * self.num_ctas_per_tok |
| 223 | if bid < num_mqa_ctas: |
| 224 | self.mqa( |
| 225 | positions, |
| 226 | q_pe, |
| 227 | q_pe_rope_cache, |
| 228 | ql_nope, |
| 229 | q_scale, |
| 230 | mqa_output, |
| 231 | bid // self.num_ctas_per_tok, |
| 232 | bid % self.num_ctas_per_tok, |
| 233 | ) |
| 234 | else: |
| 235 | bid -= num_mqa_ctas |
| 236 | self.indexer( |
| 237 | positions, |
| 238 | idx_q, |
| 239 | idx_q_rope_cache, |
| 240 | idx_weights, |
| 241 | idx_q_fp8, |
| 242 | idx_weights_out, |
| 243 | weight_scale, |
| 244 | bid // self.num_ctas_per_idx_tok, |
| 245 | bid % self.num_ctas_per_idx_tok, |
| 246 | ) |
| 247 | |
| 248 | @cute.jit |
| 249 | def mqa( |
| 250 | self, |
| 251 | positions: cute.Tensor, |
| 252 | q_pe: cute.Tensor, |
| 253 | q_pe_rope_cache: cute.Tensor, |
| 254 | ql_nope: cute.Tensor, |
| 255 | q_scale: cute.Tensor, |
| 256 | mqa_output: cute.Tensor, |
| 257 | token_id, |
| 258 | group_id, |
| 259 | ): |
| 260 | tid, _, _ = cute.arch.thread_idx() |
| 261 | warp_id = cute.arch.make_warp_uniform(tid // 32) |
| 262 | lane_id = tid % 32 |
| 263 | head_id = group_id * self.num_warps + warp_id |
| 264 | |
| 265 | cute.arch.griddepcontrol_wait() |
| 266 | |
| 267 | pos = positions[token_id] |
| 268 | # no need to do Q calculation for MHA |
| 269 | if cutlass.const_expr(self.nope_dim == 0): |
| 270 | q0, q1 = cvt.bf16x2_to_fp32x2( |
| 271 | cute.recast_tensor(q_pe, Uint32)[token_id, head_id, lane_id] |
| 272 | ) |
| 273 | cos = q_pe_rope_cache[pos, lane_id].to(Float32) |
| 274 | sin = q_pe_rope_cache[pos, 32 + lane_id].to(Float32) |
| 275 | result = cvt.fp32x2_to_bf16x2(q0 * cos - q1 * sin, q1 * cos + q0 * sin) |
| 276 | cute.arch.griddepcontrol_launch_dependents() |
| 277 | cute.recast_tensor(mqa_output, Uint32)[token_id, head_id, lane_id] = result |
| 278 | return |
| 279 | inv_scale = 1.0 / q_scale[0] |
| 280 | |
| 281 | cp_op = cute.nvgpu.CopyUniversalOp() |
| 282 | cp_32B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256) |
| 283 | cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) |
| 284 | cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32) |
| 285 | cp_2B = cute.make_copy_atom(cp_op, Uint8, num_bits_per_copy=16) |
| 286 | |
| 287 | ##### issue all loads asap ##### |
| 288 | rQ_nope_bf16 = cute.make_rmem_tensor(16, BFloat16) |
| 289 | rQ_rope_bf16 = cute.make_rmem_tensor(2, BFloat16) |
| 290 | |
| 291 | src_ql_nope = cute.local_tile( |
| 292 | ql_nope[token_id, head_id, None], (16,), (lane_id,) |
| 293 | ) |
| 294 | src_q_rope = cute.local_tile(q_pe[token_id, head_id, None], (2,), (lane_id,)) |
| 295 | cute.copy(cp_32B, src_ql_nope, rQ_nope_bf16) |
| 296 | cute.copy(cp_4B, src_q_rope, rQ_rope_bf16) |
| 297 | |
| 298 | rCos_raw = q_pe_rope_cache[pos, 0 + lane_id] |
| 299 | rSin_raw = q_pe_rope_cache[pos, 32 + lane_id] |
| 300 | |
| 301 | ##### process NoPE ##### |
| 302 | rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16).load() * inv_scale |
| 303 | rQ_nope_f8 = cute.make_rmem_tensor(16, Float8E4M3FN) |
| 304 | rQ_nope_f8.store(rQ_nope_f32.to(Float8E4M3FN)) |
| 305 | dst_Q_nope = cute.local_tile( |
| 306 | mqa_output[token_id, head_id, None], (16,), (lane_id,) |
| 307 | ) |
| 308 | cute.copy(cp_16B, rQ_nope_f8, dst_Q_nope) |
| 309 | |
| 310 | ##### process RoPE ###### |
| 311 | rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16) |
| 312 | rCos = rCos_raw.to(Float32) |
| 313 | rSin = rSin_raw.to(Float32) |
| 314 | r0 = (rQ_rope_f32[0] * rCos - rQ_rope_f32[1] * rSin) * inv_scale |
| 315 | r1 = (rQ_rope_f32[1] * rCos + rQ_rope_f32[0] * rSin) * inv_scale |
| 316 | |
| 317 | cute.arch.griddepcontrol_launch_dependents() |
| 318 | |
| 319 | # TensorSSA fp32->fp8 cvt has a bug. rely on direct PTX |
| 320 | rQ_rope_f8 = cute.make_rmem_tensor(2, Float8E4M3FN) |
| 321 | cute.recast_tensor(rQ_rope_f8, Uint16)[0] = cvt.fp32x2_to_fp8x2(r0, r1) |
| 322 | dst_Q_rope = cute.local_tile( |
| 323 | mqa_output[token_id, head_id, None], (2,), (256 + lane_id,) |
| 324 | ) |
| 325 | cute.copy(cp_2B, rQ_rope_f8, dst_Q_rope) |
| 326 | |
| 327 | @cute.jit |
| 328 | def indexer( |
| 329 | self, |
| 330 | positions: cute.Tensor, |
| 331 | idx_q: cute.Tensor, |
| 332 | idx_q_rope_cache: cute.Tensor, |
| 333 | idx_weights: cute.Tensor, |
| 334 | idx_q_fp8: cute.Tensor, |
| 335 | idx_weights_out: cute.Tensor, |
| 336 | weight_scale: Float32, |
| 337 | token_id, |
| 338 | group_id, |
| 339 | ): |
| 340 | tid, _, _ = cute.arch.thread_idx() |
| 341 | subwarp_id = tid // 8 |
| 342 | sublane_id = tid % 8 |
| 343 | head_id = group_id * (4 * self.num_warps) + subwarp_id |
| 344 | cute.arch.griddepcontrol_wait() |
| 345 | |
| 346 | pos = positions[token_id] |
| 347 | |
| 348 | cp_op = cute.nvgpu.CopyUniversalOp() |
| 349 | cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) |
| 350 | cp_8B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) |
| 351 | cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32) |
| 352 | |
| 353 | ##### issue all loads first ##### |
| 354 | rQ_rope_bf16 = cute.make_rmem_tensor(8, BFloat16) |
| 355 | if cutlass.const_expr(self.index_rope_interleave): |
| 356 | src_idx_q_rope = cute.local_tile( |
| 357 | idx_q[token_id, head_id, None], (8,), (sublane_id,) |
| 358 | ) |
| 359 | cute.copy(cp_16B, src_idx_q_rope, rQ_rope_bf16) |
| 360 | else: |
| 361 | src_idx_q_rope = cute.zipped_divide( |
| 362 | idx_q[token_id, head_id, None], (4,) |
| 363 | ) # (4,32) |
| 364 | cute.copy( |
| 365 | cp_8B, |
| 366 | src_idx_q_rope[None, 0 + sublane_id], |
| 367 | cute.local_tile(rQ_rope_bf16, (4,), (0,)), |
| 368 | ) |
| 369 | cute.copy( |
| 370 | cp_8B, |
| 371 | src_idx_q_rope[None, 8 + sublane_id], |
| 372 | cute.local_tile(rQ_rope_bf16, (4,), (1,)), |
| 373 | ) |
| 374 | |
| 375 | rQ_nope_bf16 = cute.make_rmem_tensor(8, BFloat16) |
| 376 | src_idx_q_nope = cute.local_tile( |
| 377 | idx_q[token_id, head_id, None], (8,), (8 + sublane_id,) |
| 378 | ) |
| 379 | cute.copy(cp_16B, src_idx_q_nope, rQ_nope_bf16) |
| 380 | |
| 381 | rCos_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type) |
| 382 | rSin_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type) |
| 383 | rope_cache_view = cute.zipped_divide(idx_q_rope_cache[pos, None], (4,)) |
| 384 | |
| 385 | if cutlass.const_expr(idx_q_rope_cache.element_type == Float32): |
| 386 | cute.copy(cp_16B, rope_cache_view[None, sublane_id], rCos_raw) |
| 387 | cute.copy(cp_16B, rope_cache_view[None, 8 + sublane_id], rSin_raw) |
| 388 | elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16): |
| 389 | cute.copy(cp_8B, rope_cache_view[None, sublane_id], rCos_raw) |
| 390 | cute.copy(cp_8B, rope_cache_view[None, 8 + sublane_id], rSin_raw) |
| 391 | |
| 392 | # unpack to FP32 |
| 393 | rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16) |
| 394 | rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16) |
| 395 | if cutlass.const_expr(idx_q_rope_cache.element_type == Float32): |
| 396 | rCos = rCos_raw |
| 397 | rSin = rSin_raw |
| 398 | elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16): |
| 399 | rCos = cvt.bf16x2_to_fp32x2(rCos_raw) |
| 400 | rSin = cvt.bf16x2_to_fp32x2(rSin_raw) |
| 401 | |
| 402 | # apply rope |
| 403 | for i in cutlass.range_constexpr(4): |
| 404 | if cutlass.const_expr(self.index_rope_interleave): |
| 405 | r0 = rQ_rope_f32[i * 2 + 0] * rCos[i] - rQ_rope_f32[i * 2 + 1] * rSin[i] |
| 406 | r1 = rQ_rope_f32[i * 2 + 1] * rCos[i] + rQ_rope_f32[i * 2 + 0] * rSin[i] |
| 407 | rQ_rope_f32[i * 2 + 0] = r0 |
| 408 | rQ_rope_f32[i * 2 + 1] = r1 |
| 409 | else: |
| 410 | r0 = rQ_rope_f32[0 + i] * rCos[i] - rQ_rope_f32[4 + i] * rSin[i] |
| 411 | r1 = rQ_rope_f32[4 + i] * rCos[i] + rQ_rope_f32[0 + i] * rSin[i] |
| 412 | rQ_rope_f32[0 + i] = r0 |
| 413 | rQ_rope_f32[4 + i] = r1 |
| 414 | |
| 415 | # amax |
| 416 | amax = Float32(1e-4) |
| 417 | for i in cutlass.range_constexpr(8): |
| 418 | amax = cute.arch.fmax(amax, cute.math.absf(rQ_rope_f32[i])) |
| 419 | amax = cute.arch.fmax(amax, cute.math.absf(rQ_nope_f32[i])) |
| 420 | |
| 421 | # warp reduction among 8 lanes |
| 422 | for i in cutlass.range_constexpr(3): |
| 423 | other = cute.arch.shuffle_sync_bfly(amax, 1 << i) |
| 424 | amax = cute.arch.fmax(amax, other) |
| 425 | |
| 426 | # compute scale from amax |
| 427 | # exp2(ceil(log2(scale))) via bit manipulation |
| 428 | scale = amax * (1.0 / 448.0) |
| 429 | bits = scale.bitcast(Uint32) |
| 430 | exp_bits = (bits + Uint32(0x007FFFFF)) & Uint32(0x7F800000) |
| 431 | scale = exp_bits.bitcast(Float32) |
| 432 | inv_scale = (Uint32(0x7F000000) - exp_bits).bitcast(Float32) |
| 433 | |
| 434 | for i in cutlass.range_constexpr(8): |
| 435 | rQ_nope_f32[i] *= inv_scale |
| 436 | rQ_rope_f32[i] *= inv_scale |
| 437 | |
| 438 | cute.arch.griddepcontrol_launch_dependents() |
| 439 | |
| 440 | # quantize and store |
| 441 | rQ_nope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN) |
| 442 | rQ_nope_f8.store(rQ_nope_f32.load().to(Float8E4M3FN)) |
| 443 | dst_idx_q_nope = cute.local_tile( |
| 444 | idx_q_fp8[token_id, head_id, None], (8,), (8 + sublane_id,) |
| 445 | ) |
| 446 | cute.copy(cp_8B, rQ_nope_f8, dst_idx_q_nope) |
| 447 | |
| 448 | rQ_rope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN) |
| 449 | rQ_rope_f8.store(rQ_rope_f32.load().to(Float8E4M3FN)) |
| 450 | if cutlass.const_expr(self.index_rope_interleave): |
| 451 | dst_idx_q_rope = cute.local_tile( |
| 452 | idx_q_fp8[token_id, head_id, None], (8,), (sublane_id,) |
| 453 | ) |
| 454 | cute.copy(cp_8B, rQ_rope_f8, dst_idx_q_rope) |
| 455 | else: |
| 456 | dst_idx_q_rope = cute.zipped_divide( |
| 457 | idx_q_fp8[token_id, head_id, None], (4,) |
| 458 | ) # (4,32) |
| 459 | cute.copy( |
| 460 | cp_4B, |
| 461 | cute.local_tile(rQ_rope_f8, (4,), (0,)), |
| 462 | dst_idx_q_rope[None, 0 + sublane_id], |
| 463 | ) |
| 464 | cute.copy( |
| 465 | cp_4B, |
| 466 | cute.local_tile(rQ_rope_f8, (4,), (1,)), |
| 467 | dst_idx_q_rope[None, 8 + sublane_id], |
| 468 | ) |
| 469 | |
| 470 | # scale indexer weights |
| 471 | if sublane_id == 0: |
| 472 | w = idx_weights[token_id, head_id].to(Float32) |
| 473 | idx_weights_out[token_id, head_id] = w * scale * weight_scale |
| 474 | |
| 475 | @staticmethod |
| 476 | @cache |
| 477 | def compile( |
| 478 | rope_dim: int, |
| 479 | nope_dim: int, |
| 480 | num_heads: int, |
| 481 | rope_type: type[cutlass.Numeric], |
| 482 | idx_dim: int, |
| 483 | num_idx_heads: int, |
| 484 | idx_rope_type: type[cutlass.Numeric] | None, |
| 485 | idx_weights_type: type[cutlass.Numeric] | None, |
| 486 | index_rope_interleave: bool, |
| 487 | ): |
| 488 | num_tokens = cute.sym_int() |
| 489 | max_pos = cute.sym_int() |
| 490 | |
| 491 | positions = _make_fake_tensor(Int64, (num_tokens,), divisibility=1) |
| 492 | q_pe = _make_fake_tensor( |
| 493 | BFloat16, (num_tokens, num_heads, rope_dim), divisibility=16 |
| 494 | ) |
| 495 | rope_cache = _make_fake_tensor(rope_type, (max_pos, rope_dim), divisibility=8) |
| 496 | ql_nope = ( |
| 497 | _make_fake_tensor( |
| 498 | BFloat16, (num_tokens, num_heads, nope_dim), divisibility=16 |
| 499 | ) |
| 500 | if nope_dim |
| 501 | else None |
| 502 | ) |
| 503 | q_scale = _make_fake_tensor(Float32, (1,), divisibility=4) if nope_dim else None |
| 504 | mqa_output = _make_fake_tensor( |
| 505 | Float8E4M3FN if nope_dim else BFloat16, |
| 506 | (num_tokens, num_heads, nope_dim + rope_dim), |
| 507 | divisibility=16, |
| 508 | ) |
| 509 | |
| 510 | if idx_rope_type is not None: |
| 511 | index_q = _make_fake_tensor( |
| 512 | BFloat16, (num_tokens, num_idx_heads, idx_dim), divisibility=16 |
| 513 | ) |
| 514 | index_rope_cache = _make_fake_tensor( |
| 515 | idx_rope_type, (max_pos, rope_dim), divisibility=8 |
| 516 | ) |
| 517 | index_weights = _make_fake_tensor( |
| 518 | idx_weights_type, (num_tokens, num_idx_heads), divisibility=8 |
| 519 | ) |
| 520 | index_q_fp8 = _make_fake_tensor( |
| 521 | Float8E4M3FN, (num_tokens, num_idx_heads, idx_dim), divisibility=16 |
| 522 | ) |
| 523 | index_weights_out = _make_fake_tensor( |
| 524 | Float32, (num_tokens, num_idx_heads), divisibility=4 |
| 525 | ) |
| 526 | else: |
| 527 | index_q = index_rope_cache = index_q_fp8 = None |
| 528 | index_weights = index_weights_out = None |
| 529 | |
| 530 | kernel = FusedQKernel( |
| 531 | rope_dim, |
| 532 | nope_dim, |
| 533 | num_heads, |
| 534 | idx_dim, |
| 535 | num_idx_heads, |
| 536 | index_rope_interleave, |
| 537 | ) |
| 538 | stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) |
| 539 | return cute.compile( |
| 540 | kernel, |
| 541 | positions, |
| 542 | q_pe, |
| 543 | rope_cache, |
| 544 | ql_nope, |
| 545 | q_scale, |
| 546 | mqa_output, |
| 547 | index_q, |
| 548 | index_rope_cache, |
| 549 | index_weights, |
| 550 | index_q_fp8, |
| 551 | index_weights_out, |
| 552 | Float32(0.0), |
| 553 | stream, |
| 554 | options="--enable-tvm-ffi", |
| 555 | ) |