[Perf][Pooling] Avoid blocking seq_lens GPU-to-CPU copy for pooling in FlashInfer metadata builder (#57214)
Signed-off-by: frankwang28 <frank.wbb@hotmail.com> Co-authored-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Frank Wang committed
Sep 27, 2026 at 09:06 UTC
0376f81530336d90b3bba85cdea3284ff2ebde64
2 files changed
+97
-3
tests/v1/attention/test_attention_backends.py
+87
-1
@@ -19,7 +19,12 @@ from tests.v1.attention.utils import (
19
try_backend_includes_kv_cache_update,
20
try_get_attention_backend,
21
)
22
-from vllm.config import ModelConfig, set_current_vllm_config
22
+from vllm.config import (
23
+ DiffusionConfig,
24
+ ModelConfig,
25
+ SpeculativeConfig,
26
+ set_current_vllm_config,
27
+)
28
from vllm.platforms import current_platform
29
from vllm.utils.math_utils import cdiv
30
from vllm.utils.torch_utils import (
@@ -1069,6 +1074,87 @@ def test_flashinfer_trtllm_gen_padded_decode_uses_varlen_offsets(
1074
assert q_cu_seq_lens.tolist() == qo_indptr_values
1075
1076
1077
+@pytest.mark.skipif(
1078
+ not current_platform.is_cuda()
1079
+ or AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
1080
+ reason="FlashInfer with CUDA is required.",
1081
+)
1082
+@pytest.mark.parametrize(
1083
+ "runner_type,spec_kind,has_upper_bound,expected_copies",
1084
+ [
1085
+ ("generate", None, True, 0),
1086
+ ("pooling", None, True, 0),
1087
+ ("generate", None, False, 1),
1088
+ ("generate", "spec_decode", True, 1),
1089
+ ("generate", "diffusion", True, 1),
1090
+ ],
1091
+)
1092
+def test_flashinfer_avoids_seq_lens_copy_without_spec_tokens(
1093
+ monkeypatch, runner_type, spec_kind, has_upper_bound, expected_copies
1094
+):
1095
+ """Planning reuses the host seq_lens bound only when it is exact.
1096
+
1097
+ Without speculative token accounting (no spec-decode drafts, no dLLM
1098
+ canvas tokens) scheduled tokens always land in KV, so the CPU upper bound
1099
+ is exact regardless of runner type. Any speculative accounting keeps the
1100
+ bound optimistic (rejected tokens are in-flight) and forces the copy.
1101
+ """
1102
+ import unittest.mock
1103
+
1104
+ from vllm.v1.attention.backends import flashinfer as flashinfer_backend
1105
+ from vllm.v1.attention.backends.utils import PerLayerParameters
1106
+
1107
+ monkeypatch.setattr(
1108
+ flashinfer_backend,
1109
+ "get_per_layer_parameters",
1110
+ lambda *args: {
1111
+ "layer.0": PerLayerParameters(
1112
+ window_left=-1, logits_soft_cap=0.0, sm_scale=0.125
1113
+ )
1114
+ },
1115
+ )
1116
+
1117
+ config = create_vllm_config(
1118
+ model_name="Qwen/Qwen3-0.6B", max_model_len=64, max_num_seqs=3
1119
+ )
1120
+ config.model_config.runner_type = runner_type
1121
+ config.attention_config.use_trtllm_attention = False
1122
+ if spec_kind == "spec_decode":
1123
+ config.speculative_config = SpeculativeConfig(
1124
+ method="ngram", num_speculative_tokens=3
1125
+ )
1126
+ elif spec_kind == "diffusion":
1127
+ config.diffusion_config = DiffusionConfig(canvas_length=8)
1128
+ device = torch.device(f"{DEVICE_TYPE}:0")
1129
+ # Include a one-token final chunk and both partial and full KV pages.
1130
+ common_attn_metadata = create_common_attn_metadata(
1131
+ BatchSpec(seq_lens=[17, 20, 32], query_lens=[1, 4, 16]), 16, device
1132
+ )
1133
+ if not has_upper_bound:
1134
+ common_attn_metadata.seq_lens_cpu_upper_bound = None
1135
+ elif spec_kind is not None:
1136
+ # Simulate rejected speculative tokens: CPU says 18, GPU says 17.
1137
+ assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
1138
+ common_attn_metadata.seq_lens_cpu_upper_bound[0] += 1
1139
+
1140
+ with set_current_vllm_config(config):
1141
+ builder = flashinfer_backend.FlashInferMetadataBuilder(
1142
+ create_standard_kv_cache_spec(config), ["layer.0"], config, device
1143
+ )
1144
+ with unittest.mock.patch.object(
1145
+ common_attn_metadata.seq_lens,
1146
+ "cpu",
1147
+ wraps=common_attn_metadata.seq_lens.cpu,
1148
+ ) as copy_to_cpu:
1149
+ metadata = builder.build(0, common_attn_metadata)
1150
+
1151
+ assert copy_to_cpu.call_count == expected_copies
1152
+ assert metadata.num_decodes == 1
1153
+ assert metadata.num_prefills == 2
1154
+ assert builder.paged_kv_indptr.cpu[:4].tolist() == [0, 2, 4, 6]
1155
+ assert builder.paged_kv_last_page_len.cpu[:3].tolist() == [1, 4, 16]
1156
+
1157
+
1158
@pytest.mark.skipif(
1159
AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
1160
reason="FlashInfer is not available.",
vllm/v1/attention/backends/flashinfer.py
+10
-2
@@ -700,6 +700,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
700
self.cache_config = vllm_config.cache_config
701
self.model_config = vllm_config.model_config
702
self.attention_config = vllm_config.attention_config
703
+ self._num_speculative_tokens = vllm_config.num_speculative_tokens
704
self._workspace_buffer = None
705
self._prefill_wrapper: (
706
BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None
@@ -1446,8 +1447,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1447
# (block_tables, seq_lens) directly.
1448
needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm
1449
if needs_seq_lens_cpu:
1449
- with gpu_sync_allowed():
1450
- seq_lens_cpu = common_attn_metadata.seq_lens.cpu()
1450
+ if (
1451
+ self._num_speculative_tokens == 0
1452
+ and common_attn_metadata.seq_lens_cpu_upper_bound is not None
1453
+ ):
1454
+ # No speculative token accounting, so the upper bound is exact.
1455
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
1456
+ else:
1457
+ with gpu_sync_allowed():
1458
+ seq_lens_cpu = common_attn_metadata.seq_lens.cpu()
1459
else:
1460
seq_lens_cpu = None
1461