@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,020 lines 30.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import inspect
4
5 import torch
6
7 # this import will also register the custom ops
8 # import vllm.model_executor.kernels.mhc # noqa: F401
9 import vllm.model_executor.kernels.mhc as mhc_kernels
10 from vllm._aiter_ops import is_aiter_found_and_supported
11 from vllm.model_executor.custom_op import CustomOp
12 from vllm.platforms import current_platform
13 from vllm.utils.import_utils import has_tilelang
14
15
16 def _has_tilelang_mhc() -> bool:
17 if not has_tilelang():
18 return False
19 if current_platform.is_cuda():
20 return True
21 if current_platform.is_rocm():
22 from vllm.platforms.rocm import on_gfx942
23
24 # TileLang MHC currently produces incorrect results on gfx942. Keep
25 # gfx942 on the existing torch/triton fallbacks until that path is fixed.
26 return not on_gfx942()
27 return False
28
29
30 HAS_TILELANG_MHC = _has_tilelang_mhc()
31 HAS_AITER_MHC = is_aiter_found_and_supported()
32
33
34 def _has_aiter_mhc_fused() -> bool:
35 if not HAS_AITER_MHC:
36 return False
37 try:
38 from aiter.ops.mhc import mhc_fused_post_pre
39 except Exception:
40 return False
41 return callable(mhc_fused_post_pre)
42
43
44 def _aiter_mhc_op_accepts_norm(op_name: str) -> bool:
45 if not HAS_AITER_MHC:
46 return False
47 try:
48 from aiter.ops import mhc as aiter_mhc
49
50 for candidate_name in (op_name, f"{op_name}_fake"):
51 op = getattr(aiter_mhc, candidate_name, None)
52 if op is None:
53 continue
54 parameters = inspect.signature(op).parameters
55 if "norm_weight" in parameters and "norm_eps" in parameters:
56 return True
57 except (AttributeError, ImportError, TypeError, ValueError):
58 pass
59 return False
60
61
62 def _has_aiter_mhc_fused_post_pre_delayed_rms_norm() -> bool:
63 if not HAS_AITER_MHC:
64 return False
65 from vllm.platforms.rocm import on_gfx950
66
67 # The fused delayed seam kernel is only validated on gfx950.
68 return on_gfx950()
69
70
71 HAS_AITER_MHC_FUSED = _has_aiter_mhc_fused()
72 HAS_AITER_MHC_PRE_NORM = _aiter_mhc_op_accepts_norm("mhc_pre")
73 HAS_AITER_MHC_FUSED_NORM = _aiter_mhc_op_accepts_norm("mhc_fused_post_pre")
74 HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM = (
75 _has_aiter_mhc_fused_post_pre_delayed_rms_norm()
76 )
77
78
79 def _aiter_mhc_supported(
80 residual: torch.Tensor,
81 norm_weight: torch.Tensor | None,
82 *,
83 supports_norm: bool,
84 ) -> bool:
85 hidden_size = residual.shape[-1]
86 hc_mult = residual.shape[-2]
87 return (
88 HAS_AITER_MHC
89 and hidden_size % 256 == 0
90 and hc_mult == 4
91 and (norm_weight is None or supports_norm)
92 )
93
94
95 def _apply_mhc_norm(
96 layer_input: torch.Tensor,
97 norm_weight: torch.Tensor | None,
98 norm_eps: float,
99 ) -> torch.Tensor:
100 if norm_weight is None:
101 return layer_input
102 from vllm import ir
103
104 return ir.ops.rms_norm(layer_input, norm_weight, norm_eps)
105
106
107 # --8<-- [start:mhc_pre]
108 @CustomOp.register("mhc_pre")
109 class MHCPreOp(CustomOp):
110 """MHC pre block.
111
112 Computes mix logits from RMS-normalized HC residual streams, then
113 returns post_mix, comb_mix, and
114 layer_input = sum_i pre_mix_i * residual_i.
115 """
116
117 # --8<-- [end:mhc_pre]
118 @classmethod
119 def enabled(cls) -> bool:
120 return True
121
122 def forward_cuda(
123 self,
124 residual: torch.Tensor,
125 fn: torch.Tensor,
126 hc_scale: torch.Tensor,
127 hc_base: torch.Tensor,
128 rms_eps: float,
129 hc_pre_eps: float,
130 hc_sinkhorn_eps: float,
131 hc_post_mult_value: float,
132 sinkhorn_repeat: int,
133 n_splits: int = 1,
134 norm_weight: torch.Tensor | None = None,
135 norm_eps: float = 0.0,
136 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
137 return torch.ops.vllm.mhc_pre_tilelang(
138 residual,
139 fn,
140 hc_scale,
141 hc_base,
142 rms_eps,
143 hc_pre_eps,
144 hc_sinkhorn_eps,
145 hc_post_mult_value,
146 sinkhorn_repeat,
147 n_splits,
148 norm_weight,
149 norm_eps,
150 )
151
152 def forward_hip(
153 self,
154 residual: torch.Tensor,
155 fn: torch.Tensor,
156 hc_scale: torch.Tensor,
157 hc_base: torch.Tensor,
158 rms_eps: float,
159 hc_pre_eps: float,
160 hc_sinkhorn_eps: float,
161 hc_post_mult_value: float,
162 sinkhorn_repeat: int,
163 n_splits: int = 1,
164 norm_weight: torch.Tensor | None = None,
165 norm_eps: float = 0.0,
166 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
167 if _aiter_mhc_supported(
168 residual,
169 norm_weight,
170 supports_norm=HAS_AITER_MHC_PRE_NORM,
171 ):
172 return torch.ops.vllm.mhc_pre_aiter(
173 residual,
174 fn,
175 hc_scale,
176 hc_base,
177 rms_eps,
178 hc_pre_eps,
179 hc_sinkhorn_eps,
180 hc_post_mult_value,
181 sinkhorn_repeat,
182 n_splits,
183 norm_weight,
184 norm_eps,
185 )
186 elif HAS_TILELANG_MHC:
187 return torch.ops.vllm.mhc_pre_tilelang(
188 residual,
189 fn,
190 hc_scale,
191 hc_base,
192 rms_eps,
193 hc_pre_eps,
194 hc_sinkhorn_eps,
195 hc_post_mult_value,
196 sinkhorn_repeat,
197 n_splits,
198 norm_weight,
199 norm_eps,
200 )
201 else:
202 post_mix, comb_mix, layer_input = self.forward_native(
203 residual,
204 fn,
205 hc_scale,
206 hc_base,
207 rms_eps,
208 hc_pre_eps,
209 hc_sinkhorn_eps,
210 hc_post_mult_value,
211 sinkhorn_repeat,
212 n_splits,
213 norm_weight,
214 norm_eps,
215 )
216 return (
217 post_mix,
218 comb_mix,
219 _apply_mhc_norm(layer_input, norm_weight, norm_eps),
220 )
221
222 def forward_native(
223 self,
224 residual: torch.Tensor,
225 fn: torch.Tensor,
226 hc_scale: torch.Tensor,
227 hc_base: torch.Tensor,
228 rms_eps: float,
229 hc_pre_eps: float,
230 hc_sinkhorn_eps: float,
231 hc_post_mult_value: float,
232 sinkhorn_repeat: int,
233 n_splits: int = 1,
234 norm_weight: torch.Tensor | None = None,
235 norm_eps: float = 0.0,
236 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
237 return mhc_kernels.mhc_pre_torch(
238 residual,
239 fn,
240 hc_scale,
241 hc_base,
242 rms_eps,
243 hc_pre_eps,
244 hc_sinkhorn_eps,
245 hc_post_mult_value,
246 sinkhorn_repeat,
247 )
248
249 def forward_xpu(
250 self,
251 residual: torch.Tensor,
252 fn: torch.Tensor,
253 hc_scale: torch.Tensor,
254 hc_base: torch.Tensor,
255 rms_eps: float,
256 hc_pre_eps: float,
257 hc_sinkhorn_eps: float,
258 hc_post_mult_value: float,
259 sinkhorn_repeat: int,
260 n_splits: int = 1,
261 norm_weight: torch.Tensor | None = None,
262 norm_eps: float = 0.0,
263 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
264 return torch.ops._xpu_C.mhc_pre(
265 residual,
266 fn,
267 hc_scale,
268 hc_base,
269 rms_eps,
270 hc_pre_eps,
271 hc_sinkhorn_eps,
272 hc_post_mult_value,
273 sinkhorn_repeat,
274 )
275
276 def forward_cpu(
277 self,
278 residual: torch.Tensor,
279 fn: torch.Tensor,
280 hc_scale: torch.Tensor,
281 hc_base: torch.Tensor,
282 rms_eps: float,
283 hc_pre_eps: float,
284 hc_sinkhorn_eps: float,
285 hc_post_mult_value: float,
286 sinkhorn_repeat: int,
287 n_splits: int = 1,
288 norm_weight: torch.Tensor | None = None,
289 norm_eps: float = 0.0,
290 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
291 return mhc_kernels.mhc_pre_cpu(
292 residual,
293 fn,
294 hc_scale,
295 hc_base,
296 rms_eps,
297 hc_pre_eps,
298 hc_sinkhorn_eps,
299 hc_post_mult_value,
300 sinkhorn_repeat,
301 n_splits,
302 norm_weight,
303 norm_eps,
304 )
305
306
307 # --8<-- [start:mhc_pre_delayed]
308 @CustomOp.register("mhc_pre_delayed")
309 class MHCPreDelayedOp(CustomOp):
310 """MHC pre block using the pre-mix carried from the previous sublayer.
311
312 Same gates as :class:`MHCPreOp`, but the stream collapse applies the
313 caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next
314 sublayer seam.
315
316 Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` also
317 applies the preceding post block, which lets AITER fold it into the pre
318 projection and saves a launch. Returns residual, post_mix, comb_mix,
319 layer_input, next_pre_mix; ``residual`` is the caller's own tensor when
320 no post is requested.
321 """
322
323 # --8<-- [end:mhc_pre_delayed]
324 @classmethod
325 def enabled(cls) -> bool:
326 return True
327
328 def __init__(self) -> None:
329 super().__init__()
330 # Built here, not lazily: CustomOp resolves dispatch against the
331 # current vLLM config, which is only set during model construction.
332 self._post = MHCPostOp()
333
334 def _maybe_post(
335 self,
336 residual: torch.Tensor,
337 sublayer_out: torch.Tensor | None,
338 post_layer_mix: torch.Tensor | None,
339 comb_res_mix: torch.Tensor | None,
340 ) -> torch.Tensor:
341 if sublayer_out is None:
342 return residual
343 assert post_layer_mix is not None and comb_res_mix is not None
344 return self._post(sublayer_out, residual, post_layer_mix, comb_res_mix)
345
346 def forward_cuda(
347 self,
348 residual: torch.Tensor,
349 fn: torch.Tensor,
350 hc_scale: torch.Tensor,
351 hc_base: torch.Tensor,
352 rms_eps: float,
353 hc_pre_eps: float,
354 hc_sinkhorn_eps: float,
355 hc_post_mult_value: float,
356 sinkhorn_repeat: int,
357 pre_mix: torch.Tensor | None = None,
358 x: torch.Tensor | None = None,
359 norm_weight: torch.Tensor | None = None,
360 norm_eps: float = 1e-6,
361 sublayer_out: torch.Tensor | None = None,
362 post_layer_mix: torch.Tensor | None = None,
363 comb_res_mix: torch.Tensor | None = None,
364 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
365 residual = self._maybe_post(
366 residual, sublayer_out, post_layer_mix, comb_res_mix
367 )
368 return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
369 residual,
370 fn,
371 hc_scale,
372 hc_base,
373 rms_eps,
374 hc_pre_eps,
375 hc_sinkhorn_eps,
376 hc_post_mult_value,
377 sinkhorn_repeat,
378 pre_mix,
379 x,
380 norm_weight,
381 norm_eps,
382 )
383
384 def forward_hip(
385 self,
386 residual: torch.Tensor,
387 fn: torch.Tensor,
388 hc_scale: torch.Tensor,
389 hc_base: torch.Tensor,
390 rms_eps: float,
391 hc_pre_eps: float,
392 hc_sinkhorn_eps: float,
393 hc_post_mult_value: float,
394 sinkhorn_repeat: int,
395 pre_mix: torch.Tensor | None = None,
396 x: torch.Tensor | None = None,
397 norm_weight: torch.Tensor | None = None,
398 norm_eps: float = 1e-6,
399 sublayer_out: torch.Tensor | None = None,
400 post_layer_mix: torch.Tensor | None = None,
401 comb_res_mix: torch.Tensor | None = None,
402 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
403 # aiter's fused delayed seam: post-mix, gate projection, and the
404 # collapse with the carried pre-mix and its RMSNorm. It always applies
405 # the norm and projects `residual` itself.
406 if (
407 HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM
408 and x is None
409 and norm_weight is not None
410 and _aiter_mhc_supported(residual, norm_weight, supports_norm=True)
411 ):
412 residual_out = (
413 torch.empty_like(residual) if sublayer_out is not None else None
414 )
415 rest = torch.ops.vllm.mhc_fused_post_pre_delayed_rms_norm_aiter(
416 residual,
417 fn,
418 hc_scale,
419 hc_base,
420 rms_eps,
421 hc_pre_eps,
422 hc_sinkhorn_eps,
423 hc_post_mult_value,
424 sinkhorn_repeat,
425 pre_mix,
426 sublayer_out,
427 post_layer_mix,
428 comb_res_mix,
429 norm_weight,
430 norm_eps,
431 residual_out,
432 )
433 return (residual if residual_out is None else residual_out), *rest
434 # The unfused AITER delayed path drives mhc_pre_gemm_sqrsum against
435 # `residual` and folds no RMSNorm, so it cannot serve the model-entry
436 # broadcast (which projects a narrower `x`) or a fused norm the branch
437 # above did not take. Both are handled by TileLang, or by the reference
438 # when TileLang is unavailable.
439 if (
440 x is None
441 and norm_weight is None
442 and _aiter_mhc_supported(residual, None, supports_norm=False)
443 ):
444 from vllm._aiter_ops import rocm_aiter_ops
445
446 num_tokens = residual.numel() // (residual.shape[-1] * residual.shape[-2])
447 # Folding the post in only pays while the residual still fits in
448 # cache; past that AITER's heuristic wants the separate kernels,
449 # which can ask for a non-temporal store and a large-m split-k.
450 if sublayer_out is not None and (
451 rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(num_tokens)
452 ):
453 residual = self._maybe_post(
454 residual, sublayer_out, post_layer_mix, comb_res_mix
455 )
456 sublayer_out = None
457 # The op writes the folded post's residual into this buffer rather
458 # than returning it, since on the unfused path it has none of its
459 # own and an op must not hand back one of its inputs.
460 residual_out = (
461 torch.empty_like(residual) if sublayer_out is not None else None
462 )
463 rest = torch.ops.vllm.mhc_pre_delayed_aiter(
464 residual,
465 fn,
466 hc_scale,
467 hc_base,
468 rms_eps,
469 hc_pre_eps,
470 hc_sinkhorn_eps,
471 hc_post_mult_value,
472 sinkhorn_repeat,
473 pre_mix,
474 sublayer_out,
475 post_layer_mix,
476 comb_res_mix,
477 residual_out,
478 )
479 return (residual if residual_out is None else residual_out), *rest
480 residual = self._maybe_post(
481 residual, sublayer_out, post_layer_mix, comb_res_mix
482 )
483 if HAS_TILELANG_MHC:
484 return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang(
485 residual,
486 fn,
487 hc_scale,
488 hc_base,
489 rms_eps,
490 hc_pre_eps,
491 hc_sinkhorn_eps,
492 hc_post_mult_value,
493 sinkhorn_repeat,
494 pre_mix,
495 x,
496 norm_weight,
497 norm_eps,
498 )
499 # The post is already applied, so it must not be requested again.
500 return self.forward_native(
501 residual,
502 fn,
503 hc_scale,
504 hc_base,
505 rms_eps,
506 hc_pre_eps,
507 hc_sinkhorn_eps,
508 hc_post_mult_value,
509 sinkhorn_repeat,
510 pre_mix,
511 x,
512 norm_weight,
513 norm_eps,
514 )
515
516 def forward_native(
517 self,
518 residual: torch.Tensor,
519 fn: torch.Tensor,
520 hc_scale: torch.Tensor,
521 hc_base: torch.Tensor,
522 rms_eps: float,
523 hc_pre_eps: float,
524 hc_sinkhorn_eps: float,
525 hc_post_mult_value: float,
526 sinkhorn_repeat: int,
527 pre_mix: torch.Tensor | None = None,
528 x: torch.Tensor | None = None,
529 norm_weight: torch.Tensor | None = None,
530 norm_eps: float = 1e-6,
531 sublayer_out: torch.Tensor | None = None,
532 post_layer_mix: torch.Tensor | None = None,
533 comb_res_mix: torch.Tensor | None = None,
534 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
535 residual = self._maybe_post(
536 residual, sublayer_out, post_layer_mix, comb_res_mix
537 )
538 post_mix, comb_mix, layer_input, next_pre_mix = (
539 mhc_kernels.mhc_pre_delayed_torch(
540 residual,
541 fn,
542 hc_scale,
543 hc_base,
544 rms_eps,
545 hc_pre_eps,
546 hc_sinkhorn_eps,
547 hc_post_mult_value,
548 sinkhorn_repeat,
549 pre_mix=pre_mix,
550 x=x,
551 )
552 )
553 return (
554 residual,
555 post_mix,
556 comb_mix,
557 _apply_mhc_norm(layer_input, norm_weight, norm_eps),
558 next_pre_mix,
559 )
560
561
562 # --8<-- [start:mhc_post]
563 @CustomOp.register("mhc_post")
564 class MHCPostOp(CustomOp):
565 """MHC post block.
566
567 Combines the layer output with the HC residual streams:
568 out_j = post_layer_mix_j * x + sum_i comb_res_mix_ij * residual_i.
569 """
570
571 # --8<-- [end:mhc_post]
572
573 @classmethod
574 def enabled(cls) -> bool:
575 return True
576
577 def forward_cuda(
578 self,
579 x: torch.Tensor,
580 residual: torch.Tensor,
581 post_layer_mix: torch.Tensor,
582 comb_res_mix: torch.Tensor,
583 ) -> torch.Tensor:
584 return torch.ops.vllm.mhc_post_tilelang(
585 x, residual, post_layer_mix, comb_res_mix
586 )
587
588 def forward_hip(
589 self,
590 x: torch.Tensor,
591 residual: torch.Tensor,
592 post_layer_mix: torch.Tensor,
593 comb_res_mix: torch.Tensor,
594 ) -> torch.Tensor:
595 if _aiter_mhc_supported(residual, None, supports_norm=True):
596 return torch.ops.vllm.mhc_post_aiter(
597 x,
598 residual,
599 post_layer_mix,
600 comb_res_mix,
601 )
602 if HAS_TILELANG_MHC:
603 return torch.ops.vllm.mhc_post_tilelang(
604 x, residual, post_layer_mix, comb_res_mix
605 )
606 else:
607 return self.forward_native(x, residual, post_layer_mix, comb_res_mix)
608
609 def forward_native(
610 self,
611 x: torch.Tensor,
612 residual: torch.Tensor,
613 post_layer_mix: torch.Tensor,
614 comb_res_mix: torch.Tensor,
615 ) -> torch.Tensor:
616 return mhc_kernels.mhc_post_torch(
617 x,
618 residual,
619 post_layer_mix,
620 comb_res_mix,
621 )
622
623 def forward_xpu(
624 self,
625 x: torch.Tensor,
626 residual: torch.Tensor,
627 post_layer_mix: torch.Tensor,
628 comb_res_mix: torch.Tensor,
629 ) -> torch.Tensor:
630 return torch.ops._xpu_C.mhc_post(
631 x,
632 residual,
633 post_layer_mix,
634 comb_res_mix,
635 )
636
637 def forward_cpu(
638 self,
639 x: torch.Tensor,
640 residual: torch.Tensor,
641 post_layer_mix: torch.Tensor,
642 comb_res_mix: torch.Tensor,
643 ) -> torch.Tensor:
644 return mhc_kernels.mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix)
645
646
647 # --8<-- [start:hc_head]
648 @CustomOp.register("hc_head")
649 class HCHeadOp(CustomOp):
650 """HC head reduction for DeepSeek V4.
651
652 Computes gates from the RMS-normalized flattened HC residual and
653 returns out = sum_i gate_i * residual_i, collapsing hc_mult streams
654 to one.
655 """
656
657 # --8<-- [end:hc_head]
658 @classmethod
659 def enabled(cls) -> bool:
660 return True
661
662 def forward_cuda(
663 self,
664 hidden_states: torch.Tensor,
665 hc_fn: torch.Tensor,
666 hc_scale: torch.Tensor,
667 hc_base: torch.Tensor,
668 rms_norm_eps: float,
669 hc_eps: float,
670 ) -> torch.Tensor:
671 hc_mult, hidden_size = hidden_states.shape[-2:]
672 outer_shape = hidden_states.shape[:-2]
673 hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
674 out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
675 hs_flat,
676 hc_fn,
677 hc_scale,
678 hc_base,
679 rms_norm_eps,
680 hc_eps,
681 )
682 return out.view(*outer_shape, hidden_size)
683
684 def forward_hip(
685 self,
686 hidden_states: torch.Tensor,
687 hc_fn: torch.Tensor,
688 hc_scale: torch.Tensor,
689 hc_base: torch.Tensor,
690 rms_norm_eps: float,
691 hc_eps: float,
692 ) -> torch.Tensor:
693 hc_mult, hidden_size = hidden_states.shape[-2:]
694 outer_shape = hidden_states.shape[:-2]
695 hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
696
697 if HAS_TILELANG_MHC:
698 out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
699 hs_flat,
700 hc_fn,
701 hc_scale,
702 hc_base,
703 rms_norm_eps,
704 hc_eps,
705 )
706 else:
707 num_tokens = hs_flat.shape[0]
708 out = torch.empty(
709 num_tokens,
710 hidden_size,
711 dtype=torch.bfloat16,
712 device=hidden_states.device,
713 )
714 torch.ops.vllm.hc_head_triton(
715 hs_flat,
716 hc_fn,
717 hc_scale,
718 hc_base,
719 out,
720 hidden_size,
721 rms_norm_eps,
722 hc_eps,
723 hc_mult,
724 )
725
726 return out.view(*outer_shape, hidden_size)
727
728 def forward_native(self, *args, **kwargs):
729 raise NotImplementedError("Native implementation of hc_head is not available")
730
731 def forward_xpu(
732 self,
733 hidden_states: torch.Tensor,
734 hc_fn: torch.Tensor,
735 hc_scale: torch.Tensor,
736 hc_base: torch.Tensor,
737 rms_norm_eps: float,
738 hc_eps: float,
739 ) -> torch.Tensor:
740 hc_mult, hidden_size = hidden_states.shape[-2:]
741 outer_shape = hidden_states.shape[:-2]
742 hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
743 num_tokens = hs_flat.shape[0]
744
745 out = torch.empty(
746 num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
747 )
748 torch.ops._xpu_C.hc_head_fused(
749 hs_flat, hc_fn, hc_scale, hc_base, out, rms_norm_eps, hc_eps
750 )
751 return out.view(*outer_shape, hidden_size)
752
753 def forward_cpu(
754 self,
755 hidden_states: torch.Tensor,
756 hc_fn: torch.Tensor,
757 hc_scale: torch.Tensor,
758 hc_base: torch.Tensor,
759 rms_norm_eps: float,
760 hc_eps: float,
761 ) -> torch.Tensor:
762 return mhc_kernels.hc_head_fused_cpu(
763 hidden_states, hc_fn, hc_scale, hc_base, rms_norm_eps, hc_eps
764 )
765
766
767 # --8<-- [start:mhc_fused_post_pre]
768 @CustomOp.register("mhc_fused_post_pre")
769 class MHCFusedPostPreOp(CustomOp):
770 """Fused MHC post block followed by the next MHC pre block.
771
772 Equivalent to applying MHCPostOp and then MHCPreOp to the updated
773 residual streams, returning residual_cur, post_mix_cur, comb_mix_cur,
774 and layer_input_cur.
775 """
776
777 # --8<-- [end:mhc_fused_post_pre]
778 @classmethod
779 def enabled(cls) -> bool:
780 return True
781
782 def forward_cuda(
783 self,
784 x: torch.Tensor,
785 residual: torch.Tensor,
786 post_layer_mix: torch.Tensor,
787 comb_res_mix: torch.Tensor,
788 fn: torch.Tensor,
789 hc_scale: torch.Tensor,
790 hc_base: torch.Tensor,
791 rms_eps: float,
792 hc_pre_eps: float,
793 hc_sinkhorn_eps: float,
794 hc_post_mult_value: float,
795 sinkhorn_repeat: int,
796 n_splits: int = 1,
797 tile_n: int = 1,
798 norm_weight: torch.Tensor | None = None,
799 norm_eps: float = 0.0,
800 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
801 return torch.ops.vllm.mhc_fused_post_pre_tilelang(
802 x,
803 residual,
804 post_layer_mix,
805 comb_res_mix,
806 fn,
807 hc_scale,
808 hc_base,
809 rms_eps,
810 hc_pre_eps,
811 hc_sinkhorn_eps,
812 hc_post_mult_value,
813 sinkhorn_repeat,
814 n_splits,
815 tile_n,
816 norm_weight,
817 norm_eps,
818 )
819
820 def forward_hip(
821 self,
822 x: torch.Tensor,
823 residual: torch.Tensor,
824 post_layer_mix: torch.Tensor,
825 comb_res_mix: torch.Tensor,
826 fn: torch.Tensor,
827 hc_scale: torch.Tensor,
828 hc_base: torch.Tensor,
829 rms_eps: float,
830 hc_pre_eps: float,
831 hc_sinkhorn_eps: float,
832 hc_post_mult_value: float,
833 sinkhorn_repeat: int,
834 n_splits: int = 1,
835 tile_n: int = 1,
836 norm_weight: torch.Tensor | None = None,
837 norm_eps: float = 0.0,
838 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
839 if HAS_AITER_MHC_FUSED and _aiter_mhc_supported(
840 residual,
841 norm_weight,
842 supports_norm=HAS_AITER_MHC_FUSED_NORM,
843 ):
844 return torch.ops.vllm.mhc_fused_post_pre_aiter(
845 x,
846 residual,
847 post_layer_mix,
848 comb_res_mix,
849 fn,
850 hc_scale,
851 hc_base,
852 rms_eps,
853 hc_pre_eps,
854 hc_sinkhorn_eps,
855 hc_post_mult_value,
856 sinkhorn_repeat,
857 n_splits,
858 tile_n,
859 norm_weight,
860 norm_eps,
861 )
862 if HAS_TILELANG_MHC:
863 return torch.ops.vllm.mhc_fused_post_pre_tilelang(
864 x,
865 residual,
866 post_layer_mix,
867 comb_res_mix,
868 fn,
869 hc_scale,
870 hc_base,
871 rms_eps,
872 hc_pre_eps,
873 hc_sinkhorn_eps,
874 hc_post_mult_value,
875 sinkhorn_repeat,
876 n_splits,
877 tile_n,
878 norm_weight,
879 norm_eps,
880 )
881 residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = self.forward_native(
882 x,
883 residual,
884 post_layer_mix,
885 comb_res_mix,
886 fn,
887 hc_scale,
888 hc_base,
889 rms_eps,
890 hc_pre_eps,
891 hc_sinkhorn_eps,
892 hc_post_mult_value,
893 sinkhorn_repeat,
894 n_splits,
895 tile_n,
896 norm_weight,
897 norm_eps,
898 )
899 return (
900 residual_cur,
901 post_mix_cur,
902 comb_mix_cur,
903 _apply_mhc_norm(layer_input_cur, norm_weight, norm_eps),
904 )
905
906 def forward_native(
907 self,
908 x: torch.Tensor,
909 residual: torch.Tensor,
910 post_layer_mix: torch.Tensor,
911 comb_res_mix: torch.Tensor,
912 fn: torch.Tensor,
913 hc_scale: torch.Tensor,
914 hc_base: torch.Tensor,
915 rms_eps: float,
916 hc_pre_eps: float,
917 hc_sinkhorn_eps: float,
918 hc_post_mult_value: float,
919 sinkhorn_repeat: int,
920 n_splits: int = 1,
921 tile_n: int = 1,
922 norm_weight: torch.Tensor | None = None,
923 norm_eps: float = 0.0,
924 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
925 # Decompose into post + pre (no fused kernel available).
926 residual_cur = mhc_kernels.mhc_post_torch(
927 x, residual, post_layer_mix, comb_res_mix
928 )
929 post_mix_cur, comb_mix_cur, layer_input_cur = mhc_kernels.mhc_pre_torch(
930 residual_cur,
931 fn,
932 hc_scale,
933 hc_base,
934 rms_eps,
935 hc_pre_eps,
936 hc_sinkhorn_eps,
937 hc_post_mult_value,
938 sinkhorn_repeat,
939 )
940 return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
941
942 def forward_xpu(
943 self,
944 x: torch.Tensor,
945 residual: torch.Tensor,
946 post_layer_mix: torch.Tensor,
947 comb_res_mix: torch.Tensor,
948 fn: torch.Tensor,
949 hc_scale: torch.Tensor,
950 hc_base: torch.Tensor,
951 rms_eps: float,
952 hc_pre_eps: float,
953 hc_sinkhorn_eps: float,
954 hc_post_mult_value: float,
955 sinkhorn_repeat: int,
956 n_splits: int = 1,
957 tile_n: int = 1,
958 norm_weight: torch.Tensor | None = None,
959 norm_eps: float = 0.0,
960 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
961 return torch.ops._xpu_C.mhc_fused_post_pre(
962 x,
963 residual,
964 post_layer_mix,
965 comb_res_mix,
966 fn,
967 hc_scale,
968 hc_base,
969 rms_eps,
970 hc_pre_eps,
971 hc_sinkhorn_eps,
972 hc_post_mult_value,
973 sinkhorn_repeat,
974 )
975
976 def forward_cpu(
977 self,
978 x: torch.Tensor,
979 residual: torch.Tensor,
980 post_layer_mix: torch.Tensor,
981 comb_res_mix: torch.Tensor,
982 fn: torch.Tensor,
983 hc_scale: torch.Tensor,
984 hc_base: torch.Tensor,
985 rms_eps: float,
986 hc_pre_eps: float,
987 hc_sinkhorn_eps: float,
988 hc_post_mult_value: float,
989 sinkhorn_repeat: int,
990 n_splits: int = 1,
991 tile_n: int = 1,
992 norm_weight: torch.Tensor | None = None,
993 norm_eps: float = 0.0,
994 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
995 # Decompose into post + pre (no fused kernel available).
996 residual_cur = mhc_kernels.mhc_post_cpu(
997 x, residual, post_layer_mix, comb_res_mix
998 )
999 post_mix_cur, comb_mix_cur, layer_input_cur = mhc_kernels.mhc_pre_cpu(
1000 residual_cur,
1001 fn,
1002 hc_scale,
1003 hc_base,
1004 rms_eps,
1005 hc_pre_eps,
1006 hc_sinkhorn_eps,
1007 hc_post_mult_value,
1008 sinkhorn_repeat,
1009 )
1010 return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
1011
1012
1013 def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor:
1014 """[s, hidden_size] -> [s, n * hidden_size] by replication."""
1015 return x.unsqueeze(1).expand(-1, n, -1).contiguous()
1016
1017
1018 def hc_contract(x: torch.Tensor, n: int) -> torch.Tensor:
1019 """[s, n * hidden_size] -> [s, hidden_size] by averaging."""
1020 return x.mean(dim=1)