A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
minor: zero workspace buffer init for flashinfer trtllm-gen attn (#22603)
eigen committed
Aug 15, 2025 at 17:38 UTC
1723ef1aae749929c1cbddd964ab3ffd96452a70
3 files changed
+4
-4
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+2
-2
@@ -113,7 +113,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
113
kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
114
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
115
116
- workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
116
+ workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
117
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
118
workspace_buffer,
119
kv_layout,
@@ -247,7 +247,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
247
kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
248
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
249
250
- workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8)
250
+ workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
251
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
252
workspace_buffer, kv_layout)
253
wrapper.plan(q_indptr,
vllm/attention/backends/flashinfer.py
+1
-1
@@ -203,7 +203,7 @@ class FlashInferState(AttentionState):
203
204
def _get_workspace_buffer(self):
205
if self._workspace_buffer is None:
206
- self._workspace_buffer = torch.empty(
206
+ self._workspace_buffer = torch.zeros(
207
FLASHINFER_WORKSPACE_BUFFER_SIZE,
208
dtype=torch.uint8,
209
device=self.runner.device)
vllm/v1/attention/backends/flashinfer.py
+1
-1
@@ -252,7 +252,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
252
253
def _get_workspace_buffer(self):
254
if self._workspace_buffer is None:
255
- self._workspace_buffer = torch.empty(
255
+ self._workspace_buffer = torch.zeros(
256
FLASHINFER_WORKSPACE_BUFFER_SIZE,
257
dtype=torch.uint8,
258
device=self.device)