[Perf][DSv4.1] Fuse small-batch WO-A with inverse RoPE and MXFP8 quant on SM100/SM103 (#58634)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Juntian Liu committed
Sep 26, 2026 at 21:16 UTC
a9eafde59cbd55182dc2265cc398b7186f0a0eaa
6 files changed
+759
-55
tests/kernels/test_fused_inv_rope_fp8_quant.py
+48
@@ -21,7 +21,12 @@ Usage:
21
import pytest
22
import torch
23
24
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
25
+ deepgemm_post_process_fp8_weight_block,
26
+)
27
from vllm.models.deepseek_v4.common.ops import fused_inv_rope_fp8_quant
28
+from vllm.platforms import current_platform
29
+from vllm.utils.deep_gemm import fp8_einsum
30
31
# -- Default dimensions matching DeepSeek V3/V4 --------------------------
32
HEAD_DIM = 512
@@ -33,6 +38,49 @@ FP8_DTYPE = torch.float8_e4m3fn
38
EPS = 1e-10
39
40
41
+@pytest.mark.skipif(
42
+ not current_platform.is_device_capability_family(100), reason="Requires SM100/SM103"
43
+)
44
+@pytest.mark.parametrize("groups", [2, 4])
45
+@pytest.mark.parametrize("tokens", [1, 6, 18, 32])
46
+def test_dsv41_fused_wo_a_matches_deepgemm(groups, tokens):
47
+ """Fused WO-A matches inverse RoPE + fp8_einsum + MXFP8 quantize."""
48
+ from flashinfer import mxfp8_quantize
49
+
50
+ from vllm.models.deepseek_v41.nvidia.ops.fused_wo_a import _FUSED_WO_A_KERNEL
51
+
52
+ torch.manual_seed(0)
53
+ n = groups * 1024
54
+ x = torch.randn(tokens, 64, 512, device="cuda", dtype=torch.bfloat16)[:, : n // 128]
55
+ positions = torch.randint(0, 128, (tokens,), device="cuda")
56
+ rope = make_cos_sin_cache(128, device="cuda")
57
+ wq = torch.randn(n, 4096, device="cuda").to(FP8_DTYPE)
58
+ scale = torch.exp2(torch.randint(-8, 0, (n, 128), device="cuda").float())
59
+ wq, ws = deepgemm_post_process_fp8_weight_block(
60
+ wq, scale, (1, 32), False, True, groups
61
+ )
62
+ q, sf = _FUSED_WO_A_KERNEL(
63
+ x=x, positions=positions, rope=rope, weight=wq, weight_scale=ws
64
+ )
65
+
66
+ a, sa = fused_inv_rope_fp8_quant(
67
+ x, positions, rope, groups, 8, quant_group_size=32, tma_aligned_scales=True
68
+ )
69
+ y = torch.empty(tokens, groups, 1024, device="cuda", dtype=torch.bfloat16)
70
+ fp8_einsum("bhr,hdr->bhd", (a, sa), (wq, ws), y, recipe=(1, 1, 32))
71
+ ref, ref_sf = mxfp8_quantize(y.flatten(1), backend="cute-dsl")
72
+
73
+ def rows(scales): # F8_128x4 bytes -> [128 rows, n / 32 blocks]
74
+ return scales.view(n // 128, 32, 4, 4).permute(2, 1, 0, 3).reshape(128, -1)
75
+
76
+ actual, expected = (
77
+ d.float().view(tokens, -1, 32) * torch.exp2(rows(s)[:tokens, :, None] - 127.0)
78
+ for d, s in ((q, sf), (ref, ref_sf))
79
+ )
80
+ assert (actual - expected).norm() / expected.norm() < 3e-3
81
+ assert not rows(sf)[tokens:].any(), "padded scale rows must be zero"
82
+
83
+
84
# =========================================================================
85
# Helpers
86
# =========================================================================
vllm/models/deepseek_v41/attention.py
+2
-2
@@ -697,7 +697,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
697
else:
698
gemm_rs.warn_incompatible_projection()
699
700
- def _wo_b_proj(self, z: torch.Tensor) -> torch.Tensor:
700
+ def _wo_b_proj(self, z: torch.Tensor | QuantizedActivation) -> torch.Tensor:
701
"""Apply ``wo_b``; with GEMM-RS bound, also reduce-scatter the result.
702
703
Every ``_o_proj`` implementation projects through this so the decoder
@@ -707,7 +707,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
707
"""
708
if self.gemm_rs is None:
709
return self.wo_b(z)
710
- if self.gemm_rs.should_run(z):
710
+ if isinstance(z, torch.Tensor) and self.gemm_rs.should_run(z):
711
return self.gemm_rs.apply(z, self.wo_b)
712
# Small batches stay on the unfused path, which is faster there.
713
return sp_reduce_scatter(self.wo_b(z))
vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
+8
-34
@@ -9,15 +9,16 @@ import torch
9
from vllm.config import VllmConfig
10
from vllm.config.cache import CacheDType
11
from vllm.forward_context import get_forward_context
12
-from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
13
- compute_fp8_einsum_recipe,
14
- deep_gemm_fp8_o_proj,
15
-)
12
+from vllm.models.deepseek_v4.nvidia.ops.o_proj import compute_fp8_einsum_recipe
13
from vllm.models.deepseek_v41.attention import DeepseekV4Attention
14
from vllm.models.deepseek_v41.common.ops import (
15
build_flashinfer_mixed_sparse_indices,
16
compute_global_topk_indices_and_lens,
17
)
18
+from vllm.models.deepseek_v41.nvidia.ops.o_proj import (
19
+ dsv41_o_proj,
20
+ register_dsv41_o_proj_warmup,
21
+)
22
from vllm.models.deepseek_v41.sparse_mla import (
23
DeepseekV4FlashMLAMetadata,
24
DeepseekV4SparseMLABackend,
@@ -242,27 +243,14 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
243
return _pad_to_supported_q_heads(num_heads)
244
245
def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
245
- o = attn_out[:, : self.n_local_heads, :]
246
- return deep_gemm_fp8_o_proj(
247
- o,
248
- positions,
249
- self.rotary_emb.cos_sin_cache,
250
- self.wo_a,
251
- self._wo_b_proj,
252
- n_groups=self.n_local_groups,
253
- heads_per_group=self.n_local_heads // self.n_local_groups,
254
- nope_dim=self.nope_head_dim,
255
- rope_dim=self.rope_head_dim,
256
- o_lora_rank=self.o_lora_rank,
257
- einsum_recipe=self._einsum_recipe,
258
- tma_aligned_scales=self._tma_aligned_scales,
259
- )
246
+ return dsv41_o_proj(self, attn_out, positions)
247
248
def __init__(self, *args, **kwargs) -> None:
249
super().__init__(*args, **kwargs)
250
self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
251
self._o_proj_block_size
252
)
253
+ register_dsv41_o_proj_warmup(self)
254
# Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
255
# per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
256
if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
@@ -623,21 +611,7 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
611
return _pad_to_supported_q_heads(num_heads)
612
613
def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
626
- o = attn_out[:, : self.n_local_heads, :]
627
- return deep_gemm_fp8_o_proj(
628
- o,
629
- positions,
630
- self.rotary_emb.cos_sin_cache,
631
- self.wo_a,
632
- self._wo_b_proj,
633
- n_groups=self.n_local_groups,
634
- heads_per_group=self.n_local_heads // self.n_local_groups,
635
- nope_dim=self.nope_head_dim,
636
- rope_dim=self.rope_head_dim,
637
- o_lora_rank=self.o_lora_rank,
638
- einsum_recipe=self._einsum_recipe,
639
- tma_aligned_scales=self._tma_aligned_scales,
640
- )
614
+ return dsv41_o_proj(self, attn_out, positions)
615
616
def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
617
super().__init__(vllm_config, *args, **kwargs)
vllm/models/deepseek_v41/nvidia/flashmla.py
+7
-19
@@ -6,16 +6,17 @@ from typing import TYPE_CHECKING, ClassVar, cast
6
import torch
7
8
from vllm.forward_context import get_forward_context
9
-from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
10
- compute_fp8_einsum_recipe,
11
- deep_gemm_fp8_o_proj,
12
-)
9
+from vllm.models.deepseek_v4.nvidia.ops.o_proj import compute_fp8_einsum_recipe
10
from vllm.models.deepseek_v41.attention import DeepseekV4Attention
11
from vllm.models.deepseek_v41.common.ops import (
12
combine_topk_swa_indices,
13
compute_global_topk_indices_and_lens,
14
dequantize_and_gather_k_cache,
15
)
16
+from vllm.models.deepseek_v41.nvidia.ops.o_proj import (
17
+ dsv41_o_proj,
18
+ register_dsv41_o_proj_warmup,
19
+)
20
from vllm.models.deepseek_v41.sparse_mla import (
21
DeepseekV4FlashMLABackend,
22
DeepseekV4FlashMLAMetadata,
@@ -57,23 +58,10 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
58
self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe(
59
self._o_proj_block_size
60
)
61
+ register_dsv41_o_proj_warmup(self)
62
63
def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
62
- o = attn_out[:, : self.n_local_heads, :]
63
- return deep_gemm_fp8_o_proj(
64
- o,
65
- positions,
66
- self.rotary_emb.cos_sin_cache,
67
- self.wo_a,
68
- self._wo_b_proj,
69
- n_groups=self.n_local_groups,
70
- heads_per_group=self.n_local_heads // self.n_local_groups,
71
- nope_dim=self.nope_head_dim,
72
- rope_dim=self.rope_head_dim,
73
- o_lora_rank=self.o_lora_rank,
74
- einsum_recipe=self._einsum_recipe,
75
- tma_aligned_scales=self._tma_aligned_scales,
76
- )
64
+ return dsv41_o_proj(self, attn_out, positions)
65
66
@classmethod
67
def get_padded_num_q_heads(cls, num_heads: int) -> int:
vllm/models/deepseek_v41/nvidia/ops/fused_wo_a.py
new
+599
@@ -0,0 +1,599 @@
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()
vllm/models/deepseek_v41/nvidia/ops/o_proj.py
new
+95
@@ -0,0 +1,95 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""DSV4.1 output projection with a small-batch SM100/SM103 fusion."""
4
+
5
+import torch
6
+from torch import nn
7
+
8
+from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
9
+ FlashInferCutedslMxfp8LinearKernel,
10
+)
11
+from vllm.model_executor.layers.fusion.quant_activation import (
12
+ QuantizedActivation,
13
+ get_input_quant_key,
14
+)
15
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
16
+from vllm.models.deepseek_v4.nvidia.ops.o_proj import deep_gemm_fp8_o_proj
17
+from vllm.platforms import current_platform
18
+
19
+
20
+def _fused_wo_a_max_tokens(layer: nn.Module) -> int:
21
+ max_tokens = getattr(layer, "_fused_wo_a_token_limit", None)
22
+ if max_tokens is None:
23
+ # Additional token tiles only pay off for smaller grids.
24
+ ctas_per_tile = layer.n_local_heads * (layer.o_lora_rank // 128)
25
+ max_tokens = 96 if ctas_per_tile <= 128 else 32
26
+ layer._fused_wo_a_token_limit = max_tokens
27
+ return max_tokens
28
+
29
+
30
+def _can_fuse_wo_a(layer: nn.Module) -> bool:
31
+ """Whether the attention layer matches the fused WO-A kernel's layout."""
32
+ method = getattr(layer.wo_b, "scheme", layer.wo_b.quant_method)
33
+ return (
34
+ current_platform.is_device_capability_family(100)
35
+ # WO-A is FP8 with per-32 scales, and WO-B takes MXFP8 input directly.
36
+ and layer._einsum_recipe == (1, 1, 32)
37
+ and get_input_quant_key(layer.wo_b) == kMxfp8Dynamic
38
+ and isinstance(
39
+ getattr(method, "kernel", None), FlashInferCutedslMxfp8LinearKernel
40
+ )
41
+ # A group's heads form one portable (<= 8 CTA) cluster.
42
+ and layer.n_local_heads // layer.n_local_groups <= 8
43
+ and (layer.nope_head_dim, layer.rope_head_dim) == (448, 64)
44
+ and layer.o_lora_rank % 128 == 0
45
+ )
46
+
47
+
48
+def register_dsv41_o_proj_warmup(layer: nn.Module) -> None:
49
+ """Warm the fused WO-A token counts ``dsv41_o_proj`` can dispatch."""
50
+ # Draft models load outside the warmup registry; their layers share the
51
+ # target's attention shapes, so the target's registration covers them too.
52
+ if _can_fuse_wo_a(layer):
53
+ from .fused_wo_a import _FUSED_WO_A_KERNEL
54
+
55
+ _FUSED_WO_A_KERNEL.register_warmup(
56
+ max_tokens=_fused_wo_a_max_tokens(layer),
57
+ n_groups=layer.n_local_groups,
58
+ heads_per_group=layer.n_local_heads // layer.n_local_groups,
59
+ o_lora_rank=layer.o_lora_rank,
60
+ )
61
+
62
+
63
+def dsv41_o_proj(
64
+ layer: nn.Module, attn_out: torch.Tensor, positions: torch.Tensor
65
+) -> torch.Tensor:
66
+ """``deep_gemm_fp8_o_proj`` with WO-A fused for small SM100/SM103 batches."""
67
+ o = attn_out[:, : layer.n_local_heads, :]
68
+ cos_sin_cache = layer.rotary_emb.cos_sin_cache
69
+ if 1 <= o.shape[0] <= _fused_wo_a_max_tokens(layer) and _can_fuse_wo_a(layer):
70
+ from .fused_wo_a import _FUSED_WO_A_KERNEL
71
+
72
+ q, scales = _FUSED_WO_A_KERNEL(
73
+ x=o,
74
+ positions=positions,
75
+ rope=cos_sin_cache,
76
+ weight=layer.wo_a.weight,
77
+ weight_scale=layer.wo_a.weight_scale,
78
+ )
79
+ return layer._wo_b_proj(
80
+ QuantizedActivation(q, scales, o.dtype, q.shape, kMxfp8Dynamic)
81
+ )
82
+ return deep_gemm_fp8_o_proj(
83
+ o,
84
+ positions,
85
+ cos_sin_cache,
86
+ layer.wo_a,
87
+ layer._wo_b_proj,
88
+ n_groups=layer.n_local_groups,
89
+ heads_per_group=layer.n_local_heads // layer.n_local_groups,
90
+ nope_dim=layer.nope_head_dim,
91
+ rope_dim=layer.rope_head_dim,
92
+ o_lora_rank=layer.o_lora_rank,
93
+ einsum_recipe=layer._einsum_recipe,
94
+ tma_aligned_scales=layer._tma_aligned_scales,
95
+ )