fix comment

Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>

Author: elvischenv

.buildkite/test-pipeline.yaml

@@ -664,7 +664,7 @@ steps:
     # Attention
     # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
     - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
-    - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_decode_attention.py
+    - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
     - pytest -v -s tests/kernels/test_cutlass_mla_decode.py
     # Quantization
     - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'

benchmarks/kernels/benchmark_trtllm_decode_attention.py

@@ -41,7 +41,6 @@ def benchmark_decode(
     device = "cuda"
     torch.manual_seed(0)
 
-    # Currently only HEAD_GRP_SIZE == 8 is supported
     HEAD_GRP_SIZE = 8
     MAX_SEQ_LEN = max_seq_len
 

benchmarks/kernels/benchmark_trtllm_prefill_attention.py

@@ -40,7 +40,6 @@ def benchmark_prefill(
     torch.set_default_device("cuda")
     torch.manual_seed(0)
 
-    # Currently only HEAD_GRP_SIZE == 8 is supported
     HEAD_GRP_SIZE = 8
     MAX_SEQ_LEN = max_seq_len
 

vllm/utils/flashinfer.py

@@ -159,7 +159,6 @@ def use_trtllm_attention(
 
     # Check if the dimensions are supported by TRTLLM decode attention
     if (attn_head_size is None or num_qo_heads is None or num_kv_heads is None
-            or num_qo_heads // num_kv_heads > 8
             or num_qo_heads % num_kv_heads != 0 or attn_head_size != 128):
         return False
 

vllm/v1/attention/backends/flashinfer.py

@@ -523,14 +523,16 @@ def build(self,
         head_dim = self.kv_cache_spec.head_size
 
         # currently prefill trtllm attention does not support fp8 kv cache
-        prefill_use_trtllm = not cache_dtype.startswith(
-            "fp8") and use_trtllm_attention(num_prefill_tokens, max_seq_len,
-                                            cache_dtype, num_qo_heads,
-                                            num_kv_heads, head_dim)
-        decode_use_trtllm = use_trtllm_attention(num_decode_tokens,
-                                                 max_seq_len, cache_dtype,
-                                                 num_qo_heads, num_kv_heads,
-                                                 head_dim)
+        # trtllm may not support sliding window
+        prefill_use_trtllm = (self.global_hyperparameters.window_left == -1
+                              and not cache_dtype.startswith("fp8")
+                              and use_trtllm_attention(
+                                num_prefill_tokens, max_seq_len, cache_dtype,
+                                num_qo_heads, num_kv_heads, head_dim))
+        decode_use_trtllm = (self.global_hyperparameters.window_left == -1
+                             and use_trtllm_attention(
+                                num_decode_tokens, max_seq_len, cache_dtype,
+                                num_qo_heads, num_kv_heads, head_dim))
 
         attn_metadata = FlashInferMetadata(
             num_actual_tokens=num_actual_tokens,