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
Remove cuda hard-code in compute_causal_conv1d_metadata (#25555)
Signed-off-by: Icey <1790571317@qq.com>
Icey committed
Sep 26, 2025 at 16:19 UTC
dd70437a4f36f08468bb572e3ab368c4d3b05c9a
1 file changed
+3
-2
vllm/v1/attention/backends/utils.py
+3
-2
@@ -947,6 +947,7 @@ def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
947
nums_dict = {} # type: ignore
948
batch_ptr = None
949
token_chunk_offset_ptr = None
950
+ device = query_start_loc_p.device
951
for BLOCK_M in [8]: # cover all BLOCK_M values
952
nums = -(-seqlens // BLOCK_M)
953
nums_dict[BLOCK_M] = {}
@@ -968,11 +969,11 @@ def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
969
batch_ptr = torch.full((MAX_NUM_PROGRAMS, ),
970
PAD_SLOT_ID,
971
dtype=torch.int32,
971
- device='cuda')
972
+ device=device)
973
token_chunk_offset_ptr = torch.full((MAX_NUM_PROGRAMS, ),
974
PAD_SLOT_ID,
975
dtype=torch.int32,
975
- device='cuda')
976
+ device=device)
977
else:
978
if batch_ptr.nelement() < MAX_NUM_PROGRAMS:
979
batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)