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

[Kernel][HY V4] Add Triton iHC pre/post fallback (#55059)

Signed-off-by: linitra24 <renshuang.zhou@daocloud.io> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

linitra24 committed Sep 5, 2026 at 23:12 UTC 7985444339e2ad7e249b88a50081e16e34637dfd
3 files changed +531 -3
benchmarks/kernels/benchmark_hy_v4_ihc.py new
+199
@@ -0,0 +1,199 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark the HY V4 Triton iHC pre/post kernels against eager PyTorch."""
4 +
5 +import os
6 +import subprocess
7 +from functools import partial
8 +from importlib.metadata import version
9 +from statistics import median
10 +
11 +import torch
12 +import torch.nn.functional as F
13 +
14 +import vllm
15 +from vllm.models.hy_v4.nvidia.triton_ihc import (
16 + triton_ihc_post,
17 + triton_ihc_pre,
18 +)
19 +from vllm.triton_utils import triton
20 +from vllm.utils.argparse_utils import FlexibleArgumentParser
21 +
22 +
23 +def eager_pre(
24 + x: torch.Tensor,
25 + weight: torch.Tensor,
26 + scale: torch.Tensor,
27 + base: torch.Tensor,
28 + magnitude: float,
29 + hc_eps: float,
30 + norm_eps: float,
31 +) -> tuple[torch.Tensor, torch.Tensor]:
32 + num_tokens, hc_mult, hidden_size = x.shape
33 + x_flat = x.flatten(1).float()
34 + reciprocal_rms = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + norm_eps)
35 + mixes = F.linear(x_flat, weight) * reciprocal_rms
36 + pre = torch.sigmoid(mixes[:, :hc_mult] * scale[0] + base[:hc_mult])
37 + post = torch.sigmoid(mixes[:, hc_mult:] * scale[1] + base[hc_mult:])
38 + pre = pre + hc_eps
39 + post = magnitude * post + hc_eps
40 + output = torch.sum(pre.unsqueeze(-1) * x.float(), dim=1)
41 + return output.to(x.dtype).reshape(num_tokens, hidden_size), post
42 +
43 +
44 +def eager_post(
45 + x: torch.Tensor,
46 + residual: torch.Tensor,
47 + post: torch.Tensor,
48 +) -> torch.Tensor:
49 + return (post.float().unsqueeze(-1) * x.float().unsqueeze(-2) + residual.float()).to(
50 + x.dtype
51 + )
52 +
53 +
54 +def _timer(method: str):
55 + if method == "cupti":
56 + from flashinfer.testing import bench_gpu_time_with_cupti
57 +
58 + return partial(
59 + bench_gpu_time_with_cupti,
60 + use_cuda_graph=True,
61 + cold_l2_cache=True,
62 + )
63 + if method == "cudagraph":
64 + from flashinfer.testing import bench_gpu_time_with_cudagraph
65 +
66 + return partial(bench_gpu_time_with_cudagraph, cold_l2_cache=True)
67 + raise ValueError(f"unknown timing method: {method}")
68 +
69 +
70 +@torch.inference_mode()
71 +def run_benchmark(
72 + token_counts: list[int],
73 + hidden_size: int,
74 + dtype: torch.dtype,
75 + method: str,
76 +) -> None:
77 + hc_mult = 4
78 + device = torch.device("cuda")
79 + torch.manual_seed(0)
80 + weight = torch.randn(
81 + 2 * hc_mult,
82 + hc_mult * hidden_size,
83 + device=device,
84 + dtype=torch.float32,
85 + )
86 + scale = torch.randn(2, device=device, dtype=torch.float32) * 0.01
87 + base = torch.randn(2 * hc_mult, device=device, dtype=torch.float32)
88 + timer = _timer(method)
89 +
90 + properties = torch.cuda.get_device_properties(device)
91 + git_branch = subprocess.run(
92 + ["git", "branch", "--show-current"],
93 + check=False,
94 + capture_output=True,
95 + text=True,
96 + ).stdout.strip()
97 + git_commit = subprocess.run(
98 + ["git", "rev-parse", "--short", "HEAD"],
99 + check=False,
100 + capture_output=True,
101 + text=True,
102 + ).stdout.strip()
103 + print(f"device: {properties.name}")
104 + print(f"branch: {git_branch}; commit: {git_commit}")
105 + print(
106 + f"vllm: {vllm.__version__}; torch: {torch.__version__}; "
107 + f"CUDA: {torch.version.cuda}"
108 + )
109 + print(
110 + f"triton: {triton.__version__}; flashinfer: {version('flashinfer-python')}; "
111 + f"dtype: {dtype}; method: {method}; cache: cold L2"
112 + )
113 + print(
114 + "env: VLLM_ENABLE_HPC_OPS="
115 + f"{os.getenv('VLLM_ENABLE_HPC_OPS', '<unset>')}; "
116 + "VLLM_BATCH_INVARIANT="
117 + f"{os.getenv('VLLM_BATCH_INVARIANT', '<unset>')}"
118 + )
119 + print(f"hidden_size: {hidden_size}; hc_mult: {hc_mult}")
120 + print(
121 + f"{'tokens':>8} {'op':>6} {'eager (us)':>12} "
122 + f"{'triton (us)':>12} {'speedup':>9} {'GiB':>8} "
123 + f"{'eager GB/s':>12} {'triton GB/s':>13}"
124 + )
125 +
126 + for num_tokens in token_counts:
127 + x = torch.randn(
128 + num_tokens,
129 + hc_mult,
130 + hidden_size,
131 + device=device,
132 + dtype=dtype,
133 + )
134 + block_output = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
135 + residual = torch.randn_like(x)
136 + eager_output, post = eager_pre(x, weight, scale, base, 2.0, 1e-6, 1e-5)
137 + triton_output, triton_post = triton_ihc_pre(
138 + x, weight, scale, base, 2.0, 1e-6, 1e-5
139 + )
140 + torch.testing.assert_close(triton_output, eager_output, atol=2e-2, rtol=1e-2)
141 + torch.testing.assert_close(triton_post, post, atol=2e-5, rtol=1e-5)
142 + torch.testing.assert_close(
143 + triton_ihc_post(block_output, residual, post),
144 + eager_post(block_output, residual, post),
145 + atol=0,
146 + rtol=0,
147 + )
148 +
149 + benchmarks = (
150 + (
151 + "pre",
152 + partial(eager_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5),
153 + partial(triton_ihc_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5),
154 + x.nbytes
155 + + weight.nbytes
156 + + scale.nbytes
157 + + base.nbytes
158 + + eager_output.nbytes
159 + + post.nbytes,
160 + ),
161 + (
162 + "post",
163 + partial(eager_post, block_output, residual, post),
164 + partial(triton_ihc_post, block_output, residual, post),
165 + block_output.nbytes + residual.nbytes + post.nbytes + residual.nbytes,
166 + ),
167 + )
168 + for op_name, eager_fn, triton_fn, logical_bytes in benchmarks:
169 + eager_us = median(timer(eager_fn)) * 1e3
170 + triton_us = median(timer(triton_fn)) * 1e3
171 + speedup = eager_us / triton_us
172 + logical_gib = logical_bytes / 2**30
173 + eager_gbps = logical_bytes / eager_us / 1e3
174 + triton_gbps = logical_bytes / triton_us / 1e3
175 + print(
176 + f"{num_tokens:>8} {op_name:>6} {eager_us:>12.1f} "
177 + f"{triton_us:>12.1f} {speedup:>8.2f}x {logical_gib:>8.3f} "
178 + f"{eager_gbps:>12.1f} {triton_gbps:>13.1f}"
179 + )
180 +
181 +
182 +if __name__ == "__main__":
183 + parser = FlexibleArgumentParser(description=__doc__)
184 + parser.add_argument(
185 + "--token-counts",
186 + type=int,
187 + nargs="+",
188 + default=[1, 2, 4, 8, 16, 64, 256, 1024, 4096, 8192],
189 + )
190 + parser.add_argument("--hidden-size", type=int, choices=[4096, 6144], default=6144)
191 + parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="bfloat16")
192 + parser.add_argument("--method", choices=["cupti", "cudagraph"], default="cupti")
193 + args = parser.parse_args()
194 + run_benchmark(
195 + args.token_counts,
196 + args.hidden_size,
197 + getattr(torch, args.dtype),
198 + args.method,
199 + )
vllm/models/hy_v4/nvidia/hc.py
+19 -3
@@ -9,9 +9,8 @@ scatters the result back over the channels (``HYV4HCPostLayer``). The final
9 ``HYV4HCHeadLayer`` merges the channels before the model's output norm.
10
11 NOTE: Each of the three steps has an optional single-kernel HPC replacement
12 -(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). They are only constructed
13 -when the hpc package is installed, ``VLLM_ENABLE_HPC_OPS=1`` and the shape /
14 -device constraints hold; otherwise the eager path below runs unchanged.
12 +(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). Pre and post fall back to
13 +in-tree Triton kernels on CUDA when HPC is unavailable, then to the eager path.
14 TODO: port the cross-layer post+pre fusion (``HpcIHCPostPre``) as well; it
15 requires restructuring the decoder-layer forward scheduling.
16 """
@@ -22,6 +21,11 @@ from transformers import PretrainedConfig
21
22 from vllm.model_executor.layers.hpc import HpcIHCHead, HpcIHCPost, HpcIHCPre
23 from vllm.model_executor.layers.linear import ReplicatedLinear
24 +from vllm.models.hy_v4.nvidia.triton_ihc import (
25 + triton_ihc_post,
26 + triton_ihc_pre,
27 + triton_ihc_supported,
28 +)
29
30
31 class HYV4HCPreLayer(nn.Module):
@@ -106,6 +110,16 @@ class HYV4HCPreLayer(nn.Module):
110 """
111 if self.hpc_op is not None:
112 return self.hpc_op(x)
113 + if triton_ihc_supported(x):
114 + return triton_ihc_pre(
115 + x,
116 + self.hc_fn.weight,
117 + self.hc_scale,
118 + self.hc_base,
119 + self.magnitude,
120 + self.hc_eps,
121 + self.layernorm_epsilon,
122 + )
123
124 shape = x.size() # [num_tokens, hc, d]
125 hc = self.hc_mult
@@ -175,6 +189,8 @@ class HYV4HCPostLayer(nn.Module):
189 """
190 if self.hpc_op is not None:
191 return self.hpc_op(x, residual, post)
192 + if triton_ihc_supported(x):
193 + return triton_ihc_post(x, residual, post)
194
195 dtype = x.dtype
196 x = x.float()
vllm/models/hy_v4/nvidia/triton_ihc.py new
+313
@@ -0,0 +1,313 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton iHC pre/post kernels for HY V4.
4 +
5 +Adapted from the SGLang HY V4 implementation:
6 +https://github.com/sgl-project/sglang/pull/36805
7 +"""
8 +
9 +import torch
10 +
11 +from vllm import envs
12 +from vllm.platforms import current_platform
13 +from vllm.triton_utils import HAS_TRITON, tl, triton
14 +
15 +_BLOCK_K = 1024
16 +_BLOCK_D = 1024
17 +
18 +
19 +def triton_ihc_supported(x: torch.Tensor) -> bool:
20 + """Return whether the in-tree Triton path can run for this input."""
21 + return (
22 + HAS_TRITON
23 + and current_platform.is_cuda()
24 + and x.is_cuda
25 + and x.dtype in (torch.float16, torch.bfloat16)
26 + and not envs.VLLM_BATCH_INVARIANT
27 + )
28 +
29 +
30 +@triton.jit
31 +def _ihc_pre_stage1(
32 + x_ptr,
33 + weight_ptr,
34 + partial_ptr,
35 + K_TOTAL: tl.constexpr,
36 + HC_MULT: tl.constexpr,
37 + HC_POW2: tl.constexpr,
38 + NUM_SPLITS: tl.constexpr,
39 + BLOCK_K: tl.constexpr,
40 + PARTIAL_STRIDE: tl.constexpr,
41 +):
42 + token_idx = tl.program_id(0).to(tl.int64)
43 + split_idx = tl.program_id(1)
44 + hc_idx = tl.arange(0, HC_POW2)
45 + hc_mask = hc_idx < HC_MULT
46 +
47 + k_offsets = split_idx * BLOCK_K + tl.arange(0, BLOCK_K)
48 + k_mask = k_offsets < K_TOTAL
49 + x = tl.load(
50 + x_ptr + token_idx * K_TOTAL + k_offsets,
51 + mask=k_mask,
52 + other=0.0,
53 + ).to(tl.float32)
54 + sum_squares = tl.sum(x * x, axis=0)
55 +
56 + weight_offsets = hc_idx[:, None] * K_TOTAL + k_offsets[None, :]
57 + weight_mask = hc_mask[:, None] & k_mask[None, :]
58 + pre_mix = tl.sum(
59 + tl.load(weight_ptr + weight_offsets, mask=weight_mask, other=0.0) * x[None, :],
60 + axis=1,
61 + )
62 + post_mix = tl.sum(
63 + tl.load(
64 + weight_ptr + HC_MULT * K_TOTAL + weight_offsets,
65 + mask=weight_mask,
66 + other=0.0,
67 + )
68 + * x[None, :],
69 + axis=1,
70 + )
71 +
72 + partial = partial_ptr + (token_idx * NUM_SPLITS + split_idx) * PARTIAL_STRIDE
73 + tl.store(partial, sum_squares)
74 + tl.store(partial + 1 + hc_idx, pre_mix, mask=hc_mask)
75 + tl.store(partial + 1 + HC_POW2 + hc_idx, post_mix, mask=hc_mask)
76 +
77 +
78 +@triton.jit
79 +def _ihc_pre_stage2(
80 + x_ptr,
81 + partial_ptr,
82 + scale_ptr,
83 + base_ptr,
84 + output_ptr,
85 + post_ptr,
86 + HIDDEN_SIZE: tl.constexpr,
87 + K_TOTAL: tl.constexpr,
88 + HC_MULT: tl.constexpr,
89 + HC_POW2: tl.constexpr,
90 + NUM_SPLITS: tl.constexpr,
91 + PARTIAL_STRIDE: tl.constexpr,
92 + BLOCK_D: tl.constexpr,
93 + MAGNITUDE: tl.constexpr,
94 + NORM_EPS: tl.constexpr,
95 + HC_EPS: tl.constexpr,
96 +):
97 + token_idx = tl.program_id(0).to(tl.int64)
98 + hidden_block_idx = tl.program_id(1)
99 + hc_idx = tl.arange(0, HC_POW2)
100 + hc_mask = hc_idx < HC_MULT
101 +
102 + partial_row = partial_ptr + token_idx * NUM_SPLITS * PARTIAL_STRIDE
103 + sum_squares = tl.zeros((), dtype=tl.float32)
104 + pre_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
105 + post_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
106 + for split_idx in tl.static_range(NUM_SPLITS):
107 + partial = partial_row + split_idx * PARTIAL_STRIDE
108 + sum_squares += tl.load(partial)
109 + pre_mix += tl.load(partial + 1 + hc_idx, mask=hc_mask, other=0.0)
110 + post_mix += tl.load(
111 + partial + 1 + HC_POW2 + hc_idx,
112 + mask=hc_mask,
113 + other=0.0,
114 + )
115 +
116 + reciprocal_rms = tl.rsqrt(sum_squares / K_TOTAL + NORM_EPS)
117 + pre_scale = tl.load(scale_ptr)
118 + post_scale = tl.load(scale_ptr + 1)
119 + pre_base = tl.load(base_ptr + hc_idx, mask=hc_mask, other=0.0)
120 + post_base = tl.load(base_ptr + HC_MULT + hc_idx, mask=hc_mask, other=0.0)
121 + pre = tl.sigmoid(pre_mix * reciprocal_rms * pre_scale + pre_base) + HC_EPS
122 +
123 + if hidden_block_idx == 0:
124 + post = (
125 + MAGNITUDE * tl.sigmoid(post_mix * reciprocal_rms * post_scale + post_base)
126 + + HC_EPS
127 + )
128 + tl.store(
129 + post_ptr + token_idx * HC_MULT + hc_idx,
130 + post,
131 + mask=hc_mask,
132 + )
133 +
134 + hidden_offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
135 + hidden_mask = hidden_offsets < HIDDEN_SIZE
136 + x_row = x_ptr + token_idx * K_TOTAL
137 + output = tl.zeros((BLOCK_D,), dtype=tl.float32)
138 + for channel_idx in tl.static_range(HC_MULT):
139 + channel = tl.load(
140 + x_row + channel_idx * HIDDEN_SIZE + hidden_offsets,
141 + mask=hidden_mask,
142 + other=0.0,
143 + ).to(tl.float32)
144 + channel_gate = tl.sum(tl.where(hc_idx == channel_idx, pre, 0.0), axis=0)
145 + output += channel_gate * channel
146 +
147 + tl.store(
148 + output_ptr + token_idx * HIDDEN_SIZE + hidden_offsets,
149 + output.to(output_ptr.dtype.element_ty),
150 + mask=hidden_mask,
151 + )
152 +
153 +
154 +@triton.jit
155 +def _ihc_post_kernel(
156 + x_ptr,
157 + residual_ptr,
158 + post_ptr,
159 + output_ptr,
160 + HIDDEN_SIZE: tl.constexpr,
161 + HC_MULT: tl.constexpr,
162 + HC_POW2: tl.constexpr,
163 + BLOCK_D: tl.constexpr,
164 +):
165 + token_idx = tl.program_id(0).to(tl.int64)
166 + hidden_block_idx = tl.program_id(1)
167 + hc_idx = tl.arange(0, HC_POW2)
168 + hc_mask = hc_idx < HC_MULT
169 + post = tl.load(
170 + post_ptr + token_idx * HC_MULT + hc_idx,
171 + mask=hc_mask,
172 + other=0.0,
173 + )
174 +
175 + hidden_offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
176 + hidden_mask = hidden_offsets < HIDDEN_SIZE
177 + x = tl.load(
178 + x_ptr + token_idx * HIDDEN_SIZE + hidden_offsets,
179 + mask=hidden_mask,
180 + other=0.0,
181 + ).to(tl.float32)
182 + residual_row = residual_ptr + token_idx * HC_MULT * HIDDEN_SIZE
183 + output_row = output_ptr + token_idx * HC_MULT * HIDDEN_SIZE
184 + for channel_idx in tl.static_range(HC_MULT):
185 + residual = tl.load(
186 + residual_row + channel_idx * HIDDEN_SIZE + hidden_offsets,
187 + mask=hidden_mask,
188 + other=0.0,
189 + ).to(tl.float32)
190 + channel_gate = tl.sum(tl.where(hc_idx == channel_idx, post, 0.0), axis=0)
191 + output = channel_gate * x + residual
192 + tl.store(
193 + output_row + channel_idx * HIDDEN_SIZE + hidden_offsets,
194 + output.to(output_ptr.dtype.element_ty),
195 + mask=hidden_mask,
196 + )
197 +
198 +
199 +def triton_ihc_pre(
200 + x: torch.Tensor,
201 + weight: torch.Tensor,
202 + scale: torch.Tensor,
203 + base: torch.Tensor,
204 + magnitude: float,
205 + hc_eps: float,
206 + norm_eps: float,
207 +) -> tuple[torch.Tensor, torch.Tensor]:
208 + """Reduce iHC channels and produce the post gates."""
209 + assert x.dim() == 3
210 + assert x.is_cuda and weight.is_cuda and scale.is_cuda and base.is_cuda
211 + assert weight.dtype == torch.float32
212 + assert scale.dtype == torch.float32 and base.dtype == torch.float32
213 +
214 + x = x.contiguous()
215 + assert weight.is_contiguous()
216 + scale = scale.contiguous()
217 + base = base.contiguous()
218 + num_tokens, hc_mult, hidden_size = x.shape
219 + k_total = hc_mult * hidden_size
220 + assert weight.shape == (2 * hc_mult, k_total)
221 + assert scale.shape == (2,)
222 + assert base.shape == (2 * hc_mult,)
223 +
224 + if num_tokens == 0:
225 + return (
226 + torch.empty((0, hidden_size), dtype=x.dtype, device=x.device),
227 + torch.empty((0, hc_mult), dtype=torch.float32, device=x.device),
228 + )
229 +
230 + output = torch.empty((num_tokens, hidden_size), dtype=x.dtype, device=x.device)
231 + post = torch.empty((num_tokens, hc_mult), dtype=torch.float32, device=x.device)
232 + hc_pow2 = triton.next_power_of_2(hc_mult)
233 + num_splits = triton.cdiv(k_total, _BLOCK_K)
234 + partial_stride = 1 + 2 * hc_pow2
235 + partial = torch.empty(
236 + (num_tokens, num_splits, partial_stride),
237 + dtype=torch.float32,
238 + device=x.device,
239 + )
240 +
241 + _ihc_pre_stage1[(num_tokens, num_splits)](
242 + x,
243 + weight,
244 + partial,
245 + K_TOTAL=k_total,
246 + HC_MULT=hc_mult,
247 + HC_POW2=hc_pow2,
248 + NUM_SPLITS=num_splits,
249 + BLOCK_K=_BLOCK_K,
250 + PARTIAL_STRIDE=partial_stride,
251 + num_warps=8,
252 + enable_fp_fusion=False,
253 + )
254 + _ihc_pre_stage2[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))](
255 + x,
256 + partial,
257 + scale,
258 + base,
259 + output,
260 + post,
261 + HIDDEN_SIZE=hidden_size,
262 + K_TOTAL=k_total,
263 + HC_MULT=hc_mult,
264 + HC_POW2=hc_pow2,
265 + NUM_SPLITS=num_splits,
266 + PARTIAL_STRIDE=partial_stride,
267 + BLOCK_D=_BLOCK_D,
268 + MAGNITUDE=magnitude,
269 + NORM_EPS=norm_eps,
270 + HC_EPS=hc_eps,
271 + num_warps=4,
272 + enable_fp_fusion=False,
273 + )
274 + return output, post
275 +
276 +
277 +def triton_ihc_post(
278 + x: torch.Tensor,
279 + residual: torch.Tensor,
280 + post: torch.Tensor,
281 +) -> torch.Tensor:
282 + """Scatter a sub-block output back over the iHC residual channels."""
283 + assert x.dim() == 2
284 + assert x.is_cuda and residual.is_cuda and post.is_cuda
285 + assert post.dtype == torch.float32
286 +
287 + x = x.contiguous()
288 + residual = residual.contiguous()
289 + post = post.contiguous()
290 + num_tokens, hidden_size = x.shape
291 + hc_mult = post.shape[-1]
292 + assert residual.shape == (num_tokens, hc_mult, hidden_size)
293 + assert post.shape == (num_tokens, hc_mult)
294 +
295 + if num_tokens == 0:
296 + return torch.empty((0, hc_mult, hidden_size), dtype=x.dtype, device=x.device)
297 +
298 + output = torch.empty(
299 + (num_tokens, hc_mult, hidden_size), dtype=x.dtype, device=x.device
300 + )
301 + _ihc_post_kernel[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))](
302 + x,
303 + residual,
304 + post,
305 + output,
306 + HIDDEN_SIZE=hidden_size,
307 + HC_MULT=hc_mult,
308 + HC_POW2=triton.next_power_of_2(hc_mult),
309 + BLOCK_D=_BLOCK_D,
310 + num_warps=4,
311 + enable_fp_fusion=False,
312 + )
313 + return output