[Perf] dsv41: Fuse request-window layout construction

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,207 @@
+"""Build query and history metadata after request-window group scans."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _window_layout_queries(
+    Req,
+    Pos,
+    Floor,
+    Group,
+    First,
+    Last,
+    Starts,
+    OutReq,
+    OutPos,
+    Write,
+    Indices,
+    Lengths,
+    Commit,
+    N,
+    PAD,
+    HISTORY,
+    GROUPS,
+    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)
+    # Starts holds num_groups slots; window_layout's invariant reports a run past them.
+    tl.store(Starts + group, row, live & (row == first) & (group < GROUPS))
+    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 _window_layout_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 of a contiguous request run carries the same first position and
+    # floor, so its last row stands for the run.
+    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))
+    _window_layout_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,
+        groups,
+        window,
+        capacity,
+        *strides,
+        floor is not None,
+        4,
+        triton.next_power_of_2(window),
+        num_warps=4,
+    )
+    _window_layout_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,
+    )
+    # Keyed by WindowLayout's fields.
+    return dict(
+        req=out_req,
+        pos=out_pos,
+        write_loc=write,
+        indices=indices,
+        lengths=lengths,
+        history_req=hist_req,
+        history_pos=hist_pos,
+        history_loc=hist_loc,
+        history_valid=hist_valid,
+        history_src=hist_src,
+        commit_mask=commit,
+        size=history_rows + padded,
+    )

python/sglang/srt/mem_cache/dsv41_request_window.py

@@ -10,13 +10,20 @@
 )
 from sglang.srt.environ import InvariantCheckLevel
 from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode
+from sglang.srt.utils import is_hip
 from sglang.srt.utils.invariants import Bucket, Invariant, IsTrue, expect, resolve_level
 
+_is_hip = is_hip()
+
 # A window row the replay expects but no longer holds; that request's attention
 # would read the wrong rows.
 _SWA_HISTORY_PRESENT = Invariant(
     "dsv41.swa_history_present", Bucket.FATAL_CONTAINABLE, IsTrue()
 )
+# More contiguous request runs than num_groups: the runs past it get no history.
+_WINDOW_GROUPS_FIT = Invariant(
+    "dsv41.window_groups_fit", Bucket.FATAL_CONTAINABLE, IsTrue()
+)
 
 
 class WindowLayout(msgspec.Struct, frozen=True):
@@ -99,6 +106,31 @@ def window_layout(
     n = pos.numel()
     if n == 0:
         raise ValueError("request-window layout needs at least one query")
+    if req.is_cuda and not _is_hip:
+        from sglang.kernels.ops.kvcache.dsv41_window_layout import build_window_layout
+
+        group, first, last = _first_row_offsets(req)
+        if (
+            num_groups is not None
+            and resolve_level() >= InvariantCheckLevel.WARN
+            and not _capturing()
+        ):
+            expect(_WINDOW_GROUPS_FIT, group[-1:] < num_groups, msg="request runs")
+        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