From 386f3ae0cbf38719dcd1ed39cf1d8bdab688559d Mon Sep 17 00:00:00 2001 From: Jin Pan Date: Thu, 20 Aug 2026 08:45:58 +0000 Subject: [PATCH] fix(cache): route null decode pages to dummy slot Keep fused target and DFlash decode writes consistent with the general cache-location contract so null or reclaimed pages cannot scatter writes across page zero. Signed-off-by: Jin Pan --- .../runtime/execution/cache_loc_kernel.py | 6 +- test/runtime/test_decode_cache_loc.py | 174 ++++++++++++++++++ 2 files changed, 177 insertions(+), 3 deletions(-) create mode 100644 test/runtime/test_decode_cache_loc.py diff --git a/python/tokenspeed/runtime/execution/cache_loc_kernel.py b/python/tokenspeed/runtime/execution/cache_loc_kernel.py index 49aa2e3dc0..724b5faae0 100644 --- a/python/tokenspeed/runtime/execution/cache_loc_kernel.py +++ b/python/tokenspeed/runtime/execution/cache_loc_kernel.py @@ -221,8 +221,8 @@ def fused_decode_input_prep_kernel( page_ptrs = page_table_ptr + req_idx * max_pages + page_indices page_ids = tl.load(page_ptrs, mask=mask, other=0) cache_locs = page_ids * page_size + offsets_in_page - # Route overflow tokens to slot 0 (fixed safe dummy target). - cache_locs = tl.where(overflow, 0, cache_locs) + # Null/hole pages share the same fixed dummy target as overflow. + cache_locs = tl.where(overflow | (page_ids <= 0), 0, cache_locs) tl.store( out_cache_loc_ptr + output_offset + token_offsets, @@ -318,7 +318,7 @@ def dflash_prepare_decode_kernel( page_ptrs = page_table_ptr + req_idx * max_pages + page_indices page_ids = tl.load(page_ptrs, mask=mask, other=0) cache_locs = page_ids * page_size + offsets_in_page - cache_locs = tl.where(overflow, 0, cache_locs) + cache_locs = tl.where(overflow | (page_ids <= 0), 0, cache_locs) tl.store( out_cache_loc_ptr + req_idx * draft_query_width + offsets, cache_locs, diff --git a/test/runtime/test_decode_cache_loc.py b/test/runtime/test_decode_cache_loc.py new file mode 100644 index 0000000000..d786284d11 --- /dev/null +++ b/test/runtime/test_decode_cache_loc.py @@ -0,0 +1,174 @@ +"""Decode write-location safety for null, boundary, and overflow pages.""" + +from __future__ import annotations + +import os +import sys + +import torch + +_TEST_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _TEST_DIR) +sys.path.insert(0, os.path.dirname(_TEST_DIR)) + +from ci_system.ci_register import register_cuda_ci + +register_cuda_ci(est_time=10, suite="runtime-1gpu") + +from test.runtime.conftest import requires_cuda + +from tokenspeed.runtime.execution.cache_loc_kernel import ( + dflash_prepare_decode, + fused_decode_input_prep, +) + +_PAGE_SIZE = 16 + + +@requires_cuda +def test_fused_decode_routes_null_negative_and_overflow_pages_to_zero() -> None: + req_pool_indices = torch.tensor([1, 2, 3, 4], device="cuda", dtype=torch.int64) + valid_cache_lengths = torch.tensor( + [0, 15, 17, 2, 48], device="cuda", dtype=torch.int32 + ) + page_table = torch.tensor( + [ + [7, 8, 9], + [11, 0, 13], + [-3, 12, 13], + [21, 22, 23], + ], + device="cuda", + dtype=torch.int32, + ) + out_cache_loc = torch.full((8,), -1, device="cuda", dtype=torch.int32) + positions = torch.full((8,), -1, device="cuda", dtype=torch.int64) + seq_lens = torch.full((4,), -1, device="cuda", dtype=torch.int32) + + fused_decode_input_prep( + out_cache_loc, + positions, + seq_lens, + req_pool_indices, + valid_cache_lengths, + uniform_input_length=2, + page_table=page_table, + page_size=_PAGE_SIZE, + ) + + assert out_cache_loc.view(4, 2).tolist() == [ + [7 * _PAGE_SIZE + 15, 8 * _PAGE_SIZE], + [0, 0], + [0, 0], + [0, 0], + ] + assert positions.view(4, 2).tolist() == [[15, 16], [17, 18], [2, 3], [48, 49]] + assert seq_lens.tolist() == [17, 19, 4, 50] + + +def _run_dflash( + page_table: torch.Tensor, + valid_cache_lengths: torch.Tensor, + *, + max_draft_prefix: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = page_table.shape[0] + verify_width = 8 + draft_query_width = 3 + output_tokens = torch.arange( + 100, + 100 + batch_size * verify_width, + device="cuda", + dtype=torch.int32, + ).view(batch_size, verify_width) + accept_lengths = torch.ones(batch_size, device="cuda", dtype=torch.int32) + req_pool_indices = torch.arange(1, batch_size + 1, device="cuda", dtype=torch.int64) + draft_seq_lens = torch.full((batch_size,), -1, device="cuda", dtype=torch.int32) + block_ids = torch.full( + (batch_size, draft_query_width), -1, device="cuda", dtype=torch.int32 + ) + block_positions = torch.full( + (batch_size, draft_query_width), -1, device="cuda", dtype=torch.int64 + ) + out_cache_loc = torch.full( + (batch_size * draft_query_width,), -1, device="cuda", dtype=torch.int32 + ) + + dflash_prepare_decode( + output_tokens=output_tokens, + accept_lengths=accept_lengths, + req_pool_indices=req_pool_indices, + valid_cache_lengths=valid_cache_lengths, + page_table=page_table, + draft_seq_lens=draft_seq_lens, + block_ids=block_ids, + block_positions=block_positions, + out_cache_loc=out_cache_loc, + verify_width=verify_width, + draft_query_width=draft_query_width, + page_size=_PAGE_SIZE, + max_draft_prefix=max_draft_prefix, + ) + return draft_seq_lens, block_ids, block_positions, out_cache_loc + + +@requires_cuda +def test_dflash_routes_null_and_negative_pages_to_zero_across_boundary() -> None: + page_table = torch.tensor( + [[7, 0, 9], [-2, 11, 12]], device="cuda", dtype=torch.int32 + ) + valid_cache_lengths = torch.tensor([0, 14, 0], device="cuda", dtype=torch.int32) + + seq_lens, block_ids, positions, cache_locs = _run_dflash( + page_table, + valid_cache_lengths, + max_draft_prefix=45, + ) + + assert seq_lens.tolist() == [15, 1] + assert block_ids[:, 0].tolist() == [100, 108] + assert positions.tolist() == [[15, 16, 17], [1, 2, 3]] + assert cache_locs.view(2, 3).tolist() == [ + [7 * _PAGE_SIZE + 15, 0, 0], + [0, 0, 0], + ] + + +@requires_cuda +def test_dflash_routes_out_of_table_positions_to_zero() -> None: + page_table = torch.tensor([[7, 8, 9]], device="cuda", dtype=torch.int32) + valid_cache_lengths = torch.tensor([0, 47], device="cuda", dtype=torch.int32) + + _, _, positions, cache_locs = _run_dflash( + page_table, + valid_cache_lengths, + max_draft_prefix=64, + ) + + assert positions.tolist() == [[48, 49, 50]] + assert cache_locs.tolist() == [0, 0, 0] + + +@requires_cuda +def test_dflash_capture_replay_observes_null_page_update() -> None: + page_table = torch.tensor([[7, 8]], device="cuda", dtype=torch.int32) + valid_cache_lengths = torch.tensor([0, 0], device="cuda", dtype=torch.int32) + + # Warm up Triton before capture. + _, _, _, cache_locs = _run_dflash( + page_table, + valid_cache_lengths, + max_draft_prefix=29, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _, _, _, cache_locs = _run_dflash( + page_table, + valid_cache_lengths, + max_draft_prefix=29, + ) + + page_table[0, 0] = 0 + graph.replay() + torch.cuda.synchronize() + assert cache_locs.tolist() == [0, 0, 0]