| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FlashInfer GQA builder: reorder threshold under DCP with spec decode.""" |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm.platforms import current_platform |
| 8 | |
| 9 | if not current_platform.is_cuda(): |
| 10 | pytest.skip("FlashInfer backend requires a CUDA platform.", allow_module_level=True) |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from tests.v1.attention.utils import create_vllm_config |
| 15 | from vllm.config import SpeculativeConfig, set_current_vllm_config |
| 16 | from vllm.v1.attention.backend import AttentionCGSupport |
| 17 | from vllm.v1.attention.backends import flashinfer as flashinfer_backend |
| 18 | from vllm.v1.attention.backends.flashinfer import ( |
| 19 | FlashInferDecodeKernel, |
| 20 | FlashInferMetadataBuilder, |
| 21 | ) |
| 22 | from vllm.v1.attention.backends.utils import PerLayerParameters |
| 23 | from vllm.v1.kv_cache_interface import FullAttentionSpec |
| 24 | |
| 25 | |
| 26 | def test_flashinfer_gqa_dcp_spec_decode_clamps_reorder_threshold(monkeypatch): |
| 27 | """trtllm-gen decode receives no cp_rank/global-seq-len information, so its |
| 28 | end-aligned causal mask is wrong for q_len > 1 over the DCP-interleaved |
| 29 | local KV shard. The builder must keep reorder_batch_threshold at 1 under |
| 30 | DCP so spec queries take the (DCP-aware) prefill path instead. |
| 31 | """ |
| 32 | vllm_config = create_vllm_config(max_model_len=1024) |
| 33 | vllm_config.parallel_config.decode_context_parallel_size = 2 |
| 34 | vllm_config.speculative_config = SpeculativeConfig( |
| 35 | method="ngram", num_speculative_tokens=3 |
| 36 | ) |
| 37 | |
| 38 | monkeypatch.setattr( |
| 39 | flashinfer_backend, "can_use_trtllm_attention", lambda *args, **kwargs: True |
| 40 | ) |
| 41 | monkeypatch.setattr( |
| 42 | FlashInferMetadataBuilder, |
| 43 | "_get_flashinfer_trtllm_api_decode_kernel", |
| 44 | staticmethod(lambda: FlashInferDecodeKernel.TRTLLM_GEN), |
| 45 | ) |
| 46 | monkeypatch.setattr( |
| 47 | flashinfer_backend, |
| 48 | "get_per_layer_parameters", |
| 49 | lambda *args, **kwargs: { |
| 50 | "layer.0": PerLayerParameters( |
| 51 | window_left=-1, logits_soft_cap=None, sm_scale=0.1, has_sinks=False |
| 52 | ) |
| 53 | }, |
| 54 | ) |
| 55 | |
| 56 | kv_cache_spec = FullAttentionSpec( |
| 57 | block_size=16, |
| 58 | num_kv_heads=vllm_config.model_config.get_num_kv_heads( |
| 59 | vllm_config.parallel_config |
| 60 | ), |
| 61 | head_size=vllm_config.model_config.get_head_size(), |
| 62 | dtype=vllm_config.model_config.dtype, |
| 63 | ) |
| 64 | with set_current_vllm_config(vllm_config): |
| 65 | builder = FlashInferMetadataBuilder( |
| 66 | kv_cache_spec, |
| 67 | ["layer.0"], |
| 68 | vllm_config, |
| 69 | torch.device("cpu"), |
| 70 | ) |
| 71 | |
| 72 | # Guard against passing vacuously with the kernel disabled. |
| 73 | assert ( |
| 74 | builder.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.TRTLLM_GEN |
| 75 | ) |
| 76 | assert builder.reorder_batch_threshold == 1 |
| 77 | assert ( |
| 78 | FlashInferMetadataBuilder.get_cudagraph_support(vllm_config, kv_cache_spec) |
| 79 | == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE |
| 80 | ) |