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 3, 09:15 UTC
add cutedsl dsv4 indexer fp8 kernel (#42899)
Signed-off-by: george <george@inferact.ai> Co-authored-by: george <george@inferact.ai>
gnovack committed
May 18, 2026 at 21:17 UTC
3ca8db2ef88ec5a6686e62ee3ac899afae85c7af
4 files changed
+411
-60
tests/kernels/test_fused_indexer_q_rope_quant.py
+30
-3
@@ -13,6 +13,9 @@ the DeepseekV4 indexer in model_tracking:
13
Expects bit-exact equality on both q_fp8 and weights_out.
14
"""
15
16
+import contextlib
17
+from unittest import mock
18
+
19
import pytest
20
import torch
21
@@ -20,6 +23,7 @@ from vllm import _custom_ops as ops
23
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
24
per_token_group_quant_fp8,
25
)
26
+from vllm.utils.import_utils import has_cutedsl
27
from vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q import (
28
fused_indexer_q_rope_quant,
29
)
@@ -125,8 +129,14 @@ def _reference(
129
@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023])
130
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
131
@pytest.mark.parametrize("use_fp4", [False, True])
132
+@pytest.mark.parametrize("use_cutedsl", [False, True])
133
@torch.inference_mode()
129
-def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
134
+def test_fused_indexer_q_rope_quant_matches_unfused(
135
+ num_tokens, cache_dtype, use_fp4, use_cutedsl
136
+):
137
+ if use_cutedsl and not has_cutedsl():
138
+ pytest.skip("cutedsl (cutlass) not installed")
139
+
140
device = "cuda"
141
torch.manual_seed(0)
142
@@ -142,9 +152,26 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use
152
q_quant_ref, weights_ref = _reference(
153
positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
154
)
145
- q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
146
- positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
155
+ # use_cutedsl=False: force the triton path even when cutedsl is installed
156
+ # by patching the dispatcher's has_cutedsl() binding to return False.
157
+ cutedsl_patch = (
158
+ mock.patch(
159
+ "vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q.has_cutedsl",
160
+ return_value=False,
161
+ )
162
+ if not use_cutedsl
163
+ else contextlib.nullcontext()
164
)
165
+ with cutedsl_patch:
166
+ q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
167
+ positions,
168
+ q.clone(),
169
+ cos_sin_cache,
170
+ weights,
171
+ softmax_scale,
172
+ head_scale,
173
+ use_fp4,
174
+ )
175
176
if use_fp4:
177
q_quant_ref, q_scale_ref = q_quant_ref
vllm/v1/attention/ops/deepseek_v4_ops/cutedsl_utils.py
+33
@@ -117,6 +117,39 @@ def _fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
117
return cute.TensorSSA(vec, 2, Uint32)
118
119
120
+@dsl_user_op
121
+def _fp32x4_to_fp8x4(
122
+ a0: Float32,
123
+ a1: Float32,
124
+ a2: Float32,
125
+ a3: Float32,
126
+ *,
127
+ loc=None,
128
+ ip=None,
129
+) -> Uint32:
130
+ # Pack four FP32 values into one b32 of four e4m3 bytes, byte order
131
+ # {a0, a1, a2, a3} from low to high address.
132
+ out = llvm.inline_asm(
133
+ T.i32(),
134
+ [
135
+ a0.ir_value(loc=loc, ip=ip),
136
+ a1.ir_value(loc=loc, ip=ip),
137
+ a2.ir_value(loc=loc, ip=ip),
138
+ a3.ir_value(loc=loc, ip=ip),
139
+ ],
140
+ "{\n\t"
141
+ ".reg .b16 t0, t1;\n\t"
142
+ "cvt.rn.satfinite.e4m3x2.f32 t0, $2, $1;\n\t"
143
+ "cvt.rn.satfinite.e4m3x2.f32 t1, $4, $3;\n\t"
144
+ "mov.b32 $0, {t0, t1};\n\t"
145
+ "}\n",
146
+ "=r,f,f,f,f",
147
+ has_side_effects=False,
148
+ is_align_stack=False,
149
+ )
150
+ return Uint32(out)
151
+
152
+
153
@dsl_user_op
154
def _fp32x8_to_fp4x8(
155
vals: cute.Tensor,
vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q.py
+37
-20
@@ -398,24 +398,41 @@ def fused_indexer_q_rope_quant(
398
), index_weights_out
399
400
index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
401
- _fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)](
402
- positions,
403
- index_q,
404
- index_q.stride(0),
405
- index_q.stride(1),
406
- index_q_cos_sin_cache,
407
- index_q_cos_sin_cache.stride(0),
408
- index_q_cos_sin_cache.shape[-1] // 2,
409
- index_q_fp8,
410
- index_q_fp8.stride(0),
411
- index_q_fp8.stride(1),
412
- index_q_head_dim,
413
- index_weights,
414
- index_weights.stride(0),
415
- index_weights_softmax_scale,
416
- index_weights_head_scale,
417
- index_weights_out,
418
- index_weights_out.stride(0),
419
- num_warps=1, # TODO: Tune this
420
- )
401
+ if has_cutedsl():
402
+ # lazily import, otherwise some tests fail due to CUDA driver init failure.
403
+ from .fused_indexer_q_cutedsl import (
404
+ fused_indexer_q_rope_quant_fp8_cutedsl,
405
+ )
406
+
407
+ fused_indexer_q_rope_quant_fp8_cutedsl(
408
+ positions,
409
+ index_q,
410
+ index_q_cos_sin_cache,
411
+ index_weights,
412
+ index_weights_softmax_scale,
413
+ index_weights_head_scale,
414
+ index_q_fp8,
415
+ index_weights_out,
416
+ )
417
+ else:
418
+ _fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)](
419
+ positions,
420
+ index_q,
421
+ index_q.stride(0),
422
+ index_q.stride(1),
423
+ index_q_cos_sin_cache,
424
+ index_q_cos_sin_cache.stride(0),
425
+ index_q_cos_sin_cache.shape[-1] // 2,
426
+ index_q_fp8,
427
+ index_q_fp8.stride(0),
428
+ index_q_fp8.stride(1),
429
+ index_q_head_dim,
430
+ index_weights,
431
+ index_weights.stride(0),
432
+ index_weights_softmax_scale,
433
+ index_weights_head_scale,
434
+ index_weights_out,
435
+ index_weights_out.stride(0),
436
+ num_warps=1, # TODO: Tune this
437
+ )
438
return index_q_fp8, index_weights_out
vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q_cutedsl.py
+311
-37
@@ -14,6 +14,7 @@ from vllm.v1.attention.ops.deepseek_v4_ops.cutedsl_utils import (
14
_bf16x2_max,
15
_bf16x2_to_fp32,
16
_fp32x2_to_bf16x2,
17
+ _fp32x4_to_fp8x4,
18
_fp32x8_to_fp4x8,
19
_recast_val,
20
)
@@ -65,8 +66,48 @@ def fused_indexer_q_rope_quant_mxfp4_cutedsl(
66
)
67
68
68
-class IndexerQMxFp4Kernel:
69
- """Eight-thread subwarps process one ``(token, head)`` row."""
69
+def fused_indexer_q_rope_quant_fp8_cutedsl(
70
+ positions: torch.Tensor,
71
+ index_q: torch.Tensor,
72
+ index_q_cos_sin_cache: torch.Tensor,
73
+ index_weights: torch.Tensor,
74
+ index_weights_softmax_scale: float,
75
+ index_weights_head_scale: float,
76
+ index_q_fp8: torch.Tensor,
77
+ index_weights_out: torch.Tensor,
78
+) -> None:
79
+ num_tokens, num_heads, head_dim = index_q.shape
80
+ rope_dim = index_q_cos_sin_cache.shape[-1]
81
+ rope_type = _TORCH_TO_CUTE[index_q_cos_sin_cache.dtype]
82
+
83
+ for coarsen in (1, 4):
84
+ IndexerQFp8Kernel.compile(head_dim, rope_dim, num_heads, rope_type, coarsen)
85
+
86
+ coarsen = 1 if num_tokens < 512 else 4
87
+ compiled = IndexerQFp8Kernel.compile(
88
+ head_dim, rope_dim, num_heads, rope_type, coarsen
89
+ )
90
+ scale = float(index_weights_softmax_scale * index_weights_head_scale)
91
+ # The cute kernel treats the FP8 buffer as raw bytes (Uint8).
92
+ compiled(
93
+ positions,
94
+ index_q,
95
+ index_q_cos_sin_cache,
96
+ index_weights,
97
+ index_q_fp8.view(torch.uint8),
98
+ index_weights_out,
99
+ scale,
100
+ )
101
+
102
+
103
+class IndexerQRopeQuantKernel:
104
+ """Shared infrastructure for indexer-Q RoPE+quant fused kernels.
105
+
106
+ Subclasses implement ``kernel`` for a particular Q quantization scheme
107
+ (MXFP4, FP8 e4m3, …). The base class owns the launch geometry and the
108
+ common preamble: thread/token addressing, the BF16 Q load, and the
109
+ interleaved-RoPE pass over the trailing ``rope_dim`` lanes.
110
+ """
111
112
def __init__(
113
self,
@@ -94,47 +135,27 @@ class IndexerQMxFp4Kernel:
135
self.threads_per_token = (self.num_heads // self.coarsen) * self.subwarp_size
136
137
@cute.jit
97
- def __call__(
138
+ def _load_q_and_rope(
139
self,
140
positions: cute.Tensor,
141
q: cute.Tensor,
142
cos_sin_cache: cute.Tensor,
102
- weights: cute.Tensor,
103
- q_fp4: cute.Tensor,
104
- q_scale: cute.Tensor,
105
- weights_out: cute.Tensor,
106
- scale: Float32,
107
- stream: CUstream,
108
- ):
109
- total_threads = q.shape[0] * self.threads_per_token
110
- grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
111
- self.kernel(
112
- positions,
113
- q,
114
- cos_sin_cache,
115
- weights,
116
- q_fp4,
117
- q_scale,
118
- weights_out,
119
- scale,
120
- ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
121
-
122
- @cute.kernel
123
- def kernel(
124
- self,
125
- positions: cute.Tensor,
126
- q: cute.Tensor,
127
- cos_sin_cache: cute.Tensor,
128
- weights: cute.Tensor,
129
- q_fp4: cute.Tensor,
130
- q_scale: cute.Tensor,
131
- weights_out: cute.Tensor,
132
- scale: Float32,
143
):
144
+ """Compute thread indices, load Q (BF16), and apply interleaved RoPE.
145
+
146
+ Returns a tuple
147
+ (q_bf16x2, tid, global_tid, sublane, token_id, head_tile_id,
148
+ head_start, in_bounds, num_token_heads)
149
+ where ``q_bf16x2`` is a (coarsen, 8) rmem tile of Uint32 packed
150
+ bf16x2 pairs covering the 16 BF16 lanes owned by this thread for
151
+ each of ``coarsen`` heads. RoPE is applied in place to the
152
+ trailing ``rope_dim`` lanes; the leading nope lanes pass through.
153
+ """
154
block_id, _, _ = cute.arch.block_idx()
155
tid, _, _ = cute.arch.thread_idx()
156
137
- num_token_heads = q.shape[0] * self.num_heads
157
+ num_tokens = q.shape[0]
158
+ num_token_heads = num_tokens * self.num_heads
159
global_tid = block_id * self.tb_size + tid
160
161
global_subwarp_id = global_tid // self.subwarp_size
@@ -150,7 +171,7 @@ class IndexerQMxFp4Kernel:
171
# must_in_bounds is constexpr, True when 1 threadblock fit within 1 token
172
# position. the compiler will remove bounds check when that happens.
173
must_in_bounds = cutlass.const_expr(self.tb_size % self.threads_per_token == 0)
153
- in_bounds = must_in_bounds or (token_id < q.shape[0])
174
+ in_bounds = must_in_bounds or (token_id < num_tokens)
175
176
cp_op = cute.nvgpu.CopyUniversalOp()
177
@@ -219,9 +240,77 @@ class IndexerQMxFp4Kernel:
240
# convert back to BF16 to match numerics
241
q_bf16x2[i, j] = _fp32x2_to_bf16x2(rot0, rot1)
242
243
+ return (
244
+ q_bf16x2,
245
+ tid,
246
+ global_tid,
247
+ sublane,
248
+ token_id,
249
+ head_tile_id,
250
+ head_start,
251
+ in_bounds,
252
+ num_token_heads,
253
+ )
254
+
255
+
256
+class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
257
+ """Eight-thread subwarps process one ``(token, head)`` row."""
258
+
259
+ @cute.jit
260
+ def __call__(
261
+ self,
262
+ positions: cute.Tensor,
263
+ q: cute.Tensor,
264
+ cos_sin_cache: cute.Tensor,
265
+ weights: cute.Tensor,
266
+ q_quant: cute.Tensor,
267
+ q_scale: cute.Tensor,
268
+ weights_out: cute.Tensor,
269
+ scale: Float32,
270
+ stream: CUstream,
271
+ ):
272
+ total_threads = q.shape[0] * self.threads_per_token
273
+ grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
274
+ self.kernel(
275
+ positions,
276
+ q,
277
+ cos_sin_cache,
278
+ weights,
279
+ q_quant,
280
+ q_scale,
281
+ weights_out,
282
+ scale,
283
+ ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
284
+
285
+ @cute.kernel
286
+ def kernel(
287
+ self,
288
+ positions: cute.Tensor,
289
+ q: cute.Tensor,
290
+ cos_sin_cache: cute.Tensor,
291
+ weights: cute.Tensor,
292
+ q_quant: cute.Tensor,
293
+ q_scale: cute.Tensor,
294
+ weights_out: cute.Tensor,
295
+ scale: Float32,
296
+ ):
297
+ (
298
+ q_bf16x2,
299
+ tid,
300
+ global_tid,
301
+ sublane,
302
+ token_id,
303
+ head_tile_id,
304
+ head_start,
305
+ in_bounds,
306
+ num_token_heads,
307
+ ) = self._load_q_and_rope(positions, q, cos_sin_cache)
308
+
309
+ cp_op = cute.nvgpu.CopyUniversalOp()
310
+
311
# layout: [coarsen, 8]
312
q_fp4_tile = cute.local_tile(
224
- q_fp4[token_id, None, None],
313
+ q_quant[token_id, None, None],
314
tiler=(self.coarsen, 8),
315
coord=(head_tile_id, sublane),
316
)
@@ -337,3 +426,188 @@ class IndexerQMxFp4Kernel:
426
stream,
427
options="--enable-tvm-ffi",
428
)
429
+
430
+
431
+class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
432
+ """Eight-thread subwarps process one ``(token, head)`` row and emit
433
+ float8 e4m3fn with a single per-(token, head) scalar scale folded
434
+ into the per-token weight (mirrors ``_fused_indexer_q_rope_quant_kernel``).
435
+ """
436
+
437
+ def __init__(
438
+ self,
439
+ head_dim: int = 128,
440
+ rope_dim: int = 64,
441
+ num_heads: int = 64,
442
+ cos_sin_dtype: type[cutlass.Numeric] = Float32,
443
+ coarsen: int = 4,
444
+ ):
445
+ super().__init__(head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen)
446
+ # Each subwarp owns `coarsen` heads; we use the first `coarsen`
447
+ # threads of the subwarp to write the per-head weights using the
448
+ # fp8 scale computed in the matching loop iteration.
449
+ assert self.coarsen <= self.subwarp_size, (
450
+ f"FP8 kernel requires coarsen ({self.coarsen}) <= "
451
+ f"subwarp_size ({self.subwarp_size}) for the weight-fold step"
452
+ )
453
+
454
+ @cute.jit
455
+ def __call__(
456
+ self,
457
+ positions: cute.Tensor,
458
+ q: cute.Tensor,
459
+ cos_sin_cache: cute.Tensor,
460
+ weights: cute.Tensor,
461
+ q_fp8: cute.Tensor,
462
+ weights_out: cute.Tensor,
463
+ scale: Float32,
464
+ stream: CUstream,
465
+ ):
466
+ total_threads = q.shape[0] * self.threads_per_token
467
+ grid = (cute.ceil_div(total_threads, self.tb_size), 1, 1)
468
+ self.kernel(
469
+ positions,
470
+ q,
471
+ cos_sin_cache,
472
+ weights,
473
+ q_fp8,
474
+ weights_out,
475
+ scale,
476
+ ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
477
+
478
+ @cute.kernel
479
+ def kernel(
480
+ self,
481
+ positions: cute.Tensor,
482
+ q: cute.Tensor,
483
+ cos_sin_cache: cute.Tensor,
484
+ weights: cute.Tensor,
485
+ q_fp8: cute.Tensor,
486
+ weights_out: cute.Tensor,
487
+ scale: Float32,
488
+ ):
489
+ (
490
+ q_bf16x2,
491
+ _tid,
492
+ _global_tid,
493
+ sublane,
494
+ token_id,
495
+ head_tile_id,
496
+ head_start,
497
+ in_bounds,
498
+ _num_token_heads,
499
+ ) = self._load_q_and_rope(positions, q, cos_sin_cache)
500
+
501
+ cp_op = cute.nvgpu.CopyUniversalOp()
502
+
503
+ # layout: [coarsen, 16] bytes (one e4m3fn per element).
504
+ q_fp8_tile = cute.local_tile(
505
+ q_fp8[token_id, None, None],
506
+ tiler=(self.coarsen, 16),
507
+ coord=(head_tile_id, sublane),
508
+ )
509
+
510
+ for i in cutlass.range_constexpr(self.coarsen):
511
+ # Reduce amax across the full head_dim: each thread already holds
512
+ # the max over its 16 lanes; a width=subwarp_size warp shuffle
513
+ # spreads the head-wide max to every lane in the subwarp.
514
+ amax_bf16x2 = _bf16x2_abs(q_bf16x2[i, 0])
515
+ for j in cutlass.range_constexpr(1, 8):
516
+ amax_bf16x2 = _bf16x2_max(amax_bf16x2, _bf16x2_abs(q_bf16x2[i, j]))
517
+ amax_bf16x2 = cute_utils.warp_reduce(
518
+ amax_bf16x2,
519
+ _bf16x2_max,
520
+ width=self.subwarp_size,
521
+ )
522
+ amax_pair = _bf16x2_to_fp32(amax_bf16x2)
523
+ amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
524
+
525
+ # scale = max(amax, eps) / fp8_max, then rounded UP to the next
526
+ # power of two. Adding the mantissa mask before shifting out the
527
+ # mantissa bumps the exponent whenever s isn't a pure pow2.
528
+ fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0)
529
+ bits = _recast_val(fp32_scale, Uint32)
530
+ scale_exp = cute_utils.shr_u32(
531
+ bits + Uint32(0x7FFFFF), Uint32(23)
532
+ ) & Uint32(0xFF)
533
+
534
+ # rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23
535
+ fp8_scale_bits = scale_exp << Uint32(23)
536
+ fp8_scale = _recast_val(fp8_scale_bits, Float32)
537
+ # inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23
538
+ inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23)
539
+ inv_fp8_scale = _recast_val(inv_scale_bits, Float32)
540
+
541
+ # Weight fold: weights_out = weights * q_scale * scale_combined.
542
+ # All threads in the subwarp share the same fp8_scale after the
543
+ # warp_reduce above, so we let thread `sublane == i` write the
544
+ # weight for head `head_start + i`.
545
+ if in_bounds and sublane == i:
546
+ head_id = head_start + i
547
+ weights_out[token_id, head_id] = (
548
+ weights[token_id, head_id].to(Float32) * scale * fp8_scale
549
+ )
550
+
551
+ if in_bounds:
552
+ # 16 BF16 → 16 e4m3 bytes per thread, packed into 4 b32s
553
+ # (one cp.async-shaped 128-bit store per row).
554
+ packed = cute.make_rmem_tensor((4,), Uint32)
555
+ for j in cutlass.range_constexpr(4):
556
+ q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j * 2])
557
+ q2, q3 = _bf16x2_to_fp32(q_bf16x2[i, j * 2 + 1])
558
+ packed[j] = _fp32x4_to_fp8x4(
559
+ q0 * inv_fp8_scale,
560
+ q1 * inv_fp8_scale,
561
+ q2 * inv_fp8_scale,
562
+ q3 * inv_fp8_scale,
563
+ )
564
+
565
+ dst = q_fp8_tile[i, None]
566
+ cp_u32x4 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
567
+ cute.copy(cp_u32x4, packed, cute.recast_tensor(dst, Uint32))
568
+
569
+ @cache
570
+ @staticmethod
571
+ def compile(
572
+ head_dim: int = 128,
573
+ rope_dim: int = 64,
574
+ num_heads: int = 64,
575
+ cos_sin_dtype: type[cutlass.Numeric] = Float32,
576
+ coarsen: int = 4,
577
+ ):
578
+ num_tokens = cute.sym_int()
579
+ max_pos = cute.sym_int()
580
+
581
+ q = make_fake_tensor(
582
+ BFloat16, (num_tokens, num_heads, head_dim), divisibility=16
583
+ )
584
+ positions = make_fake_tensor(Int64, (num_tokens,), divisibility=1)
585
+ cos_sin_cache = make_fake_tensor(
586
+ cos_sin_dtype,
587
+ (max_pos, rope_dim),
588
+ divisibility=8,
589
+ )
590
+ weights = make_fake_tensor(BFloat16, (num_tokens, num_heads), divisibility=8)
591
+ q_fp8 = make_fake_tensor(
592
+ Uint8,
593
+ (num_tokens, num_heads, head_dim),
594
+ divisibility=16,
595
+ )
596
+ weights_out = make_fake_tensor(Float32, (num_tokens, num_heads), divisibility=4)
597
+
598
+ kernel = IndexerQFp8Kernel(
599
+ head_dim, rope_dim, num_heads, cos_sin_dtype, coarsen
600
+ )
601
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
602
+ return cute.compile(
603
+ kernel,
604
+ positions,
605
+ q,
606
+ cos_sin_cache,
607
+ weights,
608
+ q_fp8,
609
+ weights_out,
610
+ Float32(0.0),
611
+ stream,
612
+ options="--enable-tvm-ffi",
613
+ )