@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d6cce94fd4

[Perf][Qwen4Exp] Fuse HC down projection and SiLU on NVIDIA (#58957)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Sep 28, 2026 at 18:51 UTC d6cce94fd4221e6f5547e1c00ca5e0c92ab52d6c
7 files changed +1175 -23
tests/models/qwen4_exp/test_hc_ops.py
+30
@@ -4,11 +4,13 @@
4 import pytest
5 import torch
6
7 +from vllm.models.qwen4_exp.nvidia.ops.cute_dsl.hc_down_silu import hc_down_silu
8 from vllm.models.qwen4_exp.nvidia.ops.hc import (
9 grouped_gemma_rmsnorm,
10 hc_combine,
11 hc_combine_norm,
12 hc_gate_mix,
13 + hc_silu,
14 )
15 from vllm.platforms import current_platform
16 from vllm.triton_utils import HAS_TRITON
@@ -22,6 +24,13 @@ HC = 4
24 HIDDEN_SIZE = 2560
25 HYPER_HIDDEN_SIZE = HC * HIDDEN_SIZE
26 EPS = 1e-6
27 +LORA_RANK = 320
28 +DOWN_N = LORA_RANK + HC + 12 # merged down+inject weight, 16-row padded
29 +
30 +requires_sm90 = pytest.mark.skipif(
31 + not current_platform.has_device_capability(90),
32 + reason="fused HC down+SiLU requires SM90+",
33 +)
34
35
36 def test_grouped_gemma_rmsnorm() -> None:
@@ -125,3 +134,24 @@ def test_hc_combine_norm_unit_injection(num_tokens: int) -> None:
134 expected_norm = grouped_gemma_rmsnorm(expected, weight, EPS, HC)
135 assert torch.equal(actual, expected)
136 torch.testing.assert_close(actual_norm, expected_norm)
137 +
138 +
139 +@requires_sm90
140 +@pytest.mark.parametrize("num_tokens", [1, 3, 5, 17, 48])
141 +def test_hc_down_silu_fused(num_tokens: int) -> None:
142 + # Compare computed columns with the unfused ll_bf16 + hc_silu reference.
143 + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm
144 +
145 + torch.manual_seed(0)
146 + x = torch.randn(num_tokens, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
147 + weight = torch.randn(DOWN_N, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
148 +
149 + lora, injection = hc_down_silu(x, weight, LORA_RANK, HC)
150 +
151 + down = ll_bf16_gemm(x, weight).to(torch.bfloat16)
152 + torch.testing.assert_close(
153 + lora, hc_silu(down[:, :LORA_RANK], HC), rtol=0.01, atol=0.01
154 + )
155 + torch.testing.assert_close(
156 + injection, down[:, LORA_RANK : LORA_RANK + HC], rtol=0.01, atol=0.01
157 + )
vllm/models/qwen4_exp/nvidia/hyperconnection.py
+38 -23
@@ -25,16 +25,19 @@ Typical usage inside a transformer decoder layer::
25 import torch
26 from torch import nn
27
28 +import vllm.envs as envs
29 from vllm.model_executor.layers.linear import (
30 MergedColumnParallelLinear,
31 ReplicatedLinear,
32 )
33 from vllm.model_executor.models.utils import maybe_prefix
34 +from vllm.platforms import current_platform
35
36 from ..common.hyperconnection import (
37 GroupedGemmaRMSNorm,
38 HyperConnectionConfig,
39 )
40 +from .ops.cute_dsl.hc_down_silu import MAX_FUSED_M, hc_down_silu
41 from .ops.hc import (
42 grouped_gemma_rmsnorm,
43 hc_combine,
@@ -58,9 +61,9 @@ class GatedResidual(nn.Module):
61 injection.
62
63 Weights: the norm owns the grouped GemmaRMSNorm affine; the projections
61 - are vLLM Linear modules (merged replicated linear for down+inject), so
62 - GEMM dispatch (e.g. the low-latency skinny GEMM) applies through the
63 - standard quant_method mechanism.
64 + are vLLM Linear modules (merged replicated linear for down+inject).
65 + Eligible down+inject projections use the fused SiLU GEMM; other projections
66 + use their Linear module's GEMM dispatch.
67 """
68
69 def __init__(
@@ -105,6 +108,12 @@ class GatedResidual(nn.Module):
108 return_bias=False,
109 disable_tp=True,
110 )
111 + weight = self.input_mix_weight_down_block_inject.weight
112 + self._use_hc_down_silu = (
113 + weight.shape[1] % 8 == 0
114 + and weight.dtype == torch.bfloat16
115 + and current_platform.has_device_capability(90)
116 + )
117 else:
118 self.input_mix_weight_down = ReplicatedLinear(
119 self.hyper_hidden_size,
@@ -125,6 +134,30 @@ class GatedResidual(nn.Module):
134 return_bias=False,
135 )
136
137 + def _down_and_inject(
138 + self, xn: torch.Tensor
139 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
140 + """Down projection + SiLU; also returns injection logits if combined."""
141 + if not self.use_combine:
142 + return hc_silu(self.input_mix_weight_down(xn), self.hc_count), None
143 +
144 + use_fused = (
145 + self._use_hc_down_silu
146 + and not envs.VLLM_BATCH_INVARIANT
147 + and 1 <= xn.shape[0] <= MAX_FUSED_M
148 + )
149 + if use_fused:
150 + return hc_down_silu(
151 + xn,
152 + self.input_mix_weight_down_block_inject.weight,
153 + self.lora_rank,
154 + self.hc_count,
155 + )
156 + down_and_injection = self.input_mix_weight_down_block_inject(xn)
157 + split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
158 + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
159 + return hc_silu(lora, self.hc_count), injection
160 +
161 def mix(
162 self, hidden_states: torch.Tensor
163 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
@@ -135,16 +168,7 @@ class GatedResidual(nn.Module):
168 self.hc_count,
169 )
170
138 - if self.use_combine:
139 - # produce injection logits for combine
140 - split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
141 - down_and_injection = self.input_mix_weight_down_block_inject(xn)
142 - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
143 - else:
144 - lora = self.input_mix_weight_down(xn)
145 - injection = None
146 -
147 - lora = hc_silu(lora, self.hc_count)
171 + lora, injection = self._down_and_inject(xn)
172 gate = self.input_mix_weight_up(lora) # [M, D]
173 block_input = hc_gate_mix(xn, gate, self.hc_count)
174
@@ -172,16 +196,7 @@ class GatedResidual(nn.Module):
196 self.hc_count,
197 )
198
175 - if self.use_combine:
176 - # produce injection logits for combine
177 - split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
178 - down_and_injection = self.input_mix_weight_down_block_inject(xn)
179 - lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
180 - else:
181 - lora = self.input_mix_weight_down(xn)
182 - injection = None
183 -
184 - lora = hc_silu(lora, self.hc_count)
199 + lora, injection = self._down_and_inject(xn)
200 gate = self.input_mix_weight_up(lora) # [M, D]
201 block_input = hc_gate_mix(xn, gate, self.hc_count)
202
vllm/models/qwen4_exp/nvidia/model.py
+11
@@ -78,6 +78,7 @@ from vllm.v1.kv_cache_interface import MambaSpec
78 from ..config import ATTENTION_LAYER_TYPES, QSA_LAYER_TYPE, Qwen4ExpConfig
79 from .hyperconnection import GatedResidual, HyperConnectionConfig
80 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
81 +from .ops.cute_dsl.hc_down_silu import request_hc_down_silu_warmup
82 from .ple_layer import Qwen4ExpPLELayer
83 from .qsa import Qwen4ExpQSAAttention
84
@@ -686,6 +687,16 @@ class Qwen4ExpForCausalLM(
687 )
688 self.set_moe_parameters(self.model.layers)
689 enable_qwen4_exp_low_latency_gemm(self, self.model_config.dtype)
690 + if self.model_config.dtype == torch.bfloat16:
691 + # Precompile the fused HC down+SiLU kernels for every CUDA-graph
692 + # capture size in the fused dispatch range, so no CuTe-DSL JIT
693 + # happens during graph capture.
694 + request_hc_down_silu_warmup(
695 + vllm_config.compilation_config.cudagraph_capture_sizes or (),
696 + self.config.hc_lowrank,
697 + self.config.hc_count,
698 + self.config.hidden_size * self.config.hc_count,
699 + )
700
701 @staticmethod
702 def get_model_state_cls():
vllm/models/qwen4_exp/nvidia/ops/cute_dsl/__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/qwen4_exp/nvidia/ops/cute_dsl/_hc_down_silu_fma.py new
+334
@@ -0,0 +1,334 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Adapted from vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py.
4 +
5 +Fuses the Qwen4Exp mHC SiLU epilogue (SiLU(x/HC) on columns < rank, passthrough
6 +elsewhere, with the production bf16 rounding boundary) into the FMA kernel.
7 +Output is bf16 instead of fp32.
8 +"""
9 +
10 +import cutlass
11 +import cutlass.cute as cute
12 +from cuda.bindings.driver import CUstream
13 +from cutlass import const_expr
14 +
15 +
16 +def _sigmoid_f32(x):
17 + return cute.math.rcp(1.0 + cute.math.exp(-x), approx=True)
18 +
19 +
20 +class HcDownSiluFma:
21 + """BF16 GEMM kernel based on CTA-local FMA dot products, with the mHC SiLU
22 + epilogue fused into the final store.
23 +
24 + Computes C[M, N] = A[M, K] @ B[N, K]^T for bf16 inputs, rounds the fp32
25 + accumulator to bf16 (production rounding boundary), then applies
26 + SiLU(x/HC) to columns < ``rank`` and passes columns >= ``rank`` through.
27 + Output is bf16 [M, N].
28 +
29 + See the base class docstring in
30 + vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py for the
31 + GEMM structure; only the epilogue differs. The compile key is
32 + ``(M, K, threadblock_size, rank, hc)``.
33 + """
34 +
35 + def __init__(
36 + self,
37 + k: int,
38 + threadblock_size: int = 128,
39 + main_vec_width: int = 8,
40 + tail_vec_width: int = 4,
41 + prefetch_pdl_weights: bool = False,
42 + rank: int = 320,
43 + hc: int = 4,
44 + ):
45 + self.threadblock_size = threadblock_size
46 + self.main_vec_width = main_vec_width
47 + self.tail_vec_width = tail_vec_width
48 + self.prefetch_pdl_weights = prefetch_pdl_weights
49 + self.rank = rank
50 + self.hc = hc
51 + self.num_warps = threadblock_size // cute.arch.WARP_SIZE
52 + self._init_k_tiles(k)
53 + self.main_prefetch_tiles = min(self.main_tiles, 8)
54 +
55 + def _vectorized_elems(self, k_extent: int, vec_width: int) -> int:
56 + vector_tile = vec_width * self.threadblock_size
57 + return (k_extent // vector_tile) * vector_tile
58 +
59 + def _init_k_tiles(self, k: int) -> None:
60 + """Split K into vector loops, scalar rounds, and ragged tail."""
61 + self.k_main_elems = self._vectorized_elems(k, self.main_vec_width)
62 + self.k_after_main = k - self.k_main_elems
63 + self.k_tail_elems = self._vectorized_elems(
64 + self.k_after_main, self.tail_vec_width
65 + )
66 + self.k_done_all = self.k_main_elems + self.k_tail_elems
67 + self.scalar_rem = k - self.k_done_all
68 + self.ks_full = self.scalar_rem // self.threadblock_size
69 + self.ks_part = self.scalar_rem % self.threadblock_size
70 + self.k_scalar_full = self.ks_full * self.threadblock_size
71 + self.k_part_offset = self.k_done_all + self.k_scalar_full
72 + self.main_tiles = self.k_main_elems // (
73 + self.main_vec_width * self.threadblock_size
74 + )
75 + self.tail_tiles = self.k_tail_elems // (
76 + self.tail_vec_width * self.threadblock_size
77 + )
78 +
79 + @cute.jit
80 + def _vector_fma(
81 + self,
82 + acc: cute.Tensor,
83 + tA: cute.Tensor,
84 + tB: cute.Tensor,
85 + M: cutlass.Constexpr,
86 + num_tiles: cutlass.Constexpr,
87 + align_bytes: cutlass.Constexpr,
88 + ):
89 + for tile in cutlass.range_constexpr(num_tiles):
90 + bt = tB[None, tile]
91 + br = cute.make_rmem_tensor_like(bt)
92 + cute.autovec_copy(bt, br)
93 + br_f32 = br.load().to(cutlass.Float32)
94 +
95 + for m in cutlass.range_constexpr(M):
96 + at = tA[m, None, tile]
97 + ar = cute.make_rmem_tensor_like(at)
98 + cute.autovec_copy(at, ar)
99 + vec_width: cutlass.Constexpr = cute.size(ar)
100 + for v in cutlass.range_constexpr(vec_width):
101 + acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v]
102 +
103 + @cute.jit
104 + def _vector_fma_prefetched(
105 + self,
106 + acc: cute.Tensor,
107 + tA: cute.Tensor,
108 + tB: cute.Tensor,
109 + rB: cute.Tensor,
110 + M: cutlass.Constexpr,
111 + num_tiles: cutlass.Constexpr,
112 + prefetch_tiles: cutlass.Constexpr,
113 + ):
114 + for tile in cutlass.range_constexpr(num_tiles):
115 + if const_expr(tile < prefetch_tiles):
116 + br_f32 = rB[tile, None].load().to(cutlass.Float32)
117 + else:
118 + bt = tB[None, tile]
119 + br = cute.make_rmem_tensor_like(bt)
120 + cute.autovec_copy(bt, br)
121 + br_f32 = br.load().to(cutlass.Float32)
122 + for m in cutlass.range_constexpr(M):
123 + at = tA[m, None, tile]
124 + ar = cute.make_rmem_tensor_like(at)
125 + cute.autovec_copy(at, ar)
126 + vec_width: cutlass.Constexpr = cute.size(ar)
127 + for v in cutlass.range_constexpr(vec_width):
128 + acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v]
129 +
130 + def _make_thread_vector_slice(
131 + self,
132 + gA_vec: cute.Tensor,
133 + gB_vec: cute.Tensor,
134 + tidx: cutlass.Int32,
135 + n_idx: cutlass.Int32,
136 + threadblock_size: cutlass.Constexpr,
137 + ):
138 + # (M/N, K_TILE, K_LANE, K_VEC); tidx selects K_LANE.
139 + tA = cute.logical_divide(gA_vec, (None, (None, threadblock_size)))
140 + tB = cute.logical_divide(gB_vec, (None, (None, threadblock_size)))
141 + return tA[None, (None, (tidx, None))], tB[n_idx, (None, (tidx, None))]
142 +
143 + def _make_k_slice(
144 + self,
145 + gX: cute.Tensor,
146 + k_offset: cutlass.Constexpr,
147 + k_extent: cutlass.Constexpr,
148 + ):
149 + k_layout_extent: cutlass.Constexpr = (
150 + 1 if const_expr(k_extent == 0) else k_extent
151 + )
152 +
153 + if const_expr(k_offset == 0):
154 + return cute.local_tile(
155 + gX, (cute.size(gX, mode=[0]), k_layout_extent), (0, 0)
156 + )
157 + return cute.local_tile(
158 + cute.domain_offset((0, k_offset), gX),
159 + (cute.size(gX, mode=[0]), k_layout_extent),
160 + (0, 0),
161 + )
162 +
163 + @cute.jit
164 + def __call__(
165 + self,
166 + gA: cute.Tensor,
167 + gB: cute.Tensor,
168 + gC: cute.Tensor,
169 + M: cutlass.Constexpr,
170 + K_dim: cutlass.Constexpr,
171 + N_dim: cutlass.Int32,
172 + stream: CUstream,
173 + ):
174 + self.kernel(
175 + gA,
176 + gB,
177 + gC,
178 + M,
179 + self.main_vec_width,
180 + self.tail_vec_width,
181 + self.threadblock_size,
182 + self.num_warps,
183 + self.k_main_elems,
184 + self.k_tail_elems,
185 + self.k_done_all,
186 + self.ks_full,
187 + self.ks_part,
188 + self.k_scalar_full,
189 + self.k_part_offset,
190 + self.main_tiles,
191 + self.tail_tiles,
192 + self.rank,
193 + self.hc,
194 + ).launch(
195 + grid=[N_dim, 1, 1],
196 + block=[self.threadblock_size, 1, 1],
197 + smem=M * 4 * self.num_warps,
198 + stream=stream,
199 + use_pdl=True,
200 + min_blocks_per_mp=1,
201 + )
202 +
203 + @cute.kernel
204 + def kernel(
205 + self,
206 + gA: cute.Tensor,
207 + gB: cute.Tensor,
208 + gC: cute.Tensor,
209 + M: cutlass.Constexpr,
210 + main_vec_width: cutlass.Constexpr,
211 + tail_vec_width: cutlass.Constexpr,
212 + threadblock_size: cutlass.Constexpr,
213 + num_warps: cutlass.Constexpr,
214 + k_main_elems: cutlass.Constexpr,
215 + k_tail_elems: cutlass.Constexpr,
216 + k_done_all: cutlass.Constexpr,
217 + ks_full: cutlass.Constexpr,
218 + ks_part: cutlass.Constexpr,
219 + k_scalar_full: cutlass.Constexpr,
220 + k_part_offset: cutlass.Constexpr,
221 + main_tiles: cutlass.Constexpr,
222 + tail_tiles: cutlass.Constexpr,
223 + rank: cutlass.Constexpr,
224 + hc: cutlass.Constexpr,
225 + ):
226 + tidx, _, _ = cute.arch.thread_idx()
227 + n_idx, _, _ = cute.arch.block_idx()
228 + wid = cute.arch.warp_idx()
229 +
230 + # One FP32 accumulator per token.
231 + acc = cute.make_rmem_tensor((M,), cutlass.Float32)
232 + acc.fill(0.0)
233 +
234 + if const_expr(k_main_elems > 0):
235 + gA_main = self._make_k_slice(gA, 0, k_main_elems)
236 + gB_main = self._make_k_slice(gB, 0, k_main_elems)
237 + gA_vec = cute.logical_divide(gA_main, (None, main_vec_width))
238 + gB_vec = cute.logical_divide(gB_main, (None, main_vec_width))
239 + tA, tB = self._make_thread_vector_slice(
240 + gA_vec, gB_vec, tidx, n_idx, threadblock_size
241 + )
242 + if const_expr(self.prefetch_pdl_weights):
243 + prefetch_tiles: cutlass.Constexpr = self.main_prefetch_tiles
244 + prefetched_b = cute.make_rmem_tensor(
245 + cute.make_layout(
246 + (prefetch_tiles, main_vec_width),
247 + stride=(main_vec_width, 1),
248 + ),
249 + cutlass.BFloat16,
250 + )
251 + for tile in cutlass.range_constexpr(prefetch_tiles):
252 + cute.autovec_copy(tB[None, tile], prefetched_b[tile, None])
253 +
254 + cute.arch.griddepcontrol_wait()
255 +
256 + # 128-bit vectorized main loop
257 + if const_expr(k_main_elems > 0):
258 + if const_expr(self.prefetch_pdl_weights):
259 + self._vector_fma_prefetched(
260 + acc,
261 + tA,
262 + tB,
263 + prefetched_b,
264 + M,
265 + main_tiles,
266 + prefetch_tiles,
267 + )
268 + else:
269 + self._vector_fma(acc, tA, tB, M, main_tiles, 16)
270 +
271 + # 64-bit vectorized tail (K remainder after main loop)
272 + if const_expr(k_tail_elems > 0):
273 + gA_tail = self._make_k_slice(gA, k_main_elems, k_tail_elems)
274 + gB_tail = self._make_k_slice(gB, k_main_elems, k_tail_elems)
275 + gA_tail_vec = cute.logical_divide(gA_tail, (None, tail_vec_width))
276 + gB_tail_vec = cute.logical_divide(gB_tail, (None, tail_vec_width))
277 + tA_t, tB_t = self._make_thread_vector_slice(
278 + gA_tail_vec, gB_tail_vec, tidx, n_idx, threadblock_size
279 + )
280 + self._vector_fma(acc, tA_t, tB_t, M, tail_tiles, 8)
281 +
282 + # Full scalar rounds use CuTe width-1 tiles; KS_PART is the ragged tail.
283 + if const_expr(ks_full > 0):
284 + gA_scalar = self._make_k_slice(gA, k_done_all, k_scalar_full)
285 + gB_scalar = self._make_k_slice(gB, k_done_all, k_scalar_full)
286 + gA_scalar_vec = cute.logical_divide(gA_scalar, (None, 1))
287 + gB_scalar_vec = cute.logical_divide(gB_scalar, (None, 1))
288 + tA_s, tB_s = self._make_thread_vector_slice(
289 + gA_scalar_vec, gB_scalar_vec, tidx, n_idx, threadblock_size
290 + )
291 + self._vector_fma(acc, tA_s, tB_s, M, ks_full, 2)
292 +
293 + # Only threads below KS_PART load the ragged tail.
294 + if const_expr(ks_part > 0):
295 + gA_part = self._make_k_slice(gA, k_part_offset, ks_part)
296 + gB_part = self._make_k_slice(gB, k_part_offset, ks_part)
297 + if tidx < ks_part:
298 + bv2 = gB_part[n_idx, tidx].to(cutlass.Float32)
299 + for m in cutlass.range_constexpr(M):
300 + acc[m] = acc[m] + gA_part[m, tidx].to(cutlass.Float32) * bv2
301 +
302 + # Intra-warp shuffle reduction
303 + for m in cutlass.range_constexpr(M):
304 + acc[m] = cute.arch.warp_reduction_sum(acc[m])
305 +
306 + # Cross-warp reduction via shared memory
307 + smem_red_layout = cute.make_layout((M, num_warps), stride=(num_warps, 1))
308 + smem = cutlass.utils.SmemAllocator()
309 + sm = smem.allocate_tensor(cutlass.Float32, smem_red_layout, byte_alignment=16)
310 + with cute.arch.elect_one():
311 + for m in cutlass.range_constexpr(M):
312 + sm[m, wid] = acc[m]
313 +
314 + # Final reduction, fused mHC SiLU epilogue, and output.
315 + cute.arch.sync_threads()
316 + if tidx == 0:
317 + for m in cutlass.range_constexpr(M):
318 + partials = sm[m, None].load()
319 + total = cutlass.Float32(
320 + partials.reduce(
321 + cute.ReductionOp.ADD,
322 + init_val=cutlass.Float32(0.0),
323 + reduction_profile=0,
324 + )
325 + )
326 + # Production rounding boundary: GEMM output is materialized as
327 + # bf16 before _hc_silu_kernel reads it.
328 + x_bf16 = total.to(cutlass.BFloat16)
329 + if n_idx < rank:
330 + z = x_bf16.to(cutlass.Float32) * (1.0 / hc)
331 + gC[m, n_idx] = (z * _sigmoid_f32(z)).to(cutlass.BFloat16)
332 + else:
333 + gC[m, n_idx] = x_bf16
334 + cute.arch.griddepcontrol_launch_dependents()
vllm/models/qwen4_exp/nvidia/ops/cute_dsl/_hc_down_silu_mma.py new
+536
@@ -0,0 +1,536 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Adapted from vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py.
4 +
5 +Fuses the Qwen4Exp mHC SiLU epilogue (SiLU(x/HC) on columns < rank, passthrough
6 +elsewhere, with the production bf16 rounding boundary) into the split-K kernel,
7 +applied after the cluster reduction where the final fp32 result materializes.
8 +Output is bf16 instead of fp32.
9 +"""
10 +
11 +import math
12 +
13 +import cutlass
14 +import cutlass.cute as cute
15 +from cuda.bindings.driver import CUstream
16 +from cutlass._mlir import ir as _ir
17 +from cutlass._mlir.dialects import llvm as _llvm
18 +from cutlass.cutlass_dsl import dsl_user_op
19 +from cutlass.pipeline import sm90 as pipeline
20 +
21 +
22 +# Map a local smem address to the same peer CTA DSMEM offset.
23 +@dsl_user_op
24 +def set_block_rank(smem_ptr, peer_rank, *, loc=None, ip=None):
25 + dsmem_ptr = cute.arch.map_dsmem_ptr(smem_ptr, peer_rank, loc=loc, ip=ip)
26 + return cutlass.Int32(dsmem_ptr.toint(loc=loc, ip=ip))
27 +
28 +
29 +# Plain DSMEM store; mbarrier helpers do not model this reduction.
30 +@dsl_user_op
31 +def st_shared_remote_f32(remote_addr, val, *, loc=None, ip=None):
32 + i32 = _ir.IntegerType.get_signless(32)
33 + addr_ir = remote_addr.ir_value(loc=loc, ip=ip)
34 + val_ir = val.ir_value(loc=loc, ip=ip)
35 + _llvm.inline_asm(
36 + i32,
37 + [addr_ir, val_ir],
38 + "st.shared::cluster.f32 [$0], $1; mov.u32 $2, 0;",
39 + "r,f,=r",
40 + has_side_effects=True, # keep the inline store ordered
41 + loc=loc,
42 + ip=ip,
43 + )
44 +
45 +
46 +def _sigmoid_f32(x):
47 + return cute.math.rcp(1.0 + cute.math.exp(-x), approx=True)
48 +
49 +
50 +class HcDownSiluMma:
51 + """BF16 GEMM kernel based on clustered split-K MMA, with the mHC SiLU
52 + epilogue fused into the final store.
53 +
54 + Computes C[M, N] = A[M, K] @ B[N, K]^T for bf16 inputs, rounds the reduced
55 + fp32 accumulator to bf16 (production rounding boundary), then applies
56 + SiLU(x/HC) to columns < ``rank`` and passes columns >= ``rank`` through.
57 + Output is bf16 [M, N].
58 +
59 + See the base class docstring in
60 + vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py for the
61 + GEMM structure; only the epilogue differs. The compile key is
62 + ``(split_k, num_stages, tile_n, rank, hc)``.
63 + """
64 +
65 + def __init__(
66 + self,
67 + ab_dtype=cutlass.BFloat16,
68 + acc_dtype=cutlass.Float32,
69 + tile_n: int = 16,
70 + tile_k: int = 256,
71 + num_stages: int = 2,
72 + num_dma_warps: int = 4,
73 + split_k: int = 8,
74 + rank: int = 320,
75 + hc: int = 4,
76 + ):
77 + self.ab_dtype = ab_dtype
78 + self.acc_dtype = acc_dtype
79 + self.rank = rank
80 + self.hc = hc
81 + self.tile_m = 16
82 + self.tile_n = tile_n
83 + self.tile_k = tile_k
84 + self.copy_bits = 128
85 + self.num_stages = num_stages
86 + self.split_k = split_k
87 + self.mma_shape = (16, 8, 16) # mma.sync.aligned.m16n8k16
88 + self.atom_layout = (1, 1, 1) # one MMA atom per warp
89 + self.num_dma_warps = num_dma_warps
90 + self.num_mma_warps = 4
91 + self.num_dma_threads = self.num_dma_warps * cute.arch.WARP_SIZE
92 + self.num_mma_threads = self.num_mma_warps * cute.arch.WARP_SIZE
93 + self.num_threads = self.num_dma_threads + self.num_mma_threads
94 + self.num_epilogue_elems = self.tile_m * self.tile_n
95 + self.epilogue_elems_per_thread = self.num_epilogue_elems // self.num_mma_threads
96 +
97 + def _make_smem_layout_AB(self, dtype, copy_bits, smem_tiler):
98 + """Build the staged swizzled SMEM layout for A or B tiles."""
99 + major_size = min(smem_tiler[1], 64)
100 + # Match swizzle span to contiguous K bytes, capped by CuTe 3-bit swizzle.
101 + swizzle_bits = int(math.log2(major_size * dtype.width // copy_bits))
102 + swizzle_bits = min(swizzle_bits, 3)
103 + # Tile the swizzled atom across (M_or_N, K, stages).
104 + layout_atom_outer = cute.make_layout((8, major_size), stride=(major_size, 1))
105 + layout_atom = cute.make_composed_layout(
106 + cute.make_swizzle(swizzle_bits, 3, 3), 0, layout_atom_outer
107 + )
108 + return cute.tile_to_shape(layout_atom, smem_tiler, (0, 1, 2))
109 +
110 + def _make_gmem_tiled_copy(self, atom_copy, dtype, copy_bits, num_threads):
111 + """Build the per-thread cp.async vector-copy layout."""
112 + # Lay threads across K so each lane issues one vector copy.
113 + copy_elems = copy_bits // dtype.width
114 + k_threads = cute.size(self.tile_k) // copy_elems # threads along K
115 + thread_layout = cute.make_layout(
116 + (num_threads // k_threads, k_threads), stride=(k_threads, 1)
117 + )
118 + value_layout = cute.make_layout((1, copy_elems))
119 + return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout)
120 +
121 + @cute.jit
122 + def _fill_pred(self, pred_flat, coord_tensor, k_tile, dim_limit, K_total):
123 + # Predicate one K tile and keep the stage-broadcast view in sync.
124 + coord_ktile = coord_tensor[None, None, 0, k_tile]
125 + num_vec = pred_flat.shape[0]
126 + num_mn = pred_flat.shape[1]
127 + for v in cutlass.range_constexpr(num_vec):
128 + # pred_flat is (K_VEC, M/N) for one K_TILE.
129 + for j in cutlass.range_constexpr(num_mn):
130 + pred_flat[v, j] = cute.elem_less(
131 + coord_ktile[(0, v), j], (dim_limit, K_total)
132 + )
133 +
134 + def _make_pred(self, tXcX, k_tile, dim_limit, K_total):
135 + # pred_flat is (K_VEC, M/N); pred is (K_VEC, M/N, STAGE).
136 + num_vec = tXcX.shape[0][1]
137 + num_mn = cute.size(tXcX, mode=[1])
138 + pred_flat = cute.make_rmem_tensor(
139 + cute.make_layout((num_vec, num_mn), stride=(num_mn, 1)),
140 + cutlass.Boolean,
141 + )
142 + pred = cute.make_tensor(
143 + pred_flat.iterator,
144 + cute.make_layout(
145 + (num_vec, num_mn, cute.size(tXcX, mode=[2])),
146 + stride=(num_mn, 1, 0),
147 + ),
148 + )
149 + return pred_flat, pred
150 +
151 + @cute.jit
152 + def __call__(
153 + self,
154 + mA: cute.Tensor,
155 + mB: cute.Tensor,
156 + mC: cute.Tensor,
157 + stream: CUstream,
158 + ):
159 + bM, bN, bK = self.tile_m, self.tile_n, self.tile_k
160 + copy_bits: cutlass.Constexpr = self.copy_bits
161 + sA_layout = self._make_smem_layout_AB(
162 + mA.element_type, copy_bits, (bM, bK, self.num_stages)
163 + )
164 + sB_layout = self._make_smem_layout_AB(
165 + mB.element_type, copy_bits, (bN, bK, self.num_stages)
166 + )
167 +
168 + @cute.struct
169 + class SharedStorage:
170 + a: cute.struct.Align[
171 + cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)], 16
172 + ]
173 + b: cute.struct.Align[
174 + cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)], 16
175 + ]
176 + mbar: cute.struct.Align[
177 + cute.struct.MemRange[cutlass.Int64, self.num_stages * 2], 8
178 + ]
179 +
180 + atom_g2s = cute.make_copy_atom(
181 + cute.nvgpu.cpasync.CopyG2SOp(
182 + cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL
183 + ),
184 + mA.element_type,
185 + num_bits_per_copy=copy_bits,
186 + ) # cp.async GMEM -> SMEM, bypassing L1
187 + tiled_copy_A = self._make_gmem_tiled_copy(
188 + atom_g2s, mA.element_type, copy_bits, self.num_dma_threads
189 + )
190 + tiled_copy_B = self._make_gmem_tiled_copy(
191 + atom_g2s, mB.element_type, copy_bits, self.num_dma_threads
192 + )
193 + op = cute.nvgpu.warp.MmaF16BF16Op(self.ab_dtype, self.acc_dtype, self.mma_shape)
194 + # Repeat the m16n8k16 atom along N to cover the CTA output tile.
195 + perm_mnk = (
196 + self.atom_layout[0] * self.mma_shape[0],
197 + self.atom_layout[1] * self.mma_shape[1] * (self.tile_n // 8),
198 + self.atom_layout[2] * self.mma_shape[2],
199 + )
200 + tiled_mma = cute.make_tiled_mma(
201 + op, cute.make_layout(self.atom_layout), permutation_mnk=perm_mnk
202 + )
203 + tiler_mn = (bM, bN)
204 + grid_m, grid_n = cute.ceil_div(mC.shape, tiler_mn)
205 + self.kernel(
206 + mA,
207 + mB,
208 + mC,
209 + sA_layout,
210 + sB_layout,
211 + tiled_copy_A,
212 + tiled_copy_B,
213 + tiled_mma,
214 + SharedStorage,
215 + self.rank,
216 + self.hc,
217 + ).launch(
218 + grid=[
219 + cute.size(grid_m),
220 + cute.size(grid_n),
221 + self.split_k,
222 + ],
223 + block=[self.num_threads, 1, 1],
224 + cluster=[
225 + 1,
226 + 1,
227 + self.split_k,
228 + ], # split-K CTAs form one cluster
229 + stream=stream,
230 + use_pdl=True,
231 + )
232 +
233 + @cute.kernel
234 + def kernel(
235 + self,
236 + mA,
237 + mB,
238 + mC,
239 + sA_layout: cute.ComposedLayout,
240 + sB_layout: cute.ComposedLayout,
241 + tiled_copy_A: cute.TiledCopy,
242 + tiled_copy_B: cute.TiledCopy,
243 + tiled_mma: cute.TiledMma,
244 + shared_storage: cutlass.Constexpr,
245 + rank: cutlass.Constexpr,
246 + hc: cutlass.Constexpr,
247 + ):
248 + bM, bN, bK = self.tile_m, self.tile_n, self.tile_k
249 + num_stages = self.num_stages
250 + tidx, _, _ = cute.arch.thread_idx()
251 + bid_m, bid_n, bid_z = cute.arch.block_idx()
252 + warp_idx = cute.arch.warp_idx()
253 + lane_id = cute.arch.lane_idx()
254 + num_dma_warps: cutlass.Constexpr = self.num_dma_warps
255 + is_dma_warp = warp_idx < num_dma_warps
256 + dma_tidx = tidx
257 + mma_tidx = tidx - self.num_dma_threads
258 + mma_warp_idx = warp_idx - num_dma_warps
259 +
260 + cta_tiler = (bM, bN, bK)
261 + coord = (bid_m, bid_n, None) # all K tiles
262 + # CTA-local tiles.
263 + gA = cute.local_tile(
264 + mA, tiler=cta_tiler, coord=coord, proj=(1, None, 1)
265 + ) # skip N
266 + gB = cute.local_tile(
267 + mB, tiler=cta_tiler, coord=coord, proj=(None, 1, 1)
268 + ) # skip M
269 + gC = cute.local_tile(
270 + mC, tiler=cta_tiler, coord=coord, proj=(1, 1, None)
271 + ) # skip K
272 +
273 + mcA = cute.make_identity_tensor(mA.layout.shape)
274 + mcB = cute.make_identity_tensor(mB.layout.shape)
275 + mcC = cute.make_identity_tensor(mC.layout.shape)
276 + # Coordinate modes: cA=(M,K,k_tile), cB=(N,K,k_tile), cC=(M,N).
277 + cA = cute.local_tile(mcA, tiler=cta_tiler, coord=coord, proj=(1, None, 1))
278 + cB = cute.local_tile(mcB, tiler=cta_tiler, coord=coord, proj=(None, 1, 1))
279 + cC = cute.local_tile(mcC, tiler=cta_tiler, coord=coord, proj=(1, 1, None))
280 +
281 + # 128-bit cp.async copies require 16-byte aligned GMEM views.
282 + gA = cute.make_tensor(gA.iterator.align(16), gA.layout)
283 + gB = cute.make_tensor(gB.iterator.align(16), gB.layout)
284 +
285 + smem = cutlass.utils.SmemAllocator()
286 + storage_ptr = smem.allocate(shared_storage.size_in_bytes(), byte_alignment=16) # type: ignore[attr-defined]
287 + storage = shared_storage(storage_ptr) # type: ignore[call-arg]
288 + sA = storage.a.get_tensor(sA_layout)
289 + sB = storage.b.get_tensor(sB_layout)
290 +
291 + # Pipeline cp.async producers into MMA consumers.
292 + producer_group = pipeline.CooperativeGroup(
293 + pipeline.Agent.Thread, self.num_dma_threads
294 + )
295 + consumer_group = pipeline.CooperativeGroup(
296 + pipeline.Agent.Thread, self.num_mma_threads
297 + )
298 + mainloop_pipeline = pipeline.PipelineCpAsync.create(
299 + barrier_storage=storage.mbar.data_ptr(),
300 + num_stages=num_stages,
301 + producer_group=producer_group,
302 + consumer_group=consumer_group,
303 + )
304 +
305 + # Round-robin split-K: split z handles tiles z, z+split_k, ...
306 + K_total = cute.size(mA, mode=[1])
307 + k_tile_count = cute.size(gA, mode=[2])
308 + k_start = bid_z
309 + num_k_tiles = cute.ceil_div(k_tile_count - k_start, self.split_k)
310 +
311 + if is_dma_warp:
312 + # DMA warps trade registers for copy throughput.
313 + cute.arch.setmaxregister_decrease(40)
314 + thr_A = tiled_copy_A.get_slice(dma_tidx)
315 + thr_B = tiled_copy_B.get_slice(dma_tidx)
316 + tAgA = thr_A.partition_S(gA)
317 + tAsA = thr_A.partition_D(sA)
318 + tBgB = thr_B.partition_S(gB)
319 + tBsB = thr_B.partition_D(sB)
320 + tAcA = thr_A.partition_S(cA)
321 + tBcB = thr_B.partition_S(cB)
322 +
323 + # Build M/K and N/K predicates, broadcast across K-tile copies.
324 + tApA_flat, tApA = self._make_pred(tAcA, k_start, mA.shape[0], K_total)
325 + tBpB_flat, tBpB = self._make_pred(tBcB, k_start, mB.shape[0], K_total)
326 + self._fill_pred(tApA_flat, tAcA, k_start, mA.shape[0], K_total)
327 + self._fill_pred(tBpB_flat, tBcB, k_start, mB.shape[0], K_total)
328 +
329 + producer_state = pipeline.make_pipeline_state(
330 + pipeline.PipelineUserType.Producer, num_stages
331 + )
332 +
333 + # Prime the first pipeline stage.
334 + mainloop_pipeline.producer_acquire(producer_state)
335 + cute.copy(
336 + tiled_copy_B,
337 + tBgB[None, None, None, k_start],
338 + tBsB[None, None, None, producer_state.index],
339 + pred=tBpB,
340 + )
341 + cute.arch.griddepcontrol_wait()
342 + cute.copy(
343 + tiled_copy_A,
344 + tAgA[None, None, None, k_start],
345 + tAsA[None, None, None, producer_state.index],
346 + pred=tApA,
347 + )
348 + mainloop_pipeline.producer_commit(producer_state)
349 + producer_state.advance()
350 +
351 + for k_tile in cutlass.range(
352 + k_start + self.split_k, k_tile_count, self.split_k, unroll=1
353 + ):
354 + self._fill_pred(tApA_flat, tAcA, k_tile, mA.shape[0], K_total)
355 + self._fill_pred(tBpB_flat, tBcB, k_tile, mB.shape[0], K_total)
356 + mainloop_pipeline.producer_acquire(producer_state)
357 + cute.copy(
358 + tiled_copy_A,
359 + tAgA[None, None, None, k_tile],
360 + tAsA[None, None, None, producer_state.index],
361 + pred=tApA,
362 + )
363 + cute.copy(
364 + tiled_copy_B,
365 + tBgB[None, None, None, k_tile],
366 + tBsB[None, None, None, producer_state.index],
367 + pred=tBpB,
368 + )
369 + mainloop_pipeline.producer_commit(producer_state)
370 + producer_state.advance()
371 +
372 + mainloop_pipeline.producer_tail(producer_state)
373 +
374 + else:
375 + # MMA warps with k-phase interleaving
376 + cute.arch.setmaxregister_increase(232) # large MMA fragments
377 +
378 + num_mma_warps: cutlass.Constexpr = self.num_mma_warps
379 +
380 + thr_mma = tiled_mma.get_slice(lane_id)
381 + tCsA = thr_mma.partition_A(sA)
382 + tCsB = thr_mma.partition_B(sB)
383 + tCgC = thr_mma.partition_C(gC)
384 +
385 + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
386 + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
387 + tCrC = tiled_mma.make_fragment_C(tCgC)
388 + tCrC.fill(0.0)
389 +
390 + # ldmatrix moves SMEM fragments into MMA registers.
391 + atom_s2r_A = cute.make_copy_atom(
392 + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), mA.element_type
393 + ) # non-transposed, x4
394 + atom_s2r_B = cute.make_copy_atom(
395 + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), mB.element_type
396 + )
397 +
398 + # SMEM -> register copy path.
399 + tiled_s2r_A = cute.make_tiled_copy_A(atom_s2r_A, tiled_mma)
400 + tiled_s2r_B = cute.make_tiled_copy_B(atom_s2r_B, tiled_mma)
401 + thr_s2r_A = tiled_s2r_A.get_slice(lane_id)
402 + thr_s2r_B = tiled_s2r_B.get_slice(lane_id)
403 + tCsA_v = thr_s2r_A.partition_S(sA) # views, not copies
404 + tCrA_v = thr_s2r_A.retile(tCrA)
405 + tCsB_v = thr_s2r_B.partition_S(sB)
406 + tCrB_v = thr_s2r_B.retile(tCrB)
407 + # Split the MMA K-fragments across the MMA warps.
408 + tCsA_warp_v = cute.logical_divide(tCsA_v, (None, None, num_mma_warps, None))
409 + tCsB_warp_v = cute.logical_divide(tCsB_v, (None, None, num_mma_warps, None))
410 +
411 + num_k_blocks = cute.size(tCrA, mode=[2])
412 + k_blocks_per_warp: cutlass.Constexpr = num_k_blocks // num_mma_warps
413 +
414 + consumer_state = pipeline.make_pipeline_state(
415 + pipeline.PipelineUserType.Consumer, num_stages
416 + )
417 +
418 + # Shape-dynamic split-K count, so this stays a runtime range.
419 + for _ in cutlass.range(num_k_tiles, unroll_full=True):
420 + mainloop_pipeline.consumer_wait(consumer_state)
421 + for ki in cutlass.range_constexpr(k_blocks_per_warp):
422 + cute.copy(
423 + tiled_s2r_A,
424 + tCsA_warp_v[
425 + None, None, (mma_warp_idx, ki), consumer_state.index
426 + ],
427 + tCrA_v[None, None, 0],
428 + ) # ldmatrix
429 + cute.copy(
430 + tiled_s2r_B,
431 + tCsB_warp_v[
432 + None, None, (mma_warp_idx, ki), consumer_state.index
433 + ],
434 + tCrB_v[None, None, 0],
435 + )
436 + cute.gemm(
437 + tiled_mma, tCrC, tCrA[None, None, 0], tCrB[None, None, 0], tCrC
438 + ) # mma.sync
439 + mainloop_pipeline.consumer_release(consumer_state)
440 + consumer_state.advance()
441 +
442 + # Cluster reduction epilogue.
443 + # Reduce per-warp accumulators within this CTA.
444 + num_elems: cutlass.Constexpr = self.num_epilogue_elems
445 + elems_per_thread: cutlass.Constexpr = self.epilogue_elems_per_thread
446 + # Map MMA threads to linear CTA output elements.
447 + epilogue_thread_layout = cute.make_layout(
448 + (elems_per_thread, self.num_mma_threads),
449 + stride=(self.num_mma_threads, 1),
450 + )
451 + epilogue_slots = cute.make_tensor(0, epilogue_thread_layout)
452 + epilogue_slot_coords = cute.make_identity_tensor((bN, bM))
453 + # Layout: (mma_warp, linear MN element).
454 + smem_red = cute.make_tensor(
455 + cute.arch.alloc_smem(
456 + cutlass.Float32, num_elems * num_mma_warps, alignment=16
457 + ),
458 + cute.make_layout((num_mma_warps, num_elems), stride=(num_elems, 1)),
459 + )
460 + smem_warp = cute.make_tensor(
461 + cute.domain_offset((mma_warp_idx, 0), smem_red).iterator,
462 + cute.make_layout((bM, bN), stride=(bN, 1)),
463 + )
464 + tCsC_partial = thr_mma.partition_C(smem_warp)
465 + cute.autovec_copy(tCrC, tCsC_partial)
466 + cute.arch.sync_threads()
467 +
468 + # Layout: (split-K rank, linear MN element).
469 + partials = cute.make_tensor(
470 + cute.arch.alloc_smem(
471 + cutlass.Float32, num_elems * self.split_k, alignment=16
472 + ),
473 + cute.make_layout((self.split_k, num_elems), stride=(num_elems, 1)),
474 + )
475 + cta_rank = cute.arch.block_idx_in_cluster()
476 +
477 + for ei in cutlass.range_constexpr(elems_per_thread):
478 + elem_idx = epilogue_slots[ei, mma_tidx]
479 + local_coord = cute.select(epilogue_slot_coords[elem_idx], mode=[1, 0])
480 + total = cutlass.Float32(0.0)
481 + if cute.elem_less(local_coord, gC.shape):
482 + total = (
483 + smem_red[None, elem_idx]
484 + .load()
485 + .reduce(
486 + cute.ReductionOp.ADD,
487 + init_val=cutlass.Float32(0.0),
488 + reduction_profile=0,
489 + )
490 + )
491 + partials[cta_rank, elem_idx] = total
492 +
493 + cute.arch.sync_threads()
494 +
495 + # Broadcast this CTA's partials to peer DSMEM.
496 + for ei in cutlass.range_constexpr(elems_per_thread):
497 + elem_idx = epilogue_slots[ei, mma_tidx]
498 + my_slot = cute.domain_offset((cta_rank, elem_idx), partials).iterator
499 + my_val = partials[cta_rank, elem_idx]
500 + for peer in cutlass.range_constexpr(self.split_k):
501 + remote = set_block_rank(my_slot, cutlass.Int32(peer))
502 + st_shared_remote_f32(remote, my_val)
503 +
504 + cute.arch.cluster_arrive()
505 + cute.arch.cluster_wait() # peer DSMEM stores are now visible
506 +
507 + if mma_tidx == 0:
508 + cute.arch.griddepcontrol_launch_dependents()
509 + cute.arch.sync_threads()
510 +
511 + # Reduce split-K partials, apply the fused mHC SiLU epilogue, and
512 + # write the bf16 global output.
513 + for ei in cutlass.range_constexpr(elems_per_thread):
514 + elem_idx = epilogue_slots[ei, mma_tidx]
515 + local_coord = cute.select(epilogue_slot_coords[elem_idx], mode=[1, 0])
516 + global_coord = cC[local_coord]
517 + if cute.elem_less(global_coord, mC.shape):
518 + acc = cutlass.Float32(
519 + partials[None, elem_idx]
520 + .load()
521 + .reduce(
522 + cute.ReductionOp.ADD,
523 + init_val=cutlass.Float32(0.0),
524 + reduction_profile=0,
525 + )
526 + )
527 + # Production rounding boundary: the GEMM result is
528 + # materialized as bf16 before _hc_silu_kernel reads it.
529 + x_bf16 = acc.to(cutlass.BFloat16)
530 + if global_coord[1] < rank:
531 + z = x_bf16.to(cutlass.Float32) * (1.0 / hc)
532 + gC[local_coord] = (z * _sigmoid_f32(z)).to(cutlass.BFloat16)
533 + else:
534 + gC[local_coord] = x_bf16
535 +
536 + cute.arch.sync_threads()
vllm/models/qwen4_exp/nvidia/ops/cute_dsl/hc_down_silu.py new
+224
@@ -0,0 +1,224 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Adapted from vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py.
4 +
5 +Fused Qwen4Exp mHC down projection + SiLU epilogue, where
6 +columns < lora_rank get ``silu(bf16(acc) / hc_count)``, the hc_count
7 +injection-logit columns pass through as ``bf16(acc)``, and the pad columns
8 +are never computed. Output is bf16 (ll_bf16's fp32-output bonus is given up
9 +to preserve the production rounding boundary).
10 +
11 +The GEMM uses an FMA backend for M <= 4 and a split-K MMA backend beyond,
12 +both with PDL. The model uses its Linear module above M=48.
13 +"""
14 +
15 +import logging
16 +from collections.abc import Iterable
17 +from functools import lru_cache, partial
18 +from typing import Any, Literal
19 +
20 +import torch
21 +
22 +from vllm.platforms import current_platform
23 +
24 +logger = logging.getLogger(__name__)
25 +
26 +# The fused kernel stops winning past M ~ 64 on Qwen3.8-Next-Flash.
27 +MAX_FUSED_M = 48
28 +
29 +# Split-K configs tuned on SM100f for Qwen3.8-Flash-Next (K=10240, N=324).
30 +_SM100F_TUNED_SPLITK: dict[int, tuple[int, int, int]] = {
31 + **{m: (6, 5, 8) for m in range(5, 17)},
32 + **{m: (6, 5, 16) for m in range(17, 33)},
33 + **{m: (6, 4, 16) for m in range(33, MAX_FUSED_M + 1)},
34 +}
35 +
36 +# FMA: (backend, M, K, threadblock_size); MMA: (backend, split_k, stages, tile_n).
37 +CompileKey = tuple[Literal["fma", "mma"], int, int, int]
38 +
39 +
40 +class HcDownSiluGemm:
41 + """Dispatch/compile cache for the fused mHC down+SiLU GEMM."""
42 +
43 + def __init__(
44 + self, rank: int, hc: int, k: int, *, prefetch_pdl_weights: bool = False
45 + ) -> None:
46 + self.rank = rank
47 + self.hc = hc
48 + self.k = k
49 + self._prefetch_pdl_weights = prefetch_pdl_weights
50 + self._compiled: dict[CompileKey, Any] = {}
51 + self._warmup_m: set[int] = set()
52 + self._warmup_registered = False
53 +
54 + def dispatch(self, m: int) -> CompileKey:
55 + if m <= 4 or self.k < 2048:
56 + return ("fma", m, self.k, 128)
57 + return ("mma", *_SM100F_TUNED_SPLITK[m])
58 +
59 + def compile(self, compile_key: CompileKey) -> None:
60 + if compile_key in self._compiled:
61 + return
62 +
63 + import cutlass.cute as cute
64 + from cutlass import BFloat16
65 + from quack.compile_utils import make_fake_tensor
66 +
67 + N = cute.sym_int()
68 + gemm: Any
69 + extra_args: tuple[int, ...]
70 + if compile_key[0] == "mma":
71 + from ._hc_down_silu_mma import HcDownSiluMma
72 +
73 + _, split_k, num_stages, tile_n = compile_key
74 + M, K = cute.sym_int(), cute.sym_int()
75 + gemm = HcDownSiluMma(
76 + tile_n=tile_n,
77 + num_stages=num_stages,
78 + split_k=split_k,
79 + rank=self.rank,
80 + hc=self.hc,
81 + )
82 + extra_args = ()
83 + options = "--enable-tvm-ffi"
84 + else:
85 + from ._hc_down_silu_fma import HcDownSiluFma
86 +
87 + _, M, K, threadblock_size = compile_key
88 + gemm = HcDownSiluFma(
89 + k=K,
90 + threadblock_size=threadblock_size,
91 + prefetch_pdl_weights=self._prefetch_pdl_weights,
92 + rank=self.rank,
93 + hc=self.hc,
94 + )
95 + extra_args = (M, K, 1) # runtime N placeholder
96 + options = "--enable-tvm-ffi --ptxas-options -maxrregcount=64"
97 +
98 + hidden_states = make_fake_tensor(BFloat16, (M, K), divisibility=8)
99 + router_weight = make_fake_tensor(BFloat16, (N, K), divisibility=8)
100 + output = make_fake_tensor(BFloat16, (M, N), divisibility=1)
101 + self._compiled[compile_key] = cute.compile(
102 + gemm,
103 + hidden_states,
104 + router_weight,
105 + output,
106 + *extra_args,
107 + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
108 + options=options,
109 + )
110 + logger.debug("Compiled hc_down_silu: %s", compile_key)
111 +
112 + def request_warmup(self, m_values: Iterable[int]) -> None:
113 + m_values = set(m_values)
114 + if not m_values:
115 + return
116 + self._warmup_m.update(m_values)
117 + if self._warmup_registered:
118 + return
119 + from vllm.model_executor.warmup.cutedsl_warmup import (
120 + register_cutedsl_warmup_provider,
121 + )
122 +
123 + register_cutedsl_warmup_provider(self)
124 + self._warmup_registered = True
125 +
126 + def get_cutedsl_warmup_compile_units(self):
127 + from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit
128 +
129 + return tuple(
130 + CuTeDSLCompileUnit(
131 + name="Qwen4Exp HC down+SiLU GEMM",
132 + key=(
133 + "qwen4-exp-hc-down-silu",
134 + self.rank,
135 + self.hc,
136 + self.k,
137 + self._prefetch_pdl_weights,
138 + compile_key,
139 + ),
140 + compile=partial(self.compile, compile_key),
141 + )
142 + for compile_key in dict.fromkeys(
143 + self.dispatch(m) for m in sorted(self._warmup_m)
144 + )
145 + )
146 +
147 + def __call__(
148 + self,
149 + hidden_states: torch.Tensor, # [M, K] bf16
150 + router_weight: torch.Tensor, # [N, K] bf16
151 + ) -> torch.Tensor: # [M, N] bf16
152 + M = hidden_states.shape[0]
153 + N = router_weight.shape[0]
154 + n_compute = self.rank + self.hc
155 + w_gemm = router_weight[:n_compute]
156 + compile_key = self.dispatch(M)
157 + self.compile(compile_key)
158 + kernel = self._compiled[compile_key]
159 +
160 + output = torch.empty(M, N, dtype=torch.bfloat16, device=hidden_states.device)
161 + out_gemm = output[:, :n_compute]
162 + if compile_key[0] == "mma":
163 + kernel(hidden_states, w_gemm, out_gemm)
164 + else:
165 + kernel(hidden_states, w_gemm, out_gemm, n_compute)
166 + return output
167 +
168 +
169 +@lru_cache
170 +def _get_kernel(
171 + rank: int, hc: int, k: int, prefetch_pdl_weights: bool = False
172 +) -> HcDownSiluGemm:
173 + return HcDownSiluGemm(rank, hc, k, prefetch_pdl_weights=prefetch_pdl_weights)
174 +
175 +
176 +def hc_down_silu(
177 + x: torch.Tensor,
178 + weight: torch.Tensor,
179 + rank: int,
180 + hc: int,
181 +) -> tuple[torch.Tensor, torch.Tensor]:
182 + """Fused Qwen4Exp HC down projection + SiLU.
183 +
184 + Args:
185 + x: Normalized hyper-hidden input, [M, K].
186 + weight: Merged down+inject weight, [N, K].
187 + rank: Number of low-rank output columns.
188 + hc: Number of injection-logit output columns.
189 +
190 + Returns:
191 + LoRA activations [M, rank] and injection logits [M, hc].
192 +
193 + """
194 + kernel = _get_kernel(rank, hc, weight.shape[1], x.shape[0] == 1)
195 + output = kernel(x, weight)
196 + return output[:, :rank], output[:, rank : rank + hc]
197 +
198 +
199 +def request_hc_down_silu_warmup(
200 + m_values: Iterable[int], rank: int, hc: int, k: int
201 +) -> None:
202 + """Precompile the fused kernels for CUDA-graph capture token counts.
203 +
204 + Args:
205 + m_values: Token counts the model may capture CUDA graphs for; values
206 + outside the fused dispatch range are ignored.
207 + rank: Number of low-rank output columns.
208 + hc: Number of injection-logit output columns.
209 + k: Input feature size.
210 +
211 + """
212 + if not current_platform.has_device_capability(90) or k % 8 != 0:
213 + return
214 + m_set = {int(m) for m in m_values if 1 <= m <= MAX_FUSED_M}
215 + # M=1 dispatches to the weight-prefetching PDL variant.
216 + _get_kernel(rank, hc, k, True).request_warmup(m_set & {1})
217 + _get_kernel(rank, hc, k).request_warmup(m_set - {1})
218 +
219 +
220 +__all__ = [
221 + "MAX_FUSED_M",
222 + "hc_down_silu",
223 + "request_hc_down_silu_warmup",
224 +]