@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3756bf1e2b

[Kernel][SM100] Add a CuTeDSL fused query kernel (#49792)

Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

xiaozhoupy committed Aug 5, 2026 at 06:23 UTC 3756bf1e2ba0909500fbe77e03fe02eafb7f7080
6 files changed +699 -1
benchmarks/kernels/benchmark_fused_q_cutedsl.py new
+95
@@ -0,0 +1,95 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Compare the CuTeDSL and Triton implementations of the DSA fused query op.
4 +
5 +``fused_q`` dispatches to the CuTeDSL kernel on SM100 when the dtypes and
6 +shapes are supported and falls back to Triton otherwise. This benchmark forces
7 +each path so the two can be compared directly.
8 +"""
9 +
10 +import argparse
11 +from unittest.mock import patch
12 +
13 +import torch
14 +
15 +import vllm.models.deepseek_v32.common.kernels as K
16 +from vllm.triton_utils import triton
17 +
18 +# GLM-5.2 / DeepSeek-V3.2 DSA dimensions.
19 +NOPE_DIM = 512
20 +ROPE_DIM = 64
21 +INDEX_HEADS = 32
22 +INDEX_HEAD_DIM = 128
23 +MAX_POS = 16384
24 +
25 +# num_heads is per-rank: 128 attention heads over TP 8/4/2.
26 +NUM_HEADS = [16, 32, 64]
27 +# MTP=5 makes a concurrency-N decode batch 6*N tokens.
28 +NUM_TOKENS = [1, 6, 24, 96, 384, 1536, 4096]
29 +
30 +
31 +def make_inputs(num_tokens: int, num_heads: int):
32 + dev = "cuda"
33 + bf16 = torch.bfloat16
34 + return dict(
35 + positions=torch.randint(0, MAX_POS, (num_tokens,), device=dev),
36 + q_pe=torch.randn(num_tokens, num_heads, ROPE_DIM, device=dev, dtype=bf16),
37 + q_pe_cos_sin_cache=torch.randn(MAX_POS, ROPE_DIM, device=dev, dtype=bf16),
38 + index_q=torch.randn(
39 + num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=bf16
40 + ),
41 + index_q_cos_sin_cache=torch.randn(MAX_POS, ROPE_DIM, device=dev, dtype=bf16),
42 + ql_nope=torch.randn(num_tokens, num_heads, NOPE_DIM, device=dev, dtype=bf16),
43 + q_scale=torch.tensor([0.5], device=dev, dtype=torch.float32),
44 + index_weights=torch.randn(
45 + num_tokens, INDEX_HEADS, device=dev, dtype=torch.float32
46 + ),
47 + index_weights_softmax_scale=INDEX_HEAD_DIM**-0.5,
48 + index_weights_head_scale=INDEX_HEADS**-0.5,
49 + )
50 +
51 +
52 +def run(kwargs, *, cutedsl: bool):
53 + # is_fused_q_cutedsl_supported gates the dispatch inside fused_q.
54 + from vllm.models.deepseek_v32.nvidia.ops import fused_q_cutedsl as C
55 +
56 + with patch.object(C, "is_fused_q_cutedsl_supported", lambda *_, **__: cutedsl):
57 + K.fused_q(**kwargs)
58 +
59 +
60 +@triton.testing.perf_report(
61 + [
62 + triton.testing.Benchmark(
63 + x_names=["num_tokens"],
64 + x_vals=NUM_TOKENS,
65 + line_arg="provider",
66 + line_vals=["triton", "cutedsl"],
67 + line_names=["Triton", "CuTeDSL"],
68 + styles=[("blue", "-"), ("green", "-")],
69 + ylabel="us",
70 + plot_name=f"fused_q-num_heads-{nh}",
71 + args={"num_heads": nh},
72 + )
73 + for nh in NUM_HEADS
74 + ]
75 +)
76 +def benchmark(num_tokens, num_heads, provider):
77 + kwargs = make_inputs(num_tokens, num_heads)
78 + fn = lambda: run(kwargs, cutedsl=provider == "cutedsl") # noqa: E731
79 + fn() # warm up JIT / kernel compilation before timing
80 + ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
81 + return 1000 * ms, 1000 * max_ms, 1000 * min_ms
82 +
83 +
84 +if __name__ == "__main__":
85 + parser = argparse.ArgumentParser(description=__doc__)
86 + parser.add_argument("--save-path", type=str, default=None)
87 + args = parser.parse_args()
88 +
89 + from vllm.platforms import current_platform
90 + from vllm.utils.import_utils import has_cutedsl
91 +
92 + if not (current_platform.has_device_capability(100) and has_cutedsl()):
93 + raise SystemExit("fused_q CuTeDSL requires an SM100 device and cutlass-dsl")
94 +
95 + benchmark.run(print_data=True, show_plots=False, save_path=args.save_path)
vllm/cute_utils/__init__.py
+1
@@ -19,6 +19,7 @@ from cutlass.cutlass_dsl import T, dsl_user_op
19 _TORCH_TO_CUTE_DTYPE = {
20 torch.bfloat16: BFloat16,
21 torch.float8_e4m3fn: Float8E4M3FN,
22 + torch.float32: Float32,
23 }
24
25 _CUTE_TO_PTX_DTYPE = {
vllm/cute_utils/cvt.py
+22 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from cutlass import Constexpr, Float32, Uint32, cute
3 +from cutlass import BFloat16, Constexpr, Float32, Uint16, Uint32, cute
4 from cutlass._mlir import ir
5 from cutlass._mlir.dialects import llvm, vector
6 from cutlass.cutlass_dsl import T, dsl_user_op
@@ -38,6 +38,11 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
38 )
39
40 elif isinstance(data, (cute.Tensor, cute.TensorSSA)):
41 + # recast to 32-bit registers if needed
42 + if data.element_type == BFloat16:
43 + data = cute.recast_tensor(data, Uint32)
44 + assert data.element_type == Uint32
45 +
46 # NOTE: the output is always 1D
47 size = cute.size(data.shape)
48 out = cute.make_rmem_tensor(size * 2, Float32)
@@ -84,6 +89,22 @@ def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
89 return cute.TensorSSA(vec, 2, Uint32)
90
91
92 +@dsl_user_op
93 +def fp32x2_to_fp8x2(a0: Float32, a1: Float32, *, loc=None, ip=None) -> Uint16:
94 + out = llvm.inline_asm(
95 + T.i16(),
96 + [
97 + a0.ir_value(loc=loc, ip=ip),
98 + a1.ir_value(loc=loc, ip=ip),
99 + ],
100 + "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;",
101 + "=h,f,f",
102 + has_side_effects=False,
103 + is_align_stack=False,
104 + )
105 + return Uint16(out)
106 +
107 +
108 @dsl_user_op
109 def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
110 out = llvm.inline_asm(
vllm/models/deepseek_v32/common/kernels.py
+42
@@ -1,7 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +
5 import torch
6
7 +from vllm.platforms import current_platform
8 from vllm.triton_utils import tl, triton
9
10 # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the
@@ -773,10 +776,12 @@ def fused_q(
776 ``(ql_nope, q_pe)`` tuple the backend expects.
777 """
778 assert positions.ndim == 1
779 + assert positions.dtype == torch.int64
780 assert q_pe.ndim == 3
781 assert q_pe_cos_sin_cache.ndim == 2
782 assert ql_nope.ndim == 3
783 assert ql_nope.shape[:2] == q_pe.shape[:2]
784 + assert q_scale.dtype == torch.float32 and q_scale.numel() == 1
785
786 num_tokens = positions.shape[0]
787 num_q_heads = q_pe.shape[1]
@@ -794,6 +799,23 @@ def fused_q(
799 assert index_weights is not None
800 num_index_q_heads = index_q.shape[1]
801 index_q_head_dim = index_q.shape[2]
802 + # fused_q is shared with the ROCm path, and the CuTeDSL module imports
803 + # cutlass at module scope, so only reach for it on CUDA.
804 + cutedsl_kernel: Callable[..., None] | None = None
805 + if current_platform.is_cuda():
806 + from vllm.models.deepseek_v32.nvidia.ops.fused_q_cutedsl import (
807 + fused_q_cutedsl,
808 + is_fused_q_cutedsl_supported,
809 + )
810 +
811 + if is_fused_q_cutedsl_supported(
812 + q_pe,
813 + index_q,
814 + ql_nope,
815 + has_indexer=has_indexer,
816 + quantize_mqa=quantize_mqa,
817 + ):
818 + cutedsl_kernel = fused_q_cutedsl
819 grid_heads = max(mqa_grid_heads, num_index_q_heads)
820 if quantize_mqa:
821 # fp8 path: pack [ql_nope; q_pe] into a single fp8 tensor.
@@ -815,6 +837,26 @@ def fused_q(
837
838 index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
839 index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
840 + if cutedsl_kernel is not None:
841 + cutedsl_kernel(
842 + positions,
843 + q_pe,
844 + q_pe_cos_sin_cache,
845 + ql_nope,
846 + q_scale,
847 + mqa_q,
848 + index_q,
849 + index_q_cos_sin_cache,
850 + index_weights,
851 + index_weights_softmax_scale,
852 + index_weights_head_scale,
853 + index_q_fp8,
854 + index_weights_out,
855 + has_indexer=has_indexer,
856 + index_rope_interleave=index_rope_interleave,
857 + )
858 + return index_q_fp8, index_weights_out, mqa_q
859 +
860 _fused_q_kernel[(3, num_tokens, grid_heads)](
861 positions,
862 q_pe,
vllm/models/deepseek_v32/nvidia/ops/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py new
+537
@@ -0,0 +1,537 @@
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 quantize_mqa
39 + and q_pe.dtype == ql_nope.dtype == torch.bfloat16
40 + and q_pe.shape[-1] == 64
41 + and ql_nope.shape[-1] == 512
42 + # One warp per head group; a non-multiple would trip the kernel's own
43 + # assert instead of falling back to Triton.
44 + and q_pe.shape[1] % 4 == 0
45 + ):
46 + return False
47 + return not has_indexer or (
48 + index_q is not None
49 + and index_q.dtype == torch.bfloat16
50 + and index_q.shape[1] % 16 == 0
51 + and index_q.shape[-1] == 128
52 + )
53 +
54 +
55 +def fused_q_cutedsl(
56 + positions: torch.Tensor,
57 + q_pe: torch.Tensor,
58 + rope_cache: torch.Tensor,
59 + ql_nope: torch.Tensor,
60 + q_scale: torch.Tensor,
61 + mqa_output: torch.Tensor,
62 + idx_q: torch.Tensor,
63 + idx_rope_cache: torch.Tensor,
64 + idx_weights: torch.Tensor,
65 + idx_weights_softmax_scale: float,
66 + idx_weights_head_scale: float,
67 + idx_q_fp8: torch.Tensor,
68 + idx_weights_out: torch.Tensor,
69 + has_indexer: bool = True,
70 + index_rope_interleave: bool = True,
71 +) -> None:
72 + _, num_heads, rope_dim = q_pe.shape
73 + _, _, nope_dim = ql_nope.shape
74 + _, num_idx_heads, idx_dim = idx_q.shape
75 +
76 + if has_indexer:
77 + idx_rope_type = _TORCH_TO_CUTE_DTYPE[idx_rope_cache.dtype]
78 + idx_weights_type = _TORCH_TO_CUTE_DTYPE[idx_weights.dtype]
79 + else:
80 + idx_dim = num_idx_heads = 0
81 + idx_q = idx_rope_cache = idx_q_fp8 = None
82 + idx_weights = idx_weights_out = None
83 + idx_rope_type = idx_weights_type = None
84 +
85 + rope_type = _TORCH_TO_CUTE_DTYPE[rope_cache.dtype]
86 + compiled = FusedQKernel.compile(
87 + rope_dim,
88 + nope_dim,
89 + num_heads,
90 + rope_type,
91 + idx_dim,
92 + num_idx_heads,
93 + idx_rope_type,
94 + idx_weights_type,
95 + index_rope_interleave,
96 + )
97 + compiled(
98 + positions,
99 + q_pe,
100 + rope_cache,
101 + ql_nope,
102 + q_scale.view(1),
103 + mqa_output,
104 + idx_q,
105 + idx_rope_cache,
106 + idx_weights,
107 + idx_q_fp8,
108 + idx_weights_out,
109 + float(idx_weights_softmax_scale * idx_weights_head_scale),
110 + )
111 +
112 +
113 +class FusedQKernel:
114 + def __init__(
115 + self,
116 + rope_dim: int,
117 + nope_dim: int,
118 + num_heads: int,
119 + idx_dim: int,
120 + num_idx_heads: int,
121 + index_rope_interleave: bool,
122 + ) -> None:
123 + assert rope_dim == 64
124 + assert nope_dim == 512
125 + assert idx_dim in (128, 0)
126 +
127 + self.rope_dim = rope_dim
128 + self.nope_dim = nope_dim
129 + self.num_heads = num_heads
130 + self.idx_dim = idx_dim
131 + self.num_idx_heads = num_idx_heads
132 + self.index_rope_interleave = index_rope_interleave
133 +
134 + # mqa: rope_dim=64, nope_dim=512, num_heads=64/TP
135 + # indexer: rope_dim=64, nope_dim=64, num_heads=32
136 +
137 + self.num_warps = 4
138 + assert num_heads % self.num_warps == 0
139 + assert num_idx_heads % (4 * self.num_warps) == 0
140 + self.num_ctas_per_tok = num_heads // self.num_warps
141 + self.num_ctas_per_idx_tok = num_idx_heads // (4 * self.num_warps)
142 +
143 + @cute.jit
144 + def __call__(
145 + self,
146 + positions: cute.Tensor,
147 + q_pe: cute.Tensor,
148 + rope_cache: cute.Tensor,
149 + ql_nope: cute.Tensor,
150 + q_scale: cute.Tensor,
151 + mqa_output: cute.Tensor,
152 + idx_q: cute.Tensor,
153 + idx_rope_cache: cute.Tensor,
154 + idx_weights: cute.Tensor,
155 + idx_q_fp8: cute.Tensor,
156 + idx_weights_out: cute.Tensor,
157 + weight_scale: Float32,
158 + stream: CUstream,
159 + ):
160 + num_tokens = positions.shape[0]
161 + if cutlass.const_expr(self.idx_dim == 0):
162 + grid = (num_tokens, self.num_ctas_per_tok, 1)
163 + else:
164 + num_mqa_ctas = num_tokens * self.num_ctas_per_tok
165 + num_idx_ctas = num_tokens * self.num_ctas_per_idx_tok
166 + grid = (num_mqa_ctas + num_idx_ctas, 1, 1)
167 +
168 + self.kernel(
169 + positions,
170 + q_pe,
171 + rope_cache,
172 + ql_nope,
173 + q_scale,
174 + mqa_output,
175 + idx_q,
176 + idx_rope_cache,
177 + idx_weights,
178 + idx_q_fp8,
179 + idx_weights_out,
180 + weight_scale,
181 + ).launch(
182 + grid=grid,
183 + block=(self.num_warps * 32, 1, 1),
184 + stream=stream,
185 + use_pdl=True,
186 + )
187 +
188 + @cute.kernel
189 + def kernel(
190 + self,
191 + positions: cute.Tensor,
192 + q_pe: cute.Tensor,
193 + q_pe_rope_cache: cute.Tensor,
194 + ql_nope: cute.Tensor,
195 + q_scale: cute.Tensor,
196 + mqa_output: cute.Tensor,
197 + idx_q: cute.Tensor,
198 + idx_q_rope_cache: cute.Tensor,
199 + idx_weights: cute.Tensor,
200 + idx_q_fp8: cute.Tensor,
201 + idx_weights_out: cute.Tensor,
202 + weight_scale: Float32,
203 + ):
204 + if cutlass.const_expr(self.idx_dim == 0):
205 + token_id, group_id, _ = cute.arch.block_idx()
206 + self.mqa(
207 + positions,
208 + q_pe,
209 + q_pe_rope_cache,
210 + ql_nope,
211 + q_scale,
212 + mqa_output,
213 + token_id,
214 + group_id,
215 + )
216 + else:
217 + # CTA-specialization
218 + bid, _, _ = cute.arch.block_idx()
219 + num_mqa_ctas = positions.shape[0] * self.num_ctas_per_tok
220 + if bid < num_mqa_ctas:
221 + self.mqa(
222 + positions,
223 + q_pe,
224 + q_pe_rope_cache,
225 + ql_nope,
226 + q_scale,
227 + mqa_output,
228 + bid // self.num_ctas_per_tok,
229 + bid % self.num_ctas_per_tok,
230 + )
231 + else:
232 + bid -= num_mqa_ctas
233 + self.indexer(
234 + positions,
235 + idx_q,
236 + idx_q_rope_cache,
237 + idx_weights,
238 + idx_q_fp8,
239 + idx_weights_out,
240 + weight_scale,
241 + bid // self.num_ctas_per_idx_tok,
242 + bid % self.num_ctas_per_idx_tok,
243 + )
244 +
245 + @cute.jit
246 + def mqa(
247 + self,
248 + positions: cute.Tensor,
249 + q_pe: cute.Tensor,
250 + q_pe_rope_cache: cute.Tensor,
251 + ql_nope: cute.Tensor,
252 + q_scale: cute.Tensor,
253 + mqa_output: cute.Tensor,
254 + token_id,
255 + group_id,
256 + ):
257 + tid, _, _ = cute.arch.thread_idx()
258 + warp_id = cute.arch.make_warp_uniform(tid // 32)
259 + lane_id = tid % 32
260 + head_id = group_id * self.num_warps + warp_id
261 +
262 + cute.arch.griddepcontrol_wait()
263 +
264 + pos = positions[token_id]
265 + inv_scale = 1.0 / q_scale[0]
266 +
267 + cp_op = cute.nvgpu.CopyUniversalOp()
268 + cp_32B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
269 + cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
270 + cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32)
271 + cp_2B = cute.make_copy_atom(cp_op, Uint8, num_bits_per_copy=16)
272 +
273 + ##### issue all loads asap #####
274 + rQ_nope_bf16 = cute.make_rmem_tensor(16, BFloat16)
275 + rQ_rope_bf16 = cute.make_rmem_tensor(2, BFloat16)
276 +
277 + src_ql_nope = cute.local_tile(
278 + ql_nope[token_id, head_id, None], (16,), (lane_id,)
279 + )
280 + src_q_rope = cute.local_tile(q_pe[token_id, head_id, None], (2,), (lane_id,))
281 + cute.copy(cp_32B, src_ql_nope, rQ_nope_bf16)
282 + cute.copy(cp_4B, src_q_rope, rQ_rope_bf16)
283 +
284 + rCos_raw = q_pe_rope_cache[pos, 0 + lane_id]
285 + rSin_raw = q_pe_rope_cache[pos, 32 + lane_id]
286 +
287 + ##### process NoPE #####
288 + rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16).load() * inv_scale
289 + rQ_nope_f8 = cute.make_rmem_tensor(16, Float8E4M3FN)
290 + rQ_nope_f8.store(rQ_nope_f32.to(Float8E4M3FN))
291 + dst_Q_nope = cute.local_tile(
292 + mqa_output[token_id, head_id, None], (16,), (lane_id,)
293 + )
294 + cute.copy(cp_16B, rQ_nope_f8, dst_Q_nope)
295 +
296 + ##### process RoPE ######
297 + rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16)
298 + rCos = rCos_raw.to(Float32)
299 + rSin = rSin_raw.to(Float32)
300 + r0 = (rQ_rope_f32[0] * rCos - rQ_rope_f32[1] * rSin) * inv_scale
301 + r1 = (rQ_rope_f32[1] * rCos + rQ_rope_f32[0] * rSin) * inv_scale
302 +
303 + cute.arch.griddepcontrol_launch_dependents()
304 +
305 + # TensorSSA fp32->fp8 cvt has a bug. rely on direct PTX
306 + rQ_rope_f8 = cute.make_rmem_tensor(2, Float8E4M3FN)
307 + cute.recast_tensor(rQ_rope_f8, Uint16)[0] = cvt.fp32x2_to_fp8x2(r0, r1)
308 + dst_Q_rope = cute.local_tile(
309 + mqa_output[token_id, head_id, None], (2,), (256 + lane_id,)
310 + )
311 + cute.copy(cp_2B, rQ_rope_f8, dst_Q_rope)
312 +
313 + @cute.jit
314 + def indexer(
315 + self,
316 + positions: cute.Tensor,
317 + idx_q: cute.Tensor,
318 + idx_q_rope_cache: cute.Tensor,
319 + idx_weights: cute.Tensor,
320 + idx_q_fp8: cute.Tensor,
321 + idx_weights_out: cute.Tensor,
322 + weight_scale: Float32,
323 + token_id,
324 + group_id,
325 + ):
326 + tid, _, _ = cute.arch.thread_idx()
327 + subwarp_id = tid // 8
328 + sublane_id = tid % 8
329 + head_id = group_id * (4 * self.num_warps) + subwarp_id
330 + cute.arch.griddepcontrol_wait()
331 +
332 + pos = positions[token_id]
333 +
334 + cp_op = cute.nvgpu.CopyUniversalOp()
335 + cp_16B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
336 + cp_8B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
337 + cp_4B = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=32)
338 +
339 + ##### issue all loads first #####
340 + rQ_rope_bf16 = cute.make_rmem_tensor(8, BFloat16)
341 + if cutlass.const_expr(self.index_rope_interleave):
342 + src_idx_q_rope = cute.local_tile(
343 + idx_q[token_id, head_id, None], (8,), (sublane_id,)
344 + )
345 + cute.copy(cp_16B, src_idx_q_rope, rQ_rope_bf16)
346 + else:
347 + src_idx_q_rope = cute.zipped_divide(
348 + idx_q[token_id, head_id, None], (4,)
349 + ) # (4,32)
350 + cute.copy(
351 + cp_8B,
352 + src_idx_q_rope[None, 0 + sublane_id],
353 + cute.local_tile(rQ_rope_bf16, (4,), (0,)),
354 + )
355 + cute.copy(
356 + cp_8B,
357 + src_idx_q_rope[None, 8 + sublane_id],
358 + cute.local_tile(rQ_rope_bf16, (4,), (1,)),
359 + )
360 +
361 + rQ_nope_bf16 = cute.make_rmem_tensor(8, BFloat16)
362 + src_idx_q_nope = cute.local_tile(
363 + idx_q[token_id, head_id, None], (8,), (8 + sublane_id,)
364 + )
365 + cute.copy(cp_16B, src_idx_q_nope, rQ_nope_bf16)
366 +
367 + rCos_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type)
368 + rSin_raw = cute.make_rmem_tensor(4, idx_q_rope_cache.element_type)
369 + rope_cache_view = cute.zipped_divide(idx_q_rope_cache[pos, None], (4,))
370 +
371 + if cutlass.const_expr(idx_q_rope_cache.element_type == Float32):
372 + cute.copy(cp_16B, rope_cache_view[None, sublane_id], rCos_raw)
373 + cute.copy(cp_16B, rope_cache_view[None, 8 + sublane_id], rSin_raw)
374 + elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16):
375 + cute.copy(cp_8B, rope_cache_view[None, sublane_id], rCos_raw)
376 + cute.copy(cp_8B, rope_cache_view[None, 8 + sublane_id], rSin_raw)
377 +
378 + # unpack to FP32
379 + rQ_rope_f32 = cvt.bf16x2_to_fp32x2(rQ_rope_bf16)
380 + rQ_nope_f32 = cvt.bf16x2_to_fp32x2(rQ_nope_bf16)
381 + if cutlass.const_expr(idx_q_rope_cache.element_type == Float32):
382 + rCos = rCos_raw
383 + rSin = rSin_raw
384 + elif cutlass.const_expr(idx_q_rope_cache.element_type == BFloat16):
385 + rCos = cvt.bf16x2_to_fp32x2(rCos_raw)
386 + rSin = cvt.bf16x2_to_fp32x2(rSin_raw)
387 +
388 + # apply rope
389 + for i in cutlass.range_constexpr(4):
390 + if cutlass.const_expr(self.index_rope_interleave):
391 + r0 = rQ_rope_f32[i * 2 + 0] * rCos[i] - rQ_rope_f32[i * 2 + 1] * rSin[i]
392 + r1 = rQ_rope_f32[i * 2 + 1] * rCos[i] + rQ_rope_f32[i * 2 + 0] * rSin[i]
393 + rQ_rope_f32[i * 2 + 0] = r0
394 + rQ_rope_f32[i * 2 + 1] = r1
395 + else:
396 + r0 = rQ_rope_f32[0 + i] * rCos[i] - rQ_rope_f32[4 + i] * rSin[i]
397 + r1 = rQ_rope_f32[4 + i] * rCos[i] + rQ_rope_f32[0 + i] * rSin[i]
398 + rQ_rope_f32[0 + i] = r0
399 + rQ_rope_f32[4 + i] = r1
400 +
401 + # amax
402 + amax = Float32(1e-4)
403 + for i in cutlass.range_constexpr(8):
404 + amax = cute.arch.fmax(amax, cute.math.absf(rQ_rope_f32[i]))
405 + amax = cute.arch.fmax(amax, cute.math.absf(rQ_nope_f32[i]))
406 +
407 + # warp reduction among 8 lanes
408 + for i in cutlass.range_constexpr(3):
409 + other = cute.arch.shuffle_sync_bfly(amax, 1 << i)
410 + amax = cute.arch.fmax(amax, other)
411 +
412 + # compute scale from amax
413 + # exp2(ceil(log2(scale))) via bit manipulation
414 + scale = amax * (1.0 / 448.0)
415 + bits = scale.bitcast(Uint32)
416 + exp_bits = (bits + Uint32(0x007FFFFF)) & Uint32(0x7F800000)
417 + scale = exp_bits.bitcast(Float32)
418 + inv_scale = (Uint32(0x7F000000) - exp_bits).bitcast(Float32)
419 +
420 + for i in cutlass.range_constexpr(8):
421 + rQ_nope_f32[i] *= inv_scale
422 + rQ_rope_f32[i] *= inv_scale
423 +
424 + cute.arch.griddepcontrol_launch_dependents()
425 +
426 + # quantize and store
427 + rQ_nope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN)
428 + rQ_nope_f8.store(rQ_nope_f32.load().to(Float8E4M3FN))
429 + dst_idx_q_nope = cute.local_tile(
430 + idx_q_fp8[token_id, head_id, None], (8,), (8 + sublane_id,)
431 + )
432 + cute.copy(cp_8B, rQ_nope_f8, dst_idx_q_nope)
433 +
434 + rQ_rope_f8 = cute.make_rmem_tensor(8, Float8E4M3FN)
435 + rQ_rope_f8.store(rQ_rope_f32.load().to(Float8E4M3FN))
436 + if cutlass.const_expr(self.index_rope_interleave):
437 + dst_idx_q_rope = cute.local_tile(
438 + idx_q_fp8[token_id, head_id, None], (8,), (sublane_id,)
439 + )
440 + cute.copy(cp_8B, rQ_rope_f8, dst_idx_q_rope)
441 + else:
442 + dst_idx_q_rope = cute.zipped_divide(
443 + idx_q_fp8[token_id, head_id, None], (4,)
444 + ) # (4,32)
445 + cute.copy(
446 + cp_4B,
447 + cute.local_tile(rQ_rope_f8, (4,), (0,)),
448 + dst_idx_q_rope[None, 0 + sublane_id],
449 + )
450 + cute.copy(
451 + cp_4B,
452 + cute.local_tile(rQ_rope_f8, (4,), (1,)),
453 + dst_idx_q_rope[None, 8 + sublane_id],
454 + )
455 +
456 + # scale indexer weights
457 + if sublane_id == 0:
458 + w = idx_weights[token_id, head_id].to(Float32)
459 + idx_weights_out[token_id, head_id] = w * scale * weight_scale
460 +
461 + @staticmethod
462 + @cache
463 + def compile(
464 + rope_dim: int,
465 + nope_dim: int,
466 + num_heads: int,
467 + rope_type: type[cutlass.Numeric],
468 + idx_dim: int,
469 + num_idx_heads: int,
470 + idx_rope_type: type[cutlass.Numeric] | None,
471 + idx_weights_type: type[cutlass.Numeric] | None,
472 + index_rope_interleave: bool,
473 + ):
474 + num_tokens = cute.sym_int()
475 + max_pos = cute.sym_int()
476 +
477 + positions = _make_fake_tensor(Int64, (num_tokens,), divisibility=1)
478 + q_pe = _make_fake_tensor(
479 + BFloat16, (num_tokens, num_heads, rope_dim), divisibility=16
480 + )
481 + rope_cache = _make_fake_tensor(rope_type, (max_pos, rope_dim), divisibility=8)
482 + ql_nope = _make_fake_tensor(
483 + BFloat16, (num_tokens, num_heads, nope_dim), divisibility=16
484 + )
485 + q_scale = _make_fake_tensor(Float32, (1,), divisibility=4)
486 + mqa_output = _make_fake_tensor(
487 + Float8E4M3FN,
488 + (num_tokens, num_heads, nope_dim + rope_dim),
489 + divisibility=16,
490 + )
491 +
492 + if idx_rope_type is not None:
493 + index_q = _make_fake_tensor(
494 + BFloat16, (num_tokens, num_idx_heads, idx_dim), divisibility=16
495 + )
496 + index_rope_cache = _make_fake_tensor(
497 + idx_rope_type, (max_pos, rope_dim), divisibility=8
498 + )
499 + index_weights = _make_fake_tensor(
500 + idx_weights_type, (num_tokens, num_idx_heads), divisibility=8
501 + )
502 + index_q_fp8 = _make_fake_tensor(
503 + Float8E4M3FN, (num_tokens, num_idx_heads, idx_dim), divisibility=16
504 + )
505 + index_weights_out = _make_fake_tensor(
506 + Float32, (num_tokens, num_idx_heads), divisibility=4
507 + )
508 + else:
509 + index_q = index_rope_cache = index_q_fp8 = None
510 + index_weights = index_weights_out = None
511 +
512 + kernel = FusedQKernel(
513 + rope_dim,
514 + nope_dim,
515 + num_heads,
516 + idx_dim,
517 + num_idx_heads,
518 + index_rope_interleave,
519 + )
520 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
521 + return cute.compile(
522 + kernel,
523 + positions,
524 + q_pe,
525 + rope_cache,
526 + ql_nope,
527 + q_scale,
528 + mqa_output,
529 + index_q,
530 + index_rope_cache,
531 + index_weights,
532 + index_q_fp8,
533 + index_weights_out,
534 + Float32(0.0),
535 + stream,
536 + options="--enable-tvm-ffi",
537 + )