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, 09:00 UTC
Revert "[Kernels][FI] Skip trtllm attention when num_kv_heads=1 (#308… (#31617)
Signed-off-by: shyeh25 <206795756+shyeh25@users.noreply.github.com>
shyeh25 committed
Jan 11, 2026 at 04:39 UTC
1c46dea0017a7689c2d9fbc0f7d1f1357ca6c7de
2 files changed
+1
-56
tests/kernels/attention/test_flashinfer_trtllm_attention.py
-35
@@ -456,38 +456,3 @@ def test_flashinfer_trtllm_prefill_with_baseline(
456
torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol),
457
f"{torch.max(torch.abs(output - output_trtllm))}",
458
)
459
-
460
-
461
-def test_trtllm_attention_rejects_num_kv_heads_1(default_vllm_config) -> None:
462
- """Test that TRTLLM attention correctly rejects num_kv_heads=1.
463
-
464
- When num_kv_heads=1 (MQA), the KV cache strides become degenerate
465
- (stride_heads == stride_batch), which causes CUDA's cuTensorMapEncodeTiled
466
- to fail because TMA descriptors cannot handle degenerate 4D tensors with
467
- singleton dimensions.
468
-
469
- This test verifies that can_use_trtllm_attention returns False for
470
- num_kv_heads=1 configurations.
471
- """
472
- from vllm.utils.flashinfer import can_use_trtllm_attention
473
-
474
- # num_kv_heads=1 should be rejected
475
- assert not can_use_trtllm_attention(num_qo_heads=64, num_kv_heads=1), (
476
- "can_use_trtllm_attention should return False for num_kv_heads=1"
477
- )
478
- assert not can_use_trtllm_attention(num_qo_heads=32, num_kv_heads=1), (
479
- "can_use_trtllm_attention should return False for num_kv_heads=1"
480
- )
481
-
482
- # num_kv_heads > 1 should be accepted (if platform supports it)
483
- # Note: This may return False on non-Blackwell platforms, which is fine
484
- result_kv8 = can_use_trtllm_attention(num_qo_heads=64, num_kv_heads=8)
485
- result_kv1 = can_use_trtllm_attention(num_qo_heads=64, num_kv_heads=1)
486
-
487
- # Even if platform doesn't support TRTLLM, num_kv_heads=1 should never
488
- # return True when num_kv_heads > 1 returns True
489
- if result_kv8:
490
- assert not result_kv1, (
491
- "If TRTLLM is supported for num_kv_heads=8, "
492
- "it must be rejected for num_kv_heads=1"
493
- )
vllm/utils/flashinfer.py
+1
-21
@@ -305,18 +305,7 @@ 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
- # 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)
308
+ return has_trtllm and (num_qo_heads % num_kv_heads == 0)
309
310
311
def use_trtllm_attention(
@@ -366,15 +355,6 @@ def use_trtllm_attention(
355
)
356
return False
357
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
-
358
if has_spec and not is_prefill:
359
# Speculative decoding requires TRTLLM attention for decodes
360
logger.info_once("Using TRTLLM attention (enabled for speculative decoding).")