inference/kernel.py

"""Generated TileLang kernel export. Do not edit; regenerate it."""

import linecache
from types import SimpleNamespace

EXPORT_FORMAT_VERSION = 3
SOURCE_SHA256 = {'kernel_tilelang_state.py': 'a3597160cb8d7317151200d05b241af7edcad2928c3731687764d914c28bf21a', 'kernel_tilelang_decode.py': 'b582151a6d2bd345cc5f87415543370e027c8f943a4dfe57a18300f08560bd73'}

def _build_state_namespace():
    source = '# pyright: reportInvalidTypeForm=false\nfrom functools import lru_cache\nfrom typing import Any\n\nimport torch\n\nFLOAT16 = torch.float16  # type: ignore[attr-defined]\nBFLOAT16 = torch.bfloat16  # type: ignore[attr-defined]\nFLOAT32 = torch.float32  # type: ignore[attr-defined]\nTORCH_WHERE = torch.where  # type: ignore[attr-defined]\nTORCH_ZEROS_LIKE = torch.zeros_like  # type: ignore[attr-defined]\nTORCH_STACK = torch.stack  # type: ignore[attr-defined]\nCUDA_GET_DEVICE_CAPABILITY = torch.cuda.get_device_capability\nIS_GRAD_ENABLED = torch.is_grad_enabled  # type: ignore[attr-defined]\n\n\ndef cuda_arch_key(device: Any | None = None) -> str:\n    major, minor = CUDA_GET_DEVICE_CAPABILITY(device)\n    return f"sm_{major}{minor}"\n\n\nEXACT_FUSED_STATE_PROJECTION_CAPABILITIES = frozenset({(8, 9)})\n\n\ndef exact_fused_state_projection_supported(device: Any | None = None) -> bool:\n    """Return whether padded narrow projection is bit-exact on this GPU."""\n    return (\n        tuple(CUDA_GET_DEVICE_CAPABILITY(device))\n        in EXACT_FUSED_STATE_PROJECTION_CAPABILITIES\n    )\n\n\ndef _require_exact_fused_state_projection(device: Any | None = None) -> None:\n    if not exact_fused_state_projection_supported(device):\n        capability = CUDA_GET_DEVICE_CAPABILITY(device)\n        raise RuntimeError(\n            "Exact fused state projection is not validated for CUDA capability "\n            f"{capability}; use tilelang_state_update for exact fallback"\n        )\n\n\ndef build_state_program(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n):\n    """Build exact pointwise RWKV state finalization as a TileLang PrimFunc."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        decay: T.Tensor((batch_size, num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        anti_update: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        value_key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        next_state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=256) as (batch, head):\n            # Preserve eager PyTorch\'s FP32 rounding boundaries exactly:\n            # multiply -> store, add anti update -> store, add value/key -> store.\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = state[\n                    batch, head, row, column\n                ] * T.cast(decay[batch, head, column], "float32")\n            T.sync_threads()\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = (\n                    next_state[batch, head, row, column]\n                    + anti_update[batch, head, row, column]\n                )\n            T.sync_threads()\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = (\n                    next_state[batch, head, row, column]\n                    + value_key[batch, head, row, column]\n                )\n\n    return kernel\n\n\ndef build_state_backward_program(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n):\n    """Build first-order gradients for exact pointwise state finalization."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        grad_output: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        decay: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        grad_state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=256) as (batch, head):\n            for row, column in T.Parallel(head_size, head_size):\n                grad_state[batch, head, row, column] = grad_output[\n                    batch, head, row, column\n                ] * T.cast(decay[batch, head, column], "float32")\n\n    return kernel\n\n\ndef build_low_precision_state_program(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n):\n    """Build FP32-compute finalization with FP16 or BF16 storage."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), input_dtype\n        ),\n        decay: T.Tensor((batch_size, num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        anti_update: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        value_key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        next_state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), input_dtype\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=256) as (batch, head):\n            workspace = T.alloc_shared((head_size, head_size), "float32")\n            for row, column in T.Parallel(head_size, head_size):\n                workspace[row, column] = T.cast(\n                    state[batch, head, row, column], "float32"\n                ) * T.cast(decay[batch, head, column], "float32")\n            T.sync_threads()\n            for row, column in T.Parallel(head_size, head_size):\n                workspace[row, column] = (\n                    workspace[row, column]\n                    + anti_update[batch, head, row, column]\n                )\n            T.sync_threads()\n            for row, column in T.Parallel(head_size, head_size):\n                workspace[row, column] = (\n                    workspace[row, column]\n                    + value_key[batch, head, row, column]\n                )\n            T.sync_threads()\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = T.cast(\n                    workspace[row, column], input_dtype\n                )\n\n    return kernel\n\n\ndef build_state_projection_program(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n):\n    """Build padded Tensor Core state-by-receptance projection."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    padded_columns = 16\n\n    @T.prim_func\n    def kernel(\n        state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), input_dtype\n        ),\n        receptance: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        mixed: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=128) as (batch, head):\n            state_shared = T.alloc_shared(\n                (head_size, head_size), input_dtype\n            )\n            receptance_shared = T.alloc_shared(\n                (head_size, padded_columns), input_dtype\n            )\n            mixed_fragment = T.alloc_fragment(\n                (head_size, padded_columns), "float32"\n            )\n            mixed_shared = T.alloc_shared(\n                (head_size, padded_columns), input_dtype\n            )\n            T.copy(state[batch, head, :, :], state_shared)\n            for row, column in T.Parallel(head_size, padded_columns):\n                receptance_shared[row, column] = receptance[batch, head, row]\n            T.clear(mixed_fragment)\n            T.gemm(state_shared, receptance_shared, mixed_fragment)\n            T.copy(mixed_fragment, mixed_shared)\n            T.sync_threads()\n            for row in T.Parallel(head_size):\n                mixed[batch, head, row] = mixed_shared[row, 0]\n\n    return kernel\n\n\ndef build_fused_state_update_program(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n):\n    """Build exact finalization plus padded Tensor Core projection."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    padded_columns = 16\n\n    @T.prim_func\n    def kernel(\n        state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        decay: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        anti_update: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        value_key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        receptance: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        next_state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        mixed: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=256) as (batch, head):\n            state_shared = T.alloc_shared(\n                (head_size, head_size), input_dtype\n            )\n            receptance_shared = T.alloc_shared(\n                (head_size, padded_columns), input_dtype\n            )\n            mixed_fragment = T.alloc_fragment(\n                (head_size, padded_columns), "float32"\n            )\n            mixed_shared = T.alloc_shared(\n                (head_size, padded_columns), input_dtype\n            )\n\n            # Preserve eager FP32 stores between multiply and both additions.\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = state[\n                    batch, head, row, column\n                ] * T.cast(decay[batch, head, column], "float32")\n            T.sync_threads()\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = (\n                    next_state[batch, head, row, column]\n                    + anti_update[batch, head, row, column]\n                )\n            T.sync_threads()\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = (\n                    next_state[batch, head, row, column]\n                    + value_key[batch, head, row, column]\n                )\n            T.sync_threads()\n\n            for row, column in T.Parallel(head_size, head_size):\n                state_shared[row, column] = T.cast(\n                    next_state[batch, head, row, column], input_dtype\n                )\n            for row, column in T.Parallel(head_size, padded_columns):\n                receptance_shared[row, column] = receptance[batch, head, row]\n            T.clear(mixed_fragment)\n            T.gemm(state_shared, receptance_shared, mixed_fragment)\n            T.copy(mixed_fragment, mixed_shared)\n            T.sync_threads()\n            for row in T.Parallel(head_size):\n                mixed[batch, head, row] = mixed_shared[row, 0]\n\n    return kernel\n\n\n\n\n\n\n\n\ndef build_x_mix_program(\n    batch_size: int,\n    hidden_size: int,\n    input_dtype: str,\n):\n    """Build exact six-way decode x-mix without temporary stacking."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    block_size = 256\n\n    @T.prim_func\n    def kernel(\n        x: T.Tensor((batch_size, hidden_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        previous: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, hidden_size), input_dtype\n        ),\n        mix_r: T.Tensor((hidden_size,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_w: T.Tensor((hidden_size,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_k: T.Tensor((hidden_size,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_v: T.Tensor((hidden_size,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_a: T.Tensor((hidden_size,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_g: T.Tensor((hidden_size,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mixed: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, 6, hidden_size), input_dtype\n        ),\n    ):\n        with T.Kernel(\n            T.ceildiv(hidden_size, block_size), batch_size, threads=block_size\n        ) as (block, batch):\n            delta_shared = T.alloc_shared((block_size,), input_dtype)\n            products_shared = T.alloc_shared((6, block_size), input_dtype)\n\n            for lane in T.Parallel(block_size):\n                column = block * block_size + lane\n                if column < hidden_size:\n                    delta_shared[lane] = T.cast(\n                        previous[batch, column] - x[batch, column], input_dtype\n                    )\n            T.sync_threads()\n\n            for lane in T.Parallel(block_size):\n                column = block * block_size + lane\n                if column < hidden_size:\n                    products_shared[0, lane] = T.cast(\n                        delta_shared[lane] * mix_r[column], input_dtype\n                    )\n                    products_shared[1, lane] = T.cast(\n                        delta_shared[lane] * mix_w[column], input_dtype\n                    )\n                    products_shared[2, lane] = T.cast(\n                        delta_shared[lane] * mix_k[column], input_dtype\n                    )\n                    products_shared[3, lane] = T.cast(\n                        delta_shared[lane] * mix_v[column], input_dtype\n                    )\n                    products_shared[4, lane] = T.cast(\n                        delta_shared[lane] * mix_a[column], input_dtype\n                    )\n                    products_shared[5, lane] = T.cast(\n                        delta_shared[lane] * mix_g[column], input_dtype\n                    )\n            T.sync_threads()\n\n            for lane in T.Parallel(block_size):\n                column = block * block_size + lane\n                if column < hidden_size:\n                    for mix_index in T.Serial(6):\n                        mixed[batch, mix_index, column] = T.cast(\n                            x[batch, column]\n                            + products_shared[mix_index, lane],\n                            input_dtype,\n                        )\n\n    return kernel\n\n\ndef build_post_state_program(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n):\n    """Build exact fused RKV correction and gating after native GroupNorm."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        normalized: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        receptance: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        value: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        r_k: T.Tensor((num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        gate: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n        output: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size), input_dtype\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=32) as (batch, head):\n            product_stage = T.alloc_shared((head_size,), input_dtype)\n            correction = T.alloc_shared((head_size,), input_dtype)\n            corrected = T.alloc_shared((head_size,), input_dtype)\n            rkv_sum = T.alloc_shared((1,), input_dtype)\n\n            for column in T.Parallel(head_size):\n                product_stage[column] = T.cast(\n                    receptance[batch, head, column]\n                    * key[batch, head, column],\n                    input_dtype,\n                )\n            T.sync_threads()\n\n            for column in T.Parallel(head_size):\n                product_stage[column] = T.cast(\n                    product_stage[column] * r_k[head, column], input_dtype\n                )\n            T.sync_threads()\n\n            for worker in T.Parallel(1):\n                total = T.alloc_local((1,), "float32")\n                total[0] = 0.0\n                for column in T.serial(head_size):\n                    total[0] += T.cast(product_stage[column], "float32")\n                rkv_sum[0] = T.cast(total[0], input_dtype)\n            T.sync_threads()\n\n            for column in T.Parallel(head_size):\n                correction[column] = T.cast(\n                    rkv_sum[0] * value[batch, head, column], input_dtype\n                )\n            T.sync_threads()\n\n            for column in T.Parallel(head_size):\n                corrected[column] = T.cast(\n                    normalized[batch, head, column] + correction[column],\n                    input_dtype,\n                )\n            T.sync_threads()\n\n            for column in T.Parallel(head_size):\n                output[batch, head, column] = T.cast(\n                    corrected[column] * gate[batch, head, column], input_dtype\n                )\n\n    return kernel\n\n\ndef _dtype_name(dtype) -> str:\n    if dtype == FLOAT16:\n        return "float16"\n    if dtype == BFLOAT16:\n        return "bfloat16"\n    if dtype == FLOAT32:\n        return "float32"\n    raise TypeError(f"TileLang RWKV kernel does not support dtype {dtype}")\n\n\n@lru_cache(maxsize=32)\ndef _compiled_kernel(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    program = build_state_program(\n        batch_size,\n        num_heads,\n        head_size,\n        input_dtype,\n    )\n    return tilelang.compile(\n        program,\n        out_idx=-1,\n        execution_backend="auto",\n    )\n\n\n@lru_cache(maxsize=32)\ndef _compiled_backward_kernel(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        build_state_backward_program(\n            batch_size, num_heads, head_size, input_dtype\n        ),\n        out_idx=-1,\n        execution_backend="auto",\n    )\n\n@lru_cache(maxsize=32)\ndef _compiled_low_precision_state_kernel(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        build_low_precision_state_program(\n            batch_size, num_heads, head_size, input_dtype\n        ),\n        out_idx=-1,\n        execution_backend="auto",\n    )\n\n\ndef torch_state_update(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    normalized_key: torch.Tensor,\n    gate_a: torch.Tensor,\n    value: torch.Tensor,\n    key: torch.Tensor,\n    receptance: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    anti_matrix = (-normalized_key).unsqueeze(-1) @ (normalized_key * gate_a).unsqueeze(\n        -2\n    )\n    value_key = value.unsqueeze(-1) @ key.unsqueeze(-2)\n    state_f32 = state.float()\n    next_state = state_f32 * decay.float().unsqueeze(-2)\n    next_state = next_state + state_f32 @ anti_matrix.float()\n    next_state = next_state + value_key.float()\n    if state.dtype == BFLOAT16:\n        next_state = next_state.to(BFLOAT16)\n    mixed = (next_state.to(receptance.dtype) @ receptance.unsqueeze(-1)).squeeze(-1)\n    return next_state, mixed\n\n\n@lru_cache(maxsize=32)\ndef _compiled_x_mix(\n    batch_size: int,\n    hidden_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        build_x_mix_program(batch_size, hidden_size, input_dtype),\n        out_idx=-1,\n        execution_backend="auto",\n    )\n\n\ndef tilelang_x_mix(\n    x: torch.Tensor,\n    previous: torch.Tensor,\n    mix_r: torch.Tensor,\n    mix_w: torch.Tensor,\n    mix_k: torch.Tensor,\n    mix_v: torch.Tensor,\n    mix_a: torch.Tensor,\n    mix_g: torch.Tensor,\n) -> torch.Tensor:\n    """Fuse exact six-way BF16/FP16 decode x-mix."""\n    if x.device.type != "cuda" or any(\n        tensor.device != x.device\n        for tensor in (previous, mix_r, mix_w, mix_k, mix_v, mix_a, mix_g)\n    ):\n        raise RuntimeError("TileLang x-mix requires one CUDA device")\n    if x.dtype != BFLOAT16 or any(\n        tensor.dtype != x.dtype\n        for tensor in (previous, mix_r, mix_w, mix_k, mix_v, mix_a, mix_g)\n    ):\n        raise TypeError("TileLang x-mix requires matching bfloat16 tensors")\n    if x.ndim != 2 or previous.shape != x.shape:\n        raise ValueError("x and previous must have shape [batch, hidden]")\n    batch_size, hidden_size = x.shape\n    if any(\n        tensor.shape != (hidden_size,)\n        for tensor in (mix_r, mix_w, mix_k, mix_v, mix_a, mix_g)\n    ):\n        raise ValueError("x-mix weights must have shape [hidden]")\n    kernel: Any = _compiled_x_mix(\n        batch_size,\n        hidden_size,\n        _dtype_name(x.dtype),\n        cuda_arch_key(x.device),\n    )\n    return kernel(\n        x.contiguous(),\n        previous.contiguous(),\n        mix_r.contiguous(),\n        mix_w.contiguous(),\n        mix_k.contiguous(),\n        mix_v.contiguous(),\n        mix_a.contiguous(),\n        mix_g.contiguous(),\n    )\n\n\n@lru_cache(maxsize=32)\ndef _compiled_post_state(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        build_post_state_program(batch_size, num_heads, head_size, input_dtype),\n        out_idx=-1,\n        execution_backend="auto",\n    )\n\n\ndef tilelang_post_state(\n    normalized: torch.Tensor,\n    receptance: torch.Tensor,\n    key: torch.Tensor,\n    value: torch.Tensor,\n    r_k: torch.Tensor,\n    gate: torch.Tensor,\n) -> torch.Tensor:\n    """Fuse exact BF16 RKV correction and gating after native GroupNorm."""\n    inputs = (receptance, key, value, gate)\n    if normalized.device.type != "cuda" or any(\n        tensor.device != normalized.device for tensor in (*inputs, r_k)\n    ):\n        raise RuntimeError("TileLang post-state fusion requires one CUDA device")\n    if normalized.dtype != BFLOAT16 or any(\n        tensor.dtype != normalized.dtype for tensor in (*inputs, r_k)\n    ):\n        raise TypeError("TileLang post-state fusion requires matching bfloat16")\n    if normalized.ndim != 3:\n        raise ValueError("normalized must have shape [batch, heads, head]")\n    batch_size, num_heads, head_size = normalized.shape\n    expected = (batch_size, num_heads, head_size)\n    if any(tensor.shape != expected for tensor in inputs):\n        raise ValueError("post-state vector shapes are incompatible")\n    if r_k.shape != (num_heads, head_size):\n        raise ValueError("r_k must have shape [heads, head]")\n    kernel: Any = _compiled_post_state(\n        batch_size,\n        num_heads,\n        head_size,\n        _dtype_name(normalized.dtype),\n        cuda_arch_key(normalized.device),\n    )\n    return kernel(\n        normalized.contiguous(),\n        receptance.contiguous(),\n        key.contiguous(),\n        value.contiguous(),\n        r_k.contiguous(),\n        gate.contiguous(),\n    )\n\n\n@lru_cache(maxsize=32)\ndef _compiled_fused_state_update(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        build_fused_state_update_program(\n            batch_size, num_heads, head_size, input_dtype\n        ),\n        out_idx=[-2, -1],\n        execution_backend="auto",\n    )\n\n\n\n\ndef tilelang_fused_state_update(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    anti_update: torch.Tensor,\n    value_key: torch.Tensor,\n    receptance: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    """Fuse exact FP32 state finalization and validated narrow projection."""\n    batch_size, num_heads, head_size, columns = state.shape\n    if columns != head_size or head_size % 16:\n        raise ValueError("fused state update requires square 16-aligned heads")\n    if state.device.type != "cuda" or any(\n        tensor.device != state.device\n        for tensor in (decay, anti_update, value_key, receptance)\n    ):\n        raise RuntimeError("fused state update requires one CUDA device")\n    _require_exact_fused_state_projection(state.device)\n    if state.dtype != FLOAT32 or anti_update.dtype != FLOAT32 or value_key.dtype != FLOAT32:\n        raise TypeError("fused state update requires FP32 state updates")\n    if decay.dtype != receptance.dtype or decay.dtype not in {FLOAT16, BFLOAT16}:\n        raise TypeError("fused state update requires matching float16/bfloat16 vectors")\n    kernel: Any = _compiled_fused_state_update(\n        batch_size,\n        num_heads,\n        head_size,\n        _dtype_name(decay.dtype),\n        cuda_arch_key(state.device),\n    )\n    return kernel(\n        state.contiguous(),\n        decay.contiguous(),\n        anti_update.contiguous(),\n        value_key.contiguous(),\n        receptance.contiguous(),\n    )\n\n\n\n\n\n\n\n\n@lru_cache(maxsize=32)\ndef _compiled_state_projection(\n    batch_size: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        build_state_projection_program(\n            batch_size, num_heads, head_size, input_dtype\n        ),\n        out_idx=-1,\n        execution_backend="auto",\n    )\n\n\ndef tilelang_state_projection(\n    state: torch.Tensor, receptance: torch.Tensor\n) -> torch.Tensor:\n    """Project head state on GPUs with validated exact padded MMA behavior."""\n    if state.device.type != "cuda" or receptance.device != state.device:\n        raise RuntimeError("TileLang state projection requires one CUDA device")\n    _require_exact_fused_state_projection(state.device)\n    if state.dtype != receptance.dtype or state.dtype not in {FLOAT16, BFLOAT16}:\n        raise TypeError("TileLang state projection requires matching float16/bfloat16")\n    if state.ndim != 4 or receptance.ndim != 3:\n        raise ValueError("state/receptance ranks must be 4 and 3")\n    batch_size, num_heads, head_size, columns = state.shape\n    if columns != head_size or receptance.shape != (batch_size, num_heads, head_size):\n        raise ValueError("state/receptance shapes are incompatible")\n    if head_size % 16:\n        raise ValueError("TileLang state projection requires head size divisible by 16")\n    kernel: Any = _compiled_state_projection(\n        batch_size,\n        num_heads,\n        head_size,\n        _dtype_name(state.dtype),\n        cuda_arch_key(state.device),\n    )\n    return kernel(state.contiguous(), receptance.contiguous())\n\n\ndef tilelang_state_finalize(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    anti_update: torch.Tensor,\n    value_key: torch.Tensor,\n) -> torch.Tensor:\n    """Run compiled pointwise state finalization without PyTorch dispatch."""\n    batch_size, num_heads, head_size, _ = state.shape\n    kernel: Any = _compiled_kernel(\n        batch_size,\n        num_heads,\n        head_size,\n        _dtype_name(decay.dtype),\n        cuda_arch_key(state.device),\n    )\n    return kernel(\n        state.contiguous(),\n        decay.contiguous(),\n        anti_update.contiguous(),\n        value_key.contiguous(),\n    )\n\n\ndef tilelang_state_finalize_backward(\n    grad_output: torch.Tensor, decay: torch.Tensor\n) -> torch.Tensor:\n    """Run the compiled first-order state-gradient kernel."""\n    batch_size, num_heads, head_size, _ = grad_output.shape\n    kernel: Any = _compiled_backward_kernel(\n        batch_size,\n        num_heads,\n        head_size,\n        _dtype_name(decay.dtype),\n        cuda_arch_key(grad_output.device),\n    )\n    return kernel(grad_output.contiguous(), decay.contiguous())\n\n\ndef tilelang_low_precision_state_finalize(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    anti_update: torch.Tensor,\n    value_key: torch.Tensor,\n) -> torch.Tensor:\n    """Finalize recurrent state with FP32 compute and low-precision storage."""\n    if state.dtype not in {FLOAT16, BFLOAT16} or decay.dtype != state.dtype:\n        raise TypeError(\n            "Low-precision state finalization requires matching FP16/BF16 state and decay"\n        )\n    if anti_update.dtype != FLOAT32 or value_key.dtype != FLOAT32:\n        raise TypeError("Low-precision state finalization requires FP32 updates")\n    batch_size, num_heads, head_size, columns = state.shape\n    if columns != head_size:\n        raise ValueError("Low-precision recurrent state must be square per head")\n    kernel: Any = _compiled_low_precision_state_kernel(\n        batch_size,\n        num_heads,\n        head_size,\n        _dtype_name(decay.dtype),\n        cuda_arch_key(state.device),\n    )\n    return kernel(\n        state.contiguous(),\n        decay.contiguous(),\n        anti_update.contiguous(),\n        value_key.contiguous(),\n    )\n\n\n\n\ndef tilelang_state_update(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    normalized_key: torch.Tensor,\n    gate_a: torch.Tensor,\n    value: torch.Tensor,\n    key: torch.Tensor,\n    receptance: torch.Tensor,\n    *,\n    state_finalize_op: Any | None = None,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    if state.device.type != "cuda":\n        raise RuntimeError("TileLang RWKV kernel requires CUDA tensors")\n\n    state = state.contiguous()\n    decay = decay.contiguous()\n    normalized_key = normalized_key.contiguous()\n    gate_a = gate_a.contiguous()\n    value = value.contiguous()\n    key = key.contiguous()\n    receptance = receptance.contiguous()\n\n    # Keep BF16/FP16 outer products and FP32 batched GEMM in PyTorch. Their\n    # accumulation order is observable after long recurrent sequences.\n    anti_matrix = (-normalized_key).unsqueeze(-1) @ (normalized_key * gate_a).unsqueeze(\n        -2\n    )\n    state_f32 = state.float()\n    anti_update = state_f32 @ anti_matrix.float()\n    value_key = (value.unsqueeze(-1) @ key.unsqueeze(-2)).float()\n    if state.dtype in {FLOAT16, BFLOAT16}:\n        next_state = tilelang_low_precision_state_finalize(\n            state, decay, anti_update, value_key\n        )\n        mixed = (next_state @ receptance.unsqueeze(-1)).squeeze(-1)\n        return next_state, mixed\n    if state.dtype != FLOAT32:\n        raise TypeError(\n            "TileLang recurrent state must be float32, float16, or bfloat16"\n        )\n\n    differentiable = IS_GRAD_ENABLED() and any(\n        tensor.requires_grad for tensor in (state, decay, anti_update, value_key)\n    )\n    if (\n        not differentiable\n        and receptance.dtype in {FLOAT16, BFLOAT16}\n        and state.shape[-1] % 16 == 0\n        and exact_fused_state_projection_supported(state.device)\n    ):\n        return tilelang_fused_state_update(\n            state, decay, anti_update, value_key, receptance\n        )\n    if differentiable:\n        if state_finalize_op is None:\n            raise RuntimeError(\n                "differentiable TileLang state update requires the registered custom op"\n            )\n        next_state = state_finalize_op(state, decay, anti_update, value_key)\n    else:\n        next_state = tilelang_state_finalize(state, decay, anti_update, value_key)\n    mixed = (next_state.to(receptance.dtype) @ receptance.unsqueeze(-1)).squeeze(-1)\n    return next_state, mixed\n\n\ndef torch_state_scan(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    normalized_key: torch.Tensor,\n    gate_a: torch.Tensor,\n    value: torch.Tensor,\n    key: torch.Tensor,\n    receptance: torch.Tensor,\n    *,\n    active: torch.Tensor | None = None,\n    reset: torch.Tensor | None = None,\n    output_mode: str = "full",\n    chunk_size: int = 0,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    """Exact executable specification for recurrent sequence state updates."""\n    if output_mode not in {"full", "final"}:\n        raise ValueError("output_mode must be full or final")\n    if chunk_size < 0:\n        raise ValueError("chunk_size must be non-negative")\n    if decay.ndim != 4:\n        raise ValueError("Sequence tensors must have shape [batch, time, heads, head]")\n    sequence_length = decay.shape[1]\n    if sequence_length == 0:\n        raise ValueError("State scan requires at least one token")\n    step = chunk_size or sequence_length\n    matrix = state\n    mixed_steps: list[torch.Tensor] = []\n    final_mixed: torch.Tensor | None = None\n    for chunk_start in range(0, sequence_length, step):\n        chunk_stop = min(chunk_start + step, sequence_length)\n        for token_index in range(chunk_start, chunk_stop):\n            if reset is not None:\n                reset_token = reset[:, token_index].bool().reshape(-1, 1, 1, 1)\n                matrix = TORCH_WHERE(reset_token, TORCH_ZEROS_LIKE(matrix), matrix)\n            matrix_candidate, mixed_candidate = torch_state_update(\n                matrix,\n                decay[:, token_index],\n                normalized_key[:, token_index],\n                gate_a[:, token_index],\n                value[:, token_index],\n                key[:, token_index],\n                receptance[:, token_index],\n            )\n            if active is None:\n                matrix = matrix_candidate\n                mixed = mixed_candidate\n            else:\n                active_token = active[:, token_index].bool().reshape(-1, 1, 1, 1)\n                matrix = TORCH_WHERE(active_token, matrix_candidate, matrix)\n                mixed = (\n                    matrix.to(receptance.dtype)\n                    @ receptance[:, token_index].unsqueeze(-1)\n                ).squeeze(-1)\n            if output_mode == "full":\n                mixed_steps.append(mixed)\n            final_mixed = mixed\n    if output_mode == "full":\n        return matrix, TORCH_STACK(mixed_steps, dim=1)\n    if final_mixed is None:\n        raise RuntimeError("State scan produced no output")\n    return matrix, final_mixed.unsqueeze(1)\n\n\ndef tilelang_state_scan(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    normalized_key: torch.Tensor,\n    gate_a: torch.Tensor,\n    value: torch.Tensor,\n    key: torch.Tensor,\n    receptance: torch.Tensor,\n    *,\n    active: torch.Tensor | None = None,\n    reset: torch.Tensor | None = None,\n    output_mode: str = "full",\n    chunk_size: int = 0,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    """Exact TileLang-assisted sequence scan using the validated state kernel."""\n    if state.device.type != "cuda":\n        raise RuntimeError("TileLang RWKV sequence scan requires CUDA tensors")\n    if output_mode not in {"full", "final"}:\n        raise ValueError("output_mode must be full or final")\n    if chunk_size < 0:\n        raise ValueError("chunk_size must be non-negative")\n    sequence_length = decay.shape[1]\n    if sequence_length == 0:\n        raise ValueError("State scan requires at least one token")\n    step = chunk_size or sequence_length\n    matrix = state\n    mixed_steps: list[torch.Tensor] = []\n    final_mixed: torch.Tensor | None = None\n    for chunk_start in range(0, sequence_length, step):\n        chunk_stop = min(chunk_start + step, sequence_length)\n        for token_index in range(chunk_start, chunk_stop):\n            if reset is not None:\n                reset_token = reset[:, token_index].bool().reshape(-1, 1, 1, 1)\n                matrix = TORCH_WHERE(reset_token, TORCH_ZEROS_LIKE(matrix), matrix)\n            candidate, candidate_mixed = tilelang_state_update(\n                matrix,\n                decay[:, token_index],\n                normalized_key[:, token_index],\n                gate_a[:, token_index],\n                value[:, token_index],\n                key[:, token_index],\n                receptance[:, token_index],\n            )\n            if active is None:\n                matrix = candidate\n                mixed = candidate_mixed\n            else:\n                active_token = active[:, token_index].bool().reshape(-1, 1, 1, 1)\n                matrix = TORCH_WHERE(active_token, candidate, matrix)\n                mixed = (\n                    matrix.to(receptance.dtype)\n                    @ receptance[:, token_index].unsqueeze(-1)\n                ).squeeze(-1)\n            if output_mode == "full":\n                mixed_steps.append(mixed)\n            final_mixed = mixed\n    if output_mode == "full":\n        return matrix, TORCH_STACK(mixed_steps, dim=1)\n    if final_mixed is None:\n        raise RuntimeError("State scan produced no output")\n    return matrix, final_mixed.unsqueeze(1)\n\n\ndef build_fast_state_scan_program(\n    batch_size: int,\n    sequence_length: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n):\n    """Build experimental persistent recurrent scan with one block per head."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        decay: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), input_dtype\n        ),\n        normalized_key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), input_dtype\n        ),\n        gate_a: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), input_dtype\n        ),\n        value: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), input_dtype\n        ),\n        key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), input_dtype\n        ),\n        receptance: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), input_dtype\n        ),\n        next_state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float32"\n        ),\n        mixed: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), input_dtype\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=256) as (batch, head):\n            current = T.alloc_shared((head_size, head_size), "float32")\n            following = T.alloc_shared((head_size, head_size), "float32")\n            for row, column in T.Parallel(head_size, head_size):\n                current[row, column] = state[batch, head, row, column]\n            T.sync_threads()\n\n            for token in T.serial(sequence_length):\n                for row, column in T.Parallel(head_size, head_size):\n                    anti_projection = T.alloc_local((1,), "float32")\n                    anti_projection[0] = 0.0\n                    for inner in T.serial(head_size):\n                        left = T.cast(\n                            -normalized_key[batch, token, head, inner], input_dtype\n                        )\n                        right = T.cast(\n                            normalized_key[batch, token, head, column]\n                            * gate_a[batch, token, head, column],\n                            input_dtype,\n                        )\n                        anti_element = T.cast(left * right, input_dtype)\n                        anti_projection[0] += current[row, inner] * T.cast(\n                            anti_element, "float32"\n                        )\n                    value_key = T.cast(\n                        T.cast(\n                            value[batch, token, head, row]\n                            * key[batch, token, head, column],\n                            input_dtype,\n                        ),\n                        "float32",\n                    )\n                    following[row, column] = (\n                        current[row, column]\n                        * T.cast(decay[batch, token, head, column], "float32")\n                        + anti_projection[0]\n                        + value_key\n                    )\n                T.sync_threads()\n                for row, column in T.Parallel(head_size, head_size):\n                    current[row, column] = following[row, column]\n                T.sync_threads()\n\n                for row in T.Parallel(head_size):\n                    projection = T.alloc_local((1,), "float32")\n                    projection[0] = 0.0\n                    for column in T.serial(head_size):\n                        projection[0] += (\n                            T.cast(current[row, column], input_dtype)\n                            * receptance[batch, token, head, column]\n                        )\n                    mixed[batch, token, head, row] = T.cast(projection[0], input_dtype)\n                T.sync_threads()\n\n            for row, column in T.Parallel(head_size, head_size):\n                next_state[batch, head, row, column] = current[row, column]\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_fast_state_scan(\n    batch_size: int,\n    sequence_length: int,\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        build_fast_state_scan_program(\n            batch_size,\n            sequence_length,\n            num_heads,\n            head_size,\n            input_dtype,\n        ),\n        out_idx=[-2, -1],\n        execution_backend="auto",\n    )\n\n\ndef tilelang_fast_state_scan(\n    state: torch.Tensor,\n    decay: torch.Tensor,\n    normalized_key: torch.Tensor,\n    gate_a: torch.Tensor,\n    value: torch.Tensor,\n    key: torch.Tensor,\n    receptance: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    """Run experimental approximate persistent sequence scan."""\n    if state.device.type != "cuda":\n        raise RuntimeError("TileLang fast state scan requires CUDA tensors")\n    batch_size, sequence_length, num_heads, head_size = decay.shape\n    kernel: Any = _compiled_fast_state_scan(\n        batch_size,\n        sequence_length,\n        num_heads,\n        head_size,\n        _dtype_name(receptance.dtype),\n        cuda_arch_key(state.device),\n    )\n    return kernel(\n        state.contiguous(),\n        decay.contiguous(),\n        normalized_key.contiguous(),\n        gate_a.contiguous(),\n        value.contiguous(),\n        key.contiguous(),\n        receptance.contiguous(),\n    )\n\n\ndef clear_tilelang_state_kernel_caches() -> None:\n    """Drop bounded Python references to compiled recurrent-state kernels."""\n    for compiler in (\n        _compiled_kernel,\n        _compiled_backward_kernel,\n        _compiled_low_precision_state_kernel,\n        _compiled_x_mix,\n        _compiled_post_state,\n        _compiled_fused_state_update,\n        _compiled_state_projection,\n        _compiled_fast_state_scan,\n    ):\n        compiler.cache_clear()\n'
    filename = '<rwkv7_state_a3597160cb8d7317151200d05b241af7edcad2928c3731687764d914c28bf21a>'
    linecache.cache[filename] = (
        len(source), None, source.splitlines(keepends=True), filename
    )
    namespace = {'__name__': 'inference.kernel.state'}
    exec(compile(source, filename, 'exec'), namespace, namespace)  # noqa: S102

    return SimpleNamespace(
        Any=namespace['Any'],
        BFLOAT16=namespace['BFLOAT16'],
        CUDA_GET_DEVICE_CAPABILITY=namespace['CUDA_GET_DEVICE_CAPABILITY'],
        EXACT_FUSED_STATE_PROJECTION_CAPABILITIES=namespace['EXACT_FUSED_STATE_PROJECTION_CAPABILITIES'],
        FLOAT16=namespace['FLOAT16'],
        FLOAT32=namespace['FLOAT32'],
        IS_GRAD_ENABLED=namespace['IS_GRAD_ENABLED'],
        TORCH_STACK=namespace['TORCH_STACK'],
        TORCH_WHERE=namespace['TORCH_WHERE'],
        TORCH_ZEROS_LIKE=namespace['TORCH_ZEROS_LIKE'],
        _compiled_backward_kernel=namespace['_compiled_backward_kernel'],
        _compiled_fast_state_scan=namespace['_compiled_fast_state_scan'],
        _compiled_fused_state_update=namespace['_compiled_fused_state_update'],
        _compiled_kernel=namespace['_compiled_kernel'],
        _compiled_low_precision_state_kernel=namespace['_compiled_low_precision_state_kernel'],
        _compiled_post_state=namespace['_compiled_post_state'],
        _compiled_state_projection=namespace['_compiled_state_projection'],
        _compiled_x_mix=namespace['_compiled_x_mix'],
        _dtype_name=namespace['_dtype_name'],
        _require_exact_fused_state_projection=namespace['_require_exact_fused_state_projection'],
        build_fast_state_scan_program=namespace['build_fast_state_scan_program'],
        build_fused_state_update_program=namespace['build_fused_state_update_program'],
        build_low_precision_state_program=namespace['build_low_precision_state_program'],
        build_post_state_program=namespace['build_post_state_program'],
        build_state_backward_program=namespace['build_state_backward_program'],
        build_state_program=namespace['build_state_program'],
        build_state_projection_program=namespace['build_state_projection_program'],
        build_x_mix_program=namespace['build_x_mix_program'],
        clear_tilelang_state_kernel_caches=namespace['clear_tilelang_state_kernel_caches'],
        cuda_arch_key=namespace['cuda_arch_key'],
        exact_fused_state_projection_supported=namespace['exact_fused_state_projection_supported'],
        lru_cache=namespace['lru_cache'],
        tilelang_fast_state_scan=namespace['tilelang_fast_state_scan'],
        tilelang_fused_state_update=namespace['tilelang_fused_state_update'],
        tilelang_low_precision_state_finalize=namespace['tilelang_low_precision_state_finalize'],
        tilelang_post_state=namespace['tilelang_post_state'],
        tilelang_state_finalize=namespace['tilelang_state_finalize'],
        tilelang_state_finalize_backward=namespace['tilelang_state_finalize_backward'],
        tilelang_state_projection=namespace['tilelang_state_projection'],
        tilelang_state_scan=namespace['tilelang_state_scan'],
        tilelang_state_update=namespace['tilelang_state_update'],
        tilelang_x_mix=namespace['tilelang_x_mix'],
        torch=namespace['torch'],
        torch_state_scan=namespace['torch_state_scan'],
        torch_state_update=namespace['torch_state_update'],
    )


def _build_decode_namespace(state_cuda_arch_key):
    source = 'from __future__ import annotations\n\n\nfrom functools import lru_cache\nfrom typing import Any\n\nimport torch\n\n\n\n_HEAD_SIZE = 64\n_TWO_NEG_41 = 4.547473508864641e-13\n_NEXP_HALF_LOG2_E = -0.8750387749145276\n_NLOG2_E = -1.4426950408889634\n_ROTATOR1_SIGNED = -1640531527\n\n\ndef _build_wkv_program(\n    batch_size: int,\n    sequence_length: int,\n    num_heads: int,\n    precise: bool = False,\n):\n    """Build a fixed-shape FP16 WKV scan with selectable recurrence precision."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    head_size = _HEAD_SIZE\n    current_dtype = "float32" if precise else "float16"\n\n    @T.prim_func\n    def kernel(\n        state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, num_heads, head_size, head_size), "float16"\n        ),\n        receptance: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), "float16"\n        ),\n        decay_raw: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), "float16"\n        ),\n        key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), "float16"\n        ),\n        value: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), "float16"\n        ),\n        gate_a: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), "float16"\n        ),\n        gate_b: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), "float16"\n        ),\n        elapsed: T.Tensor((batch_size,), "int32"),  # type: ignore[reportInvalidTypeForm]\n        output: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (batch_size, sequence_length, num_heads, head_size), "float16"\n        ),\n    ):\n        with T.Kernel(batch_size, num_heads, threads=head_size) as (batch, head):\n            current = T.alloc_local((head_size,), current_dtype)\n            r_shared = T.alloc_shared((2, head_size), "float16")\n            w_shared = T.alloc_shared((2, head_size), "float16")\n            k_shared = T.alloc_shared((2, head_size), "float16")\n            v_shared = T.alloc_shared((2, head_size), "float16")\n            a_shared = T.alloc_shared((2, head_size), "float16")\n            b_shared = T.alloc_shared((2, head_size), "float16")\n\n            row = T.get_thread_binding()\n            for column in T.serial(head_size):\n                current[column] = T.cast(\n                    state[batch, head, row, column], current_dtype\n                )\n\n            T.async_copy(\n                receptance[batch, 0, head, 0:head_size],\n                r_shared[0, 0:head_size],\n                coalesced_width=8,\n            )\n            T.async_copy(\n                decay_raw[batch, 0, head, 0:head_size],\n                w_shared[0, 0:head_size],\n                coalesced_width=8,\n            )\n            T.async_copy(\n                key[batch, 0, head, 0:head_size],\n                k_shared[0, 0:head_size],\n                coalesced_width=8,\n            )\n            T.async_copy(\n                value[batch, 0, head, 0:head_size],\n                v_shared[0, 0:head_size],\n                coalesced_width=8,\n            )\n            T.async_copy(\n                gate_a[batch, 0, head, 0:head_size],\n                a_shared[0, 0:head_size],\n                coalesced_width=8,\n            )\n            T.async_copy(\n                gate_b[batch, 0, head, 0:head_size],\n                b_shared[0, 0:head_size],\n                coalesced_width=8,\n            )\n\n            for token in T.serial(sequence_length):\n                current_buffer = token % 2\n                T.ptx_wait_group(0)\n                T.sync_threads()\n                if not precise:\n                    for column in T.Parallel(head_size):\n                        raw = T.cast(\n                            w_shared[current_buffer, column], "float32"\n                        )\n                        phase = T.cast(\n                            elapsed[batch] + head * head_size + column + token,\n                            "int32",\n                        )\n                        rotation = T.cast(\n                            phase * T.cast(_ROTATOR1_SIGNED, "int32"), "float32"\n                        ) * _TWO_NEG_41\n                        transformed = T.exp2(\n                            _NEXP_HALF_LOG2_E\n                            / (1.0 + T.exp2(_NLOG2_E * raw))\n                        ) - 1.0 + rotation\n                        w_shared[current_buffer, column] = T.cast(\n                            transformed, "float16"\n                        )\n                    T.sync_threads()\n\n                if token + 1 < sequence_length:\n                    next_buffer = (token + 1) % 2\n                    T.async_copy(\n                        receptance[batch, token + 1, head, 0:head_size],\n                        r_shared[next_buffer, 0:head_size],\n                        coalesced_width=8,\n                    )\n                    T.async_copy(\n                        decay_raw[batch, token + 1, head, 0:head_size],\n                        w_shared[next_buffer, 0:head_size],\n                        coalesced_width=8,\n                    )\n                    T.async_copy(\n                        key[batch, token + 1, head, 0:head_size],\n                        k_shared[next_buffer, 0:head_size],\n                        coalesced_width=8,\n                    )\n                    T.async_copy(\n                        value[batch, token + 1, head, 0:head_size],\n                        v_shared[next_buffer, 0:head_size],\n                        coalesced_width=8,\n                    )\n                    T.async_copy(\n                        gate_a[batch, token + 1, head, 0:head_size],\n                        a_shared[next_buffer, 0:head_size],\n                        coalesced_width=8,\n                    )\n                    T.async_copy(\n                        gate_b[batch, token + 1, head, 0:head_size],\n                        b_shared[next_buffer, 0:head_size],\n                        coalesced_width=8,\n                    )\n\n                if precise:\n                    # Match the PyTorch recurrent store boundary: accumulate the\n                    # factored rank-one update in FP32, then round state once.\n                    state_projection_f32 = T.alloc_local((1,), "float32")\n                    state_projection_f32[0] = T.cast(0.0, "float32")\n                    for column in T.serial(head_size):\n                        state_projection_f32[0] += current[column] * T.cast(\n                            a_shared[current_buffer, column], "float32"\n                        )\n                    rounded_state = T.alloc_local((head_size,), "float16")\n                    for column in T.serial(head_size):\n                        updated_f32 = T.alloc_local((1,), "float32")\n                        updated_f32[0] = current[column] * T.cast(\n                            w_shared[current_buffer, column], "float32"\n                        )\n                        updated_f32[0] = (\n                            updated_f32[0]\n                            + state_projection_f32[0]\n                            * T.cast(b_shared[current_buffer, column], "float32")\n                        )\n                        value_key = T.cast(\n                            k_shared[current_buffer, column]\n                            * v_shared[current_buffer, row],\n                            "float16",\n                        )\n                        updated_f32[0] = updated_f32[0] + T.cast(\n                            value_key, "float32"\n                        )\n                        rounded = T.cast(updated_f32[0], "float16")\n                        current[column] = T.cast(rounded, "float32")\n                        rounded_state[column] = rounded\n                    # Fixed two-lane FP16 projection preserves recurrent output\n                    # order while avoiding a second long FP32 dependency chain.\n                    output_pair = T.alloc_local((2,), "float16")\n                    for lane in T.vectorized(2):\n                        output_pair[lane] = T.cast(0.0, "float16")\n                    for pair in T.serial(head_size // 2):\n                        for lane in T.vectorized(2):\n                            column = pair * 2 + lane\n                            output_pair[lane] = T.cast(\n                                rounded_state[column]\n                                * r_shared[current_buffer, column]\n                                + output_pair[lane],\n                                "float16",\n                            )\n                    output[batch, token, head, row] = T.cast(\n                        output_pair[0] + output_pair[1], "float16"\n                    )\n                else:\n                    projection_pair = T.alloc_local((2,), "float16")\n                    for lane in T.vectorized(2):\n                        projection_pair[lane] = T.cast(0.0, "float16")\n                    for pair in T.serial(head_size // 2):\n                        for lane in T.vectorized(2):\n                            column = pair * 2 + lane\n                            projection_pair[lane] = T.cast(\n                                a_shared[current_buffer, column] * current[column]\n                                + projection_pair[lane],\n                                "float16",\n                            )\n                    state_projection = T.cast(\n                        projection_pair[0] + projection_pair[1], "float16"\n                    )\n                    output_pair = T.alloc_local((2,), "float16")\n                    for lane in T.vectorized(2):\n                        output_pair[lane] = T.cast(0.0, "float16")\n                    for pair in T.serial(head_size // 2):\n                        for lane in T.vectorized(2):\n                            column = pair * 2 + lane\n                            updated = T.cast(\n                                current[column]\n                                * w_shared[current_buffer, column]\n                                + T.cast(\n                                    k_shared[current_buffer, column]\n                                    * v_shared[current_buffer, row]\n                                    + T.cast(\n                                        state_projection\n                                        * b_shared[current_buffer, column]\n                                        + current[column],\n                                        "float16",\n                                    ),\n                                    "float16",\n                                ),\n                                "float16",\n                            )\n                            current[column] = updated\n                            output_pair[lane] = T.cast(\n                                updated * r_shared[current_buffer, column]\n                                + output_pair[lane],\n                                "float16",\n                            )\n                    output[batch, token, head, row] = T.cast(\n                        output_pair[0] + output_pair[1], "float16"\n                    )\n                # Precise mode synchronizes at the next wait/barrier before a\n                # double buffer is reused; the legacy transform needs this barrier.\n                if not precise:\n                    T.sync_threads()\n\n            for column in T.serial(head_size):\n                state[batch, head, row, column] = current[column]\n\n    return kernel\n\n\n\n\n@lru_cache(maxsize=16)\ndef _compiled_wkv(\n    batch_size: int,\n    sequence_length: int,\n    num_heads: int,\n    device_arch: str,\n    precise: bool = False,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_wkv_program(\n            batch_size,\n            sequence_length,\n            num_heads,\n            precise,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _build_wkv_w0_t1_program(num_heads: int):\n    """Build specialized B1T1 FP16 WKV with fused decay bias."""\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    head_size = _HEAD_SIZE\n\n    @T.prim_func\n    def kernel(\n        state: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, num_heads, head_size, head_size), "float16"\n        ),\n        receptance: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, 1, num_heads, head_size), "float16"\n        ),\n        decay: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, 1, num_heads, head_size), "float16"\n        ),\n        decay_bias: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads, head_size), "float16"\n        ),\n        key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, 1, num_heads, head_size), "float16"\n        ),\n        value: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, 1, num_heads, head_size), "float16"\n        ),\n        gate_a: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, 1, num_heads, head_size), "float16"\n        ),\n        gate_b: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, 1, num_heads, head_size), "float16"\n        ),\n        elapsed: T.Tensor((1,), "int32"),  # type: ignore[reportInvalidTypeForm]\n        output: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (1, 1, num_heads, head_size), "float16"\n        ),\n    ):\n        with T.Kernel(num_heads, threads=head_size) as head:\n            row = T.get_thread_binding(0)\n            current = T.alloc_local((head_size,), "float16")\n            state_vector = T.alloc_local((8,), "float16")\n            state_shared = T.alloc_shared(\n                (head_size, head_size), "float16"\n            )\n            r_shared = T.alloc_shared((head_size,), "float16")\n            w_shared = T.alloc_shared((head_size,), "float16")\n            k_shared = T.alloc_shared((head_size,), "float16")\n            v_shared = T.alloc_shared((head_size,), "float16")\n            a_shared = T.alloc_shared((head_size,), "float16")\n            b_shared = T.alloc_shared((head_size,), "float16")\n            for chunk in T.serial(head_size // 8):\n                linear_start = (chunk * head_size + row) * 8\n                for vector_lane in T.vectorized(8):\n                    linear = linear_start + vector_lane\n                    source_row = linear // head_size\n                    source_column = linear % head_size\n                    state_vector[vector_lane] = state[\n                        0, head, source_row, source_column\n                    ]\n                for vector_lane in T.serial(8):\n                    linear = linear_start + vector_lane\n                    source_row = linear // head_size\n                    source_column = linear % head_size\n                    swizzled_column = (\n                        ((source_row % 32) ^ (source_column // 2)) * 2\n                        + source_column % 2\n                    )\n                    state_shared[source_row, swizzled_column] = (\n                        state_vector[vector_lane]\n                    )\n            T.sync_threads()\n            for column in T.serial(head_size):\n                swizzled_column = (\n                    ((row % 32) ^ (column // 2)) * 2 + column % 2\n                )\n                current[column] = state_shared[row, swizzled_column]\n            r_shared[row] = receptance[0, 0, head, row]\n            k_shared[row] = key[0, 0, head, row]\n            v_shared[row] = value[0, 0, head, row]\n            a_shared[row] = gate_a[0, 0, head, row]\n            b_shared[row] = gate_b[0, 0, head, row]\n            raw = T.cast(\n                decay[0, 0, head, row] + decay_bias[head, row],\n                "float32",\n            )\n            phase = T.cast(\n                elapsed[0] + head * head_size + row, "int32"\n            )\n            rotation = T.cast(\n                phase * T.cast(_ROTATOR1_SIGNED, "int32"), "float32"\n            ) * _TWO_NEG_41\n            transformed = T.exp2(\n                _NEXP_HALF_LOG2_E\n                / (1.0 + T.exp2(_NLOG2_E * raw))\n            ) - 1.0 + rotation\n            w_shared[row] = T.cast(transformed, "float16")\n            T.sync_threads()\n            projection_pair = T.alloc_local((2,), "float16")\n            for pair_lane in T.vectorized(2):\n                projection_pair[pair_lane] = T.cast(0.0, "float16")\n            for pair in T.serial(head_size // 2):\n                for pair_lane in T.vectorized(2):\n                    column = pair * 2 + pair_lane\n                    projection_pair[pair_lane] = T.cast(\n                        a_shared[column] * current[column]\n                        + projection_pair[pair_lane],\n                        "float16",\n                    )\n            state_projection = T.cast(\n                projection_pair[0] + projection_pair[1], "float16"\n            )\n            output_pair = T.alloc_local((2,), "float16")\n            for pair_lane in T.vectorized(2):\n                output_pair[pair_lane] = T.cast(0.0, "float16")\n            for pair in T.serial(head_size // 2):\n                for pair_lane in T.vectorized(2):\n                    column = pair * 2 + pair_lane\n                    updated = T.cast(\n                        current[column] * w_shared[column]\n                        + T.cast(\n                            k_shared[column] * v_shared[row]\n                            + T.cast(\n                                state_projection * b_shared[column]\n                                + current[column],\n                                "float16",\n                            ),\n                            "float16",\n                        ),\n                        "float16",\n                    )\n                    current[column] = updated\n                    output_pair[pair_lane] = T.cast(\n                        updated * r_shared[column]\n                        + output_pair[pair_lane],\n                        "float16",\n                    )\n            output[0, 0, head, row] = T.cast(\n                output_pair[0] + output_pair[1], "float16"\n            )\n            for column in T.serial(head_size):\n                swizzled_column = (\n                    ((row % 32) ^ (column // 2)) * 2 + column % 2\n                )\n                state_shared[row, swizzled_column] = current[column]\n            T.sync_threads()\n            for chunk in T.serial(head_size // 8):\n                linear_start = (chunk * head_size + row) * 8\n                for vector_lane in T.serial(8):\n                    linear = linear_start + vector_lane\n                    target_row = linear // head_size\n                    target_column = linear % head_size\n                    swizzled_column = (\n                        ((target_row % 32) ^ (target_column // 2)) * 2\n                        + target_column % 2\n                    )\n                    state_vector[vector_lane] = state_shared[\n                        target_row, swizzled_column\n                    ]\n                for vector_lane in T.vectorized(8):\n                    linear = linear_start + vector_lane\n                    target_row = linear // head_size\n                    target_column = linear % head_size\n                    state[0, head, target_row, target_column] = (\n                        state_vector[vector_lane]\n                    )\n\n    return kernel\n\n\n@lru_cache(maxsize=8)\ndef _compiled_wkv_w0_t1(num_heads: int, device_arch: str):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_wkv_w0_t1_program(num_heads),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _wkv_w0_t1_out(\n    state: torch.Tensor,\n    receptance: torch.Tensor,\n    decay: torch.Tensor,\n    decay_bias: torch.Tensor,\n    key: torch.Tensor,\n    value: torch.Tensor,\n    gate_a: torch.Tensor,\n    gate_b: torch.Tensor,\n    elapsed: torch.Tensor,\n    output: torch.Tensor,\n) -> None:\n    """Run specialized B1T1 WKV with fused decay bias."""\n    batch, sequence, heads, head_size = receptance.shape\n    if (batch, sequence, head_size) != (1, 1, _HEAD_SIZE):\n        raise ValueError("fused-bias WKV requires B1T1 with head size 64")\n    tensors = (\n        state,\n        receptance,\n        decay,\n        decay_bias,\n        key,\n        value,\n        gate_a,\n        gate_b,\n        output,\n    )\n    if any(tensor.device.type != "cuda" for tensor in tensors):\n        raise RuntimeError("fused-bias WKV requires CUDA")\n    if any(tensor.dtype != torch.float16 for tensor in tensors):\n        raise TypeError("fused-bias WKV requires FP16")\n    if any(not tensor.is_contiguous() for tensor in tensors):\n        raise ValueError("fused-bias WKV requires contiguous tensors")\n    if tuple(state.shape) != (1, heads, head_size, head_size):\n        raise ValueError("fused-bias WKV state shape mismatch")\n    if tuple(decay_bias.shape) != (heads, head_size):\n        raise ValueError("fused-bias WKV decay-bias shape mismatch")\n    if any(\n        tuple(tensor.shape) != tuple(receptance.shape)\n        for tensor in (decay, key, value, gate_a, gate_b, output)\n    ):\n        raise ValueError("fused-bias WKV vector shape mismatch")\n    if elapsed.dtype != torch.int32 or tuple(elapsed.shape) != (1,):\n        raise TypeError("elapsed must be int32 [1]")\n    kernel: Any = _compiled_wkv_w0_t1(\n        heads, cuda_arch_key(state.device)\n    )\n    kernel(\n        state,\n        receptance,\n        decay,\n        decay_bias,\n        key,\n        value,\n        gate_a,\n        gate_b,\n        elapsed,\n        output,\n    )\n\n\ndef _wkv_kernel_out(\n    state: torch.Tensor,\n    receptance: torch.Tensor,\n    decay: torch.Tensor,\n    key: torch.Tensor,\n    value: torch.Tensor,\n    gate_a: torch.Tensor,\n    gate_b: torch.Tensor,\n    elapsed: torch.Tensor,\n    output: torch.Tensor,\n    *,\n    precise: bool,\n) -> None:\n    if state.device.type != "cuda" or output.device.type != "cuda":\n        raise RuntimeError("TileLang FP16 WKV requires CUDA tensors")\n    batch_size, sequence_length, num_heads, head_size = receptance.shape\n    expected_state = (batch_size, num_heads, head_size, head_size)\n    if head_size != _HEAD_SIZE or tuple(state.shape) != expected_state:\n        raise ValueError("TileLang WKV requires [B,H,64,64] state")\n    if not state.is_contiguous() or not output.is_contiguous():\n        raise ValueError("TileLang WKV state/output must be contiguous")\n    vectors = (receptance, decay, key, value, gate_a, gate_b)\n    if any(vector.shape != receptance.shape for vector in vectors):\n        raise ValueError("All WKV vectors must have the same [B,T,H,N] shape")\n    if output.shape != receptance.shape:\n        raise ValueError("WKV output must match receptance shape")\n    if any(vector.dtype != torch.float16 for vector in (*vectors, state, output)):\n        raise TypeError("TileLang WKV tensors must use float16")\n    if (\n        elapsed.dtype != torch.int32\n        or tuple(elapsed.shape) != (batch_size,)\n        or not elapsed.is_contiguous()\n    ):\n        raise TypeError("elapsed must be contiguous int32 [B]")\n    kernel: Any = _compiled_wkv(\n        batch_size,\n        sequence_length,\n        num_heads,\n        cuda_arch_key(state.device),\n        precise,\n    )\n    kernel(\n        state,\n        receptance.contiguous(),\n        decay.contiguous(),\n        key.contiguous(),\n        value.contiguous(),\n        gate_a.contiguous(),\n        gate_b.contiguous(),\n        elapsed,\n        output,\n    )\n\n\ndef _wkv_out(\n    state: torch.Tensor,\n    receptance: torch.Tensor,\n    decay_raw: torch.Tensor,\n    key: torch.Tensor,\n    value: torch.Tensor,\n    gate_a: torch.Tensor,\n    gate_b: torch.Tensor,\n    elapsed: torch.Tensor,\n    output: torch.Tensor,\n) -> None:\n    """Run Albatross-compatible WKV into caller-owned state/output."""\n    _wkv_kernel_out(\n        state,\n        receptance,\n        decay_raw,\n        key,\n        value,\n        gate_a,\n        gate_b,\n        elapsed,\n        output,\n        precise=False,\n    )\n\n\ndef _wkv_precise_out(\n    state: torch.Tensor,\n    receptance: torch.Tensor,\n    decay: torch.Tensor,\n    key: torch.Tensor,\n    value: torch.Tensor,\n    gate_a: torch.Tensor,\n    gate_b: torch.Tensor,\n    elapsed: torch.Tensor,\n    output: torch.Tensor,\n) -> None:\n    """Run FP32-accumulating WKV with FP16 recurrent storage."""\n    _wkv_kernel_out(\n        state,\n        receptance,\n        decay,\n        key,\n        value,\n        gate_a,\n        gate_b,\n        elapsed,\n        output,\n        precise=True,\n    )\n\n\n\n\n\ndef _build_gemv_program(\n    input_rows: int,\n    output_rows: int,\n    input_dtype: str = "float16",\n    out_tile: int = 2,\n    reduce_threads: int = 128,\n    clear_input_sized_output: bool = False,\n):\n    """Build output-tiled FP32-accumulating inference GEMV."""\n    if input_rows <= 0 or output_rows <= 0 or output_rows % out_tile:\n        raise ValueError("GEMV dimensions must be positive and output tiled")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("GEMV dtype must be float16 or bfloat16")\n    vector_width = 8\n    block_k = reduce_threads * vector_width\n    if input_rows % block_k:\n        raise ValueError("GEMV input must divide reduction tile")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        value: T.Tensor((input_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        weight: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (output_rows, input_rows), input_dtype\n        ),\n        output: T.Tensor((output_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        clear_output: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (input_rows,), input_dtype\n        ),\n    ):\n        with T.Kernel(\n            output_rows // out_tile, threads=reduce_threads\n        ) as block:\n            thread = T.get_thread_binding(0)\n            row = block * out_tile\n            value_local = T.alloc_local((vector_width,), input_dtype)\n            weight_local = T.alloc_local(\n                (out_tile, vector_width), input_dtype\n            )\n            accumulator0 = T.alloc_local((1,), "float32")\n            accumulator1 = T.alloc_local((1,), "float32")\n            reduced0 = T.alloc_local((1,), "float32")\n            reduced1 = T.alloc_local((1,), "float32")\n            T.clear(accumulator0)\n            T.clear(accumulator1)\n            for chunk in T.serial(input_rows // block_k):\n                for lane in T.vectorized(vector_width):\n                    column = (\n                        chunk * block_k + thread * vector_width + lane\n                    )\n                    value_local[lane] = value[column]\n                    weight_local[0, lane] = weight[row, column]\n                    weight_local[1, lane] = weight[row + 1, column]\n                for lane in T.serial(vector_width):\n                    current = T.cast(value_local[lane], "float32")\n                    accumulator0[0] += current * T.cast(\n                        weight_local[0, lane], "float32"\n                    )\n                    accumulator1[0] += current * T.cast(\n                        weight_local[1, lane], "float32"\n                    )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        accumulator0[0],\n                        True,\n                        reduced0[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        accumulator1[0],\n                        True,\n                        reduced1[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            if thread == 0:\n                output[row] = T.cast(reduced0[0], input_dtype)\n                output[row + 1] = T.cast(reduced1[0], input_dtype)\n                if clear_input_sized_output and row < input_rows:\n                    clear_output[row] = T.cast(0.0, input_dtype)\n                    clear_output[row + 1] = T.cast(0.0, input_dtype)\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_gemv(\n    input_rows: int,\n    output_rows: int,\n    input_dtype: str,\n    out_tile: int,\n    reduce_threads: int,\n    clear_input_sized_output: bool,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_gemv_program(\n            input_rows,\n            output_rows,\n            input_dtype,\n            out_tile,\n            reduce_threads,\n            clear_input_sized_output,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\ndef _build_ffn_program(\n    channels: int,\n    ffn_rows: int,\n    block_rows: int = 2,\n    reduce_threads: int = 64,\n    input_dtype: str = "float16",\n):\n    """Build B1T1 FFN key GEMV fused with ReLU-square."""\n    if channels <= 0 or ffn_rows <= 0 or ffn_rows % block_rows:\n        raise ValueError("dimensions must be positive and rows divisible by block_rows")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("FFN dtype must be float16 or bfloat16")\n    vector_width = 8\n    block_k = reduce_threads * vector_width\n    if channels % block_k:\n        raise ValueError("channels must be divisible by reduce_threads * 8")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        mixed: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        weight: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (ffn_rows, channels), input_dtype\n        ),\n        output: T.Tensor((ffn_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(\n            ffn_rows // block_rows,\n            threads=(block_rows, reduce_threads),\n        ) as block:\n            row_lane = T.get_thread_binding(0)\n            reduce_lane = T.get_thread_binding(1)\n            row = block * block_rows + row_lane\n            mixed_local = T.alloc_local((vector_width,), input_dtype)\n            weight_local = T.alloc_local((vector_width,), input_dtype)\n            accumulator = T.alloc_local((1,), "float32")\n            reduced = T.alloc_local((1,), "float32")\n            T.clear(accumulator)\n            for chunk in T.serial(channels // block_k):\n                for lane in T.vectorized(vector_width):\n                    column = (\n                        chunk * block_k\n                        + reduce_lane * vector_width\n                        + lane\n                    )\n                    mixed_local[lane] = mixed[column]\n                    weight_local[lane] = weight[row, column]\n                for lane in T.serial(vector_width):\n                    accumulator[0] += T.cast(\n                        mixed_local[lane], "float32"\n                    ) * T.cast(weight_local[lane], "float32")\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        accumulator[0],\n                        True,\n                        reduced[0],\n                        reduce_lane,\n                        dtype="handle",\n                    )\n                )\n            if reduce_lane == 0:\n                projected = T.cast(reduced[0], input_dtype)\n                activated = T.max(projected, T.cast(0, input_dtype))\n                output[row] = T.cast(activated * activated, input_dtype)\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_ffn(\n    channels: int,\n    ffn_rows: int,\n    block_rows: int,\n    reduce_threads: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_ffn_program(\n            channels, ffn_rows, block_rows, reduce_threads, input_dtype\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _build_tmix_layernorm_mix6_program(\n    channels: int,\n    input_dtype: str = "float16",\n    epsilon: float = 1e-5,\n    threads: int = 256,\n):\n    """Build fused LayerNorm and six shifted time-mix vectors."""\n    if channels <= 0 or channels % threads:\n        raise ValueError("channels must be divisible by LayerNorm threads")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("time-mix dtype must be float16 or bfloat16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        residual: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        previous: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_weight: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_bias: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_weights: T.Tensor((6, channels), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        normalized: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mixed: T.Tensor((6, channels), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(1, threads=threads):\n            thread = T.get_thread_binding(0)\n            local_sum = T.alloc_local((1,), "float32")\n            local_square = T.alloc_local((1,), "float32")\n            reduced_sum = T.alloc_local((1,), "float32")\n            reduced_square = T.alloc_local((1,), "float32")\n            T.clear(local_sum)\n            T.clear(local_square)\n            for chunk in T.serial(channels // threads):\n                channel = chunk * threads + thread\n                value = T.cast(residual[channel], "float32")\n                local_sum[0] += value\n                local_square[0] += value * value\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        local_sum[0],\n                        True,\n                        reduced_sum[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        local_square[0],\n                        True,\n                        reduced_square[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            mean = reduced_sum[0] / channels\n            variance = T.max(\n                reduced_square[0] / channels - mean * mean, 0.0\n            )\n            inverse_std = 1.0 / T.sqrt(variance + epsilon)\n            for chunk in T.serial(channels // threads):\n                channel = chunk * threads + thread\n                current = T.cast(\n                    (\n                        (T.cast(residual[channel], "float32") - mean)\n                        * inverse_std\n                        * T.cast(norm_weight[channel], "float32")\n                        + T.cast(norm_bias[channel], "float32")\n                    ),\n                    input_dtype,\n                )\n                previous_value = previous[channel]\n                normalized[channel] = current\n                delta = T.cast(previous_value - current, input_dtype)\n                for stream in T.serial(6):\n                    mixed[stream, channel] = T.cast(\n                        current + delta * mix_weights[stream, channel],\n                        input_dtype,\n                    )\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_tmix_layernorm_mix6(\n    channels: int,\n    input_dtype: str,\n    epsilon: float,\n    threads: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_tmix_layernorm_mix6_program(\n            channels, input_dtype, epsilon, threads\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\ndef _build_cmix_layernorm_mix_program(\n    channels: int,\n    input_dtype: str = "float16",\n    epsilon: float = 1e-5,\n    threads: int = 256,\n):\n    """Build fused final residual LayerNorm and channel-mix input."""\n    if channels <= 0 or channels % threads:\n        raise ValueError("channels must be divisible by LayerNorm threads")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("channel-mix dtype must be float16 or bfloat16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        residual: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        previous: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_weight: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_bias: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_weight: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        normalized: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mixed: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(1, threads=threads):\n            thread = T.get_thread_binding(0)\n            local_sum = T.alloc_local((1,), "float32")\n            local_square = T.alloc_local((1,), "float32")\n            reduced_sum = T.alloc_local((1,), "float32")\n            reduced_square = T.alloc_local((1,), "float32")\n            T.clear(local_sum)\n            T.clear(local_square)\n            for chunk in T.serial(channels // threads):\n                channel = chunk * threads + thread\n                value = T.cast(residual[channel], "float32")\n                local_sum[0] += value\n                local_square[0] += value * value\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        local_sum[0],\n                        True,\n                        reduced_sum[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        local_square[0],\n                        True,\n                        reduced_square[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            mean = reduced_sum[0] / channels\n            variance = T.max(\n                reduced_square[0] / channels - mean * mean, 0.0\n            )\n            inverse_std = 1.0 / T.sqrt(variance + epsilon)\n            for chunk in T.serial(channels // threads):\n                channel = chunk * threads + thread\n                current = T.cast(\n                    (\n                        (T.cast(residual[channel], "float32") - mean)\n                        * inverse_std\n                        * T.cast(norm_weight[channel], "float32")\n                        + T.cast(norm_bias[channel], "float32")\n                    ),\n                    input_dtype,\n                )\n                normalized[channel] = current\n                delta = T.cast(previous[channel] - current, input_dtype)\n                mixed[channel] = T.cast(\n                    current + delta * mix_weight[channel], input_dtype\n                )\n\n    return kernel\n\n\ndef _build_cmix_add_layernorm_mix_program(\n    channels: int,\n    input_dtype: str = "float16",\n    epsilon: float = 1e-5,\n    threads: int = 256,\n):\n    """Build fused residual add, LayerNorm, and channel-mix input."""\n    if channels <= 0 or channels % threads:\n        raise ValueError("channels must be divisible by LayerNorm threads")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("channel-mix dtype must be float16 or bfloat16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        residual: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        update: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        previous: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_weight: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_bias: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mix_weight: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        combined: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        normalized: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        mixed: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(1, threads=threads):\n            thread = T.get_thread_binding(0)\n            local_sum = T.alloc_local((1,), "float32")\n            local_square = T.alloc_local((1,), "float32")\n            reduced_sum = T.alloc_local((1,), "float32")\n            reduced_square = T.alloc_local((1,), "float32")\n            T.clear(local_sum)\n            T.clear(local_square)\n            for chunk in T.serial(channels // threads):\n                channel = chunk * threads + thread\n                value_f16 = T.cast(\n                    residual[channel] + update[channel], input_dtype\n                )\n                combined[channel] = value_f16\n                value = T.cast(value_f16, "float32")\n                local_sum[0] += value\n                local_square[0] += value * value\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        local_sum[0],\n                        True,\n                        reduced_sum[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        local_square[0],\n                        True,\n                        reduced_square[0],\n                        thread,\n                        dtype="handle",\n                    )\n                )\n            mean = reduced_sum[0] / channels\n            variance = T.max(\n                reduced_square[0] / channels - mean * mean, 0.0\n            )\n            inverse_std = 1.0 / T.sqrt(variance + epsilon)\n            for chunk in T.serial(channels // threads):\n                channel = chunk * threads + thread\n                current = T.cast(\n                    (\n                        (T.cast(combined[channel], "float32") - mean)\n                        * inverse_std\n                        * T.cast(norm_weight[channel], "float32")\n                        + T.cast(norm_bias[channel], "float32")\n                    ),\n                    input_dtype,\n                )\n                previous_value = previous[channel]\n                normalized[channel] = current\n                delta = T.cast(previous_value - current, input_dtype)\n                mixed[channel] = T.cast(\n                    current + delta * mix_weight[channel], input_dtype\n                )\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_cmix_add_layernorm_mix(\n    channels: int,\n    input_dtype: str,\n    epsilon: float,\n    threads: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_add_layernorm_mix_program(\n            channels, input_dtype, epsilon, threads\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\n@lru_cache(maxsize=32)\ndef _compiled_cmix_layernorm_mix(\n    channels: int,\n    input_dtype: str,\n    epsilon: float,\n    threads: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_layernorm_mix_program(\n            channels, input_dtype, epsilon, threads\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _dtype_name(dtype: torch.dtype) -> str:\n    if dtype == torch.float16:\n        return "float16"\n    if dtype == torch.bfloat16:\n        return "bfloat16"\n    raise TypeError("TileLang decode tensors must use float16 or bfloat16")\n\n\ndef _require_contiguous_cuda(\n    tensors: tuple[torch.Tensor, ...], name: str\n) -> str:\n    if torch.is_grad_enabled() and any(\n        tensor.requires_grad for tensor in tensors\n    ):\n        raise RuntimeError(f"TileLang {name} is inference-only")\n    if any(tensor.device.type != "cuda" for tensor in tensors):\n        raise RuntimeError(f"TileLang {name} requires CUDA tensors")\n    dtype = tensors[0].dtype\n    input_dtype = _dtype_name(dtype)\n    if any(tensor.dtype != dtype for tensor in tensors):\n        raise TypeError(f"TileLang {name} tensors must share one dtype")\n    if any(not tensor.is_contiguous() for tensor in tensors):\n        raise ValueError(f"TileLang {name} requires contiguous tensors")\n    if len({tensor.device for tensor in tensors}) != 1:\n        raise ValueError(f"TileLang {name} tensors must share a device")\n    return input_dtype\n\n\ndef _ffn_out(\n    mixed: torch.Tensor,\n    weight: torch.Tensor,\n    output: torch.Tensor,\n    *,\n    block_rows: int = 1,\n    reduce_threads: int = 128,\n) -> None:\n    """Run FFN key projection and ReLU-square into caller-owned output."""\n    input_dtype = _require_contiguous_cuda(\n        (mixed, weight, output), "FFN key"\n    )\n    if mixed.dim() != 1 or weight.dim() != 2:\n        raise ValueError("FFN key expects mixed [C] and weight [F,C]")\n    channels = mixed.numel()\n    ffn_rows = weight.size(0)\n    if weight.size(1) != channels or tuple(output.shape) != (ffn_rows,):\n        raise ValueError("FFN key weight/output shape mismatch")\n    kernel: Any = _compiled_ffn(\n        channels,\n        ffn_rows,\n        block_rows,\n        reduce_threads,\n        input_dtype,\n        cuda_arch_key(mixed.device),\n    )\n    kernel(mixed, weight, output)\n\n\ndef _build_cmix_value_program(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str = "float16",\n    block_rows: int = 64,\n    reduce_threads: int = 8,\n):\n    """Build sparse FFN value projection over transposed packed weights."""\n    if channels <= 0 or ffn_rows <= 0:\n        raise ValueError("dimensions must be positive")\n    if channels % block_rows or ffn_rows % reduce_threads:\n        raise ValueError("dimensions must divide block and reduction sizes")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("FFN dtype must be float16 or bfloat16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        hidden: T.Tensor((ffn_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        weight: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (ffn_rows, channels), input_dtype\n        ),\n        output: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(\n            channels // block_rows,\n            threads=(reduce_threads, block_rows),\n        ) as block:\n            reduce_lane = T.get_thread_binding(0)\n            output_lane = T.get_thread_binding(1)\n            row = block * block_rows + output_lane\n            accumulator = T.alloc_local((1,), "float32")\n            reduced = T.alloc_local((1,), "float32")\n            T.clear(accumulator)\n            for chunk in T.serial(ffn_rows // reduce_threads):\n                hidden_row = chunk * reduce_threads + reduce_lane\n                activation = hidden[hidden_row]\n                if activation != T.cast(0, input_dtype):\n                    accumulator[0] += T.cast(\n                        activation, "float32"\n                    ) * T.cast(weight[hidden_row, row], "float32")\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        accumulator[0],\n                        True,\n                        reduced[0],\n                        reduce_lane,\n                        dtype="handle",\n                    )\n                )\n            if reduce_lane == 0:\n                output[row] = T.cast(reduced[0], input_dtype)\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_cmix_value(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str,\n    block_rows: int,\n    reduce_threads: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_value_program(\n            channels, ffn_rows, input_dtype, block_rows, reduce_threads\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _build_cmix_sparse_atomic_program(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str = "float16",\n    ffn_tile: int = 128,\n    threads: int = 128,\n):\n    """Build tiled exact-zero FFN down projection with FP16 atomics."""\n    output_tile = threads * 2\n    if (\n        channels <= 0\n        or ffn_rows <= 0\n        or ffn_rows % ffn_tile\n        or channels % output_tile\n    ):\n        raise ValueError("FFN dimensions must divide sparse tiles")\n    if input_dtype != "float16":\n        raise ValueError("atomic sparse FFN currently requires float16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        hidden: T.Tensor((ffn_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        weight: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (ffn_rows, channels), input_dtype\n        ),\n        output: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(\n            ffn_rows // ffn_tile,\n            channels // output_tile,\n            threads=threads,\n        ) as (ffn_block, channel_block):\n            thread = T.get_thread_binding(0)\n            activation_shared = T.alloc_shared((ffn_tile,), input_dtype)\n            nonzero_ids = T.alloc_shared((ffn_tile,), "int32")\n            nonzero_count = T.alloc_shared((1,), "int32")\n            warp_counts = T.alloc_shared((ffn_tile // 32,), "int32")\n            warp_prefix = T.alloc_shared((ffn_tile // 32,), "int32")\n            accumulators = T.alloc_local((2,), input_dtype)\n            ffn_row = ffn_block * ffn_tile + thread\n            preactivation = T.cast(hidden[ffn_row], "float32")\n            activated = T.max(preactivation, 0.0)\n            activation_shared[thread] = T.cast(\n                activated * activated, input_dtype\n            )\n            if thread == 0:\n                nonzero_count[0] = 0\n            T.sync_threads()\n            lane = thread % 32\n            warp = thread // 32\n            nonzero = activation_shared[thread] != T.cast(0, input_dtype)\n            mask = T.ballot_sync(nonzero)\n            local_position = T.popcount(\n                mask & ((T.uint32(1) << lane) - T.uint32(1))\n            )\n            if lane == 0:\n                warp_counts[warp] = T.popcount(mask)\n            T.sync_threads()\n            if thread == 0:\n                for warp_index in T.serial(ffn_tile // 32):\n                    warp_prefix[warp_index] = nonzero_count[0]\n                    nonzero_count[0] += warp_counts[warp_index]\n            T.sync_threads()\n            if nonzero:\n                nonzero_ids[warp_prefix[warp] + local_position] = thread\n            T.sync_threads()\n            T.clear(accumulators)\n            for index in T.serial(ffn_tile):\n                if index < nonzero_count[0]:\n                    local_row = nonzero_ids[index]\n                    actual_row = ffn_block * ffn_tile + local_row\n                    activation = activation_shared[local_row]\n                    for pair_lane in T.serial(2):\n                        channel = (\n                            channel_block * output_tile\n                            + thread * 2\n                            + pair_lane\n                        )\n                        accumulators[pair_lane] = T.cast(\n                            activation * weight[actual_row, channel]\n                            + accumulators[pair_lane],\n                            input_dtype,\n                        )\n            channel = channel_block * output_tile + thread * 2\n            T.atomic_addx2(\n                output[channel : channel + 2], accumulators[0:2]\n            )\n\n    return kernel\n\n\n@lru_cache(maxsize=16)\ndef _compiled_cmix_sparse_atomic(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str,\n    ffn_tile: int,\n    threads: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_sparse_atomic_program(\n            channels,\n            ffn_rows,\n            input_dtype,\n            ffn_tile,\n            threads,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _build_cmix_sparse_binned_program(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str = "float16",\n    ffn_tile: int = 128,\n    threads: int = 128,\n    exponent_bins: int = 6,\n):\n    """Build deterministic exponent-binned exact-zero FFN accumulation."""\n    output_tile = threads * 4\n    if (\n        channels <= 0\n        or ffn_rows <= 0\n        or ffn_rows % ffn_tile\n        or channels % output_tile\n        or exponent_bins != 6\n    ):\n        raise ValueError("binned FFN dimensions must divide sparse tiles")\n    if input_dtype != "float16":\n        raise ValueError("binned sparse FFN currently requires float16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n    from tvm.tirx import op as tirx_op  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        hidden: T.Tensor((ffn_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        weight: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (ffn_rows, channels), input_dtype\n        ),\n        output: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (exponent_bins, channels), "float32"\n        ),\n    ):\n        with T.Kernel(\n            ffn_rows // ffn_tile,\n            channels // output_tile,\n            threads=threads,\n        ) as (ffn_block, channel_block):\n            thread = T.get_thread_binding(0)\n            activation_shared = T.alloc_shared((ffn_tile,), input_dtype)\n            nonzero_ids = T.alloc_shared((ffn_tile,), "int32")\n            nonzero_count = T.alloc_shared((1,), "int32")\n            warp_counts = T.alloc_shared((ffn_tile // 32,), "int32")\n            warp_prefix = T.alloc_shared((ffn_tile // 32,), "int32")\n            accumulators = T.alloc_local((4,), input_dtype)\n            accumulator_bins = T.alloc_local((4,), "int32")\n            remaining_bins = T.alloc_local((4,), "int32")\n            vector_values = T.alloc_local((4,), "float32")\n            ffn_row = ffn_block * ffn_tile + thread\n            preactivation = T.cast(hidden[ffn_row], "float32")\n            activated = T.max(preactivation, 0.0)\n            activation_shared[thread] = T.cast(\n                activated * activated, input_dtype\n            )\n            if thread == 0:\n                nonzero_count[0] = 0\n            T.sync_threads()\n            lane = thread % 32\n            warp = thread // 32\n            nonzero = activation_shared[thread] != T.cast(0, input_dtype)\n            mask = T.ballot_sync(nonzero)\n            local_position = T.popcount(\n                mask & ((T.uint32(1) << lane) - T.uint32(1))\n            )\n            if lane == 0:\n                warp_counts[warp] = T.popcount(mask)\n            T.sync_threads()\n            if thread == 0:\n                for warp_index in T.serial(ffn_tile // 32):\n                    warp_prefix[warp_index] = nonzero_count[0]\n                    nonzero_count[0] += warp_counts[warp_index]\n            T.sync_threads()\n            if nonzero:\n                nonzero_ids[warp_prefix[warp] + local_position] = thread\n            T.sync_threads()\n            T.clear(accumulators)\n            for index in T.serial(ffn_tile):\n                if index < nonzero_count[0]:\n                    local_row = nonzero_ids[index]\n                    actual_row = ffn_block * ffn_tile + local_row\n                    activation = activation_shared[local_row]\n                    for output_lane in T.serial(4):\n                        channel = (\n                            channel_block * output_tile\n                            + thread * 4\n                            + output_lane\n                        )\n                        accumulators[output_lane] = T.cast(\n                            activation * weight[actual_row, channel]\n                            + accumulators[output_lane],\n                            input_dtype,\n                        )\n            channel = channel_block * output_tile + thread * 4\n            for output_lane in T.serial(4):\n                value = accumulators[output_lane]\n                bits = T.reinterpret(value, "uint16")\n                exponent = T.cast(\n                    T.bitwise_and(T.shift_right(bits, 10), 31),\n                    "int32",\n                )\n                accumulator_bins[output_lane] = T.if_then_else(\n                    exponent == 0,\n                    0,\n                    T.min(1 + (exponent - 1) // 6, exponent_bins - 1),\n                )\n            common_bin = T.max(\n                T.max(accumulator_bins[0], accumulator_bins[1]),\n                T.max(accumulator_bins[2], accumulator_bins[3]),\n            )\n            for output_lane in T.serial(4):\n                vector_values[output_lane] = T.if_then_else(\n                    accumulator_bins[output_lane] == common_bin,\n                    T.cast(accumulators[output_lane], "float32"),\n                    0.0,\n                )\n            T.call_intrin(\n                "float4",\n                tirx_op.Op.get("tl.atomic_addx4_elem_op"),\n                T.access_ptr(\n                    output[common_bin, channel : channel + 4], "rw"\n                ),\n                T.access_ptr(vector_values[0:4], "r"),\n            )\n            for output_lane in T.serial(4):\n                remaining_bins[output_lane] = T.if_then_else(\n                    accumulator_bins[output_lane] != common_bin\n                    and accumulators[output_lane] != T.cast(0, input_dtype),\n                    accumulator_bins[output_lane],\n                    -1,\n                )\n            second_bin = T.max(\n                T.max(remaining_bins[0], remaining_bins[1]),\n                T.max(remaining_bins[2], remaining_bins[3]),\n            )\n            if second_bin >= 0:\n                for output_lane in T.serial(4):\n                    vector_values[output_lane] = T.if_then_else(\n                        accumulator_bins[output_lane] == second_bin,\n                        T.cast(accumulators[output_lane], "float32"),\n                        0.0,\n                    )\n                T.call_intrin(\n                    "float4",\n                    tirx_op.Op.get("tl.atomic_addx4_elem_op"),\n                    T.access_ptr(\n                        output[second_bin, channel : channel + 4], "rw"\n                    ),\n                    T.access_ptr(vector_values[0:4], "r"),\n                )\n            for output_lane in T.serial(4):\n                if (\n                    accumulator_bins[output_lane] != common_bin\n                    and accumulator_bins[output_lane] != second_bin\n                    and accumulators[output_lane] != T.cast(0, input_dtype)\n                ):\n                    T.atomic_add(\n                        output[accumulator_bins[output_lane], channel + output_lane],\n                        T.cast(accumulators[output_lane], "float32"),\n                    )\n    return kernel\n\n\n@lru_cache(maxsize=16)\ndef _compiled_cmix_sparse_binned(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str,\n    ffn_tile: int,\n    threads: int,\n    exponent_bins: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_sparse_binned_program(\n            channels,\n            ffn_rows,\n            input_dtype,\n            ffn_tile,\n            threads,\n            exponent_bins,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _build_cmix_binned_finalize_program(\n    channels: int,\n    exponent_bins: int = 6,\n    input_dtype: str = "float16",\n    threads: int = 256,\n):\n    """Build fixed-order bin reduction plus FP16 residual finalization."""\n    if channels <= 0 or channels % threads or exponent_bins != 6:\n        raise ValueError("binned finalization dimensions are unsupported")\n    if input_dtype != "float16":\n        raise ValueError("binned finalization currently requires float16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        bins: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (exponent_bins, channels), "float32"\n        ),\n        residual: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        output: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(channels // threads, threads=threads) as block:\n            channel = block * threads + T.get_thread_binding(0)\n            total = T.alloc_local((1,), input_dtype)\n            total[0] = T.cast(bins[0, channel], input_dtype)\n            bins[0, channel] = 0.0\n            for exponent in T.serial(1, exponent_bins):\n                total[0] = T.cast(\n                    total[0] + T.cast(bins[exponent, channel], input_dtype),\n                    input_dtype,\n                )\n                bins[exponent, channel] = 0.0\n            contribution = total[0]\n            output[channel] = T.cast(\n                residual[channel] + contribution, input_dtype\n            )\n\n    return kernel\n\n\n@lru_cache(maxsize=16)\ndef _compiled_cmix_binned_finalize(\n    channels: int,\n    exponent_bins: int,\n    input_dtype: str,\n    threads: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_binned_finalize_program(\n            channels, exponent_bins, input_dtype, threads\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\ndef _build_cmix_sparse_split_program(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str = "float16",\n    ffn_tile: int = 128,\n    threads: int = 128,\n    splits: int = 8,\n):\n    """Build split exact-zero FFN down projection without global atomics."""\n    output_tile = threads * 2\n    if (\n        channels <= 0\n        or ffn_rows <= 0\n        or ffn_rows % (ffn_tile * splits)\n        or channels % output_tile\n        or input_dtype != "float16"\n    ):\n        raise ValueError("FP16 sparse split dimensions must divide tiles")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n\n    @T.prim_func\n    def kernel(\n        hidden: T.Tensor((ffn_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        weight: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (ffn_rows, channels), input_dtype\n        ),\n        partials: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (splits, channels), input_dtype\n        ),\n    ):\n        with T.Kernel(\n            splits,\n            channels // output_tile,\n            threads=threads,\n        ) as (split, channel_block):\n            thread = T.get_thread_binding(0)\n            activation_shared = T.alloc_shared((ffn_tile,), input_dtype)\n            nonzero_ids = T.alloc_shared((ffn_tile,), "int32")\n            nonzero_count = T.alloc_shared((1,), "int32")\n            warp_counts = T.alloc_shared((ffn_tile // 32,), "int32")\n            warp_prefix = T.alloc_shared((ffn_tile // 32,), "int32")\n            accumulators = T.alloc_local((2,), input_dtype)\n            tile_accumulators = T.alloc_local((2,), input_dtype)\n            T.clear(accumulators)\n\n            for local_tile in T.serial(\n                ffn_rows // (ffn_tile * splits)\n            ):\n                ffn_block = (\n                    split * (ffn_rows // (ffn_tile * splits)) + local_tile\n                )\n                ffn_row = ffn_block * ffn_tile + thread\n                preactivation = T.cast(hidden[ffn_row], "float32")\n                activated = T.max(preactivation, 0.0)\n                activation_shared[thread] = T.cast(\n                    activated * activated, input_dtype\n                )\n                if thread == 0:\n                    nonzero_count[0] = 0\n                T.sync_threads()\n                lane = thread % 32\n                warp = thread // 32\n                nonzero = activation_shared[thread] != T.cast(0, input_dtype)\n                mask = T.ballot_sync(nonzero)\n                local_position = T.popcount(\n                    mask & ((T.uint32(1) << lane) - T.uint32(1))\n                )\n                if lane == 0:\n                    warp_counts[warp] = T.popcount(mask)\n                T.sync_threads()\n                if thread == 0:\n                    for warp_index in T.serial(ffn_tile // 32):\n                        warp_prefix[warp_index] = nonzero_count[0]\n                        nonzero_count[0] += warp_counts[warp_index]\n                T.sync_threads()\n                if nonzero:\n                    nonzero_ids[warp_prefix[warp] + local_position] = thread\n                T.sync_threads()\n                T.clear(tile_accumulators)\n                for index in T.serial(ffn_tile):\n                    if index < nonzero_count[0]:\n                        local_row = nonzero_ids[index]\n                        actual_row = ffn_block * ffn_tile + local_row\n                        activation = activation_shared[local_row]\n                        for pair_lane in T.serial(2):\n                            channel = (\n                                channel_block * output_tile\n                                + thread * 2\n                                + pair_lane\n                            )\n                            tile_accumulators[pair_lane] = T.cast(\n                                activation * weight[actual_row, channel]\n                                + tile_accumulators[pair_lane],\n                                input_dtype,\n                            )\n                for pair_lane in T.serial(2):\n                    accumulators[pair_lane] = T.cast(\n                        accumulators[pair_lane]\n                        + tile_accumulators[pair_lane],\n                        input_dtype,\n                    )\n                T.sync_threads()\n\n            channel = channel_block * output_tile + thread * 2\n            partials[split, channel] = accumulators[0]\n            partials[split, channel + 1] = accumulators[1]\n\n    return kernel\n\n\n@lru_cache(maxsize=16)\ndef _compiled_cmix_sparse_split(\n    channels: int,\n    ffn_rows: int,\n    input_dtype: str,\n    ffn_tile: int,\n    threads: int,\n    splits: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_sparse_split_program(\n            channels,\n            ffn_rows,\n            input_dtype,\n            ffn_tile,\n            threads,\n            splits,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\ndef _build_cmix_finalize_program(\n    channels: int, splits: int, input_dtype: str = "float16"\n):\n    """Build deterministic split-output finalization."""\n    if channels <= 0 or channels % 256 or splits <= 0:\n        raise ValueError("channels must be positive and divisible by 256")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("FFN dtype must be float16 or bfloat16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        partials: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (splits, channels), input_dtype\n        ),\n        output: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(channels // 256, threads=256) as block:\n            row = block * 256 + T.get_thread_binding(0)\n            value = T.alloc_local((1,), input_dtype)\n            value[0] = partials[0, row]\n            for split in T.serial(1, splits):\n                value[0] = T.cast(\n                    value[0] + partials[split, row], input_dtype\n                )\n            output[row] = value[0]\n\n    return kernel\n\n\n@lru_cache(maxsize=16)\ndef _compiled_cmix_finalize(\n    channels: int, splits: int, input_dtype: str, device_arch: str\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_cmix_finalize_program(channels, splits, input_dtype),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _cmix_value_out(\n    hidden: torch.Tensor,\n    weight: torch.Tensor,\n    partials: torch.Tensor,\n    output: torch.Tensor,\n    *,\n    splits: int = 4,\n    block_rows: int = 64,\n    reduce_threads: int = 8,\n) -> None:\n    """Run sparse split FFN value projection with fixed graph workspaces."""\n    input_dtype = _require_contiguous_cuda(\n        (hidden, weight, partials, output), "FFN value"\n    )\n    if hidden.dim() != 1 or weight.dim() != 2:\n        raise ValueError("FFN value expects hidden [F] and weight [F,C]")\n    ffn_rows, channels = weight.shape\n    if (\n        hidden.numel() != ffn_rows\n        or ffn_rows % splits\n        or tuple(partials.shape) != (splits, channels)\n        or tuple(output.shape) != (channels,)\n    ):\n        raise ValueError("FFN value split/workspace shape mismatch")\n    split_rows = ffn_rows // splits\n    kernel: Any = _compiled_cmix_value(\n        channels,\n        split_rows,\n        input_dtype,\n        block_rows,\n        reduce_threads,\n        cuda_arch_key(hidden.device),\n    )\n    for split in range(splits):\n        start = split * split_rows\n        stop = start + split_rows\n        kernel(\n            hidden[start:stop],\n            weight[start:stop],\n            partials[split],\n        )\n    finalize: Any = _compiled_cmix_finalize(\n        channels, splits, input_dtype, cuda_arch_key(hidden.device)\n    )\n    finalize(partials, output)\n\n\ndef _build_rankout_program(\n    channels: int,\n    rank_w: int,\n    rank_a: int,\n    rank_g: int,\n    rank_v: int,\n    use_value_mix: bool,\n    input_dtype: str = "float16",\n):\n    """Build fused W/A/G/V rank-out with raw decay and pointwise gates."""\n    if channels <= 0 or min(rank_w, rank_a, rank_g, rank_v) < 0:\n        raise ValueError("rank-out dimensions must be non-negative")\n    if channels % 256:\n        raise ValueError("rank-out channels must be divisible by 256")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("rank-out dtype must be float16 or bfloat16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        decay_rank: T.Tensor((rank_w,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        a_rank: T.Tensor((rank_a,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        g_rank: T.Tensor((rank_g,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        value_rank: T.Tensor((rank_v,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        w2: T.Tensor((rank_w, channels), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        a2: T.Tensor((rank_a, channels), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        g2: T.Tensor((rank_g, channels), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        v2: T.Tensor((rank_v, channels), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        w0: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        a0: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        v0: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        value_base: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        first_value: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        decay: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        gate_a: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        gate_g: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        value: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(channels // 256, threads=256) as block:\n            channel = block * 256 + T.get_thread_binding(0)\n            decay_acc = T.alloc_local((1,), "float32")\n            a_acc = T.alloc_local((1,), "float32")\n            g_acc = T.alloc_local((1,), "float32")\n            v_acc = T.alloc_local((1,), "float32")\n            T.clear(decay_acc)\n            T.clear(a_acc)\n            T.clear(g_acc)\n            T.clear(v_acc)\n            for rank in T.serial(rank_w):\n                transformed = T.cast(\n                    T.tanh(T.cast(decay_rank[rank], "float32")),\n                    input_dtype,\n                )\n                decay_acc[0] += T.cast(\n                    transformed, "float32"\n                ) * T.cast(w2[rank, channel], "float32")\n            for rank in T.serial(rank_a):\n                a_acc[0] += T.cast(\n                    a_rank[rank], "float32"\n                ) * T.cast(a2[rank, channel], "float32")\n            for rank in T.serial(rank_g):\n                transformed = T.cast(\n                    1.0\n                    / (\n                        1.0\n                        + T.exp(-T.cast(g_rank[rank], "float32"))\n                    ),\n                    input_dtype,\n                )\n                g_acc[0] += T.cast(\n                    transformed, "float32"\n                ) * T.cast(g2[rank, channel], "float32")\n            if use_value_mix:\n                for rank in T.serial(rank_v):\n                    v_acc[0] += T.cast(\n                        value_rank[rank], "float32"\n                    ) * T.cast(v2[rank, channel], "float32")\n\n            decay[channel] = T.cast(decay_acc[0], input_dtype)\n\n            a_raw = T.cast(\n                a0[channel] + T.cast(a_acc[0], input_dtype),\n                input_dtype,\n            )\n            gate_a[channel] = T.cast(\n                1.0 / (1.0 + T.exp(-T.cast(a_raw, "float32"))),\n                input_dtype,\n            )\n            gate_g[channel] = T.cast(g_acc[0], input_dtype)\n            if use_value_mix:\n                v_raw = T.cast(\n                    v0[channel] + T.cast(v_acc[0], input_dtype),\n                    input_dtype,\n                )\n                v_gate = T.cast(\n                    1.0 / (1.0 + T.exp(-T.cast(v_raw, "float32"))),\n                    input_dtype,\n                )\n                delta = T.cast(\n                    first_value[channel] - value_base[channel],\n                    input_dtype,\n                )\n                value[channel] = T.cast(\n                    value_base[channel]\n                    + T.cast(delta * v_gate, input_dtype),\n                    input_dtype,\n                )\n            else:\n                value[channel] = value_base[channel]\n\n    return kernel\n\n\ndef _build_rankout_reduced_program(\n    channels: int,\n    rank_w: int,\n    rank_a: int,\n    rank_g: int,\n    rank_v: int,\n    use_value_mix: bool,\n    input_dtype: str = "float16",\n    out_tile: int = 4,\n    reduce_threads: int = 128,\n):\n    """Build rank-parallel W/A/G/V rank-out from transposed weights."""\n    if channels <= 0 or channels % out_tile:\n        raise ValueError("rank-out channels must be divisible by output tile")\n    if min(rank_w, rank_a, rank_g, rank_v) < 0:\n        raise ValueError("rank-out dimensions must be non-negative")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("rank-out dtype must be float16 or bfloat16")\n    if reduce_threads & (reduce_threads - 1):\n        raise ValueError("rank-out reduction threads must be a power of two")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        decay_rank: T.Tensor((rank_w,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        a_rank: T.Tensor((rank_a,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        g_rank: T.Tensor((rank_g,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        value_rank: T.Tensor((rank_v,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        w2: T.Tensor((channels, rank_w), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        a2: T.Tensor((channels, rank_a), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        g2: T.Tensor((channels, rank_g), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        v2: T.Tensor((channels, rank_v), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        a0: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        v0: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        value_base: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        first_value: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        decay: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        gate_a: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        gate_g: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        value: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(\n            channels // out_tile, threads=reduce_threads\n        ) as block:\n            thread = T.get_thread_binding(0)\n            channel_start = block * out_tile\n            accumulator = T.alloc_local((4, out_tile), "float32")\n            partial = T.alloc_shared(\n                (4, out_tile, reduce_threads), "float32"\n            )\n            T.clear(accumulator)\n\n            for chunk in T.serial((rank_w + reduce_threads - 1) // reduce_threads):\n                rank = chunk * reduce_threads + thread\n                if rank < rank_w:\n                    rank_value = T.tanh(T.cast(decay_rank[rank], "float32"))\n                    for output_lane in T.serial(out_tile):\n                        accumulator[0, output_lane] += rank_value * T.cast(\n                            w2[channel_start + output_lane, rank], "float32"\n                        )\n            for chunk in T.serial((rank_a + reduce_threads - 1) // reduce_threads):\n                rank = chunk * reduce_threads + thread\n                if rank < rank_a:\n                    rank_value = T.cast(a_rank[rank], "float32")\n                    for output_lane in T.serial(out_tile):\n                        accumulator[1, output_lane] += rank_value * T.cast(\n                            a2[channel_start + output_lane, rank], "float32"\n                        )\n            for chunk in T.serial((rank_g + reduce_threads - 1) // reduce_threads):\n                rank = chunk * reduce_threads + thread\n                if rank < rank_g:\n                    rank_value = 1.0 / (\n                        1.0 + T.exp(-T.cast(g_rank[rank], "float32"))\n                    )\n                    for output_lane in T.serial(out_tile):\n                        accumulator[2, output_lane] += rank_value * T.cast(\n                            g2[channel_start + output_lane, rank], "float32"\n                        )\n            if use_value_mix:\n                for chunk in T.serial(\n                    (rank_v + reduce_threads - 1) // reduce_threads\n                ):\n                    rank = chunk * reduce_threads + thread\n                    if rank < rank_v:\n                        rank_value = T.cast(value_rank[rank], "float32")\n                        for output_lane in T.serial(out_tile):\n                            accumulator[3, output_lane] += rank_value * T.cast(\n                                v2[channel_start + output_lane, rank], "float32"\n                            )\n\n            for role in T.serial(4):\n                for output_lane in T.serial(out_tile):\n                    partial[role, output_lane, thread] = accumulator[\n                        role, output_lane\n                    ]\n            T.sync_threads()\n            for reduction_step in T.serial(7):\n                stride = reduce_threads >> (reduction_step + 1)\n                if thread < stride:\n                    for role in T.serial(4):\n                        for output_lane in T.serial(out_tile):\n                            partial[role, output_lane, thread] += partial[\n                                role, output_lane, thread + stride\n                            ]\n                T.sync_threads()\n\n            if thread == 0:\n                for output_lane in T.serial(out_tile):\n                    channel = channel_start + output_lane\n                    decay[channel] = T.cast(\n                        partial[0, output_lane, 0], input_dtype\n                    )\n                    a_raw = T.cast(\n                        a0[channel]\n                        + T.cast(partial[1, output_lane, 0], input_dtype),\n                        input_dtype,\n                    )\n                    gate_a[channel] = T.cast(\n                        1.0 / (1.0 + T.exp(-T.cast(a_raw, "float32"))),\n                        input_dtype,\n                    )\n                    gate_g[channel] = T.cast(\n                        partial[2, output_lane, 0], input_dtype\n                    )\n                    if use_value_mix:\n                        v_raw = T.cast(\n                            v0[channel]\n                            + T.cast(partial[3, output_lane, 0], input_dtype),\n                            input_dtype,\n                        )\n                        v_gate = T.cast(\n                            1.0 / (1.0 + T.exp(-T.cast(v_raw, "float32"))),\n                            input_dtype,\n                        )\n                        delta = T.cast(\n                            first_value[channel] - value_base[channel],\n                            input_dtype,\n                        )\n                        value[channel] = T.cast(\n                            value_base[channel]\n                            + T.cast(delta * v_gate, input_dtype),\n                            input_dtype,\n                        )\n                    else:\n                        value[channel] = value_base[channel]\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_rankout_reduced(\n    channels: int,\n    rank_w: int,\n    rank_a: int,\n    rank_g: int,\n    rank_v: int,\n    use_value_mix: bool,\n    input_dtype: str,\n    out_tile: int,\n    reduce_threads: int,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_rankout_reduced_program(\n            channels,\n            rank_w,\n            rank_a,\n            rank_g,\n            rank_v,\n            use_value_mix,\n            input_dtype,\n            out_tile,\n            reduce_threads,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n@lru_cache(maxsize=32)\ndef _compiled_rankout(\n    channels: int,\n    rank_w: int,\n    rank_a: int,\n    rank_g: int,\n    rank_v: int,\n    use_value_mix: bool,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_rankout_program(\n            channels,\n            rank_w,\n            rank_a,\n            rank_g,\n            rank_v,\n            use_value_mix,\n            input_dtype,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _rankout_out(\n    ranks: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],\n    weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],\n    vectors: tuple[torch.Tensor, torch.Tensor, torch.Tensor],\n    value_base: torch.Tensor,\n    first_value: torch.Tensor,\n    outputs: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],\n    *,\n    use_value_mix: bool,\n) -> None:\n    """Run fused rank-out and pointwise finalization into fixed outputs."""\n    tensors = (*ranks, *weights, *vectors, value_base, first_value, *outputs)\n    input_dtype = _require_contiguous_cuda(tensors, "rank-out")\n    channels = value_base.numel()\n    rank_sizes = tuple(rank.numel() for rank in ranks)\n    if any(tuple(rank.shape) != (size,) for rank, size in zip(ranks, rank_sizes)):\n        raise ValueError("rank-out inputs must be one-dimensional")\n    for weight, size in zip(weights, rank_sizes, strict=True):\n        if tuple(weight.shape) != (size, channels):\n            raise ValueError("rank-out weight shape mismatch")\n    if any(tuple(vector.shape) != (channels,) for vector in vectors):\n        raise ValueError("rank-out vectors must have shape [C]")\n    if tuple(first_value.shape) != (channels,) or any(\n        tuple(output.shape) != (channels,) for output in outputs\n    ):\n        raise ValueError("rank-out value/output shape mismatch")\n    kernel: Any = _compiled_rankout(\n        channels,\n        *rank_sizes,\n        use_value_mix,\n        input_dtype,\n        cuda_arch_key(value_base.device),\n    )\n    kernel(*ranks, *weights, *vectors, value_base, first_value, *outputs)\n\n\ndef _build_key_gate_program(\n    num_heads: int,\n    head_size: int = 64,\n    input_dtype: str = "float16",\n):\n    """Build fused per-head key normalization and recurrent gate vectors."""\n    if num_heads <= 0 or head_size != 64:\n        raise ValueError("key gate requires positive heads of size 64")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("key gate dtype must be float16 or bfloat16")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n        key_scale: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n        gate_a: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n        gate_scale: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n        normalized: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n        modified: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n        anti_key: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n        anti_gate: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (num_heads * head_size,), input_dtype\n        ),\n    ):\n        with T.Kernel(num_heads, threads=head_size) as head:\n            lane = T.get_thread_binding(0)\n            channel = head * head_size + lane\n            local_square = T.alloc_local((1,), "float32")\n            reduced_square = T.alloc_local((1,), "float32")\n            scaled = T.cast(key[channel] * key_scale[channel], input_dtype)\n            local_square[0] = T.cast(scaled, "float32") * T.cast(\n                scaled, "float32"\n            )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        local_square[0],\n                        True,\n                        reduced_square[0],\n                        lane,\n                        dtype="handle",\n                    )\n                )\n            norm = T.max(T.sqrt(reduced_square[0]), 1.0e-12)\n            normalized_value = T.cast(\n                T.cast(scaled, "float32") / norm, input_dtype\n            )\n            gate_delta = T.cast(gate_a[channel] - 1.0, input_dtype)\n            gate_factor = T.cast(\n                1.0\n                + T.cast(gate_delta * gate_scale[channel], input_dtype),\n                input_dtype,\n            )\n            modified_value = T.cast(\n                key[channel] * gate_factor, input_dtype\n            )\n            normalized[channel] = normalized_value\n            modified[channel] = modified_value\n            anti_key[channel] = T.cast(-normalized_value, input_dtype)\n            anti_gate[channel] = T.cast(\n                normalized_value * gate_a[channel], input_dtype\n            )\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_key_gate(\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_key_gate_program(num_heads, head_size, input_dtype),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\ndef _build_post_state_program(\n    num_heads: int,\n    head_size: int,\n    input_dtype: str = "float16",\n    epsilon: float = 6.4e-4,\n):\n    """Build fused GroupNorm, RKV residual, and gate finalization."""\n    if num_heads <= 0 or head_size <= 0 or head_size > 1024:\n        raise ValueError("invalid post-state dimensions")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("post-state dtype must be float16 or bfloat16")\n    channels = num_heads * head_size\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    @T.prim_func\n    def kernel(\n        projected: T.Tensor((num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        receptance: T.Tensor((num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        key: T.Tensor((num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        value: T.Tensor((num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        r_k: T.Tensor((num_heads, head_size), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        gate: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_weight: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        norm_bias: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        output: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(\n            num_heads, threads=channels // num_heads\n        ) as head:\n            lane = T.get_thread_binding(0)\n            channel = head * head_size + lane\n            sample = T.cast(projected[head, lane], "float32")\n            sample_sum = T.alloc_local((1,), "float32")\n            square_sum = T.alloc_local((1,), "float32")\n            rkv_sum = T.alloc_local((1,), "float32")\n            reduced_sum = T.alloc_local((1,), "float32")\n            reduced_square = T.alloc_local((1,), "float32")\n            reduced_rkv = T.alloc_local((1,), "float32")\n            sample_sum[0] = sample\n            square_sum[0] = sample * sample\n            rkv_sum[0] = (\n                T.cast(receptance[head, lane], "float32")\n                * T.cast(key[head, lane], "float32")\n                * T.cast(r_k[head, lane], "float32")\n            )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        sample_sum[0],\n                        True,\n                        reduced_sum[0],\n                        lane,\n                        dtype="handle",\n                    )\n                )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        square_sum[0],\n                        True,\n                        reduced_square[0],\n                        lane,\n                        dtype="handle",\n                    )\n                )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        rkv_sum[0],\n                        True,\n                        reduced_rkv[0],\n                        lane,\n                        dtype="handle",\n                    )\n                )\n            mean = reduced_sum[0] / head_size\n            variance = T.max(\n                reduced_square[0] / head_size - mean * mean, 0.0\n            )\n            normalized = (sample - mean) / T.sqrt(variance + epsilon)\n            affine = (\n                normalized * T.cast(norm_weight[channel], "float32")\n                + T.cast(norm_bias[channel], "float32")\n            )\n            residual = reduced_rkv[0] * T.cast(\n                value[head, lane], "float32"\n            )\n            output[channel] = T.cast(\n                (affine + residual) * T.cast(gate[channel], "float32"),\n                input_dtype,\n            )\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_post_state(\n    num_heads: int,\n    head_size: int,\n    input_dtype: str,\n    epsilon: float,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_post_state_program(\n            num_heads, head_size, input_dtype, epsilon\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _post_state_out(\n    projected: torch.Tensor,\n    receptance: torch.Tensor,\n    key: torch.Tensor,\n    value: torch.Tensor,\n    r_k: torch.Tensor,\n    gate: torch.Tensor,\n    norm_weight: torch.Tensor,\n    norm_bias: torch.Tensor,\n    output: torch.Tensor,\n    *,\n    epsilon: float,\n) -> None:\n    """Run fused post-state normalization and residual finalization."""\n    tensors = (\n        projected,\n        receptance,\n        key,\n        value,\n        r_k,\n        gate,\n        norm_weight,\n        norm_bias,\n        output,\n    )\n    input_dtype = _require_contiguous_cuda(tensors, "post-state")\n    if projected.dim() != 2:\n        raise ValueError("post-state projected input must have shape [H,N]")\n    num_heads, head_size = projected.shape\n    channels = num_heads * head_size\n    if any(\n        tuple(tensor.shape) != (num_heads, head_size)\n        for tensor in (receptance, key, value, r_k)\n    ):\n        raise ValueError("post-state head tensor shape mismatch")\n    if any(\n        tuple(tensor.shape) != (channels,)\n        for tensor in (gate, norm_weight, norm_bias, output)\n    ):\n        raise ValueError("post-state channel tensor shape mismatch")\n    kernel: Any = _compiled_post_state(\n        num_heads,\n        head_size,\n        input_dtype,\n        epsilon,\n        cuda_arch_key(projected.device),\n    )\n    kernel(*tensors)\n\n\ndef _build_rkv_program(\n    channels: int,\n    rank_w: int,\n    rank_a: int,\n    rank_g: int,\n    rank_v: int,\n    block_rows: int = 2,\n    reduce_threads: int = 128,\n    input_dtype: str = "float16",\n):\n    """Build output-tiled direct R/K/V plus W/A/G/V rank-input GEMV."""\n    ranks = (rank_w, rank_a, rank_g, rank_v)\n    if channels <= 0 or any(rank < 0 for rank in ranks):\n        raise ValueError("channels must be positive and ranks non-negative")\n    if input_dtype not in {"float16", "bfloat16"}:\n        raise ValueError("RKV dtype must be float16 or bfloat16")\n    if block_rows not in {1, 2}:\n        raise ValueError("RKV output tile must contain one or two rows")\n    total_rows = 3 * channels + sum(ranks)\n    if any(size % block_rows for size in (channels, *ranks)):\n        raise ValueError("every RKV segment must be divisible by block_rows")\n    vector_width = 8\n    block_k = reduce_threads * vector_width\n    if channels % block_k:\n        raise ValueError("channels must be divisible by reduce_threads * 8")\n\n    import tilelang.language as T  # type: ignore[import-not-found]\n\n    r_end = channels\n    k_end = 2 * channels\n    v_end = 3 * channels\n    w_end = v_end + rank_w\n    a_end = w_end + rank_a\n    g_end = a_end + rank_g\n\n    @T.prim_func\n    def kernel(\n        xr: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        xk: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        xv: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        xw: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        xa: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        xg: T.Tensor((channels,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n        weight: T.Tensor(  # type: ignore[reportInvalidTypeForm]\n            (total_rows, channels), input_dtype\n        ),\n        output: T.Tensor((total_rows,), input_dtype),  # type: ignore[reportInvalidTypeForm]\n    ):\n        with T.Kernel(\n            total_rows // block_rows, threads=reduce_threads\n        ) as block:\n            reduce_lane = T.get_thread_binding(0)\n            row = block * block_rows\n            mixed_local = T.alloc_local((vector_width,), input_dtype)\n            weight_local = T.alloc_local(\n                (block_rows, vector_width), input_dtype\n            )\n            accumulator0 = T.alloc_local((1,), "float32")\n            accumulator1 = T.alloc_local((1,), "float32")\n            reduced0 = T.alloc_local((1,), "float32")\n            reduced1 = T.alloc_local((1,), "float32")\n            T.clear(accumulator0)\n            T.clear(accumulator1)\n            for chunk in T.serial(channels // block_k):\n                for lane in T.vectorized(vector_width):\n                    column = (\n                        chunk * block_k\n                        + reduce_lane * vector_width\n                        + lane\n                    )\n                    mixed_local[lane] = T.if_then_else(\n                        row < r_end,\n                        xr[column],\n                        T.if_then_else(\n                            row < k_end,\n                            xk[column],\n                            T.if_then_else(\n                                row < v_end,\n                                xv[column],\n                                T.if_then_else(\n                                    row < w_end,\n                                    xw[column],\n                                    T.if_then_else(\n                                        row < a_end,\n                                        xa[column],\n                                        T.if_then_else(\n                                            row < g_end, xg[column], xv[column]\n                                        ),\n                                    ),\n                                ),\n                            ),\n                        ),\n                    )\n                    weight_local[0, lane] = weight[row, column]\n                    if block_rows == 2:\n                        weight_local[1, lane] = weight[row + 1, column]\n                for lane in T.serial(vector_width):\n                    mixed_value = T.cast(mixed_local[lane], "float32")\n                    accumulator0[0] += mixed_value * T.cast(\n                        weight_local[0, lane], "float32"\n                    )\n                    if block_rows == 2:\n                        accumulator1[0] += mixed_value * T.cast(\n                            weight_local[1, lane], "float32"\n                        )\n            with T.attr(\n                T.comm_reducer(\n                    lambda left, right: left + right,\n                    [T.cast(0, "float32")],\n                ),\n                "reduce_scope",\n                T.reinterpret(T.uint64(0), dtype="handle"),\n            ):\n                T.evaluate(\n                    T.tvm_thread_allreduce(\n                        T.uint32(1),\n                        accumulator0[0],\n                        True,\n                        reduced0[0],\n                        reduce_lane,\n                        dtype="handle",\n                    )\n                )\n            if block_rows == 2:\n                with T.attr(\n                    T.comm_reducer(\n                        lambda left, right: left + right,\n                        [T.cast(0, "float32")],\n                    ),\n                    "reduce_scope",\n                    T.reinterpret(T.uint64(0), dtype="handle"),\n                ):\n                    T.evaluate(\n                        T.tvm_thread_allreduce(\n                            T.uint32(1),\n                            accumulator1[0],\n                            True,\n                            reduced1[0],\n                            reduce_lane,\n                            dtype="handle",\n                        )\n                    )\n            if reduce_lane == 0:\n                output[row] = T.cast(reduced0[0], input_dtype)\n                if block_rows == 2:\n                    output[row + 1] = T.cast(reduced1[0], input_dtype)\n\n    return kernel\n\n\n@lru_cache(maxsize=32)\ndef _compiled_rkv(\n    channels: int,\n    rank_w: int,\n    rank_a: int,\n    rank_g: int,\n    rank_v: int,\n    block_rows: int,\n    reduce_threads: int,\n    input_dtype: str,\n    device_arch: str,\n):\n    del device_arch\n    import tilelang  # type: ignore[import-not-found]\n\n    return tilelang.compile(\n        _build_rkv_program(\n            channels,\n            rank_w,\n            rank_a,\n            rank_g,\n            rank_v,\n            block_rows,\n            reduce_threads,\n            input_dtype,\n        ),\n        out_idx=[],\n        execution_backend="auto",\n    )\n\n\ndef _rkv_out(\n    inputs: tuple[torch.Tensor, ...],\n    weight: torch.Tensor,\n    output: torch.Tensor,\n    ranks: tuple[int, int, int, int],\n    *,\n    block_rows: int = 1,\n    reduce_threads: int = 128,\n) -> None:\n    """Run direct-input packed RKV and rank-input projections."""\n    if len(inputs) != 6:\n        raise ValueError("RKV expects xr, xk, xv, xw, xa, and xg")\n    input_dtype = _require_contiguous_cuda(\n        (*inputs, weight, output), "RKV"\n    )\n    if len(ranks) != 4 or any(rank < 0 for rank in ranks):\n        raise ValueError("ranks must contain four non-negative values")\n    channels = inputs[0].numel()\n    if any(tuple(value.shape) != (channels,) for value in inputs):\n        raise ValueError("every RKV input must have shape [C]")\n    total_rows = 3 * channels + sum(ranks)\n    if tuple(weight.shape) != (total_rows, channels):\n        raise ValueError("packed RKV low-rank weight shape mismatch")\n    if tuple(output.shape) != (total_rows,):\n        raise ValueError("packed RKV low-rank output shape mismatch")\n    kernel: Any = _compiled_rkv(\n        channels,\n        *ranks,\n        block_rows,\n        reduce_threads,\n        input_dtype,\n        cuda_arch_key(inputs[0].device),\n    )\n    kernel(*inputs, weight, output)\n\n\ndef clear_tilelang_kernel_caches() -> None:\n    """Drop bounded Python references to compiled decode kernels."""\n    for compiler in (\n        _compiled_wkv,\n        _compiled_wkv_w0_t1,\n        _compiled_gemv,\n        _compiled_ffn,\n        _compiled_tmix_layernorm_mix6,\n        _compiled_cmix_add_layernorm_mix,\n        _compiled_cmix_layernorm_mix,\n        _compiled_cmix_value,\n        _compiled_cmix_sparse_atomic,\n        _compiled_cmix_sparse_binned,\n        _compiled_cmix_binned_finalize,\n        _compiled_cmix_sparse_split,\n        _compiled_cmix_finalize,\n        _compiled_rankout_reduced,\n        _compiled_rankout,\n        _compiled_key_gate,\n        _compiled_post_state,\n        _compiled_rkv,\n    ):\n        compiler.cache_clear()\n'
    filename = '<rwkv7_decode_b582151a6d2bd345cc5f87415543370e027c8f943a4dfe57a18300f08560bd73>'
    linecache.cache[filename] = (
        len(source), None, source.splitlines(keepends=True), filename
    )
    namespace = {'__name__': 'inference.kernel.decode'}
    namespace['cuda_arch_key'] = state_cuda_arch_key
    exec(compile(source, filename, 'exec'), namespace, namespace)  # noqa: S102

    return SimpleNamespace(
        Any=namespace['Any'],
        _HEAD_SIZE=namespace['_HEAD_SIZE'],
        _NEXP_HALF_LOG2_E=namespace['_NEXP_HALF_LOG2_E'],
        _NLOG2_E=namespace['_NLOG2_E'],
        _ROTATOR1_SIGNED=namespace['_ROTATOR1_SIGNED'],
        _TWO_NEG_41=namespace['_TWO_NEG_41'],
        _build_cmix_add_layernorm_mix_program=namespace['_build_cmix_add_layernorm_mix_program'],
        _build_cmix_binned_finalize_program=namespace['_build_cmix_binned_finalize_program'],
        _build_cmix_finalize_program=namespace['_build_cmix_finalize_program'],
        _build_cmix_layernorm_mix_program=namespace['_build_cmix_layernorm_mix_program'],
        _build_cmix_sparse_atomic_program=namespace['_build_cmix_sparse_atomic_program'],
        _build_cmix_sparse_binned_program=namespace['_build_cmix_sparse_binned_program'],
        _build_cmix_sparse_split_program=namespace['_build_cmix_sparse_split_program'],
        _build_cmix_value_program=namespace['_build_cmix_value_program'],
        _build_ffn_program=namespace['_build_ffn_program'],
        _build_gemv_program=namespace['_build_gemv_program'],
        _build_key_gate_program=namespace['_build_key_gate_program'],
        _build_post_state_program=namespace['_build_post_state_program'],
        _build_rankout_program=namespace['_build_rankout_program'],
        _build_rankout_reduced_program=namespace['_build_rankout_reduced_program'],
        _build_rkv_program=namespace['_build_rkv_program'],
        _build_tmix_layernorm_mix6_program=namespace['_build_tmix_layernorm_mix6_program'],
        _build_wkv_program=namespace['_build_wkv_program'],
        _build_wkv_w0_t1_program=namespace['_build_wkv_w0_t1_program'],
        _cmix_value_out=namespace['_cmix_value_out'],
        _compiled_cmix_add_layernorm_mix=namespace['_compiled_cmix_add_layernorm_mix'],
        _compiled_cmix_binned_finalize=namespace['_compiled_cmix_binned_finalize'],
        _compiled_cmix_finalize=namespace['_compiled_cmix_finalize'],
        _compiled_cmix_layernorm_mix=namespace['_compiled_cmix_layernorm_mix'],
        _compiled_cmix_sparse_atomic=namespace['_compiled_cmix_sparse_atomic'],
        _compiled_cmix_sparse_binned=namespace['_compiled_cmix_sparse_binned'],
        _compiled_cmix_sparse_split=namespace['_compiled_cmix_sparse_split'],
        _compiled_cmix_value=namespace['_compiled_cmix_value'],
        _compiled_ffn=namespace['_compiled_ffn'],
        _compiled_gemv=namespace['_compiled_gemv'],
        _compiled_key_gate=namespace['_compiled_key_gate'],
        _compiled_post_state=namespace['_compiled_post_state'],
        _compiled_rankout=namespace['_compiled_rankout'],
        _compiled_rankout_reduced=namespace['_compiled_rankout_reduced'],
        _compiled_rkv=namespace['_compiled_rkv'],
        _compiled_tmix_layernorm_mix6=namespace['_compiled_tmix_layernorm_mix6'],
        _compiled_wkv=namespace['_compiled_wkv'],
        _compiled_wkv_w0_t1=namespace['_compiled_wkv_w0_t1'],
        _dtype_name=namespace['_dtype_name'],
        _ffn_out=namespace['_ffn_out'],
        _post_state_out=namespace['_post_state_out'],
        _rankout_out=namespace['_rankout_out'],
        _require_contiguous_cuda=namespace['_require_contiguous_cuda'],
        _rkv_out=namespace['_rkv_out'],
        _wkv_kernel_out=namespace['_wkv_kernel_out'],
        _wkv_out=namespace['_wkv_out'],
        _wkv_precise_out=namespace['_wkv_precise_out'],
        _wkv_w0_t1_out=namespace['_wkv_w0_t1_out'],
        clear_tilelang_kernel_caches=namespace['clear_tilelang_kernel_caches'],
        cuda_arch_key=namespace['cuda_arch_key'],
        lru_cache=namespace['lru_cache'],
        torch=namespace['torch'],
    )


state = _build_state_namespace()
decode = _build_decode_namespace(state.cuda_arch_key)

__all__ = ["state", "decode"]