xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-14 18:45:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

Remove cuda hard-code in compute_causal_conv1d_metadata (#25555)

Browse Source 
Signed-off-by: Icey <1790571317@qq.com>
This commit is contained in:
Icey 2025-09-26 16:19:20 +08:00 committed by GitHub
parent 99b3a504c5
commit dd70437a4f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/v1/attention/backends/utils.py
View File

@ -947,6 +947,7 @@ def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
nums_dict = {} # type: ignore nums_dict = {} # type: ignore
batch_ptr = None batch_ptr = None
token_chunk_offset_ptr = None token_chunk_offset_ptr = None
device = query_start_loc_p.device
for BLOCK_M in [8]: # cover all BLOCK_M values for BLOCK_M in [8]: # cover all BLOCK_M values
nums = -(-seqlens // BLOCK_M) nums = -(-seqlens // BLOCK_M)
nums_dict[BLOCK_M] = {} nums_dict[BLOCK_M] = {}
@ -968,11 +969,11 @@ def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
batch_ptr = torch.full((MAX_NUM_PROGRAMS, ), batch_ptr = torch.full((MAX_NUM_PROGRAMS, ),
PAD_SLOT_ID, PAD_SLOT_ID,
dtype=torch.int32, dtype=torch.int32,
device='cuda') device=device)
token_chunk_offset_ptr = torch.full((MAX_NUM_PROGRAMS, ), token_chunk_offset_ptr = torch.full((MAX_NUM_PROGRAMS, ),
PAD_SLOT_ID, PAD_SLOT_ID,
dtype=torch.int32, dtype=torch.int32,
device='cuda') device=device)
else: else:
if batch_ptr.nelement() < MAX_NUM_PROGRAMS: if batch_ptr.nelement() < MAX_NUM_PROGRAMS:
batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID) batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)