A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 01:15 UTC
[DSv4] Improved fused Indexer Q quant kernel (#41428)
Thien Tran committed
May 9, 2026 at 15:20 UTC
530d37130278d38b07b089f936850537aa1ea5e6
4 files changed
+474
-25
tests/kernels/test_fused_indexer_q_rope_quant.py
+1
-1
@@ -122,7 +122,7 @@ def _reference(
122
return q_fp8, weights_out
123
124
125
-@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
125
+@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023])
126
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
127
@pytest.mark.parametrize("use_fp4", [False, True])
128
@torch.inference_mode()
vllm/utils/import_utils.py
+5
@@ -469,3 +469,8 @@ def has_mori() -> bool:
469
def has_fbgemm_gpu() -> bool:
470
"""Whether the optional `fbgemm_gpu` package is available."""
471
return _has_module("fbgemm_gpu")
472
+
473
+
474
+def has_cutedsl() -> bool:
475
+ """Whether the optional `cutelass` package is available."""
476
+ return _has_module("cutlass")
vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q.py
+45
-24
@@ -1,8 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
import torch
5
6
from vllm.triton_utils import tl, triton
7
+from vllm.utils.import_utils import has_cutedsl
8
9
# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale.
10
MXFP4_BLOCK_SIZE = 32
@@ -342,30 +344,49 @@ def fused_indexer_q_rope_quant(
344
dtype=torch.uint8,
345
device=index_q.device,
346
)
345
- _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)](
346
- positions,
347
- index_q,
348
- index_q.stride(0),
349
- index_q.stride(1),
350
- index_q_cos_sin_cache,
351
- index_q_cos_sin_cache.stride(0),
352
- index_q_cos_sin_cache.shape[-1] // 2,
353
- index_q_packed,
354
- index_q_packed.stride(0),
355
- index_q_packed.stride(1),
356
- index_q_scale,
357
- index_q_scale.stride(0),
358
- index_q_scale.stride(1),
359
- index_q_head_dim,
360
- MXFP4_BLOCK_SIZE,
361
- index_weights,
362
- index_weights.stride(0),
363
- index_weights_softmax_scale,
364
- index_weights_head_scale,
365
- index_weights_out,
366
- index_weights_out.stride(0),
367
- num_warps=1, # TODO: Tune this
368
- )
347
+ if has_cutedsl():
348
+ # lazily import, otherwise some tests fail due to CUDA driver init failure.
349
+ from .fused_indexer_q_cutedsl import (
350
+ fused_indexer_q_rope_quant_mxfp4_cutedsl,
351
+ )
352
+
353
+ fused_indexer_q_rope_quant_mxfp4_cutedsl(
354
+ positions,
355
+ index_q,
356
+ index_q_cos_sin_cache,
357
+ index_weights,
358
+ index_weights_softmax_scale,
359
+ index_weights_head_scale,
360
+ index_q_packed,
361
+ index_q_scale,
362
+ index_weights_out,
363
+ )
364
+ else:
365
+ _fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)](
366
+ positions,
367
+ index_q,
368
+ index_q.stride(0),
369
+ index_q.stride(1),
370
+ index_q_cos_sin_cache,
371
+ index_q_cos_sin_cache.stride(0),
372
+ index_q_cos_sin_cache.shape[-1] // 2,
373
+ index_q_packed,
374
+ index_q_packed.stride(0),
375
+ index_q_packed.stride(1),
376
+ index_q_scale,
377
+ index_q_scale.stride(0),
378
+ index_q_scale.stride(1),
379
+ index_q_head_dim,
380
+ MXFP4_BLOCK_SIZE,
381
+ index_weights,
382
+ index_weights.stride(0),
383
+ index_weights_softmax_scale,
384
+ index_weights_head_scale,
385
+ index_weights_out,
386
+ index_weights_out.stride(0),
387
+ num_warps=1, # TODO: Tune this
388
+ )
389
+
390
# Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0
391
# bytes per (token, head) reinterpreted as one int32, then squeezed
392
# from (T, H, 1) to (T, H) to match DeepGEMM's expected q_sf rank
vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q_cutedsl.py
new
+423
@@ -0,0 +1,423 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# once we have more CuteDSL kernels in vLLM, we can refactor small helper functions
4
+# to a separate file
5
+from functools import cache
6
+
7
+import cutlass
8
+import cutlass.cute as cute
9
+import torch
10
+from cuda.bindings.driver import CUstream
11
+from cutlass import BFloat16, Float32, Int64, Uint8, Uint32, const_expr
12
+from cutlass._mlir.dialects import llvm
13
+from cutlass.cutlass_dsl import T, dsl_user_op
14
+from quack.compile_utils import make_fake_tensor
15
+
16
+from vllm.vllm_flash_attn.cute import utils as cute_utils
17
+
18
+# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale.
19
+MXFP4_BLOCK_SIZE = 32
20
+
21
+_TORCH_TO_CUTE = {
22
+ torch.bfloat16: BFloat16,
23
+ torch.float32: Float32,
24
+}
25
+
26
+
27
+def fused_indexer_q_rope_quant_mxfp4_cutedsl(
28
+ positions: torch.Tensor,
29
+ index_q: torch.Tensor,
30
+ index_q_cos_sin_cache: torch.Tensor,
31
+ index_weights: torch.Tensor,
32
+ index_weights_softmax_scale: float,
33
+ index_weights_head_scale: float,
34
+ index_q_packed: torch.Tensor,
35
+ index_q_scale: torch.Tensor,
36
+ index_weights_out: torch.Tensor,
37
+) -> None:
38
+ num_tokens, num_heads, head_dim = index_q.shape
39
+ rope_dim = index_q_cos_sin_cache.shape[-1]
40
+ rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype]
41
+
42
+ # compile all variants at first invocation
43
+ for coarsen in (1, 4):
44
+ IndexerQMxFp4Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen)
45
+
46
+ # heuristic
47
+ coarsen = 1 if num_tokens < 512 else 4
48
+ compiled = IndexerQMxFp4Kernel.compile(
49
+ head_dim, rope_dim, num_heads, rope_type, coarsen
50
+ )
51
+ scale = float(index_weights_softmax_scale * index_weights_head_scale)
52
+ compiled(
53
+ positions,
54
+ index_q,
55
+ index_q_cos_sin_cache,
56
+ index_weights,
57
+ index_q_packed,
58
+ index_q_scale,
59
+ index_weights_out,
60
+ scale,
61
+ )
62
+
63
+
64
+@dsl_user_op
65
+def _recast_val(x, dtype, *, loc=None, ip=None):
66
+ return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
67
+
68
+
69
+@dsl_user_op
70
+def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
71
+ out = llvm.inline_asm(
72
+ T.i32(),
73
+ [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
74
+ "cvt.rn.bf16x2.f32 $0, $2, $1;",
75
+ "=r,f,f",
76
+ has_side_effects=False,
77
+ is_align_stack=False,
78
+ )
79
+ return Uint32(out)
80
+
81
+
82
+@dsl_user_op
83
+def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]:
84
+ out = llvm.inline_asm(
85
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
86
+ [data.ir_value(loc=loc, ip=ip)],
87
+ "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n",
88
+ "=f,=f,r",
89
+ has_side_effects=False,
90
+ is_align_stack=False,
91
+ )
92
+ return (
93
+ Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
94
+ Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
95
+ )
96
+
97
+
98
+@dsl_user_op
99
+def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
100
+ out = llvm.inline_asm(
101
+ T.i32(),
102
+ [a.ir_value(loc=loc, ip=ip)],
103
+ "abs.bf16x2 $0, $1;",
104
+ "=r,r",
105
+ has_side_effects=False,
106
+ is_align_stack=False,
107
+ )
108
+ return Uint32(out)
109
+
110
+
111
+@dsl_user_op
112
+def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
113
+ out = llvm.inline_asm(
114
+ T.i32(),
115
+ [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
116
+ "max.bf16x2 $0, $1, $2;",
117
+ "=r,r,r",
118
+ has_side_effects=False,
119
+ is_align_stack=False,
120
+ )
121
+ return Uint32(out)
122
+
123
+
124
+@dsl_user_op
125
+def _fp32x8_to_fp4x8(
126
+ vals: cute.Tensor,
127
+ offset: cutlass.Constexpr[int],
128
+ *,
129
+ loc=None,
130
+ ip=None,
131
+) -> Uint32:
132
+ # Pack eight scaled FP32 values into four E2M1x2 bytes, returned as one b32.
133
+ assert vals.element_type is Float32
134
+ out = llvm.inline_asm(
135
+ T.i32(),
136
+ [vals[offset + i].ir_value(loc=loc, ip=ip) for i in range(8)],
137
+ "{\n\t"
138
+ ".reg .b8 x0, x1, x2, x3;\n\t"
139
+ "cvt.rn.satfinite.e2m1x2.f32 x0, $2, $1;\n\t"
140
+ "cvt.rn.satfinite.e2m1x2.f32 x1, $4, $3;\n\t"
141
+ "cvt.rn.satfinite.e2m1x2.f32 x2, $6, $5;\n\t"
142
+ "cvt.rn.satfinite.e2m1x2.f32 x3, $8, $7;\n\t"
143
+ "mov.b32 $0, {x0, x1, x2, x3};\n\t"
144
+ "}\n",
145
+ "=r,f,f,f,f,f,f,f,f",
146
+ has_side_effects=False,
147
+ is_align_stack=False,
148
+ )
149
+ return Uint32(out)
150
+
151
+
152
+class IndexerQMxFp4Kernel:
153
+ """Eight-thread subwarps process one ``(token, head)`` row."""
154
+
155
+ def __init__(
156
+ self,
157
+ head_dim: int = 128,
158
+ rope_dim: int = 64,
159
+ num_heads: int = 64,
160
+ cos_sin_dtype: type[cutlass.Numeric] = Float32,
161
+ coarsen: int = 4,
162
+ ):
163
+ self.head_dim = head_dim
164
+ self.rope_dim = rope_dim
165
+ self.nope_dim = head_dim - rope_dim
166
+ self.num_heads = num_heads
167
+ self.cos_sin_dtype = cos_sin_dtype
168
+
169
+ # process multiple heads at the same time to armotize RoPE load costs
170
+ assert num_heads % coarsen == 0
171
+ self.coarsen = coarsen
172
+
173
+ # later we will use 32B load = 16 BF16 elems
174
+ # thus, head_dim=128 requires 8 threads to handle.
175
+ # let's call subwarp = 8 threads.
176
+ self.subwarp_size = head_dim // 16
177
+ self.tb_size = 128
178
+ self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size
179
+
180
+ @cute.jit
181
+ def __call__(
182
+ self,
183
+ positions: cute.Tensor,
184
+ q: cute.Tensor,
185
+ cos_sin_cache: cute.Tensor,
186
+ weights: cute.Tensor,
187
+ q_fp4: cute.Tensor,
188
+ q_scale: cute.Tensor,
189
+ weights_out: cute.Tensor,
190
+ scale: Float32,
191
+ stream: CUstream,
192
+ ):
193
+ total_threads = q.shape[0] * self.threads_per_token
194
+ grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
195
+ self.kernel(
196
+ positions,
197
+ q,
198
+ cos_sin_cache,
199
+ weights,
200
+ q_fp4,
201
+ q_scale,
202
+ weights_out,
203
+ scale,
204
+ ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
205
+
206
+ @cute.kernel
207
+ def kernel(
208
+ self,
209
+ positions: cute.Tensor,
210
+ q: cute.Tensor,
211
+ cos_sin_cache: cute.Tensor,
212
+ weights: cute.Tensor,
213
+ q_fp4: cute.Tensor,
214
+ q_scale: cute.Tensor,
215
+ weights_out: cute.Tensor,
216
+ scale: Float32,
217
+ ):
218
+ block_id, _, _ = cute.arch.block_idx()
219
+ tid, _, _ = cute.arch.thread_idx()
220
+
221
+ num_token_heads = q.shape[0] * self.num_heads
222
+ global_tid = block_id * self.tb_size + tid
223
+
224
+ global_subwarp_id = global_tid // self.subwarp_size
225
+ sublane = tid % self.subwarp_size
226
+
227
+ token_id = global_subwarp_id // (self.num_heads // self.coarsen)
228
+ head_tile_id = global_subwarp_id % (self.num_heads // self.coarsen)
229
+ head_start = head_tile_id * self.coarsen
230
+
231
+ # NOTE: token_id may exceed bounds, hence we need to add load/store guards
232
+ # we can't do early exit because CuteDSL doesn't support it. and we also need
233
+ # all threads in a warp to be active since we utilize warp shuffle later.
234
+ # must_in_bounds is constexpr, True when 1 threadblock fit within 1 token
235
+ # position. the compiler will remove bounds check when that happens.
236
+ must_in_bounds = cutlass.const_expr(self.tb_size % self.threads_per_token == 0)
237
+ in_bounds = must_in_bounds or (token_id < q.shape[0])
238
+
239
+ cp_op = cute.nvgpu.CopyUniversalOp()
240
+
241
+ _layout = cute.make_layout((self.coarsen, 8), stride=(8, 1))
242
+ q_bf16x2 = cute.make_rmem_tensor(_layout, Uint32)
243
+
244
+ if in_bounds:
245
+ # we can't do cute.copy() on the whole 2D tile directly because
246
+ # cute.copy() wants the 1st mode to be covered by the copy atom,
247
+ # and other modes as for loop. there is no fast way to
248
+ # "transpose" the tensor view.
249
+ q_tile = cute.local_tile(
250
+ q[token_id, None, None],
251
+ tiler=(self.coarsen, 16),
252
+ coord=(head_tile_id, sublane),
253
+ )
254
+ cp_u32x8 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=256)
255
+ for i in cutlass.range_constexpr(self.coarsen):
256
+ src = cute.recast_tensor(q_tile[i, None], Uint32)
257
+ cute.copy(cp_u32x8, src, q_bf16x2[i, None])
258
+
259
+ # RoPE applies only to the trailing rope_dim values. We keep the rounded
260
+ # BF16 result in q_bits so the later amax and quantization see BF16.
261
+ # cos_sin_cache layout: [max_pos, rope_dim]
262
+ if in_bounds and sublane * 16 >= self.nope_dim:
263
+ cos_vals = cute.make_rmem_tensor((8,), Float32)
264
+ sin_vals = cute.make_rmem_tensor((8,), Float32)
265
+
266
+ pos = positions[token_id]
267
+
268
+ # select 8 elems from cos and sin
269
+ cos_id = sublane - self.nope_dim // 16
270
+ sin_id = cos_id + self.rope_dim // 16
271
+ cos_src = cute.local_tile(
272
+ cos_sin_cache[pos, None], tiler=(8,), coord=(cos_id,)
273
+ )
274
+ sin_src = cute.local_tile(
275
+ cos_sin_cache[pos, None], tiler=(8,), coord=(sin_id,)
276
+ )
277
+
278
+ cp_f32x8 = cute.make_copy_atom(cp_op, Float32, num_bits_per_copy=256)
279
+ cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
280
+
281
+ if const_expr(self.cos_sin_dtype is Float32):
282
+ cute.copy(cp_f32x8, cos_src, cos_vals)
283
+ cute.copy(cp_f32x8, sin_src, sin_vals)
284
+ else:
285
+ cos_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
286
+ sin_bf16x2 = cute.make_rmem_tensor((4,), Uint32)
287
+ cute.copy(cp_u32x4, cute.recast_tensor(cos_src, Uint32), cos_bf16x2)
288
+ cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2)
289
+
290
+ for i in cutlass.range_constexpr(4):
291
+ cos0, cos1 = _bf16x2_to_fp32(cos_bf16x2[i])
292
+ sin0, sin1 = _bf16x2_to_fp32(sin_bf16x2[i])
293
+ cos_vals[i * 2] = cos0
294
+ cos_vals[i * 2 + 1] = cos1
295
+ sin_vals[i * 2] = sin0
296
+ sin_vals[i * 2 + 1] = sin1
297
+
298
+ for i in cutlass.range_constexpr(self.coarsen):
299
+ for j in cutlass.range_constexpr(8):
300
+ q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j])
301
+ rot0 = q0 * cos_vals[j] - q1 * sin_vals[j]
302
+ rot1 = q0 * sin_vals[j] + q1 * cos_vals[j]
303
+ # convert back to BF16 to match numerics
304
+ q_bf16x2[i, j] = _fp32x2_to_bf16x2(rot0, rot1)
305
+
306
+ # layout: [coarsen, 8]
307
+ q_fp4_tile = cute.local_tile(
308
+ q_fp4[token_id, None, None],
309
+ tiler=(self.coarsen, 8),
310
+ coord=(head_tile_id, sublane),
311
+ )
312
+
313
+ for i in cutlass.range_constexpr(self.coarsen):
314
+ # compute amax in packed bf16x2 to save instructions
315
+ # Each thread holds 16 elems. Two adjacent threads form one 32-elem
316
+ # MXFP4 block, so a width-2 shuffle gives the block amax.
317
+ amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
318
+ for j in cutlass.range_constexpr(1, 8):
319
+ amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
320
+ amax_bf16x2 = cute_utils.warp_reduce(
321
+ amax_bf16x2,
322
+ _bf16x2_max,
323
+ width=MXFP4_BLOCK_SIZE // 16,
324
+ )
325
+ amax_pair = _bf16x2_to_fp32(amax_bf16x2)
326
+ amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
327
+
328
+ if in_bounds:
329
+ # compute block scale with bit manipulation
330
+ # UE8M0 stores ceil(log2(fp4_scale)) + 127. Adding the mantissa mask
331
+ # increments the exponent whenever fp4_scale is not exactly a power of 2
332
+ eps = cutlass.const_expr(float.fromhex("0x6p-126"))
333
+ fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0)
334
+ bits = _recast_val(fp4_scale, Uint32)
335
+ ue8m0 = cute_utils.shr_u32(
336
+ bits + Uint32(0x7FFFFF), Uint32(23)
337
+ ) & Uint32(0xFF)
338
+
339
+ # Only one of the two threads in an MXFP4 block writes the shared scale.
340
+ if tid % 2 == 0:
341
+ mx_block = sublane // 2
342
+ q_scale[token_id, head_start + i, mx_block] = Uint8(ue8m0)
343
+
344
+ # If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent
345
+ # -A + 127 = 254 - ue8m0.
346
+ inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23)
347
+ inv_fp4_scale = _recast_val(inv_scale_bits, Float32)
348
+
349
+ vals = cute.make_rmem_tensor(16, Float32)
350
+ for j in cutlass.range_constexpr(8):
351
+ q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j])
352
+ vals[j * 2] = q0 * inv_fp4_scale
353
+ vals[j * 2 + 1] = q1 * inv_fp4_scale
354
+
355
+ # pack to FP4
356
+ packed = cute.make_rmem_tensor((2,), Uint32)
357
+ packed[0] = _fp32x8_to_fp4x8(vals, 0)
358
+ packed[1] = _fp32x8_to_fp4x8(vals, 8)
359
+
360
+ dst = q_fp4_tile[i, None]
361
+ cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
362
+ cute.copy(cp_u32x2, packed, cute.recast_tensor(dst, Uint32))
363
+
364
+ # Weight scaling is independent of the Q subwarp work. The first
365
+ # num_tokens * num_heads logical threads cover one weight each.
366
+ if global_tid < num_token_heads:
367
+ weight_token_id = global_tid // self.num_heads
368
+ weight_head_id = global_tid % self.num_heads
369
+ weights_out[weight_token_id, weight_head_id] = (
370
+ weights[weight_token_id, weight_head_id].to(Float32) * scale
371
+ )
372
+
373
+ @cache
374
+ @staticmethod
375
+ def compile(
376
+ head_dim: int = 128,
377
+ rope_dim: int = 64,
378
+ num_heads: int = 64,
379
+ cos_sin_dtype: type[cutlass.Numeric] = Float32,
380
+ coarsen: int = 4,
381
+ ):
382
+ num_tokens = cute.sym_int()
383
+ max_pos = cute.sym_int()
384
+
385
+ q = make_fake_tensor(
386
+ BFloat16, (num_tokens, num_heads, head_dim), divisibility=16
387
+ )
388
+ positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1)
389
+ cos_sin_cache = make_fake_tensor(
390
+ cos_sin_dtype,
391
+ (max_pos, rope_dim),
392
+ divisibility=8,
393
+ )
394
+ weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8)
395
+ q_fp4 = make_fake_tensor(
396
+ Uint8,
397
+ (num_tokens, num_heads, head_dim // 2),
398
+ divisibility=16,
399
+ )
400
+ q_scale = make_fake_tensor(
401
+ Uint8,
402
+ (num_tokens, num_heads, head_dim // MXFP4_BLOCK_SIZE),
403
+ divisibility=4,
404
+ )
405
+ weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4)
406
+
407
+ kernel = IndexerQMxFp4Kernel(
408
+ head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen
409
+ )
410
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
411
+ return cute.compile(
412
+ kernel,
413
+ positions,
414
+ q,
415
+ cos_sin_cache,
416
+ weights,
417
+ q_fp4,
418
+ q_scale,
419
+ weights_out,
420
+ Float32(0.0),
421
+ stream,
422
+ options="--enable-tvm-ffi",
423
+ )