@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 313 lines 9.24 KB
Raw
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