Skip to content

vllm.model_executor.kernels.mhc

Modules:

Functions:

hc_head_fused_cpu(hidden_states, hc_fn, hc_scale, hc_base, rms_norm_eps, hc_eps)

CPU-ported HC head reduction (see test_hc_head_cpu in tests/kernels/mhc/test_mhc_kernels.py for the eager reference this is tested against).

The ported kernel's C++ signature takes (hc_eps, norm_eps) -- the opposite order from this wrapper's (rms_norm_eps, hc_eps), which matches the real call site in models/deepseek_v4/cpu/model.py.

Source code in vllm/model_executor/kernels/mhc/cpu.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def hc_head_fused_cpu(
    hidden_states: torch.Tensor,
    hc_fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_norm_eps: float,
    hc_eps: float,
) -> torch.Tensor:
    """CPU-ported HC head reduction (see `test_hc_head_cpu` in
    tests/kernels/mhc/test_mhc_kernels.py for the eager reference this is tested
    against).

    The ported kernel's C++ signature takes `(hc_eps, norm_eps)` -- the
    opposite order from this wrapper's `(rms_norm_eps, hc_eps)`, which matches
    the real call site in `models/deepseek_v4/cpu/model.py`.
    """
    hc_mult, hidden_size = hidden_states.shape[-2:]
    outer_shape = hidden_states.shape[:-2]
    hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size)

    out = ops.hc_head_fused_cpu(hs_flat, hc_fn, hc_scale, hc_base, hc_eps, rms_norm_eps)
    return out.view(*outer_shape, hidden_size)

mhc_fused_post_pre_aiter(x, residual, post_layer_mix, comb_res_mix, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, n_splits=1, tile_n=1, norm_weight=None, norm_eps=0.0)

Fused mHC post + next mHC pre on ROCm via AITER.

Returns residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur.

Source code in vllm/model_executor/kernels/mhc/aiter.py 
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
def mhc_fused_post_pre_aiter(
    x: torch.Tensor,
    residual: torch.Tensor,
    post_layer_mix: torch.Tensor,
    comb_res_mix: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 1,
    tile_n: int = 1,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Fused mHC post + next mHC pre on ROCm via AITER.

    Returns residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur.
    """
    hidden_size = residual.shape[-1]
    assert hidden_size % 256 == 0
    from vllm._aiter_ops import rocm_aiter_ops

    return rocm_aiter_ops.mhc_fused_post_pre(
        x,
        residual,
        post_layer_mix,
        comb_res_mix,
        fn,
        hc_scale,
        hc_base,
        rms_eps,
        hc_pre_eps,
        hc_sinkhorn_eps,
        hc_post_mult_value,
        sinkhorn_repeat,
        norm_weight,
        norm_eps,
    )

mhc_fused_post_pre_delayed_rms_norm_aiter(residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, pre_mix, sublayer_out, post_layer_mix, comb_res_mix, norm_weight, norm_eps, residual_out)

:func:mhc_pre_delayed_aiter with the RMSNorm of the collapse folded in: layer_input comes back normalised with norm_weight / norm_eps.

Source code in vllm/model_executor/kernels/mhc/aiter.py 
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def mhc_fused_post_pre_delayed_rms_norm_aiter(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    pre_mix: torch.Tensor | None,
    sublayer_out: torch.Tensor | None,
    post_layer_mix: torch.Tensor | None,
    comb_res_mix: torch.Tensor | None,
    norm_weight: torch.Tensor,
    norm_eps: float,
    residual_out: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """:func:`mhc_pre_delayed_aiter` with the RMSNorm of the collapse folded in:
    ``layer_input`` comes back normalised with ``norm_weight`` / ``norm_eps``."""
    from vllm._aiter_ops import rocm_aiter_ops

    return rocm_aiter_ops.mhc_fused_post_pre_delayed_rms_norm(
        residual,
        fn,
        hc_scale,
        hc_base,
        rms_eps,
        hc_pre_eps,
        hc_sinkhorn_eps,
        hc_post_mult_value,
        sinkhorn_repeat,
        pre_mix,
        sublayer_out,
        post_layer_mix,
        comb_res_mix,
        norm_weight,
        norm_eps,
        residual_out,
    )

mhc_fused_post_pre_tilelang(x, residual, post_layer_mix, comb_res_mix, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, n_splits=1, tile_n=1, norm_weight=None, norm_eps=1e-06)

See _mhc_fused_post_pre_tilelang_impl. Never writes to its inputs.

Source code in vllm/model_executor/kernels/mhc/tilelang.py 
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
def mhc_fused_post_pre_tilelang(
    x: torch.Tensor,
    residual: torch.Tensor,
    post_layer_mix: torch.Tensor,
    comb_res_mix: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 1,
    tile_n: int = 1,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """See ``_mhc_fused_post_pre_tilelang_impl``. Never writes to its inputs."""
    return _mhc_fused_post_pre_tilelang_impl(
        x,
        residual,
        post_layer_mix,
        comb_res_mix,
        fn,
        hc_scale,
        hc_base,
        rms_eps,
        hc_pre_eps,
        hc_sinkhorn_eps,
        hc_post_mult_value,
        sinkhorn_repeat,
        norm_weight=norm_weight,
        norm_eps=norm_eps,
    )

mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix)

CPU-ported mHC post block (see mhc_post_torch for the eager reference).

Source code in vllm/model_executor/kernels/mhc/cpu.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def mhc_post_cpu(
    x: torch.Tensor,
    residual: torch.Tensor,
    post_layer_mix: torch.Tensor,
    comb_res_mix: torch.Tensor,
) -> torch.Tensor:
    """CPU-ported mHC post block (see `mhc_post_torch` for the eager reference)."""
    hc_mult, hidden_size = residual.shape[-2:]
    outer_shape = residual.shape[:-2]

    x_flat = x.reshape(-1, hidden_size)
    residual_flat = residual.reshape(-1, hc_mult, hidden_size)
    post_flat = post_layer_mix.reshape(-1, hc_mult).float()
    comb_flat = comb_res_mix.reshape(-1, hc_mult, hc_mult).float()

    out = ops.hc_post_fused_cpu(x_flat, residual_flat, post_flat, comb_flat)
    return out.view(*outer_shape, hc_mult, hidden_size)

mhc_pre_aiter(residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, n_splits=1, norm_weight=None, norm_eps=0.0)

Forward pass for mHC pre block.

Parameters:

  • residual

    (Tensor) –

    shape (..., hc_mult, hidden_size), dtype torch.bfloat16

  • fn

    (Tensor) –

    shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32

  • hc_scale

    (Tensor) –

    shape (3,), dtype torch.float32

  • hc_base

    (Tensor) –

    shape (hc_mult3,), dtype torch.float32

  • rms_eps

    (float) –

    RMS normalization epsilon

  • hc_pre_eps

    (float) –

    pre-mix epsilon

  • hc_sinkhorn_eps

    (float) –

    sinkhorn epsilon

  • hc_post_mult_value

    (float) –

    post-mix multiplier value

  • sinkhorn_repeat

    (int) –

    number of sinkhorn iterations

  • n_splits

    (int, default: 1 ) –

    split-k factor;

  • norm_weight

    (Tensor | None, default: None ) –

    optional RMSNorm weight fused into the pre kernel

  • norm_eps

    (float, default: 0.0 ) –

    epsilon for the fused RMSNorm when norm_weight is set

Returns:

  • post_mix ( Tensor ) –

    shape (..., hc_mult), dtype torch.float32

  • comb_mix ( Tensor ) –

    shape (..., hc_mult, hc_mult), dtype torch.float32

  • layer_input ( Tensor ) –

    shape (..., hidden_size), dtype torch.bfloat16

Source code in vllm/model_executor/kernels/mhc/aiter.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def mhc_pre_aiter(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 1,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward pass for mHC pre block.

    Args:
        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
        hc_scale: shape (3,), dtype torch.float32
        hc_base: shape (hc_mult3,), dtype torch.float32
        rms_eps: RMS normalization epsilon
        hc_pre_eps: pre-mix epsilon
        hc_sinkhorn_eps: sinkhorn epsilon
        hc_post_mult_value: post-mix multiplier value
        sinkhorn_repeat: number of sinkhorn iterations
        n_splits: split-k factor;
        norm_weight: optional RMSNorm weight fused into the pre kernel
        norm_eps: epsilon for the fused RMSNorm when norm_weight is set

    Returns:
        post_mix: shape (..., hc_mult), dtype torch.float32
        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
        layer_input: shape (..., hidden_size), dtype torch.bfloat16

    """
    hidden_size = residual.shape[-1]
    assert hidden_size % 256 == 0
    from vllm._aiter_ops import rocm_aiter_ops

    return rocm_aiter_ops.mhc_pre(
        residual,
        fn,
        hc_scale,
        hc_base,
        rms_eps,
        hc_pre_eps,
        hc_sinkhorn_eps,
        hc_post_mult_value,
        sinkhorn_repeat,
        norm_weight,
        norm_eps,
    )

mhc_pre_cpu(residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, n_splits=1, norm_weight=None, norm_eps=0.0)

CPU-ported mHC pre block (see mhc_pre_torch for the eager reference).

The ported kernel (hc_pre_fused_cpu) only exposes one merged hc_eps (used for both hc_pre_eps/hc_sinkhorn_eps) and hardcodes the post-mix multiplier to 2.0 -- true of every real call site in models/deepseek_v4/cpu/model.py, so this is not a capability loss here.

Source code in vllm/model_executor/kernels/mhc/cpu.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def mhc_pre_cpu(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 1,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """CPU-ported mHC pre block (see `mhc_pre_torch` for the eager reference).

    The ported kernel (`hc_pre_fused_cpu`) only exposes one merged `hc_eps`
    (used for both `hc_pre_eps`/`hc_sinkhorn_eps`) and hardcodes the post-mix
    multiplier to 2.0 -- true of every real call site in
    `models/deepseek_v4/cpu/model.py`, so this is not a capability loss here.
    """
    assert n_splits == 1, "mhc_pre_cpu does not support n_splits != 1"
    assert hc_pre_eps == hc_sinkhorn_eps, (
        "mhc_pre_cpu requires hc_pre_eps == hc_sinkhorn_eps (single merged hc_eps)"
    )
    assert hc_post_mult_value == 2.0, "mhc_pre_cpu hardcodes the post-mix multiplier"

    hc_mult, hidden_size = residual.shape[-2:]
    outer_shape = residual.shape[:-2]
    x_flat = residual.reshape(-1, hc_mult, hidden_size)

    layer_input, post, comb = ops.hc_pre_fused_cpu(
        x_flat,
        fn,
        hc_scale,
        hc_base,
        hc_mult,
        sinkhorn_repeat,
        rms_eps,
        hc_pre_eps,
    )

    post_mix = post.view(*outer_shape, hc_mult, 1)
    comb_mix = comb.view(*outer_shape, hc_mult, hc_mult)
    layer_input = layer_input.view(*outer_shape, hidden_size)
    return post_mix, comb_mix, layer_input

mhc_pre_delayed_aiter(residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, pre_mix=None, sublayer_out=None, post_layer_mix=None, comb_res_mix=None, residual_out=None)

MHC pre with the pre-mix carried in from the previous sublayer.

Matches mhc_pre_delayed_torch: the stream collapse uses pre_mix rather than the gate computed here, and that gate is returned as the pre-mix for the next sublayer seam.

Parameters:

  • residual

    (Tensor) –

    shape (..., hc_mult, hidden_size), dtype torch.bfloat16

  • fn

    (Tensor) –

    shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32

  • hc_scale

    (Tensor) –

    shape (3,), dtype torch.float32

  • hc_base

    (Tensor) –

    shape (hc_mult3,), dtype torch.float32

  • rms_eps

    (float) –

    RMS normalization epsilon

  • hc_pre_eps

    (float) –

    pre-mix epsilon

  • hc_sinkhorn_eps

    (float) –

    sinkhorn epsilon

  • hc_post_mult_value

    (float) –

    post-mix multiplier value

  • sinkhorn_repeat

    (int) –

    number of sinkhorn iterations

  • pre_mix

    (Tensor | None, default: None ) –

    shape (..., hc_mult) from the previous sublayer, or None at model entry to select residual stream zero.

  • sublayer_out

    (Tensor | None, default: None ) –

    attention or FFN output, shape (..., hidden_size). When given with the two mixes below, the preceding post block is applied here so AITER can fold it into the pre projection.

  • post_layer_mix

    (Tensor | None, default: None ) –

    shape (..., hc_mult, 1), post gate for that block.

  • comb_res_mix

    (Tensor | None, default: None ) –

    shape (..., hc_mult, hc_mult), residual comb for it.

  • residual_out

    (Tensor | None, default: None ) –

    shape (..., hc_mult, hidden_size), written with the post block's new residual. Required exactly when the post is folded in.

Returns:

  • post_mix ( Tensor ) –

    shape (..., hc_mult, 1), dtype torch.float32

  • comb_mix ( Tensor ) –

    shape (..., hc_mult, hc_mult), dtype torch.float32

  • layer_input ( Tensor ) –

    shape (..., hidden_size), dtype torch.bfloat16

  • next_pre_mix ( Tensor ) –

    shape (..., hc_mult), dtype torch.float32

Source code in vllm/model_executor/kernels/mhc/aiter.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def mhc_pre_delayed_aiter(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    pre_mix: torch.Tensor | None = None,
    sublayer_out: torch.Tensor | None = None,
    post_layer_mix: torch.Tensor | None = None,
    comb_res_mix: torch.Tensor | None = None,
    residual_out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """MHC pre with the pre-mix carried in from the previous sublayer.

    Matches ``mhc_pre_delayed_torch``: the stream collapse uses *pre_mix*
    rather than the gate computed here, and that gate is returned as the
    pre-mix for the next sublayer seam.

    Args:
        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
        hc_scale: shape (3,), dtype torch.float32
        hc_base: shape (hc_mult3,), dtype torch.float32
        rms_eps: RMS normalization epsilon
        hc_pre_eps: pre-mix epsilon
        hc_sinkhorn_eps: sinkhorn epsilon
        hc_post_mult_value: post-mix multiplier value
        sinkhorn_repeat: number of sinkhorn iterations
        pre_mix: shape (..., hc_mult) from the previous sublayer, or None at
            model entry to select residual stream zero.
        sublayer_out: attention or FFN output, shape (..., hidden_size). When
            given with the two mixes below, the preceding post block is
            applied here so AITER can fold it into the pre projection.
        post_layer_mix: shape (..., hc_mult, 1), post gate for that block.
        comb_res_mix: shape (..., hc_mult, hc_mult), residual comb for it.
        residual_out: shape (..., hc_mult, hidden_size), written with the post
            block's new residual. Required exactly when the post is folded in.

    Returns:
        post_mix: shape (..., hc_mult, 1), dtype torch.float32
        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
        layer_input: shape (..., hidden_size), dtype torch.bfloat16
        next_pre_mix: shape (..., hc_mult), dtype torch.float32

    """
    hidden_size = residual.shape[-1]
    assert hidden_size % 256 == 0
    from vllm._aiter_ops import rocm_aiter_ops

    return rocm_aiter_ops.mhc_pre_delayed(
        residual,
        fn,
        hc_scale,
        hc_base,
        rms_eps,
        hc_pre_eps,
        hc_sinkhorn_eps,
        hc_post_mult_value,
        sinkhorn_repeat,
        pre_mix,
        sublayer_out,
        post_layer_mix,
        comb_res_mix,
        residual_out,
    )

mhc_pre_delayed_torch(residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, pre_mix=None, x=None)

Reference for mHC pre using coefficients from the previous sublayer.

Source code in vllm/model_executor/kernels/mhc/torch.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def mhc_pre_delayed_torch(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    pre_mix: torch.Tensor | None = None,
    x: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Reference for mHC pre using coefficients from the previous sublayer."""
    hc_mult = residual.shape[1]
    x = (residual.flatten(1) if x is None else x).float()
    mixes = (x @ fn.t()) * torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_eps)
    pre = (
        torch.sigmoid(mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]) + hc_pre_eps
    )
    post = (
        torch.sigmoid(
            mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1]
            + hc_base[hc_mult : 2 * hc_mult]
        )
        * hc_post_mult_value
    )
    comb = mixes[:, 2 * hc_mult :].view(-1, hc_mult, hc_mult) * hc_scale[2]
    comb = comb + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
    comb = torch.softmax(comb, dim=-1) + hc_sinkhorn_eps
    comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
    for _ in range(sinkhorn_repeat - 1):
        comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
        comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
    layer_input = (
        residual[:, 0]
        if pre_mix is None
        else (pre_mix.unsqueeze(-1) * residual.float()).sum(dim=1).to(residual.dtype)
    )
    return post.unsqueeze(-1), comb, layer_input, pre

mhc_pre_mix_triton(gemm_out, sqrsum, hc_scale, hc_base, hc_mult, hc_hidden_size, rms_eps, hc_pre_eps)

Pre-mix gate for the delayed mHC pre, from AITER's split-k GEMM output.

AITER's mhc_pre_big_fuse consumes the unreduced [splitk, tokens, hc_mult3] GEMM output and the matching row square-sums, but only returns the post and comb gates. The delayed formulation also needs the pre gate, to carry into the next sublayer seam. It is the same slice of the same numbers, so recover it here rather than repeating the projection.

Source code in vllm/model_executor/kernels/mhc/triton.py 
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
def mhc_pre_mix_triton(
    gemm_out: Tensor,
    sqrsum: Tensor,
    hc_scale: Tensor,
    hc_base: Tensor,
    hc_mult: int,
    hc_hidden_size: int,
    rms_eps: float,
    hc_pre_eps: float,
) -> Tensor:
    """Pre-mix gate for the delayed mHC pre, from AITER's split-k GEMM output.

    AITER's ``mhc_pre_big_fuse`` consumes the unreduced ``[splitk, tokens,
    hc_mult3]`` GEMM output and the matching row square-sums, but only returns
    the post and comb gates. The delayed formulation also needs the pre gate,
    to carry into the next sublayer seam. It is the same slice of the same
    numbers, so recover it here rather than repeating the projection.
    """
    assert gemm_out.ndim == 3 and gemm_out.dtype == torch.float32
    assert sqrsum.ndim == 2 and sqrsum.dtype == torch.float32
    splitk, num_tokens = gemm_out.shape[0], gemm_out.shape[1]
    assert sqrsum.shape == (splitk, num_tokens)

    out = torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=gemm_out.device)
    if num_tokens == 0:
        return out

    _mhc_pre_mix_kernel[(num_tokens,)](
        gemm_out,
        sqrsum,
        hc_scale,
        hc_base,
        out,
        splitk,
        hc_mult,
        gemm_out.stride(0),
        gemm_out.stride(1),
        gemm_out.stride(2),
        sqrsum.stride(0),
        sqrsum.stride(1),
        out.stride(0),
        out.stride(1),
        1.0 / hc_hidden_size,
        rms_eps,
        hc_pre_eps,
        SPLITK_BLOCK=triton.next_power_of_2(splitk),
        HC_BLOCK=triton.next_power_of_2(hc_mult),
        num_warps=1,
        # Match the separate FP32 multiply and add in the Torch reference.
        enable_fp_fusion=False,
    )
    return out

mhc_pre_tilelang(residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, n_splits=1, norm_weight=None, norm_eps=1e-06)

Forward pass for mHC pre block.

Parameters:

  • residual

    (Tensor) –

    shape (..., hc_mult, hidden_size), dtype torch.bfloat16

  • fn

    (Tensor) –

    shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32

  • hc_scale

    (Tensor) –

    shape (3,), dtype torch.float32

  • hc_base

    (Tensor) –

    shape (hc_mult3,), dtype torch.float32

  • rms_eps

    (float) –

    RMS normalization epsilon

  • hc_pre_eps

    (float) –

    pre-mix epsilon

  • hc_sinkhorn_eps

    (float) –

    sinkhorn epsilon

  • hc_post_mult_value

    (float) –

    post-mix multiplier value

  • sinkhorn_repeat

    (int) –

    number of sinkhorn iterations

  • n_splits

    (int, default: 1 ) –

    retained for the shared MHC operator API; the active GEMM backend selects its split factor internally.

  • norm_weight

    (Tensor | None, default: None ) –

    optional RMSNorm weight, shape (hidden_size,), dtype torch.bfloat16. When provided, RMSNorm is fused into the layer_input write path of the big_fuse kernel.

  • norm_eps

    (float, default: 1e-06 ) –

    epsilon for the fused RMSNorm; only consulted when norm_weight is given.

Returns:

  • post_mix ( Tensor ) –

    shape (..., hc_mult), dtype torch.float32

  • comb_mix ( Tensor ) –

    shape (..., hc_mult, hc_mult), dtype torch.float32

  • layer_input ( Tensor ) –

    shape (..., hidden_size), dtype torch.bfloat16

Source code in vllm/model_executor/kernels/mhc/tilelang.py 
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
def mhc_pre_tilelang(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 1,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward pass for mHC pre block.

    Args:
        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
        hc_scale: shape (3,), dtype torch.float32
        hc_base: shape (hc_mult3,), dtype torch.float32
        rms_eps: RMS normalization epsilon
        hc_pre_eps: pre-mix epsilon
        hc_sinkhorn_eps: sinkhorn epsilon
        hc_post_mult_value: post-mix multiplier value
        sinkhorn_repeat: number of sinkhorn iterations
        n_splits: retained for the shared MHC operator API; the active GEMM
            backend selects its split factor internally.
        norm_weight: optional RMSNorm weight, shape (hidden_size,), dtype
            torch.bfloat16. When provided, RMSNorm is fused into the
            layer_input write path of the big_fuse kernel.
        norm_eps: epsilon for the fused RMSNorm; only consulted when
            norm_weight is given.

    Returns:
        post_mix: shape (..., hc_mult), dtype torch.float32
        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
        layer_input: shape (..., hidden_size), dtype torch.bfloat16

    """
    from vllm.model_executor.kernels.mhc.tilelang_kernels import (
        _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
    )

    assert residual.dtype == torch.bfloat16
    assert fn.dtype == torch.float32
    assert hc_scale.dtype == torch.float32
    assert hc_base.dtype == torch.float32

    hc_mult = residual.shape[-2]
    hidden_size = residual.shape[-1]
    hc_mult2 = hc_mult * hc_mult
    hc_mult3 = hc_mult * 2 + hc_mult2

    hc_hidden_size = hc_mult * hidden_size
    assert fn.shape[0] == hc_mult3
    assert fn.shape[1] == hc_hidden_size
    assert hc_scale.shape == (3,)
    assert hc_base.shape == (hc_mult3,)

    if norm_weight is not None:
        assert norm_weight.shape == (hidden_size,)
        if norm_weight.dtype != torch.bfloat16:
            norm_weight = norm_weight.to(torch.bfloat16)
        if not norm_weight.is_contiguous():
            norm_weight = norm_weight.contiguous()

    outer_shape = residual.shape[:-2]

    residual_flat = residual.view(-1, hc_mult, hidden_size)
    num_tokens = residual_flat.shape[0]

    post_mix = torch.empty(
        num_tokens, hc_mult, dtype=torch.float32, device=residual.device
    )
    comb_mix = torch.empty(
        num_tokens, hc_mult2, dtype=torch.float32, device=residual.device
    )
    layer_input = torch.empty(
        num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
    )

    gemm_out_mul, gemm_out_sqrsum = _hc_prenorm_gemm_outputs(
        residual_flat.view(num_tokens, hc_mult * hidden_size),
        fn,
        hidden_size=hidden_size,
        hc_mult=hc_mult,
    )
    _MHC_PRE_BIG_FUSE_TILELANG_KERNEL(
        gemm_out_mul,
        gemm_out_sqrsum,
        hc_scale,
        hc_base,
        residual_flat,
        post_mix,
        comb_mix,
        layer_input,
        rms_eps,
        hc_pre_eps,
        hc_sinkhorn_eps,
        hc_post_mult_value,
        sinkhorn_repeat,
        norm_weight=norm_weight,
        norm_eps=norm_eps,
    )

    return (
        post_mix.view(*outer_shape, hc_mult, 1),
        comb_mix.view(*outer_shape, hc_mult, hc_mult),
        layer_input.view(*outer_shape, hidden_size),
    )

mhc_pre_torch(residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, n_splits=1)

Forward pass for mHC pre block.

Parameters:

  • residual

    (Tensor) –

    shape (..., hc_mult, hidden_size), dtype torch.bfloat16

  • fn

    (Tensor) –

    shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32

  • hc_scale

    (Tensor) –

    shape (3,), dtype torch.float32

  • hc_base

    (Tensor) –

    shape (hc_mult3,), dtype torch.float32

  • rms_eps

    (float) –

    RMS normalization epsilon

  • hc_pre_eps

    (float) –

    pre-mix epsilon

  • hc_sinkhorn_eps

    (float) –

    sinkhorn epsilon

  • hc_post_mult_value

    (float) –

    post-mix multiplier value

  • sinkhorn_repeat

    (int) –

    number of sinkhorn iterations

  • n_splits

    (int, default: 1 ) –

    split-k factor;

Returns:

  • post_mix ( Tensor ) –

    shape (..., hc_mult), dtype torch.float32

  • comb_mix ( Tensor ) –

    shape (..., hc_mult, hc_mult), dtype torch.float32

  • layer_input ( Tensor ) –

    shape (..., hidden_size), dtype torch.bfloat16

Source code in vllm/model_executor/kernels/mhc/torch.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def mhc_pre_torch(
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 1,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward pass for mHC pre block.

    Args:
        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
        hc_scale: shape (3,), dtype torch.float32
        hc_base: shape (hc_mult3,), dtype torch.float32
        rms_eps: RMS normalization epsilon
        hc_pre_eps: pre-mix epsilon
        hc_sinkhorn_eps: sinkhorn epsilon
        hc_post_mult_value: post-mix multiplier value
        sinkhorn_repeat: number of sinkhorn iterations
        n_splits: split-k factor;

    Returns:
        post_mix: shape (..., hc_mult), dtype torch.float32
        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
        layer_input: shape (..., hidden_size), dtype torch.bfloat16

    """
    # Validate shapes
    assert residual.dtype == torch.bfloat16
    assert fn.dtype == torch.float32
    assert hc_scale.dtype == torch.float32
    assert hc_base.dtype == torch.float32

    hc_mult = residual.shape[-2]
    hidden_size = residual.shape[-1]
    hc_mult2 = hc_mult * hc_mult
    hc_mult3 = hc_mult * 2 + hc_mult2

    hc_hidden_size = hc_mult * hidden_size
    assert fn.shape[0] == hc_mult3
    assert fn.shape[1] == hc_hidden_size
    assert hc_scale.shape == (3,)
    assert hc_base.shape == (hc_mult3,)

    outer_shape = residual.shape[:-2]

    residual_flat = residual.view(-1, hc_mult, hidden_size)
    num_tokens = residual_flat.shape[0]
    fn_flat = fn

    x = residual_flat.view(num_tokens, hc_mult * hidden_size).to(torch.float32)
    mixes = torch.matmul(x, fn_flat.t())
    sqrsum = x.square().sum(dim=-1, keepdim=True)
    mixes = mixes * torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)

    pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]
    pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps

    post_logits = (
        mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1] + hc_base[hc_mult : 2 * hc_mult]
    )
    post_mix = torch.sigmoid(post_logits) * hc_post_mult_value

    comb_logits = mixes[:, 2 * hc_mult :].view(num_tokens, hc_mult, hc_mult) * hc_scale[
        2
    ] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
    comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps
    comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
    for _ in range(sinkhorn_repeat - 1):
        comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
        comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)

    layer_input = torch.sum(
        pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1
    ).to(torch.bfloat16)
    return (
        post_mix.view(*outer_shape, hc_mult, 1),
        comb_mix.view(*outer_shape, hc_mult, hc_mult),
        layer_input.view(*outer_shape, hidden_size),
    )