DeepInfer/scripts/sparse_mla_port/ab_vs_tilelang.py

"""FlagGems' `flash_mla_sparse_fwd` vs DeepInfer's TileLang sparse MLA, on the engine's own contract.

Why this pairing and not the DSA kernel: `attn_sm90.prefill_attn_varlen_sparse` calls the TileLang
kernel with `is_causal=False`, documented in the source -- the indexer only scores positions before
each query, so the indices carry causality, and bounding on the row index instead collapses the
layer to NaN (that bug already happened once). `flash_mla_sparse_fwd` has no causal parameter for
the same reason. The DSA kernel (`triton_sparse_mla_fwd_interface`) hardcodes causal and needs the
V3.2 576-wide split layout, which V4.1 does not have, so it is not a candidate here.

Argument correspondence, which is closer than expected:

    prefill_attn_varlen_sparse(query, kv, index_topk, softmax_scale,
                               attn_sink, index_length)
    flash_mla_sparse_fwd(q,     kv, indices,    sm_scale,
                         attn_sink=..., topk_length=...)

Both return (out, ...). `index_length` is a per-query valid count applied as -1 padding upstream;
`topk_length` is the same thing applied inside the kernel.
"""

import sys
from pathlib import Path

import torch

# The vendored copy imports `flag_gems.utils...`; its stub sits next to it, so put that on the
# path here rather than making the caller set PYTHONPATH for a layout detail.
sys.path.insert(0, str(Path(__file__).resolve().parent / "vendor"))

from vendor.fg_flashmla_sparse import flash_mla_sparse_fwd
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn

DEV = "cuda"
HQ, HKV, D, TOPK = 64, 1, 512, 512
FAILED = []


def tl_call(q3, kv3, idx4, scale):
    """DeepInfer's shape convention: (S,H,D)/(S,1,D)/(S,1,topk) -> pad a batch axis."""
    s_q = q3.shape[0]
    s_kv = kv3.shape[0]
    out, lse = sparse_mla_fwd_attn(
        q3.view(1, s_q, HQ, D),
        kv3.view(1, s_kv, 1, D),
        idx4.view(1, s_q, 1, -1),
        sm_scale=scale,
        is_causal=False,
    )
    return out.view(s_q, HQ, D), lse.view(s_q, HQ)


def compare(tag, a, b, extra=""):
    af, bf = a.float(), b.float()
    if af.shape != bf.shape:
        print(f"  FAIL {tag}: shape {tuple(af.shape)} vs {tuple(bf.shape)}")
        FAILED.append(tag)
        return
    diff = (af - bf).abs()
    scale = af.abs().max().item()
    tol = 2.0**-8 * max(scale, 1e-9) * 2
    finite = bool(torch.isfinite(bf).all())
    ok = diff.max().item() <= tol and finite
    print(
        f"  {'ok  ' if ok else 'FAIL'} {tag}: max|ref| {scale:.4f}  max diff {diff.max().item():.3e}"
        f"  mean {diff.mean().item():.3e}  tol {tol:.3e}  finite {finite}{extra}"
    )
    if not ok:
        FAILED.append(tag)


def make(SQ, SKV, seed, distinct):
    torch.manual_seed(seed)
    q = (torch.randn((SQ, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    kv = (torch.randn((SKV, HKV, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    if distinct and SKV > TOPK:
        rows = [torch.randperm(SKV, device=DEV)[:TOPK] for _ in range(SQ)]
        idx = torch.stack(rows).view(SQ, 1, TOPK).to(torch.int32)
    else:
        idx = torch.randint(0, SKV, (SQ, 1, TOPK), device=DEV, dtype=torch.int32)
    return q, kv, idx.contiguous()


def case(tag, SQ, SKV, seed, distinct=True, use_length=False, use_sink=False):
    q, kv, idx = make(SQ, SKV, seed, distinct)
    scale = D**-0.5
    print(f"\n[{tag}] SQ={SQ} SKV={SKV} topk={TOPK} distinct={distinct}"
          f"{' +topk_length' if use_length else ''}{' +attn_sink' if use_sink else ''}")

    # DeepInfer applies `index_length` by masking the tail of the index row to -1 upstream
    # (`_mask_beyond_length`), so the reference call is the same kernel with a padded index.
    idx_ref = idx
    tl_len = None
    if use_length:
        torch.manual_seed(seed + 100)
        tl_len = torch.randint(64, TOPK, (SQ,), device=DEV, dtype=torch.int32)
        idx_ref = idx.clone()
        ar = torch.arange(TOPK, device=DEV).view(1, 1, TOPK)
        idx_ref[(ar >= tl_len.view(SQ, 1, 1)).expand(SQ, 1, TOPK)] = -1
        idx_ref = idx_ref.contiguous()

    tl_out, _ = tl_call(q, kv, idx_ref, scale)
    torch.cuda.synchronize()

    sink = None
    if use_sink:
        torch.manual_seed(seed + 200)
        sink = (torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5)
        # DeepInfer's sink is folded into the output after the kernel; feed the kernel no sink and
        # compare the raw attention, then check the sink formula separately below.
    fg_out, fg_max, fg_lse = flash_mla_sparse_fwd(
        q, kv, idx_ref, sm_scale=scale, d_v=D,
        attn_sink=sink,
        topk_length=tl_len if use_length else None,
    )
    torch.cuda.synchronize()

    if use_sink:
        # Compare against the reference with the sink applied the same way the kernel documents:
        # out * exp(lse) / (exp(lse) + exp(sink)).
        tl_lse = sparse_mla_fwd_attn(
            q.view(1, SQ, HQ, D), kv.view(1, SKV, 1, D), idx_ref.view(1, SQ, 1, -1),
            sm_scale=scale, is_causal=False,
        )[1].view(SQ, HQ)
        import math
        w = torch.exp(tl_lse.float() * math.log(2))  # TileLang lse is log2 units
        # tl_out is (SQ, HQ, D), so the weight must gain a trailing axis to broadcast over D.
        s = sink.view(1, HQ, 1)
        gate = w.unsqueeze(-1) / (w.unsqueeze(-1) + torch.exp(s))
        applied = tl_out.float() * gate
        compare("sink applied", applied, fg_out)
    else:
        compare("out", tl_out, fg_out)


CASES = [
    ("prefill, distinct top-k", dict(SQ=64, SKV=2048, seed=0)),
    ("prefill, duplicate top-k", dict(SQ=32, SKV=512, seed=1, distinct=False)),
    ("decode (SQ=1)", dict(SQ=1, SKV=4096, seed=2)),
    ("odd SQ", dict(SQ=7, SKV=1024, seed=3)),
    ("short KV (< topk)", dict(SQ=16, SKV=128, seed=4, distinct=False)),
    ("with topk_length", dict(SQ=32, SKV=2048, seed=5, use_length=True)),
    ("with attn_sink", dict(SQ=16, SKV=1024, seed=6, use_sink=True)),
]

if __name__ == "__main__":
    # Run from the repo root: PYTHONPATH=. python scripts/sparse_mla_port/ab_vs_tilelang.py
    for tag, kw in CASES:
        case(tag, **kw)
    print()
    if FAILED:
        print(f"MISMATCH {len(FAILED)}: {FAILED}")
        sys.exit(1)
    print(f"ALL_MATCH ({len(CASES)} cases)")