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, 19:00 UTC
[Bugfix] [ROCm] [DSV4] [Perf] Add aiter mhc support (#41946)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
TJian committed
May 13, 2026 at 21:43 UTC
a8887c208f34c04c3b021cf3949ed6545d77bb01
12 files changed
+1920
-1033
requirements/rocm.txt
-3
@@ -21,6 +21,3 @@ timm>=1.0.17
21
# amd-quark: required for Quark quantization on ROCm
22
# To be consistent with test_quark.py
23
amd-quark>=0.8.99
24
-# tilelang has to be installed for mhc module to be
25
-# imported correctly.
26
-tilelang==0.1.9
tests/kernels/test_mhc_kernels.py
+58
-2
@@ -3,7 +3,7 @@
3
import pytest
4
import torch
5
6
-import vllm.model_executor.layers.mhc as mhc_ops # noqa: F401
6
+import vllm.model_executor.kernels.mhc # noqa: F401
7
from vllm.platforms import current_platform
8
from vllm.utils.torch_utils import set_random_seed
9
@@ -74,6 +74,23 @@ def mhc_post_ref(
74
return (x.float().unsqueeze(-2) * post_layer_mix + term2).bfloat16()
75
76
77
+def hc_head_ref(
78
+ residual: torch.Tensor,
79
+ fn: torch.Tensor,
80
+ hc_scale: torch.Tensor,
81
+ hc_base: torch.Tensor,
82
+ rms_eps: float,
83
+ hc_eps: float,
84
+) -> torch.Tensor:
85
+ residual_flat = residual.flatten(-2).float()
86
+ residual_norm = residual_flat * torch.rsqrt(
87
+ residual_flat.square().mean(dim=-1, keepdim=True) + rms_eps
88
+ )
89
+ pre_mix = torch.nn.functional.linear(residual_norm, fn)
90
+ pre_mix = torch.sigmoid(pre_mix * hc_scale + hc_base) + hc_eps
91
+ return torch.sum(pre_mix.unsqueeze(-1) * residual.float(), dim=-2).bfloat16()
92
+
93
+
94
@pytest.mark.skipif(
95
not current_platform.is_cuda(),
96
reason="CUDA required",
@@ -121,7 +138,7 @@ def test_mhc_fused_post_pre(num_tokens, hidden_size, hc_mult):
138
139
residual_ref, post_mix_ref, res_mix_ref, layer_input_ref = run_ref()
140
124
- residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre(
141
+ residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre_tilelang(
142
x,
143
residual,
144
post_layer_mix,
@@ -140,3 +157,42 @@ def test_mhc_fused_post_pre(num_tokens, hidden_size, hc_mult):
157
torch.testing.assert_close(post_mix, post_mix_ref, atol=1e-2, rtol=1e-2)
158
torch.testing.assert_close(res_mix, res_mix_ref, atol=1e-2, rtol=1e-2)
159
torch.testing.assert_close(x, layer_input_ref, atol=1e-2, rtol=1e-2)
160
+
161
+
162
+@pytest.mark.skipif(
163
+ not current_platform.is_rocm(),
164
+ reason="ROCm required",
165
+)
166
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
167
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
168
+@pytest.mark.parametrize("hc_mult", [4])
169
+def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
170
+ torch.set_default_device(DEVICE)
171
+ set_random_seed(0)
172
+
173
+ residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
174
+ fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
175
+ hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
176
+ hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
177
+ rms_eps = hc_eps = 1e-6
178
+
179
+ out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)
180
+ out.fill_(float("nan"))
181
+
182
+ result = torch.ops.vllm.hc_head_triton(
183
+ residual,
184
+ fn,
185
+ hc_scale,
186
+ hc_base,
187
+ out,
188
+ hidden_size,
189
+ rms_eps,
190
+ hc_eps,
191
+ hc_mult,
192
+ )
193
+
194
+ assert result is None
195
+ assert not torch.isnan(out).any()
196
+
197
+ out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
198
+ torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
vllm/_aiter_ops.py
+178
@@ -2395,5 +2395,183 @@ class rocm_aiter_ops:
2395
kv_cache_dtype=kv_cache_dtype,
2396
)
2397
2398
+ @staticmethod
2399
+ def mhc_pre(
2400
+ residual: torch.Tensor,
2401
+ fn: torch.Tensor,
2402
+ hc_scale: torch.Tensor,
2403
+ hc_base: torch.Tensor,
2404
+ rms_eps: float,
2405
+ hc_pre_eps: float,
2406
+ hc_sinkhorn_eps: float,
2407
+ hc_post_mult_value: float,
2408
+ sinkhorn_repeat: int,
2409
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
2410
+ """
2411
+ Forward pass for mHC pre block.
2412
+
2413
+ Args:
2414
+ residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
2415
+ fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
2416
+ hc_scale: shape (3,), dtype torch.float32
2417
+ hc_base: shape (hc_mult3,), dtype torch.float32
2418
+ rms_eps: RMS normalization epsilon
2419
+ hc_pre_eps: pre-mix epsilon
2420
+ hc_sinkhorn_eps: sinkhorn epsilon
2421
+ hc_post_mult_value: post-mix multiplier value
2422
+ sinkhorn_repeat: number of sinkhorn iterations
2423
+ n_splits: split-k factor;
2424
+
2425
+ Returns:
2426
+ post_mix: shape (..., hc_mult), dtype torch.float32
2427
+ comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
2428
+ layer_input: shape (..., hidden_size), dtype torch.bfloat16
2429
+ """
2430
+ from aiter.ops.mhc import mhc_pre
2431
+
2432
+ # Validate shapes
2433
+ assert residual.dtype == torch.bfloat16
2434
+ assert fn.dtype == torch.float32
2435
+ assert hc_scale.dtype == torch.float32
2436
+ assert hc_base.dtype == torch.float32
2437
+
2438
+ hc_mult = residual.shape[-2]
2439
+ hidden_size = residual.shape[-1]
2440
+ hc_mult2 = hc_mult * hc_mult
2441
+ hc_mult3 = hc_mult * 2 + hc_mult2
2442
+
2443
+ hc_hidden_size = hc_mult * hidden_size
2444
+ assert fn.shape[0] == hc_mult3
2445
+ assert fn.shape[1] == hc_hidden_size
2446
+ assert hc_scale.shape == (3,)
2447
+ assert hc_base.shape == (hc_mult3,)
2448
+
2449
+ outer_shape = residual.shape[:-2]
2450
+
2451
+ residual_flat = residual.view(-1, hc_mult, hidden_size)
2452
+
2453
+ num_tokens = residual_flat.shape[0]
2454
+ if num_tokens == 0:
2455
+ return (
2456
+ torch.empty(
2457
+ num_tokens,
2458
+ hc_mult,
2459
+ 1,
2460
+ dtype=torch.float32,
2461
+ device=residual_flat.device,
2462
+ ),
2463
+ torch.empty(
2464
+ num_tokens,
2465
+ hc_mult,
2466
+ hc_mult,
2467
+ dtype=torch.float32,
2468
+ device=residual_flat.device,
2469
+ ),
2470
+ torch.empty(
2471
+ num_tokens,
2472
+ hidden_size,
2473
+ dtype=torch.bfloat16,
2474
+ device=residual_flat.device,
2475
+ ),
2476
+ )
2477
+
2478
+ # AITER's Python wrapper allocates intermediate/output tensors without
2479
+ # explicit device arguments, so run it under the residual tensor's device.
2480
+ with torch.device(residual_flat.device):
2481
+ post_mix, comb_mix, layer_input = mhc_pre(
2482
+ residual_flat,
2483
+ fn,
2484
+ hc_scale,
2485
+ hc_base,
2486
+ rms_eps,
2487
+ hc_pre_eps,
2488
+ hc_sinkhorn_eps,
2489
+ hc_post_mult_value,
2490
+ sinkhorn_repeat,
2491
+ )
2492
+ return (
2493
+ post_mix.view(*outer_shape, hc_mult, 1),
2494
+ comb_mix.view(*outer_shape, hc_mult, hc_mult),
2495
+ layer_input.view(*outer_shape, hidden_size),
2496
+ )
2497
+
2498
+ @staticmethod
2499
+ def hc_head(
2500
+ hs_flat: torch.Tensor,
2501
+ fn: torch.Tensor,
2502
+ hc_scale: torch.Tensor,
2503
+ hc_base: torch.Tensor,
2504
+ out: torch.Tensor,
2505
+ hidden_size: int,
2506
+ rms_eps: float,
2507
+ hc_eps: float,
2508
+ hc_mult: int,
2509
+ ) -> None:
2510
+ """Run hc_head through AITER mhc_pre and write the result to out."""
2511
+ assert hs_flat.dtype == torch.bfloat16
2512
+ assert fn.dtype == torch.float32
2513
+ assert hc_scale.dtype == torch.float32
2514
+ assert hc_base.dtype == torch.float32
2515
+ assert hs_flat.shape[-2:] == (hc_mult, hidden_size)
2516
+ assert fn.shape == (hc_mult, hc_mult * hidden_size)
2517
+ assert hc_scale.shape == (1,)
2518
+ assert hc_base.shape == (hc_mult,)
2519
+
2520
+ num_tokens = hs_flat.shape[0]
2521
+ if num_tokens == 0:
2522
+ return
2523
+
2524
+ hc_mult3 = hc_mult * 2 + hc_mult * hc_mult
2525
+
2526
+ full_fn = torch.zeros(
2527
+ hc_mult3,
2528
+ hc_mult * hidden_size,
2529
+ dtype=fn.dtype,
2530
+ device=fn.device,
2531
+ )
2532
+ full_fn[:hc_mult] = fn
2533
+
2534
+ full_base = torch.zeros(hc_mult3, dtype=hc_base.dtype, device=hc_base.device)
2535
+ full_base[:hc_mult] = hc_base
2536
+
2537
+ full_scale = torch.zeros(3, dtype=hc_scale.dtype, device=hc_scale.device)
2538
+ full_scale[0] = hc_scale[0]
2539
+
2540
+ _, _, layer_input = rocm_aiter_ops.mhc_pre(
2541
+ hs_flat,
2542
+ full_fn,
2543
+ full_scale,
2544
+ full_base,
2545
+ rms_eps,
2546
+ hc_eps,
2547
+ 0.0,
2548
+ 1.0,
2549
+ 0,
2550
+ )
2551
+ out.copy_(layer_input)
2552
+
2553
+ @staticmethod
2554
+ def mhc_post(
2555
+ x: torch.Tensor,
2556
+ residual: torch.Tensor,
2557
+ post_layer_mix: torch.Tensor,
2558
+ comb_res_mix: torch.Tensor,
2559
+ ) -> torch.Tensor:
2560
+ from aiter.ops.mhc import mhc_post
2561
+
2562
+ hc_mult = residual.shape[-2]
2563
+ hidden_size = residual.shape[-1]
2564
+ residual_flat = residual.view(-1, hc_mult, hidden_size)
2565
+ num_tokens = residual_flat.shape[0]
2566
+ out = torch.empty_like(residual_flat)
2567
+ mhc_post(
2568
+ out,
2569
+ x.view(num_tokens, hidden_size),
2570
+ residual_flat,
2571
+ post_layer_mix.view(num_tokens, hc_mult, 1),
2572
+ comb_res_mix.view(num_tokens, hc_mult, hc_mult),
2573
+ )
2574
+ return out.view_as(residual)
2575
+
2576
2577
rocm_aiter_ops.register_ops_once()
vllm/_tilelang_ops.py
new
+462
@@ -0,0 +1,462 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
+from functools import cache
5
+from typing import TYPE_CHECKING
6
+
7
+import torch
8
+
9
+from vllm.platforms import current_platform
10
+from vllm.utils.import_utils import has_tilelang
11
+from vllm.utils.math_utils import cdiv
12
+
13
+# tilelang is only available on CUDA platforms
14
+if TYPE_CHECKING or current_platform.is_cuda():
15
+ if not has_tilelang():
16
+ raise ImportError(
17
+ "tilelang is required for mhc but is not installed. Install it with "
18
+ "`pip install tilelang`."
19
+ )
20
+ import tilelang
21
+ import tilelang.language as T
22
+else:
23
+ tilelang = None # type: ignore[assignment]
24
+ T = None # type: ignore[assignment]
25
+
26
+
27
+@cache
28
+def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
29
+ device_props = torch.cuda.get_device_properties(0)
30
+ n_sms = device_props.multi_processor_count
31
+ split_k = n_sms // grid_size
32
+ if k is not None:
33
+ # avoid split_k for small k
34
+ num_block_k = cdiv(k, block_k)
35
+ split_k = min(split_k, num_block_k // 4)
36
+ split_k = max(split_k, 1)
37
+ return split_k
38
+
39
+
40
+@tilelang.jit(
41
+ pass_configs={
42
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
43
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
44
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
45
+ },
46
+)
47
+def mhc_pre_big_fuse_tilelang(
48
+ gemm_out_mul,
49
+ gemm_out_sqrsum,
50
+ hc_scale,
51
+ hc_base,
52
+ residual,
53
+ post_mix,
54
+ comb_mix,
55
+ layer_input,
56
+ hidden_size: int,
57
+ rms_eps: float,
58
+ hc_pre_eps: float,
59
+ hc_sinkhorn_eps: float,
60
+ hc_post_mult_value: float,
61
+ sinkhorn_repeat: int,
62
+ n_splits: int = 16,
63
+ hc_mult: int = 4,
64
+):
65
+ """Deeply fused kernels, everything other than gemm & sqrsum in mHC pre block."""
66
+ num_tokens = T.dynamic("num_tokens")
67
+ hc_mult3 = hc_mult * (2 + hc_mult)
68
+ hidden_block = math.gcd(512, hidden_size)
69
+
70
+ gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
71
+ gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type]
72
+ hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type]
73
+ hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
74
+ residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
75
+ # outputs
76
+ post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
77
+ comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type]
78
+ layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
79
+
80
+ with T.Kernel(num_tokens, threads=96) as i:
81
+ T.pdl_sync()
82
+ ##################################################################
83
+ # _pre_norm_fn_fwd_norm
84
+ rms = T.alloc_fragment(1, T.float32)
85
+ mixes = T.alloc_fragment(hc_mult3, T.float32)
86
+ T.clear(mixes)
87
+ rms[0] = 0
88
+ for i_split in T.serial(n_splits):
89
+ rms[0] += gemm_out_sqrsum[i_split, i]
90
+ rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
91
+ for j in T.Parallel(hc_mult3):
92
+ mixes[j] = 0
93
+ for i_split in T.serial(n_splits):
94
+ mixes[j] += gemm_out_mul[i_split, i, j]
95
+ mixes[j] *= rms[0]
96
+ mixes_shared = T.alloc_shared(hc_mult3, T.float32)
97
+ T.copy(mixes, mixes_shared)
98
+
99
+ if T.get_thread_binding() < 32:
100
+ ##################################################################
101
+ # _pre_split_mixes_fwd (post & comb)
102
+ cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
103
+ for j in T.Parallel(hc_mult):
104
+ post_mix[i, j] = (
105
+ T.sigmoid(
106
+ mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
107
+ )
108
+ * hc_post_mult_value
109
+ )
110
+ for j, k in T.Parallel(hc_mult, hc_mult):
111
+ cm[j, k] = (
112
+ mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
113
+ + hc_base[j * hc_mult + k + hc_mult * 2]
114
+ )
115
+
116
+ ##################################################################
117
+ # _sinkhorn_fwd
118
+ row_sum = T.alloc_fragment(hc_mult, T.float32)
119
+ col_sum = T.alloc_fragment(hc_mult, T.float32)
120
+
121
+ # comb = comb.softmax(-1) + eps
122
+ row_max = T.alloc_fragment(hc_mult, T.float32)
123
+ T.reduce_max(cm, row_max, dim=1)
124
+ for j, k in T.Parallel(hc_mult, hc_mult):
125
+ cm[j, k] = T.exp(cm[j, k] - row_max[j])
126
+ T.reduce_sum(cm, row_sum, dim=1)
127
+ for j, k in T.Parallel(hc_mult, hc_mult):
128
+ cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
129
+
130
+ # comb = comb / (comb.sum(-2) + eps)
131
+ T.reduce_sum(cm, col_sum, dim=0)
132
+ for j, k in T.Parallel(hc_mult, hc_mult):
133
+ cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
134
+
135
+ for _ in T.serial(sinkhorn_repeat - 1):
136
+ # comb = comb / (comb.sum(-1) + eps)
137
+ T.reduce_sum(cm, row_sum, dim=1)
138
+ for j, k in T.Parallel(hc_mult, hc_mult):
139
+ cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
140
+
141
+ # comb = comb / (comb.sum(-2) + eps)
142
+ T.reduce_sum(cm, col_sum, dim=0)
143
+ for j, k in T.Parallel(hc_mult, hc_mult):
144
+ cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
145
+
146
+ # save comb_mix to global memory
147
+ for j, k in T.Parallel(hc_mult, hc_mult):
148
+ comb_mix[i, j * hc_mult + k] = cm[j, k]
149
+ else:
150
+ ##################################################################
151
+ # _pre_split_mixes_fwd (pre)
152
+ pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
153
+ for j in T.Parallel(hc_mult):
154
+ pre_mix_shared[j] = (
155
+ T.sigmoid(
156
+ mixes_shared[j] * hc_scale[0] + hc_base[j],
157
+ )
158
+ + hc_pre_eps
159
+ )
160
+ ###################################################################
161
+ # _pre_apply_mix_fwd
162
+ for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
163
+ xs = T.alloc_shared((hc_mult, hidden_block), T.float32)
164
+ xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
165
+ T.copy(residual[i, 0, i0_h * hidden_block], xs)
166
+ T.copy(xs, xl)
167
+
168
+ ol = T.alloc_fragment(hidden_block, T.float32)
169
+ T.clear(ol)
170
+
171
+ for i_hc in T.serial(hc_mult):
172
+ pre = pre_mix_shared[i_hc]
173
+ for i1_h in T.Parallel(hidden_block):
174
+ ol[i1_h] += pre * xl[i_hc, i1_h]
175
+
176
+ T.copy(ol, layer_input[i, i0_h * hidden_block])
177
+ T.pdl_trigger()
178
+
179
+
180
+@tilelang.jit(
181
+ pass_configs={
182
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
183
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
184
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
185
+ },
186
+)
187
+def mhc_fused_tilelang(
188
+ comb_mix,
189
+ residual_in,
190
+ post_mix,
191
+ x_in,
192
+ weight_t,
193
+ yp_out,
194
+ rp_out,
195
+ residual_out,
196
+ hc: int,
197
+ hidden: int,
198
+ n_out: int,
199
+ n_thr: int = 256,
200
+ h_blk: int = 256,
201
+ tile_n: int = 1,
202
+ split_k: int = 1,
203
+) -> tilelang.JITKernel:
204
+ """Fused mhc post-mapping + pre-norm GEMM FMA"""
205
+ m = T.dynamic("num_tokens")
206
+ split_k = T.dynamic("split_k")
207
+ h = hidden
208
+ h_blk = math.gcd(hidden, h_blk)
209
+ h_per_split = h // split_k
210
+ n_tiles = n_out // tile_n
211
+
212
+ comb_mix: T.Tensor((m, hc, hc), T.float32) # type: ignore[no-redef, valid-type]
213
+ residual_in: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
214
+ post_mix: T.Tensor((m, hc), T.float32) # type: ignore[no-redef, valid-type]
215
+ x_in: T.Tensor((m, h), T.bfloat16) # type: ignore[no-redef, valid-type]
216
+ weight_t: T.Tensor((n_out, hc, h), T.float32) # type: ignore[no-redef, valid-type]
217
+ yp_out: T.Tensor((split_k, m, n_out), T.float32) # type: ignore[no-redef, valid-type]
218
+ rp_out: T.Tensor((split_k, m), T.float32) # type: ignore[no-redef, valid-type]
219
+ residual_out: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
220
+
221
+ h_iters = h_per_split // n_thr
222
+ num_warps = n_thr // 32
223
+
224
+ with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
225
+ tid = T.get_thread_binding()
226
+ warp_id = T.get_warp_idx()
227
+ lane = T.get_lane_idx()
228
+
229
+ s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
230
+ s_post = T.alloc_shared((hc,), T.float32)
231
+ s_comb = T.alloc_shared((hc, hc), T.float32)
232
+
233
+ pm = T.alloc_local((hc,), T.float32)
234
+ cm = T.alloc_local((hc, hc), T.float32)
235
+ acc = T.alloc_local((tile_n,), T.float32)
236
+ sqr = T.alloc_local((1,), T.float32)
237
+ new_r = T.alloc_local((hc,), T.float32)
238
+
239
+ T.clear(acc)
240
+ T.clear(sqr)
241
+ h_split_start = i_ks * h_per_split
242
+
243
+ T.pdl_sync()
244
+
245
+ T.copy(post_mix[i_n, 0], s_post)
246
+ T.copy(comb_mix[i_n, 0, 0], s_comb)
247
+
248
+ for j in T.unroll(hc):
249
+ pm[j] = s_post[j]
250
+ for j in T.unroll(hc):
251
+ for k in T.unroll(hc):
252
+ cm[k, j] = s_comb[k, j]
253
+
254
+ # Each thread owns h_iters elements of the k-split's h slice.
255
+ for it in T.serial(h_iters):
256
+ h_idx = h_split_start + it * n_thr + tid
257
+
258
+ # Compute new residual from layer output and past residual
259
+ for j in T.unroll(hc):
260
+ new_r[j] = pm[j] * x_in[i_n, h_idx]
261
+ for k in T.unroll(hc):
262
+ new_r[j] += cm[k, j] * residual_in[i_n, k, h_idx]
263
+
264
+ # populate residual_out and compute sqr sum
265
+ if i_nt == 0:
266
+ for j in T.unroll(hc):
267
+ residual_out[i_n, j, h_idx] = new_r[j]
268
+ sqr[0] += new_r[j] * new_r[j]
269
+
270
+ # Per-thread FMA into acc[n]
271
+ for n in T.unroll(tile_n):
272
+ for j in T.unroll(hc):
273
+ acc[n] += weight_t[i_nt * tile_n + n, j, h_idx] * new_r[j]
274
+
275
+ for n in T.unroll(tile_n):
276
+ acc[n] = T.warp_reduce_sum(acc[n])
277
+ if i_nt == 0:
278
+ sqr[0] = T.warp_reduce_sum(sqr[0])
279
+
280
+ # Cross-warp reduce via shared mem
281
+ if lane == 0:
282
+ for n in T.unroll(tile_n):
283
+ s_warp[warp_id, n] = acc[n]
284
+ if i_nt == 0:
285
+ s_warp[warp_id, tile_n] = sqr[0]
286
+ T.sync_threads()
287
+
288
+ # Warp 0 does the final cross-warp sum and writes outputs
289
+ if warp_id == 0:
290
+ if lane < tile_n:
291
+ v = T.alloc_var(T.float32, init=0.0)
292
+ for w in T.unroll(num_warps):
293
+ v += s_warp[w, lane]
294
+ yp_out[i_ks, i_n, i_nt * tile_n + lane] = v
295
+
296
+ if i_nt == 0 and lane == 0:
297
+ v2 = T.alloc_var(T.float32, init=0.0)
298
+ for w in T.unroll(num_warps):
299
+ v2 += s_warp[w, tile_n]
300
+ rp_out[i_ks, i_n] = v2
301
+
302
+ T.pdl_trigger()
303
+
304
+
305
+@tilelang.jit(
306
+ pass_configs={
307
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
308
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
309
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
310
+ },
311
+)
312
+def mhc_post_tilelang(
313
+ a,
314
+ b,
315
+ c,
316
+ d,
317
+ x,
318
+ hc: int,
319
+ hidden: int,
320
+ n_thr: int = 128,
321
+ h_blk: int = 1024,
322
+) -> tilelang.JITKernel:
323
+ # rename for shorter code
324
+ n = T.dynamic("num_tokens")
325
+ h = hidden
326
+
327
+ h_blk = math.gcd(hidden, h_blk)
328
+ a: T.Tensor((n, hc, hc), T.float32) # type: ignore[no-redef, valid-type]
329
+ b: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
330
+ c: T.Tensor((n, hc), T.float32) # type: ignore[no-redef, valid-type]
331
+ d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type]
332
+ x: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
333
+ with T.Kernel(n, threads=n_thr) as i_n:
334
+ x_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
335
+ b_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
336
+ d_shared = T.alloc_shared(h_blk, T.bfloat16)
337
+
338
+ x_local = T.alloc_fragment((hc, h_blk), T.float32)
339
+ b_local = T.alloc_fragment((hc, h_blk), T.float32)
340
+ d_local = T.alloc_fragment(h_blk, T.float32)
341
+
342
+ a_local = T.alloc_fragment((hc, hc), T.float32)
343
+ c_local = T.alloc_fragment(hc, T.float32)
344
+ T.pdl_sync()
345
+ T.copy(a[i_n, 0, 0], a_local)
346
+ T.copy(c[i_n, 0], c_local)
347
+
348
+ for i0_h in T.Pipelined(T.ceildiv(h, h_blk), num_stages=2):
349
+ T.copy(b[i_n, 0, i0_h * h_blk], b_shared)
350
+ T.copy(d[i_n, i0_h * h_blk], d_shared)
351
+
352
+ T.copy(b_shared, b_local)
353
+ T.copy(d_shared, d_local)
354
+ for i_hco, i1_h in T.Parallel(hc, h_blk):
355
+ x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h]
356
+ for i_hci in T.serial(hc):
357
+ x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
358
+ T.copy(x_local, x_shared)
359
+
360
+ T.copy(x_shared, x[i_n, 0, i0_h * h_blk])
361
+ T.pdl_trigger()
362
+
363
+
364
+@tilelang.jit(
365
+ pass_configs={
366
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
367
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
368
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
369
+ },
370
+)
371
+def hc_head_fuse_tilelang(
372
+ residual,
373
+ fn,
374
+ hc_scale,
375
+ hc_base,
376
+ out,
377
+ hidden_size: int,
378
+ rms_eps: float,
379
+ hc_eps: float,
380
+ hc_mult: int = 4,
381
+ n_thr: int = 128,
382
+ h_blk: int = 1024,
383
+):
384
+ """Two-pass fused kernel for hc_head.
385
+
386
+ Pass 1: accumulate per-token squared sum and hc_mult dot-products
387
+ (projections onto fn rows) using cross-thread reducers.
388
+ Pass 2: apply sigmoid-gated weighted sum of residual channels to output.
389
+
390
+ Avoids materialising mixes / rsqrt / pre tensors to global memory.
391
+ """
392
+ num_tokens = T.dynamic("num_tokens")
393
+ hc_dim = hc_mult * hidden_size
394
+ h_block = math.gcd(h_blk, hidden_size)
395
+ n_h = hidden_size // h_block
396
+
397
+ residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]
398
+ fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type]
399
+ hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type]
400
+ hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type]
401
+ out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]
402
+
403
+ with T.Kernel(num_tokens, threads=n_thr) as i:
404
+ T.pdl_sync()
405
+
406
+ # ------------------------------------------------------------------
407
+ # Pass 1 – for each residual channel m_c and h_block:
408
+ # • accumulate squared sum (for RMS norm denominator)
409
+ # • accumulate hc_mult dot-products with fn rows
410
+ # ------------------------------------------------------------------
411
+ sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all")
412
+ mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all")
413
+ T.fill(sqrsum_r, 0.0)
414
+ T.fill(mixes_r, 0.0)
415
+
416
+ for m_c in T.serial(hc_mult):
417
+ for i_h in T.serial(n_h):
418
+ x_local = T.alloc_fragment(h_block, T.float32)
419
+ T.copy(residual[i, m_c, i_h * h_block], x_local)
420
+
421
+ for k in T.Parallel(h_block):
422
+ sqrsum_r[0] += x_local[k] * x_local[k]
423
+
424
+ for m_m in T.unroll(hc_mult):
425
+ fn_local = T.alloc_fragment(h_block, T.float32)
426
+ T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local)
427
+ for k in T.Parallel(h_block):
428
+ mixes_r[m_m] += x_local[k] * fn_local[k]
429
+
430
+ T.finalize_reducer(sqrsum_r)
431
+ T.finalize_reducer(mixes_r)
432
+
433
+ # ------------------------------------------------------------------
434
+ # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps
435
+ # ------------------------------------------------------------------
436
+ pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
437
+ rsqrt_val = T.alloc_fragment(1, T.float32)
438
+ rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps)
439
+ for m in T.Parallel(hc_mult):
440
+ pre_mix_shared[m] = (
441
+ T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps
442
+ )
443
+
444
+ # ------------------------------------------------------------------
445
+ # Pass 2 – apply_mix: pipelined weighted sum over residual channels
446
+ # ------------------------------------------------------------------
447
+ for i0_h in T.Pipelined(n_h, num_stages=2):
448
+ xs = T.alloc_shared((hc_mult, h_block), T.bfloat16)
449
+ xl = T.alloc_fragment((hc_mult, h_block), T.float32)
450
+ T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True)
451
+ T.copy(xs, xl)
452
+
453
+ ol = T.alloc_fragment(h_block, T.float32)
454
+ T.clear(ol)
455
+ for i_hc in T.serial(hc_mult):
456
+ pre = pre_mix_shared[i_hc]
457
+ for i1_h in T.Parallel(h_block):
458
+ ol[i1_h] += pre * xl[i_hc, i1_h]
459
+
460
+ T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
461
+
462
+ T.pdl_trigger()
vllm/model_executor/kernels/mhc/__init__.py
new
+29
@@ -0,0 +1,29 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from .aiter import *
4
+from .tilelang import *
5
+from .torch import *
6
+from .triton import *
7
+
8
+__all__ = [
9
+ "mhc_pre_cuda",
10
+ "mhc_post_cuda",
11
+ "mhc_fused_post_pre_cuda",
12
+ "hc_head_fused_kernel_cuda",
13
+ "mhc_pre_aiter",
14
+ "mhc_post_aiter",
15
+ "mhc_fused_post_pre_aiter",
16
+ "hc_head_fused_aiter",
17
+ "mhc_pre_tilelang",
18
+ "mhc_post_tilelang",
19
+ "mhc_fused_post_pre_tilelang",
20
+ "hc_head_fused_tilelang",
21
+ "mhc_pre_torch",
22
+ "mhc_post_torch",
23
+ "mhc_fused_post_pre_torch",
24
+ "hc_head_fused_torch",
25
+ "mhc_pre_triton",
26
+ "mhc_post_triton",
27
+ "mhc_fused_post_pre_triton",
28
+ "hc_head_fused_triton",
29
+]
vllm/model_executor/kernels/mhc/aiter.py
new
+138
@@ -0,0 +1,138 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.utils.torch_utils import direct_register_custom_op
6
+
7
+
8
+def mhc_pre_aiter(
9
+ residual: torch.Tensor,
10
+ fn: torch.Tensor,
11
+ hc_scale: torch.Tensor,
12
+ hc_base: torch.Tensor,
13
+ rms_eps: float,
14
+ hc_pre_eps: float,
15
+ hc_sinkhorn_eps: float,
16
+ hc_post_mult_value: float,
17
+ sinkhorn_repeat: int,
18
+ n_splits: int = 1,
19
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
20
+ """
21
+ Forward pass for mHC pre block.
22
+
23
+ Args:
24
+ residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
25
+ fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
26
+ hc_scale: shape (3,), dtype torch.float32
27
+ hc_base: shape (hc_mult3,), dtype torch.float32
28
+ rms_eps: RMS normalization epsilon
29
+ hc_pre_eps: pre-mix epsilon
30
+ hc_sinkhorn_eps: sinkhorn epsilon
31
+ hc_post_mult_value: post-mix multiplier value
32
+ sinkhorn_repeat: number of sinkhorn iterations
33
+ n_splits: split-k factor;
34
+
35
+ Returns:
36
+ post_mix: shape (..., hc_mult), dtype torch.float32
37
+ comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
38
+ layer_input: shape (..., hidden_size), dtype torch.bfloat16
39
+ """
40
+
41
+ hidden_size = residual.shape[-1]
42
+ assert hidden_size % 256 == 0
43
+ from vllm._aiter_ops import rocm_aiter_ops
44
+
45
+ return rocm_aiter_ops.mhc_pre(
46
+ residual,
47
+ fn,
48
+ hc_scale,
49
+ hc_base,
50
+ rms_eps,
51
+ hc_pre_eps,
52
+ hc_sinkhorn_eps,
53
+ hc_post_mult_value,
54
+ sinkhorn_repeat,
55
+ )
56
+
57
+
58
+def _mhc_pre_aiter_fake(
59
+ residual: torch.Tensor,
60
+ fn: torch.Tensor,
61
+ hc_scale: torch.Tensor,
62
+ hc_base: torch.Tensor,
63
+ rms_eps: float,
64
+ hc_pre_eps: float,
65
+ hc_sinkhorn_eps: float,
66
+ hc_post_mult_value: float,
67
+ sinkhorn_repeat: int,
68
+ n_splits: int = 1,
69
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
70
+ hc_mult = residual.shape[-2]
71
+ hidden_size = residual.shape[-1]
72
+ outer_shape = residual.shape[:-2]
73
+
74
+ # Create empty tensors with correct shapes for meta device / shape inference
75
+ post_mix = torch.empty(
76
+ *outer_shape,
77
+ hc_mult,
78
+ 1,
79
+ dtype=torch.float32,
80
+ device=residual.device,
81
+ )
82
+ comb_mix = torch.empty(
83
+ *outer_shape,
84
+ hc_mult,
85
+ hc_mult,
86
+ dtype=torch.float32,
87
+ device=residual.device,
88
+ )
89
+ layer_input = torch.empty(
90
+ *outer_shape,
91
+ hidden_size,
92
+ dtype=torch.bfloat16,
93
+ device=residual.device,
94
+ )
95
+
96
+ return post_mix, comb_mix, layer_input
97
+
98
+
99
+def mhc_post_aiter(
100
+ x: torch.Tensor,
101
+ residual: torch.Tensor,
102
+ post_layer_mix: torch.Tensor,
103
+ comb_res_mix: torch.Tensor,
104
+) -> torch.Tensor:
105
+ hidden_size = residual.shape[-1]
106
+
107
+ assert hidden_size % 256 == 0
108
+ from vllm._aiter_ops import rocm_aiter_ops
109
+
110
+ return rocm_aiter_ops.mhc_post(
111
+ x,
112
+ residual,
113
+ post_layer_mix,
114
+ comb_res_mix,
115
+ )
116
+
117
+
118
+def _mhc_post_aiter_fake(
119
+ x: torch.Tensor,
120
+ residual: torch.Tensor,
121
+ post_layer_mix: torch.Tensor,
122
+ comb_res_mix: torch.Tensor,
123
+) -> torch.Tensor:
124
+ return torch.empty_like(residual)
125
+
126
+
127
+direct_register_custom_op(
128
+ op_name="mhc_pre_aiter",
129
+ op_func=mhc_pre_aiter,
130
+ mutates_args=[],
131
+ fake_impl=_mhc_pre_aiter_fake,
132
+)
133
+direct_register_custom_op(
134
+ op_name="mhc_post_aiter",
135
+ op_func=mhc_post_aiter,
136
+ mutates_args=[],
137
+ fake_impl=_mhc_post_aiter_fake,
138
+)
vllm/model_executor/kernels/mhc/tilelang.py
new
+468
@@ -0,0 +1,468 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.utils.torch_utils import direct_register_custom_op
6
+
7
+
8
+def mhc_pre_tilelang(
9
+ residual: torch.Tensor,
10
+ fn: torch.Tensor,
11
+ hc_scale: torch.Tensor,
12
+ hc_base: torch.Tensor,
13
+ rms_eps: float,
14
+ hc_pre_eps: float,
15
+ hc_sinkhorn_eps: float,
16
+ hc_post_mult_value: float,
17
+ sinkhorn_repeat: int,
18
+ n_splits: int = 1,
19
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
20
+ """
21
+ Forward pass for mHC pre block.
22
+
23
+ Args:
24
+ residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
25
+ fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
26
+ hc_scale: shape (3,), dtype torch.float32
27
+ hc_base: shape (hc_mult3,), dtype torch.float32
28
+ rms_eps: RMS normalization epsilon
29
+ hc_pre_eps: pre-mix epsilon
30
+ hc_sinkhorn_eps: sinkhorn epsilon
31
+ hc_post_mult_value: post-mix multiplier value
32
+ sinkhorn_repeat: number of sinkhorn iterations
33
+ n_splits: split-k factor;
34
+
35
+ Returns:
36
+ post_mix: shape (..., hc_mult), dtype torch.float32
37
+ comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
38
+ layer_input: shape (..., hidden_size), dtype torch.bfloat16
39
+ """
40
+ from vllm._tilelang_ops import (
41
+ compute_num_split,
42
+ mhc_pre_big_fuse_tilelang,
43
+ )
44
+ from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
45
+ from vllm.utils.math_utils import cdiv
46
+
47
+ assert residual.dtype == torch.bfloat16
48
+ assert fn.dtype == torch.float32
49
+ assert hc_scale.dtype == torch.float32
50
+ assert hc_base.dtype == torch.float32
51
+
52
+ hc_mult = residual.shape[-2]
53
+ hidden_size = residual.shape[-1]
54
+ hc_mult2 = hc_mult * hc_mult
55
+ hc_mult3 = hc_mult * 2 + hc_mult2
56
+
57
+ hc_hidden_size = hc_mult * hidden_size
58
+ assert fn.shape[0] == hc_mult3
59
+ assert fn.shape[1] == hc_hidden_size
60
+ assert hc_scale.shape == (3,)
61
+ assert hc_base.shape == (hc_mult3,)
62
+
63
+ outer_shape = residual.shape[:-2]
64
+
65
+ residual_flat = residual.view(-1, hc_mult, hidden_size)
66
+ num_tokens = residual_flat.shape[0]
67
+
68
+ # these numbers are from deepgemm kernel impl
69
+ block_k = 64
70
+ block_m = 64
71
+ n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
72
+
73
+ post_mix = torch.empty(
74
+ num_tokens, hc_mult, dtype=torch.float32, device=residual.device
75
+ )
76
+ comb_mix = torch.empty(
77
+ num_tokens, hc_mult2, dtype=torch.float32, device=residual.device
78
+ )
79
+ layer_input = torch.empty(
80
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
81
+ )
82
+
83
+ gemm_out_mul = torch.empty(
84
+ n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
85
+ )
86
+ gemm_out_sqrsum = torch.empty(
87
+ n_splits, num_tokens, dtype=torch.float32, device=residual.device
88
+ )
89
+
90
+ tf32_hc_prenorm_gemm(
91
+ residual_flat.view(num_tokens, hc_mult * hidden_size),
92
+ fn,
93
+ gemm_out_mul,
94
+ gemm_out_sqrsum,
95
+ n_splits,
96
+ )
97
+
98
+ mhc_pre_big_fuse_tilelang(
99
+ gemm_out_mul,
100
+ gemm_out_sqrsum,
101
+ hc_scale,
102
+ hc_base,
103
+ residual_flat,
104
+ post_mix,
105
+ comb_mix,
106
+ layer_input,
107
+ hidden_size,
108
+ rms_eps,
109
+ hc_pre_eps,
110
+ hc_sinkhorn_eps,
111
+ hc_post_mult_value,
112
+ sinkhorn_repeat,
113
+ n_splits,
114
+ hc_mult,
115
+ )
116
+
117
+ return (
118
+ post_mix.view(*outer_shape, hc_mult, 1),
119
+ comb_mix.view(*outer_shape, hc_mult, hc_mult),
120
+ layer_input.view(*outer_shape, hidden_size),
121
+ )
122
+
123
+
124
+def _mhc_pre_tilelang_fake(
125
+ residual: torch.Tensor,
126
+ fn: torch.Tensor,
127
+ hc_scale: torch.Tensor,
128
+ hc_base: torch.Tensor,
129
+ rms_eps: float,
130
+ hc_pre_eps: float,
131
+ hc_sinkhorn_eps: float,
132
+ hc_post_mult_value: float,
133
+ sinkhorn_repeat: int,
134
+ n_splits: int = 1,
135
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
136
+ hc_mult = residual.shape[-2]
137
+ hidden_size = residual.shape[-1]
138
+ outer_shape = residual.shape[:-2]
139
+
140
+ # Create empty tensors with correct shapes for meta device / shape inference
141
+ post_mix = torch.empty(
142
+ *outer_shape,
143
+ hc_mult,
144
+ 1,
145
+ dtype=torch.float32,
146
+ device=residual.device,
147
+ )
148
+ comb_mix = torch.empty(
149
+ *outer_shape,
150
+ hc_mult,
151
+ hc_mult,
152
+ dtype=torch.float32,
153
+ device=residual.device,
154
+ )
155
+ layer_input = torch.empty(
156
+ *outer_shape,
157
+ hidden_size,
158
+ dtype=torch.bfloat16,
159
+ device=residual.device,
160
+ )
161
+
162
+ return post_mix, comb_mix, layer_input
163
+
164
+
165
+def mhc_post_tilelang(
166
+ x: torch.Tensor,
167
+ residual: torch.Tensor,
168
+ post_layer_mix: torch.Tensor,
169
+ comb_res_mix: torch.Tensor,
170
+) -> torch.Tensor:
171
+ from vllm._tilelang_ops import mhc_post_tilelang as _mhc_post_kernel
172
+
173
+ out = torch.empty_like(residual)
174
+ _mhc_post_kernel(
175
+ comb_res_mix,
176
+ residual,
177
+ post_layer_mix.squeeze(-1),
178
+ x,
179
+ out,
180
+ residual.shape[-2],
181
+ residual.shape[-1],
182
+ )
183
+ return out
184
+
185
+
186
+def mhc_fused_post_pre_tilelang(
187
+ x: torch.Tensor,
188
+ residual: torch.Tensor,
189
+ post_layer_mix: torch.Tensor,
190
+ comb_res_mix: torch.Tensor,
191
+ fn: torch.Tensor,
192
+ hc_scale: torch.Tensor,
193
+ hc_base: torch.Tensor,
194
+ rms_eps: float,
195
+ hc_pre_eps: float,
196
+ hc_sinkhorn_eps: float,
197
+ hc_post_mult_value: float,
198
+ sinkhorn_repeat: int,
199
+ n_splits: int = 1,
200
+ tile_n: int = 1,
201
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
202
+ """
203
+ Run one MHC post block followed by the next MHC pre block.
204
+
205
+ Returns:
206
+ residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
207
+ post_mix_cur: shape (..., hc_mult, 1)
208
+ comb_mix_cur: shape (..., hc_mult, hc_mult)
209
+ layer_input_cur: shape (..., hidden_size)
210
+ """
211
+
212
+ from vllm._tilelang_ops import (
213
+ compute_num_split,
214
+ mhc_fused_tilelang,
215
+ mhc_post_tilelang,
216
+ mhc_pre_big_fuse_tilelang,
217
+ )
218
+ from vllm.utils.math_utils import cdiv
219
+
220
+ assert residual.dtype == torch.bfloat16
221
+ assert x.dtype == torch.bfloat16
222
+ assert post_layer_mix.dtype == torch.float32
223
+ assert comb_res_mix.dtype == torch.float32
224
+ assert fn.dtype == torch.float32
225
+ assert hc_scale.dtype == torch.float32
226
+ assert hc_base.dtype == torch.float32
227
+
228
+ hc_mult = residual.shape[-2]
229
+ hidden_size = residual.shape[-1]
230
+ hc_mult2 = hc_mult * hc_mult
231
+ hc_mult3 = hc_mult * 2 + hc_mult2
232
+ hc_hidden_size = hc_mult * hidden_size
233
+ outer_shape = residual.shape[:-2]
234
+
235
+ assert x.shape == (*outer_shape, hidden_size)
236
+ assert post_layer_mix.shape in (
237
+ (*outer_shape, hc_mult, 1),
238
+ (*outer_shape, hc_mult),
239
+ )
240
+ assert comb_res_mix.shape == (*outer_shape, hc_mult, hc_mult)
241
+ assert fn.shape == (hc_mult3, hc_hidden_size)
242
+ assert hc_scale.shape == (3,)
243
+ assert hc_base.shape == (hc_mult3,)
244
+
245
+ assert n_splits in (1, 2, 4, 8)
246
+ assert hidden_size % n_splits == 0
247
+
248
+ residual_flat = residual.view(-1, hc_mult, hidden_size)
249
+ num_tokens = residual_flat.shape[0]
250
+ x_flat = x.view(num_tokens, hidden_size)
251
+ post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
252
+ comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
253
+
254
+ fma_token_threshold = 16
255
+ if num_tokens <= fma_token_threshold:
256
+ # TODO(gnovack): investigate autotuning these heuristics
257
+ tile_n = 2 if num_tokens < 8 else 3
258
+ n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
259
+ else:
260
+ # these number are from deepgemm kernel impl
261
+ block_k = 64
262
+ block_m = 64
263
+ n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
264
+
265
+ gemm_out_mul = torch.empty(
266
+ n_splits,
267
+ num_tokens,
268
+ hc_mult3,
269
+ dtype=torch.float32,
270
+ device=residual.device,
271
+ )
272
+ gemm_out_sqrsum = torch.empty(
273
+ n_splits,
274
+ num_tokens,
275
+ dtype=torch.float32,
276
+ device=residual.device,
277
+ )
278
+ residual_cur = torch.empty_like(residual_flat)
279
+ post_mix_cur = torch.empty(
280
+ num_tokens,
281
+ hc_mult,
282
+ dtype=torch.float32,
283
+ device=residual.device,
284
+ )
285
+ comb_mix_cur = torch.empty(
286
+ num_tokens,
287
+ hc_mult2,
288
+ dtype=torch.float32,
289
+ device=residual.device,
290
+ )
291
+ layer_input_cur = torch.empty(
292
+ num_tokens,
293
+ hidden_size,
294
+ dtype=torch.bfloat16,
295
+ device=residual.device,
296
+ )
297
+
298
+ if num_tokens <= fma_token_threshold:
299
+ mhc_fused_tilelang(
300
+ comb_res_mix_flat,
301
+ residual_flat,
302
+ post_layer_mix_flat,
303
+ x_flat,
304
+ fn.view(hc_mult3, hc_mult, hidden_size),
305
+ gemm_out_mul,
306
+ gemm_out_sqrsum,
307
+ residual_cur,
308
+ hc_mult,
309
+ hidden_size,
310
+ hc_mult3,
311
+ tile_n=tile_n,
312
+ n_splits=n_splits,
313
+ )
314
+ else:
315
+ mhc_post_tilelang(
316
+ comb_res_mix_flat,
317
+ residual_flat,
318
+ post_layer_mix_flat,
319
+ x_flat,
320
+ residual_cur,
321
+ residual.shape[-2],
322
+ residual.shape[-1],
323
+ )
324
+
325
+ from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
326
+
327
+ tf32_hc_prenorm_gemm(
328
+ residual_cur.view(num_tokens, hc_mult * hidden_size),
329
+ fn,
330
+ gemm_out_mul,
331
+ gemm_out_sqrsum,
332
+ n_splits,
333
+ )
334
+
335
+ mhc_pre_big_fuse_tilelang(
336
+ gemm_out_mul,
337
+ gemm_out_sqrsum,
338
+ hc_scale,
339
+ hc_base,
340
+ residual_cur,
341
+ post_mix_cur,
342
+ comb_mix_cur,
343
+ layer_input_cur,
344
+ hidden_size,
345
+ rms_eps,
346
+ hc_pre_eps,
347
+ hc_sinkhorn_eps,
348
+ hc_post_mult_value,
349
+ sinkhorn_repeat,
350
+ n_splits,
351
+ hc_mult,
352
+ )
353
+
354
+ return (
355
+ residual_cur.view(*outer_shape, hc_mult, hidden_size),
356
+ post_mix_cur.view(*outer_shape, hc_mult, 1),
357
+ comb_mix_cur.view(*outer_shape, hc_mult, hc_mult),
358
+ layer_input_cur.view(*outer_shape, hidden_size),
359
+ )
360
+
361
+
362
+def _mhc_fused_post_pre_tilelang_fake(
363
+ x: torch.Tensor,
364
+ residual: torch.Tensor,
365
+ post_layer_mix: torch.Tensor,
366
+ comb_res_mix: torch.Tensor,
367
+ fn: torch.Tensor,
368
+ hc_scale: torch.Tensor,
369
+ hc_base: torch.Tensor,
370
+ rms_eps: float,
371
+ hc_pre_eps: float,
372
+ hc_sinkhorn_eps: float,
373
+ hc_post_mult_value: float,
374
+ sinkhorn_repeat: int,
375
+ n_splits: int = 1,
376
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
377
+ hc_mult = residual.shape[-2]
378
+ hidden_size = residual.shape[-1]
379
+ outer_shape = residual.shape[:-2]
380
+
381
+ residual_cur = torch.empty_like(residual)
382
+ post_mix_cur = torch.empty(
383
+ *outer_shape,
384
+ hc_mult,
385
+ 1,
386
+ dtype=torch.float32,
387
+ device=residual.device,
388
+ )
389
+ comb_mix_cur = torch.empty(
390
+ *outer_shape,
391
+ hc_mult,
392
+ hc_mult,
393
+ dtype=torch.float32,
394
+ device=residual.device,
395
+ )
396
+ layer_input_cur = torch.empty(
397
+ *outer_shape,
398
+ hidden_size,
399
+ dtype=torch.bfloat16,
400
+ device=residual.device,
401
+ )
402
+
403
+ return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
404
+
405
+
406
+def _mhc_post_tilelang_fake(
407
+ x: torch.Tensor,
408
+ residual: torch.Tensor,
409
+ post_layer_mix: torch.Tensor,
410
+ comb_res_mix: torch.Tensor,
411
+) -> torch.Tensor:
412
+ return torch.empty_like(residual)
413
+
414
+
415
+def _hc_head_fused_kernel_tilelang(
416
+ hs_flat: torch.Tensor,
417
+ fn: torch.Tensor,
418
+ hc_scale: torch.Tensor,
419
+ hc_base: torch.Tensor,
420
+ out: torch.Tensor,
421
+ hidden_size: int,
422
+ rms_eps: float,
423
+ hc_eps: float,
424
+ hc_mult: int,
425
+) -> None:
426
+ """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
427
+ if hs_flat.shape[0] == 0:
428
+ return
429
+ from vllm._tilelang_ops import hc_head_fuse_tilelang
430
+
431
+ hc_head_fuse_tilelang(
432
+ hs_flat,
433
+ fn,
434
+ hc_scale,
435
+ hc_base,
436
+ out,
437
+ hidden_size,
438
+ rms_eps,
439
+ hc_eps,
440
+ hc_mult,
441
+ )
442
+
443
+
444
+direct_register_custom_op(
445
+ op_name="mhc_pre_tilelang",
446
+ op_func=mhc_pre_tilelang,
447
+ mutates_args=[],
448
+ fake_impl=_mhc_pre_tilelang_fake,
449
+)
450
+direct_register_custom_op(
451
+ op_name="mhc_post_tilelang",
452
+ op_func=mhc_post_tilelang,
453
+ mutates_args=[],
454
+ fake_impl=_mhc_post_tilelang_fake,
455
+)
456
+
457
+direct_register_custom_op(
458
+ op_name="mhc_fused_post_pre_tilelang",
459
+ op_func=mhc_fused_post_pre_tilelang,
460
+ mutates_args=[],
461
+ fake_impl=_mhc_fused_post_pre_tilelang_fake,
462
+)
463
+
464
+direct_register_custom_op(
465
+ op_name="hc_head_fused_kernel_tilelang",
466
+ op_func=_hc_head_fused_kernel_tilelang,
467
+ mutates_args=["out"],
468
+)
vllm/model_executor/kernels/mhc/torch.py
new
+106
@@ -0,0 +1,106 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+
6
+def mhc_pre_torch(
7
+ residual: torch.Tensor,
8
+ fn: torch.Tensor,
9
+ hc_scale: torch.Tensor,
10
+ hc_base: torch.Tensor,
11
+ rms_eps: float,
12
+ hc_pre_eps: float,
13
+ hc_sinkhorn_eps: float,
14
+ hc_post_mult_value: float,
15
+ sinkhorn_repeat: int,
16
+ n_splits: int = 1,
17
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
18
+ """
19
+ Forward pass for mHC pre block.
20
+
21
+ Args:
22
+ residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
23
+ fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
24
+ hc_scale: shape (3,), dtype torch.float32
25
+ hc_base: shape (hc_mult3,), dtype torch.float32
26
+ rms_eps: RMS normalization epsilon
27
+ hc_pre_eps: pre-mix epsilon
28
+ hc_sinkhorn_eps: sinkhorn epsilon
29
+ hc_post_mult_value: post-mix multiplier value
30
+ sinkhorn_repeat: number of sinkhorn iterations
31
+ n_splits: split-k factor;
32
+
33
+ Returns:
34
+ post_mix: shape (..., hc_mult), dtype torch.float32
35
+ comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
36
+ layer_input: shape (..., hidden_size), dtype torch.bfloat16
37
+ """
38
+
39
+ # Validate shapes
40
+ assert residual.dtype == torch.bfloat16
41
+ assert fn.dtype == torch.float32
42
+ assert hc_scale.dtype == torch.float32
43
+ assert hc_base.dtype == torch.float32
44
+
45
+ hc_mult = residual.shape[-2]
46
+ hidden_size = residual.shape[-1]
47
+ hc_mult2 = hc_mult * hc_mult
48
+ hc_mult3 = hc_mult * 2 + hc_mult2
49
+
50
+ hc_hidden_size = hc_mult * hidden_size
51
+ assert fn.shape[0] == hc_mult3
52
+ assert fn.shape[1] == hc_hidden_size
53
+ assert hc_scale.shape == (3,)
54
+ assert hc_base.shape == (hc_mult3,)
55
+
56
+ outer_shape = residual.shape[:-2]
57
+
58
+ residual_flat = residual.view(-1, hc_mult, hidden_size)
59
+ num_tokens = residual_flat.shape[0]
60
+ fn_flat = fn
61
+
62
+ x = residual_flat.view(num_tokens, hc_mult * hidden_size).to(torch.float32)
63
+ mixes = torch.matmul(x, fn_flat.t())
64
+ sqrsum = x.square().sum(dim=-1, keepdim=True)
65
+ mixes = mixes * torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
66
+
67
+ pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]
68
+ pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps
69
+
70
+ post_logits = (
71
+ mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1] + hc_base[hc_mult : 2 * hc_mult]
72
+ )
73
+ post_mix = torch.sigmoid(post_logits) * hc_post_mult_value
74
+
75
+ comb_logits = mixes[:, 2 * hc_mult :].view(num_tokens, hc_mult, hc_mult) * hc_scale[
76
+ 2
77
+ ] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
78
+ comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps
79
+ comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
80
+ for _ in range(sinkhorn_repeat - 1):
81
+ comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
82
+ comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
83
+
84
+ layer_input = torch.sum(
85
+ pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1
86
+ ).to(torch.bfloat16)
87
+ return (
88
+ post_mix.view(*outer_shape, hc_mult, 1),
89
+ comb_mix.view(*outer_shape, hc_mult, hc_mult),
90
+ layer_input.view(*outer_shape, hidden_size),
91
+ )
92
+
93
+
94
+def mhc_post_torch(
95
+ x: torch.Tensor,
96
+ residual: torch.Tensor,
97
+ post_layer_mix: torch.Tensor,
98
+ comb_res_mix: torch.Tensor,
99
+) -> torch.Tensor:
100
+ mixed_residual = torch.einsum(
101
+ "...ij,...ih->...jh",
102
+ comb_res_mix.to(torch.float32),
103
+ residual.to(torch.float32),
104
+ )
105
+ post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32)
106
+ return (mixed_residual + post_term).to(residual.dtype)
vllm/model_executor/kernels/mhc/triton.py
new
+174
@@ -0,0 +1,174 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+import torch.nn.functional as F
5
+from torch import Tensor
6
+
7
+from vllm.triton_utils import tl, triton
8
+from vllm.utils.torch_utils import direct_register_custom_op
9
+
10
+
11
+@triton.jit
12
+def _rmsnorm_nw_kernel(
13
+ x_ptr,
14
+ out_ptr,
15
+ stride_row,
16
+ D,
17
+ eps,
18
+ RBLOCK: tl.constexpr,
19
+):
20
+ """Weight-free RMSNorm Triton kernel: out = x * rsqrt(mean(x², -1) + eps)."""
21
+ row = tl.program_id(0)
22
+ cols = tl.arange(0, RBLOCK)
23
+ mask = cols < D
24
+
25
+ x = tl.load(
26
+ x_ptr + row * stride_row + cols,
27
+ mask=mask,
28
+ other=0.0,
29
+ eviction_policy="evict_first",
30
+ ).to(tl.float32)
31
+
32
+ var = tl.sum(x * x, 0) / D
33
+ rstd = tl.rsqrt(var + eps)
34
+
35
+ out = (x * rstd).to(out_ptr.dtype.element_ty)
36
+ tl.store(out_ptr + row * D + cols, out, mask=mask, eviction_policy="evict_first")
37
+
38
+
39
+def rmsnorm_nw(x: Tensor, eps: float) -> Tensor:
40
+ """Weight-free RMSNorm over the last dimension.
41
+
42
+ Treats *x* as ``[num_rows, D]`` where ``num_rows = product(shape[:-1])``.
43
+ Returns a contiguous tensor with the same shape and dtype as *x*.
44
+ """
45
+ orig_shape = x.shape
46
+ D = orig_shape[-1]
47
+ x_2d = x.reshape(-1, D)
48
+ num_rows = x_2d.shape[0]
49
+
50
+ out = torch.empty_like(x_2d)
51
+ RBLOCK = triton.next_power_of_2(D)
52
+
53
+ _rmsnorm_nw_kernel[(num_rows,)](
54
+ x_2d,
55
+ out,
56
+ x_2d.stride(0),
57
+ D,
58
+ eps,
59
+ RBLOCK=RBLOCK,
60
+ num_warps=1 if RBLOCK <= 512 else (4 if RBLOCK <= 4096 else 8),
61
+ )
62
+ return out.view(orig_shape)
63
+
64
+
65
+@triton.jit
66
+def _hc_head_reduce_store_kernel(
67
+ pre_ptr,
68
+ x_ptr,
69
+ out_ptr,
70
+ hidden_size: tl.constexpr,
71
+ hc_mult: tl.constexpr,
72
+ pre_stride_t: tl.constexpr,
73
+ pre_stride_m: tl.constexpr,
74
+ x_stride_t: tl.constexpr,
75
+ x_stride_m: tl.constexpr,
76
+ x_stride_h: tl.constexpr,
77
+ out_stride_t: tl.constexpr,
78
+ out_stride_h: tl.constexpr,
79
+ BLOCK_H: tl.constexpr,
80
+):
81
+ token_idx = tl.program_id(0)
82
+ block_idx = tl.program_id(1)
83
+ offsets = block_idx * BLOCK_H + tl.arange(0, BLOCK_H)
84
+ mask = offsets < hidden_size
85
+
86
+ acc = tl.zeros((BLOCK_H,), dtype=tl.float32)
87
+ for mix_idx in tl.static_range(0, hc_mult):
88
+ pre = tl.load(pre_ptr + token_idx * pre_stride_t + mix_idx * pre_stride_m).to(
89
+ tl.float32
90
+ )
91
+ x = tl.load(
92
+ x_ptr
93
+ + token_idx * x_stride_t
94
+ + mix_idx * x_stride_m
95
+ + offsets * x_stride_h,
96
+ mask=mask,
97
+ other=0.0,
98
+ ).to(tl.float32)
99
+ acc += pre * x
100
+
101
+ tl.store(
102
+ out_ptr + token_idx * out_stride_t + offsets * out_stride_h,
103
+ acc,
104
+ mask=mask,
105
+ )
106
+
107
+
108
+def hc_head_reduce_triton_kernel(
109
+ x: torch.Tensor,
110
+ hc_fn: torch.Tensor,
111
+ hc_scale: torch.Tensor,
112
+ hc_base: torch.Tensor,
113
+ out: torch.Tensor,
114
+ norm_eps: float,
115
+ hc_eps: float,
116
+) -> None:
117
+ x_flat = x.flatten(-2)
118
+ x_normed = rmsnorm_nw(x_flat, norm_eps)
119
+ mixes = F.linear(x_normed.float(), hc_fn)
120
+ pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
121
+
122
+ hidden_size = x.shape[-1]
123
+ hc_mult = x.shape[-2]
124
+ block_h = 1024
125
+ _hc_head_reduce_store_kernel[(x.shape[0], (hidden_size + block_h - 1) // block_h)](
126
+ pre,
127
+ x,
128
+ out,
129
+ hidden_size,
130
+ hc_mult,
131
+ pre.stride(0),
132
+ pre.stride(1),
133
+ x.stride(0),
134
+ x.stride(1),
135
+ x.stride(2),
136
+ out.stride(0),
137
+ out.stride(1),
138
+ BLOCK_H=block_h,
139
+ num_warps=4,
140
+ )
141
+
142
+
143
+def _hc_head_triton(
144
+ hs_flat: torch.Tensor,
145
+ fn: torch.Tensor,
146
+ hc_scale: torch.Tensor,
147
+ hc_base: torch.Tensor,
148
+ out: torch.Tensor,
149
+ hidden_size: int,
150
+ rms_eps: float,
151
+ hc_eps: float,
152
+ hc_mult: int,
153
+) -> None:
154
+ """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
155
+ if hs_flat.shape[0] == 0:
156
+ return
157
+
158
+ hc_head_reduce_triton_kernel(
159
+ hs_flat,
160
+ fn,
161
+ hc_scale,
162
+ hc_base,
163
+ out,
164
+ rms_eps,
165
+ hc_eps,
166
+ )
167
+ return
168
+
169
+
170
+direct_register_custom_op(
171
+ op_name="hc_head_triton",
172
+ op_func=_hc_head_triton,
173
+ mutates_args=["out"],
174
+)
vllm/model_executor/layers/mhc.py
+244
-988
@@ -1,1030 +1,286 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import math
4
-from functools import cache
5
-from typing import TYPE_CHECKING
6
-
3
import torch
4
5
+# this import will also register the custom ops
6
+import vllm.model_executor.kernels.mhc as mhc_kernels
7
+from vllm.model_executor.custom_op import CustomOp
8
from vllm.platforms import current_platform
10
-from vllm.utils.import_utils import has_tilelang
11
-from vllm.utils.math_utils import cdiv
12
-from vllm.utils.torch_utils import direct_register_custom_op
13
-
14
-# tilelang is only available on CUDA platforms
15
-if TYPE_CHECKING or current_platform.is_cuda_alike():
16
- if not has_tilelang():
17
- raise ImportError(
18
- "tilelang is required for mhc but is not installed. Install it with "
19
- "`pip install tilelang`."
20
- )
21
- import tilelang
22
- import tilelang.language as T
23
-else:
24
- tilelang = None # type: ignore[assignment]
25
- T = None # type: ignore[assignment]
26
-
27
-
28
-@cache
29
-def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
30
- device_props = torch.cuda.get_device_properties(0)
31
- n_sms = device_props.multi_processor_count
32
- split_k = n_sms // grid_size
33
- if k is not None:
34
- # avoid split_k for small k
35
- num_block_k = cdiv(k, block_k)
36
- split_k = min(split_k, num_block_k // 4)
37
- split_k = max(split_k, 1)
38
- return split_k
39
-
40
-
41
-@tilelang.jit(
42
- pass_configs={
43
- tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
44
- tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
45
- tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
46
- },
47
-)
48
-def mhc_pre_big_fuse_tilelang(
49
- gemm_out_mul,
50
- gemm_out_sqrsum,
51
- hc_scale,
52
- hc_base,
53
- residual,
54
- post_mix,
55
- comb_mix,
56
- layer_input,
57
- hidden_size: int,
58
- rms_eps: float,
59
- hc_pre_eps: float,
60
- hc_sinkhorn_eps: float,
61
- hc_post_mult_value: float,
62
- sinkhorn_repeat: int,
63
- n_splits: int = 16,
64
- hc_mult: int = 4,
65
-):
66
- """Deeply fused kernels, everything other than gemm & sqrsum in mHC pre block."""
67
- num_tokens = T.dynamic("num_tokens")
68
- hc_mult3 = hc_mult * (2 + hc_mult)
69
- hidden_block = math.gcd(512, hidden_size)
70
-
71
- gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
72
- gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type]
73
- hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type]
74
- hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
75
- residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
76
- # outputs
77
- post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
78
- comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type]
79
- layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
80
-
81
- with T.Kernel(num_tokens, threads=96) as i:
82
- T.pdl_sync()
83
- ##################################################################
84
- # _pre_norm_fn_fwd_norm
85
- rms = T.alloc_fragment(1, T.float32)
86
- mixes = T.alloc_fragment(hc_mult3, T.float32)
87
- T.clear(mixes)
88
- rms[0] = 0
89
- for i_split in T.serial(n_splits):
90
- rms[0] += gemm_out_sqrsum[i_split, i]
91
- rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
92
- for j in T.Parallel(hc_mult3):
93
- mixes[j] = 0
94
- for i_split in T.serial(n_splits):
95
- mixes[j] += gemm_out_mul[i_split, i, j]
96
- mixes[j] *= rms[0]
97
- mixes_shared = T.alloc_shared(hc_mult3, T.float32)
98
- T.copy(mixes, mixes_shared)
9
100
- if T.get_thread_binding() < 32:
101
- ##################################################################
102
- # _pre_split_mixes_fwd (post & comb)
103
- cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
104
- for j in T.Parallel(hc_mult):
105
- post_mix[i, j] = (
106
- T.sigmoid(
107
- mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
108
- )
109
- * hc_post_mult_value
110
- )
111
- for j, k in T.Parallel(hc_mult, hc_mult):
112
- cm[j, k] = (
113
- mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
114
- + hc_base[j * hc_mult + k + hc_mult * 2]
115
- )
10
117
- ##################################################################
118
- # _sinkhorn_fwd
119
- row_sum = T.alloc_fragment(hc_mult, T.float32)
120
- col_sum = T.alloc_fragment(hc_mult, T.float32)
11
+# --8<-- [start:mhc_pre]
12
+@CustomOp.register("mhc_pre")
13
+class MHCPreOp(CustomOp):
14
+ """MHC pre block.
15
122
- # comb = comb.softmax(-1) + eps
123
- row_max = T.alloc_fragment(hc_mult, T.float32)
124
- T.reduce_max(cm, row_max, dim=1)
125
- for j, k in T.Parallel(hc_mult, hc_mult):
126
- cm[j, k] = T.exp(cm[j, k] - row_max[j])
127
- T.reduce_sum(cm, row_sum, dim=1)
128
- for j, k in T.Parallel(hc_mult, hc_mult):
129
- cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
130
-
131
- # comb = comb / (comb.sum(-2) + eps)
132
- T.reduce_sum(cm, col_sum, dim=0)
133
- for j, k in T.Parallel(hc_mult, hc_mult):
134
- cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
135
-
136
- for _ in T.serial(sinkhorn_repeat - 1):
137
- # comb = comb / (comb.sum(-1) + eps)
138
- T.reduce_sum(cm, row_sum, dim=1)
139
- for j, k in T.Parallel(hc_mult, hc_mult):
140
- cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
16
+ Computes mix logits from RMS-normalized HC residual streams, then
17
+ returns post_mix, comb_mix, and
18
+ layer_input = sum_i pre_mix_i * residual_i.
19
+ """
20
142
- # comb = comb / (comb.sum(-2) + eps)
143
- T.reduce_sum(cm, col_sum, dim=0)
144
- for j, k in T.Parallel(hc_mult, hc_mult):
145
- cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
21
+ # --8<-- [end:mhc_pre]
22
+ @classmethod
23
+ def enabled(cls) -> bool:
24
+ return True
25
+
26
+ def forward_cuda(
27
+ self,
28
+ residual: torch.Tensor,
29
+ fn: torch.Tensor,
30
+ hc_scale: torch.Tensor,
31
+ hc_base: torch.Tensor,
32
+ rms_eps: float,
33
+ hc_pre_eps: float,
34
+ hc_sinkhorn_eps: float,
35
+ hc_post_mult_value: float,
36
+ sinkhorn_repeat: int,
37
+ n_splits: int = 1,
38
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
39
+ return torch.ops.vllm.mhc_pre_tilelang(
40
+ residual,
41
+ fn,
42
+ hc_scale,
43
+ hc_base,
44
+ rms_eps,
45
+ hc_pre_eps,
46
+ hc_sinkhorn_eps,
47
+ hc_post_mult_value,
48
+ sinkhorn_repeat,
49
+ n_splits,
50
+ )
51
147
- # save comb_mix to global memory
148
- for j, k in T.Parallel(hc_mult, hc_mult):
149
- comb_mix[i, j * hc_mult + k] = cm[j, k]
52
+ def forward_hip(
53
+ self,
54
+ residual: torch.Tensor,
55
+ fn: torch.Tensor,
56
+ hc_scale: torch.Tensor,
57
+ hc_base: torch.Tensor,
58
+ rms_eps: float,
59
+ hc_pre_eps: float,
60
+ hc_sinkhorn_eps: float,
61
+ hc_post_mult_value: float,
62
+ sinkhorn_repeat: int,
63
+ n_splits: int = 1,
64
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
65
+ hidden_size = residual.shape[-1]
66
+ if hidden_size % 256 == 0:
67
+ return torch.ops.vllm.mhc_pre_aiter(
68
+ residual,
69
+ fn,
70
+ hc_scale,
71
+ hc_base,
72
+ rms_eps,
73
+ hc_pre_eps,
74
+ hc_sinkhorn_eps,
75
+ hc_post_mult_value,
76
+ sinkhorn_repeat,
77
+ )
78
else:
151
- ##################################################################
152
- # _pre_split_mixes_fwd (pre)
153
- pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
154
- for j in T.Parallel(hc_mult):
155
- pre_mix_shared[j] = (
156
- T.sigmoid(
157
- mixes_shared[j] * hc_scale[0] + hc_base[j],
158
- )
159
- + hc_pre_eps
160
- )
161
- ###################################################################
162
- # _pre_apply_mix_fwd
163
- for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
164
- xs = T.alloc_shared((hc_mult, hidden_block), T.float32)
165
- xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
166
- T.copy(residual[i, 0, i0_h * hidden_block], xs)
167
- T.copy(xs, xl)
168
-
169
- ol = T.alloc_fragment(hidden_block, T.float32)
170
- T.clear(ol)
79
+ return mhc_kernels.mhc_pre_torch(
80
+ residual,
81
+ fn,
82
+ hc_scale,
83
+ hc_base,
84
+ rms_eps,
85
+ hc_pre_eps,
86
+ hc_sinkhorn_eps,
87
+ hc_post_mult_value,
88
+ sinkhorn_repeat,
89
+ )
90
172
- for i_hc in T.serial(hc_mult):
173
- pre = pre_mix_shared[i_hc]
174
- for i1_h in T.Parallel(hidden_block):
175
- ol[i1_h] += pre * xl[i_hc, i1_h]
91
+ def forward_native(self, *args, **kwargs):
92
+ raise NotImplementedError("Native implementation of mhc_pre is not available")
93
177
- T.copy(ol, layer_input[i, i0_h * hidden_block])
178
- T.pdl_trigger()
94
95
+# --8<-- [start:mhc_post]
96
+@CustomOp.register("mhc_post")
97
+class MHCPostOp(CustomOp):
98
+ """MHC post block.
99
181
-def mhc_pre(
182
- residual: torch.Tensor,
183
- fn: torch.Tensor,
184
- hc_scale: torch.Tensor,
185
- hc_base: torch.Tensor,
186
- rms_eps: float,
187
- hc_pre_eps: float,
188
- hc_sinkhorn_eps: float,
189
- hc_post_mult_value: float,
190
- sinkhorn_repeat: int,
191
- n_splits: int = 1,
192
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
100
+ Combines the layer output with the HC residual streams:
101
+ out_j = post_layer_mix_j * x + sum_i comb_res_mix_ij * residual_i.
102
"""
194
- Forward pass for mHC pre block.
195
-
196
- Args:
197
- residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
198
- fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
199
- hc_scale: shape (3,), dtype torch.float32
200
- hc_base: shape (hc_mult3,), dtype torch.float32
201
- rms_eps: RMS normalization epsilon
202
- hc_pre_eps: pre-mix epsilon
203
- hc_sinkhorn_eps: sinkhorn epsilon
204
- hc_post_mult_value: post-mix multiplier value
205
- sinkhorn_repeat: number of sinkhorn iterations
206
- n_splits: split-k factor;
207
-
208
- Returns:
209
- post_mix: shape (..., hc_mult), dtype torch.float32
210
- comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
211
- layer_input: shape (..., hidden_size), dtype torch.bfloat16
212
- """
213
-
214
- # Validate shapes
215
- assert residual.dtype == torch.bfloat16
216
- assert fn.dtype == torch.float32
217
- assert hc_scale.dtype == torch.float32
218
- assert hc_base.dtype == torch.float32
219
-
220
- hc_mult = residual.shape[-2]
221
- hidden_size = residual.shape[-1]
222
- hc_mult2 = hc_mult * hc_mult
223
- hc_mult3 = hc_mult * 2 + hc_mult2
224
-
225
- hc_hidden_size = hc_mult * hidden_size
226
- assert fn.shape[0] == hc_mult3
227
- assert fn.shape[1] == hc_hidden_size
228
- assert hc_scale.shape == (3,)
229
- assert hc_base.shape == (hc_mult3,)
103
231
- outer_shape = residual.shape[:-2]
232
-
233
- residual_flat = residual.view(-1, hc_mult, hidden_size)
234
- num_tokens = residual_flat.shape[0]
235
- fn_flat = fn
236
-
237
- if current_platform.is_rocm():
238
- x = residual_flat.view(num_tokens, hc_mult * hidden_size).to(torch.float32)
239
- mixes = torch.matmul(x, fn_flat.t())
240
- sqrsum = x.square().sum(dim=-1, keepdim=True)
241
- mixes = mixes * torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
242
-
243
- pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]
244
- pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps
245
-
246
- post_logits = (
247
- mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1]
248
- + hc_base[hc_mult : 2 * hc_mult]
104
+ # --8<-- [end:mhc_post]
105
+
106
+ @classmethod
107
+ def enabled(cls) -> bool:
108
+ return True
109
+
110
+ def forward_cuda(
111
+ self,
112
+ x: torch.Tensor,
113
+ residual: torch.Tensor,
114
+ post_layer_mix: torch.Tensor,
115
+ comb_res_mix: torch.Tensor,
116
+ ) -> torch.Tensor:
117
+ return torch.ops.vllm.mhc_post_tilelang(
118
+ x, residual, post_layer_mix, comb_res_mix
119
)
250
- post_mix = torch.sigmoid(post_logits) * hc_post_mult_value
251
-
252
- comb_logits = mixes[:, 2 * hc_mult :].view(
253
- num_tokens, hc_mult, hc_mult
254
- ) * hc_scale[2] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
255
- comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps
256
- comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
257
- for _ in range(sinkhorn_repeat - 1):
258
- comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
259
- comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
260
-
261
- layer_input = torch.sum(
262
- pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1
263
- ).to(torch.bfloat16)
264
- return (
265
- post_mix.view(*outer_shape, hc_mult, 1),
266
- comb_mix.view(*outer_shape, hc_mult, hc_mult),
267
- layer_input.view(*outer_shape, hidden_size),
268
- )
269
-
270
- # these number are from deepgemm kernel impl
271
- block_k = 64
272
- block_m = 64
273
- n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
274
-
275
- post_mix = torch.empty(
276
- num_tokens,
277
- hc_mult,
278
- dtype=torch.float32,
279
- device=residual.device,
280
- )
281
- comb_mix = torch.empty(
282
- num_tokens,
283
- hc_mult2,
284
- dtype=torch.float32,
285
- device=residual.device,
286
- )
287
- layer_input = torch.empty(
288
- num_tokens,
289
- hidden_size,
290
- dtype=torch.bfloat16,
291
- device=residual.device,
292
- )
293
-
294
- gemm_out_mul = torch.empty(
295
- n_splits,
296
- num_tokens,
297
- hc_mult3,
298
- dtype=torch.float32,
299
- device=residual.device,
300
- )
301
- gemm_out_sqrsum = torch.empty(
302
- n_splits,
303
- num_tokens,
304
- dtype=torch.float32,
305
- device=residual.device,
306
- )
307
-
308
- from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
309
-
310
- tf32_hc_prenorm_gemm(
311
- residual_flat.view(num_tokens, hc_mult * hidden_size),
312
- fn_flat,
313
- gemm_out_mul,
314
- gemm_out_sqrsum,
315
- n_splits,
316
- )
317
-
318
- mhc_pre_big_fuse_tilelang(
319
- gemm_out_mul,
320
- gemm_out_sqrsum,
321
- hc_scale,
322
- hc_base,
323
- residual_flat,
324
- post_mix,
325
- comb_mix,
326
- layer_input,
327
- hidden_size,
328
- rms_eps,
329
- hc_pre_eps,
330
- hc_sinkhorn_eps,
331
- hc_post_mult_value,
332
- sinkhorn_repeat,
333
- n_splits,
334
- hc_mult,
335
- )
336
-
337
- post_mix = post_mix.view(*outer_shape, hc_mult, 1)
338
- comb_mix = comb_mix.view(*outer_shape, hc_mult, hc_mult)
339
- layer_input = layer_input.view(*outer_shape, hidden_size)
340
-
341
- return post_mix, comb_mix, layer_input
342
-
343
-
344
-def _mhc_pre_fake(
345
- residual: torch.Tensor,
346
- fn: torch.Tensor,
347
- hc_scale: torch.Tensor,
348
- hc_base: torch.Tensor,
349
- rms_eps: float,
350
- hc_pre_eps: float,
351
- hc_sinkhorn_eps: float,
352
- hc_post_mult_value: float,
353
- sinkhorn_repeat: int,
354
- n_splits: int = 1,
355
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
356
- hc_mult = residual.shape[-2]
357
- hidden_size = residual.shape[-1]
358
- outer_shape = residual.shape[:-2]
359
-
360
- # Create empty tensors with correct shapes for meta device / shape inference
361
- post_mix = torch.empty(
362
- *outer_shape,
363
- hc_mult,
364
- 1,
365
- dtype=torch.float32,
366
- device=residual.device,
367
- )
368
- comb_mix = torch.empty(
369
- *outer_shape,
370
- hc_mult,
371
- hc_mult,
372
- dtype=torch.float32,
373
- device=residual.device,
374
- )
375
- layer_input = torch.empty(
376
- *outer_shape,
377
- hidden_size,
378
- dtype=torch.bfloat16,
379
- device=residual.device,
380
- )
381
-
382
- return post_mix, comb_mix, layer_input
383
-
384
-
385
-@tilelang.jit(
386
- pass_configs={
387
- tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
388
- tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
389
- tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
390
- },
391
-)
392
-def mhc_post_tilelang(
393
- a,
394
- b,
395
- c,
396
- d,
397
- x,
398
- hc: int,
399
- hidden: int,
400
- n_thr: int = 128,
401
- h_blk: int = 1024,
402
-) -> tilelang.JITKernel:
403
- # rename for shorter code
404
- n = T.dynamic("num_tokens")
405
- h = hidden
406
-
407
- h_blk = math.gcd(hidden, h_blk)
408
- a: T.Tensor((n, hc, hc), T.float32) # type: ignore[no-redef, valid-type]
409
- b: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
410
- c: T.Tensor((n, hc), T.float32) # type: ignore[no-redef, valid-type]
411
- d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type]
412
- x: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
413
- with T.Kernel(n, threads=n_thr) as i_n:
414
- x_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
415
- b_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
416
- d_shared = T.alloc_shared(h_blk, T.bfloat16)
417
-
418
- x_local = T.alloc_fragment((hc, h_blk), T.float32)
419
- b_local = T.alloc_fragment((hc, h_blk), T.float32)
420
- d_local = T.alloc_fragment(h_blk, T.float32)
421
-
422
- a_local = T.alloc_fragment((hc, hc), T.float32)
423
- c_local = T.alloc_fragment(hc, T.float32)
424
- T.pdl_sync()
425
- T.copy(a[i_n, 0, 0], a_local)
426
- T.copy(c[i_n, 0], c_local)
427
-
428
- for i0_h in T.Pipelined(T.ceildiv(h, h_blk), num_stages=2):
429
- T.copy(b[i_n, 0, i0_h * h_blk], b_shared)
430
- T.copy(d[i_n, i0_h * h_blk], d_shared)
431
-
432
- T.copy(b_shared, b_local)
433
- T.copy(d_shared, d_local)
434
- for i_hco, i1_h in T.Parallel(hc, h_blk):
435
- x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h]
436
- for i_hci in T.serial(hc):
437
- x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
438
- T.copy(x_local, x_shared)
439
-
440
- T.copy(x_shared, x[i_n, 0, i0_h * h_blk])
441
- T.pdl_trigger()
442
-
443
-
444
-@tilelang.jit(
445
- pass_configs={
446
- tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
447
- tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
448
- tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
449
- },
450
-)
451
-def mhc_fused_tilelang(
452
- comb_mix,
453
- residual_in,
454
- post_mix,
455
- x_in,
456
- weight_t,
457
- yp_out,
458
- rp_out,
459
- residual_out,
460
- hc: int,
461
- hidden: int,
462
- n_out: int,
463
- n_thr: int = 256,
464
- h_blk: int = 256,
465
- tile_n: int = 1,
466
- split_k: int = 1,
467
-) -> tilelang.JITKernel:
468
- """Fused mhc post-mapping + pre-norm GEMM FMA"""
469
- m = T.dynamic("num_tokens")
470
- split_k = T.dynamic("split_k")
471
- h = hidden
472
- h_blk = math.gcd(hidden, h_blk)
473
- h_per_split = h // split_k
474
- n_tiles = n_out // tile_n
475
-
476
- comb_mix: T.Tensor((m, hc, hc), T.float32) # type: ignore[no-redef, valid-type]
477
- residual_in: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
478
- post_mix: T.Tensor((m, hc), T.float32) # type: ignore[no-redef, valid-type]
479
- x_in: T.Tensor((m, h), T.bfloat16) # type: ignore[no-redef, valid-type]
480
- weight_t: T.Tensor((n_out, hc, h), T.float32) # type: ignore[no-redef, valid-type]
481
- yp_out: T.Tensor((split_k, m, n_out), T.float32) # type: ignore[no-redef, valid-type]
482
- rp_out: T.Tensor((split_k, m), T.float32) # type: ignore[no-redef, valid-type]
483
- residual_out: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
484
-
485
- h_iters = h_per_split // n_thr
486
- num_warps = n_thr // 32
487
-
488
- with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
489
- tid = T.get_thread_binding()
490
- warp_id = T.get_warp_idx()
491
- lane = T.get_lane_idx()
492
-
493
- s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
494
- s_post = T.alloc_shared((hc,), T.float32)
495
- s_comb = T.alloc_shared((hc, hc), T.float32)
496
-
497
- pm = T.alloc_local((hc,), T.float32)
498
- cm = T.alloc_local((hc, hc), T.float32)
499
- acc = T.alloc_local((tile_n,), T.float32)
500
- sqr = T.alloc_local((1,), T.float32)
501
- new_r = T.alloc_local((hc,), T.float32)
502
-
503
- T.clear(acc)
504
- T.clear(sqr)
505
- h_split_start = i_ks * h_per_split
506
-
507
- T.pdl_sync()
508
-
509
- T.copy(post_mix[i_n, 0], s_post)
510
- T.copy(comb_mix[i_n, 0, 0], s_comb)
511
-
512
- for j in T.unroll(hc):
513
- pm[j] = s_post[j]
514
- for j in T.unroll(hc):
515
- for k in T.unroll(hc):
516
- cm[k, j] = s_comb[k, j]
517
-
518
- # Each thread owns h_iters elements of the k-split's h slice.
519
- for it in T.serial(h_iters):
520
- h_idx = h_split_start + it * n_thr + tid
521
-
522
- # Compute new residual from layer output and past residual
523
- for j in T.unroll(hc):
524
- new_r[j] = pm[j] * x_in[i_n, h_idx]
525
- for k in T.unroll(hc):
526
- new_r[j] += cm[k, j] * residual_in[i_n, k, h_idx]
527
-
528
- # populate residual_out and compute sqr sum
529
- if i_nt == 0:
530
- for j in T.unroll(hc):
531
- residual_out[i_n, j, h_idx] = new_r[j]
532
- sqr[0] += new_r[j] * new_r[j]
533
-
534
- # Per-thread FMA into acc[n]
535
- for n in T.unroll(tile_n):
536
- for j in T.unroll(hc):
537
- acc[n] += weight_t[i_nt * tile_n + n, j, h_idx] * new_r[j]
538
-
539
- for n in T.unroll(tile_n):
540
- acc[n] = T.warp_reduce_sum(acc[n])
541
- if i_nt == 0:
542
- sqr[0] = T.warp_reduce_sum(sqr[0])
543
-
544
- # Cross-warp reduce via shared mem
545
- if lane == 0:
546
- for n in T.unroll(tile_n):
547
- s_warp[warp_id, n] = acc[n]
548
- if i_nt == 0:
549
- s_warp[warp_id, tile_n] = sqr[0]
550
- T.sync_threads()
551
-
552
- # Warp 0 does the final cross-warp sum and writes outputs
553
- if warp_id == 0:
554
- if lane < tile_n:
555
- v = T.alloc_var(T.float32, init=0.0)
556
- for w in T.unroll(num_warps):
557
- v += s_warp[w, lane]
558
- yp_out[i_ks, i_n, i_nt * tile_n + lane] = v
120
560
- if i_nt == 0 and lane == 0:
561
- v2 = T.alloc_var(T.float32, init=0.0)
562
- for w in T.unroll(num_warps):
563
- v2 += s_warp[w, tile_n]
564
- rp_out[i_ks, i_n] = v2
565
-
566
- T.pdl_trigger()
121
+ def forward_hip(
122
+ self,
123
+ x: torch.Tensor,
124
+ residual: torch.Tensor,
125
+ post_layer_mix: torch.Tensor,
126
+ comb_res_mix: torch.Tensor,
127
+ ) -> torch.Tensor:
128
+ hidden_size = residual.shape[-1]
129
+ if hidden_size % 256 == 0:
130
+ return torch.ops.vllm.mhc_post_aiter(
131
+ x,
132
+ residual,
133
+ post_layer_mix,
134
+ comb_res_mix,
135
+ )
136
+ else:
137
+ return mhc_kernels.mhc_post_torch(
138
+ x,
139
+ residual,
140
+ post_layer_mix,
141
+ comb_res_mix,
142
+ )
143
144
+ def forward_native(self, *args, **kwargs):
145
+ raise NotImplementedError("Native implementation of mhc_post is not available")
146
569
-def mhc_post(
570
- x: torch.Tensor,
571
- residual: torch.Tensor,
572
- post_layer_mix: torch.Tensor,
573
- comb_res_mix: torch.Tensor,
574
-) -> torch.Tensor:
575
- if current_platform.is_rocm():
576
- mixed_residual = torch.einsum(
577
- "...ij,...ih->...jh",
578
- comb_res_mix.to(torch.float32),
579
- residual.to(torch.float32),
580
- )
581
- post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32)
582
- return (mixed_residual + post_term).to(residual.dtype)
583
- out = torch.empty_like(residual)
584
- mhc_post_tilelang(
585
- comb_res_mix,
586
- residual,
587
- post_layer_mix.squeeze(-1),
588
- x,
589
- out,
590
- residual.shape[-2],
591
- residual.shape[-1],
592
- )
593
- return out
147
148
+# --8<-- [start:hc_head]
149
+@CustomOp.register("hc_head")
150
+class HCHeadOp(CustomOp):
151
+ """HC head reduction for DeepSeek V4.
152
596
-def mhc_fused_post_pre(
597
- x: torch.Tensor,
598
- residual: torch.Tensor,
599
- post_layer_mix: torch.Tensor,
600
- comb_res_mix: torch.Tensor,
601
- fn: torch.Tensor,
602
- hc_scale: torch.Tensor,
603
- hc_base: torch.Tensor,
604
- rms_eps: float,
605
- hc_pre_eps: float,
606
- hc_sinkhorn_eps: float,
607
- hc_post_mult_value: float,
608
- sinkhorn_repeat: int,
609
- n_splits: int = 1,
610
- tile_n: int = 1,
611
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
153
+ Computes gates from the RMS-normalized flattened HC residual and
154
+ returns out = sum_i gate_i * residual_i, collapsing hc_mult streams
155
+ to one.
156
"""
613
- Run one MHC post block followed by the next MHC pre block.
614
-
615
- Returns:
616
- residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
617
- post_mix_cur: shape (..., hc_mult, 1)
618
- comb_mix_cur: shape (..., hc_mult, hc_mult)
619
- layer_input_cur: shape (..., hidden_size)
620
- """
621
-
622
- assert residual.dtype == torch.bfloat16
623
- assert x.dtype == torch.bfloat16
624
- assert post_layer_mix.dtype == torch.float32
625
- assert comb_res_mix.dtype == torch.float32
626
- assert fn.dtype == torch.float32
627
- assert hc_scale.dtype == torch.float32
628
- assert hc_base.dtype == torch.float32
629
-
630
- hc_mult = residual.shape[-2]
631
- hidden_size = residual.shape[-1]
632
- hc_mult2 = hc_mult * hc_mult
633
- hc_mult3 = hc_mult * 2 + hc_mult2
634
- hc_hidden_size = hc_mult * hidden_size
635
- outer_shape = residual.shape[:-2]
157
637
- assert x.shape == (*outer_shape, hidden_size)
638
- assert post_layer_mix.shape in (
639
- (*outer_shape, hc_mult, 1),
640
- (*outer_shape, hc_mult),
641
- )
642
- assert comb_res_mix.shape == (*outer_shape, hc_mult, hc_mult)
643
- assert fn.shape == (hc_mult3, hc_hidden_size)
644
- assert hc_scale.shape == (3,)
645
- assert hc_base.shape == (hc_mult3,)
646
-
647
- assert n_splits in (1, 2, 4, 8)
648
- assert hidden_size % n_splits == 0
649
-
650
- residual_flat = residual.view(-1, hc_mult, hidden_size)
651
- num_tokens = residual_flat.shape[0]
652
- x_flat = x.view(num_tokens, hidden_size)
653
- post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
654
- comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
655
-
656
- fma_token_threshold = 16
657
- if num_tokens <= fma_token_threshold:
658
- # TODO(gnovack): investigate autotuning these heuristics
659
- tile_n = 2 if num_tokens < 8 else 3
660
- n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
661
- else:
662
- # these number are from deepgemm kernel impl
663
- block_k = 64
664
- block_m = 64
665
- n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
666
-
667
- gemm_out_mul = torch.empty(
668
- n_splits,
669
- num_tokens,
670
- hc_mult3,
671
- dtype=torch.float32,
672
- device=residual.device,
673
- )
674
- gemm_out_sqrsum = torch.empty(
675
- n_splits,
676
- num_tokens,
677
- dtype=torch.float32,
678
- device=residual.device,
679
- )
680
- residual_cur = torch.empty_like(residual_flat)
681
- post_mix_cur = torch.empty(
682
- num_tokens,
683
- hc_mult,
684
- dtype=torch.float32,
685
- device=residual.device,
686
- )
687
- comb_mix_cur = torch.empty(
688
- num_tokens,
689
- hc_mult2,
690
- dtype=torch.float32,
691
- device=residual.device,
692
- )
693
- layer_input_cur = torch.empty(
694
- num_tokens,
695
- hidden_size,
696
- dtype=torch.bfloat16,
697
- device=residual.device,
698
- )
699
-
700
- if num_tokens <= fma_token_threshold:
701
- mhc_fused_tilelang(
702
- comb_res_mix_flat,
703
- residual_flat,
704
- post_layer_mix_flat,
705
- x_flat,
706
- fn.view(hc_mult3, hc_mult, hidden_size),
707
- gemm_out_mul,
708
- gemm_out_sqrsum,
709
- residual_cur,
710
- hc_mult,
158
+ # --8<-- [end:hc_head]
159
+ @classmethod
160
+ def enabled(cls) -> bool:
161
+ return True
162
+
163
+ def forward_cuda(
164
+ self,
165
+ hidden_states: torch.Tensor,
166
+ hc_fn: torch.Tensor,
167
+ hc_scale: torch.Tensor,
168
+ hc_base: torch.Tensor,
169
+ rms_norm_eps: float,
170
+ hc_eps: float,
171
+ ) -> torch.Tensor:
172
+ hc_mult, hidden_size = hidden_states.shape[-2:]
173
+ outer_shape = hidden_states.shape[:-2]
174
+ hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
175
+ num_tokens = hs_flat.shape[0]
176
+
177
+ out = torch.empty(
178
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
179
+ )
180
+ torch.ops.vllm.hc_head_fused_kernel_tilelang(
181
+ hs_flat,
182
+ hc_fn,
183
+ hc_scale,
184
+ hc_base,
185
+ out,
186
hidden_size,
712
- hc_mult3,
713
- tile_n=tile_n,
714
- n_splits=n_splits,
187
+ rms_norm_eps,
188
+ hc_eps,
189
+ hc_mult,
190
)
716
- else:
717
- mhc_post_tilelang(
718
- comb_res_mix_flat,
719
- residual_flat,
720
- post_layer_mix_flat,
721
- x_flat,
722
- residual_cur,
723
- residual.shape[-2],
724
- residual.shape[-1],
191
+ return out.view(*outer_shape, hidden_size)
192
+
193
+ # This @torch.compile is necessary for accuracy as well as performance.
194
+ @torch.compile(backend=current_platform.simple_compile_backend)
195
+ def forward_hip(
196
+ self,
197
+ hidden_states: torch.Tensor,
198
+ hc_fn: torch.Tensor,
199
+ hc_scale: torch.Tensor,
200
+ hc_base: torch.Tensor,
201
+ rms_norm_eps: float,
202
+ hc_eps: float,
203
+ ) -> torch.Tensor:
204
+ hc_mult, hidden_size = hidden_states.shape[-2:]
205
+ outer_shape = hidden_states.shape[:-2]
206
+ hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
207
+ num_tokens = hs_flat.shape[0]
208
+
209
+ out = torch.empty(
210
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
211
)
726
-
727
- from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
728
-
729
- tf32_hc_prenorm_gemm(
730
- residual_cur.view(num_tokens, hc_mult * hidden_size),
731
- fn,
732
- gemm_out_mul,
733
- gemm_out_sqrsum,
734
- n_splits,
212
+ torch.ops.vllm.hc_head_triton(
213
+ hs_flat,
214
+ hc_fn,
215
+ hc_scale,
216
+ hc_base,
217
+ out,
218
+ hidden_size,
219
+ rms_norm_eps,
220
+ hc_eps,
221
+ hc_mult,
222
)
223
+ return out.view(*outer_shape, hidden_size)
224
737
- mhc_pre_big_fuse_tilelang(
738
- gemm_out_mul,
739
- gemm_out_sqrsum,
740
- hc_scale,
741
- hc_base,
742
- residual_cur,
743
- post_mix_cur,
744
- comb_mix_cur,
745
- layer_input_cur,
746
- hidden_size,
747
- rms_eps,
748
- hc_pre_eps,
749
- hc_sinkhorn_eps,
750
- hc_post_mult_value,
751
- sinkhorn_repeat,
752
- n_splits,
753
- hc_mult,
754
- )
755
-
756
- return (
757
- residual_cur.view(*outer_shape, hc_mult, hidden_size),
758
- post_mix_cur.view(*outer_shape, hc_mult, 1),
759
- comb_mix_cur.view(*outer_shape, hc_mult, hc_mult),
760
- layer_input_cur.view(*outer_shape, hidden_size),
761
- )
762
-
763
-
764
-def _mhc_fused_post_pre_fake(
765
- x: torch.Tensor,
766
- residual: torch.Tensor,
767
- post_layer_mix: torch.Tensor,
768
- comb_res_mix: torch.Tensor,
769
- fn: torch.Tensor,
770
- hc_scale: torch.Tensor,
771
- hc_base: torch.Tensor,
772
- rms_eps: float,
773
- hc_pre_eps: float,
774
- hc_sinkhorn_eps: float,
775
- hc_post_mult_value: float,
776
- sinkhorn_repeat: int,
777
- n_splits: int = 1,
778
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
779
- hc_mult = residual.shape[-2]
780
- hidden_size = residual.shape[-1]
781
- outer_shape = residual.shape[:-2]
225
+ def forward_native(self, *args, **kwargs):
226
+ raise NotImplementedError("Native implementation of hc_head is not available")
227
783
- residual_cur = torch.empty_like(residual)
784
- post_mix_cur = torch.empty(
785
- *outer_shape,
786
- hc_mult,
787
- 1,
788
- dtype=torch.float32,
789
- device=residual.device,
790
- )
791
- comb_mix_cur = torch.empty(
792
- *outer_shape,
793
- hc_mult,
794
- hc_mult,
795
- dtype=torch.float32,
796
- device=residual.device,
797
- )
798
- layer_input_cur = torch.empty(
799
- *outer_shape,
800
- hidden_size,
801
- dtype=torch.bfloat16,
802
- device=residual.device,
803
- )
228
805
- return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
229
+# --8<-- [start:mhc_fused_post_pre]
230
+@CustomOp.register("mhc_fused_post_pre")
231
+class MHCFusedPostPreOp(CustomOp):
232
+ """Fused MHC post block followed by the next MHC pre block.
233
807
-
808
-def _mhc_post_fake(
809
- x: torch.Tensor,
810
- residual: torch.Tensor,
811
- post_layer_mix: torch.Tensor,
812
- comb_res_mix: torch.Tensor,
813
-) -> torch.Tensor:
814
- return torch.empty_like(residual)
815
-
816
-
817
-direct_register_custom_op(
818
- op_name="mhc_pre",
819
- op_func=mhc_pre,
820
- mutates_args=[],
821
- fake_impl=_mhc_pre_fake,
822
-)
823
-direct_register_custom_op(
824
- op_name="mhc_post",
825
- op_func=mhc_post,
826
- mutates_args=[],
827
- fake_impl=_mhc_post_fake,
828
-)
829
-direct_register_custom_op(
830
- op_name="mhc_fused_post_pre",
831
- op_func=mhc_fused_post_pre,
832
- mutates_args=[],
833
- fake_impl=_mhc_fused_post_pre_fake,
834
-)
835
-
836
-
837
-@tilelang.jit(
838
- pass_configs={
839
- tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
840
- tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
841
- tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
842
- },
843
-)
844
-def hc_head_fuse_tilelang(
845
- residual,
846
- fn,
847
- hc_scale,
848
- hc_base,
849
- out,
850
- hidden_size: int,
851
- rms_eps: float,
852
- hc_eps: float,
853
- hc_mult: int = 4,
854
- n_thr: int = 128,
855
- h_blk: int = 1024,
856
-):
857
- """Two-pass fused kernel for hc_head.
858
-
859
- Pass 1: accumulate per-token squared sum and hc_mult dot-products
860
- (projections onto fn rows) using cross-thread reducers.
861
- Pass 2: apply sigmoid-gated weighted sum of residual channels to output.
862
-
863
- Avoids materialising mixes / rsqrt / pre tensors to global memory.
234
+ Equivalent to applying MHCPostOp and then MHCPreOp to the updated
235
+ residual streams, returning residual_cur, post_mix_cur, comb_mix_cur,
236
+ and layer_input_cur.
237
"""
865
- num_tokens = T.dynamic("num_tokens")
866
- hc_dim = hc_mult * hidden_size
867
- h_block = math.gcd(h_blk, hidden_size)
868
- n_h = hidden_size // h_block
869
-
870
- residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]
871
- fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type]
872
- hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type]
873
- hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type]
874
- out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]
875
-
876
- with T.Kernel(num_tokens, threads=n_thr) as i:
877
- T.pdl_sync()
878
-
879
- # ------------------------------------------------------------------
880
- # Pass 1 – for each residual channel m_c and h_block:
881
- # • accumulate squared sum (for RMS norm denominator)
882
- # • accumulate hc_mult dot-products with fn rows
883
- # ------------------------------------------------------------------
884
- sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all")
885
- mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all")
886
- T.fill(sqrsum_r, 0.0)
887
- T.fill(mixes_r, 0.0)
888
-
889
- for m_c in T.serial(hc_mult):
890
- for i_h in T.serial(n_h):
891
- x_local = T.alloc_fragment(h_block, T.float32)
892
- T.copy(residual[i, m_c, i_h * h_block], x_local)
893
-
894
- for k in T.Parallel(h_block):
895
- sqrsum_r[0] += x_local[k] * x_local[k]
896
-
897
- for m_m in T.unroll(hc_mult):
898
- fn_local = T.alloc_fragment(h_block, T.float32)
899
- T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local)
900
- for k in T.Parallel(h_block):
901
- mixes_r[m_m] += x_local[k] * fn_local[k]
902
-
903
- T.finalize_reducer(sqrsum_r)
904
- T.finalize_reducer(mixes_r)
905
-
906
- # ------------------------------------------------------------------
907
- # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps
908
- # ------------------------------------------------------------------
909
- pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
910
- rsqrt_val = T.alloc_fragment(1, T.float32)
911
- rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps)
912
- for m in T.Parallel(hc_mult):
913
- pre_mix_shared[m] = (
914
- T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps
915
- )
916
-
917
- # ------------------------------------------------------------------
918
- # Pass 2 – apply_mix: pipelined weighted sum over residual channels
919
- # ------------------------------------------------------------------
920
- for i0_h in T.Pipelined(n_h, num_stages=2):
921
- xs = T.alloc_shared((hc_mult, h_block), T.bfloat16)
922
- xl = T.alloc_fragment((hc_mult, h_block), T.float32)
923
- T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True)
924
- T.copy(xs, xl)
925
-
926
- ol = T.alloc_fragment(h_block, T.float32)
927
- T.clear(ol)
928
- for i_hc in T.serial(hc_mult):
929
- pre = pre_mix_shared[i_hc]
930
- for i1_h in T.Parallel(h_block):
931
- ol[i1_h] += pre * xl[i_hc, i1_h]
238
933
- T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
934
-
935
- T.pdl_trigger()
936
-
937
-
938
-def _hc_head_fused_reference(
939
- hs_flat: torch.Tensor,
940
- fn: torch.Tensor,
941
- hc_scale: torch.Tensor,
942
- hc_base: torch.Tensor,
943
- out: torch.Tensor,
944
- hidden_size: int,
945
- rms_eps: float,
946
- hc_eps: float,
947
- hc_mult: int,
948
-) -> None:
949
- """Pure-PyTorch reference for `hc_head_fuse_tilelang`.
950
-
951
- Used on platforms where the tilelang HIP/CUDA backend is not available
952
- (e.g. ROCm builds shipping a tilelang wheel without `target.build.tilelang_hip`).
953
- Mirrors the math of the tilelang kernel exactly:
954
-
955
- x = hs_flat.flatten(-2, -1) # (T, hc_mult * H), fp32
956
- mixes = x @ fn.T # (T, hc_mult)
957
- rsqrt = 1 / sqrt(||x||^2 / (hc_mult * H) + rms_eps)
958
- pre[m] = sigmoid(mixes[m] * rsqrt * hc_scale[0] + hc_base[m]) + hc_eps
959
- out = sum_m pre[m] * hs_flat[:, m, :] # cast back to bf16
960
-
961
- `out` is mutated in place to keep the same op contract
962
- (`mutates_args=["out"]`).
963
- """
964
- num_tokens = hs_flat.shape[0]
965
- if num_tokens == 0:
966
- return
967
- x = hs_flat.reshape(num_tokens, hc_mult * hidden_size).to(torch.float32)
968
- # fn: (hc_mult, hc_mult * hidden_size) → mixes: (T, hc_mult)
969
- mixes = torch.matmul(x, fn.t())
970
- sqrsum = x.square().sum(dim=-1, keepdim=True)
971
- rsqrt = torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
972
- # hc_scale has shape (1,); hc_base has shape (hc_mult,)
973
- pre_mix = torch.sigmoid(mixes * rsqrt * hc_scale[0] + hc_base) + hc_eps
974
- # weighted sum over the hc_mult channel dim
975
- result = torch.sum(pre_mix.unsqueeze(-1) * hs_flat.to(torch.float32), dim=1).to(
976
- out.dtype
977
- )
978
- out.copy_(result)
979
-
980
-
981
-def _hc_head_fused_kernel(
982
- hs_flat: torch.Tensor,
983
- fn: torch.Tensor,
984
- hc_scale: torch.Tensor,
985
- hc_base: torch.Tensor,
986
- out: torch.Tensor,
987
- hidden_size: int,
988
- rms_eps: float,
989
- hc_eps: float,
990
- hc_mult: int,
991
-) -> None:
992
- """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
993
- if hs_flat.shape[0] == 0:
994
- return
995
- if current_platform.is_rocm():
996
- # tilelang ships only the CUDA codegen in upstream wheels, so the HIP
997
- # FFI target (`target.build.tilelang_hip`) is missing and the JIT call
998
- # would raise `ValueError: Cannot find global function ...`. Use a
999
- # numerically equivalent torch fallback instead. `mhc_pre` and
1000
- # `mhc_post` already follow this same pattern above.
1001
- _hc_head_fused_reference(
1002
- hs_flat,
239
+ # --8<-- [end:mhc_fused_post_pre]
240
+ @classmethod
241
+ def enabled(cls) -> bool:
242
+ return True
243
+
244
+ def forward_cuda(
245
+ self,
246
+ x: torch.Tensor,
247
+ residual: torch.Tensor,
248
+ post_layer_mix: torch.Tensor,
249
+ comb_res_mix: torch.Tensor,
250
+ fn: torch.Tensor,
251
+ hc_scale: torch.Tensor,
252
+ hc_base: torch.Tensor,
253
+ rms_eps: float,
254
+ hc_pre_eps: float,
255
+ hc_sinkhorn_eps: float,
256
+ hc_post_mult_value: float,
257
+ sinkhorn_repeat: int,
258
+ n_splits: int = 1,
259
+ tile_n: int = 1,
260
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
261
+ return torch.ops.vllm.mhc_fused_post_pre_tilelang(
262
+ x,
263
+ residual,
264
+ post_layer_mix,
265
+ comb_res_mix,
266
fn,
267
hc_scale,
268
hc_base,
1006
- out,
1007
- hidden_size,
269
rms_eps,
1009
- hc_eps,
1010
- hc_mult,
270
+ hc_pre_eps,
271
+ hc_sinkhorn_eps,
272
+ hc_post_mult_value,
273
+ sinkhorn_repeat,
274
+ n_splits,
275
+ tile_n,
276
)
1012
- return
1013
- hc_head_fuse_tilelang(
1014
- hs_flat,
1015
- fn,
1016
- hc_scale,
1017
- hc_base,
1018
- out,
1019
- hidden_size,
1020
- rms_eps,
1021
- hc_eps,
1022
- hc_mult,
1023
- )
277
278
+ def forward_hip(self, *args, **kwargs):
279
+ raise NotImplementedError(
280
+ "Hip implementation of mhc_fused_post_pre is not available"
281
+ )
282
1026
-direct_register_custom_op(
1027
- op_name="hc_head_fused_kernel",
1028
- op_func=_hc_head_fused_kernel,
1029
- mutates_args=["out"],
1030
-)
283
+ def forward_native(self, *args, **kwargs):
284
+ raise NotImplementedError(
285
+ "Native implementation of mhc_fused_post_pre is not available"
286
+ )
vllm/model_executor/models/deepseek_v4.py
+59
-38
@@ -35,6 +35,12 @@ from vllm.model_executor.layers.linear import (
35
RowParallelLinear,
36
)
37
from vllm.model_executor.layers.logits_processor import LogitsProcessor
38
+from vllm.model_executor.layers.mhc import (
39
+ HCHeadOp,
40
+ MHCFusedPostPreOp,
41
+ MHCPostOp,
42
+ MHCPreOp,
43
+)
44
from vllm.model_executor.layers.quantization import (
45
QuantizationConfig,
46
QuantizationMethods,
@@ -1168,6 +1174,9 @@ class DeepseekV4DecoderLayer(nn.Module):
1174
),
1175
requires_grad=False,
1176
)
1177
+ self.mhc_pre = MHCPreOp()
1178
+ self.mhc_post = MHCPostOp()
1179
+ self.mhc_fused_post_pre = MHCFusedPostPreOp()
1180
1181
def hc_pre(
1182
self,
@@ -1176,7 +1185,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1185
hc_scale: torch.Tensor,
1186
hc_base: torch.Tensor,
1187
):
1179
- post_mix, res_mix, layer_input = torch.ops.vllm.mhc_pre(
1188
+ post_mix, res_mix, layer_input = self.mhc_pre(
1189
residual=x,
1190
fn=hc_fn,
1191
hc_scale=hc_scale,
@@ -1196,9 +1205,9 @@ class DeepseekV4DecoderLayer(nn.Module):
1205
post: torch.Tensor,
1206
comb: torch.Tensor,
1207
):
1199
- return torch.ops.vllm.mhc_post(x, residual, post, comb)
1208
+ return self.mhc_post(x, residual, post, comb)
1209
1201
- def forward(
1210
+ def _forward_cuda(
1211
self,
1212
x: torch.Tensor,
1213
positions: torch.Tensor,
@@ -1214,7 +1223,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1223
x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1224
)
1225
else:
1217
- residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre(
1226
+ residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
1227
x,
1228
residual,
1229
post_mix,
@@ -1232,7 +1241,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1241
x = self.attn_norm(x)
1242
x = self.attn(positions, x, None)
1243
1235
- residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre(
1244
+ residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
1245
x,
1246
residual,
1247
post_mix,
@@ -1251,6 +1260,48 @@ class DeepseekV4DecoderLayer(nn.Module):
1260
x = self.ffn(x, input_ids)
1261
return x, residual, post_mix, res_mix
1262
1263
+ def _forward_rocm(
1264
+ self,
1265
+ x: torch.Tensor,
1266
+ positions: torch.Tensor,
1267
+ input_ids: torch.Tensor | None,
1268
+ post_mix: torch.Tensor | None,
1269
+ res_mix: torch.Tensor | None,
1270
+ residual: torch.Tensor | None,
1271
+ ) -> torch.Tensor:
1272
+ residual = x
1273
+ x, post, comb = self.hc_pre(
1274
+ x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1275
+ )
1276
+ x = self.attn_norm(x)
1277
+ x = self.attn(positions, x, None)
1278
+ x = self.hc_post(x, residual, post, comb)
1279
+
1280
+ residual = x
1281
+ x, post, comb = self.hc_pre(
1282
+ x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
1283
+ )
1284
+ x = self.ffn_norm(x)
1285
+ x = self.ffn(x, input_ids)
1286
+ x = self.hc_post(x, residual, post, comb)
1287
+ return x, None, None, None
1288
+
1289
+ def forward(
1290
+ self,
1291
+ x: torch.Tensor,
1292
+ positions: torch.Tensor,
1293
+ input_ids: torch.Tensor | None,
1294
+ post_mix: torch.Tensor | None,
1295
+ res_mix: torch.Tensor | None,
1296
+ residual: torch.Tensor | None,
1297
+ ) -> torch.Tensor:
1298
+ if current_platform.is_rocm():
1299
+ return self._forward_rocm(
1300
+ x, positions, input_ids, post_mix, res_mix, residual
1301
+ )
1302
+
1303
+ return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)
1304
+
1305
1306
@support_torch_compile
1307
class DeepseekV4Model(nn.Module):
@@ -1340,7 +1391,7 @@ class DeepseekV4Model(nn.Module):
1391
torch.empty(1, dtype=torch.float32),
1392
requires_grad=False,
1393
)
1343
-
1394
+ self.hc_head_op = HCHeadOp()
1395
# Pre-hc_head residual stream buffer for the MTP draft. Stable
1396
# address (outside the cudagraph pool) so the copy_ in forward()
1397
# refreshes it correctly across captured shapes.
@@ -1410,7 +1461,7 @@ class DeepseekV4Model(nn.Module):
1461
res_mix,
1462
residual,
1463
)
1413
- else:
1464
+ if layer is not None and current_platform.is_cuda():
1465
hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
1466
1467
if not get_pp_group().is_last_rank:
@@ -1420,7 +1471,7 @@ class DeepseekV4Model(nn.Module):
1471
num_tokens = hidden_states.shape[0]
1472
self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1473
1423
- hidden_states = hc_head(
1474
+ hidden_states = self.hc_head_op(
1475
hidden_states,
1476
self.hc_head_fn,
1477
self.hc_head_scale,
@@ -1549,36 +1600,6 @@ class DeepseekV4Model(nn.Module):
1600
layer.ffn.finalize_mega_moe_weights()
1601
1602
1552
-@torch.compile(backend=current_platform.simple_compile_backend)
1553
-def hc_head(
1554
- hidden_states: torch.Tensor,
1555
- hc_fn: torch.Tensor,
1556
- hc_scale: torch.Tensor,
1557
- hc_base: torch.Tensor,
1558
- rms_norm_eps: float,
1559
- hc_eps: float,
1560
-) -> torch.Tensor:
1561
- hc_mult, hidden_size = hidden_states.shape[-2:]
1562
- outer_shape = hidden_states.shape[:-2]
1563
- hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
1564
- num_tokens = hs_flat.shape[0]
1565
- out = torch.empty(
1566
- num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
1567
- )
1568
- torch.ops.vllm.hc_head_fused_kernel(
1569
- hs_flat,
1570
- hc_fn,
1571
- hc_scale,
1572
- hc_base,
1573
- out,
1574
- hidden_size,
1575
- rms_norm_eps,
1576
- hc_eps,
1577
- hc_mult,
1578
- )
1579
- return out.view(*outer_shape, hidden_size)
1580
-
1581
-
1603
def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1604
if expert_dtype == "fp4":
1605
# MXFP4 experts use Mxfp4MoEMethod, which registers scales as
vllm/model_executor/models/deepseek_v4_mtp.py
+4
-2
@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe import FusedMoE
29
from vllm.model_executor.layers.layernorm import RMSNorm
30
from vllm.model_executor.layers.linear import ReplicatedLinear
31
from vllm.model_executor.layers.logits_processor import LogitsProcessor
32
+from vllm.model_executor.layers.mhc import HCHeadOp
33
from vllm.model_executor.layers.vocab_parallel_embedding import (
34
VocabParallelEmbedding,
35
)
@@ -40,7 +41,6 @@ from .deepseek_mtp import SharedHead
41
from .deepseek_v2 import get_spec_layer_idx_from_weight_name
42
from .deepseek_v4 import (
43
DeepseekV4DecoderLayer,
43
- hc_head,
44
make_deepseek_v4_expert_params_mapping,
45
)
46
from .utils import maybe_prefix
@@ -119,6 +119,8 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
119
aux_stream_list=aux_stream_list,
120
)
121
122
+ self.hc_head_op = HCHeadOp()
123
+
124
def forward(
125
self,
126
input_ids: torch.Tensor,
@@ -230,7 +232,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
232
hidden_states = hidden_states.view(
233
-1, mtp_layer.hc_mult, mtp_layer.config.hidden_size
234
)
233
- hidden_states = hc_head(
235
+ hidden_states = self.hc_head_op(
236
hidden_states,
237
mtp_layer.hc_head_fn,
238
mtp_layer.hc_head_scale,