A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 22:45 UTC
add fused mhc_post_pre kernel (#41536)
Signed-off-by: george <george@inferact.ai> Co-authored-by: george <george@inferact.ai>
gnovack committed
May 10, 2026 at 19:56 UTC
171019ab192362c2c9f4ce9897f0a35ddcaff506
3 files changed
+533
-11
tests/kernels/test_mhc_kernels.py
new
+142
@@ -0,0 +1,142 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+import vllm.model_executor.layers.mhc as mhc_ops # noqa: F401
7
+from vllm.platforms import current_platform
8
+from vllm.utils.torch_utils import set_random_seed
9
+
10
+DEVICE = current_platform.device_type
11
+
12
+
13
+def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor:
14
+ x = x.softmax(-1) + eps
15
+ x = x / (x.sum(-2, keepdim=True) + eps)
16
+ for _ in range(repeat - 1):
17
+ x = x / (x.sum(-1, keepdim=True) + eps)
18
+ x = x / (x.sum(-2, keepdim=True) + eps)
19
+ return x
20
+
21
+
22
+def mhc_pre_ref(
23
+ residual: torch.Tensor,
24
+ fn: torch.Tensor,
25
+ hc_scale: torch.Tensor,
26
+ hc_base: torch.Tensor,
27
+ rms_eps: float,
28
+ hc_pre_eps: float,
29
+ hc_sinkhorn_eps: float,
30
+ hc_post_mult_value: float,
31
+ sinkhorn_repeat: int,
32
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
33
+ """mHC pre reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_pre.py#L303"""
34
+ hc_mult = residual.shape[-2]
35
+
36
+ residual_flat = residual.flatten(-2, -1).float()
37
+ sqrsum = residual_flat.square().sum(-1)
38
+ mixes = (
39
+ residual_flat @ fn.T * (sqrsum.unsqueeze(-1) / fn.shape[-1] + rms_eps).rsqrt()
40
+ )
41
+
42
+ hc_scale = torch.cat(
43
+ [
44
+ hc_scale[0].expand(hc_mult),
45
+ hc_scale[1].expand(hc_mult),
46
+ hc_scale[2].expand(hc_mult * hc_mult),
47
+ ],
48
+ )
49
+ mixes = mixes * hc_scale + hc_base
50
+
51
+ pre_mix = mixes[:, :hc_mult].sigmoid().unsqueeze(-1) + hc_pre_eps
52
+ post_mix = (
53
+ mixes[:, hc_mult : 2 * hc_mult].sigmoid() * hc_post_mult_value
54
+ ).unsqueeze(-1)
55
+ res_mix = mixes[:, 2 * hc_mult :].view(-1, hc_mult, hc_mult)
56
+
57
+ res_mix = sinkhorn_normalize_ref(
58
+ res_mix, repeat=sinkhorn_repeat, eps=hc_sinkhorn_eps
59
+ )
60
+
61
+ layer_input = (residual * pre_mix).sum(-2).bfloat16()
62
+
63
+ return post_mix, res_mix, layer_input
64
+
65
+
66
+def mhc_post_ref(
67
+ x: torch.Tensor,
68
+ residual: torch.Tensor,
69
+ post_layer_mix: torch.Tensor,
70
+ comb_res_mix: torch.Tensor,
71
+) -> torch.Tensor:
72
+ """mHC post reference kernel from tilelang repo: https://github.com/tile-ai/tilelang/blob/d135bd1cd2d2eee74fbb41dd0a0831a427194c86/examples/deepseek_mhc/example_mhc_post.py#L68"""
73
+ term2 = torch.bmm(comb_res_mix.mT, residual.float())
74
+ return (x.float().unsqueeze(-2) * post_layer_mix + term2).bfloat16()
75
+
76
+
77
+@pytest.mark.skipif(
78
+ not current_platform.is_cuda(),
79
+ reason="CUDA required",
80
+)
81
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
82
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
83
+@pytest.mark.parametrize("hc_mult", [4])
84
+def test_mhc_fused_post_pre(num_tokens, hidden_size, hc_mult):
85
+ torch.set_default_device(DEVICE)
86
+ set_random_seed(0)
87
+
88
+ x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
89
+ residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
90
+ post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
91
+ comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
92
+
93
+ hc_mult2 = hc_mult * hc_mult
94
+ hc_mult3 = hc_mult * 2 + hc_mult2
95
+ fn = (
96
+ torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float)
97
+ * 1e-4
98
+ * (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1))
99
+ ).flatten(1, 2)
100
+ hc_scale = torch.randn((3,), dtype=torch.float) * 0.1
101
+ hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1
102
+
103
+ hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6
104
+ sinkhorn_repeat = 20
105
+ hc_post_alpha = 1.0
106
+
107
+ def run_ref():
108
+ residual_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix)
109
+ post_mix_ref, res_mix_ref, layer_input_ref = mhc_pre_ref(
110
+ residual_ref,
111
+ fn,
112
+ hc_scale,
113
+ hc_base,
114
+ rms_eps,
115
+ hc_pre_eps,
116
+ hc_sinkhorn_eps,
117
+ hc_post_alpha,
118
+ sinkhorn_repeat,
119
+ )
120
+ return residual_ref, post_mix_ref, res_mix_ref, layer_input_ref
121
+
122
+ residual_ref, post_mix_ref, res_mix_ref, layer_input_ref = run_ref()
123
+
124
+ residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre(
125
+ x,
126
+ residual,
127
+ post_layer_mix,
128
+ comb_res_mix,
129
+ fn,
130
+ hc_scale,
131
+ hc_base,
132
+ rms_eps,
133
+ hc_pre_eps,
134
+ hc_sinkhorn_eps,
135
+ hc_post_alpha,
136
+ sinkhorn_repeat,
137
+ )
138
+
139
+ torch.testing.assert_close(residual, residual_ref, atol=1e-2, rtol=1e-2)
140
+ torch.testing.assert_close(post_mix, post_mix_ref, atol=1e-2, rtol=1e-2)
141
+ torch.testing.assert_close(res_mix, res_mix_ref, atol=1e-2, rtol=1e-2)
142
+ torch.testing.assert_close(x, layer_input_ref, atol=1e-2, rtol=1e-2)
vllm/model_executor/layers/mhc.py
+343
@@ -441,6 +441,131 @@ def mhc_post_tilelang(
441
T.pdl_trigger()
442
443
444
+@tilelang.jit(
445
+ pass_configs={
446
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
447
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
448
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
449
+ },
450
+)
451
+def mhc_fused_tilelang(
452
+ comb_mix,
453
+ residual_in,
454
+ post_mix,
455
+ x_in,
456
+ weight_t,
457
+ yp_out,
458
+ rp_out,
459
+ residual_out,
460
+ hc: int,
461
+ hidden: int,
462
+ n_out: int,
463
+ n_thr: int = 256,
464
+ h_blk: int = 256,
465
+ tile_n: int = 1,
466
+ split_k: int = 1,
467
+) -> tilelang.JITKernel:
468
+ """Fused mhc post-mapping + pre-norm GEMM FMA"""
469
+ m = T.dynamic("num_tokens")
470
+ split_k = T.dynamic("split_k")
471
+ h = hidden
472
+ h_blk = math.gcd(hidden, h_blk)
473
+ h_per_split = h // split_k
474
+ n_tiles = n_out // tile_n
475
+
476
+ comb_mix: T.Tensor((m, hc, hc), T.float32) # type: ignore[no-redef, valid-type]
477
+ residual_in: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
478
+ post_mix: T.Tensor((m, hc), T.float32) # type: ignore[no-redef, valid-type]
479
+ x_in: T.Tensor((m, h), T.bfloat16) # type: ignore[no-redef, valid-type]
480
+ weight_t: T.Tensor((n_out, hc, h), T.float32) # type: ignore[no-redef, valid-type]
481
+ yp_out: T.Tensor((split_k, m, n_out), T.float32) # type: ignore[no-redef, valid-type]
482
+ rp_out: T.Tensor((split_k, m), T.float32) # type: ignore[no-redef, valid-type]
483
+ residual_out: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
484
+
485
+ h_iters = h_per_split // n_thr
486
+ num_warps = n_thr // 32
487
+
488
+ with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
489
+ tid = T.get_thread_binding()
490
+ warp_id = T.get_warp_idx()
491
+ lane = T.get_lane_idx()
492
+
493
+ s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
494
+ s_post = T.alloc_shared((hc,), T.float32)
495
+ s_comb = T.alloc_shared((hc, hc), T.float32)
496
+
497
+ pm = T.alloc_local((hc,), T.float32)
498
+ cm = T.alloc_local((hc, hc), T.float32)
499
+ acc = T.alloc_local((tile_n,), T.float32)
500
+ sqr = T.alloc_local((1,), T.float32)
501
+ new_r = T.alloc_local((hc,), T.float32)
502
+
503
+ T.clear(acc)
504
+ T.clear(sqr)
505
+ h_split_start = i_ks * h_per_split
506
+
507
+ T.pdl_sync()
508
+
509
+ T.copy(post_mix[i_n, 0], s_post)
510
+ T.copy(comb_mix[i_n, 0, 0], s_comb)
511
+
512
+ for j in T.unroll(hc):
513
+ pm[j] = s_post[j]
514
+ for j in T.unroll(hc):
515
+ for k in T.unroll(hc):
516
+ cm[k, j] = s_comb[k, j]
517
+
518
+ # Each thread owns h_iters elements of the k-split's h slice.
519
+ for it in T.serial(h_iters):
520
+ h_idx = h_split_start + it * n_thr + tid
521
+
522
+ # Compute new residual from layer output and past residual
523
+ for j in T.unroll(hc):
524
+ new_r[j] = pm[j] * x_in[i_n, h_idx]
525
+ for k in T.unroll(hc):
526
+ new_r[j] += cm[k, j] * residual_in[i_n, k, h_idx]
527
+
528
+ # populate residual_out and compute sqr sum
529
+ if i_nt == 0:
530
+ for j in T.unroll(hc):
531
+ residual_out[i_n, j, h_idx] = new_r[j]
532
+ sqr[0] += new_r[j] * new_r[j]
533
+
534
+ # Per-thread FMA into acc[n]
535
+ for n in T.unroll(tile_n):
536
+ for j in T.unroll(hc):
537
+ acc[n] += weight_t[i_nt * tile_n + n, j, h_idx] * new_r[j]
538
+
539
+ for n in T.unroll(tile_n):
540
+ acc[n] = T.warp_reduce_sum(acc[n])
541
+ if i_nt == 0:
542
+ sqr[0] = T.warp_reduce_sum(sqr[0])
543
+
544
+ # Cross-warp reduce via shared mem
545
+ if lane == 0:
546
+ for n in T.unroll(tile_n):
547
+ s_warp[warp_id, n] = acc[n]
548
+ if i_nt == 0:
549
+ s_warp[warp_id, tile_n] = sqr[0]
550
+ T.sync_threads()
551
+
552
+ # Warp 0 does the final cross-warp sum and writes outputs
553
+ if warp_id == 0:
554
+ if lane < tile_n:
555
+ v = T.alloc_var(T.float32, init=0.0)
556
+ for w in T.unroll(num_warps):
557
+ v += s_warp[w, lane]
558
+ yp_out[i_ks, i_n, i_nt * tile_n + lane] = v
559
+
560
+ if i_nt == 0 and lane == 0:
561
+ v2 = T.alloc_var(T.float32, init=0.0)
562
+ for w in T.unroll(num_warps):
563
+ v2 += s_warp[w, tile_n]
564
+ rp_out[i_ks, i_n] = v2
565
+
566
+ T.pdl_trigger()
567
+
568
+
569
def mhc_post(
570
x: torch.Tensor,
571
residual: torch.Tensor,
@@ -468,6 +593,218 @@ def mhc_post(
593
return out
594
595
596
+def mhc_fused_post_pre(
597
+ x: torch.Tensor,
598
+ residual: torch.Tensor,
599
+ post_layer_mix: torch.Tensor,
600
+ comb_res_mix: torch.Tensor,
601
+ fn: torch.Tensor,
602
+ hc_scale: torch.Tensor,
603
+ hc_base: torch.Tensor,
604
+ rms_eps: float,
605
+ hc_pre_eps: float,
606
+ hc_sinkhorn_eps: float,
607
+ hc_post_mult_value: float,
608
+ sinkhorn_repeat: int,
609
+ n_splits: int = 1,
610
+ tile_n: int = 1,
611
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
612
+ """
613
+ Run one MHC post block followed by the next MHC pre block.
614
+
615
+ Returns:
616
+ residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
617
+ post_mix_cur: shape (..., hc_mult, 1)
618
+ comb_mix_cur: shape (..., hc_mult, hc_mult)
619
+ layer_input_cur: shape (..., hidden_size)
620
+ """
621
+
622
+ assert residual.dtype == torch.bfloat16
623
+ assert x.dtype == torch.bfloat16
624
+ assert post_layer_mix.dtype == torch.float32
625
+ assert comb_res_mix.dtype == torch.float32
626
+ assert fn.dtype == torch.float32
627
+ assert hc_scale.dtype == torch.float32
628
+ assert hc_base.dtype == torch.float32
629
+
630
+ hc_mult = residual.shape[-2]
631
+ hidden_size = residual.shape[-1]
632
+ hc_mult2 = hc_mult * hc_mult
633
+ hc_mult3 = hc_mult * 2 + hc_mult2
634
+ hc_hidden_size = hc_mult * hidden_size
635
+ outer_shape = residual.shape[:-2]
636
+
637
+ assert x.shape == (*outer_shape, hidden_size)
638
+ assert post_layer_mix.shape in (
639
+ (*outer_shape, hc_mult, 1),
640
+ (*outer_shape, hc_mult),
641
+ )
642
+ assert comb_res_mix.shape == (*outer_shape, hc_mult, hc_mult)
643
+ assert fn.shape == (hc_mult3, hc_hidden_size)
644
+ assert hc_scale.shape == (3,)
645
+ assert hc_base.shape == (hc_mult3,)
646
+
647
+ assert n_splits in (1, 2, 4, 8)
648
+ assert hidden_size % n_splits == 0
649
+
650
+ residual_flat = residual.view(-1, hc_mult, hidden_size)
651
+ num_tokens = residual_flat.shape[0]
652
+ x_flat = x.view(num_tokens, hidden_size)
653
+ post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
654
+ comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
655
+
656
+ fma_token_threshold = 16
657
+ if num_tokens <= fma_token_threshold:
658
+ # TODO(gnovack): investigate autotuning these heuristics
659
+ tile_n = 2 if num_tokens < 8 else 3
660
+ n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
661
+ else:
662
+ # these number are from deepgemm kernel impl
663
+ block_k = 64
664
+ block_m = 64
665
+ n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
666
+
667
+ gemm_out_mul = torch.empty(
668
+ n_splits,
669
+ num_tokens,
670
+ hc_mult3,
671
+ dtype=torch.float32,
672
+ device=residual.device,
673
+ )
674
+ gemm_out_sqrsum = torch.empty(
675
+ n_splits,
676
+ num_tokens,
677
+ dtype=torch.float32,
678
+ device=residual.device,
679
+ )
680
+ residual_cur = torch.empty_like(residual_flat)
681
+ post_mix_cur = torch.empty(
682
+ num_tokens,
683
+ hc_mult,
684
+ dtype=torch.float32,
685
+ device=residual.device,
686
+ )
687
+ comb_mix_cur = torch.empty(
688
+ num_tokens,
689
+ hc_mult2,
690
+ dtype=torch.float32,
691
+ device=residual.device,
692
+ )
693
+ layer_input_cur = torch.empty(
694
+ num_tokens,
695
+ hidden_size,
696
+ dtype=torch.bfloat16,
697
+ device=residual.device,
698
+ )
699
+
700
+ if num_tokens <= fma_token_threshold:
701
+ mhc_fused_tilelang(
702
+ comb_res_mix_flat,
703
+ residual_flat,
704
+ post_layer_mix_flat,
705
+ x_flat,
706
+ fn.view(hc_mult3, hc_mult, hidden_size),
707
+ gemm_out_mul,
708
+ gemm_out_sqrsum,
709
+ residual_cur,
710
+ hc_mult,
711
+ hidden_size,
712
+ hc_mult3,
713
+ tile_n=tile_n,
714
+ n_splits=n_splits,
715
+ )
716
+ else:
717
+ mhc_post_tilelang(
718
+ comb_res_mix_flat,
719
+ residual_flat,
720
+ post_layer_mix_flat,
721
+ x_flat,
722
+ residual_cur,
723
+ residual.shape[-2],
724
+ residual.shape[-1],
725
+ )
726
+
727
+ from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
728
+
729
+ tf32_hc_prenorm_gemm(
730
+ residual_cur.view(num_tokens, hc_mult * hidden_size),
731
+ fn,
732
+ gemm_out_mul,
733
+ gemm_out_sqrsum,
734
+ n_splits,
735
+ )
736
+
737
+ mhc_pre_big_fuse_tilelang(
738
+ gemm_out_mul,
739
+ gemm_out_sqrsum,
740
+ hc_scale,
741
+ hc_base,
742
+ residual_cur,
743
+ post_mix_cur,
744
+ comb_mix_cur,
745
+ layer_input_cur,
746
+ hidden_size,
747
+ rms_eps,
748
+ hc_pre_eps,
749
+ hc_sinkhorn_eps,
750
+ hc_post_mult_value,
751
+ sinkhorn_repeat,
752
+ n_splits,
753
+ hc_mult,
754
+ )
755
+
756
+ return (
757
+ residual_cur.view(*outer_shape, hc_mult, hidden_size),
758
+ post_mix_cur.view(*outer_shape, hc_mult, 1),
759
+ comb_mix_cur.view(*outer_shape, hc_mult, hc_mult),
760
+ layer_input_cur.view(*outer_shape, hidden_size),
761
+ )
762
+
763
+
764
+def _mhc_fused_post_pre_fake(
765
+ x: torch.Tensor,
766
+ residual: torch.Tensor,
767
+ post_layer_mix: torch.Tensor,
768
+ comb_res_mix: torch.Tensor,
769
+ fn: torch.Tensor,
770
+ hc_scale: torch.Tensor,
771
+ hc_base: torch.Tensor,
772
+ rms_eps: float,
773
+ hc_pre_eps: float,
774
+ hc_sinkhorn_eps: float,
775
+ hc_post_mult_value: float,
776
+ sinkhorn_repeat: int,
777
+ n_splits: int = 1,
778
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
779
+ hc_mult = residual.shape[-2]
780
+ hidden_size = residual.shape[-1]
781
+ outer_shape = residual.shape[:-2]
782
+
783
+ residual_cur = torch.empty_like(residual)
784
+ post_mix_cur = torch.empty(
785
+ *outer_shape,
786
+ hc_mult,
787
+ 1,
788
+ dtype=torch.float32,
789
+ device=residual.device,
790
+ )
791
+ comb_mix_cur = torch.empty(
792
+ *outer_shape,
793
+ hc_mult,
794
+ hc_mult,
795
+ dtype=torch.float32,
796
+ device=residual.device,
797
+ )
798
+ layer_input_cur = torch.empty(
799
+ *outer_shape,
800
+ hidden_size,
801
+ dtype=torch.bfloat16,
802
+ device=residual.device,
803
+ )
804
+
805
+ return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
806
+
807
+
808
def _mhc_post_fake(
809
x: torch.Tensor,
810
residual: torch.Tensor,
@@ -489,6 +826,12 @@ direct_register_custom_op(
826
mutates_args=[],
827
fake_impl=_mhc_post_fake,
828
)
829
+direct_register_custom_op(
830
+ op_name="mhc_fused_post_pre",
831
+ op_func=mhc_fused_post_pre,
832
+ mutates_args=[],
833
+ fake_impl=_mhc_fused_post_pre_fake,
834
+)
835
836
837
@tilelang.jit(
vllm/model_executor/models/deepseek_v4.py
+48
-11
@@ -1203,23 +1203,53 @@ class DeepseekV4DecoderLayer(nn.Module):
1203
x: torch.Tensor,
1204
positions: torch.Tensor,
1205
input_ids: torch.Tensor | None,
1206
+ post_mix: torch.Tensor | None,
1207
+ res_mix: torch.Tensor | None,
1208
+ residual: torch.Tensor | None,
1209
) -> torch.Tensor:
1207
- residual = x
1208
- x, post, comb = self.hc_pre(
1209
- x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1210
- )
1210
+ if residual is None:
1211
+ # Run standalone hc_pre on first layer
1212
+ residual = x
1213
+ x, post_mix, res_mix = self.hc_pre(
1214
+ x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1215
+ )
1216
+ else:
1217
+ residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre(
1218
+ x,
1219
+ residual,
1220
+ post_mix,
1221
+ res_mix,
1222
+ self.hc_attn_fn,
1223
+ self.hc_attn_scale,
1224
+ self.hc_attn_base,
1225
+ self.rms_norm_eps,
1226
+ self.hc_eps,
1227
+ self.hc_eps,
1228
+ self.hc_post_alpha,
1229
+ self.hc_sinkhorn_iters,
1230
+ )
1231
+
1232
x = self.attn_norm(x)
1233
x = self.attn(positions, x, None)
1213
- x = self.hc_post(x, residual, post, comb)
1234
1215
- residual = x
1216
- x, post, comb = self.hc_pre(
1217
- x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
1235
+ residual, post_mix, res_mix, x = torch.ops.vllm.mhc_fused_post_pre(
1236
+ x,
1237
+ residual,
1238
+ post_mix,
1239
+ res_mix,
1240
+ self.hc_ffn_fn,
1241
+ self.hc_ffn_scale,
1242
+ self.hc_ffn_base,
1243
+ self.rms_norm_eps,
1244
+ self.hc_eps,
1245
+ self.hc_eps,
1246
+ self.hc_post_alpha,
1247
+ self.hc_sinkhorn_iters,
1248
)
1249
+
1250
x = self.ffn_norm(x)
1251
x = self.ffn(x, input_ids)
1221
- x = self.hc_post(x, residual, post, comb)
1222
- return x
1252
+ return x, residual, post_mix, res_mix
1253
1254
1255
@support_torch_compile
@@ -1369,12 +1399,19 @@ class DeepseekV4Model(nn.Module):
1399
1400
if self.use_mega_moe:
1401
input_ids = input_ids.to(torch.int64)
1402
+
1403
+ residual, post_mix, res_mix = None, None, None
1404
for layer in islice(self.layers, self.start_layer, self.end_layer):
1373
- hidden_states = layer(
1405
+ hidden_states, residual, post_mix, res_mix = layer(
1406
hidden_states,
1407
positions,
1408
input_ids,
1409
+ post_mix,
1410
+ res_mix,
1411
+ residual,
1412
)
1413
+ else:
1414
+ hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
1415
1416
if not get_pp_group().is_last_rank:
1417
return IntermediateTensors({"hidden_states": hidden_states})