recipes/docker-images/glm53-kindling/dev/kernel-tests/_cudagraph_replay_shapes_test.py

#!/usr/bin/env python3
"""T6: FULL-cudagraph replay checks for the two decode kernels whose plan is
built host-side: the SM90 FlashInfer MLA wrapper (_SM90State) and the
DeepGEMM paged MQA logits + per_row top-k of the sparse indexer.

vLLM captures decode graphs with a dummy batch whose seq_len equals the query
length (InputBatch.make_dummy: seq_lens = ceil(num_tokens/num_reqs), positions
zero), then replans every step and replays. This mirrors that: capture with
8 rows of lens 1..8, replay with real lens, compare against a reference.
"""
import sys
import traceback

import torch

torch.manual_seed(0)
dev = torch.device("cuda")
results = {}


def section(name):
    def deco(fn):
        def run():
            print(f"\n===== {name} =====", flush=True)
            try:
                fn()
                results[name] = "PASS"
            except Exception as e:  # noqa: BLE001
                results[name] = f"FAIL: {e}"
                traceback.print_exc()
        return run
    return deco


def t6a_impl(R, cap_lens, replays):
    from vllm import _custom_ops as ops
    from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 import _SM90State

    heads, d = 16, 512
    sm_scale = 1.0 / 16.0
    topk_width = 2176
    max_tokens = 256
    N = 2300
    bs = 64
    num_blocks = N // bs + 3
    kv_c = torch.randn(N, d, device=dev, dtype=torch.bfloat16)
    k_pe = torch.empty(N, 0, device=dev, dtype=torch.bfloat16)
    kv_cache = torch.zeros(num_blocks, bs, d, device=dev, dtype=torch.uint8)
    slots = (torch.arange(N, device=dev) + 37).to(torch.int64)
    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slots, "fp8_e4m3", torch.tensor(1.0, device=dev))
    deq = kv_cache.view(torch.float8_e4m3fn).reshape(-1, d).float()
    flat = kv_cache.view(torch.float8_e4m3fn).reshape(-1, 1, d)
    ckv, kpe = flat[..., :d], flat[..., d:]

    state = _SM90State(dev, heads, torch.float8_e4m3fn, max_tokens, topk_width,
                       kv_lora_rank=d, qk_rope_head_dim=0, sm_scale=sm_scale)
    # static input buffers (what the captured graph reads)
    q_nope = torch.randn(R, heads, d, device=dev, dtype=torch.bfloat16)
    q_pe = torch.empty(R, heads, 0, device=dev, dtype=torch.bfloat16)
    out_buf = torch.empty(R, heads, d, device=dev, dtype=torch.bfloat16)

    def set_rows(lens):
        state.kv_indices.fill_(0)
        for r, L in enumerate(lens):
            state.kv_indices[r * topk_width: r * topk_width + L] = slots[:L].to(torch.int32)
        state.plan(R, torch.tensor(lens, dtype=torch.int32))

    def reference(lens):
        ref = torch.empty(R, heads, d, device=dev)
        for r, L in enumerate(lens):
            K = deq[slots[:L]]
            s = (q_nope[r].float() @ K.T) * sm_scale
            ref[r] = torch.softmax(s, dim=-1) @ K
        return ref

    def run_once():
        out = state.wrapper.run(q_nope, q_pe, ckv, kpe, ckv_scale=1.0, kpe_scale=1.0)
        out_buf.copy_(out)

    # --- capture with the dummy lens vLLM uses ---
    set_rows(cap_lens)
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(2):
            run_once()
    torch.cuda.current_stream().wait_stream(s)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        run_once()
    torch.cuda.synchronize()

    worst = 0.0
    for lens in replays:
        q_nope.copy_(torch.randn_like(q_nope))
        set_rows(lens)
        graph.replay()
        torch.cuda.synchronize()
        ref = reference(lens)
        err = (out_buf.float() - ref).abs()
        rel = err.max().item() / max(ref.abs().max().item(), 1e-6)
        row_err = [round(v, 4) for v in err.flatten(1).max(dim=1).values.tolist()]
        nan = torch.isnan(out_buf).any().item()
        print(f"  replay lens={lens[0]}..{lens[-1]}: rel err {rel:.4g} nan={nan} per-row {row_err}")
        worst = max(worst, rel)
        assert not nan
    print(f"T6a R={R} cap={cap_lens} worst replay rel err {worst:.4g}")
    assert worst < 0.1


@section("T6a SM90 MLA wrapper: capture at dummy lens, replay at real lens")
def t6a():
    t6a_impl(8, list(range(1, 9)), [list(range(1, 9)), list(range(100, 108)), list(range(693, 701)), list(range(2044, 2052))])
    # spec-off production shape: one request, one token, dummy seq_len 1
    t6a_impl(1, [1], [[1], [100], [700], [2051], [100], [37], [2600 // 4 * 4 + 3], [5]])
    # MTP k=3 shape: one request, 4 tokens, dummy seq_lens 1..4
    t6a_impl(4, [1, 2, 3, 4], [[1, 2, 3, 4], [100, 101, 102, 103], [700, 701, 702, 703], [2049, 2050, 2051, 2052], [40, 41, 42, 43]])


@section("T6b indexer paged MQA logits + per_row topk: capture short, replay long")
def t6b():
    from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits, get_paged_mqa_logits_metadata
    from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
    from vllm import _custom_ops as ops
    from vllm.utils.platform_utils import num_compute_units

    # Indexer cache: [num_blocks, pool_block_size, 1, head_dim+4] uint8 (fp8 + fp32 scale)
    head_dim = 128
    kpool = 4
    index_heads = 32  # index_n_heads
    pool_bs = 64  # production indexer kernel page: 64 pools = 256 tokens
    num_blocks = 80
    max_pool_len = 524288 // kpool
    fp8 = torch.float8_e4m3fn
    k_pool = torch.randn(num_blocks * pool_bs, head_dim, device=dev).clamp(-4, 4)
    kv_cache = torch.zeros(num_blocks, pool_bs, 1, head_dim + 4, device=dev, dtype=torch.uint8)
    kv_cache[..., :head_dim] = k_pool.to(fp8).view(torch.uint8).reshape(num_blocks, pool_bs, 1, head_dim)
    kv_cache[..., head_dim:] = torch.ones(num_blocks * pool_bs, 1, 1, device=dev).view(torch.uint8).reshape(num_blocks, pool_bs, 1, 4)
    num_sms = num_compute_units(dev.index if dev.index is not None else 0)

    R = 8
    # per-token expanded block table (flattened path): every row -> request 0's blocks [1..num_blocks-1]
    block_table = torch.zeros(R, num_blocks, device=dev, dtype=torch.int32)
    block_table[:, : num_blocks - 1] = torch.arange(1, num_blocks, device=dev, dtype=torch.int32)
    q = torch.randn(R, 1, index_heads, head_dim, device=dev).to(fp8)
    weights = torch.rand(R, index_heads, device=dev, dtype=torch.float32)
    seq_lens = torch.zeros(R, 1, device=dev, dtype=torch.int32)  # pool-granular, refreshed
    sched_buf = torch.empty((num_sms + 1, 2), dtype=torch.int32, device=dev)
    select_k = 2048 // kpool
    pool_topk = torch.full((R, select_k), -1, dtype=torch.int32, device=dev)
    logits_holder = {}

    def set_lens(pool_lens):
        seq_lens[:, 0] = torch.tensor(pool_lens, device=dev, dtype=torch.int32)
        meta = get_paged_mqa_logits_metadata(seq_lens, pool_bs, num_sms)
        sched = sched_buf[: meta.shape[0]]
        sched[:] = meta
        return sched

    def run_once(sched):
        logits = fp8_fp4_paged_mqa_logits((q, None), kv_cache, weights, seq_lens, block_table,
                                          sched, max_model_len=max_pool_len, clean_logits=False)
        logits_holder["l"] = logits
        pool_topk.fill_(-1)
        ops.top_k_per_row_decode(logits, 1, seq_lens, pool_topk, R, logits.stride(0), logits.stride(1), select_k)

    def reference(pool_lens):
        # logits[r, p] = sum_h w[r,h] * relu(q[r,h] . k[p]) (DeepSeek indexer semantics)
        out = []
        for r, L in enumerate(pool_lens):
            blocks = block_table[r, : (L + pool_bs - 1) // pool_bs]
            K = k_pool.view(num_blocks, pool_bs, head_dim)[blocks.long()].reshape(-1, head_dim)[:L]
            Kq = K.to(fp8).float()
            s = torch.relu(q[r, 0].float() @ Kq.T)  # [heads, L]
            out.append((weights[r][:, None] * s).sum(0))
        return out

    cap_lens = [max(1, (i + 1) // kpool) for i in range(1, R + 1)]  # dummy: seq_len = query_len
    sched = set_lens(cap_lens)
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(2):
            run_once(sched)
    torch.cuda.current_stream().wait_stream(s)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        run_once(sched)
    torch.cuda.synchronize()

    bad = 0
    for tok_lens in (list(range(100, 108)), list(range(693, 701)), list(range(1500, 1508)), list(range(2600, 2608)), list(range(4600, 4608))):
        pool_lens = [t // kpool for t in tok_lens]
        # the schedule buffer is refreshed in place; its slot count must be stable
        meta = get_paged_mqa_logits_metadata(seq_lens.new_tensor([[p] for p in pool_lens]), pool_bs, num_sms)
        assert meta.shape[0] == sched.shape[0], f"schedule slot count changed {meta.shape} vs {sched.shape}"
        set_lens(pool_lens)
        graph.replay()
        torch.cuda.synchronize()
        logits = logits_holder["l"]
        refs = reference(pool_lens)
        for r, L in enumerate(pool_lens):
            got = logits[r, :L].float()
            ref = refs[r]
            nan = torch.isnan(got).any().item()
            err = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-6)
            sel = pool_topk[r][pool_topk[r] >= 0].tolist()
            if L <= select_k:
                ok_sel = set(sel) == set(range(L)) and len(sel) == L
            else:
                ok_sel = set(sel) == set(ref.topk(select_k).indices.tolist())
            if nan or err > 0.05 or not ok_sel:
                bad += 1
                miss = sorted(set(range(min(L, select_k))) - set(sel))[:8] if L <= select_k else []
                print(f"  BAD replay tok_len={tok_lens[r]} pools={L}: nan={nan} rel_err={err:.3g} "
                      f"sel_ok={ok_sel} nsel={len(sel)} missing_pools={miss}")
        print(f"  replay tok_lens {tok_lens[0]}..{tok_lens[-1]}: checked")
    # eager comparison
    pool_lens = [t // kpool for t in range(693, 701)]
    sched = set_lens(pool_lens)
    run_once(sched)
    torch.cuda.synchronize()
    refs = reference(pool_lens)
    e = max(((logits_holder["l"][r, :L].float() - refs[r]).abs().max() / refs[r].abs().max()).item() for r, L in enumerate(pool_lens))
    print(f"  eager 693..700 rel err {e:.3g}")
    print(f"T6b bad rows: {bad}")
    assert bad == 0


for fn in (t6a,):
    fn()
print("\n===== SUMMARY =====")
for k, v in results.items():
    print(f"{k}: {v}")
sys.exit(0 if all(v == "PASS" for v in results.values()) else 1)