recipes/docker-images/glm53-kindling/experimental/fixes/tilelang_kernels.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

import math
from dataclasses import dataclass
from functools import cache
from typing import Any

import torch

from vllm.model_executor.warmup.jit_warmup import WarmupIntRange, zip_inputs
from vllm.model_executor.warmup.jit_warmup_tilelang_helper import (
    TileLangLaunchSpec,
    VllmTileLangJitKernel,
    kernel_launcher,
    make_tilelang_warmup_tensor,
)
from vllm.platforms import current_platform
from vllm.tilelang_utils import T, tilelang, tilelang_jit
from vllm.utils.math_utils import cdiv

ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()

# TileLang's HIP AllReduce shuffle phase ignores a nonzero reduction-group
# offset, so the 64-thread RMSNorm reduction must be wave64-aligned and use
# threads 0..63.
_MIX_ON_TRAILING_THREADS = current_platform.is_rocm()


def _is_mix_thread():
    if _MIX_ON_TRAILING_THREADS:
        return T.get_thread_binding() >= 64
    return T.get_thread_binding() < 32


def _is_collapse_thread():
    if _MIX_ON_TRAILING_THREADS:
        return T.get_thread_binding() < 64
    return T.get_thread_binding() >= 32


@cache
def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
    device_props = torch.cuda.get_device_properties(0)
    n_sms = device_props.multi_processor_count
    split_k = n_sms // grid_size
    if k is not None:
        # avoid split_k for small k
        num_block_k = cdiv(k, block_k)
        split_k = min(split_k, num_block_k // 4)
    split_k = max(split_k, 1)
    return split_k


# ``mhc_fused_tilelang`` splits the hidden size across a fixed thread block, so
# every (tile_n, n_thr) pair is one more kernel to compile. Tuned on GB300
# (hc_mult 4, hidden_size 5120 and 7168) against the separate post + split-k
# GEMM it replaces; the bands are coarse because the loss against a per-shape
# optimum is ~1-2%, while one fixed config gives up to 18% at the top of the
# range.
_FUSED_POST_PRE_N_THR = 128
_FUSED_POST_PRE_MAX_TOKENS = int(__import__("os").environ.get("VLLM_MHC_FUSED_MAX_TOKENS", "32"))


def mhc_fused_post_pre_split_config(
    num_tokens: int, hidden_size: int, hc_mult: int
) -> tuple[int, int, int] | None:
    """Pick ``(tile_n, n_splits, n_thr)`` for the fused post + pre-norm GEMM.

    Returns None when a separate post kernel followed by a split-k GEMM is the
    better choice, or when the hidden size does not divide evenly across the
    fused kernel's block. One source of truth for the dispatch decision, the
    compile key and the launch, which must agree.
    """
    if num_tokens > _FUSED_POST_PRE_MAX_TOKENS:
        return None
    n_thr = _FUSED_POST_PRE_N_THR
    if hidden_size % n_thr:
        return None
    # The projection tiles supply the parallelism a small token count cannot,
    # so the tile grows with the batch: 1.3x at one token and 1.03x at 32
    # against a split-k GEMM given the split its own estimator computes.
    if num_tokens < 16:
        tile_n, n_splits = 2, 8
    else:
        tile_n, n_splits = 6, 8
    # The kernel drops projection tiles and k-slices it cannot fill evenly.
    mix_size = hc_mult * (hc_mult + 2)
    while tile_n > 1 and mix_size % tile_n:
        tile_n -= 1
    while n_splits > 1 and hidden_size % (n_splits * n_thr):
        n_splits //= 2
    return tile_n, n_splits, n_thr


def require_fused_post_pre_config(
    num_tokens: int, hidden_size: int, hc_mult: int
) -> tuple[int, int, int]:
    """The config for a shape the caller has already committed to fusing."""
    config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
    if config is None:
        raise ValueError(
            "the fused mHC post + pre-norm GEMM does not cover num_tokens="
            f"{num_tokens} at hidden_size={hidden_size}"
        )
    return config


def mhc_fused_post_pre_splits(hidden_size: int, hc_mult: int) -> tuple[int, ...]:
    """Every split-k factor the fused post + pre-norm GEMM path can pick."""
    configs = (
        mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
        for num_tokens in range(1, _FUSED_POST_PRE_MAX_TOKENS + 1)
    )
    return tuple(sorted({config[1] for config in configs if config is not None}))


@tilelang_jit
def mhc_pre_big_fuse_tilelang(
    gemm_out_mul,
    gemm_out_sqrsum,
    hc_scale,
    hc_base,
    residual,
    post_mix,
    comb_mix,
    layer_input,
    pre_mix_in,
    pre_mix_out,
    aux_out,
    hidden_size: int,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 16,
    hc_mult: int = 4,
    use_pre_mix_in: bool = False,
    save_pre_mix: bool = False,
    rms_numel: int = 0,
    write_aux: bool = False,
):
    """Fuse coefficient generation and residual collapse after the projection.

    With save_pre_mix, store the new pre-mix and collapse with pre_mix_in,
    or select stream zero when use_pre_mix_in is false.
    """
    num_tokens = T.dynamic("num_tokens")
    hc_mult3 = hc_mult * (2 + hc_mult)
    hidden_block = math.gcd(512, hidden_size)
    if rms_numel == 0:
        rms_numel = hc_mult * hidden_size

    gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32]  # type: ignore[no-redef, valid-type]
    gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32]  # type: ignore[no-redef, valid-type]
    hc_scale: T.Tensor[[3], T.float32]  # type: ignore[no-redef, valid-type]
    hc_base: T.Tensor[[hc_mult3], T.float32]  # type: ignore[no-redef, valid-type]
    residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
    # outputs
    post_mix: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]

    pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]

    with T.Kernel(num_tokens, threads=96) as i:
        if ENABLE_PDL:
            T.pdl_sync()
        ##################################################################
        # _pre_norm_fn_fwd_norm
        rms = T.alloc_fragment(1, T.float32)
        mixes = T.alloc_fragment(hc_mult3, T.float32)
        T.clear(mixes)
        rms[0] = 0
        for i_split in T.serial(n_splits):
            rms[0] += gemm_out_sqrsum[i_split, i]
        rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)
        for j in T.Parallel(hc_mult3):
            mixes[j] = 0
            for i_split in T.serial(n_splits):
                mixes[j] += gemm_out_mul[i_split, i, j]
            mixes[j] *= rms[0]
        mixes_shared = T.alloc_shared(hc_mult3, T.float32)
        T.copy(mixes, mixes_shared)

        if T.get_thread_binding() < 32:
            ##################################################################
            # _pre_split_mixes_fwd (post & comb)
            cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
            for j in T.Parallel(hc_mult):
                if save_pre_mix:
                    pre_mix_out[i, j] = (
                        T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
                        + hc_pre_eps
                    )
                post_mix[i, j] = (
                    T.sigmoid(
                        mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
                    )
                    * hc_post_mult_value
                )
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = (
                    mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
                    + hc_base[j * hc_mult + k + hc_mult * 2]
                )

            ##################################################################
            # _sinkhorn_fwd
            row_sum = T.alloc_fragment(hc_mult, T.float32)
            col_sum = T.alloc_fragment(hc_mult, T.float32)

            # comb = comb.softmax(-1) + eps
            row_max = T.alloc_fragment(hc_mult, T.float32)
            T.reduce_max(cm, row_max, dim=1)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = T.exp(cm[j, k] - row_max[j])
            T.reduce_sum(cm, row_sum, dim=1)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps

            # comb = comb / (comb.sum(-2) + eps)
            T.reduce_sum(cm, col_sum, dim=0)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)

            for _ in T.serial(sinkhorn_repeat - 1):
                # comb = comb / (comb.sum(-1) + eps)
                T.reduce_sum(cm, row_sum, dim=1)
                for j, k in T.Parallel(hc_mult, hc_mult):
                    cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)

                # comb = comb / (comb.sum(-2) + eps)
                T.reduce_sum(cm, col_sum, dim=0)
                for j, k in T.Parallel(hc_mult, hc_mult):
                    cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)

            # save comb_mix to global memory
            for j, k in T.Parallel(hc_mult, hc_mult):
                comb_mix[i, j * hc_mult + k] = cm[j, k]
        else:
            ##################################################################
            # _pre_split_mixes_fwd (pre)
            pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
            for j in T.Parallel(hc_mult):
                if use_pre_mix_in:
                    pre_mix_shared[j] = pre_mix_in[i, j]
                elif save_pre_mix:
                    pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0)
                else:
                    pre_mix_shared[j] = (
                        T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
                        + hc_pre_eps
                    )
            ###################################################################
            # _pre_apply_mix_fwd
            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
                xs = T.alloc_shared((hc_mult, hidden_block), T.float32)
                xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
                T.copy(residual[i, 0, i0_h * hidden_block], xs)
                T.copy(xs, xl)

                ol = T.alloc_fragment(hidden_block, T.float32)
                T.clear(ol)
                if write_aux:
                    # Aux consumers (draft models) take the plain stream mean
                    # of the same residual this collapse reads.
                    aux = T.alloc_fragment(hidden_block, T.float32)
                    T.clear(aux)

                for i_hc in T.serial(hc_mult):
                    pre = pre_mix_shared[i_hc]
                    for i1_h in T.Parallel(hidden_block):
                        ol[i1_h] += pre * xl[i_hc, i1_h]
                        if write_aux:
                            aux[i1_h] += xl[i_hc, i1_h]

                if write_aux:
                    for i1_h in T.Parallel(hidden_block):
                        aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16(
                            aux[i1_h] / hc_mult
                        )

                T.copy(ol, layer_input[i, i0_h * hidden_block])

        if ENABLE_PDL:
            T.pdl_trigger()


# Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478


@tilelang_jit
def mhc_pre_big_fuse_with_norm_tilelang(
    gemm_out_mul,
    gemm_out_sqrsum,
    hc_scale,
    hc_base,
    residual,
    post_mix,
    comb_mix,
    layer_input,
    norm_weight,
    pre_mix_in,
    pre_mix_out,
    aux_out,
    hidden_size: int,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    norm_eps: float,
    n_splits: int = 16,
    hc_mult: int = 4,
    gemm_last_dim: int = -1,
    use_pre_mix_in: bool = False,
    save_pre_mix: bool = False,
    rms_numel: int = 0,
    split_mode: str = "fused",
    write_aux: bool = False,
):
    # Split modes require shifted mHC: input collapse uses a carried pre-mix.
    assert split_mode in ("fused", "stats", "input")
    assert split_mode == "fused" or save_pre_mix
    num_tokens = T.dynamic("num_tokens")
    hc_mult3 = hc_mult * (2 + hc_mult)
    if gemm_last_dim < 0:
        gemm_last_dim = hc_mult3
    if rms_numel == 0:
        rms_numel = hc_mult * hidden_size
    hidden_block = math.gcd(1024, hidden_size)

    gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32]  # type: ignore[no-redef, valid-type]
    gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32]  # type: ignore[no-redef, valid-type]
    hc_scale: T.Tensor[[3], T.float32]  # type: ignore[no-redef, valid-type]
    hc_base: T.Tensor[[hc_mult3], T.float32]  # type: ignore[no-redef, valid-type]
    residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
    post_mix: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
    norm_weight: T.Tensor[[hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]

    pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]

    with T.Kernel(num_tokens, threads=96) as i:
        if ENABLE_PDL:
            T.pdl_sync()
        mixes_shared = T.alloc_shared(hc_mult3, T.float32)
        if split_mode != "input":
            rms = T.alloc_fragment(1, T.float32)
            mixes = T.alloc_fragment(hc_mult3, T.float32)
            T.clear(mixes)
            rms[0] = 0

            for i_split in T.serial(n_splits):
                rms[0] += gemm_out_sqrsum[i_split, i]
            rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)
            for j in T.Parallel(hc_mult3):
                mixes[j] = 0
                for i_split in T.serial(n_splits):
                    mixes[j] += gemm_out_mul[i_split, i, j]
                mixes[j] *= rms[0]
            T.copy(mixes, mixes_shared)

        if split_mode != "input" and _is_mix_thread():
            cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
            for j in T.Parallel(hc_mult):
                if save_pre_mix:
                    pre_mix_out[i, j] = (
                        T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
                        + hc_pre_eps
                    )
                post_mix[i, j] = (
                    T.sigmoid(
                        mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
                    )
                    * hc_post_mult_value
                )
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = (
                    mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
                    + hc_base[j * hc_mult + k + hc_mult * 2]
                )

            row_sum = T.alloc_fragment(hc_mult, T.float32)
            col_sum = T.alloc_fragment(hc_mult, T.float32)

            row_max = T.alloc_fragment(hc_mult, T.float32)
            T.reduce_max(cm, row_max, dim=1)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = T.exp(cm[j, k] - row_max[j])
            T.reduce_sum(cm, row_sum, dim=1)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps

            T.reduce_sum(cm, col_sum, dim=0)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)

            for _ in T.serial(sinkhorn_repeat - 1):
                T.reduce_sum(cm, row_sum, dim=1)
                for j, k in T.Parallel(hc_mult, hc_mult):
                    cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)

                T.reduce_sum(cm, col_sum, dim=0)
                for j, k in T.Parallel(hc_mult, hc_mult):
                    cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)

            for j, k in T.Parallel(hc_mult, hc_mult):
                comb_mix[i, j * hc_mult + k] = cm[j, k]
        elif split_mode != "stats" and _is_collapse_thread():
            pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
            for j in T.Parallel(hc_mult):
                if use_pre_mix_in:
                    pre_mix_shared[j] = pre_mix_in[i, j]
                elif save_pre_mix:
                    pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0)
                else:
                    pre_mix_shared[j] = (
                        T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j])
                        + hc_pre_eps
                    )

            # Pass 1: stash unnormalized weighted-sum output in shared memory
            # as bf16 (matches the rounding that RMSNorm would see) while
            # accumulating the per-position squared sum.
            output_shared = T.alloc_shared(hidden_size, T.bfloat16)
            sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32)
            T.clear(sumsq_per_pos)

            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
                xs = T.alloc_shared((hc_mult, hidden_block), T.bfloat16)
                xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
                T.copy(residual[i, 0, i0_h * hidden_block], xs)
                T.copy(xs, xl)

                ol = T.alloc_fragment(hidden_block, T.float32)
                T.clear(ol)
                if write_aux:
                    # Aux consumers (draft models) take the plain stream mean
                    # of the same residual this collapse reads.
                    aux = T.alloc_fragment(hidden_block, T.float32)
                    T.clear(aux)

                for i_hc in T.serial(hc_mult):
                    pre = pre_mix_shared[i_hc]
                    for i1_h in T.Parallel(hidden_block):
                        ol[i1_h] += pre * xl[i_hc, i1_h]
                        if write_aux:
                            aux[i1_h] += xl[i_hc, i1_h]

                if write_aux:
                    for i1_h in T.Parallel(hidden_block):
                        aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16(
                            aux[i1_h] / hc_mult
                        )

                if save_pre_mix:
                    # Keep the BF16 boundary before the delayed input RMSNorm.
                    rounded = T.alloc_fragment(hidden_block, T.bfloat16)
                    T.copy(ol, rounded)
                    for i1_h in T.Parallel(hidden_block):
                        value = T.float32(rounded[i1_h])
                        sumsq_per_pos[i1_h] += value * value
                        output_shared[i0_h * hidden_block + i1_h] = rounded[i1_h]
                else:
                    for i1_h in T.Parallel(hidden_block):
                        sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
                        output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])

            sumsq = T.alloc_fragment(1, T.float32)
            T.reduce_sum(sumsq_per_pos, sumsq, dim=0)
            rsqrt_norm = T.alloc_fragment(1, T.float32)
            rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps)

            # Pass 2: scale by rsqrt * norm_weight and write the result to HBM.
            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
                w_shared = T.alloc_shared(hidden_block, T.bfloat16)
                w_local = T.alloc_fragment(hidden_block, T.float32)
                T.copy(norm_weight[i0_h * hidden_block], w_shared)
                T.copy(w_shared, w_local)

                ol = T.alloc_fragment(hidden_block, T.float32)
                for i1_h in T.Parallel(hidden_block):
                    ol[i1_h] = (
                        output_shared[i0_h * hidden_block + i1_h]
                        * rsqrt_norm[0]
                        * w_local[i1_h]
                    )

                T.copy(ol, layer_input[i, i0_h * hidden_block])

        if ENABLE_PDL:
            T.pdl_trigger()


@tilelang_jit
def mhc_pre_big_fuse_broadcast_with_norm_tilelang(
    gemm_out_mul,
    gemm_out_sqrsum,
    hc_scale,
    hc_base,
    residual,
    residual_out,
    post_mix,
    comb_mix,
    layer_input,
    norm_weight,
    hidden_size: int,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    norm_eps: float,
    n_splits: int = 1,
    hc_mult: int = 4,
    gemm_last_dim: int = -1,
):
    num_tokens = T.dynamic("num_tokens")
    hc_mult3 = hc_mult * (2 + hc_mult)
    if gemm_last_dim < 0:
        gemm_last_dim = hc_mult3
    hidden_block = math.gcd(1024, hidden_size)

    gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32]  # type: ignore[no-redef, valid-type]
    gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32]  # type: ignore[no-redef, valid-type]
    hc_scale: T.Tensor[[3], T.float32]  # type: ignore[no-redef, valid-type]
    hc_base: T.Tensor[[hc_mult3], T.float32]  # type: ignore[no-redef, valid-type]
    residual: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
    residual_out: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
    post_mix: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
    layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
    norm_weight: T.Tensor[[hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]

    with T.Kernel(num_tokens, threads=96) as i:
        rms = T.alloc_fragment(1, T.float32)
        mixes = T.alloc_fragment(hc_mult3, T.float32)
        T.clear(mixes)
        rms[0] = 0

        if ENABLE_PDL:
            T.pdl_sync()

        for i_split in T.serial(n_splits):
            rms[0] += gemm_out_sqrsum[i_split, i]
        # rms scale
        rms[0] = T.rsqrt(rms[0] / hidden_size + rms_eps)
        for j in T.Parallel(hc_mult3):
            mixes[j] = 0
            for i_split in T.serial(n_splits):
                mixes[j] += gemm_out_mul[i_split, i, j]
            mixes[j] *= rms[0]
        mixes_shared = T.alloc_shared(hc_mult3, T.float32)
        T.copy(mixes, mixes_shared)

        if _is_mix_thread():
            cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
            for j in T.Parallel(hc_mult):
                post_mix[i, j] = (
                    T.sigmoid(
                        mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
                    )
                    * hc_post_mult_value
                )
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = (
                    mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
                    + hc_base[j * hc_mult + k + hc_mult * 2]
                )

            row_sum = T.alloc_fragment(hc_mult, T.float32)
            col_sum = T.alloc_fragment(hc_mult, T.float32)

            row_max = T.alloc_fragment(hc_mult, T.float32)
            T.reduce_max(cm, row_max, dim=1)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = T.exp(cm[j, k] - row_max[j])
            T.reduce_sum(cm, row_sum, dim=1)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps

            T.reduce_sum(cm, col_sum, dim=0)
            for j, k in T.Parallel(hc_mult, hc_mult):
                cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)

            for _ in T.serial(sinkhorn_repeat - 1):
                T.reduce_sum(cm, row_sum, dim=1)
                for j, k in T.Parallel(hc_mult, hc_mult):
                    cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)

                T.reduce_sum(cm, col_sum, dim=0)
                for j, k in T.Parallel(hc_mult, hc_mult):
                    cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)

            for j, k in T.Parallel(hc_mult, hc_mult):
                comb_mix[i, j * hc_mult + k] = cm[j, k]
        else:
            pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
            for j in T.Parallel(hc_mult):
                pre_mix_shared[j] = (
                    T.sigmoid(
                        mixes_shared[j] * hc_scale[0] + hc_base[j],
                    )
                    + hc_pre_eps
                )

            output_shared = T.alloc_shared(hidden_size, T.bfloat16)
            sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32)
            T.clear(sumsq_per_pos)

            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
                xs = T.alloc_shared(hidden_block, T.bfloat16)
                xl = T.alloc_fragment(hidden_block, T.float32)
                T.copy(residual[i, i0_h * hidden_block], xs)
                T.copy(xs, xl)

                for i_hc, i1_h in T.Parallel(hc_mult, hidden_block):
                    residual_out[i, i_hc, i0_h * hidden_block + i1_h] = xs[i1_h]

                ol = T.alloc_fragment(hidden_block, T.float32)
                T.clear(ol)

                for i_hc in T.serial(hc_mult):
                    pre = pre_mix_shared[i_hc]
                    for i1_h in T.Parallel(hidden_block):
                        ol[i1_h] += pre * xl[i1_h]

                for i1_h in T.Parallel(hidden_block):
                    sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
                    output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])

            sumsq = T.alloc_fragment(1, T.float32)
            T.reduce_sum(sumsq_per_pos, sumsq, dim=0)
            rsqrt_norm = T.alloc_fragment(1, T.float32)
            rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps)

            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
                w_shared = T.alloc_shared(hidden_block, T.bfloat16)
                w_local = T.alloc_fragment(hidden_block, T.float32)
                T.copy(norm_weight[i0_h * hidden_block], w_shared)
                T.copy(w_shared, w_local)

                ol = T.alloc_fragment(hidden_block, T.float32)
                for i1_h in T.Parallel(hidden_block):
                    ol[i1_h] = (
                        output_shared[i0_h * hidden_block + i1_h]
                        * rsqrt_norm[0]
                        * w_local[i1_h]
                    )

                T.copy(ol, layer_input[i, i0_h * hidden_block])

        if ENABLE_PDL:
            T.pdl_trigger()


@tilelang_jit
def mhc_fused_tilelang(
    comb_mix,
    residual_in,
    post_mix,
    x_in,
    weight_t,
    yp_out,
    rp_out,
    residual_out,
    hc: int,
    hidden: int,
    n_out: int,
    n_thr: int = 256,
    h_blk: int = 256,
    tile_n: int = 1,
    split_k: int = 1,
) -> tilelang.JITKernel:
    """Fused mhc post-mapping + pre-norm GEMM FMA."""
    m = T.dynamic("num_tokens")
    split_k = T.dynamic("split_k")
    h = hidden
    h_blk = math.gcd(hidden, h_blk)
    h_per_split = h // split_k
    n_tiles = n_out // tile_n

    comb_mix: T.Tensor((m, hc, hc), T.float32)  # type: ignore[no-redef, valid-type]
    residual_in: T.Tensor((m, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
    post_mix: T.Tensor((m, hc), T.float32)  # type: ignore[no-redef, valid-type]
    x_in: T.Tensor((m, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
    weight_t: T.Tensor((n_out, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
    yp_out: T.Tensor((split_k, m, n_out), T.float32)  # type: ignore[no-redef, valid-type]
    rp_out: T.Tensor((split_k, m), T.float32)  # type: ignore[no-redef, valid-type]
    residual_out: T.Tensor((m, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]

    h_iters = h_per_split // n_thr
    num_warps = n_thr // 32

    with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
        tid = T.get_thread_binding()
        warp_id = tid // 32
        lane = tid % 32

        s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
        s_post = T.alloc_shared((hc,), T.float32)
        s_comb = T.alloc_shared((hc, hc), T.float32)

        pm = T.alloc_local((hc,), T.float32)
        cm = T.alloc_local((hc, hc), T.float32)
        acc = T.alloc_local((tile_n,), T.float32)
        sqr = T.alloc_local((1,), T.float32)
        new_r = T.alloc_local((hc,), T.float32)

        T.clear(acc)
        T.clear(sqr)
        h_split_start = i_ks * h_per_split

        if ENABLE_PDL:
            T.pdl_sync()

        T.copy(post_mix[i_n, 0], s_post)
        T.copy(comb_mix[i_n, 0, 0], s_comb)

        for j in T.unroll(hc):
            pm[j] = s_post[j]
        for j in T.unroll(hc):
            for k in T.unroll(hc):
                cm[k, j] = s_comb[k, j]

        # Each thread owns h_iters elements of the k-split's h slice.
        for it in T.serial(h_iters):
            h_idx = h_split_start + it * n_thr + tid

            # Compute new residual from layer output and past residual
            for j in T.unroll(hc):
                new_r[j] = pm[j] * x_in[i_n, h_idx]
                for k in T.unroll(hc):
                    new_r[j] += cm[k, j] * residual_in[i_n, k, h_idx]

            # populate residual_out and compute sqr sum
            if i_nt == 0:
                for j in T.unroll(hc):
                    residual_out[i_n, j, h_idx] = new_r[j]
                    sqr[0] += new_r[j] * new_r[j]

            # Per-thread FMA into acc[n]
            for n in T.unroll(tile_n):
                for j in T.unroll(hc):
                    acc[n] += T.cast(weight_t[i_nt * tile_n + n, j, h_idx], T.float32) * new_r[j]

        for n in T.unroll(tile_n):
            acc[n] = T.warp_reduce_sum(acc[n])
        if i_nt == 0:
            sqr[0] = T.warp_reduce_sum(sqr[0])

        # Cross-warp reduce via shared mem
        if lane == 0:
            for n in T.unroll(tile_n):
                s_warp[warp_id, n] = acc[n]
            if i_nt == 0:
                s_warp[warp_id, tile_n] = sqr[0]
        T.sync_threads()

        # Warp 0 does the final cross-warp sum and writes outputs
        if warp_id == 0:
            if lane < tile_n:
                v = T.alloc_var(T.float32, init=0.0)
                for w in T.unroll(num_warps):
                    v += s_warp[w, lane]
                yp_out[i_ks, i_n, i_nt * tile_n + lane] = v

            if i_nt == 0 and lane == 0:
                v2 = T.alloc_var(T.float32, init=0.0)
                for w in T.unroll(num_warps):
                    v2 += s_warp[w, tile_n]
                rp_out[i_ks, i_n] = v2

        if ENABLE_PDL:
            T.pdl_trigger()


@tilelang_jit
def mhc_post_tilelang(
    a,
    b,
    c,
    d,
    x,
    hc: int,
    hidden: int,
    n_thr: int = 128,
    h_blk: int = 1024,
) -> tilelang.JITKernel:
    # rename for shorter code
    n = T.dynamic("num_tokens")
    h = hidden

    h_blk = math.gcd(hidden, h_blk)
    a: T.Tensor((n, hc, hc), T.float32)  # type: ignore[no-redef, valid-type]
    b: T.Tensor((n, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
    c: T.Tensor((n, hc), T.float32)  # type: ignore[no-redef, valid-type]
    d: T.Tensor((n, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
    x: T.Tensor((n, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
    with T.Kernel(n, threads=n_thr) as i_n:
        b_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
        d_shared = T.alloc_shared(h_blk, T.bfloat16)

        x_local = T.alloc_fragment((hc, h_blk), T.float32)
        b_local = T.alloc_fragment((hc, h_blk), T.float32)
        d_local = T.alloc_fragment(h_blk, T.float32)

        a_local = T.alloc_fragment((hc, hc), T.float32)
        c_local = T.alloc_fragment(hc, T.float32)
        if ENABLE_PDL:
            T.pdl_sync()
        T.copy(a[i_n, 0, 0], a_local)
        T.copy(c[i_n, 0], c_local)

        for i0_h in T.Serial(T.ceildiv(h, h_blk)):
            T.copy(b[i_n, 0, i0_h * h_blk], b_shared)
            T.copy(d[i_n, i0_h * h_blk], d_shared)

            T.copy(b_shared, b_local)
            T.copy(d_shared, d_local)
            for i_hco, i1_h in T.Parallel(hc, h_blk):
                x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h]
                for i_hci in T.vectorized(hc):
                    x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]

            T.copy(x_local, x[i_n, 0, i0_h * h_blk])
        if ENABLE_PDL:
            T.pdl_trigger()


@tilelang_jit
def hc_prenorm_gemm_tilelang(
    x,
    fn,
    out,
    sqrsum,
    hidden_size: int,
    hc_mult: int = 4,
    n_out: int = 24,
    n_thr: int = 512,
    tile_n: int = 12,
    n_splits: int = 1,
) -> tilelang.JITKernel:
    num_tokens = T.dynamic("num_tokens")
    hc_hidden_size = hc_mult * hidden_size
    k_per_split = hc_hidden_size // n_splits
    k_iters = k_per_split // n_thr
    n_tiles = T.ceildiv(n_out, tile_n)

    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
    out: T.Tensor((n_splits, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
    sqrsum: T.Tensor((n_splits, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]

    with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as (
        i_n,
        i_t,
        i_s,
    ):
        tid = T.get_thread_binding()
        acc = T.alloc_local((tile_n,), T.float32)
        sqr = T.alloc_local((1,), T.float32)
        T.clear(acc)
        T.clear(sqr)

        if ENABLE_PDL:
            T.pdl_sync()

        for it in T.serial(k_iters):
            i_k = i_s * k_per_split + it * n_thr + tid
            x_val = T.cast(x[i_n, i_k], T.float32)
            for i_o in T.unroll(tile_n):
                out_idx = i_t * tile_n + i_o
                if out_idx < n_out:
                    acc[i_o] += x_val * fn[out_idx, i_k]
            if i_t == 0:
                sqr[0] += x_val * x_val

        for i_o in T.unroll(tile_n):
            acc[i_o] = T.warp_reduce_sum(acc[i_o])
        if i_t == 0:
            sqr[0] = T.warp_reduce_sum(sqr[0])

        lane = tid % 32
        warp_id = tid // 32
        num_warps = n_thr // 32
        warp_acc = T.alloc_shared((num_warps, tile_n), T.float32)
        warp_sqr = T.alloc_shared(num_warps, T.float32)

        if lane == 0:
            for i_o in T.unroll(tile_n):
                warp_acc[warp_id, i_o] = acc[i_o]
            if i_t == 0:
                warp_sqr[warp_id] = sqr[0]
        T.sync_threads()

        if warp_id == 0:
            if lane < tile_n:
                reduced_acc = T.alloc_var(T.float32, init=0.0)
                for i_w in T.unroll(num_warps):
                    reduced_acc += warp_acc[i_w, lane]
                out_idx = i_t * tile_n + lane
                if out_idx < n_out:
                    out[i_s, i_n, out_idx] = reduced_acc
            if lane == 0 and i_t == 0:
                reduced_sqr = T.alloc_var(T.float32, init=0.0)
                for i_w in T.unroll(num_warps):
                    reduced_sqr += warp_sqr[i_w]
                sqrsum[i_s, i_n] = reduced_sqr

        if ENABLE_PDL:
            T.pdl_trigger()


@tilelang_jit
def hc_prenorm_gemm_block_m_tilelang(
    x,
    fn,
    out,
    sqrsum,
    hidden_size: int,
    hc_mult: int = 4,
    n_out: int = 24,
    n_thr: int = 512,
    tile_n: int = 12,
    block_m: int = 2,
) -> tilelang.JITKernel:
    num_tokens = T.dynamic("num_tokens")
    hc_hidden_size = hc_mult * hidden_size
    k_iters = hc_hidden_size // n_thr
    n_tiles = T.ceildiv(n_out, tile_n)
    m_tiles = T.ceildiv(num_tokens, block_m)

    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
    out: T.Tensor((1, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
    sqrsum: T.Tensor((1, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]

    with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t):
        tid = T.get_thread_binding()
        acc = T.alloc_local((block_m, tile_n), T.float32)
        sqr = T.alloc_local((block_m,), T.float32)
        T.clear(acc)
        T.clear(sqr)

        if ENABLE_PDL:
            T.pdl_sync()

        for it in T.serial(k_iters):
            i_k = it * n_thr + tid
            fn_val = T.alloc_local((tile_n,), T.float32)
            for i_o in T.unroll(tile_n):
                out_idx = i_t * tile_n + i_o
                if out_idx < n_out:
                    fn_val[i_o] = fn[out_idx, i_k]
                else:
                    fn_val[i_o] = 0.0
            for i_m in T.unroll(block_m):
                token_idx = i_mt * block_m + i_m
                if token_idx < num_tokens:
                    x_val = T.cast(x[token_idx, i_k], T.float32)
                    for i_o in T.unroll(tile_n):
                        acc[i_m, i_o] += x_val * fn_val[i_o]
                    if i_t == 0:
                        sqr[i_m] += x_val * x_val

        for i_m in T.unroll(block_m):
            for i_o in T.unroll(tile_n):
                acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o])
            if i_t == 0:
                sqr[i_m] = T.warp_reduce_sum(sqr[i_m])

        lane = tid % 32
        warp_id = tid // 32
        num_warps = n_thr // 32
        warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32)
        warp_sqr = T.alloc_shared((num_warps, block_m), T.float32)

        if lane == 0:
            for i_m in T.unroll(block_m):
                for i_o in T.unroll(tile_n):
                    warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o]
                if i_t == 0:
                    warp_sqr[warp_id, i_m] = sqr[i_m]
        T.sync_threads()

        if warp_id == 0:
            reduced_acc = T.alloc_local((block_m,), T.float32)
            reduced_sqr = T.alloc_local((block_m,), T.float32)
            T.clear(reduced_acc)
            T.clear(reduced_sqr)
            for i_m in T.unroll(block_m):
                token_idx = i_mt * block_m + i_m
                if token_idx < num_tokens:
                    if lane < tile_n:
                        for i_w in T.unroll(num_warps):
                            reduced_acc[i_m] += warp_acc[i_w, i_m, lane]
                        out_idx = i_t * tile_n + lane
                        if out_idx < n_out:
                            out[0, token_idx, out_idx] = reduced_acc[i_m]
                    if lane == 0 and i_t == 0:
                        for i_w in T.unroll(num_warps):
                            reduced_sqr[i_m] += warp_sqr[i_w, i_m]
                        sqrsum[0, token_idx] = reduced_sqr[i_m]

        if ENABLE_PDL:
            T.pdl_trigger()


@tilelang_jit
def hc_head_fuse_tilelang(
    residual,
    fn,
    hc_scale,
    hc_base,
    out,
    hidden_size: int,
    rms_eps: float,
    hc_eps: float,
    hc_mult: int = 4,
    n_thr: int = 128,
    h_blk: int = 1024,
):
    """Two-pass fused kernel for hc_head.

    Pass 1: accumulate per-token squared sum and hc_mult dot-products
            (projections onto fn rows) using cross-thread reducers.
    Pass 2: apply sigmoid-gated weighted sum of residual channels to output.

    Avoids materialising mixes / rsqrt / pre tensors to global memory.
    """
    num_tokens = T.dynamic("num_tokens")
    hc_dim = hc_mult * hidden_size
    h_block = math.gcd(h_blk, hidden_size)
    n_h = hidden_size // h_block

    residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]
    fn: T.Tensor[[hc_mult, hc_dim], T.float32]  # type: ignore[no-redef,valid-type]
    hc_scale: T.Tensor[[1], T.float32]  # type: ignore[no-redef,valid-type]
    hc_base: T.Tensor[[hc_mult], T.float32]  # type: ignore[no-redef,valid-type]
    out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]

    with T.Kernel(num_tokens, threads=n_thr) as i:
        if ENABLE_PDL:
            T.pdl_sync()

        # ------------------------------------------------------------------
        # Pass 1 – for each residual channel m_c and h_block:
        #   • accumulate squared sum (for RMS norm denominator)
        #   • accumulate hc_mult dot-products with fn rows
        # ------------------------------------------------------------------
        sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all")
        mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all")
        T.fill(sqrsum_r, 0.0)
        T.fill(mixes_r, 0.0)

        for m_c in T.serial(hc_mult):
            for i_h in T.serial(n_h):
                x_local = T.alloc_fragment(h_block, T.float32)
                T.copy(residual[i, m_c, i_h * h_block], x_local)

                for k in T.Parallel(h_block):
                    sqrsum_r[0] += x_local[k] * x_local[k]

                for m_m in T.unroll(hc_mult):
                    fn_local = T.alloc_fragment(h_block, T.float32)
                    T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local)
                    for k in T.Parallel(h_block):
                        mixes_r[m_m] += x_local[k] * fn_local[k]

        T.finalize_reducer(sqrsum_r)
        T.finalize_reducer(mixes_r)

        # ------------------------------------------------------------------
        # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps
        # ------------------------------------------------------------------
        pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
        rsqrt_val = T.alloc_fragment(1, T.float32)
        rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps)
        for m in T.Parallel(hc_mult):
            pre_mix_shared[m] = (
                T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps
            )

        # ------------------------------------------------------------------
        # Pass 2 – apply_mix: pipelined weighted sum over residual channels
        # ------------------------------------------------------------------
        for i0_h in T.Pipelined(n_h, num_stages=2):
            xs = T.alloc_shared((hc_mult, h_block), T.bfloat16)
            xl = T.alloc_fragment((hc_mult, h_block), T.float32)
            T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True)
            T.copy(xs, xl)

            ol = T.alloc_fragment(h_block, T.float32)
            T.clear(ol)
            for i_hc in T.serial(hc_mult):
                pre = pre_mix_shared[i_hc]
                for i1_h in T.Parallel(h_block):
                    ol[i1_h] += pre * xl[i_hc, i1_h]

            T.copy(ol, out[i, i0_h * h_block], disable_tma=True)

        if ENABLE_PDL:
            T.pdl_trigger()


class HcPrenormGemmTileLangKernel(
    VllmTileLangJitKernel["HcPrenormGemmTileLangKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        hidden_size: int
        hc_mult: int
        n_out: int
        n_thr: int
        tile_n: int
        n_splits: int
        use_block_m: bool
        block_m: int

    @staticmethod
    def kernel(compile_key: CompileKey) -> Any:
        if compile_key.use_block_m:
            return hc_prenorm_gemm_block_m_tilelang
        return hc_prenorm_gemm_tilelang

    def dispatch(  # type: ignore[override]
        self,
        *,
        num_tokens: int,
        hc_hidden_size: int,
        n_thr: int = 512,
        tile_n: int = 12,
        n_splits: int = 1,
        **compile_key_fields: int,
    ) -> CompileKey:
        use_default_config = n_splits == 1 and tile_n == 12 and n_thr == 512
        use_block_m = use_default_config and num_tokens >= 1024
        use_small_tile = (
            use_default_config
            and not use_block_m
            and num_tokens < 128
            and hc_hidden_size % 1024 == 0
        )
        effective_n_thr = 1024 if use_small_tile else n_thr
        effective_tile_n = 4 if use_small_tile else tile_n
        return self.CompileKey(
            **compile_key_fields,
            n_thr=effective_n_thr,
            tile_n=effective_tile_n,
            n_splits=n_splits,
            use_block_m=use_block_m,
            block_m=2 if use_block_m else 1,
        )

    def get_warmup_keys(
        self,
        vllm_config: Any,
        *,
        hidden_size: int,
        hc_mult: int,
        n_out: int,
    ) -> list[CompileKey]:
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        if max_tokens <= 0:
            return []
        return self._trace_dispatch(self.dispatch)(
            num_tokens=WarmupIntRange(
                1,
                max_tokens + 1,
                advance=lambda value: (
                    128 if value < 128 else 1024 if value < 1024 else max_tokens + 1
                ),
            ),
            hc_hidden_size=hidden_size * hc_mult,
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            n_out=n_out,
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        num_tokens = 128
        if compile_key.use_block_m:
            num_tokens = 1024
        elif compile_key.n_thr == 1024:
            num_tokens = 1
        hc_hidden_size = compile_key.hidden_size * compile_key.hc_mult
        return dict(
            x=make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hc_hidden_size),
            fn=make_tilelang_warmup_tensor(
                torch.float32, compile_key.n_out, hc_hidden_size
            ),
            out=make_tilelang_warmup_tensor(
                torch.float32,
                compile_key.n_splits,
                num_tokens,
                compile_key.n_out,
            ),
            sqrsum=make_tilelang_warmup_tensor(
                torch.float32, compile_key.n_splits, num_tokens
            ),
            hidden_size=compile_key.hidden_size,
            hc_mult=compile_key.hc_mult,
            tile_n=compile_key.tile_n,
            n_thr=compile_key.n_thr,
            n_splits=compile_key.n_splits,
        )

    @kernel_launcher
    def __call__(
        self,
        x: Any,
        fn: Any,
        out: Any,
        sqrsum: Any,
        hidden_size: int,
        hc_mult: int,
        tile_n: int = 12,
        n_thr: int = 512,
        n_splits: int = 1,
    ) -> TileLangLaunchSpec:
        assert out.shape[0] == n_splits
        assert sqrsum.shape[0] == n_splits
        assert x.shape[1] == hc_mult * hidden_size
        assert x.shape[1] % n_splits == 0
        assert (x.shape[1] // n_splits) % n_thr == 0
        compile_key = self.dispatch(
            num_tokens=x.shape[0],
            hc_hidden_size=x.shape[1],
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            n_out=fn.shape[0],
            n_thr=n_thr,
            tile_n=tile_n,
            n_splits=n_splits,
        )
        kernel_arg = (
            compile_key.block_m if compile_key.use_block_m else compile_key.n_splits
        )
        return (compile_key,), (
            x,
            fn,
            out,
            sqrsum,
            hidden_size,
            hc_mult,
            fn.shape[0],
            compile_key.n_thr,
            compile_key.tile_n,
            kernel_arg,
        )


class MhcPreBigFuseTileLangKernel(
    VllmTileLangJitKernel["MhcPreBigFuseTileLangKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        hidden_size: int
        hc_mult: int
        n_splits: int
        use_norm_weight: bool
        is_broadcast: bool
        rms_eps: float
        hc_pre_eps: float
        hc_sinkhorn_eps: float
        hc_post_mult_value: float
        sinkhorn_repeat: int
        norm_eps: float

    @staticmethod
    def kernel(compile_key: CompileKey) -> Any:
        if compile_key.is_broadcast:
            return mhc_pre_big_fuse_broadcast_with_norm_tilelang
        if compile_key.use_norm_weight:
            return mhc_pre_big_fuse_with_norm_tilelang
        return mhc_pre_big_fuse_tilelang

    @staticmethod
    def _kernel_args(
        compile_key: CompileKey,
        *,
        gemm_out_mul: Any,
        gemm_out_sqrsum: Any,
        hc_scale: Any,
        hc_base: Any,
        residual: Any,
        residual_out: Any | None,
        post_mix: Any,
        comb_mix: Any,
        layer_input: Any,
        norm_weight: Any | None,
    ) -> tuple[Any, ...]:
        common_args = (
            gemm_out_mul,
            gemm_out_sqrsum,
            hc_scale,
            hc_base,
            residual,
        )
        output_args = (post_mix, comb_mix, layer_input)
        eps_args = (
            compile_key.rms_eps,
            compile_key.hc_pre_eps,
            compile_key.hc_sinkhorn_eps,
            compile_key.hc_post_mult_value,
            compile_key.sinkhorn_repeat,
        )
        middle_args: tuple[Any, ...]
        if compile_key.is_broadcast:
            assert residual_out is not None
            assert norm_weight is not None
            middle_args = (residual_out, *output_args, norm_weight)
        elif compile_key.use_norm_weight:
            assert norm_weight is not None
            # layer_input stands in for the aux buffer on both epilogues: same
            # shape and dtype, and write_aux is off here so it is never written.
            middle_args = (
                *output_args,
                norm_weight,
                post_mix,
                post_mix,
                layer_input,
            )
        else:
            middle_args = (*output_args, post_mix, post_mix, layer_input)
        norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else ()
        return (
            *common_args,
            *middle_args,
            compile_key.hidden_size,
            *eps_args,
            *norm_eps_args,
            compile_key.n_splits,
            compile_key.hc_mult,
        )

    def dispatch(  # type: ignore[override]
        self,
        *,
        hidden_size: int,
        hc_mult: int,
        n_splits: int,
        is_broadcast: bool,
        use_norm_weight: bool,
        use_pre_gemm_splits: bool = False,
        pre_gemm_k: int = 0,
        sinkhorn_repeat: int,
        norm_eps: float,
        broadcast_norm_eps: float = 0.0,
        num_tokens: int = 1,
        use_fused_tilelang: bool = False,
        **compile_key_fields: float,
    ) -> CompileKey:
        pre_gemm_n_splits = (
            compute_num_split(64, pre_gemm_k, (num_tokens + 63) // 64)
            if use_pre_gemm_splits
            else n_splits
        )
        # The epilogue reduces over whatever the fused kernel split the GEMM
        # into, so both have to read the same config.
        fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
        actual_n_splits = (
            fused_config[1]
            if use_fused_tilelang and fused_config is not None
            else pre_gemm_n_splits
        )
        actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps
        actual_use_norm_weight = use_norm_weight or is_broadcast
        return self.CompileKey(
            **compile_key_fields,
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            n_splits=actual_n_splits,
            use_norm_weight=actual_use_norm_weight,
            is_broadcast=is_broadcast,
            sinkhorn_repeat=sinkhorn_repeat,
            norm_eps=actual_norm_eps if actual_use_norm_weight else 0.0,
        )

    def get_warmup_keys(
        self,
        vllm_config: Any,
        *,
        hidden_size: int,
        hc_mult: int,
        use_norm_weight: bool,
        rms_eps: float,
        hc_pre_eps: float,
        hc_sinkhorn_eps: float,
        hc_post_mult_value: float,
        sinkhorn_repeat: int,
        norm_eps: float | tuple[float, ...],
        broadcast_norm_eps: float | tuple[float, ...] | None = None,
        include_pre_gemm_splits: bool = False,
        include_broadcast_splits: bool = False,
    ) -> list[CompileKey]:
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        if max_tokens <= 0:
            return []

        warmup_cases = zip_inputs(
            dict(
                is_broadcast=False,
                use_pre_gemm_splits=False,
                pre_gemm_k=0,
                use_fused_tilelang=False,
            ),
            dict(
                is_broadcast=False,
                use_pre_gemm_splits=include_pre_gemm_splits,
                pre_gemm_k=hc_mult * hidden_size,
                use_fused_tilelang=False,
            ),
            dict(
                is_broadcast=include_broadcast_splits,
                use_pre_gemm_splits=include_broadcast_splits,
                pre_gemm_k=hidden_size,
                use_fused_tilelang=False,
            ),
            dict(
                is_broadcast=False,
                use_pre_gemm_splits=False,
                pre_gemm_k=0,
                use_fused_tilelang=True,
            ),
        )
        return self._trace_dispatch(self.dispatch)(
            warmup_cases,
            num_tokens=WarmupIntRange(
                1,
                max_tokens + 1,
                advance=lambda value: 8 if value < 8 else cdiv(value, 64) * 64 + 1,
            ),
            n_splits=1,
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            use_norm_weight=use_norm_weight,
            rms_eps=rms_eps,
            hc_pre_eps=hc_pre_eps,
            hc_sinkhorn_eps=hc_sinkhorn_eps,
            hc_post_mult_value=hc_post_mult_value,
            sinkhorn_repeat=sinkhorn_repeat,
            norm_eps=norm_eps,
            broadcast_norm_eps=broadcast_norm_eps
            if broadcast_norm_eps is not None
            else norm_eps,
            _when=lambda *, use_fused_tilelang, num_tokens: (
                not use_fused_tilelang or num_tokens <= 16
            ),
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        hidden_size = compile_key.hidden_size
        hc_mult = compile_key.hc_mult
        hc_mult3 = hc_mult * (2 + hc_mult)
        num_tokens = 1

        gemm_out_mul = make_tilelang_warmup_tensor(
            torch.float32, compile_key.n_splits, num_tokens, hc_mult3
        )
        gemm_out_sqrsum = make_tilelang_warmup_tensor(
            torch.float32, compile_key.n_splits, num_tokens
        )
        hc_scale = make_tilelang_warmup_tensor(torch.float32, 3)
        hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult3)
        post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
        comb_mix = make_tilelang_warmup_tensor(
            torch.float32, num_tokens, hc_mult * hc_mult
        )
        layer_input = make_tilelang_warmup_tensor(
            torch.bfloat16, num_tokens, hidden_size
        )
        residual_shape = (
            (num_tokens, hidden_size)
            if compile_key.is_broadcast
            else (num_tokens, hc_mult, hidden_size)
        )
        residual = make_tilelang_warmup_tensor(torch.bfloat16, *residual_shape)
        residual_out = (
            make_tilelang_warmup_tensor(
                torch.bfloat16, num_tokens, hc_mult, hidden_size
            )
            if compile_key.is_broadcast
            else None
        )
        norm_weight = (
            make_tilelang_warmup_tensor(torch.bfloat16, hidden_size)
            if compile_key.is_broadcast or compile_key.use_norm_weight
            else None
        )
        return dict(
            gemm_out_mul=gemm_out_mul,
            gemm_out_sqrsum=gemm_out_sqrsum,
            hc_scale=hc_scale,
            hc_base=hc_base,
            residual=residual,
            post_mix=post_mix,
            comb_mix=comb_mix,
            layer_input=layer_input,
            rms_eps=compile_key.rms_eps,
            hc_pre_eps=compile_key.hc_pre_eps,
            hc_sinkhorn_eps=compile_key.hc_sinkhorn_eps,
            hc_post_mult_value=compile_key.hc_post_mult_value,
            sinkhorn_repeat=compile_key.sinkhorn_repeat,
            residual_out=residual_out,
            norm_weight=norm_weight,
            norm_eps=compile_key.norm_eps,
        )

    @kernel_launcher
    def __call__(
        self,
        gemm_out_mul: Any,
        gemm_out_sqrsum: Any,
        hc_scale: Any,
        hc_base: Any,
        residual: Any,
        post_mix: Any,
        comb_mix: Any,
        layer_input: Any,
        rms_eps: float,
        hc_pre_eps: float,
        hc_sinkhorn_eps: float,
        hc_post_mult_value: float,
        sinkhorn_repeat: int,
        *,
        residual_out: Any | None = None,
        norm_weight: Any | None = None,
        norm_eps: float = 0.0,
    ) -> TileLangLaunchSpec:
        is_broadcast = residual_out is not None
        use_norm_weight = norm_weight is not None
        n_splits = gemm_out_mul.shape[0]
        if is_broadcast:
            assert residual_out is not None
            hidden_size = residual.shape[-1]
            hc_mult = residual_out.shape[-2]
        else:
            hc_mult = residual.shape[-2]
            hidden_size = residual.shape[-1]
        compile_key = self.dispatch(
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            n_splits=n_splits,
            is_broadcast=is_broadcast,
            use_norm_weight=use_norm_weight,
            rms_eps=rms_eps,
            hc_pre_eps=hc_pre_eps,
            hc_sinkhorn_eps=hc_sinkhorn_eps,
            hc_post_mult_value=hc_post_mult_value,
            sinkhorn_repeat=sinkhorn_repeat,
            norm_eps=norm_eps,
            broadcast_norm_eps=norm_eps,
        )
        return (
            (compile_key,),
            self._kernel_args(
                compile_key,
                gemm_out_mul=gemm_out_mul,
                gemm_out_sqrsum=gemm_out_sqrsum,
                hc_scale=hc_scale,
                hc_base=hc_base,
                residual=residual,
                residual_out=residual_out,
                post_mix=post_mix,
                comb_mix=comb_mix,
                layer_input=layer_input,
                norm_weight=norm_weight,
            ),
        )


class MhcPostTileLangKernel(VllmTileLangJitKernel["MhcPostTileLangKernel.CompileKey"]):
    @dataclass(frozen=True)
    class CompileKey:
        hidden_size: int
        hc_mult: int

    @staticmethod
    def kernel() -> Any:
        return mhc_post_tilelang

    def dispatch(  # type: ignore[override]
        self,
        **compile_key_fields: int,
    ) -> CompileKey:
        return self.CompileKey(**compile_key_fields)

    def get_warmup_keys(
        self,
        *,
        hidden_size: int,
        hc_mult: int,
    ) -> list[CompileKey]:
        return self._trace_dispatch(self.dispatch)(
            hidden_size=hidden_size,
            hc_mult=hc_mult,
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        num_tokens = 1
        hidden_size = compile_key.hidden_size
        hc_mult = compile_key.hc_mult
        comb_mix = make_tilelang_warmup_tensor(
            torch.float32, num_tokens, hc_mult, hc_mult
        )
        residual = make_tilelang_warmup_tensor(
            torch.bfloat16, num_tokens, hc_mult, hidden_size
        )
        post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
        layer_input = make_tilelang_warmup_tensor(
            torch.bfloat16, num_tokens, hidden_size
        )
        out = make_tilelang_warmup_tensor(
            torch.bfloat16, num_tokens, hc_mult, hidden_size
        )
        return dict(
            comb_mix=comb_mix,
            residual=residual,
            post_mix=post_mix,
            layer_input=layer_input,
            out=out,
            hc_mult=hc_mult,
            hidden_size=hidden_size,
        )

    @kernel_launcher
    def __call__(
        self,
        comb_mix: Any,
        residual: Any,
        post_mix: Any,
        layer_input: Any,
        out: Any,
        hc_mult: int,
        hidden_size: int,
    ) -> TileLangLaunchSpec:
        return (), (
            comb_mix,
            residual,
            post_mix,
            layer_input,
            out,
            hc_mult,
            hidden_size,
        )


class MhcFusedTileLangKernel(
    VllmTileLangJitKernel["MhcFusedTileLangKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        hidden_size: int
        hc_mult: int
        n_splits: int
        tile_n: int
        n_thr: int

    @staticmethod
    def kernel() -> Any:
        return mhc_fused_tilelang

    def dispatch(  # type: ignore[override]
        self,
        *,
        num_tokens: int,
        hidden_size: int,
        hc_mult: int,
    ) -> CompileKey:
        # Subscripts rather than an unpack: the warmup tracer parses this body
        # and allows only plain assignments before the return.
        config = require_fused_post_pre_config(num_tokens, hidden_size, hc_mult)
        return self.CompileKey(
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            n_splits=config[1],
            tile_n=config[0],
            n_thr=config[2],
        )

    def get_warmup_keys(
        self,
        vllm_config: Any,
        *,
        hidden_size: int,
        hc_mult: int,
    ) -> list[CompileKey]:
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        return self._trace_dispatch(self.dispatch)(
            num_tokens=WarmupIntRange(1, max_tokens + 1),
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            _when=lambda *, num_tokens: num_tokens <= _FUSED_POST_PRE_MAX_TOKENS,
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        num_tokens = 1 if compile_key.tile_n == 2 else 16
        hidden_size = compile_key.hidden_size
        hc_mult = compile_key.hc_mult
        hc_mult3 = hc_mult * (2 + hc_mult)
        comb_mix = make_tilelang_warmup_tensor(
            torch.float32, num_tokens, hc_mult, hc_mult
        )
        residual_in = make_tilelang_warmup_tensor(
            torch.bfloat16, num_tokens, hc_mult, hidden_size
        )
        post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult)
        x_in = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size)
        weight_t = make_tilelang_warmup_tensor(
            torch.bfloat16, hc_mult3, hc_mult, hidden_size
        )
        return dict(
            comb_mix=comb_mix,
            residual_in=residual_in,
            post_mix=post_mix,
            x_in=x_in,
            weight_t=weight_t,
            hc_mult=hc_mult,
            hidden_size=hidden_size,
            hc_mult3=hc_mult3,
        )

    @kernel_launcher
    def __call__(
        self,
        comb_mix: Any,
        residual_in: Any,
        post_mix: Any,
        x_in: Any,
        weight_t: Any,
        hc_mult: int,
        hidden_size: int,
        hc_mult3: int,
    ) -> TileLangLaunchSpec:
        num_tokens = residual_in.shape[0]
        tile_n, n_splits, n_thr = require_fused_post_pre_config(
            num_tokens, hidden_size, hc_mult
        )
        yp_out = residual_in.new_empty(
            (n_splits, num_tokens, hc_mult3), dtype=torch.float32
        )
        rp_out = residual_in.new_empty((n_splits, num_tokens), dtype=torch.float32)
        residual_out = residual_in.new_empty(residual_in.shape)
        return (
            (),
            (
                comb_mix,
                residual_in,
                post_mix,
                x_in,
                weight_t,
                yp_out,
                rp_out,
                residual_out,
                hc_mult,
                hidden_size,
                hc_mult3,
            ),
            dict(n_thr=n_thr, tile_n=tile_n),
            (yp_out, rp_out, residual_out),
        )


class HcHeadFusedTileLangKernel(
    VllmTileLangJitKernel["HcHeadFusedTileLangKernel.CompileKey"]
):
    @dataclass(frozen=True)
    class CompileKey:
        hidden_size: int
        hc_mult: int
        rms_eps: float
        hc_eps: float

    @staticmethod
    def kernel() -> Any:
        return hc_head_fuse_tilelang

    def dispatch(  # type: ignore[override]
        self,
        *,
        hidden_size: int,
        hc_mult: int,
        rms_eps: float,
        hc_eps: float,
    ) -> CompileKey:
        return self.CompileKey(
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            rms_eps=rms_eps,
            hc_eps=hc_eps,
        )

    def get_warmup_keys(
        self,
        *,
        hidden_size: int,
        hc_mult: int,
        rms_eps: float,
        hc_eps: float,
    ) -> list[CompileKey]:
        return self._trace_dispatch(self.dispatch)(
            hidden_size=hidden_size,
            hc_mult=hc_mult,
            rms_eps=rms_eps,
            hc_eps=hc_eps,
        )

    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
        num_tokens = 1
        hidden_size = compile_key.hidden_size
        hc_mult = compile_key.hc_mult
        residual = make_tilelang_warmup_tensor(
            torch.bfloat16, num_tokens, hc_mult, hidden_size
        )
        fn = make_tilelang_warmup_tensor(torch.float32, hc_mult, hc_mult * hidden_size)
        hc_scale = make_tilelang_warmup_tensor(torch.float32, 1)
        hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult)
        out = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size)
        return dict(
            residual=residual,
            fn=fn,
            hc_scale=hc_scale,
            hc_base=hc_base,
            out=out,
            hidden_size=hidden_size,
            rms_eps=compile_key.rms_eps,
            hc_eps=compile_key.hc_eps,
            hc_mult=hc_mult,
        )

    @kernel_launcher
    def __call__(
        self,
        residual: Any,
        fn: Any,
        hc_scale: Any,
        hc_base: Any,
        out: Any,
        hidden_size: int,
        rms_eps: float,
        hc_eps: float,
        hc_mult: int,
    ) -> TileLangLaunchSpec:
        return (), (
            residual,
            fn,
            hc_scale,
            hc_base,
            out,
            hidden_size,
            rms_eps,
            hc_eps,
            hc_mult,
        )


_HC_PRENORM_GEMM_TILELANG_KERNEL = HcPrenormGemmTileLangKernel()
_MHC_PRE_BIG_FUSE_TILELANG_KERNEL = MhcPreBigFuseTileLangKernel()
_MHC_POST_TILELANG_KERNEL = MhcPostTileLangKernel()
_MHC_FUSED_TILELANG_KERNEL = MhcFusedTileLangKernel()
_HC_HEAD_FUSED_TILELANG_KERNEL = HcHeadFusedTileLangKernel()