DeepInfer/scripts/sparse_mla_port/ab_decode.py

"""Decode-path A/B: does the vendored kernel serve the decode call shape at all?

Decode differs from prefill in a way that matters for this kernel. Prefill passes one KV buffer
shared by every query and lets the indices pick per-query top-k. Decode gathers first, so each query
owns a *different* KV window. The vendored kernel's `kv` operand is 3-D with no batch axis, so it
cannot be handed a different KV per query.

It does not need to. Concatenate every query's window into one KV buffer and give query q the slice
of index space `[q*W, q*W+W)`. That is one launch, and the buffer is the tensor decode already
materialises -- no extra memory.

This test exercises that adapter against the current TileLang call, end to end from an fp8 paged
cache in the V41 layout (512 fp8 values + 16 UE8M0 scale bytes per token), including the SWA + full
two-cache concatenation that `decode_mla_mqa_sparse` supports.
"""

import sys

import torch

from deepinfer.h20_ops.attn_sm90 import (
    CACHE_ELEMENTS,
    CACHE_ELEMENTS_PER_SF,
    UE8M0_BIAS,
    _gather_dequant,
)
from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn
from deepinfer.h20_ops.vendor.flashmla_sparse import flash_mla_sparse_fwd

DEV = "cuda"
HQ = 64
D = 512
FAILED = []


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 in a narrow band around the bias so scales stay near 1.
    exponents = torch.randint(-3, 4, (rows, n_sf), device=DEV, dtype=torch.int32)
    sf = (exponents + UE8M0_BIAS).to(torch.uint8)
    packed = torch.cat([fp8, sf], dim=1)
    return packed.view(num_blocks, block_size, 1, -1)


def compare(tag: str, a: torch.Tensor, b: torch.Tensor) -> None:
    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
    ok = diff.max().item() <= tol and bool(torch.isfinite(bf).all())
    print(f"  {'ok  ' if ok else 'FAIL'} {tag}: max|ref| {scale:.4f}  max diff {diff.max().item():.3e}  tol {tol:.3e}")
    if not ok:
        FAILED.append(tag)


def tl_decode(q4: torch.Tensor, kv4: torch.Tensor, scale: float) -> torch.Tensor:
    """What `decode_mla_mqa_sparse` does today: one batch row per query, identity indices."""
    total_q, window, _, hd = kv4.shape
    identity = torch.arange(window, dtype=torch.int32, device=q4.device).view(1, 1, 1, -1)
    indices = identity.expand(total_q, 1, 1, window).contiguous()
    out, _ = sparse_mla_fwd_attn(q4, kv4, indices, sm_scale=scale, is_causal=False)
    return out.view(total_q, HQ, hd)


def triton_decode(q4: torch.Tensor, kv4: torch.Tensor, scale: float) -> torch.Tensor:
    """The proposed adapter: one shared KV, per-query offset index window."""
    total_q, window, _, hd = kv4.shape
    flat_kv = kv4.reshape(total_q * window, 1, hd).contiguous()
    offsets = torch.arange(window, dtype=torch.int32, device=q4.device).view(1, window)
    base = (torch.arange(total_q, dtype=torch.int32, device=q4.device) * window).view(total_q, 1)
    indices = (base + offsets).view(total_q, 1, window).contiguous()
    out, _, _ = flash_mla_sparse_fwd(q4.view(total_q, HQ, hd).contiguous(), flat_kv, indices, sm_scale=scale, d_v=hd)
    return out


def run_case(tag: str, total_q: int, topk: int, seed: int, two_caches: bool) -> None:
    print(f"\n[{tag}] total_q={total_q} window={topk} two_caches={two_caches}")
    scale = D**-0.5
    torch.manual_seed(seed)
    q = (torch.randn((total_q, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()

    cache = make_cache(8, 64, seed)  # 512 rows
    n_rows = 8 * 64
    idx = torch.randint(0, n_rows, (total_q, topk), device=DEV, dtype=torch.int32)

    gathered = _gather_dequant(cache, idx, None, D)
    if two_caches:
        cache2 = make_cache(8, 64, seed + 1000)
        idx2 = torch.randint(0, n_rows, (total_q, topk // 2), device=DEV, dtype=torch.int32)
        extra = _gather_dequant(cache2, idx2, None, D)
        gathered = torch.cat([gathered, extra], dim=1)

    window = gathered.shape[1]
    kv4 = gathered.view(total_q, window, 1, D)
    q4 = q.view(total_q, 1, HQ, D)

    tl_out = tl_decode(q4, kv4, scale)
    torch.cuda.synchronize()
    tri_out = triton_decode(q4, kv4, scale)
    torch.cuda.synchronize()
    compare("decode out", tl_out, tri_out)

    # The gathered KV itself must be identical for the two paths, or the comparison would be
    # hiding a layout bug behind a numeric one.
    if not torch.equal(gathered.view(total_q, window, D), kv4.view(total_q, window, D)):
        print("  FAIL gathered kv changed shape semantics")
        FAILED.append("gather")


CASES = [
    ("decode, topk 512", 32, 512, 0, False),
    ("decode, topk 256", 64, 256, 1, False),
    ("decode, single query", 1, 512, 2, False),
    ("decode, SWA + full (two caches)", 16, 512, 3, True),
    ("decode, topk 128 (short window)", 128, 128, 4, False),
]

if __name__ == "__main__":
    for c in CASES:
        run_case(*c)
    print()
    if FAILED:
        print(f"MISMATCH {len(FAILED)}: {FAILED}")
        sys.exit(1)
    print(f"ALL_MATCH ({len(CASES)} cases)")