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 6, 15:45 UTC
[Bugfix] Relax TRTLLM KV cache contiguity assertion for cross-layer layout (#34158)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Itay Etelis committed
Mar 16, 2026 at 17:20 UTC
5ae685c1c85bb659476a21ce7a2457eb6cccc4bb
1 file changed
+27
-2
vllm/v1/attention/backends/flashinfer.py
+27
-2
@@ -586,6 +586,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
586
# try to use fp8 q if kv cache is fp8, and will fall back to model dtype
587
# if TRTLLM attention kernel is not used when building attn metadata
588
can_use_trtllm = can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)
589
+
590
if (
591
can_use_trtllm
592
and not vllm_config.attention_config.disable_flashinfer_q_quantization
@@ -1436,7 +1437,6 @@ class FlashInferImpl(AttentionImpl):
1437
# This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
1438
assert get_kv_cache_layout() == "HND"
1439
assert is_strictly_contiguous(prefill_query)
1439
- assert is_strictly_contiguous(kv_cache_permute)
1440
assert is_strictly_contiguous(workspace_buffer)
1441
assert is_strictly_contiguous(block_tables_prefill)
1442
assert is_strictly_contiguous(seq_lens_prefill)
@@ -1461,6 +1461,20 @@ class FlashInferImpl(AttentionImpl):
1461
# and fp8 kv cache. So to enable prefill attention
1462
# with fp8 kv cache, we can construct a mock block
1463
# and mock kv cache with BF16 KV involved in the prefill
1464
+ #
1465
+ # The inner (block_size, head_size) dims must be
1466
+ # contiguous; outer dims may have non-canonical strides
1467
+ # (e.g. cross-layer unified allocation).
1468
+ # Degenerate strides on outer dims break TMA descriptors
1469
+ # (see flashinfer-ai/flashinfer#2232).
1470
+ kv_strides = kv_cache_permute.stride()
1471
+ assert (
1472
+ kv_strides[-1] == 1
1473
+ and kv_strides[-2] == kv_cache_permute.shape[-1]
1474
+ ), (
1475
+ "KV cache inner dims (block_size, head_size) must be "
1476
+ f"contiguous, got strides {kv_strides}"
1477
+ )
1478
mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
1479
kv_cache_permute,
1480
block_tables_prefill,
@@ -1549,10 +1563,21 @@ class FlashInferImpl(AttentionImpl):
1563
# This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
1564
assert get_kv_cache_layout() == "HND"
1565
assert is_strictly_contiguous(decode_query)
1552
- assert is_strictly_contiguous(kv_cache_permute)
1566
assert is_strictly_contiguous(workspace_buffer)
1567
assert is_strictly_contiguous(block_tables_decode)
1568
assert is_strictly_contiguous(seq_lens_decode)
1569
+ # kv_cache outer dims may be non-contiguous (e.g.
1570
+ # cross-layer unified allocation), but inner dims
1571
+ # (block_size, head_size) must be contiguous and
1572
+ # strides must be canonical to avoid TMA descriptor
1573
+ # failures (see flashinfer-ai/flashinfer#2232).
1574
+ kv_strides = kv_cache_permute.stride()
1575
+ assert (
1576
+ kv_strides[-1] == 1 and kv_strides[-2] == kv_cache_permute.shape[-1]
1577
+ ), (
1578
+ "KV cache inner dims (block_size, head_size) must be "
1579
+ f"contiguous, got strides {kv_strides}"
1580
+ )
1581
1582
if output.dtype == FP4_DTYPE:
1583
assert self.o_sf_scale is not None