From 9d65bdba95c0b3499dd837995d93f0a7c4e6092f Mon Sep 17 00:00:00 2001
From: leavelet <leavelet@163.com>
Date: Fri, 12 Jun 2026 08:27:19 +0000
Subject: [PATCH] Measure mqa-logits chunk pipelining: not worth it (S2b
 closed)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The bench gains a pipelined mode (topk_N event-gated on a side stream,
overlapping logits_{N+1}).  Byte-equality validation was dropped after
establishing the kernel chain is run-to-run nondeterministic even
serial-vs-serial (near-equal fp32 selection).

Result on g0033 H200: pipelining recovers only 0.3-8.9% where small
chunks cost +15-46% — fp8_mqa_logits and fast_topk_transform_fused are
both SM-saturating, so concurrent streams timeshare instead of
overlapping; the small-chunk penalty is small-M GEMM inefficiency.  No
production pipeline path; the serial loop at CHUNK_MAX_GB=2 stands.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 test/manual/bench_mqa_logits_chunking.py | 93 +++++++++++++++++++++++-
 1 file changed, 92 insertions(+), 1 deletion(-)

diff --git a/test/manual/bench_mqa_logits_chunking.py b/test/manual/bench_mqa_logits_chunking.py
index 80ab5d406..933e5789a 100644
--- a/test/manual/bench_mqa_logits_chunking.py
+++ b/test/manual/bench_mqa_logits_chunking.py
@@ -65,6 +65,73 @@ def _build(prefix: int, extend: int, device):
     )
 
 
+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
@@ -127,13 +194,37 @@ def main() -> None:
                 _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)"
+                f"({(ms / ref_ms - 1) * 100:+6.1f}% vs one-shot)" + pipe_str
             )