@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,125 lines 34.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 from vllm.utils.math_utils import cdiv
6 from vllm.utils.torch_utils import direct_register_custom_op
7
8
9 def _torch_hc_prenorm_gemm(
10 x: torch.Tensor,
11 fn: torch.Tensor,
12 out: torch.Tensor,
13 sqrsum: torch.Tensor,
14 ) -> None:
15 assert out.shape[0] == 1
16 assert sqrsum.shape[0] == 1
17 x_float = x.float()
18 out[0].copy_(x_float @ fn.t())
19 sqrsum[0].copy_(x_float.square().sum(dim=-1))
20
21
22 def _hc_prenorm_gemm_outputs(
23 x: torch.Tensor,
24 fn: torch.Tensor,
25 *,
26 hidden_size: int,
27 hc_mult: int,
28 use_tilelang_fallback: bool = True,
29 ) -> tuple[torch.Tensor, torch.Tensor]:
30 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
31 compute_num_split,
32 )
33 from vllm.utils.deep_gemm import (
34 is_deep_gemm_supported,
35 tf32_hc_prenorm_gemm,
36 )
37
38 use_deep_gemm = is_deep_gemm_supported() or not use_tilelang_fallback
39 num_tokens = x.shape[0]
40 n_splits = (
41 compute_num_split(64, x.shape[1], cdiv(num_tokens, 64)) if use_deep_gemm else 1
42 )
43 out = torch.empty(
44 n_splits,
45 num_tokens,
46 fn.shape[0],
47 dtype=torch.float32,
48 device=x.device,
49 )
50 sqrsum = torch.empty(
51 n_splits,
52 num_tokens,
53 dtype=torch.float32,
54 device=x.device,
55 )
56 if use_deep_gemm:
57 tf32_hc_prenorm_gemm(x, fn, out, sqrsum, n_splits)
58 else:
59 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
60 _HC_PRENORM_GEMM_TILELANG_KERNEL,
61 )
62
63 _HC_PRENORM_GEMM_TILELANG_KERNEL(
64 x,
65 fn,
66 out,
67 sqrsum,
68 hidden_size,
69 hc_mult,
70 )
71 return out, sqrsum
72
73
74 def mhc_pre_delayed_tilelang(
75 residual: torch.Tensor,
76 fn: torch.Tensor,
77 hc_scale: torch.Tensor,
78 hc_base: torch.Tensor,
79 rms_eps: float,
80 hc_pre_eps: float,
81 hc_sinkhorn_eps: float,
82 hc_post_mult_value: float,
83 sinkhorn_repeat: int,
84 pre_mix: torch.Tensor | None = None,
85 x: torch.Tensor | None = None,
86 norm_weight: torch.Tensor | None = None,
87 norm_eps: float = 1e-6,
88 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
89 """Run mHC pre with a carried pre-mix and return the next pre-mix.
90
91 Args:
92 residual: BF16 residual streams of shape (tokens, hc_mult, hidden_size).
93 fn: FP32 projection of shape (hc_mult * (hc_mult + 2), input_size).
94 hc_scale: FP32 scales of shape (3,).
95 hc_base: FP32 bias of shape (hc_mult * (hc_mult + 2),).
96 rms_eps: RMS normalization epsilon.
97 hc_pre_eps: Pre-mix epsilon.
98 hc_sinkhorn_eps: Sinkhorn epsilon.
99 hc_post_mult_value: Post-mix multiplier.
100 sinkhorn_repeat: Number of Sinkhorn iterations.
101 pre_mix: FP32 coefficients from the previous sublayer, or None to
102 select residual stream zero at model entry.
103 x: Optional BF16 projection input of shape (tokens, input_size), for
104 the first layer's broadcast embedding and summed projection.
105 norm_weight: Optional BF16 RMSNorm weight for the collapsed input.
106 norm_eps: RMSNorm epsilon for the collapsed input.
107
108 Returns:
109 Post and residual coefficients, optionally normalized BF16 layer input,
110 and the next FP32 pre-mix, with shapes (tokens, hc_mult, 1),
111 (tokens, hc_mult, hc_mult), (tokens, hidden_size), and (tokens, hc_mult).
112
113 """
114 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
115 _HC_PRENORM_GEMM_TILELANG_KERNEL,
116 mhc_pre_big_fuse_tilelang,
117 )
118 from vllm.model_executor.kernels.mhc.warmup import (
119 MHC_PRE_NORM_KERNEL,
120 compute_mhc_pre_num_splits,
121 )
122 from vllm.utils.deep_gemm import (
123 is_deep_gemm_supported,
124 tf32_hc_prenorm_gemm,
125 )
126
127 assert residual.ndim == 3 and residual.dtype == torch.bfloat16
128 assert residual.is_contiguous()
129 num_tokens, hc_mult, hidden_size = residual.shape
130 if x is None:
131 x = residual.view(num_tokens, hc_mult * hidden_size)
132 assert x.ndim == 2 and x.dtype == torch.bfloat16 and x.is_contiguous()
133 assert x.shape[0] == num_tokens
134 input_size = x.shape[1]
135 mix_size = hc_mult * (hc_mult + 2)
136 assert fn.shape == (mix_size, input_size) and fn.dtype == torch.float32
137 assert hc_scale.shape == (3,) and hc_scale.dtype == torch.float32
138 assert hc_base.shape == (mix_size,) and hc_base.dtype == torch.float32
139 if pre_mix is not None:
140 assert pre_mix.shape == (num_tokens, hc_mult)
141 assert pre_mix.dtype == torch.float32 and pre_mix.is_contiguous()
142
143 next_pre_mix = torch.empty(
144 num_tokens, hc_mult, dtype=torch.float32, device=residual.device
145 )
146 post = torch.empty_like(next_pre_mix)
147 comb = torch.empty(
148 num_tokens, hc_mult * hc_mult, dtype=torch.float32, device=residual.device
149 )
150 layer_input = torch.empty(
151 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
152 )
153 outputs = (
154 post.unsqueeze(-1),
155 comb.view(num_tokens, hc_mult, hc_mult),
156 layer_input,
157 next_pre_mix,
158 )
159 if num_tokens == 0:
160 return outputs
161
162 use_deep_gemm = is_deep_gemm_supported()
163 n_splits = (
164 compute_mhc_pre_num_splits(input_size, num_tokens) if use_deep_gemm else 1
165 )
166 mixes = torch.empty(
167 n_splits, num_tokens, mix_size, dtype=torch.float32, device=residual.device
168 )
169 sqrsum = torch.empty(
170 n_splits, num_tokens, dtype=torch.float32, device=residual.device
171 )
172 if use_deep_gemm:
173 tf32_hc_prenorm_gemm(x, fn, mixes, sqrsum, n_splits)
174 else:
175 _HC_PRENORM_GEMM_TILELANG_KERNEL(
176 x,
177 fn,
178 mixes,
179 sqrsum,
180 input_size,
181 1,
182 )
183 if norm_weight is not None:
184 assert norm_weight.shape == (hidden_size,)
185 assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous()
186 MHC_PRE_NORM_KERNEL(
187 mixes,
188 sqrsum,
189 hc_scale,
190 hc_base,
191 residual,
192 post,
193 comb,
194 layer_input,
195 norm_weight,
196 pre_mix if pre_mix is not None else post,
197 next_pre_mix,
198 layer_input,
199 hidden_size=hidden_size,
200 rms_eps=rms_eps,
201 hc_pre_eps=hc_pre_eps,
202 hc_sinkhorn_eps=hc_sinkhorn_eps,
203 hc_post_mult_value=hc_post_mult_value,
204 sinkhorn_repeat=sinkhorn_repeat,
205 norm_eps=norm_eps,
206 hc_mult=hc_mult,
207 use_pre_mix_in=pre_mix is not None,
208 save_pre_mix=True,
209 rms_numel=input_size,
210 )
211 return outputs
212 mhc_pre_big_fuse_tilelang(
213 mixes,
214 sqrsum,
215 hc_scale,
216 hc_base,
217 residual,
218 post,
219 comb,
220 layer_input,
221 pre_mix if pre_mix is not None else post,
222 next_pre_mix,
223 layer_input,
224 hidden_size,
225 rms_eps,
226 hc_pre_eps,
227 hc_sinkhorn_eps,
228 hc_post_mult_value,
229 sinkhorn_repeat,
230 n_splits,
231 hc_mult,
232 use_pre_mix_in=pre_mix is not None,
233 save_pre_mix=True,
234 rms_numel=input_size,
235 )
236 return outputs
237
238
239 def _mhc_pre_delayed_tilelang_fake(
240 residual: torch.Tensor,
241 fn: torch.Tensor,
242 hc_scale: torch.Tensor,
243 hc_base: torch.Tensor,
244 rms_eps: float,
245 hc_pre_eps: float,
246 hc_sinkhorn_eps: float,
247 hc_post_mult_value: float,
248 sinkhorn_repeat: int,
249 pre_mix: torch.Tensor | None = None,
250 x: torch.Tensor | None = None,
251 norm_weight: torch.Tensor | None = None,
252 norm_eps: float = 1e-6,
253 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
254 num_tokens, hc_mult, hidden_size = residual.shape
255 return (
256 torch.empty(
257 num_tokens, hc_mult, 1, dtype=torch.float32, device=residual.device
258 ),
259 torch.empty(
260 num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=residual.device
261 ),
262 torch.empty(
263 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
264 ),
265 torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=residual.device),
266 )
267
268
269 def mhc_fused_post_pre_delayed_tilelang(
270 x: torch.Tensor,
271 residual: torch.Tensor,
272 post_layer_mix: torch.Tensor,
273 comb_res_mix: torch.Tensor,
274 fn: torch.Tensor,
275 hc_scale: torch.Tensor,
276 hc_base: torch.Tensor,
277 rms_eps: float,
278 hc_pre_eps: float,
279 hc_sinkhorn_eps: float,
280 hc_post_mult_value: float,
281 sinkhorn_repeat: int,
282 pre_mix: torch.Tensor | None = None,
283 norm_weight: torch.Tensor | None = None,
284 norm_eps: float = 1e-6,
285 capture_aux: bool = False,
286 ) -> tuple[
287 torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
288 ]:
289 """Run one mHC post block followed by the next delayed mHC pre block.
290
291 Within the fused kernel's token range the post mapping is folded into the
292 pre-norm GEMM, so the updated residual streams feed the projection from
293 registers instead of a second pass over global memory. Above it this runs
294 the same post kernel and split-k GEMM as the unfused pair.
295
296 Args:
297 x: BF16 sublayer output of shape (tokens, hidden_size).
298 residual: BF16 residual streams of shape (tokens, hc_mult, hidden_size).
299 post_layer_mix: FP32 post coefficients of shape (tokens, hc_mult, 1).
300 comb_res_mix: FP32 residual coefficients, (tokens, hc_mult, hc_mult).
301 fn: FP32 projection of shape (hc_mult * (hc_mult + 2), input_size).
302 hc_scale: FP32 scales of shape (3,).
303 hc_base: FP32 bias of shape (hc_mult * (hc_mult + 2),).
304 rms_eps: RMS normalization epsilon.
305 hc_pre_eps: Pre-mix epsilon.
306 hc_sinkhorn_eps: Sinkhorn epsilon.
307 hc_post_mult_value: Post-mix multiplier.
308 sinkhorn_repeat: Number of Sinkhorn iterations.
309 pre_mix: FP32 coefficients from the previous sublayer, or None to
310 select residual stream zero.
311 norm_weight: Optional BF16 RMSNorm weight for the collapsed input.
312 norm_eps: RMSNorm epsilon for the collapsed input.
313 capture_aux: Also return the mean over the post-mapped streams, which
314 draft models consume as the target's hidden state. It is folded
315 into the collapse, which already reads those streams.
316
317 Returns:
318 The post-mapped residual streams, the post and residual coefficients,
319 the optionally normalized BF16 layer input, the next FP32 pre-mix, and
320 the BF16 stream mean (empty unless capture_aux), with shapes
321 (tokens, hc_mult, hidden_size), (tokens, hc_mult, 1),
322 (tokens, hc_mult, hc_mult), (tokens, hidden_size), (tokens, hc_mult),
323 and (tokens, hidden_size).
324
325 """
326 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
327 _HC_PRENORM_GEMM_TILELANG_KERNEL,
328 _MHC_FUSED_TILELANG_KERNEL,
329 _MHC_POST_TILELANG_KERNEL,
330 mhc_fused_post_pre_split_config,
331 mhc_pre_big_fuse_tilelang,
332 )
333 from vllm.model_executor.kernels.mhc.warmup import (
334 MHC_PRE_NORM_KERNEL,
335 compute_mhc_pre_num_splits,
336 )
337 from vllm.utils.deep_gemm import (
338 is_deep_gemm_supported,
339 tf32_hc_prenorm_gemm,
340 )
341
342 assert residual.ndim == 3 and residual.dtype == torch.bfloat16
343 assert residual.is_contiguous()
344 num_tokens, hc_mult, hidden_size = residual.shape
345 input_size = hc_mult * hidden_size
346 mix_size = hc_mult * (hc_mult + 2)
347 assert x.shape == (num_tokens, hidden_size) and x.dtype == torch.bfloat16
348 assert x.is_contiguous()
349 assert post_layer_mix.shape[:2] == (num_tokens, hc_mult)
350 assert post_layer_mix.dtype == torch.float32 and post_layer_mix.is_contiguous()
351 assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult)
352 assert comb_res_mix.dtype == torch.float32 and comb_res_mix.is_contiguous()
353 assert fn.shape == (mix_size, input_size) and fn.dtype == torch.float32
354 assert hc_scale.shape == (3,) and hc_scale.dtype == torch.float32
355 assert hc_base.shape == (mix_size,) and hc_base.dtype == torch.float32
356 if pre_mix is not None:
357 assert pre_mix.shape == (num_tokens, hc_mult)
358 assert pre_mix.dtype == torch.float32 and pre_mix.is_contiguous()
359
360 next_pre_mix = torch.empty(
361 num_tokens, hc_mult, dtype=torch.float32, device=residual.device
362 )
363 post = torch.empty_like(next_pre_mix)
364 comb = torch.empty(
365 num_tokens, hc_mult * hc_mult, dtype=torch.float32, device=residual.device
366 )
367 layer_input = torch.empty(
368 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
369 )
370 aux = torch.empty(
371 num_tokens if capture_aux else 0,
372 hidden_size,
373 dtype=torch.bfloat16,
374 device=residual.device,
375 )
376 if num_tokens == 0:
377 return (
378 torch.empty_like(residual),
379 post.unsqueeze(-1),
380 comb.view(num_tokens, hc_mult, hc_mult),
381 layer_input,
382 next_pre_mix,
383 aux,
384 )
385
386 fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
387 if fused_config is not None:
388 mixes, sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL(
389 comb_res_mix,
390 residual,
391 post_layer_mix.view(num_tokens, hc_mult),
392 x,
393 fn.view(mix_size, hc_mult, hidden_size),
394 hc_mult,
395 hidden_size,
396 mix_size,
397 )
398 else:
399 residual_cur = torch.empty_like(residual)
400 _MHC_POST_TILELANG_KERNEL(
401 comb_res_mix,
402 residual,
403 post_layer_mix.view(num_tokens, hc_mult),
404 x,
405 residual_cur,
406 hc_mult,
407 hidden_size,
408 )
409 # The delayed epilogue is compiled per bucketed split count, so the
410 # projection has to use the same bucket rather than the raw estimate.
411 use_deep_gemm = is_deep_gemm_supported()
412 n_splits = (
413 compute_mhc_pre_num_splits(input_size, num_tokens) if use_deep_gemm else 1
414 )
415 mixes = torch.empty(
416 n_splits, num_tokens, mix_size, dtype=torch.float32, device=residual.device
417 )
418 sqrsum = torch.empty(
419 n_splits, num_tokens, dtype=torch.float32, device=residual.device
420 )
421 residual_cur_2d = residual_cur.view(num_tokens, input_size)
422 if use_deep_gemm:
423 tf32_hc_prenorm_gemm(residual_cur_2d, fn, mixes, sqrsum, n_splits)
424 else:
425 _HC_PRENORM_GEMM_TILELANG_KERNEL(
426 residual_cur_2d,
427 fn,
428 mixes,
429 sqrsum,
430 input_size,
431 1,
432 )
433
434 outputs = (
435 residual_cur,
436 post.unsqueeze(-1),
437 comb.view(num_tokens, hc_mult, hc_mult),
438 layer_input,
439 next_pre_mix,
440 aux,
441 )
442 if norm_weight is not None:
443 assert norm_weight.shape == (hidden_size,)
444 assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous()
445 MHC_PRE_NORM_KERNEL(
446 mixes,
447 sqrsum,
448 hc_scale,
449 hc_base,
450 residual_cur,
451 post,
452 comb,
453 layer_input,
454 norm_weight,
455 pre_mix if pre_mix is not None else post,
456 next_pre_mix,
457 aux if capture_aux else layer_input,
458 hidden_size=hidden_size,
459 rms_eps=rms_eps,
460 hc_pre_eps=hc_pre_eps,
461 hc_sinkhorn_eps=hc_sinkhorn_eps,
462 hc_post_mult_value=hc_post_mult_value,
463 sinkhorn_repeat=sinkhorn_repeat,
464 norm_eps=norm_eps,
465 hc_mult=hc_mult,
466 use_pre_mix_in=pre_mix is not None,
467 save_pre_mix=True,
468 rms_numel=input_size,
469 write_aux=capture_aux,
470 )
471 return outputs
472 mhc_pre_big_fuse_tilelang(
473 mixes,
474 sqrsum,
475 hc_scale,
476 hc_base,
477 residual_cur,
478 post,
479 comb,
480 layer_input,
481 pre_mix if pre_mix is not None else post,
482 next_pre_mix,
483 aux if capture_aux else layer_input,
484 hidden_size,
485 rms_eps,
486 hc_pre_eps,
487 hc_sinkhorn_eps,
488 hc_post_mult_value,
489 sinkhorn_repeat,
490 mixes.shape[0],
491 hc_mult,
492 use_pre_mix_in=pre_mix is not None,
493 save_pre_mix=True,
494 rms_numel=input_size,
495 write_aux=capture_aux,
496 )
497 return outputs
498
499
500 def _mhc_fused_post_pre_delayed_tilelang_fake(
501 x: torch.Tensor,
502 residual: torch.Tensor,
503 post_layer_mix: torch.Tensor,
504 comb_res_mix: torch.Tensor,
505 fn: torch.Tensor,
506 hc_scale: torch.Tensor,
507 hc_base: torch.Tensor,
508 rms_eps: float,
509 hc_pre_eps: float,
510 hc_sinkhorn_eps: float,
511 hc_post_mult_value: float,
512 sinkhorn_repeat: int,
513 pre_mix: torch.Tensor | None = None,
514 norm_weight: torch.Tensor | None = None,
515 norm_eps: float = 1e-6,
516 capture_aux: bool = False,
517 ) -> tuple[
518 torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
519 ]:
520 num_tokens, hc_mult, hidden_size = residual.shape
521 return (
522 torch.empty_like(residual),
523 torch.empty(
524 num_tokens, hc_mult, 1, dtype=torch.float32, device=residual.device
525 ),
526 torch.empty(
527 num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=residual.device
528 ),
529 torch.empty(
530 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
531 ),
532 torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=residual.device),
533 torch.empty(
534 num_tokens if capture_aux else 0,
535 hidden_size,
536 dtype=torch.bfloat16,
537 device=residual.device,
538 ),
539 )
540
541
542 def mhc_pre_tilelang(
543 residual: torch.Tensor,
544 fn: torch.Tensor,
545 hc_scale: torch.Tensor,
546 hc_base: torch.Tensor,
547 rms_eps: float,
548 hc_pre_eps: float,
549 hc_sinkhorn_eps: float,
550 hc_post_mult_value: float,
551 sinkhorn_repeat: int,
552 n_splits: int = 1,
553 norm_weight: torch.Tensor | None = None,
554 norm_eps: float = 1e-6,
555 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
556 """Forward pass for mHC pre block.
557
558 Args:
559 residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
560 fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
561 hc_scale: shape (3,), dtype torch.float32
562 hc_base: shape (hc_mult3,), dtype torch.float32
563 rms_eps: RMS normalization epsilon
564 hc_pre_eps: pre-mix epsilon
565 hc_sinkhorn_eps: sinkhorn epsilon
566 hc_post_mult_value: post-mix multiplier value
567 sinkhorn_repeat: number of sinkhorn iterations
568 n_splits: retained for the shared MHC operator API; the active GEMM
569 backend selects its split factor internally.
570 norm_weight: optional RMSNorm weight, shape (hidden_size,), dtype
571 torch.bfloat16. When provided, RMSNorm is fused into the
572 layer_input write path of the big_fuse kernel.
573 norm_eps: epsilon for the fused RMSNorm; only consulted when
574 norm_weight is given.
575
576 Returns:
577 post_mix: shape (..., hc_mult), dtype torch.float32
578 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
579 layer_input: shape (..., hidden_size), dtype torch.bfloat16
580
581 """
582 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
583 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
584 )
585
586 assert residual.dtype == torch.bfloat16
587 assert fn.dtype == torch.float32
588 assert hc_scale.dtype == torch.float32
589 assert hc_base.dtype == torch.float32
590
591 hc_mult = residual.shape[-2]
592 hidden_size = residual.shape[-1]
593 hc_mult2 = hc_mult * hc_mult
594 hc_mult3 = hc_mult * 2 + hc_mult2
595
596 hc_hidden_size = hc_mult * hidden_size
597 assert fn.shape[0] == hc_mult3
598 assert fn.shape[1] == hc_hidden_size
599 assert hc_scale.shape == (3,)
600 assert hc_base.shape == (hc_mult3,)
601
602 if norm_weight is not None:
603 assert norm_weight.shape == (hidden_size,)
604 if norm_weight.dtype != torch.bfloat16:
605 norm_weight = norm_weight.to(torch.bfloat16)
606 if not norm_weight.is_contiguous():
607 norm_weight = norm_weight.contiguous()
608
609 outer_shape = residual.shape[:-2]
610
611 residual_flat = residual.view(-1, hc_mult, hidden_size)
612 num_tokens = residual_flat.shape[0]
613
614 post_mix = torch.empty(
615 num_tokens, hc_mult, dtype=torch.float32, device=residual.device
616 )
617 comb_mix = torch.empty(
618 num_tokens, hc_mult2, dtype=torch.float32, device=residual.device
619 )
620 layer_input = torch.empty(
621 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
622 )
623
624 gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
625 residual_flat.view(num_tokens, hc_mult * hidden_size),
626 fn,
627 hidden_size=hidden_size,
628 hc_mult=hc_mult,
629 )
630 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
631 gemm_out_mul,
632 gemm_out_sqrsum,
633 hc_scale,
634 hc_base,
635 residual_flat,
636 post_mix,
637 comb_mix,
638 layer_input,
639 rms_eps,
640 hc_pre_eps,
641 hc_sinkhorn_eps,
642 hc_post_mult_value,
643 sinkhorn_repeat,
644 norm_weight=norm_weight,
645 norm_eps=norm_eps,
646 )
647
648 return (
649 post_mix.view(*outer_shape, hc_mult, 1),
650 comb_mix.view(*outer_shape, hc_mult, hc_mult),
651 layer_input.view(*outer_shape, hidden_size),
652 )
653
654
655 def _mhc_pre_tilelang_fake(
656 residual: torch.Tensor,
657 fn: torch.Tensor,
658 hc_scale: torch.Tensor,
659 hc_base: torch.Tensor,
660 rms_eps: float,
661 hc_pre_eps: float,
662 hc_sinkhorn_eps: float,
663 hc_post_mult_value: float,
664 sinkhorn_repeat: int,
665 n_splits: int = 1,
666 norm_weight: torch.Tensor | None = None,
667 norm_eps: float = 1e-6,
668 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
669 hc_mult = residual.shape[-2]
670 hidden_size = residual.shape[-1]
671 outer_shape = residual.shape[:-2]
672
673 # Create empty tensors with correct shapes for meta device / shape inference
674 post_mix = torch.empty(
675 *outer_shape,
676 hc_mult,
677 1,
678 dtype=torch.float32,
679 device=residual.device,
680 )
681 comb_mix = torch.empty(
682 *outer_shape,
683 hc_mult,
684 hc_mult,
685 dtype=torch.float32,
686 device=residual.device,
687 )
688 layer_input = torch.empty(
689 *outer_shape,
690 hidden_size,
691 dtype=torch.bfloat16,
692 device=residual.device,
693 )
694
695 return post_mix, comb_mix, layer_input
696
697
698 def mhc_pre_broadcast_tilelang(
699 residual: torch.Tensor,
700 fn: torch.Tensor,
701 hc_scale: torch.Tensor,
702 hc_base: torch.Tensor,
703 rms_eps: float,
704 hc_pre_eps: float,
705 hc_sinkhorn_eps: float,
706 hc_post_mult_value: float,
707 sinkhorn_repeat: int,
708 n_splits: int = 1,
709 norm_weight: torch.Tensor | None = None,
710 norm_eps: float = 1e-6,
711 fn_broadcast: torch.Tensor | None = None,
712 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
713 """First-layer mHC pre for a residual broadcast from ``(T, H)``."""
714 # n_splits is retained for the shared API; split selection is internal.
715 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
716 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
717 )
718
719 assert norm_weight is not None, "broadcast mHC pre currently requires fused RMSNorm"
720 assert residual.dtype == torch.bfloat16
721 assert residual.dim() == 2
722 assert fn.dtype == torch.float32
723 assert hc_scale.dtype == torch.float32
724 assert hc_base.dtype == torch.float32
725
726 hidden_size = residual.shape[-1]
727 hc_mult = fn.shape[1] // hidden_size
728 hc_mult2 = hc_mult * hc_mult
729 hc_mult3 = hc_mult * 2 + hc_mult2
730 assert fn.shape == (hc_mult3, hc_mult * hidden_size)
731 assert hc_scale.shape == (3,)
732 assert hc_base.shape == (hc_mult3,)
733 assert fn_broadcast is not None
734 assert fn_broadcast.dtype == torch.float32
735 assert fn_broadcast.shape == (hc_mult3, hidden_size)
736
737 if norm_weight.dtype != torch.bfloat16:
738 norm_weight = norm_weight.to(torch.bfloat16)
739 if not norm_weight.is_contiguous():
740 norm_weight = norm_weight.contiguous()
741
742 residual_flat = residual
743 num_tokens = residual.shape[0]
744
745 residual_out = torch.empty(
746 num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=residual.device
747 )
748 post_mix = torch.empty(
749 num_tokens, hc_mult, dtype=torch.float32, device=residual.device
750 )
751 comb_mix = torch.empty(
752 num_tokens, hc_mult2, dtype=torch.float32, device=residual.device
753 )
754 layer_input = torch.empty(
755 num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
756 )
757
758 gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
759 residual_flat,
760 fn_broadcast,
761 hidden_size=hidden_size,
762 hc_mult=hc_mult,
763 use_tilelang_fallback=False,
764 )
765 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
766 gemm_out_mul,
767 gemm_out_sqrsum,
768 hc_scale,
769 hc_base,
770 residual_flat,
771 post_mix,
772 comb_mix,
773 layer_input,
774 rms_eps,
775 hc_pre_eps,
776 hc_sinkhorn_eps,
777 hc_post_mult_value,
778 sinkhorn_repeat,
779 residual_out=residual_out,
780 norm_weight=norm_weight,
781 norm_eps=norm_eps,
782 )
783 return (
784 residual_out,
785 post_mix.unsqueeze(-1),
786 comb_mix.view(num_tokens, hc_mult, hc_mult),
787 layer_input,
788 )
789
790
791 def mhc_post_tilelang(
792 x: torch.Tensor,
793 residual: torch.Tensor,
794 post_layer_mix: torch.Tensor,
795 comb_res_mix: torch.Tensor,
796 ) -> torch.Tensor:
797 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
798 _MHC_POST_TILELANG_KERNEL,
799 )
800
801 out = torch.empty_like(residual)
802 _MHC_POST_TILELANG_KERNEL(
803 comb_res_mix,
804 residual,
805 post_layer_mix.squeeze(-1),
806 x,
807 out,
808 residual.shape[-2],
809 residual.shape[-1],
810 )
811 return out
812
813
814 def mhc_fused_post_pre_tilelang(
815 x: torch.Tensor,
816 residual: torch.Tensor,
817 post_layer_mix: torch.Tensor,
818 comb_res_mix: torch.Tensor,
819 fn: torch.Tensor,
820 hc_scale: torch.Tensor,
821 hc_base: torch.Tensor,
822 rms_eps: float,
823 hc_pre_eps: float,
824 hc_sinkhorn_eps: float,
825 hc_post_mult_value: float,
826 sinkhorn_repeat: int,
827 n_splits: int = 1,
828 tile_n: int = 1,
829 norm_weight: torch.Tensor | None = None,
830 norm_eps: float = 1e-6,
831 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
832 """Run one MHC post block followed by the next MHC pre block.
833
834 When ``norm_weight`` is provided, the layer_input_cur output is the
835 RMSNorm'd activation (fused into the kernel); otherwise it is the
836 raw pre-norm activation as before.
837
838 ``n_splits`` and ``tile_n`` are retained for the shared MHC operator API.
839 The TileLang path selects both values internally from the runtime shape.
840
841 Returns:
842 residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
843 post_mix_cur: shape (..., hc_mult, 1)
844 comb_mix_cur: shape (..., hc_mult, hc_mult)
845 layer_input_cur: shape (..., hidden_size)
846
847 """
848 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
849 _MHC_FUSED_TILELANG_KERNEL,
850 _MHC_POST_TILELANG_KERNEL,
851 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
852 mhc_fused_post_pre_split_config,
853 )
854
855 assert residual.dtype == torch.bfloat16
856 assert x.dtype == torch.bfloat16
857 assert post_layer_mix.dtype == torch.float32
858 assert comb_res_mix.dtype == torch.float32
859 assert fn.dtype == torch.float32
860 assert hc_scale.dtype == torch.float32
861 assert hc_base.dtype == torch.float32
862
863 hc_mult = residual.shape[-2]
864 hidden_size = residual.shape[-1]
865 hc_mult2 = hc_mult * hc_mult
866 hc_mult3 = hc_mult * 2 + hc_mult2
867 hc_hidden_size = hc_mult * hidden_size
868 outer_shape = residual.shape[:-2]
869
870 assert x.shape == (*outer_shape, hidden_size)
871 assert post_layer_mix.shape in (
872 (*outer_shape, hc_mult, 1),
873 (*outer_shape, hc_mult),
874 )
875 assert comb_res_mix.shape == (*outer_shape, hc_mult, hc_mult)
876 assert fn.shape == (hc_mult3, hc_hidden_size)
877 assert hc_scale.shape == (3,)
878 assert hc_base.shape == (hc_mult3,)
879
880 if norm_weight is not None:
881 assert norm_weight.shape == (hidden_size,)
882 if norm_weight.dtype != torch.bfloat16:
883 norm_weight = norm_weight.to(torch.bfloat16)
884 if not norm_weight.is_contiguous():
885 norm_weight = norm_weight.contiguous()
886
887 residual_flat = residual.view(-1, hc_mult, hidden_size)
888 num_tokens = residual_flat.shape[0]
889 x_flat = x.view(num_tokens, hidden_size)
890 post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
891 comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
892
893 fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
894
895 post_mix_cur = torch.empty(
896 num_tokens,
897 hc_mult,
898 dtype=torch.float32,
899 device=residual.device,
900 )
901 comb_mix_cur = torch.empty(
902 num_tokens,
903 hc_mult2,
904 dtype=torch.float32,
905 device=residual.device,
906 )
907 layer_input_cur = torch.empty(
908 num_tokens,
909 hidden_size,
910 dtype=torch.bfloat16,
911 device=residual.device,
912 )
913
914 if fused_config is not None:
915 gemm_out_mul, gemm_out_sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL(
916 comb_res_mix_flat,
917 residual_flat,
918 post_layer_mix_flat,
919 x_flat,
920 fn.view(hc_mult3, hc_mult, hidden_size),
921 hc_mult,
922 hidden_size,
923 hc_mult3,
924 )
925 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
926 gemm_out_mul,
927 gemm_out_sqrsum,
928 hc_scale,
929 hc_base,
930 residual_cur,
931 post_mix_cur,
932 comb_mix_cur,
933 layer_input_cur,
934 rms_eps,
935 hc_pre_eps,
936 hc_sinkhorn_eps,
937 hc_post_mult_value,
938 sinkhorn_repeat,
939 norm_weight=norm_weight,
940 norm_eps=norm_eps,
941 )
942 else:
943 residual_cur = torch.empty_like(residual_flat)
944 _MHC_POST_TILELANG_KERNEL(
945 comb_res_mix_flat,
946 residual_flat,
947 post_layer_mix_flat,
948 x_flat,
949 residual_cur,
950 residual.shape[-2],
951 residual.shape[-1],
952 )
953 gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
954 residual_cur.view(num_tokens, hc_mult * hidden_size),
955 fn,
956 hidden_size=hidden_size,
957 hc_mult=hc_mult,
958 )
959 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
960 gemm_out_mul,
961 gemm_out_sqrsum,
962 hc_scale,
963 hc_base,
964 residual_cur,
965 post_mix_cur,
966 comb_mix_cur,
967 layer_input_cur,
968 rms_eps,
969 hc_pre_eps,
970 hc_sinkhorn_eps,
971 hc_post_mult_value,
972 sinkhorn_repeat,
973 norm_weight=norm_weight,
974 norm_eps=norm_eps,
975 )
976
977 return (
978 residual_cur.view(*outer_shape, hc_mult, hidden_size),
979 post_mix_cur.view(*outer_shape, hc_mult, 1),
980 comb_mix_cur.view(*outer_shape, hc_mult, hc_mult),
981 layer_input_cur.view(*outer_shape, hidden_size),
982 )
983
984
985 def _mhc_fused_post_pre_tilelang_fake(
986 x: torch.Tensor,
987 residual: torch.Tensor,
988 post_layer_mix: torch.Tensor,
989 comb_res_mix: torch.Tensor,
990 fn: torch.Tensor,
991 hc_scale: torch.Tensor,
992 hc_base: torch.Tensor,
993 rms_eps: float,
994 hc_pre_eps: float,
995 hc_sinkhorn_eps: float,
996 hc_post_mult_value: float,
997 sinkhorn_repeat: int,
998 n_splits: int = 1,
999 tile_n: int = 1,
1000 norm_weight: torch.Tensor | None = None,
1001 norm_eps: float = 1e-6,
1002 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1003 hc_mult = residual.shape[-2]
1004 hidden_size = residual.shape[-1]
1005 outer_shape = residual.shape[:-2]
1006
1007 residual_cur = torch.empty_like(residual)
1008 post_mix_cur = torch.empty(
1009 *outer_shape,
1010 hc_mult,
1011 1,
1012 dtype=torch.float32,
1013 device=residual.device,
1014 )
1015 comb_mix_cur = torch.empty(
1016 *outer_shape,
1017 hc_mult,
1018 hc_mult,
1019 dtype=torch.float32,
1020 device=residual.device,
1021 )
1022 layer_input_cur = torch.empty(
1023 *outer_shape,
1024 hidden_size,
1025 dtype=torch.bfloat16,
1026 device=residual.device,
1027 )
1028
1029 return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
1030
1031
1032 def _mhc_post_tilelang_fake(
1033 x: torch.Tensor,
1034 residual: torch.Tensor,
1035 post_layer_mix: torch.Tensor,
1036 comb_res_mix: torch.Tensor,
1037 ) -> torch.Tensor:
1038 return torch.empty_like(residual)
1039
1040
1041 def hc_head_fused_kernel_tilelang(
1042 hs_flat: torch.Tensor,
1043 fn: torch.Tensor,
1044 hc_scale: torch.Tensor,
1045 hc_base: torch.Tensor,
1046 rms_eps: float,
1047 hc_eps: float,
1048 ) -> torch.Tensor:
1049 """Apply the fused hc_head kernel and return the (T, H) bf16 result."""
1050 num_tokens, hc_mult, hidden_size = hs_flat.shape
1051 out = torch.empty(
1052 num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device
1053 )
1054 if num_tokens == 0:
1055 return out
1056 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
1057 _HC_HEAD_FUSED_TILELANG_KERNEL,
1058 )
1059
1060 _HC_HEAD_FUSED_TILELANG_KERNEL(
1061 hs_flat,
1062 fn,
1063 hc_scale,
1064 hc_base,
1065 out,
1066 hidden_size,
1067 rms_eps,
1068 hc_eps,
1069 hc_mult,
1070 )
1071 return out
1072
1073
1074 def _hc_head_fused_kernel_tilelang_fake(
1075 hs_flat: torch.Tensor,
1076 fn: torch.Tensor,
1077 hc_scale: torch.Tensor,
1078 hc_base: torch.Tensor,
1079 rms_eps: float,
1080 hc_eps: float,
1081 ) -> torch.Tensor:
1082 num_tokens, _, hidden_size = hs_flat.shape
1083 return torch.empty(
1084 num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device
1085 )
1086
1087
1088 direct_register_custom_op(
1089 op_name="mhc_pre_delayed_tilelang",
1090 op_func=mhc_pre_delayed_tilelang,
1091 mutates_args=[],
1092 fake_impl=_mhc_pre_delayed_tilelang_fake,
1093 )
1094 direct_register_custom_op(
1095 op_name="mhc_fused_post_pre_delayed_tilelang",
1096 op_func=mhc_fused_post_pre_delayed_tilelang,
1097 mutates_args=[],
1098 fake_impl=_mhc_fused_post_pre_delayed_tilelang_fake,
1099 )
1100 direct_register_custom_op(
1101 op_name="mhc_pre_tilelang",
1102 op_func=mhc_pre_tilelang,
1103 mutates_args=[],
1104 fake_impl=_mhc_pre_tilelang_fake,
1105 )
1106 direct_register_custom_op(
1107 op_name="mhc_post_tilelang",
1108 op_func=mhc_post_tilelang,
1109 mutates_args=[],
1110 fake_impl=_mhc_post_tilelang_fake,
1111 )
1112
1113 direct_register_custom_op(
1114 op_name="mhc_fused_post_pre_tilelang",
1115 op_func=mhc_fused_post_pre_tilelang,
1116 mutates_args=[],
1117 fake_impl=_mhc_fused_post_pre_tilelang_fake,
1118 )
1119
1120 direct_register_custom_op(
1121 op_name="hc_head_fused_kernel_tilelang",
1122 op_func=hc_head_fused_kernel_tilelang,
1123 mutates_args=[],
1124 fake_impl=_hc_head_fused_kernel_tilelang_fake,
1125 )