"""CPU correctness tests for the pure-torch deep_gemm port.

Reference implementations and data generators are ported from
DeepGEMM-Ascend/tests (generators.py, test_layout.py, test_attention.py,
test_hyperconnection.py), adapted to CPU and small shapes.
"""

import math
import os
import sys
import random

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

import torch

import deep_gemm

torch.manual_seed(0)
random.seed(0)

FAILURES = []

# bf16 outputs: the dominant error is rounding D to bf16 (~1e-6 in calc_diff terms).
BF16_TOL = 1e-5


def check(name, cond, extra=""):
    if cond:
        print(f"  [ok] {name}")
    else:
        print(f"  [FAIL] {name} {extra}")
        FAILURES.append(name)


def calc_diff(x, y):
    x, y = x.double(), y.double()
    denom = (x * x + y * y).sum().item()
    if denom == 0:
        return 0.0
    return 1 - 2 * (x * y).sum().item() / denom


def ceil_div(x, y):
    return (x + y - 1) // y


def aligned(x, a):
    return ceil_div(x, a) * a


# --------------------------------------------------------------------------- #
# generators (ported from DeepGEMM-Ascend/tests/generators.py)
# --------------------------------------------------------------------------- #

def ceil_to_ue8m0(x):
    bits = x.abs().float().view(torch.int32)
    exp = ((bits // 2 ** 23) % 256) + (bits % 2 ** 23).bool().int()
    return (exp.clamp(1, 254) << 23).view(torch.float32)


def quantize_to_fp4_raw(x):
    ax = x.abs()
    code = torch.zeros_like(x, dtype=torch.uint8)
    for boundary in (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0):
        code += (ax > boundary).to(torch.uint8)
    sign = (x < 0) & (code != 0)
    values = torch.tensor((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0), dtype=torch.float32)
    repr_f32 = values[code.long()]
    repr_f32 = torch.where(sign, -repr_f32, repr_f32)
    code = code | (sign.to(torch.uint8) * 8)
    packed = (code[..., ::2] & 0x0F) | ((code[..., 1::2] & 0x0F) << 4)
    return packed.contiguous().view(torch.int8), repr_f32


def quantize_last_dim(x, quant_dtype, gran_k=32):
    """Quantize along the last dim in gran_k blocks; returns (data, sf_fp32, ref_f32)."""
    k = x.shape[-1]
    assert k % gran_k == 0
    blocks = x.float().reshape(*x.shape[:-1], k // gran_k, gran_k)
    sf_div = 6.0 if quant_dtype == "e2m1" else 448.0
    sf = ceil_to_ue8m0(blocks.abs().amax(dim=-1).clamp(1e-4) / sf_div)
    scaled = (blocks / sf[..., None]).reshape_as(x)
    if quant_dtype == "e2m1":
        data, repr_f32 = quantize_to_fp4_raw(scaled)
    else:
        data = scaled.to(torch.float8_e4m3fn)
        repr_f32 = data.float()
    ref = (repr_f32.reshape(*x.shape[:-1], k // gran_k, gran_k) * sf[..., None]).reshape_as(x)
    return data, sf.contiguous(), ref


def pack_sf_row_major(sf_fp32):
    """Packed UE8M0 int16 in K-major order (what per_token_cast emits without the col-major view)."""
    exponent = ((sf_fp32.view(torch.int32) // (1 << 23)) & 0xFF).to(torch.uint8)
    if exponent.shape[-1] % 2:
        exponent = torch.nn.functional.pad(exponent, (0, 1))
    return exponent.contiguous().view(torch.int16)


def pack_sf_reference(sf_fp32, mn, gran_mn):
    """test_layout.py::reference - fp32 sf -> canonical int16 MN-major."""
    exponent = ((sf_fp32.view(torch.int32) // (1 << 23)) & 0xFF).to(torch.uint8)
    if exponent.shape[-1] % 2:
        exponent = torch.nn.functional.pad(exponent, (0, 1))
    sf = exponent.contiguous().view(torch.int16)
    return (
        sf.transpose(-2, -1)
        .repeat_interleave(gran_mn, -1)[..., :mn]
        .contiguous()
        .transpose(-2, -1)
    )


# --------------------------------------------------------------------------- #
# 1. SF layout
# --------------------------------------------------------------------------- #

def make_sf(shape, dtype):
    if dtype == torch.float32:
        exponent = torch.randint(1, 255, shape, dtype=torch.int32)
        return (exponent * (1 << 23)).view(torch.float32)
    return torch.randint(-32768, 32768, shape, dtype=torch.int16)


def with_major(sf, mn_major):
    if mn_major and sf.size(-2) > 1 and sf.size(-1) > 1:
        sf = sf.transpose(-2, -1).contiguous().transpose(-2, -1)
        assert sf.stride(-2) == 1
    else:
        assert sf.stride(-1) == 1
    return sf


def test_sf_layout():
    print("test_sf_layout")
    for dtype in (torch.float32, torch.int16):
        for mn_major in (False, True):
            for gran_mn in (1, 32):
                for num_groups in (None, 3):
                    mn, k = 128, 256
                    sf_k_div = 32 if dtype == torch.float32 else 64
                    shape = (ceil_div(mn, gran_mn), ceil_div(k, sf_k_div))
                    if num_groups is not None:
                        shape = (num_groups, *shape)
                    sf = with_major(make_sf(shape, dtype), mn_major)
                    out = deep_gemm.transform_sf_into_required_layout(
                        sf, mn, k, (gran_mn, gran_mn, 32), num_groups=num_groups, is_sfa=True
                    )
                    expected = pack_sf_reference(sf, mn, gran_mn) if dtype == torch.float32 else (
                        sf.transpose(-2, -1).repeat_interleave(gran_mn, -1)[..., :mn].contiguous().transpose(-2, -1)
                    )
                    tag = f"dtype={dtype}, mn_major={mn_major}, gran_mn={gran_mn}, groups={num_groups}"
                    ok = (
                        out.dtype == torch.int16
                        and out.shape == expected.shape
                        and out.stride(-2) == 1
                        and out.stride(-1) == out.size(-2)
                        and torch.equal(out, expected)
                    )
                    check(f"sf layout ({tag})", ok, f"{out.shape=} {expected.shape=} {out.stride()=}")
    # Canonical input is returned unchanged (no copy).
    mn, k = 64, 128
    canonical = torch.zeros((ceil_div(k, 64), mn), dtype=torch.int16).transpose(0, 1)
    out = deep_gemm.transform_sf_into_required_layout(canonical, mn, k, (1, 1, 32), is_sfa=True)
    check("canonical sf passthrough", out.data_ptr() == canonical.data_ptr())
    # Ascend/Neuron col-major sf is unpadded, so get_tma_aligned_size is the identity.
    check("get_tma_aligned_size (unpadded)", deep_gemm.get_tma_aligned_size(5, 2) == 5 and deep_gemm.get_tma_aligned_size(9, 4) == 9)
    # Unpadded col-major int16 sf (strides (1, MN)) is the canonical layout.
    mn2, k2 = 37, 256
    col_major = torch.zeros((ceil_div(k2, 64), mn2), dtype=torch.int16).transpose(0, 1)
    out2 = deep_gemm.transform_sf_into_required_layout(col_major, mn2, k2, (1, 1, 32), is_sfa=False)
    check("unpadded col-major sf passthrough", out2.data_ptr() == col_major.data_ptr() and out2.stride() == (1, mn2))


# --------------------------------------------------------------------------- #
# 2. dense GEMMs
# --------------------------------------------------------------------------- #

def test_bf16_gemm():
    print("test_bf16_gemm")
    m, n, k = 65, 128, 256
    a = torch.randn(m, k, dtype=torch.bfloat16)
    b = torch.randn(n, k, dtype=torch.bfloat16)
    ref = a.float() @ b.float().T
    d = torch.empty(m, n, dtype=torch.bfloat16)
    deep_gemm.bf16_gemm_nt(a, b, d, None)
    check("bf16_gemm_nt", calc_diff(d.float(), ref) < BF16_TOL, calc_diff(d.float(), ref))

    d32 = torch.empty(m, n, dtype=torch.float32)
    deep_gemm.cublaslt_gemm_nt(a, b, d32, None)
    check("cublaslt_gemm_nt == bf16 ref", calc_diff(d32, ref) < 1e-10)

    # accumulate: c aliases d
    acc = torch.randn(m, n, dtype=torch.float32)
    d_acc = acc.clone()
    deep_gemm.bf16_gemm_nt(a, b, d_acc, d_acc)
    check("bf16_gemm_nt accumulate", calc_diff(d_acc, ref + acc) < 1e-10)

    # separate [1, N] bias
    bias = torch.randn(1, n, dtype=torch.float32)
    d_bias = torch.empty(m, n, dtype=torch.float32)
    deep_gemm.bf16_gemm_nt(a, b, d_bias, bias)
    check("bf16_gemm_nt bias", calc_diff(d_bias, ref + bias) < 1e-10)

    # tn / nn variants
    at = a.T.contiguous()
    d_tn = torch.empty(m, n, dtype=torch.float32)
    deep_gemm.bf16_gemm_tn(at, b.T.contiguous(), d_tn, None)
    check("bf16_gemm_tn", calc_diff(d_tn, ref) < 1e-10)
    d_nn = torch.empty(m, n, dtype=torch.float32)
    deep_gemm.bf16_gemm_nn(a, b.T.contiguous(), d_nn, None)
    check("bf16_gemm_nn", calc_diff(d_nn, ref) < 1e-10)
    # strided output view
    storage = torch.empty(m, n + 7, dtype=torch.bfloat16)
    view = storage[:, :n]
    deep_gemm.bf16_gemm_nt(a, b, view, None)
    check("bf16_gemm_nt strided out", calc_diff(view.float(), ref) < BF16_TOL, calc_diff(view.float(), ref))


def run_fp8_fp4_gemm_case(a_fmt, b_fmt, sf_mode, m=64, n=128, k=256, acc=False, gran_b=1):
    a_bf16 = torch.randn(m, k, dtype=torch.bfloat16)
    b_bf16 = torch.randn(n, k, dtype=torch.bfloat16)
    a_data, a_sf, a_ref = quantize_last_dim(a_bf16, a_fmt)
    b_data, b_sf, b_ref = quantize_last_dim(b_bf16, b_fmt)
    if sf_mode == "packed":
        a_sf_in = pack_sf_reference(a_sf, m, 1)
        b_sf_in = pack_sf_reference(b_sf, n, 1)
    elif sf_mode == "packed_row_major":
        a_sf_in = pack_sf_row_major(a_sf)
        b_sf_in = pack_sf_row_major(b_sf)
    else:
        a_sf_in, b_sf_in = a_sf, b_sf
    ref = a_ref @ b_ref.T
    d = torch.empty(m, n, dtype=torch.bfloat16)
    c = None
    if acc:
        base = torch.randn(m, n, dtype=torch.bfloat16)
        d.copy_(base)
        c = d
        ref = ref + base.float()
    deep_gemm.fp8_fp4_gemm_nt((a_data, a_sf_in), (b_data, b_sf_in), d, c, recipe=(1, 1, 32))
    return calc_diff(d.float(), ref)


def test_fp8_fp4_gemm():
    print("test_fp8_fp4_gemm")
    for a_fmt, b_fmt in (("e4m3", "e4m3"), ("e4m3", "e2m1"), ("e2m1", "e2m1")):
        for sf_mode in ("fp32", "packed", "packed_row_major"):
            diff = run_fp8_fp4_gemm_case(a_fmt, b_fmt, sf_mode)
            check(f"fp8_fp4_gemm_nt ({a_fmt}x{b_fmt}, sf={sf_mode})", diff < BF16_TOL, f"{diff=}")
    diff = run_fp8_fp4_gemm_case("e4m3", "e2m1", "packed", acc=True)
    check("fp8_fp4_gemm_nt accumulate", diff < BF16_TOL, f"{diff=}")


# --------------------------------------------------------------------------- #
# 3. m-grouped contiguous GEMMs
# --------------------------------------------------------------------------- #

def make_psum_layout(expected_m, num_groups, alignment):
    layout, psum = [], 0
    for _ in range(num_groups):
        size = max(int(random.uniform(0.7, 1.3) * expected_m), 0)
        layout.append(psum + size)
        psum = aligned(psum + size, alignment)
    return psum, layout


def iter_psum_slice(layout, alignment):
    last = 0
    for ptr in layout:
        yield slice(last, ptr)
        last = aligned(ptr, alignment)


def test_m_grouped():
    print("test_m_grouped")
    alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout()
    check("mk alignment == 256", alignment == 256)
    deep_gemm.set_mk_alignment_for_contiguous_layout(alignment)
    num_groups, expected_m, n, k = 3, 200, 64, 128
    total_m, layout = make_psum_layout(expected_m, num_groups, alignment)
    grouped = torch.tensor(layout, dtype=torch.int32)
    spans = list(iter_psum_slice(layout, alignment))

    # bf16
    a = torch.randn(total_m, k, dtype=torch.bfloat16)
    for span in spans:  # zero the inter-group padding
        pass
    pad_mask = torch.ones(total_m, dtype=torch.bool)
    for span in spans:
        pad_mask[span] = False
    a[pad_mask] = 0
    b = torch.randn(num_groups, n, k, dtype=torch.bfloat16)
    ref = torch.zeros(total_m, n, dtype=torch.float32)
    for i, span in enumerate(spans):
        ref[span] = a[span].float() @ b[i].float().T
    d = torch.empty(total_m, n, dtype=torch.bfloat16)
    deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped, use_psum_layout=True, expected_m_for_psum_layout=expected_m)
    check("m_grouped_bf16_gemm_nt_contiguous", calc_diff(d.float(), ref) < BF16_TOL, calc_diff(d.float(), ref))
    check("m_grouped bf16 zero padding", bool((d.float()[pad_mask] == 0).all()))

    # fp8 x fp4 (the DeepInfer mxfp4 MoE path)
    a_bf16 = torch.randn(total_m, k, dtype=torch.bfloat16)
    a_bf16[pad_mask] = 0
    a_data, a_sf, a_ref = quantize_last_dim(a_bf16, "e4m3")
    b_bf16 = torch.randn(num_groups, n, k, dtype=torch.bfloat16)
    b_data, b_sf, b_ref = quantize_last_dim(b_bf16, "e2m1")
    a_sf_packed = pack_sf_reference(a_sf, total_m, 1)
    b_sf_packed = pack_sf_reference(b_sf, n, 1)
    ref = torch.zeros(total_m, n, dtype=torch.float32)
    for i, span in enumerate(spans):
        ref[span] = a_ref[span] @ b_ref[i].T
    d = torch.empty(total_m, n, dtype=torch.bfloat16)
    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
        (a_data, a_sf_packed), (b_data, b_sf_packed), d, grouped,
        recipe_a=(1, 32), recipe_b=(1, 32), use_psum_layout=True, expected_m_for_psum_layout=expected_m,
    )
    diff = calc_diff(d.float(), ref)
    check("m_grouped_fp8_fp4_gemm_nt_contiguous", diff < BF16_TOL, f"{diff=}")
    check("m_grouped fp8xfp4 zero padding", bool((d.float()[pad_mask] == 0).all()))


# --------------------------------------------------------------------------- #
# 4. einsum
# --------------------------------------------------------------------------- #

def test_einsum():
    print("test_einsum")
    b_dim, h, r, dd = 7, 4, 128, 64
    a_bf16 = torch.randn(b_dim, h, r, dtype=torch.bfloat16)
    w_bf16 = torch.randn(h, dd, r, dtype=torch.bfloat16)
    ref = torch.einsum("bhr,hdr->bhd", a_bf16.float(), w_bf16.float())
    out = torch.empty(b_dim, h, dd, dtype=torch.bfloat16)
    deep_gemm.einsum("bhr,hdr->bhd", a_bf16, w_bf16, out, None)
    check("einsum bf16 bhr,hdr->bhd", calc_diff(out.float(), ref) < BF16_TOL, calc_diff(out.float(), ref))

    a_data, a_sf, a_ref = quantize_last_dim(a_bf16, "e4m3")
    w_data, w_sf, w_ref = quantize_last_dim(w_bf16, "e4m3")
    a_sf_packed = pack_sf_reference(a_sf, a_sf.shape[-2], 1)  # [b, h, r/64] canonical per b
    w_sf_packed = pack_sf_reference(w_sf, w_sf.shape[-2], 1)
    ref = torch.einsum("bhr,hdr->bhd", a_ref, w_ref)
    out = torch.empty(b_dim, h, dd, dtype=torch.bfloat16)
    deep_gemm.fp8_einsum("bhr,hdr->bhd", (a_data, a_sf_packed), (w_data, w_sf_packed), out, recipe=(1, 1, 32))
    diff = calc_diff(out.float(), ref)
    check("fp8_einsum bhr,hdr->bhd", diff < BF16_TOL, f"{diff=}")
    # the fp32-sf input form must agree too
    out2 = torch.empty_like(out)
    deep_gemm.fp8_einsum("bhr,hdr->bhd", (a_data, a_sf), (w_data, w_sf), out2, recipe=(1, 1, 32))
    check("fp8_einsum fp32 sf == packed sf", torch.equal(out, out2))


# --------------------------------------------------------------------------- #
# 5. MQA logits (ported from tests/test_attention.py)
# --------------------------------------------------------------------------- #

def ceil_to_ue8m0_with_byte(x):
    bits = x.abs().float().view(torch.int32)
    exp = (((bits >> 23) & 0xFF) + (bits & 0x7FFFFF).bool().int()).clamp(1, 254)
    return (exp << 23).view(torch.float32), exp.to(torch.uint8)


def quantize_mqa_tensor(x, use_fp4):
    sf_div = 6.0 if use_fp4 else 448.0
    sf_shape = (*x.shape[:-1], x.shape[-1] // 32, 32)
    sf, sf_byte = ceil_to_ue8m0_with_byte(x.float().reshape(sf_shape).abs().amax(dim=-1).clamp(1e-4) / sf_div)
    x_scaled = (x.float().reshape(sf_shape) / sf[..., None]).reshape_as(x)
    if use_fp4:
        x_quant, x_repr = quantize_to_fp4_raw(x_scaled)
    else:
        x_quant = x_scaled.to(torch.float8_e4m3fn)
        x_repr = x_quant.float()
    x_ref = (x_repr.reshape(sf_shape) * sf[..., None]).reshape_as(x)
    sf_packed = sf_byte[..., ::2].to(torch.int16) | (sf_byte[..., 1::2].to(torch.int16) << 8)
    return x_quant, sf_packed.contiguous(), x_ref


def ref_mqa_logits(q, kv, weights, ks, ke):
    seq_len, num_heads, _ = q.shape
    seq_len_kv = kv.shape[0]
    score = torch.einsum("qhd,kd->qhk", q.float(), kv.float())
    out = (score.relu() * weights.float()[:, :, None]).sum(dim=1)
    pos = torch.arange(seq_len_kv)
    mask = (pos[None, :] >= ks[:, None]) & (pos[None, :] < ke[:, None])
    return out.masked_fill(~mask, float("-inf"))


def test_mqa_logits(use_fp4):
    name = "fp4" if use_fp4 else "fp8"
    seq_len, seq_len_kv, num_heads, head_dim = 9, 512, 8, 128
    q_bf16 = torch.randn(seq_len, num_heads, head_dim, dtype=torch.bfloat16)
    kv_bf16 = torch.randn(seq_len_kv, head_dim, dtype=torch.bfloat16)
    weights = torch.randn(seq_len, num_heads, dtype=torch.bfloat16)
    a = torch.randint(0, seq_len_kv + 1, (seq_len,), dtype=torch.int32)
    b = torch.randint(0, seq_len_kv + 1, (seq_len,), dtype=torch.int32)
    ks, ke = torch.minimum(a, b), torch.maximum(a, b)
    max_seqlen_k = max(1, int((ke - ks).max()))

    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
    ref = ref_mqa_logits(q_ref, kv_ref, weights, ks, ke).to(torch.bfloat16)

    out = deep_gemm.fp8_fp4_mqa_logits((q, q_sf), (kv, kv_sf), weights, ks, ke, max_seqlen_k)
    check(f"mqa logits {name} dtype/stride", out.dtype == torch.bfloat16
          and out.shape == (seq_len, max_seqlen_k)
          and out.stride(0) * out.element_size() % 512 == 0 and out.stride(1) == 1)
    dense = torch.full((seq_len, seq_len_kv), float("-inf"), dtype=torch.bfloat16)
    for i, (s, e) in enumerate(zip(ks.tolist(), ke.tolist())):
        dense[i, s:e] = out[i, : e - s]
    neg = torch.isinf(ref.float()) & (ref.float() < 0)
    check(f"mqa logits {name} -inf mask", torch.equal(torch.isinf(dense.float()) & (dense.float() < 0), neg))
    diff = calc_diff(dense.float().masked_fill(neg, 0), ref.float().masked_fill(neg, 0))
    check(f"mqa logits {name} values", diff < 5e-4, f"{diff=}")
    return q, q_sf, q_ref, kv, kv_sf, kv_ref, weights, ks, ke


def make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4):
    data_bytes = 64 if use_fp4 else 128
    sf_bytes = 4
    token_bytes = data_bytes + sf_bytes
    data = kv.contiguous().view(torch.uint8).reshape(num_pages, page_kv * data_bytes)
    sf = kv_sf.contiguous().view(torch.uint8).reshape(num_pages, page_kv * sf_bytes)
    fused = torch.cat([data, sf], dim=1)
    return fused.reshape(num_pages, page_kv, 1, token_bytes)


def ref_paged_mqa_logits(q, kv, weights, context_lens, block_table, max_context_len):
    num_q = q.shape[0]
    page_size = kv.shape[1]
    out = torch.zeros(num_q, max_context_len)
    for q_idx, (context_len,) in enumerate(context_lens.tolist()):
        pages = ceil_div(context_len, page_size)
        keys = kv[block_table[q_idx, :pages].long()].flatten(0, 1).float()
        scores = torch.einsum("hd,kd->hk", q[q_idx, 0].float(), keys)
        logits = (scores.relu() * weights[q_idx, :, None].float()).sum(0)
        out[q_idx, :context_len] = logits[:context_len]
    return out


def test_paged_mqa_logits(use_fp4):
    name = "fp4" if use_fp4 else "fp8"
    batch, num_heads, head_dim, page_kv = 3, 8, 128, 64
    context_lens_host = [70, 128, 33]
    pages_per_request = [ceil_div(c, page_kv) for c in context_lens_host]
    num_pages = sum(pages_per_request)
    q_bf16 = torch.randn(batch, 1, num_heads, head_dim, dtype=torch.bfloat16)
    kv_bf16 = torch.randn(num_pages, page_kv, head_dim, dtype=torch.bfloat16)
    weights = torch.randn(batch, num_heads, dtype=torch.bfloat16)
    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
    fused = make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4)

    perm = torch.randperm(num_pages, dtype=torch.int32)
    max_pages = max(pages_per_request)
    block_table = torch.zeros(batch, max_pages, dtype=torch.int32)
    offset = 0
    for i, count in enumerate(pages_per_request):
        block_table[i, :count] = perm[offset:offset + count]
        offset += count
    context_lens = torch.tensor(context_lens_host, dtype=torch.int32).view(-1, 1)
    indices = torch.arange(batch, dtype=torch.int32)
    max_context_len = max(context_lens_host)

    # Reference uses the (permuted) page order through the block table.
    ref = ref_paged_mqa_logits(q_ref, kv_ref, weights, context_lens, block_table, max_context_len)
    meta = deep_gemm.get_paged_mqa_logits_metadata(context_lens, num_heads, indices)
    out = deep_gemm.fp8_fp4_paged_mqa_logits(
        (q, q_sf), fused, weights, context_lens, block_table, meta, max_context_len, indices, False
    )
    valid = torch.arange(max_context_len)[None] < context_lens.view(-1, 1)
    diff = calc_diff(out.float().masked_fill(~valid, 0), ref.masked_fill(~valid, 0))
    check(f"paged mqa logits {name}", diff < 5e-4, f"{diff=}")
    check(f"paged mqa logits {name} stride", out.stride(0) * out.element_size() % 512 == 0 and out.stride(1) == 1)
    return (q, q_sf), fused, weights, context_lens, block_table, indices, max_context_len, out


def test_sparse_mqa_logits(use_fp4):
    """Sparse logits must match the dense logits at the selected tokens."""
    name = "fp4" if use_fp4 else "fp8"
    num_q, num_kv, num_heads, head_dim = 6, 512, 8, 128
    sparse_block_kv, capacity = 8, 16
    q_bf16 = torch.randn(num_q, num_heads, head_dim, dtype=torch.bfloat16)
    kv_bf16 = torch.randn(num_kv, head_dim, dtype=torch.bfloat16)
    weights = torch.randn(num_q, num_heads, dtype=torch.bfloat16)
    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)

    rng = random.Random(1)
    starts, ends, blocks = [], [], []
    for _ in range(num_q):
        start = rng.randrange(0, 4) * sparse_block_kv + rng.randrange(0, sparse_block_kv)
        end = rng.randrange(start + 1, num_kv)
        first_block = start // sparse_block_kv
        available = ceil_div(max(0, end - start), sparse_block_kv)
        count = min(available, capacity)
        chosen = sorted(rng.sample(range(first_block, first_block + available), count))
        chosen += [chosen[-1] if chosen else 0] * (capacity - count)
        starts.append(start)
        ends.append(end)
        blocks.append(chosen)
    ks = torch.tensor(starts, dtype=torch.int32)
    ke = torch.tensor(ends, dtype=torch.int32)
    sparse_blocks = torch.tensor(blocks, dtype=torch.int32)

    meta = deep_gemm.get_sparse_mqa_logits_metadata(ks, ke, num_kv, sparse_blocks, q.dtype, sparse_block_kv)
    sparse = deep_gemm.fp8_fp4_sparse_mqa_logits((q, q_sf), (kv, kv_sf), weights, meta, capacity, sparse_block_kv)
    check(f"sparse mqa {name} shape/stride",
          sparse.shape == (num_q, capacity * sparse_block_kv) and sparse.dtype == torch.bfloat16
          and sparse.stride(0) * sparse.element_size() % 512 == 0)

    full = deep_gemm.fp8_fp4_mqa_logits(
        (q, q_sf), (kv, kv_sf), weights, ks, ke, max(1, int((ke - ks).max()))
    )
    offsets = torch.arange(sparse_block_kv, dtype=torch.int32)
    token_indices = (sparse_blocks.unsqueeze(-1) * sparse_block_kv + (ks % sparse_block_kv)[:, None, None] + offsets).flatten(1).long()
    num_blocks = torch.tensor([min(capacity, ceil_div(max(0, e - s), sparse_block_kv)) for s, e in zip(starts, ends)])
    valid = torch.arange(token_indices.shape[1])[None] < (num_blocks * sparse_block_kv)[:, None]
    valid &= token_indices >= ks[:, None]
    valid &= token_indices < ke[:, None]
    full_cols = (token_indices - ks[:, None]).clamp(0, full.shape[1] - 1)
    check(f"sparse mqa {name} == dense logits",
          torch.equal(sparse[valid], full.gather(1, full_cols)[valid]))

    # Determinism / reuse
    sparse2 = deep_gemm.fp8_fp4_sparse_mqa_logits((q, q_sf), (kv, kv_sf), weights, meta, capacity, sparse_block_kv)
    check(f"sparse mqa {name} repeatable", torch.equal(sparse[valid], sparse2[valid]))


def test_paged_sparse_mqa_logits(use_fp4):
    name = "fp4" if use_fp4 else "fp8"
    num_heads, head_dim, page_kv = 8, 128, 64
    sparse_block_kv, capacity = 8, 16
    context_lens_host = [70, 128, 33]
    batch = len(context_lens_host)
    pages_per_request = [ceil_div(c, page_kv) for c in context_lens_host]
    num_pages = sum(pages_per_request)
    q_bf16 = torch.randn(batch, 1, num_heads, head_dim, dtype=torch.bfloat16)
    kv_bf16 = torch.randn(num_pages, page_kv, head_dim, dtype=torch.bfloat16)
    weights = torch.randn(batch, num_heads, dtype=torch.bfloat16)
    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
    fused = make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4)

    perm = torch.randperm(num_pages, dtype=torch.int32)
    max_pages = max(pages_per_request)
    block_table = torch.zeros(batch, max_pages, dtype=torch.int32)
    offset = 0
    for i, count in enumerate(pages_per_request):
        block_table[i, :count] = perm[offset:offset + count]
        offset += count
    context_lens = torch.tensor(context_lens_host, dtype=torch.int32)
    indices = torch.arange(batch, dtype=torch.int32)

    rng = random.Random(2)
    blocks = []
    for context_len in context_lens_host:
        available = ceil_div(context_len, sparse_block_kv)
        count = min(available, capacity)
        chosen = sorted(rng.sample(range(available), count))
        blocks.append(chosen + [chosen[-1]] * (capacity - count))
    sparse_blocks = torch.tensor(blocks, dtype=torch.int32)

    meta = deep_gemm.get_paged_sparse_mqa_logits_metadata(
        context_lens, block_table, indices, page_kv, sparse_blocks, q.dtype, sparse_block_kv
    )
    sparse = deep_gemm.fp8_fp4_paged_sparse_mqa_logits((q, q_sf), fused, weights, meta, capacity, sparse_block_kv)
    check(f"paged sparse {name} shape/stride",
          sparse.shape == (batch, capacity * sparse_block_kv) and sparse.dtype == torch.bfloat16
          and sparse.stride(0) * sparse.element_size() % 512 == 0)

    dense_meta = deep_gemm.get_paged_mqa_logits_metadata(context_lens.view(-1, 1), num_heads, indices)
    full = deep_gemm.fp8_fp4_paged_mqa_logits(
        (q, q_sf), fused, weights, context_lens.view(-1, 1), block_table, dense_meta,
        max(context_lens_host), indices, False
    )
    offsets = torch.arange(sparse_block_kv, dtype=torch.int32)
    token_indices = (sparse_blocks.unsqueeze(-1) * sparse_block_kv + offsets).flatten(1).long()
    num_blocks = torch.tensor([min(capacity, ceil_div(c, sparse_block_kv)) for c in context_lens_host])
    valid = torch.arange(token_indices.shape[1])[None] < (num_blocks * sparse_block_kv)[:, None]
    valid &= token_indices < context_lens[:, None]
    full_cols = token_indices.clamp(0, full.shape[1] - 1)
    check(f"paged sparse {name} == dense logits",
          torch.equal(sparse[valid], full.gather(1, full_cols)[valid]))

    # REUSE: the KV snapshot taken on the first call survives a KV overwrite.
    first = sparse.clone()
    fused.zero_()
    again = deep_gemm.fp8_fp4_paged_sparse_mqa_logits((q, q_sf), fused, weights, meta, capacity, sparse_block_kv)
    check(f"paged sparse {name} reuses the KV snapshot", torch.equal(first[valid], again[valid]))


# --------------------------------------------------------------------------- #
# 6. hyperconnection prenorm
# --------------------------------------------------------------------------- #

def test_hc_prenorm():
    print("test_hc_prenorm")
    m, n, k = 65, 24, 256
    a = torch.randn(m, k, dtype=torch.bfloat16)
    b = torch.randn(n, k, dtype=torch.float32)
    d = torch.empty(m, n, dtype=torch.float32)
    sqr_sum = torch.empty(m, dtype=torch.float32)
    deep_gemm.tf32_hc_prenorm_gemm(a, b, d, sqr_sum)
    ref_d = a.float() @ b.T
    ref_sqr = a.float().square().sum(-1)
    check("tf32_hc_prenorm_gemm d", calc_diff(d, ref_d) < 1e-6, calc_diff(d, ref_d))
    check("tf32_hc_prenorm_gemm sqr_sum", calc_diff(sqr_sum, ref_sqr) < 1e-6)
    d_split = torch.empty(1, m, n, dtype=torch.float32)
    sqr_split = torch.empty(1, m, dtype=torch.float32)
    deep_gemm.tf32_hc_prenorm_gemm(a, b, d_split, sqr_split, num_splits=1)
    check("tf32_hc_prenorm_gemm num_splits=1", torch.equal(d_split[0], d) and torch.equal(sqr_split[0], sqr_sum))


# --------------------------------------------------------------------------- #
# 7. API surface
# --------------------------------------------------------------------------- #

def test_api_surface():
    print("test_api_surface")
    required = [
        "SymmBuffer", "bf16_gemm_nt", "bf16_mega_gate", "bf16_mega_moe", "cublaslt_gemm_nt",
        "fp8_einsum", "fp8_fp4_gemm_nt", "fp8_fp4_mega_moe", "fp8_fp4_mqa_logits",
        "fp8_fp4_paged_mqa_logits", "fp8_fp4_paged_sparse_mqa_logits", "fp8_fp4_sparse_mqa_logits",
        "get_block_m_for_mega_moe", "get_mk_alignment_for_contiguous_layout", "get_num_sms",
        "get_paged_mqa_logits_metadata", "get_paged_sparse_mqa_logits_metadata",
        "get_sparse_mqa_logits_metadata", "get_theoretical_mk_alignment_for_contiguous_layout",
        "get_tma_aligned_size", "m_grouped_bf16_gemm_nt_contiguous",
        "m_grouped_fp8_fp4_gemm_nt_contiguous", "mega_mhc", "set_mk_alignment_for_contiguous_layout",
        "set_num_sms", "set_pdl", "transform_sf_into_required_layout",
        "transform_weights_for_mega_moe", "use_deterministic_algorithms", "tf32_hc_prenorm_gemm",
    ]
    missing = [name for name in required if not hasattr(deep_gemm, name)]
    check("all required symbols exported", not missing, f"missing={missing}")
    for name in ("bf16_mega_gate", "bf16_mega_moe", "fp8_fp4_mega_moe", "mega_mhc",
                 "transform_weights_for_mega_moe", "get_block_m_for_mega_moe"):
        try:
            getattr(deep_gemm, name)()
            check(f"{name} raises NotImplementedError", False)
        except NotImplementedError:
            check(f"{name} raises NotImplementedError", True)
        except Exception as exc:  # noqa: BLE001
            check(f"{name} raises NotImplementedError", False, repr(exc))
    prev = deep_gemm.get_num_sms()
    check("default num_sms == 8 (tile_kernels.get_device_num_sms)", prev == 8, f"{prev=}")
    deep_gemm.set_num_sms(prev + 4)
    check("set/get_num_sms", deep_gemm.get_num_sms() == prev + 4)
    deep_gemm.set_num_sms(prev)
    deep_gemm.use_deterministic_algorithms(True)
    check("use_deterministic_algorithms", deep_gemm.get_deterministic_algorithms() is True)
    deep_gemm.use_deterministic_algorithms(False)
    deep_gemm.set_pdl(True)

    class _Group:
        def size(self):
            return 2

        def rank(self):
            return 0

    buf = deep_gemm.SymmBuffer(
        _Group(), num_experts=8, num_max_tokens_per_rank=16, num_topk=4, hidden=256,
        intermediate_hidden=128, num_shared_experts=1, mma_type="fp8xfp4", activation="swiglu",
        device="cpu",
    )
    check("SymmBuffer views", buf.x.shape == (16, 256) and buf.x_sf.shape == (16, 4)
          and buf.topk_idx.shape == (16, 4) and buf.num_bytes > 0)


if __name__ == "__main__":
    test_sf_layout()
    test_bf16_gemm()
    test_fp8_fp4_gemm()
    test_m_grouped()
    test_einsum()
    print("test_mqa_logits")
    test_mqa_logits(False)
    test_mqa_logits(True)
    print("test_paged_mqa_logits")
    test_paged_mqa_logits(False)
    test_paged_mqa_logits(True)
    print("test_sparse_mqa_logits")
    test_sparse_mqa_logits(False)
    test_sparse_mqa_logits(True)
    print("test_paged_sparse_mqa_logits")
    test_paged_sparse_mqa_logits(False)
    test_paged_sparse_mqa_logits(True)
    test_hc_prenorm()
    test_api_surface()
    print()
    if FAILURES:
        print(f"FAILED ({len(FAILURES)}): {FAILURES}")
        sys.exit(1)
    print("ALL TESTS PASSED")