@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c30b405b8f

Fork 0

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 14, 05:30 UTC

[Spec Decode] Enable FlashInfer Spec Decoding (#25196)

Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai> Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Co-authored-by: lhsjohn <huashuoli@tencent.com>

Benjamin Chislett committed Sep 23, 2025 at 22:29 UTC c30b405b8f022d41f57042cad7032256d4831355

12 files changed +250 -49

tests/v1/attention/test_attention_splitting.py

+108 -1

@@ -9,7 +9,8 @@ from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata

9 from vllm.v1.attention.backends.utils import (UBatchSlice,

10 _make_metadata_with_slice,

11 slice_query_start_locs,

12 - split_attn_metadata)

12 + split_attn_metadata,

13 + split_decodes_and_prefills)

14 from vllm.v1.worker.ubatch_utils import create_ubatch_slices

@@ -158,6 +159,112 @@ def test_split_attn_metadata_decode_batch(large_decode_metadata):

159 assert torch.equal(results[1].seq_lens, torch.tensor([2048] * mid_point))

160

161

162 +def apply_split_decodes_and_prefills(query_lens: list[int],

163 + decode_threshold: int,

164 + require_uniform: bool):

165 + """Helper function to apply split_decodes_and_prefills and return

166 + the results."""

167 + device = torch.device("cpu")

168 + seq_lens = [10 * (i + 1) for i in range(len(query_lens))]

169 + common_metadata = create_common_attn_metadata(BatchSpec(

170 + seq_lens=seq_lens, query_lens=query_lens),

171 + block_size=16,

172 + device=device)

173 + return split_decodes_and_prefills(common_metadata,

174 + decode_threshold=decode_threshold,

175 + require_uniform=require_uniform)

176 +

177 +

178 +def test_split_decodes_and_prefills_nonuniform_all_ones():

179 + query_lens = [1, 1, 1]

180 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

181 + apply_split_decodes_and_prefills(query_lens, 1, False))

182 + assert num_decodes == 3

183 + assert num_prefills == 0

184 + assert num_decode_tokens == 3

185 + assert num_prefill_tokens == 0

186 +

187 +

188 +def test_split_decodes_and_prefills_nonuniform_all_short_decodes():

189 + query_lens = [1, 2, 1, 3, 2, 1, 2]

190 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

191 + apply_split_decodes_and_prefills(query_lens, 3, False))

192 + assert num_decodes == 7

193 + assert num_prefills == 0

194 + assert num_decode_tokens == sum(query_lens)

195 + assert num_prefill_tokens == 0

196 +

197 +

198 +def test_split_decodes_and_prefills_nonuniform_all_prefills():

199 + query_lens = [4, 5, 6, 7]

200 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

201 + apply_split_decodes_and_prefills(query_lens, 3, False))

202 + assert num_decodes == 0

203 + assert num_prefills == 4

204 + assert num_decode_tokens == 0

205 + assert num_prefill_tokens == sum(query_lens)

206 +

207 +

208 +def test_split_decodes_and_prefills_nonuniform_mixed_batch():

209 + query_lens = [2, 1, 3, 4, 5, 6, 7, 8]

210 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

211 + apply_split_decodes_and_prefills(query_lens, 4, False))

212 + assert num_decodes == 4 # 2, 1, 3, 4 are all <= 4

213 + assert num_prefills == 4 # 5, 6, 7, 8 are all > 4

214 + assert num_decode_tokens == 10 # 2 + 1 + 3 + 4

215 + assert num_prefill_tokens == 26 # 5 + 6 + 7 + 8

216 +

217 +

218 +def test_split_decodes_and_prefills_uniform_all_ones():

219 + query_lens = [1, 1, 1]

220 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

221 + apply_split_decodes_and_prefills(query_lens, 1, True))

222 + assert num_decodes == 3

223 + assert num_prefills == 0

224 + assert num_decode_tokens == 3

225 + assert num_prefill_tokens == 0

226 +

227 +

228 +def test_split_decodes_and_prefills_uniform_all_short_decodes():

229 + query_lens = [2, 2, 1, 3, 2, 1, 2]

230 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

231 + apply_split_decodes_and_prefills(query_lens, 3, True))

232 + assert num_decodes == 2

233 + assert num_prefills == 5

234 + assert num_decode_tokens == 4

235 + assert num_prefill_tokens == (1 + 3 + 2 + 1 + 2)

236 +

237 +

238 +def test_split_decodes_and_prefills_uniform_all_prefills():

239 + query_lens = [4, 5, 6, 7]

240 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

241 + apply_split_decodes_and_prefills(query_lens, 3, True))

242 + assert num_decodes == 0

243 + assert num_prefills == 4

244 + assert num_decode_tokens == 0

245 + assert num_prefill_tokens == sum(query_lens)

246 +

247 +

248 +def test_split_decodes_and_prefills_uniform_mixed_batch_all_uniform_decodes():

249 + query_lens = [2, 2, 2, 4, 5, 6, 7, 8]

250 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

251 + apply_split_decodes_and_prefills(query_lens, 4, True))

252 + assert num_decodes == 3 # 2, 2, 2 are all <= 4 and uniform

253 + assert num_prefills == 5 # 4, 5, 6, 7, 8 are all > 4

254 + assert num_decode_tokens == 6 # 2 + 2 + 2

255 + assert num_prefill_tokens == 30 # 4 + 5 + 6 + 7 + 8

256 +

257 +

258 +def test_split_decodes_and_prefills_uniform_mixed_batch_non_uniform_decodes():

259 + query_lens = [2, 1, 2, 4, 5, 6, 7, 8]

260 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

261 + apply_split_decodes_and_prefills(query_lens, 4, True))

262 + assert num_decodes == 1 # only the first 2 is taken as decode

263 + assert num_prefills == 7 # 1, 2, 4, 5, 6, 7, 8 are all > 4 or non-uniform

264 + assert num_decode_tokens == 2 # only the first 2

265 + assert num_prefill_tokens == (sum(query_lens) - 2) # rest of the tokens

266 +

267 +

268 @pytest.mark.parametrize(

269 "seq_lens,query_lens,split_point,expected_first_reqs,expected_second_reqs",

270 [

vllm/utils/flashinfer.py

+15

@@ -181,6 +181,12 @@ def force_use_trtllm_attention() -> Optional[bool]:

181 return _force_use_trtllm_attention(envs.VLLM_USE_TRTLLM_ATTENTION)

182

183

184 +def can_use_trtllm_attention(num_qo_heads: int, num_kv_heads: int) -> bool:

185 + """Check if the current configuration supports TRTLLM attention."""

186 + has_trtllm = supports_trtllm_attention()

187 + return has_trtllm and (num_qo_heads % num_kv_heads == 0)

188 +

189 +

190 def use_trtllm_attention(

191 num_qo_heads: int,

192 num_kv_heads: int,

@@ -188,7 +194,9 @@ def use_trtllm_attention(

194 max_seq_len: int,

195 kv_cache_dtype: str,

196 q_dtype: torch.dtype,

197 + is_prefill: bool,

198 has_sinks: bool = False,

199 + has_spec: bool = False,

200 ) -> bool:

201 """Return ``True`` if TRTLLM attention is used."""

202 force_use_trtllm = force_use_trtllm_attention()

@@ -214,6 +222,12 @@ def use_trtllm_attention(

222 )

223 return False

224

225 + if has_spec and not is_prefill:

226 + # Speculative decoding requires TRTLLM attention for decodes

227 + logger.info_once(

228 + "Using TRTLLM attention (enabled for speculative decoding).")

229 + return True

230 +

231 # Must use TRTLLM attention if query is FP8 quantized

232 if q_dtype == current_platform.fp8_dtype():

233 if has_sinks:

@@ -391,6 +405,7 @@ __all__ = [

405 "has_flashinfer_cutlass_fused_moe",

406 "has_nvidia_artifactory",

407 "supports_trtllm_attention",

408 + "can_use_trtllm_attention",

409 "use_trtllm_attention",

410 "flashinfer_disable_q_quantization",

411 "flashinfer_scaled_fp4_mm",

vllm/v1/attention/backends/flashinfer.py

+44 -11

@@ -25,7 +25,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

25 from vllm.platforms import current_platform

26 from vllm.triton_utils import tl, triton

27 from vllm.utils import cdiv, is_pin_memory_available

28 -from vllm.utils.flashinfer import (flashinfer_disable_q_quantization,

28 +from vllm.utils.flashinfer import (can_use_trtllm_attention,

29 + flashinfer_disable_q_quantization,

30 supports_trtllm_attention,

31 use_trtllm_attention)

32 from vllm.v1.attention.backends.flash_attn import use_cascade_attention

@@ -223,6 +224,7 @@ class FlashInferMetadata:

224

225 # For flashinfer trtllm batch decode

226 max_q_len: int

227 + max_q_len_prefill: int

228 max_seq_len: int

229 seq_lens: torch.Tensor

230 block_table_tensor: torch.Tensor

@@ -250,7 +252,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

252 cudagraph_support: ClassVar[AttentionCGSupport] = \

253 AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

254

253 - reorder_batch_threshold: ClassVar[int] = 1

255 + reorder_batch_threshold: int = 1

256

257 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

258 vllm_config: VllmConfig, device: torch.device):

@@ -302,6 +304,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

304 else:

305 self.q_data_type = self.model_config.dtype

306

307 + supports_spec_as_decode = \

308 + can_use_trtllm_attention(self.num_qo_heads, self.num_kv_heads)

309 + self._init_reorder_batch_threshold(1, supports_spec_as_decode)

310 +

311 self._cascade_wrapper = None # Wrapper for cascade attention

312

313 # Global hyperparameters shared by all attention layers

@@ -416,7 +422,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

422 num_actual_tokens = common_attn_metadata.num_actual_tokens

423 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens =\

424 split_decodes_and_prefills(common_attn_metadata,

419 - decode_threshold=self.reorder_batch_threshold)

425 + decode_threshold=self.reorder_batch_threshold,

426 + require_uniform=True)

427

428 page_size = self.page_size

429 max_q_len = common_attn_metadata.max_query_len

@@ -491,20 +498,25 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

498 paged_kv_last_page_len_np,

499 )

500

501 + uses_spec_reorder = self.reorder_batch_threshold > 1

502 prefill_use_trtllm = use_trtllm_attention(self.num_qo_heads,

503 self.num_kv_heads,

504 num_prefill_tokens,

505 max_seq_len,

506 self.cache_dtype,

507 self.q_data_type,

500 - has_sinks=self.has_sinks)

508 + is_prefill=True,

509 + has_sinks=self.has_sinks,

510 + has_spec=uses_spec_reorder)

511 decode_use_trtllm = use_trtllm_attention(self.num_qo_heads,

512 self.num_kv_heads,

513 num_decode_tokens,

514 max_seq_len,

515 self.cache_dtype,

516 self.q_data_type,

507 - has_sinks=self.has_sinks)

517 + is_prefill=False,

518 + has_sinks=self.has_sinks,

519 + has_spec=uses_spec_reorder)

520 if self.has_sinks and not (prefill_use_trtllm and decode_use_trtllm):

521 raise NotImplementedError(

522 "FlashInfer backend currently does not support attention "

@@ -521,6 +533,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

533 q_data_type=self.q_data_type,

534 slot_mapping=common_attn_metadata.slot_mapping,

535 max_q_len=max_q_len,

536 + max_q_len_prefill=max_q_len,

537 max_seq_len=max_seq_len,

538 seq_lens=seq_lens,

539 block_table_tensor=block_table_tensor,

@@ -577,6 +590,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

590 qo_indptr_cpu = qo_indptr_cpu[prefill_start:] - qo_indptr_cpu[

591 prefill_start]

592 paged_kv_indptr_cpu = paged_kv_indptr_cpu[prefill_start:]

593 +

594 + # Recompute max_q_len for the slice of requests we are using

595 + # for prefills. This can be different from max_q_len when

596 + # we have a non-uniform batch with some short decodes offloaded

597 + # to the prefill pathway

598 + query_lens_prefill = qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]

599 + attn_metadata.max_q_len_prefill = \

600 + int(query_lens_prefill.max().item())

601 +

602 if not attn_metadata.prefill_use_trtllm:

603 attn_metadata.prefill_wrapper.plan(

604 qo_indptr_cpu,

@@ -607,7 +629,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

629 num_decodes <= self._decode_cudagraph_max_bs)

630 if use_cudagraph:

631 num_input_tokens = (

610 - self.vllm_config.pad_for_cudagraph(num_decodes))

632 + self.vllm_config.pad_for_cudagraph(num_decode_tokens))

633 # Carefully fulfill the padding region with reasonable value

634 # on cpu.

635 # Make sure paged_kv_indptr_cpu is not decreasing

@@ -621,7 +643,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

643 num_decodes:num_input_tokens].fill_(1)

644

645 else:

624 - num_input_tokens = num_decodes

646 + num_input_tokens = num_decode_tokens

647

648 attn_metadata.decode_wrapper = self._get_decode_wrapper(

649 num_input_tokens, use_cudagraph)

@@ -842,6 +864,9 @@ class FlashInferImpl(AttentionImpl):

864 output.copy_(attn_metadata.cascade_wrapper.run(query, kv_cache))

865 return output

866

867 + # When using spec decoding, num_decodes can be < num_decode_tokens

868 + # because some decode requests may have more than one query token.

869 + num_decodes = attn_metadata.num_decodes

870 num_decode_tokens = attn_metadata.num_decode_tokens

871 num_prefill_tokens = attn_metadata.num_prefill_tokens

872

@@ -874,8 +899,8 @@ class FlashInferImpl(AttentionImpl):

899 prefill_query = prefill_query.contiguous()

900 workspace_buffer = _get_trtllm_gen_workspace_buffer()

901 block_tables_prefill = attn_metadata.block_table_tensor[

877 - num_decode_tokens:]

878 - seq_lens_prefill = attn_metadata.seq_lens[num_decode_tokens:]

902 + num_decodes:]

903 + seq_lens_prefill = attn_metadata.seq_lens[num_decodes:]

904

905 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND

906 assert get_kv_cache_layout() == "HND"

@@ -919,7 +944,7 @@ class FlashInferImpl(AttentionImpl):

944 workspace_buffer=workspace_buffer,

945 block_tables=mock_block_table,

946 seq_lens=seq_lens_prefill,

922 - max_q_len=attn_metadata.max_q_len,

947 + max_q_len=attn_metadata.max_q_len_prefill,

948 max_kv_len=attn_metadata.max_seq_len,

949 bmm1_scale=self.bmm1_scale,

950 bmm2_scale=self.bmm2_scale,

@@ -976,6 +1001,14 @@ class FlashInferImpl(AttentionImpl):

1001 assert self.o_sf_scale is None

1002 out = output[:num_decode_tokens]

1003

1004 + if num_decode_tokens % attn_metadata.num_decodes != 0:

1005 + # This gets triggered when the dummy_run forces

1006 + # attention to be initialized with q_len = 0

1007 + q_len_per_req = 1

1008 + else:

1009 + q_len_per_req = \

1010 + num_decode_tokens // attn_metadata.num_decodes

1011 +

1012 trtllm_batch_decode_with_kv_cache(

1013 query=decode_query,

1014 kv_cache=kv_cache_permute,

@@ -989,7 +1022,7 @@ class FlashInferImpl(AttentionImpl):

1022 sinks=self.sinks,

1023 o_sf_scale=self.o_sf_scale,

1024 out=out,

992 - )

1025 + q_len_per_req=q_len_per_req)

1026 return output_padded

1027

1028

vllm/v1/attention/backends/gdn_attn.py

+3 -3

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Backend for GatedDeltaNet attention."""

4 from dataclasses import dataclass

5 -from typing import ClassVar, Optional

5 +from typing import Optional

7 import torch

@@ -62,7 +62,7 @@ class GDNAttentionMetadataBuilder(

63 cudagraph_support = AttentionCGSupport.UNIFORM_BATCH

65 - reorder_batch_threshold: ClassVar[int] = 1

65 + reorder_batch_threshold: int = 1

67 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

68 vllm_config: VllmConfig, device: torch.device):

@@ -76,7 +76,7 @@ class GDNAttentionMetadataBuilder(

76 else:

77 self.num_spec = 0

78 self.use_spec_decode = self.num_spec > 0

79 - self.reorder_batch_threshold = self.num_spec + 1 # type: ignore[misc]

79 + self._init_reorder_batch_threshold(1, self.use_spec_decode)

81 self.use_full_cuda_graph = \

82 self.compilation_config.cudagraph_mode.has_full_cudagraphs()

vllm/v1/attention/backends/linear_attn.py

+1 -2

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 -from typing import ClassVar

5 import torch

@@ -35,7 +34,7 @@ class LinearAttentionMetadata:

34 class LinearAttentionMetadataBuilder(

35 AttentionMetadataBuilder[LinearAttentionMetadata]):

38 - reorder_batch_threshold: ClassVar[int] = 1

37 + reorder_batch_threshold: int = 1

39 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

40 vllm_config: VllmConfig, device: torch.device):

vllm/v1/attention/backends/mamba_attn.py

+1 -1

@@ -16,7 +16,7 @@ M = TypeVar("M")

18 class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

19 - reorder_batch_threshold: ClassVar[int] = 1

19 + reorder_batch_threshold: int = 1

20 cudagraph_support: ClassVar[AttentionCGSupport] = \

21 AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

vllm/v1/attention/backends/mla/common.py

+2 -2

@@ -190,7 +190,7 @@ return curr_o @ W_O

190 import functools

191 from abc import abstractmethod

192 from dataclasses import dataclass, field

193 -from typing import ClassVar, Generic, Optional, TypeVar, Union

193 +from typing import Generic, Optional, TypeVar, Union

194

195 import torch

196 from tqdm import tqdm

@@ -434,7 +434,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

434 NOTE: Please read the comment at the top of the file before trying to

435 understand this class

436 """

437 - reorder_batch_threshold: ClassVar[int] = 1

437 + reorder_batch_threshold: int = 1

438

439 @staticmethod

440 def determine_chunked_prefill_workspace_size(

vllm/v1/attention/backends/mla/flashattn_mla.py

+2 -2

@@ -64,7 +64,7 @@ class FlashAttnMLAMetadataBuilder(

64 cudagraph_support: ClassVar[AttentionCGSupport] = \

65 AttentionCGSupport.UNIFORM_BATCH

67 - reorder_batch_threshold: ClassVar[int] = 512

67 + reorder_batch_threshold: int = 512

69 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

70 vllm_config: VllmConfig, device: torch.device):

@@ -99,7 +99,7 @@ class FlashAttnMLAMetadataBuilder(

100 # TODO(lucas): Until we add support for the DCP custom masking we need

101 # to restrict decodes to q_len == 1 when DCP is enabled.

102 - self.__class__.reorder_batch_threshold = 1 \

102 + self.reorder_batch_threshold = 1 \

103 if get_dcp_group().world_size > 1 else self.reorder_batch_threshold

104

105 def _schedule_decode(self, num_reqs, cu_query_lens, max_query_len, seqlens,

vllm/v1/attention/backends/short_conv_attn.py

+2 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from dataclasses import dataclass

4 -from typing import ClassVar, Optional

4 +from typing import Optional

6 import torch

@@ -41,7 +41,7 @@ class ShortConvAttentionMetadata:

41 class ShortConvAttentionMetadataBuilder(

42 AttentionMetadataBuilder[ShortConvAttentionMetadata]):

44 - reorder_batch_threshold: ClassVar[int] = 1

44 + reorder_batch_threshold: int = 1

46 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

47 vllm_config: VllmConfig, device: torch.device):

vllm/v1/attention/backends/utils.py

+66 -6

@@ -236,7 +236,7 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

236 # Does this backend/builder reorder the batch?

237 # If not, set this to None. Otherwise set it to the query

238 # length that will be pulled into the front of the batch.

239 - reorder_batch_threshold: ClassVar[Optional[int]] = None

239 + reorder_batch_threshold: Optional[int] = None

240

241 @abstractmethod

242 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

@@ -246,6 +246,22 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

246 self.vllm_config = vllm_config

247 self.device = device

248

249 + def _init_reorder_batch_threshold(

250 + self,

251 + reorder_batch_threshold: int = 1,

252 + supports_spec_as_decode: bool = False) -> None:

253 + self.reorder_batch_threshold = reorder_batch_threshold

254 + if self.reorder_batch_threshold is not None \

255 + and supports_spec_as_decode:

256 + # If the backend supports spec-as-decode kernels, then we can set

257 + # the reorder_batch_threshold based on the number of speculative

258 + # tokens from the config.

259 + speculative_config = self.vllm_config.speculative_config

260 + if (speculative_config is not None

261 + and speculative_config.num_speculative_tokens is not None):

262 + self.reorder_batch_threshold = \

263 + 1 + speculative_config.num_speculative_tokens

264 +

265 @abstractmethod

266 def build(self,

267 common_prefix_len: int,

@@ -703,9 +719,9 @@ def subclass_attention_backend(

719

720

721 def split_decodes_and_prefills(

706 - common_attn_metadata: CommonAttentionMetadata,

707 - decode_threshold: int = 1,

708 -) -> tuple[int, int, int, int]:

722 + common_attn_metadata: CommonAttentionMetadata,

723 + decode_threshold: int = 1,

724 + require_uniform: bool = False) -> tuple[int, int, int, int]:

725 """

726 Assuming a reordered batch, finds the boundary between prefill and decode

727 requests.

@@ -714,6 +730,9 @@ def split_decodes_and_prefills(

730 common_attn_metadata: CommonAttentionMetadata object containing the

731 batch metadata.

732 decode_threshold: The maximum query length to be considered a decode.

733 + require_uniform: If True, requires that all decode requests have the

734 + same query length. When set, some queries may be considered prefills

735 + even if they are <= decode_threshold, in order to ensure uniformity.

736

737 Returns:

738 num_decodes: The number of decode requests.

@@ -726,11 +745,20 @@ def split_decodes_and_prefills(

745 num_tokens = common_attn_metadata.num_actual_tokens

746 query_start_loc = common_attn_metadata.query_start_loc_cpu

747

729 - if max_query_len <= decode_threshold:

748 + if max_query_len <= decode_threshold and \

749 + (not require_uniform or decode_threshold <= 1):

750 return num_reqs, 0, num_tokens, 0

751

752 query_lens = query_start_loc[1:] - query_start_loc[:-1]

733 - is_prefill = query_lens > decode_threshold

753 + if query_lens[0].item() > decode_threshold:

754 + # first request is not decode, so no decode requests

755 + return 0, num_reqs, 0, num_tokens

756 +

757 + if require_uniform:

758 + is_prefill = query_lens != query_lens[0]

759 + else:

760 + is_prefill = query_lens > decode_threshold

761 +

762 if not torch.any(is_prefill):

763 return num_reqs, 0, num_tokens, 0

764

@@ -806,6 +834,38 @@ def reorder_batch_to_split_decodes_and_prefills(

834 return modified_batch

835

836

837 +def reshape_query_for_spec_decode(query: torch.Tensor,

838 + batch_size: int) -> torch.Tensor:

839 + """

840 + Reshapes the query tensor for the specified batch size, so that

841 + it has shape (batch_size, seq_len, num_heads, head_dim).

842 + """

843 + assert query.dim() == 3, f"query must be 3D, got {query.dim()}D"

844 + total_tokens = query.shape[0]

845 + num_heads = query.shape[1]

846 + head_dim = query.shape[2]

847 + assert total_tokens % batch_size == 0, (

848 + f"{total_tokens=} is not divisible by {batch_size=}")

849 + seq_len = total_tokens // batch_size

850 + return query.view(batch_size, seq_len, num_heads, head_dim)

851 +

852 +

853 +def reshape_attn_output_for_spec_decode(

854 + attn_output: torch.Tensor) -> torch.Tensor:

855 + """

856 + Reshapes the attention output tensor, so that

857 + the batch_size and seq_len dimensions are combined.

858 + """

859 + if attn_output.dim() == 3:

860 + # Already in the correct shape

861 + return attn_output

862 + assert attn_output.dim() == 4, \

863 + f"attn_output must be 4D, got {attn_output.dim()}D"

864 + total_tokens = attn_output.shape[0] * attn_output.shape[1]

865 + return attn_output.view(total_tokens, attn_output.shape[2],

866 + attn_output.shape[3])

867 +

868 +

869 KV_SHARING_FAST_PREFILL_METADATA_FIELDS = [

870 ('logits_indices_padded', Optional[torch.Tensor], None),

871 ('num_logits_indices', int, 0),

vllm/v1/attention/backends/xformers.py

+2 -2

@@ -3,7 +3,7 @@

3 """Attention layer with XFormersAttention."""

5 from dataclasses import dataclass

6 -from typing import TYPE_CHECKING, ClassVar, Optional

6 +from typing import TYPE_CHECKING, Optional

8 import torch

@@ -197,7 +197,7 @@ class XFormersAttentionMetadata:

197 class XFormersAttentionMetadataBuilder(

198 AttentionMetadataBuilder[XFormersAttentionMetadata]):

199

200 - reorder_batch_threshold: ClassVar[int] = 1

200 + reorder_batch_threshold: int = 1

201

202 def __init__(

203 self,

vllm/v1/spec_decode/eagle.py

+4 -17

@@ -3,7 +3,7 @@

3 import ast

4 from dataclasses import replace

5 from importlib.util import find_spec

6 -from typing import Optional, Protocol

6 +from typing import Optional

8 import numpy as np

9 import torch

@@ -37,17 +37,6 @@ logger = init_logger(__name__)

37 PADDING_SLOT_ID = -1

40 -class EagleAttentionMetadata(Protocol):

41 - # Required attributes

42 - num_actual_tokens: int

43 - max_query_len: int

44 - query_start_loc: torch.Tensor

45 - max_seq_len: int

46 - seq_lens: torch.Tensor

47 - block_table: torch.Tensor

48 - slot_mapping: torch.Tensor

49 -

50 -

40 class EagleProposer:

42 def __init__(

@@ -120,7 +109,7 @@ class EagleProposer:

109 with_numpy=True)

110

111 # Determine allowed attention backends once during initialization.

123 - self.allowed_attn_types: tuple[type, ...]

112 + self.allowed_attn_types: Optional[tuple] = None

113 if current_platform.is_rocm():

114 rocm_types = [TritonAttentionMetadata, FlashAttentionMetadata]

115 # vllm.v1.attention.backends.rocm_aiter_fa is an optional backend

@@ -129,9 +118,6 @@ class EagleProposer:

118 AiterFlashAttentionMetadata)

119 rocm_types.append(AiterFlashAttentionMetadata)

120 self.allowed_attn_types = tuple(rocm_types)

132 - else:

133 - self.allowed_attn_types = (FlashAttentionMetadata,

134 - TreeAttentionMetadata)

121

122 # Parse the speculative token tree.

123 spec_token_tree = self.speculative_config.speculative_token_tree

@@ -266,7 +252,8 @@ class EagleProposer:

252

253 draft_token_ids = logits.argmax(dim=-1)

254

269 - if not isinstance(attn_metadata, self.allowed_attn_types):

255 + if self.allowed_attn_types is not None and \

256 + not isinstance(attn_metadata, self.allowed_attn_types):

257 raise ValueError(

258 f"Unsupported attention metadata type for speculative "

259 "decoding with num_speculative_tokens > 1: "