[Perf] dsv41: Fuse request-window layout construction

benchmark/kernels/bench_dsv41_window_layout.py

@@ -0,0 +1,47 @@
+"""Measure eager request-window metadata construction on synthetic inputs.
+
+Run this file with PYTHONPATH pointing to each checkout's python directory.
+Compilation and warmup are outside the measured interval.
+"""
+
+import argparse
+import json
+import time
+
+import torch
+
+from sglang.srt.mem_cache.dsv41_request_window import window_layout
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--iterations", type=int, default=100)
+    parser.add_argument("--repeats", type=int, default=5)
+    args = parser.parse_args()
+    for rows, groups in ((4096, 1), (32768, 8)):
+        req = torch.arange(groups, device="cuda").repeat_interleave(rows // groups)
+        pos = torch.arange(8192, 8192 + rows // groups, device="cuda").repeat(groups)
+        floor = torch.full_like(pos, 8192)
+
+        def run():
+            return window_layout(
+                req, pos, zero_row=65536, floor=floor, num_groups=groups
+            )
+
+        for _ in range(20):
+            run()
+        torch.cuda.synchronize()
+        samples = []
+        for _ in range(args.repeats):
+            start = time.perf_counter()
+            for _ in range(args.iterations):
+                run()
+            torch.cuda.synchronize()
+            samples.append((time.perf_counter() - start) * 1000 / args.iterations)
+        print(
+            json.dumps({"rows": rows, "groups": groups, "wall_ms": samples}), flush=True
+        )
+
+
+if __name__ == "__main__":
+    main()

python/sglang/kernels/ops/kvcache/__init__.py

@@ -54,6 +54,7 @@ def reshape_and_cache_flash(
 # Other Triton kernels migrated into this group (from attention/mem_cache
 # triton_ops); registered for inventory. Import them from their modules.
 _TRITON_KERNELS = [
+    ("dsv41_window_layout", "build_window_layout"),
     ("shuffle_5d_cache", "launch_reshape_and_cache_shuffle_5d"),
     ("shuffle_5d_cache", "launch_gather_shuffle_5d_to_linear"),
     ("pd_dcp_gather", "copy_mla_rows_into_pack"),

python/sglang/kernels/ops/kvcache/dsv41_window_layout.py

@@ -0,0 +1,203 @@
+"""Build query and history metadata after request-window group scans."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _queries(
+    Req,
+    Pos,
+    Floor,
+    Group,
+    First,
+    Last,
+    Starts,
+    OutReq,
+    OutPos,
+    Write,
+    Indices,
+    Lengths,
+    Commit,
+    N,
+    PAD,
+    HISTORY,
+    WINDOW: tl.constexpr,
+    CAPACITY: tl.constexpr,
+    RS: tl.constexpr,
+    PS: tl.constexpr,
+    FS: tl.constexpr,
+    HAS_FLOOR: tl.constexpr,
+    R: tl.constexpr,
+    W: tl.constexpr,
+):
+    row = tl.program_id(0) * R + tl.arange(0, R)
+    live = row < N
+    req = tl.load(Req + row * RS, live, 0).to(tl.int64)
+    pos = tl.load(Pos + row * PS, live, 0).to(tl.int64)
+    group = tl.load(Group + row, live, 0)
+    first = tl.load(First + row, live, 0)
+    last = tl.load(Last + row, live, 0)
+    tl.store(Starts + group, row, live & (row == first))
+    first_pos = pos - (row - first)
+    lookback = tl.arange(0, W)
+    seen = pos[:, None] - lookback[None, :]
+    valid = live[:, None] & (lookback[None, :] < WINDOW) & (seen >= 0)
+    if HAS_FLOOR:
+        floor = tl.load(Floor + row * FS, live, 0).to(tl.int64)
+        valid &= seen >= floor[:, None]
+    old_loc = group[:, None] * WINDOW + seen - (first_pos[:, None] - WINDOW)
+    new_loc = HISTORY + first[:, None] + seen - first_pos[:, None]
+    loc = tl.where(seen < first_pos[:, None], old_loc, new_loc)
+    tl.store(
+        Indices + row[:, None] * WINDOW + lookback[None, :],
+        tl.where(valid, loc, -1),
+        (row[:, None] < PAD) & (lookback[None, :] < WINDOW),
+    )
+    tl.store(OutReq + row, req, row < PAD)
+    tl.store(OutPos + row, pos, row < PAD)
+    tl.store(Write + row, HISTORY + row, row < PAD)
+    tl.store(Lengths + row, tl.sum(valid.to(tl.int32), 1), row < PAD)
+    tl.store(Commit + row, live & ((last - row) < CAPACITY), row < PAD)
+
+
+@triton.jit
+def _history(
+    Req,
+    Pos,
+    Floor,
+    Last,
+    Starts,
+    HistReq,
+    HistPos,
+    HistLoc,
+    HistValid,
+    HistSrc,
+    ROWS,
+    WINDOW: tl.constexpr,
+    CAPACITY: tl.constexpr,
+    ZERO,
+    RS: tl.constexpr,
+    PS: tl.constexpr,
+    FS: tl.constexpr,
+    HAS_FLOOR: tl.constexpr,
+    B: tl.constexpr,
+):
+    row = tl.program_id(0) * B + tl.arange(0, B)
+    mask = row < ROWS
+    start = tl.load(Starts + row // WINDOW, mask, -1)
+    live = mask & (start >= 0)
+    end = tl.load(Last + start, live, 0)
+    req = tl.load(Req + end * RS, live, 0).to(tl.int64)
+    # Every row in a contiguous request run has the same first position and
+    # floor. Use its last row, matching the CPU scatter reference.
+    first_pos = (
+        tl.load(Pos + end * PS, live, 0).to(tl.int64) - end + tl.maximum(start, 0)
+    )
+    pos = first_pos - WINDOW + row % WINDOW
+    valid = live & (pos >= 0)
+    if HAS_FLOOR:
+        floor = tl.load(Floor + end * FS, live, 0).to(tl.int64)
+        valid &= pos >= floor
+    src = tl.where(valid, req * CAPACITY + pos % CAPACITY, ZERO)
+    tl.store(HistReq + row, req, mask)
+    tl.store(HistPos + row, pos, mask)
+    tl.store(HistLoc + row, row, mask)
+    tl.store(HistValid + row, valid, mask)
+    tl.store(HistSrc + row, src, mask)
+
+
+def build_window_layout(
+    req,
+    pos,
+    group,
+    first,
+    last,
+    *,
+    zero_row,
+    window,
+    capacity,
+    floor,
+    num_groups,
+    padded_rows,
+):
+    n = pos.numel()
+    padded = max(n, padded_rows or n)
+    groups = n if num_groups is None else int(num_groups)
+    history_rows = groups * window
+    device = pos.device
+
+    def empty(shape, dtype=torch.int64):
+        return torch.empty(shape, dtype=dtype, device=device)
+
+    out_req, out_pos = empty(padded), empty(padded)
+    write = empty(padded, torch.int32)
+    indices = empty((padded, window), torch.int32)
+    lengths = empty(padded, torch.int32)
+    commit = empty(padded, torch.bool)
+    starts = torch.full((groups,), -1, dtype=torch.int64, device=device)
+    hist_req, hist_pos, hist_loc = (empty(history_rows) for _ in range(3))
+    hist_valid = empty(history_rows, torch.bool)
+    hist_src = empty(history_rows)
+    floor_arg = floor if floor is not None else pos
+    strides = (req.stride(0), pos.stride(0), floor_arg.stride(0))
+    _queries[(triton.cdiv(padded, 4),)](
+        req,
+        pos,
+        floor_arg,
+        group,
+        first,
+        last,
+        starts,
+        out_req,
+        out_pos,
+        write,
+        indices,
+        lengths,
+        commit,
+        n,
+        padded,
+        history_rows,
+        window,
+        capacity,
+        *strides,
+        floor is not None,
+        4,
+        triton.next_power_of_2(window),
+        num_warps=4,
+    )
+    _history[(triton.cdiv(history_rows, 256),)](
+        req,
+        pos,
+        floor_arg,
+        last,
+        starts,
+        hist_req,
+        hist_pos,
+        hist_loc,
+        hist_valid,
+        hist_src,
+        history_rows,
+        window,
+        capacity,
+        zero_row,
+        *strides,
+        floor is not None,
+        256,
+        num_warps=4,
+    )
+    return (
+        out_req,
+        out_pos,
+        write,
+        indices,
+        lengths,
+        hist_req,
+        hist_pos,
+        hist_loc,
+        hist_valid,
+        hist_src,
+        commit,
+        history_rows + padded,
+    )

python/sglang/srt/mem_cache/dsv41_request_window.py

@@ -99,6 +99,25 @@ def window_layout(
     n = pos.numel()
     if n == 0:
         raise ValueError("request-window layout needs at least one query")
+    if req.is_cuda and torch.version.hip is None:
+        from sglang.kernels.ops.kvcache.dsv41_window_layout import build_window_layout
+
+        group, first, last = _first_row_offsets(req)
+        return WindowLayout(
+            *build_window_layout(
+                req,
+                pos,
+                group,
+                first,
+                last,
+                zero_row=zero_row,
+                window=window,
+                capacity=capacity,
+                floor=floor,
+                num_groups=num_groups,
+                padded_rows=padded_rows,
+            )
+        )
     req = req.to(torch.int64)
     pos = pos.to(torch.int64)
     device = pos.device

test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py

@@ -0,0 +1,96 @@
+"""Fused metadata agrees with the CPU reference, including graph refreshes."""
+
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.dsv41_request_window import window_layout
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+class TestWindowLayout(unittest.TestCase):
+    def compare(self, actual, expected):
+        for name in expected.__struct_fields__:
+            a, b = getattr(actual, name), getattr(expected, name)
+            if isinstance(b, torch.Tensor):
+                self.assertEqual(a.dtype, b.dtype, name)
+                torch.testing.assert_close(a.cpu(), b, rtol=0, atol=0, msg=name)
+            else:
+                self.assertEqual(a, b, name)
+
+    def test_reference(self):
+        for sizes in ([1], [3, 1, 260, 7], [4096], [4096] * 8):
+            for window in (7, 128):
+                for dtype in (torch.int32, torch.int64):
+                    req = torch.cat(
+                        [
+                            torch.full((n,), [9, 2, 9, 5][i % 4], dtype=dtype)
+                            for i, n in enumerate(sizes)
+                        ]
+                    )
+                    pos = torch.cat(
+                        [
+                            torch.arange(n, dtype=dtype) + (0 if i % 2 else 123456)
+                            for i, n in enumerate(sizes)
+                        ]
+                    )
+                    floor = torch.cat(
+                        [
+                            torch.full((n,), 123400 if i % 2 == 0 else 0, dtype=dtype)
+                            for i, n in enumerate(sizes)
+                        ]
+                    )
+                    for use_floor in (False, True):
+                        for padding in (None, len(req) + 13):
+                            kw = dict(
+                                zero_row=12345678,
+                                window=window,
+                                capacity=256,
+                                num_groups=len(sizes) + 2,
+                                padded_rows=padding,
+                            )
+                            cpu = window_layout(
+                                req, pos, floor=floor if use_floor else None, **kw
+                            )
+
+                            # Noncontiguous input vectors exercise input strides.
+                            def gpu(t):
+                                return torch.stack([t, t], dim=1).cuda()[:, 0]
+
+                            actual = window_layout(
+                                gpu(req),
+                                gpu(pos),
+                                floor=gpu(floor) if use_floor else None,
+                                **kw,
+                            )
+                            self.compare(actual, cpu)
+
+    def test_graph_refresh(self):
+        req = torch.tensor([1, 1, 1, 2, 2, 2], device="cuda")
+        pos = torch.tensor([256, 257, 258, 0, 1, 2], device="cuda")
+        floor = torch.zeros_like(pos)
+        kw = dict(zero_row=99999, num_groups=6, padded_rows=10)
+        for _ in range(3):
+            window_layout(req, pos, floor=floor, **kw)
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = window_layout(req, pos, floor=floor, **kw)
+        for r, p, f in [
+            ([5, 5, 8, 8, 8, 5], [0, 1, 512, 513, 514, 70], [0, 0, 510, 510, 510, 0]),
+            ([3] * 6, list(range(200, 206)), [195] * 6),
+        ]:
+            req.copy_(torch.tensor(r, device="cuda"))
+            pos.copy_(torch.tensor(p, device="cuda"))
+            floor.copy_(torch.tensor(f, device="cuda"))
+            graph.replay()
+            self.compare(
+                captured, window_layout(req.cpu(), pos.cpu(), floor=floor.cpu(), **kw)
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()