xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-09 20:37:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

minor: zero workspace buffer init for flashinfer trtllm-gen attn (#22603)

Browse Source
This commit is contained in:
eigen 2025-08-15 17:38:10 -04:00 committed by GitHub
parent 00d6cba0cf
commit 1723ef1aae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/kernels/attention/test_flashinfer_trtllm_attention.py
View File

@ -113,7 +113,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
kv_indices = torch.tensor(kv_indices, dtype=torch.int32) kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8) workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
workspace_buffer, workspace_buffer,
kv_layout, kv_layout,
@ -247,7 +247,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
kv_indices = torch.tensor(kv_indices, dtype=torch.int32) kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8) workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(
workspace_buffer, kv_layout) workspace_buffer, kv_layout)
wrapper.plan(q_indptr, wrapper.plan(q_indptr,

2
vllm/attention/backends/flashinfer.py
View File

@ -203,7 +203,7 @@ class FlashInferState(AttentionState):
def _get_workspace_buffer(self): def _get_workspace_buffer(self):
if self._workspace_buffer is None: if self._workspace_buffer is None:
self._workspace_buffer = torch.empty( self._workspace_buffer = torch.zeros(
FLASHINFER_WORKSPACE_BUFFER_SIZE, FLASHINFER_WORKSPACE_BUFFER_SIZE,
dtype=torch.uint8, dtype=torch.uint8,
device=self.runner.device) device=self.runner.device)

2
vllm/v1/attention/backends/flashinfer.py
View File

@ -252,7 +252,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
def _get_workspace_buffer(self): def _get_workspace_buffer(self):
if self._workspace_buffer is None: if self._workspace_buffer is None:
self._workspace_buffer = torch.empty( self._workspace_buffer = torch.zeros(
FLASHINFER_WORKSPACE_BUFFER_SIZE, FLASHINFER_WORKSPACE_BUFFER_SIZE,
dtype=torch.uint8, dtype=torch.uint8,
device=self.device) device=self.device)