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 11, 08:45 UTC
[Kernels][FI] Skip trtllm attention when num_kv_heads=1 (#30842)
Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com>
Ye (Charlotte) Qi committed
Dec 17, 2025 at 01:54 UTC
a100152288c8ec50336aea842f0b3d8e36624024
2 files changed
+56
-1
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+35
@@ -455,3 +455,38 @@ def test_flashinfer_trtllm_prefill_with_baseline(
455
torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol),
456
f"{torch.max(torch.abs(output - output_trtllm))}",
457
)
458
+
459
+
460
+def test_trtllm_attention_rejects_num_kv_heads_1() -> None:
461
+ """Test that TRTLLM attention correctly rejects num_kv_heads=1.
462
+
463
+ When num_kv_heads=1 (MQA), the KV cache strides become degenerate
464
+ (stride_heads == stride_batch), which causes CUDA's cuTensorMapEncodeTiled
465
+ to fail because TMA descriptors cannot handle degenerate 4D tensors with
466
+ singleton dimensions.
467
+
468
+ This test verifies that can_use_trtllm_attention returns False for
469
+ num_kv_heads=1 configurations.
470
+ """
471
+ from vllm.utils.flashinfer import can_use_trtllm_attention
472
+
473
+ # num_kv_heads=1 should be rejected
474
+ assert not can_use_trtllm_attention(num_qo_heads=64, num_kv_heads=1), (
475
+ "can_use_trtllm_attention should return False for num_kv_heads=1"
476
+ )
477
+ assert not can_use_trtllm_attention(num_qo_heads=32, num_kv_heads=1), (
478
+ "can_use_trtllm_attention should return False for num_kv_heads=1"
479
+ )
480
+
481
+ # num_kv_heads > 1 should be accepted (if platform supports it)
482
+ # Note: This may return False on non-Blackwell platforms, which is fine
483
+ result_kv8 = can_use_trtllm_attention(num_qo_heads=64, num_kv_heads=8)
484
+ result_kv1 = can_use_trtllm_attention(num_qo_heads=64, num_kv_heads=1)
485
+
486
+ # Even if platform doesn't support TRTLLM, num_kv_heads=1 should never
487
+ # return True when num_kv_heads > 1 returns True
488
+ if result_kv8:
489
+ assert not result_kv1, (
490
+ "If TRTLLM is supported for num_kv_heads=8, "
491
+ "it must be rejected for num_kv_heads=1"
492
+ )
vllm/utils/flashinfer.py
+21
-1
@@ -305,7 +305,18 @@ def can_use_trtllm_attention(num_qo_heads: int, num_kv_heads: int) -> bool:
305
if force_use_trtllm_attention() is False:
306
return False
307
has_trtllm = supports_trtllm_attention()
308
- return has_trtllm and (num_qo_heads % num_kv_heads == 0)
308
+ # num_kv_heads=1 is not supported due to TMA descriptor building limitations.
309
+ # When num_kv_heads=1, the KV cache strides become degenerate (stride_heads ==
310
+ # stride_batch), which causes CUDA's cuTensorMapEncodeTiled to fail because
311
+ # TMA descriptors cannot handle degenerate 4D tensors with singleton dimensions.
312
+ # See: https://fburl.com/352mrydz
313
+ if has_trtllm and num_kv_heads == 1:
314
+ logger.warning_once(
315
+ "TRTLLM attention does not support num_kv_heads=1. "
316
+ "This configuration causes TMA descriptor building to fail due to "
317
+ "degenerate tensor strides. Falling back to FlashInfer attention."
318
+ )
319
+ return has_trtllm and (num_qo_heads % num_kv_heads == 0) and (num_kv_heads != 1)
320
321
322
def use_trtllm_attention(
@@ -355,6 +366,15 @@ def use_trtllm_attention(
366
)
367
return False
368
369
+ # num_kv_heads=1 is not supported
370
+ if num_kv_heads == 1:
371
+ if force_use_trtllm:
372
+ logger.warning_once(
373
+ "TRTLLM attention does not support num_kv_heads=1, "
374
+ "but --attention-config.use_trtllm_attention is set to 1"
375
+ )
376
+ return False
377
+
378
if has_spec and not is_prefill:
379
# Speculative decoding requires TRTLLM attention for decodes
380
logger.info_once("Using TRTLLM attention (enabled for speculative decoding).")