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 10, 00:15 UTC
[BUGFIX] Fix degenerate strides in TRTLLM query tensors for FlashInfer backend. Fixes issue #32353 (#32417)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Jan 19, 2026 at 04:57 UTC
6101a26dc95631d9e0cfe217089b1bb7f73d533e
1 file changed
+10
-4
vllm/v1/attention/backends/flashinfer.py
+10
-4
@@ -1385,8 +1385,11 @@ class FlashInferImpl(AttentionImpl):
1385
)
1386
else:
1387
assert isinstance(attn_metadata.prefill, TRTLLMPrefill)
1388
- # prefill_query may be non-contiguous
1389
- prefill_query = prefill_query.contiguous()
1388
+ # prefill_query may be non-contiguous or have degenerate strides
1389
+ # First ensure memory contiguity, then fix degenerate strides
1390
+ # with reshape. contiguous() alone doesn't fix degenerate
1391
+ # strides when a dimension has size 1.
1392
+ prefill_query = prefill_query.contiguous().reshape(prefill_query.shape)
1393
workspace_buffer = _get_trtllm_gen_workspace_buffer()
1394
block_tables_prefill = attn_metadata.prefill.block_tables
1395
seq_lens_prefill = attn_metadata.prefill.seq_lens
@@ -1495,9 +1498,12 @@ class FlashInferImpl(AttentionImpl):
1498
out=output[:num_decode_tokens],
1499
)
1500
else:
1498
- # decode_query may be non-contiguous
1501
+ # decode_query may be non-contiguous or have degenerate strides
1502
assert isinstance(attn_metadata.decode, TRTLLMDecode)
1500
- decode_query = decode_query.contiguous()
1503
+ # First ensure memory contiguity, then fix degenerate strides
1504
+ # with reshape. contiguous() alone doesn't fix degenerate
1505
+ # strides when a dimension has size 1.
1506
+ decode_query = decode_query.contiguous().reshape(decode_query.shape)
1507
workspace_buffer = _get_trtllm_gen_workspace_buffer()
1508
block_tables_decode = attn_metadata.decode.block_tables
1509
seq_lens_decode = attn_metadata.decode.seq_lens