sglang/test/manual/bench_mqa_logits_chunking.py

#!/usr/bin/env python3
"""Micro-benchmark: cost of chunking the NSA indexer mqa-logits/topk loop.

专题 S2(a) (docs_internal/perf/prefill-compute-intensity-plan.md): the 18G CP
admission budget is dominated by the SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB=8
worst-case term.  Before shrinking it, measure how much per-layer indexer
time the smaller chunks actually cost — the loop is
``deep_gemm.fp8_mqa_logits`` (logits buffer = rows x kv x 4B) followed by
``sgl_kernel.fast_topk_transform_fused`` per chunk
(nsa_indexer._mqa_logits_topk_ragged_chunked), serial on one stream.

GLM-5.1 shapes: index_n_heads=32, index_head_dim=128, index_topk=2048,
attn-cp8 in-seq split (per-rank q rows = chunk/8; the rank sees the FULL
dense kv).  Single GPU — the loop is rank-local.

Run (g0033 syh-dev-new, 1 GPU):
  cd /mnt/beegfs/syh/sglang-stable && \
  PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \
  python test/manual/bench_mqa_logits_chunking.py \
  2>&1 | tee /mnt/beegfs/syh/log/bench_mqa_chunking_$(date +%Y%m%d_%H%M%S).log
"""
from __future__ import annotations

import time

import deep_gemm
import torch
from sgl_kernel import fast_topk_transform_fused

H = 32  # index_n_heads
D = 128  # index_head_dim
TOPK = 2048
BYTES_PER_ELEM = 4  # fp32 logits

# (name, prefix_tokens, chunk_extend_tokens) — global lens; per-rank q rows =
# extend/8 (cp8 in-seq split, this rank takes the LAST slice = worst-case ke).
SCENARIOS = [
    ("cold_chunk1_64K", 0, 65536),
    ("80K_tail_chunk_14K", 65536, 14464),
    ("warm_cont_150K+4K", 150_000, 4096),
    ("warm_long_150K+60K", 150_000, 60000),
]
CHUNK_GBS = [0.0, 8.0, 4.0, 2.0, 1.0, 0.5]  # 0.0 = one-shot reference


def _build(prefix: int, extend: int, device):
    n = prefix + extend  # full dense kv this rank sees
    m = max(1, extend // 8)  # per-rank q rows
    rank_start = prefix + (extend // 8) * 7  # last slice global offset
    q = torch.randn(m, H, D, device=device).clamp(-2, 2).to(torch.float8_e4m3fn)
    kv = torch.randn(n, D, device=device).clamp(-2, 2).to(torch.float8_e4m3fn)
    kv_scales = torch.rand(n, device=device) + 0.5
    weights = torch.rand(m, H, device=device)
    ks = torch.zeros(m, dtype=torch.int32, device=device)
    ke = (
        torch.arange(m, dtype=torch.int32, device=device) + rank_start + 1
    ).clamp(max=n)
    # token-granular page table for one request (slot i = token i) + paged
    # topk inputs, mirroring the PAGED topk_transform path.
    page_table_1 = torch.arange(n, dtype=torch.int32, device=device).view(1, n)
    cu_seqlens_q = torch.tensor([0, m], dtype=torch.int32, device=device)
    return dict(
        m=m, n=n, q=q, kv=kv, kv_scales=kv_scales, weights=weights,
        ks=ks, ke=ke, page_table_1=page_table_1, cu_seqlens_q=cu_seqlens_q,
    )


def _run_chunked_pipelined(s, max_rows: int, side: torch.cuda.Stream) -> torch.Tensor:
    """Same loop with topk_N on a side stream overlapping logits_{N+1}.

    Settles whether the small-chunk penalty is serialization (pipelining
    helps) or small-M GEMM inefficiency (it cannot) — both kernels are
    SM-heavy, so concurrent streams may just timeshare.
    """
    m = s["m"]
    main = torch.cuda.current_stream()
    out = None
    bounds = []
    start = 0
    while start < m:
        end = min(start + max_rows, m)
        bounds.append((start, end))
        start = end

    def _logits(b):
        return deep_gemm.fp8_mqa_logits(
            s["q"][b[0] : b[1]],
            (s["kv"], s["kv_scales"]),
            s["weights"][b[0] : b[1]],
            s["ks"][b[0] : b[1]],
            s["ke"][b[0] : b[1]],
            clean_logits=False,
        )

    def _topk(logits, b):
        cu = torch.tensor([0, b[1] - b[0]], dtype=torch.int32, device=logits.device)
        return fast_topk_transform_fused(
            score=logits,
            lengths=s["ke"][b[0] : b[1]],
            page_table_size_1=s["page_table_1"],
            cu_seqlens_q=cu,
            topk=TOPK,
            row_starts=s["ks"][b[0] : b[1]],
        )

    prev_logits, prev_b, prev_evt = None, None, None
    for b in bounds:
        if prev_logits is not None:
            # Launch topk(N-1) on the side stream gated ONLY on its own
            # logits event, then immediately launch logits(N) on main — the
            # two overlap (or timeshare SMs; that is what we are measuring).
            side.wait_event(prev_evt)
            with torch.cuda.stream(side):
                topk = _topk(prev_logits, prev_b)
                if out is None:
                    out = topk.new_full((m, topk.shape[1]), -1)
                out[prev_b[0] : prev_b[1]] = topk
            prev_logits.record_stream(side)
        logits = _logits(b)
        evt = torch.cuda.Event()
        evt.record(main)
        prev_logits, prev_b, prev_evt = logits, b, evt
    side.wait_event(prev_evt)
    with torch.cuda.stream(side):
        topk = _topk(prev_logits, prev_b)
        if out is None:
            out = topk.new_full((m, topk.shape[1]), -1)
        out[prev_b[0] : prev_b[1]] = topk
    prev_logits.record_stream(side)
    main.wait_stream(side)
    out.record_stream(main)
    return out


def _run_chunked(s, max_rows: int) -> torch.Tensor:
    m = s["m"]
    out = None
    start = 0
    while start < m:
        end = min(start + max_rows, m)
        logits = deep_gemm.fp8_mqa_logits(
            s["q"][start:end],
            (s["kv"], s["kv_scales"]),
            s["weights"][start:end],
            s["ks"][start:end],
            s["ke"][start:end],
            clean_logits=False,
        )
        cu = torch.tensor(
            [0, end - start], dtype=torch.int32, device=logits.device
        )
        topk = fast_topk_transform_fused(
            score=logits,
            lengths=s["ke"][start:end],
            page_table_size_1=s["page_table_1"],
            cu_seqlens_q=cu,
            topk=TOPK,
            row_starts=s["ks"][start:end],
        )
        if out is None:
            out = topk.new_full((m, topk.shape[1]), -1)
        out[start:end] = topk
        start = end
    return out


def main() -> None:
    device = torch.device("cuda", 0)
    torch.cuda.set_device(device)
    print(f"device={torch.cuda.get_device_name(device)} H={H} D={D} topk={TOPK}")
    for name, prefix, extend in SCENARIOS:
        s = _build(prefix, extend, device)
        m, n = s["m"], s["n"]
        print(
            f"\n=== {name}: prefix={prefix} extend={extend} "
            f"per-rank q_rows={m} kv={n} ==="
        )
        ref_ms = None
        for gb in CHUNK_GBS:
            if gb <= 0:
                max_rows = m
            else:
                max_rows = max(1, int(gb * 1e9) // (n * BYTES_PER_ELEM))
                max_rows = min(max_rows, m)
            num_chunks = -(-m // max_rows)
            peak_mb = max_rows * n * BYTES_PER_ELEM / 1e6
            # warmup
            for _ in range(3):
                _run_chunked(s, max_rows)
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            reps = 10
            for _ in range(reps):
                _run_chunked(s, max_rows)
            torch.cuda.synchronize()
            ms = (time.perf_counter() - t0) / reps * 1000
            pipe_ms = None
            if num_chunks > 1:
                side = torch.cuda.Stream()
                # NOTE: the kernel chain is run-to-run nondeterministic even
                # serial-vs-serial (near-equal fp32 selection), so byte
                # equality is meaningless here; sanity-check shape only and
                # leave correctness to the indexer's own validation paths.
                ref = _run_chunked(s, max_rows)
                got = _run_chunked_pipelined(s, max_rows, side)
                torch.cuda.synchronize()
                assert ref.shape == got.shape
                for _ in range(3):
                    _run_chunked_pipelined(s, max_rows, side)
                torch.cuda.synchronize()
                t0 = time.perf_counter()
                for _ in range(reps):
                    _run_chunked_pipelined(s, max_rows, side)
                torch.cuda.synchronize()
                pipe_ms = (time.perf_counter() - t0) / reps * 1000
            if ref_ms is None:
                ref_ms = ms
            label = "one-shot" if gb <= 0 else f"{gb:>4.1f}GB"
            pipe_str = (
                f"  pipelined {pipe_ms:8.3f} ms ({(pipe_ms / ms - 1) * 100:+6.1f}%)"
                if pipe_ms is not None
                else ""
            )
            print(
                f"  {label}: chunks={num_chunks:>3} rows/chunk={max_rows:>6} "
                f"peak_logits={peak_mb:>8.1f}MB  {ms:8.3f} ms/layer "
                f"({(ms / ref_ms - 1) * 100:+6.1f}% vs one-shot)" + pipe_str
            )


if __name__ == "__main__":
    main()