biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[BUGFIX] Fix degenerate strides in TRTLLM query tensors for FlashInfer backend. Fixes issue #32353 (#32417)

Browse Source 
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
This commit is contained in:
Vadim Gimpelson
2026-01-19 04:57:32 +04:00
committed by GitHub
parent f5d1740030
commit 6101a26dc9
1 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1385,8 +1385,11 @@ class FlashInferImpl(AttentionImpl):
) )
else: else:
assert isinstance(attn_metadata.prefill, TRTLLMPrefill) assert isinstance(attn_metadata.prefill, TRTLLMPrefill)
# prefill_query may be non-contiguous # prefill_query may be non-contiguous or have degenerate strides
prefill_query = prefill_query.contiguous() # First ensure memory contiguity, then fix degenerate strides
# with reshape. contiguous() alone doesn't fix degenerate
# strides when a dimension has size 1.
prefill_query = prefill_query.contiguous().reshape(prefill_query.shape)
workspace_buffer = _get_trtllm_gen_workspace_buffer() workspace_buffer = _get_trtllm_gen_workspace_buffer()
block_tables_prefill = attn_metadata.prefill.block_tables block_tables_prefill = attn_metadata.prefill.block_tables
seq_lens_prefill = attn_metadata.prefill.seq_lens seq_lens_prefill = attn_metadata.prefill.seq_lens
@@ -1495,9 +1498,12 @@ class FlashInferImpl(AttentionImpl):
out=output[:num_decode_tokens], out=output[:num_decode_tokens],
) )
else: else:
# decode_query may be non-contiguous # decode_query may be non-contiguous or have degenerate strides
assert isinstance(attn_metadata.decode, TRTLLMDecode) assert isinstance(attn_metadata.decode, TRTLLMDecode)
decode_query = decode_query.contiguous() # First ensure memory contiguity, then fix degenerate strides
# with reshape. contiguous() alone doesn't fix degenerate
# strides when a dimension has size 1.
decode_query = decode_query.contiguous().reshape(decode_query.shape)
workspace_buffer = _get_trtllm_gen_workspace_buffer() workspace_buffer = _get_trtllm_gen_workspace_buffer()
block_tables_decode = attn_metadata.decode.block_tables block_tables_decode = attn_metadata.decode.block_tables
seq_lens_decode = attn_metadata.decode.seq_lens seq_lens_decode = attn_metadata.decode.seq_lens