perf: use int32 paged KV indices

- store page-table, request-row, and cache-location indices as int32
- preserve CUDA graph replay with bit-exact logits and KV cache coverage
- improve B=1 decode latency by 1-6% across 1K-32K contexts on L20
This commit is contained in:
2026-08-15 13:17:06 +08:00
parent b5afe3d7a4
commit 3fb4b8ab13
11 changed files with 135 additions and 69 deletions
+2 -2
View File
@@ -74,7 +74,7 @@ class InferenceWorkspace:
# when the Executor passes this workspace). Stable addresses make the
# decode forward CUDA-graph capturable.
self.req_pool_indices = torch.empty(
(max_batch_size,), dtype=torch.long, device=device
(max_batch_size,), dtype=torch.int32, device=device
)
self.seq_lens = torch.empty((max_batch_size,), dtype=torch.long, device=device)
self.kv_indptr = torch.empty(
@@ -85,7 +85,7 @@ class InferenceWorkspace:
)
self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device)
self.out_cache_loc = torch.empty(
(max_batch_size, 1), dtype=torch.long, device=device
(max_batch_size, 1), dtype=torch.int32, device=device
)
# Per-step position IDs (must be at a fixed address for CUDA-graph capture).