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:
@@ -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).
|
||||
|
||||
Reference in New Issue
Block a user