perf(jit_kernel/deepseek_v4): optimize paged_mqa_metadata (#25855)

Co-authored-by: SII-yangdian <yangdian@sii.edu.cn>

benchmark/kernels/bench_paged_mqa_metadata.py

@@ -0,0 +1,57 @@
+"""Benchmark paged_mqa_metadata JIT kernel.
+
+Reports per-shape median latency in µs via ``marker.do_bench`` (CUDA-graph
+timing).
+
+Shape axes:
+  - ``bs``: dense sweep from single-request decode (1) to large multi-block
+    batch (32768). Covers the three internal dispatch paths
+    (tiny ``bs<=64`` / small ``bs<=2048`` / multi-block ``bs>2048``).
+  - ``max_ctx``: two extremes (2048, 32768). The kernel is value-invariant
+    (cost is O(bs) regardless of seq_lens values); sweeping both bookends
+    makes that empirically visible.
+
+Constants: ``num_sm`` queried from the active GPU; ``page_size = 64``.
+
+Local run:
+    python benchmark/kernels/bench_paged_mqa_metadata.py
+"""
+
+import torch
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.ops.attention.dsv4 import get_paged_mqa_logits_metadata
+
+NUM_SM = (
+    torch.cuda.get_device_properties(0).multi_processor_count
+    if torch.cuda.is_available()
+    else 132
+)
+PAGE_SIZE = 64
+DEVICE = "cuda"
+
+
+def _make_seq_lens(bs: int, max_ctx: int, seed: int = 0) -> torch.Tensor:
+    g = torch.Generator(device=DEVICE).manual_seed(seed)
+    return torch.randint(
+        1, max_ctx + 1, (bs,), dtype=torch.int32, device=DEVICE, generator=g
+    )
+
+
+@marker.parametrize(
+    "bs",
+    [1, 8, 16, 32, 64, 128, 256, 384, 512, 1024, 2048, 4096, 8192, 16384, 32768],
+    [128, 2048],
+)
+@marker.parametrize("max_ctx", [2048, 32768], [8192])
+@marker.benchmark("impl", ["jit"])
+def benchmark(bs: int, max_ctx: int, impl: str):
+    seq_lens = _make_seq_lens(bs, max_ctx)
+    return marker.do_bench(
+        get_paged_mqa_logits_metadata,
+        input_args=(seq_lens, PAGE_SIZE, NUM_SM),
+    )
+
+
+if __name__ == "__main__":
+    benchmark.run()