@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 487 lines 13.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """AMD ROCm HyperConnection kernels for Qwen4Exp."""
4
5 import torch
6
7 from vllm.platforms import current_platform
8 from vllm.triton_utils import tl, triton
9 from vllm.utils.torch_utils import direct_register_custom_op
10
11
12 @triton.jit
13 def _grouped_gemma_rmsnorm_kernel(
14 x_ptr,
15 w_ptr,
16 y_ptr,
17 stride_x,
18 stride_y,
19 DIM: tl.constexpr,
20 NUM_GROUPS: tl.constexpr,
21 W_SHARED: tl.constexpr,
22 EPS: tl.constexpr,
23 launch_pdl: tl.constexpr,
24 ) -> None:
25 GROUP_DIM: tl.constexpr = DIM // NUM_GROUPS
26 BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(GROUP_DIM)
27
28 pid = tl.program_id(0)
29 group_id = pid % NUM_GROUPS
30 row = pid // NUM_GROUPS
31
32 offs_g = tl.arange(0, BLOCK_SIZE)
33 offsets = group_id * GROUP_DIM + offs_g
34 mask = offs_g < GROUP_DIM
35 # A [GROUP_DIM] affine is shared; a [DIM] affine follows the grouped
36 # checkpoint layout.
37 w_offs = offs_g if W_SHARED else offsets
38
39 if launch_pdl:
40 tl.extra.cuda.gdc_wait()
41
42 x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0).to(tl.float32)
43 w = tl.load(w_ptr + w_offs, mask, other=0.0)
44
45 rrms = tl.rsqrt(tl.sum(x * x) / GROUP_DIM + EPS)
46 # Gemma's (1 + w) affine is written this way to lower to an FMA.
47 y = x * rrms
48 y += y * w.to(tl.float32)
49
50 if launch_pdl:
51 tl.extra.cuda.gdc_launch_dependents()
52 tl.store(y_ptr + row * stride_y + offsets, y, mask)
53
54
55 def _grouped_gemma_rmsnorm(
56 x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
57 ) -> torch.Tensor:
58 N, DIM = x.shape
59 assert x.stride(1) == 1, "grouped Gemma RMSNorm requires unit inner stride"
60 assert weight.is_contiguous(), "grouped Gemma RMSNorm weight must be contiguous"
61 assert DIM % num_groups == 0
62 group_dim = DIM // num_groups
63 assert weight.numel() in (group_dim, DIM)
64
65 y = x.new_empty(x.shape)
66 _grouped_gemma_rmsnorm_kernel[(N * num_groups,)](
67 x,
68 weight,
69 y,
70 x.stride(0),
71 y.stride(0),
72 DIM,
73 num_groups,
74 W_SHARED=weight.numel() == group_dim,
75 EPS=eps,
76 launch_pdl=current_platform.is_arch_support_pdl(),
77 )
78 return y
79
80
81 @triton.jit
82 def _hc_silu_kernel(
83 x_ptr,
84 y_ptr,
85 stride_x,
86 stride_y,
87 DIM: tl.constexpr,
88 HC: tl.constexpr,
89 launch_pdl: tl.constexpr,
90 ) -> None:
91 BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(DIM)
92
93 row = tl.program_id(0)
94 offs = tl.arange(0, BLOCK_SIZE)
95 mask = offs < DIM
96
97 if launch_pdl:
98 tl.extra.cuda.gdc_wait()
99
100 x = tl.load(x_ptr + row * stride_x + offs, mask).to(tl.float32) / HC
101 y = x * tl.sigmoid(x)
102
103 if launch_pdl:
104 tl.extra.cuda.gdc_launch_dependents()
105 tl.store(y_ptr + row * stride_y + offs, y, mask)
106
107
108 def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
109 num_tokens, DIM = x.shape
110 assert x.stride(1) == 1
111
112 output = x.new_empty(x.shape)
113 _hc_silu_kernel[(num_tokens,)](
114 x,
115 output,
116 x.stride(0),
117 output.stride(0),
118 DIM=DIM,
119 HC=hc_count,
120 launch_pdl=current_platform.is_arch_support_pdl(),
121 )
122 return output
123
124
125 @triton.jit
126 def _hc_gate_mix_kernel(
127 x_ptr,
128 g_ptr,
129 y_ptr,
130 stride_x,
131 stride_g,
132 stride_y,
133 DIM: tl.constexpr,
134 HC: tl.constexpr,
135 BLOCK_SIZE: tl.constexpr,
136 launch_pdl: tl.constexpr,
137 ) -> None:
138 HC_DIM: tl.constexpr = DIM // HC
139
140 row = tl.program_id(0)
141 tile_id = tl.program_id(1)
142 offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
143 mask = offs_inner < HC_DIM
144
145 if launch_pdl:
146 tl.extra.cuda.gdc_wait()
147
148 # The constexpr loop is unrolled and keeps one stream live at a time.
149 # Materializing [HC, BLOCK_SIZE] more than doubles latency at large M.
150 acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
151 for stream in tl.static_range(HC):
152 offsets = stream * HC_DIM + offs_inner
153 g = tl.load(g_ptr + row * stride_g + offsets, mask, other=0.0)
154 x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0)
155 acc += tl.sigmoid(g.to(tl.float32)) * x.to(tl.float32)
156 acc /= HC
157
158 if launch_pdl:
159 tl.extra.cuda.gdc_launch_dependents()
160 tl.store(y_ptr + row * stride_y + offs_inner, acc, mask)
161
162
163 def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
164 N, DIM = gate.shape
165 assert x.shape == gate.shape
166 assert DIM % hc_count == 0
167 assert x.stride(1) == 1
168 assert gate.stride(1) == 1
169
170 HC_DIM = DIM // hc_count
171 out = x.new_empty(N, HC_DIM)
172 BLOCK_SIZE = 512
173 _hc_gate_mix_kernel[(N, triton.cdiv(HC_DIM, BLOCK_SIZE))](
174 x,
175 gate,
176 out,
177 x.stride(0),
178 gate.stride(0),
179 out.stride(0),
180 DIM,
181 hc_count,
182 BLOCK_SIZE,
183 launch_pdl=current_platform.is_arch_support_pdl(),
184 )
185 return out
186
187
188 @triton.jit
189 def _hc_combine_kernel(
190 block_ptr,
191 res_ptr,
192 inj_ptr,
193 out_ptr,
194 stride_block,
195 stride_res,
196 stride_inj,
197 stride_out,
198 HC_DIM: tl.constexpr,
199 HC: tl.constexpr,
200 BLOCK_SIZE: tl.constexpr,
201 launch_pdl: tl.constexpr,
202 ) -> None:
203 HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
204
205 row = tl.program_id(0)
206 tile_id = tl.program_id(1)
207
208 offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
209 mask_inner = offs_inner < HC_DIM
210 offs_hc = tl.arange(0, HC_PAD)
211 mask_hc = offs_hc < HC
212 offs = offs_hc[:, None] * HC_DIM + offs_inner[None, :]
213 mask = mask_hc[:, None] & mask_inner[None, :]
214
215 if launch_pdl:
216 tl.extra.cuda.gdc_wait()
217
218 inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)
219 block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)
220 res = tl.load(res_ptr + row * stride_res + offs, mask, other=0.0)
221
222 # Keeping HC as a broadcast dimension is faster here than four separate
223 # residual load/store sequences.
224 inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)
225 out = res.to(tl.float32) + block.to(tl.float32)[None, :] * inj[:, None]
226
227 if launch_pdl:
228 tl.extra.cuda.gdc_launch_dependents()
229 tl.store(out_ptr + row * stride_out + offs, out, mask=mask)
230
231
232 def _hc_combine(
233 residual: torch.Tensor,
234 block_output: torch.Tensor,
235 injection_logits: torch.Tensor,
236 hc_count: int,
237 ) -> torch.Tensor:
238 N, DIM = residual.shape
239 assert DIM % hc_count == 0
240 hc_dim = DIM // hc_count
241 assert block_output.shape == (N, hc_dim)
242 assert injection_logits.shape == (N, hc_count)
243 assert residual.stride(1) == 1
244 assert block_output.stride(1) == 1
245 assert injection_logits.stride(1) == 1
246
247 out = residual.new_empty(residual.shape)
248 BLOCK_SIZE = 512
249 _hc_combine_kernel[(N, triton.cdiv(hc_dim, BLOCK_SIZE))](
250 block_output,
251 residual,
252 injection_logits,
253 out,
254 block_output.stride(0),
255 residual.stride(0),
256 injection_logits.stride(0),
257 out.stride(0),
258 hc_dim,
259 hc_count,
260 BLOCK_SIZE,
261 launch_pdl=current_platform.is_arch_support_pdl(),
262 )
263 return out
264
265
266 @triton.jit
267 def _hc_combine_norm_kernel(
268 block_ptr,
269 res_ptr,
270 inj_ptr,
271 w_ptr,
272 out_ptr,
273 y_ptr,
274 stride_block,
275 stride_res,
276 stride_inj,
277 stride_out,
278 stride_y,
279 HC_DIM: tl.constexpr,
280 HC: tl.constexpr,
281 W_SHARED: tl.constexpr,
282 EPS: tl.constexpr,
283 BLOCK_SIZE: tl.constexpr,
284 launch_pdl: tl.constexpr,
285 ) -> None:
286 HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
287 NUM_TILES: tl.constexpr = triton.cdiv(HC_DIM, BLOCK_SIZE)
288 NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES)
289
290 row = tl.program_id(0)
291 stream = tl.program_id(1)
292 offs_hc = tl.arange(0, HC_PAD)
293 mask_hc = offs_hc < HC
294 tile_ids = tl.arange(0, NUM_TILES_PAD)
295 offs_inner = tile_ids[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :]
296 mask_inner = offs_inner < HC_DIM
297 offs = stream * HC_DIM + offs_inner
298 # Shared norm weights repeat across streams; per-branch weights use the
299 # same flattened HC layout as the residual.
300 w_offs = offs_inner if W_SHARED else offs
301
302 if launch_pdl:
303 tl.extra.cuda.gdc_wait()
304
305 # Start the uncached residual load first, then issue the other combine
306 # loads before consuming any of them.
307 res = tl.load(res_ptr + row * stride_res + offs, mask_inner, other=0.0)
308 inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)
309 block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)
310 inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)
311 inj = tl.sum(tl.where(offs_hc == stream, inj, 0.0))
312 # Round the materialized combine result before normalization. This matches
313 # the unfused combine -> RMSNorm boundary.
314 out = (res.to(tl.float32) + block.to(tl.float32) * inj).to(out_ptr.dtype.element_ty)
315 tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner)
316
317 out = out.to(tl.float32)
318 # Keep the two-axis reduction: flattening the padded tile is ~40% slower
319 # at decode sizes.
320 sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0)
321 rrms = tl.rsqrt(sum_sq / HC_DIM + EPS)
322
323 if launch_pdl:
324 tl.extra.cuda.gdc_launch_dependents()
325
326 # Loading the weight earlier helps decode but keeps the tile live across
327 # the reduction and regresses larger batches, so defer it to the norm.
328 w = tl.load(w_ptr + w_offs, mask_inner, other=0.0)
329 y = out * rrms
330 y += y * w.to(tl.float32)
331 tl.store(y_ptr + row * stride_y + offs, y, mask_inner)
332
333
334 def _hc_combine_norm(
335 residual: torch.Tensor,
336 block_output: torch.Tensor,
337 injection_logits: torch.Tensor,
338 norm_weight: torch.Tensor,
339 eps: float,
340 hc_count: int,
341 ) -> tuple[torch.Tensor, torch.Tensor]:
342 N, DIM = residual.shape
343 assert DIM % hc_count == 0
344 hc_dim = DIM // hc_count
345 assert block_output.shape == (N, hc_dim)
346 assert injection_logits.shape == (N, hc_count)
347 assert residual.stride(1) == 1
348 assert block_output.stride(1) == 1
349 assert injection_logits.stride(1) == 1
350 assert norm_weight.is_contiguous()
351 assert norm_weight.numel() in (hc_dim, DIM)
352
353 out = residual.new_empty(residual.shape)
354 y = residual.new_empty(residual.shape)
355 BLOCK_SIZE = 512
356 _hc_combine_norm_kernel[(N, hc_count)](
357 block_output,
358 residual,
359 injection_logits,
360 norm_weight,
361 out,
362 y,
363 block_output.stride(0),
364 residual.stride(0),
365 injection_logits.stride(0),
366 out.stride(0),
367 y.stride(0),
368 hc_dim,
369 hc_count,
370 W_SHARED=norm_weight.numel() == hc_dim,
371 EPS=eps,
372 BLOCK_SIZE=BLOCK_SIZE,
373 launch_pdl=current_platform.is_arch_support_pdl(),
374 )
375 return out, y
376
377
378 def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor:
379 return x.new_empty(x.shape)
380
381
382 def _hc_gate_mix_fake(
383 x: torch.Tensor, gate: torch.Tensor, hc_count: int
384 ) -> torch.Tensor:
385 del gate
386 return x.new_empty((x.shape[0], x.shape[1] // hc_count))
387
388
389 def _hc_combine_fake(
390 residual: torch.Tensor,
391 block_output: torch.Tensor,
392 injection_logits: torch.Tensor,
393 hc_count: int,
394 ) -> torch.Tensor:
395 del block_output, injection_logits, hc_count
396 return residual.new_empty(residual.shape)
397
398
399 def _hc_combine_norm_fake(
400 residual: torch.Tensor,
401 block_output: torch.Tensor,
402 injection_logits: torch.Tensor,
403 norm_weight: torch.Tensor,
404 eps: float,
405 hc_count: int,
406 ) -> tuple[torch.Tensor, torch.Tensor]:
407 del block_output, injection_logits, norm_weight, eps, hc_count
408 return residual.new_empty(residual.shape), residual.new_empty(residual.shape)
409
410
411 direct_register_custom_op(
412 op_name="qwen4_exp_grouped_gemma_rmsnorm",
413 op_func=_grouped_gemma_rmsnorm,
414 fake_impl=_same_shape_fake,
415 )
416 direct_register_custom_op(
417 op_name="qwen4_exp_hc_silu",
418 op_func=_hc_silu,
419 fake_impl=_same_shape_fake,
420 )
421 direct_register_custom_op(
422 op_name="qwen4_exp_hc_gate_mix",
423 op_func=_hc_gate_mix,
424 fake_impl=_hc_gate_mix_fake,
425 )
426 direct_register_custom_op(
427 op_name="qwen4_exp_hc_combine",
428 op_func=_hc_combine,
429 fake_impl=_hc_combine_fake,
430 )
431 direct_register_custom_op(
432 op_name="qwen4_exp_hc_combine_norm",
433 op_func=_hc_combine_norm,
434 fake_impl=_hc_combine_norm_fake,
435 )
436
437
438 def grouped_gemma_rmsnorm(
439 x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
440 ) -> torch.Tensor:
441 return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups)
442
443
444 def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
445 return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count)
446
447
448 def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
449 return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count)
450
451
452 def hc_combine(
453 residual: torch.Tensor,
454 block_output: torch.Tensor,
455 injection_logits: torch.Tensor,
456 hc_count: int,
457 ) -> torch.Tensor:
458 return torch.ops.vllm.qwen4_exp_hc_combine(
459 residual, block_output, injection_logits, hc_count
460 )
461
462
463 def hc_combine_norm(
464 residual: torch.Tensor,
465 block_output: torch.Tensor,
466 injection_logits: torch.Tensor,
467 norm_weight: torch.Tensor,
468 eps: float,
469 hc_count: int,
470 ) -> tuple[torch.Tensor, torch.Tensor]:
471 return torch.ops.vllm.qwen4_exp_hc_combine_norm(
472 residual,
473 block_output,
474 injection_logits,
475 norm_weight,
476 eps,
477 hc_count,
478 )
479
480
481 __all__ = [
482 "grouped_gemma_rmsnorm",
483 "hc_combine",
484 "hc_combine_norm",
485 "hc_gate_mix",
486 "hc_silu",
487 ]