From 6fade8dd2a5433797bc1c4dfa54936d8b3cda2cd Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 22 Oct 2026 10:56:51 +0800 Subject: [PATCH 1/6] [Perf] dsv41: Fuse request-window metadata construction --- python/sglang/kernels/ops/kvcache/__init__.py | 1 + .../ops/kvcache/dsv41_window_layout.py | 203 ++++++++++++++++++ .../srt/mem_cache/dsv41_request_window.py | 19 ++ .../ops/kvcache/test_dsv41_window_layout.py | 96 +++++++++ 4 files changed, 319 insertions(+) create mode 100644 python/sglang/kernels/ops/kvcache/dsv41_window_layout.py create mode 100644 test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py diff --git a/python/sglang/kernels/ops/kvcache/__init__.py b/python/sglang/kernels/ops/kvcache/__init__.py index 10e56ebe5f6d..a2b3623debd5 100644 --- a/python/sglang/kernels/ops/kvcache/__init__.py +++ b/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"), diff --git a/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py b/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py new file mode 100644 index 000000000000..f1e51c40413f --- /dev/null +++ b/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: tl.constexpr, + PAD: tl.constexpr, + HISTORY: tl.constexpr, + 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: tl.constexpr, + WINDOW: tl.constexpr, + CAPACITY: tl.constexpr, + ZERO: tl.constexpr, + 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, + ) diff --git a/python/sglang/srt/mem_cache/dsv41_request_window.py b/python/sglang/srt/mem_cache/dsv41_request_window.py index 9f8e7ed39acb..76d32f4954b2 100644 --- a/python/sglang/srt/mem_cache/dsv41_request_window.py +++ b/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 diff --git a/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py b/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py new file mode 100644 index 000000000000..478ce68b964f --- /dev/null +++ b/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], [4224], [4224] * 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() From c8ecd4c168205f036b9c3f72c9683a33f6e33b78 Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 22 Oct 2026 10:59:35 +0800 Subject: [PATCH 2/6] [Test] dsv41: Add synthetic layout construction benchmark --- .../kernels/bench_dsv41_window_layout.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 benchmark/kernels/bench_dsv41_window_layout.py diff --git a/benchmark/kernels/bench_dsv41_window_layout.py b/benchmark/kernels/bench_dsv41_window_layout.py new file mode 100644 index 000000000000..0f7e4be8ffc5 --- /dev/null +++ b/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() From 381c4c868e76091a811fa0f0ec0df56ed910e940 Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 22 Oct 2026 11:05:36 +0800 Subject: [PATCH 3/6] [Test] dsv41: Use generic power-of-two layout sizes --- test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py b/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py index 478ce68b964f..317383095d4c 100644 --- a/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py +++ b/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py @@ -22,7 +22,7 @@ def compare(self, actual, expected): self.assertEqual(a, b, name) def test_reference(self): - for sizes in ([1], [3, 1, 260, 7], [4224], [4224] * 8): + 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( From fc63b690eacd34a5448550b06eef10bc5f1a91fb Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 22 Oct 2026 11:07:40 +0800 Subject: [PATCH 4/6] [Perf] dsv41: Keep layout lengths as runtime kernel arguments --- .../sglang/kernels/ops/kvcache/dsv41_window_layout.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py b/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py index f1e51c40413f..3666d3e2179b 100644 --- a/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py +++ b/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py @@ -20,9 +20,9 @@ def _queries( Indices, Lengths, Commit, - N: tl.constexpr, - PAD: tl.constexpr, - HISTORY: tl.constexpr, + N, + PAD, + HISTORY, WINDOW: tl.constexpr, CAPACITY: tl.constexpr, RS: tl.constexpr, @@ -74,10 +74,10 @@ def _history( HistLoc, HistValid, HistSrc, - ROWS: tl.constexpr, + ROWS, WINDOW: tl.constexpr, CAPACITY: tl.constexpr, - ZERO: tl.constexpr, + ZERO, RS: tl.constexpr, PS: tl.constexpr, FS: tl.constexpr, From 5616a394f693d57bcd19e5c35a9a0ec2337c214e Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 22 Oct 2026 13:01:18 +0800 Subject: [PATCH 5/6] Remove standalone window layout test and benchmark per review --- .../kernels/bench_dsv41_window_layout.py | 47 --------- .../ops/kvcache/test_dsv41_window_layout.py | 96 ------------------- 2 files changed, 143 deletions(-) delete mode 100644 benchmark/kernels/bench_dsv41_window_layout.py delete mode 100644 test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py diff --git a/benchmark/kernels/bench_dsv41_window_layout.py b/benchmark/kernels/bench_dsv41_window_layout.py deleted file mode 100644 index 0f7e4be8ffc5..000000000000 --- a/benchmark/kernels/bench_dsv41_window_layout.py +++ /dev/null @@ -1,47 +0,0 @@ -"""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() diff --git a/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py b/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py deleted file mode 100644 index 317383095d4c..000000000000 --- a/test/registered/kernels/ops/kvcache/test_dsv41_window_layout.py +++ /dev/null @@ -1,96 +0,0 @@ -"""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() From 7418d0bbecf621b6384278d1fb6be4d0c2fb0b86 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Thu, 22 Oct 2026 00:11:52 -0700 Subject: [PATCH 6/6] mask window-layout group starts to num_groups with an invariant; name the kernels and return the layout fields by name; use the hip helper --- .../ops/kvcache/dsv41_window_layout.py | 44 ++++++++++--------- .../srt/mem_cache/dsv41_request_window.py | 17 ++++++- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py b/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py index 3666d3e2179b..aa826ea7b41c 100644 --- a/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py +++ b/python/sglang/kernels/ops/kvcache/dsv41_window_layout.py @@ -6,7 +6,7 @@ @triton.jit -def _queries( +def _window_layout_queries( Req, Pos, Floor, @@ -23,6 +23,7 @@ def _queries( N, PAD, HISTORY, + GROUPS, WINDOW: tl.constexpr, CAPACITY: tl.constexpr, RS: tl.constexpr, @@ -39,7 +40,8 @@ def _queries( 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)) + # 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, :] @@ -63,7 +65,7 @@ def _queries( @triton.jit -def _history( +def _window_layout_history( Req, Pos, Floor, @@ -90,8 +92,8 @@ def _history( 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. + # 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) ) @@ -142,7 +144,7 @@ def empty(shape, dtype=torch.int64): 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),)]( + _window_layout_queries[(triton.cdiv(padded, 4),)]( req, pos, floor_arg, @@ -159,6 +161,7 @@ def empty(shape, dtype=torch.int64): n, padded, history_rows, + groups, window, capacity, *strides, @@ -167,7 +170,7 @@ def empty(shape, dtype=torch.int64): triton.next_power_of_2(window), num_warps=4, ) - _history[(triton.cdiv(history_rows, 256),)]( + _window_layout_history[(triton.cdiv(history_rows, 256),)]( req, pos, floor_arg, @@ -187,17 +190,18 @@ def empty(shape, dtype=torch.int64): 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, + # 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, ) diff --git a/python/sglang/srt/mem_cache/dsv41_request_window.py b/python/sglang/srt/mem_cache/dsv41_request_window.py index 76d32f4954b2..47184a75782e 100644 --- a/python/sglang/srt/mem_cache/dsv41_request_window.py +++ b/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,12 +106,18 @@ 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: + 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( + **build_window_layout( req, pos, group,