DeepInfer/scripts/sparse_mla_port/test_wiring.py

"""Unit tests for the backend selection and for `attn_sm90`'s two entry points.

Two different things are under test, and they fail differently:

* the *dispatcher* (`h20_ops.sparse_mla`) -- a wrong choice means the wrong kernel runs, which no
  numerical test downstream could attribute correctly;
* the *entry points* -- both must give the same answer whichever backend is selected, including the
  work around the kernel (`_mask_beyond_length`, `_apply_sink`, the fp8 gather-and-dequant), not
  just the kernel call.

The "no TileLang installed" case runs in a subprocess: `tilelang` is importable here, and the thing
being tested is what happens when it is not. `scripts/simulate_no_tilelang` provides that, and the
engine's own end-to-end check uses the same trick.

Run: PYTHONPATH=. python scripts/sparse_mla_port/test_wiring.py
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
from collections.abc import Callable
from typing import Any

import torch

from deepinfer.h20_ops import sparse_mla, sparse_mla_triton
from deepinfer.h20_ops.attn_sm90 import (
    CACHE_ELEMENTS,
    CACHE_ELEMENTS_PER_SF,
    UE8M0_BIAS,
    decode_mla_mqa_sparse,
    prefill_attn_varlen_sparse,
)

DEV = "cuda"
HQ = 64
D = 512
FAILED: list[str] = []


def check(tag: str, cond: bool, detail: str = "") -> None:
    print(f"  {'ok  ' if cond else 'FAIL'} {tag}{'  ' + detail if detail else ''}")
    if not cond:
        FAILED.append(tag)


def close(tag: str, a: torch.Tensor, b: torch.Tensor, ulp: float = 2.0) -> None:
    af, bf = a.float(), b.float()
    if af.shape != bf.shape:
        check(tag, False, f"shape {tuple(af.shape)} vs {tuple(bf.shape)}")
        return
    diff = (af - bf).abs().max().item()
    scale = af.abs().max().item()
    tol = ulp * 2.0**-8 * max(scale, 1e-9)
    check(tag, diff <= tol and bool(torch.isfinite(bf).all()), f"max diff {diff:.3e} tol {tol:.3e}")


def make_cache(num_blocks: int, block_size: int, seed: int) -> torch.Tensor:
    """A V41-layout fp8 cache: 512 e4m3 values then 16 UE8M0 exponent bytes, per token."""
    torch.manual_seed(seed)
    rows = num_blocks * block_size
    values = torch.randn((rows, CACHE_ELEMENTS), device=DEV, dtype=torch.float32) * 0.5
    fp8 = values.to(torch.float8_e4m3fn).view(torch.uint8)
    n_sf = CACHE_ELEMENTS // CACHE_ELEMENTS_PER_SF
    exponents = torch.randint(-3, 4, (rows, n_sf), device=DEV, dtype=torch.int32)
    packed = torch.cat([fp8, (exponents + UE8M0_BIAS).to(torch.uint8)], dim=1)
    return packed.view(num_blocks, block_size, 1, -1)


def run_with_kernel(
    kernel_module: Any,
    entry: Callable[..., Any],
    args: tuple,
    kwargs: dict[str, Any],
) -> Any:
    """Run `entry` with its module-level kernel swapped to another backend.

    `attn_sm90` binds the kernel by name at module scope, so replacing that global is how the other
    backend gets exercised without an environment change and a second process. Restored in a
    `finally` so a failure here cannot corrupt the following test.
    """
    original = entry.__globals__["sparse_mla_fwd_attn"]
    entry.__globals__["sparse_mla_fwd_attn"] = kernel_module.sparse_mla_fwd_attn
    try:
        return entry(*args, **kwargs)
    finally:
        entry.__globals__["sparse_mla_fwd_attn"] = original


def test_selection_logic() -> None:
    """Explicit values win; a bad value is refused rather than defaulted."""
    print("\n[selection: explicit override]")

    os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "triton"
    sparse_mla.reset()
    check("triton forced", sparse_mla.backend() == "triton")

    os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "tilelang"
    sparse_mla.reset()
    check("tilelang forced", sparse_mla.backend() == "tilelang")

    os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "nonsense"
    sparse_mla.reset()
    try:
        sparse_mla.backend()
        check("bad value raises", False)
    except RuntimeError as exc:
        check("bad value raises", "not recognised" in str(exc))

    os.environ.pop("DEEPINFER_SPARSE_MLA_BACKEND", None)
    sparse_mla.reset()
    check("auto picks tilelang where it imports", sparse_mla.backend() == "tilelang")


def test_selection_without_tilelang() -> None:
    """Auto must fall back when TileLang is genuinely absent; forcing it must explain why."""
    print("\n[selection: TileLang unimportable]")
    code = """
import json, os
from deepinfer.h20_ops import sparse_mla
out = {"backend": sparse_mla.backend()}
try:
    os.environ["DEEPINFER_SPARSE_MLA_BACKEND"] = "tilelang"
    sparse_mla.reset()
    sparse_mla.backend()
    out["forced"] = "no error"
except RuntimeError as exc:
    out["forced"] = "raised"
    out["message_shown"] = "not importable" in str(exc)
print(json.dumps(out))
"""
    root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    env = dict(os.environ)
    env["PYTHONPATH"] = os.pathsep.join([os.path.join(root, "scripts", "simulate_no_tilelang"), root])
    env["DEEPINFER_SIMULATE_NO_TILELANG"] = "1"
    env.pop("DEEPINFER_SPARSE_MLA_BACKEND", None)
    proc = subprocess.run(
        [sys.executable, "-c", code],
        capture_output=True,
        text=True,
        env=env,
        cwd=root,
        timeout=600,
        check=False,
    )
    payload = None
    for line in reversed(proc.stdout.splitlines()):
        if line.strip().startswith("{"):
            payload = json.loads(line)
            break
    if payload is None:
        tail = proc.stderr.strip().splitlines()[-1] if proc.stderr.strip() else "no output"
        check("subprocess ran", False, tail)
        return
    check("auto falls back to triton", payload["backend"] == "triton", payload["backend"])
    check("forcing tilelang raises", payload.get("forced") == "raised")
    check("message explains why", payload.get("message_shown", False))


def test_prefill_entry_point() -> None:
    """`prefill_attn_varlen_sparse`, both backends, with and without the optional arguments.

    `index_length` and `attn_sink` are covered because they are the two places the entry point does
    work around the kernel -- `_mask_beyond_length` before it, `_apply_sink` after -- and a backend
    swap that got either wrong would look fine without them.
    """
    print("\n[attn_sm90.prefill_attn_varlen_sparse]")
    torch.manual_seed(10)
    total_q, total_k, topk = 48, 1536, 512
    q = (torch.randn((total_q, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    kv = (torch.randn((total_k, 1, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    idx = torch.stack([torch.randperm(total_k, device=DEV)[:topk] for _ in range(total_q)])
    idx = idx.to(torch.int32).contiguous()
    length = torch.randint(64, topk, (total_q,), device=DEV, dtype=torch.int32)
    sink = torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5
    kwargs = {"attn_sink": sink, "index_length": length}
    args = (q, kv, idx, D**-0.5)

    tl_out, tl_lse = prefill_attn_varlen_sparse(*args, **kwargs)
    tri_out, tri_lse = run_with_kernel(sparse_mla_triton, prefill_attn_varlen_sparse, args, kwargs)
    torch.cuda.synchronize()
    close("prefill out", tl_out, tri_out)
    close("prefill lse", tl_lse, tri_lse)

    tl_out2, tl_lse2 = prefill_attn_varlen_sparse(*args)
    tri_out2, tri_lse2 = run_with_kernel(sparse_mla_triton, prefill_attn_varlen_sparse, args, {})
    torch.cuda.synchronize()
    close("prefill out (no sink/length)", tl_out2, tri_out2)
    close("prefill lse (no sink/length)", tl_lse2, tri_lse2)


def test_decode_entry_point() -> None:
    """`decode_mla_mqa_sparse`, both backends, single cache and SWA + full two-cache."""
    print("\n[attn_sm90.decode_mla_mqa_sparse]")
    torch.manual_seed(11)
    total_q, topk = 24, 512
    q = (torch.randn((total_q, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
    cache = make_cache(8, 64, 11)
    idx = torch.randint(0, 512, (total_q, topk), device=DEV, dtype=torch.int32)
    sink = torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5
    args = (q, cache, D, idx, None, D**-0.5)

    tl_out, tl_lse = decode_mla_mqa_sparse(*args, attn_sink=sink)
    tri_out, tri_lse = run_with_kernel(sparse_mla_triton, decode_mla_mqa_sparse, args, {"attn_sink": sink})
    torch.cuda.synchronize()
    close("decode out", tl_out, tri_out)
    close("decode lse", tl_lse, tri_lse)

    cache2 = make_cache(8, 64, 111)
    idx2 = torch.randint(0, 512, (total_q, topk // 2), device=DEV, dtype=torch.int32)
    kwargs2 = {"attn_sink": sink, "key_cache_extra": cache2, "indices_in_kvcache_extra": idx2}
    tl_out2, tl_lse2 = decode_mla_mqa_sparse(*args, **kwargs2)
    tri_out2, tri_lse2 = run_with_kernel(sparse_mla_triton, decode_mla_mqa_sparse, args, kwargs2)
    torch.cuda.synchronize()
    close("decode two-cache out", tl_out2, tri_out2)
    close("decode two-cache lse", tl_lse2, tri_lse2)


TESTS = [
    test_selection_logic,
    test_selection_without_tilelang,
    test_prefill_entry_point,
    test_decode_entry_point,
]

if __name__ == "__main__":
    for test in TESTS:
        test()
    print()
    if FAILED:
        print(f"FAILED {len(FAILED)}: {FAILED}")
        sys.exit(1)
    print(f"ALL_PASSED ({len(TESTS)} tests)")