"""CPU correctness tests: Neuron pure-torch tile_kernels vs the upstream tile_kernels.torch references."""

import os
import sys

VENDOR = os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..")), "deepinfer", "neuron_ops", "vendor")
sys.path.insert(0, VENDOR)

import pytest
import torch

import tile_kernels
from tile_kernels import quant
from tile_kernels import torch as ref
from tile_kernels.quant.common import get_packed_ue8m0_torch_dtype

torch.set_num_threads(4)


def bits(t):
    if t.dtype in (torch.float8_e4m3fn,):
        return t.view(torch.uint8)
    return t


def assert_same(a, b, check_stride=True):
    assert a.dtype == b.dtype, (a.dtype, b.dtype)
    assert a.shape == b.shape, (a.shape, b.shape)
    if check_stride and a.numel() > 1:
        assert a.stride() == b.stride(), (a.stride(), b.stride())
    if a.dtype.is_floating_point and a.dtype != torch.float8_e4m3fn:
        assert torch.equal(torch.nan_to_num(a), torch.nan_to_num(b)), (a - b).abs().max()
    else:
        assert torch.equal(bits(a), bits(b))


# ---------------------------------------------------------------------------
# config / surface
# ---------------------------------------------------------------------------


def test_surface():
    assert tile_kernels.config.is_ascend() is False
    assert tile_kernels.get_num_vec_cores() == 2 * tile_kernels.get_num_sms()
    assert tile_kernels.get_max_ub_per_vector_core() > 0 and tile_kernels.config.get_max_ub_per_vector_core() > 0
    prev = tile_kernels.get_num_sms()
    tile_kernels.set_num_sms(3)
    assert tile_kernels.get_num_sms() == 3 and tile_kernels.get_num_vec_cores() == 6
    tile_kernels.set_num_sms(prev)
    tile_kernels.set_pdl(True)
    tile_kernels.set_pdl(False)
    tile_kernels.use_deterministic_algorithms(True)
    tile_kernels.set_token_alignment(64)
    assert tile_kernels.get_token_alignment() == 64
    tile_kernels.set_token_alignment(128)
    assert tile_kernels.utils.align(5, 4) == 8
    assert tile_kernels.utils.dtype_to_str(torch.bfloat16) == 'bf16'
    assert get_packed_ue8m0_torch_dtype() == torch.int16
    from tile_kernels.config import is_ascend  # noqa
    from tile_kernels.mhc.expand_kernel import expand_to_mhc_fwd  # noqa
    from tile_kernels.mhc.head_compute_mix_kernel import mhc_head_compute_mix_fwd  # noqa
    from tile_kernels.mhc.norm_fn_kernel import mhc_fn_normw_merge_fwd, mhc_reduce_partials_and_rmsnorm_fwd  # noqa
    from tile_kernels.mhc.post_kernel import mhc_post_fwd  # noqa
    from tile_kernels.mhc.pre_apply_mix_kernel import mhc_pre_apply_mix_fwd  # noqa
    from tile_kernels.mhc.pre_split_mixes_kernel import mhc_pre_split_mixes_fwd  # noqa
    from tile_kernels.mhc.sinkhorn_kernel import mhc_sinkhorn_fwd  # noqa
    from tile_kernels.mhc.pre_big_fuse_kernel import mhc_pre_big_fuse_fwd  # noqa
    from tile_kernels.moe import moe_topk_gate  # noqa
    from tile_kernels.quant import swiglu_forward_and_per_token_cast, swiglu_forward, per_token_cast, batched_transpose_weight_sf  # noqa
    from tile_kernels.transform import apply_rotary  # noqa
    from tile_kernels.torch import cast, cast_back  # noqa
    from tile_kernels import get_max_ub_per_vector_core, get_num_vec_cores  # noqa
    for name in ('engram_hash', 'engram_gate_fwd', 'fused_weight'):
        assert callable(getattr(tile_kernels.engram, name))
    assert callable(tile_kernels.torch.engram.make_offsets)
    for name in ('per_token_cast', 'per_token_cast_back', 'per_token_cast_and_cast_back'):
        assert callable(getattr(quant, name))
    import importlib

    for mod in ('benchmark', 'gpu_mem', 'random', 'tilelang_warning', 'xdist_failfast'):
        importlib.import_module(f'tile_kernels.testing.pytest.{mod}')


# ---------------------------------------------------------------------------
# per_token_cast
# ---------------------------------------------------------------------------

LAYOUTS = [
    # round_sf, packed, col_major, e4m3_sf
    (False, False, False, False),
    (True, False, False, False),
    (True, True, False, False),
    (True, True, True, False),
    (True, False, True, False),
    (False, False, True, False),
    (False, False, False, True),
]


@pytest.mark.parametrize('layout', LAYOUTS)
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
@pytest.mark.parametrize('npc', [16, 32, 64, 128])
@pytest.mark.parametrize('m,hidden', [(1, 128), (7, 256), (33, 512), (64, 192)])
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
def test_per_token_cast(layout, fmt, npc, m, hidden, dtype):
    round_sf, packed, col, e4m3_sf = layout
    if hidden % npc:
        pytest.skip('hidden % npc')
    torch.manual_seed(m * 1000 + hidden + npc)
    x = (torch.randn(m, hidden) * torch.logspace(-3, 3, hidden)).to(dtype)
    x[0, :npc] = 0  # all-zero group exercises the amax clamp
    kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
    q, sf = quant.per_token_cast(x, fmt, npc, **kw)
    rq, rsf = ref.cast(x, fmt, (1, npc), **kw)
    assert_same(q, rq)
    assert_same(sf, rsf)
    # sf_only / precomputed sf variants
    sf2 = quant.per_token_cast_with_sf_only(x, fmt, npc, use_tma_aligned_col_major_sf=col, round_sf=round_sf, use_packed_ue8m0=packed, use_e4m3_sf=e4m3_sf)
    assert_same(sf2, rsf)
    if not packed and not col:
        q3 = quant.per_token_cast_with_precomputed_sf(x, fmt, npc, rsf)
        rq3 = ref.cast(x, fmt, (1, npc), sf=rsf.float())
        assert_same(q3, rq3)
    # out= in place
    if not col:
        out_x = torch.empty_like(q)
        out_sf = torch.empty_like(sf)
        rx, rs = quant.per_token_cast(x, fmt, npc, out=(out_x, out_sf), **kw)
        assert rx.data_ptr() == out_x.data_ptr() and rs.data_ptr() == out_sf.data_ptr()
        assert_same(out_x, rq)
        assert_same(out_sf, rsf)


def test_per_token_cast_ascend_layout_bytes():
    """Packed col-major sf: byte (k % 2) of storage element [k // 2, m] is exponent of group k (Ascend get_sf_index)."""
    torch.manual_seed(0)
    m, hidden, npc = 5, 256, 32
    x = torch.randn(m, hidden, dtype=torch.bfloat16) * 10
    _, sf = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
    assert sf.dtype == torch.int16 and sf.shape == (m, hidden // npc // 2) and sf.stride() == (1, m)
    storage = sf.T.contiguous().view(torch.uint8)  # (P, 2m) bytes
    _, sf_rm = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=False)
    exp = (sf_rm.view(torch.int32) >> 23)
    for mi in range(m):
        for k in range(hidden // npc):
            assert storage[k // 2, 2 * mi + k % 2].item() == exp[mi, k].item()
    # row-major packed: byte 2p + j of row m = group 2p + j
    _, sfp = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True)
    assert sfp.shape == (m, hidden // npc // 2) and sfp.is_contiguous()
    assert torch.equal(sfp.view(torch.uint8).to(torch.int32), exp)


@pytest.mark.parametrize('in_layout', ['fp32_128', 'packed_col_32', 'packed_row_32', 'fp32_1x32'])
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
def test_per_token_cast_with_input_sf(in_layout, fmt):
    torch.manual_seed(1)
    m, hidden = 256, 256
    w = torch.randn(m, hidden) * 3
    if in_layout == 'fp32_128':
        xq, xsf = ref.cast(w, 'e4m3', (128, 128), round_sf=True)
        blk = (128, 128)
    elif in_layout == 'fp32_1x32':
        xq, xsf = ref.cast(w, 'e4m3', (1, 32))
        blk = (1, 32)
    elif in_layout == 'packed_col_32':
        xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
        blk = (1, 32)
    else:
        xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True)
        blk = (1, 32)
    kw = dict(round_sf=True, use_tma_aligned_col_major_sf=True, use_packed_ue8m0=True)
    q, sf = quant.per_token_cast((xq, xsf), fmt, 32, x_block_size=blk, **kw)
    rq, rsf = ref.cast((xq, xsf), fmt, (1, 32), x_block_size=blk, **kw)
    assert_same(q, rq)
    assert_same(sf, rsf)


def test_per_token_cast_empty():
    x = torch.empty(0, 128, dtype=torch.bfloat16)
    q, sf = quant.per_token_cast(x, 'e4m3', 32, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
    assert q.shape == (0, 128) and sf.shape == (0, 2) and sf.dtype == torch.int16


# ---------------------------------------------------------------------------
# cast_back / cast_and_cast_back / per_block_cast
# ---------------------------------------------------------------------------


@pytest.mark.parametrize('layout', LAYOUTS)
@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
@pytest.mark.parametrize('npc', [16, 32, 128])
@pytest.mark.parametrize('out_fmt', ['bf16', 'fp32'])
def test_per_token_cast_back(layout, fmt, npc, out_fmt):
    round_sf, packed, col, e4m3_sf = layout
    torch.manual_seed(npc)
    x = torch.randn(19, 256, dtype=torch.bfloat16) * 5
    kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
    q, sf = ref.cast(x, fmt, (1, npc), **kw)
    out = quant.per_token_cast_back((q, sf), out_fmt, npc)
    rout = ref.cast_back((q, sf), out_fmt, (1, npc))
    assert_same(out, rout)


@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
@pytest.mark.parametrize('npc', [16, 32, 64, 128])
@pytest.mark.parametrize('sf_mode', ['plain', 'round', 'e4m3'])
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
def test_per_token_cast_and_cast_back(fmt, npc, sf_mode, dtype):
    torch.manual_seed(npc)
    base = torch.randn(23, 384, dtype=dtype) * 7
    big = torch.zeros(23, 512, dtype=dtype)
    big[:, :384] = base
    x = big[:, :384]  # strided rows
    round_sf, e4m3_sf = sf_mode == 'round', sf_mode == 'e4m3'
    quant.per_token_cast_and_cast_back(x, fmt, npc, round_sf, e4m3_sf)
    q, sf = ref.cast(base, fmt, (1, npc), round_sf=round_sf, use_e4m3_sf=e4m3_sf)
    expected = ref.cast_back((q, sf), 'fp32', (1, npc)).to(dtype)
    assert torch.equal(x, expected)
    assert torch.equal(big[:, 384:], torch.zeros(23, 128, dtype=dtype))


@pytest.mark.parametrize('block', [(128, 128), (32, 32), (1, 128)])
@pytest.mark.parametrize('layout', LAYOUTS[:6])
def test_per_block_cast(block, layout):
    round_sf, packed, col, _ = layout
    torch.manual_seed(2)
    x = torch.randn(256, 384, dtype=torch.bfloat16)
    kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col)
    q, sf = quant.per_block_cast(x, 'e4m3', block, **kw)
    rq, rsf = ref.cast(x, 'e4m3', block, **kw)
    assert_same(q, rq)
    assert_same(sf, rsf)


# ---------------------------------------------------------------------------
# swiglu
# ---------------------------------------------------------------------------


@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
@pytest.mark.parametrize('with_weights', [False, True])
@pytest.mark.parametrize('clamp', [None, 3.0])
@pytest.mark.parametrize('with_psum', [False, True])
def test_swiglu_forward(dtype, with_weights, clamp, with_psum):
    torch.manual_seed(3)
    tile_kernels.set_token_alignment(8)
    n, hidden = 64, 256
    x = (torch.randn(n, hidden * 2) * 3).to(dtype)
    w = torch.rand(n) if with_weights else None
    rsf = 2.5 if with_weights else None
    psum = torch.tensor([5, 8, 8, 20, 37, 44], dtype=torch.int32) if with_psum else None  # aligned-prefix layout
    for fmt in ('fp32', 'bf16'):
        cnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
        rcnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
        out = quant.swiglu_forward(x, fmt, psum, w, rsf, clamp, cnt)
        rout = ref.swiglu_forward(x, psum, w, rsf, clamp, rcnt)
        assert torch.equal(out, rout.to(out.dtype))
        if cnt is not None:
            assert torch.equal(cnt, rcnt), (cnt, rcnt)
    # fused fp8 cast (deepinfer get_fp8_cast_params)
    kw = dict(round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
    for npc in (32, 128):
        q, sf = quant.swiglu_forward_and_per_token_cast(x, 'e4m3', npc, psum, w, rsf, clamp_value=clamp, **kw)
        rq, rsf_ = ref.cast(ref.swiglu_forward(x, psum, w, rsf, clamp), 'e4m3', (1, npc), **kw)
        assert_same(q, rq)
        assert_same(sf, rsf_)
    tile_kernels.set_token_alignment(128)


# ---------------------------------------------------------------------------
# batched_transpose_weight_sf
# ---------------------------------------------------------------------------


@pytest.mark.parametrize('block', [32, 128])
def test_batched_transpose_weight_sf(block):
    torch.manual_seed(4)

    def make(*shape):
        e = torch.randint(110, 140, shape, dtype=torch.int32)
        return (e << 23).view(torch.float32)

    sfs = [make(384 // block, 640 // block), make(3, 256 // block, 896 // block), make(128 // block, 128 // block)]
    flags = [False, True, True]
    outs = quant.batched_transpose_weight_sf(sfs, block, flags, outs=[None] * 3)
    refs = ref.batched_transpose_weight_sf(sfs, block, flags)
    for (a0, a1), (r0, r1) in zip(outs, refs):
        assert_same(a0, r0)
        assert_same(a1, r1)
    # in-place rewrite with new values
    sfs2 = [make(*s.shape) for s in sfs]
    outs2 = quant.batched_transpose_weight_sf(sfs2, block, flags, outs=outs)
    refs2 = ref.batched_transpose_weight_sf(sfs2, block, flags)
    for (a0, a1), (o0, o1), (r0, r1) in zip(outs2, outs, refs2):
        assert a0.data_ptr() == o0.data_ptr() and a1.data_ptr() == o1.data_ptr()
        assert_same(a0, r0)
        assert_same(a1, r1)


# ---------------------------------------------------------------------------
# rope
# ---------------------------------------------------------------------------


@pytest.mark.parametrize('rot_dim', [64, 128])
@pytest.mark.parametrize('ndim', [3, 4])
@pytest.mark.parametrize('interleaved', [False, True])
@pytest.mark.parametrize('conjugate', [False, True])
@pytest.mark.parametrize('use_positions', [False, True])
@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
def test_apply_rotary(rot_dim, ndim, interleaved, conjugate, use_positions, dtype):
    torch.manual_seed(5)
    seqlen_ro, s, h, hk, b = 64, 9, 4, 2, 2
    cache = torch.randn(seqlen_ro, rot_dim)
    cos, sin = cache[:, : rot_dim // 2], cache[:, rot_dim // 2 :]
    lead = (s,) if ndim == 3 else (b, s)
    full_q = torch.randn(*lead, h, rot_dim + 32).to(dtype)
    q = full_q[..., 32:]  # tail partial rope, strided
    k = torch.randn(*lead, hk, rot_dim).to(dtype)
    pos = torch.randint(0, seqlen_ro - 4, lead, dtype=torch.int32) if use_positions else None
    off = 3
    rq = ref.rotary_embedding_ref(q.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
    rk = ref.rotary_embedding_ref(k.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
    head = full_q[..., :32].clone()
    tile_kernels.transform.apply_rotary(q, cache, k, pos, interleaved, conjugate, off)
    tol = dict(atol=0, rtol=0) if dtype == torch.float32 else dict(atol=1e-2, rtol=1e-2)
    torch.testing.assert_close(q.float(), rq.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
    torch.testing.assert_close(k.float(), rk.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
    assert torch.equal(full_q[..., :32], head)
    # bitwise vs the Ascend op order: out1 = fma(first, cos, second * -sin), out2 = fma(first, sin, second * cos)
    def fma(a, b, c):
        return (a.double() * b.double() + c.double()).float()

    q2 = full_q[..., 32:].clone()
    q0 = torch.randn_like(q2.float()).to(dtype)
    q2.copy_(q0)
    tile_kernels.transform.apply_rotary(q2, cache, None, pos, interleaved, conjugate, off)
    idx = (pos.long() + off) if pos is not None else (torch.arange(s) + off).expand(*lead[:-1], s)
    c_, s_ = cos[idx].unsqueeze(-2), sin[idx].unsqueeze(-2)
    s_ = -s_ if conjugate else s_
    xf = q0.float()
    f, sc = (xf[..., 0::2], xf[..., 1::2]) if interleaved else (xf[..., : rot_dim // 2], xf[..., rot_dim // 2 :])
    c_, s_ = c_.expand_as(f), s_.expand_as(f)
    o1, o2 = fma(f, c_, sc * -s_), fma(f, s_, sc * c_)
    exp = torch.stack([o1, o2], -1).flatten(-2) if interleaved else torch.cat([o1, o2], -1)
    assert torch.equal(q2, exp.to(dtype))


# ---------------------------------------------------------------------------
# moe
# ---------------------------------------------------------------------------


@pytest.mark.parametrize('num_shared', [0, 1, 2])
@pytest.mark.parametrize('with_map', [False, True])
@pytest.mark.parametrize('with_mask', [False, True])
@pytest.mark.parametrize('with_image', [False, True])
@pytest.mark.parametrize('with_fix', [False, True])
def test_moe_topk_gate(num_shared, with_map, with_mask, with_image, with_fix):
    torch.manual_seed(6)
    t, e, k = 37, 64, 6
    logits = torch.randn(t, e) * 2
    logits[3, 10] = logits[3, 11]  # tie
    bias = torch.randn(e) * 0.1
    kwargs = dict(num_topk=k, use_shared_as_routed=num_shared > 0, num_shared_experts=num_shared, routed_scaling_factor=2.5, ep_rank=3, scoring_func='sqrtsoftplus', bias=bias)
    if with_mask:
        kwargs['mask'] = torch.rand(t) > 0.3
    if with_image:
        kwargs['image_bias'] = torch.randn(e) * 0.1
        kwargs['image_token_mask'] = torch.rand(t) > 0.5
    if with_map:
        nl = e + num_shared
        dup = 3
        count = torch.randint(1, dup + 1, (nl,), dtype=torch.int32)
        kwargs['to_physical_map'] = torch.randint(0, 500, (nl, dup), dtype=torch.int32)
        kwargs['logical_count'] = count
    unm = unm_ref = None
    if with_fix:
        unm = torch.stack([torch.randperm(e)[:k] for _ in range(t)]).to(torch.int64)
        unm_ref = unm.clone()
        kwargs['fix_routing_mask'] = torch.rand(t) > 0.5
    idx, w = tile_kernels.moe.moe_topk_gate(logits, unmapped_topk_idx=unm, **kwargs)
    ridx, rw = ref.moe_topk_gate_forward(logits, unmapped_topk_idx=unm_ref, **kwargs)
    assert torch.equal(idx, ridx)
    torch.testing.assert_close(w, rw, atol=1e-6, rtol=1e-6)
    if with_fix:
        assert torch.equal(unm, unm_ref)


def test_moe_topk_gate_out_scores_force_random():
    torch.manual_seed(7)
    t, e, k = 16, 32, 4
    logits = torch.randn(t, e)
    mask = torch.ones(t, dtype=torch.bool)
    mask[0] = False
    fr = torch.zeros(t, dtype=torch.bool)
    fr[1:4] = True
    out = (torch.empty(t, k, dtype=torch.int64), torch.empty(t, k))
    scores = torch.empty(t, e)
    unm = torch.empty(t, k, dtype=torch.int64)
    idx, w = tile_kernels.moe.moe_topk_gate(logits, k, False, 0, 1.0, 0, mask=mask, force_random=fr, out=out, scores=scores, unmapped_topk_idx=unm)
    assert idx.data_ptr() == out[0].data_ptr()
    assert (idx[0] == -1).all() and (w[0] == 0).all() and (scores[0] == 0).all() and (unm[0] == -1).all()
    assert (unm[1:4] == -1).all() and (scores[1:4] == 0).all()
    for r in range(1, 4):
        assert len(set(idx[r].tolist())) == k and idx[r].min() >= 0 and idx[r].max() < e
    ridx, rw = ref.moe_topk_gate_forward(logits[4:], k, False, 0, 1.0, 0, 'sqrtsoftplus')
    assert torch.equal(idx[4:], ridx)
    torch.testing.assert_close(scores[4:], torch.nn.functional.softplus(logits[4:]).sqrt())


# ---------------------------------------------------------------------------
# engram
# ---------------------------------------------------------------------------


def test_engram_hash():
    torch.manual_seed(8)
    t, n, layers, tables = 29, 4, 2, 3
    ids = torch.randint(0, 100000, (t, n), dtype=torch.int32)
    mult = torch.randint(1, 2**40, (layers, n), dtype=torch.int64)
    vs = torch.randint(1000, 50000, (layers, n - 1, tables), dtype=torch.int32)
    off = tile_kernels.torch.engram.make_offsets(vs)
    mask = torch.rand(t) > 0.7
    for m in (None, mask):
        out = tile_kernels.engram.engram_hash(ids, mult, vs, off, m)
        rout = ref.engram_hash_ref(ids, mult, vs, off, m)
        assert out.dtype == torch.int32 and torch.equal(out, rout)


@pytest.mark.parametrize('inplace', [False, True])
@pytest.mark.parametrize('with_mask', [False, True])
def test_engram_gate(inplace, with_mask):
    torch.manual_seed(9)
    t, hc, d = 13, 4, 256
    hs = torch.randn(t, hc, d, dtype=torch.bfloat16)
    kv = torch.randn(t, hc + 1, d, dtype=torch.bfloat16)
    wh = torch.randn(hc, d, dtype=torch.bfloat16)
    we = torch.randn(hc, d, dtype=torch.bfloat16)
    mask = (torch.rand(t) > 0.5) if with_mask else None
    wf = tile_kernels.engram.fused_weight(wh, we)
    assert wf.dtype == torch.float32 and torch.equal(wf, wh.float() * we.float())
    r_out, r_dot, r_g, r_rx, r_rk = ref.engram_gate_ref(hs, kv, wh, we, 1e-4, 1e-6, save_for_backward=True, image_token_mask=mask)
    hs_in = hs.clone()
    out, dot, g, rx, rk = tile_kernels.engram.engram_gate_fwd(hs_in, kv, wf, 1e-6, 1e-4, save_for_backward=True, image_token_mask=mask, inplace=inplace)
    if inplace:
        assert out.data_ptr() == hs_in.data_ptr()
    sel = ~mask if mask is not None else slice(None)
    assert (out.float() - r_out.float()).abs().max() <= 0.0625
    assert torch.equal(out[mask], hs[mask]) if mask is not None else True
    torch.testing.assert_close(dot[sel], r_dot[sel], atol=1e-3, rtol=1e-4)
    torch.testing.assert_close(g[sel], r_g[sel], atol=1e-5, rtol=1e-5)
    torch.testing.assert_close(rx[sel], r_rx[sel], atol=1e-6, rtol=1e-5)
    torch.testing.assert_close(rk[sel], r_rk[sel], atol=1e-6, rtol=1e-5)
    out2, *rest = tile_kernels.engram.engram_gate_fwd(hs.clone(), kv, wf, 1e-6, 1e-4, save_for_backward=False, image_token_mask=mask)
    assert all(r is None for r in rest)
    assert torch.equal(out2, out)


# ---------------------------------------------------------------------------
# mhc
# ---------------------------------------------------------------------------

MHC = 4


def test_mhc_expand_and_mix():
    torch.manual_seed(10)
    n, h = 11, 256
    x = torch.randn(n, h, dtype=torch.bfloat16)
    o = torch.empty(n, MHC, h, dtype=torch.bfloat16)
    tile_kernels.mhc.expand_to_mhc_fwd(x, o)
    assert torch.equal(o, ref.expand_to_mhc_ref(x, MHC))

    inp = torch.randn(n, MHC)
    scale, base = torch.randn(1), torch.randn(MHC)
    out = torch.empty_like(inp)
    tile_kernels.mhc.mhc_head_compute_mix_fwd(inp, scale, base, out, 1e-6)
    torch.testing.assert_close(out, ref.mhc_head_compute_mix_ref(inp, scale, base, 1e-6), atol=0, rtol=0)

    m3 = MHC * (MHC + 2)
    mixes = torch.randn(1, n, m3)
    scale3, base3 = torch.randn(3), torch.randn(m3)
    pre, post, comb = torch.empty(n, MHC), torch.empty(n, MHC), torch.empty(n, MHC * MHC)
    tile_kernels.mhc.mhc_pre_split_mixes_fwd(mixes.view(n, m3), scale3, base3, pre, post, comb, 2.0, 1e-6)
    rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale3, base3, MHC, 2.0, 1e-6)
    torch.testing.assert_close(pre, rpre.view(n, MHC), atol=1e-7, rtol=1e-6)
    torch.testing.assert_close(post, rpost.view(n, MHC), atol=1e-7, rtol=1e-6)
    torch.testing.assert_close(comb, rcomb.reshape(n, MHC * MHC), atol=1e-7, rtol=1e-6)

    xr = torch.randn(n, MHC, h, dtype=torch.bfloat16)
    mix = torch.rand(n, MHC)
    o2 = torch.empty(n, h, dtype=torch.bfloat16)
    tile_kernels.mhc.mhc_pre_apply_mix_fwd(xr, mix, o2)
    torch.testing.assert_close(o2.float(), ref.mhc_pre_apply_mix_ref(xr.float(), mix.unsqueeze(-1)).float(), atol=1e-2, rtol=1e-2)

    cm = torch.randn(n, MHC, MHC)
    cmo = torch.empty_like(cm)
    tile_kernels.mhc.mhc_sinkhorn_fwd(cm, cmo, 10, 1e-6)
    torch.testing.assert_close(cmo, ref.sinkhorn_normalize_ref(cm, 10, 1e-6), atol=1e-6, rtol=1e-5)


def test_mhc_post():
    torch.manual_seed(11)
    a, b, h = 2, 5, 256
    x = torch.randn(a, b, h, dtype=torch.bfloat16)
    res = torch.randn(a, b, MHC, h, dtype=torch.bfloat16)
    post = torch.rand(a, b, MHC, 1)
    comb = torch.rand(a, b, MHC, MHC)
    out = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb)
    rout = ref.mhc_post_ref(x, res, post, comb)
    assert out.dtype == torch.bfloat16 and out.shape == res.shape
    assert (out.float() - rout.float()).abs().max() <= 0.0625
    out2 = torch.empty_like(res)
    r = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb, out=out2)
    assert r.data_ptr() == out2.data_ptr() and torch.equal(out2, out)


@pytest.mark.parametrize('n_splits', [1, 4, 16])
def test_mhc_norm_fn(n_splits):
    torch.manual_seed(12)
    t, h = 9, 256
    m3 = MHC * (MHC + 2)
    residual = torch.randn(1, t, MHC, h, dtype=torch.bfloat16)
    fn = torch.randn(m3, MHC * h) * 0.02
    normw = torch.rand(MHC * h) + 0.5
    merged = torch.empty_like(fn)
    tile_kernels.mhc.mhc_fn_normw_merge_fwd(fn, normw, merged)
    assert torch.equal(merged, fn * normw)
    om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual, merged, n_splits)  # (S, T, 1, M3), (S, T, 1)
    om, sq, out = torch.empty(t, 1, m3), torch.empty(t, 1), torch.empty(t, m3)
    tile_kernels.mhc.mhc_reduce_partials_and_rmsnorm_fwd(om_s, sq_s, om, sq, out, MHC * h, 1e-6, n_splits)
    rout = ref.mhc_pre_norm_fn_ref(residual, fn, normw, 1e-6).view(t, m3)
    torch.testing.assert_close(out, rout, atol=1e-4, rtol=1e-4)
    torch.testing.assert_close(om, om_s.sum(0))
    torch.testing.assert_close(sq, sq_s.sum(0))


@pytest.mark.parametrize('n_splits', [1, 16])
def test_mhc_pre_big_fuse(n_splits):
    torch.manual_seed(13)
    t, h = 7, 256
    m3 = MHC * (MHC + 2)
    residual = torch.randn(t, MHC, h, dtype=torch.bfloat16)
    fn = torch.randn(m3, MHC * h) * 0.02
    om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual.unsqueeze(0), fn, n_splits)
    gemm_mul, gemm_sq = om_s.view(n_splits, t, m3), sq_s.view(n_splits, t)
    scale, base = torch.rand(3), torch.randn(m3) * 0.1
    post, comb, layer_in = torch.empty(t, MHC), torch.empty(t, MHC * MHC), torch.empty(t, h, dtype=torch.bfloat16)
    tile_kernels.mhc.mhc_pre_big_fuse_fwd(gemm_mul, gemm_sq, scale, base, residual, post, comb, layer_in, 1e-6, 1e-6, 1e-6, 2.0, 20, n_splits=n_splits)
    mixes = ref.mhc_pre_norm_fn_ref(residual.unsqueeze(0), fn, None, 1e-6)  # (1, T, M3)
    rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale, base, MHC, 2.0, 1e-6)
    rcomb = ref.sinkhorn_normalize_ref(rcomb, 20, 1e-6)
    rli = ref.mhc_pre_apply_mix_ref(residual.unsqueeze(0).float(), rpre)
    torch.testing.assert_close(post, rpost.view(t, MHC), atol=1e-5, rtol=1e-5)
    torch.testing.assert_close(comb, rcomb.reshape(t, MHC * MHC), atol=1e-5, rtol=1e-4)
    torch.testing.assert_close(layer_in.float(), rli.view(t, h).float(), atol=2e-2, rtol=1e-2)