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

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

[V0 Deprecation] Remove V0 FlashInfer attention backend (#22776)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Aug 18, 2025 at 19:54 UTC 14006840eacf74f83e0d486eca6a24e75cafa6d3

8 files changed +9 -1133

tests/basic_correctness/test_basic_correctness.py

+1 -8

@@ -12,7 +12,6 @@ import pytest

12 import torch

14 from vllm import LLM, envs

15 -from vllm.platforms import current_platform

15 from vllm.v1.engine.llm_engine import LLMEngine as LLMEngineV1

17 from ..conftest import HfRunner, VllmRunner

@@ -78,11 +77,7 @@ def test_models(

77 "VLLM_USE_V1") and envs.VLLM_USE_V1:

78 pytest.skip("enable_prompt_embeds is not supported in v1.")

81 - if backend == "FLASHINFER" and current_platform.is_rocm():

82 - pytest.skip("Flashinfer does not support ROCm/HIP.")

83 -

84 - if backend in ("XFORMERS",

85 - "FLASHINFER") and model == "google/gemma-2-2b-it":

80 + if backend == "XFORMERS" and model == "google/gemma-2-2b-it":

81 pytest.skip(

82 f"{backend} does not support gemma2 with full context length.")

@@ -141,8 +136,6 @@ def test_models(

136 ("meta-llama/Llama-3.2-1B-Instruct", "mp", "", "L4", {}),

137 ("distilbert/distilgpt2", "ray", "", "A100", {}),

138 ("distilbert/distilgpt2", "mp", "", "A100", {}),

144 - ("distilbert/distilgpt2", "mp", "FLASHINFER", "A100", {}),

145 - ("meta-llama/Meta-Llama-3-8B", "ray", "FLASHINFER", "A100", {}),

139 ])

140 @pytest.mark.parametrize("enable_prompt_embeds", [True, False])

141 def test_models_distributed(

tests/compile/test_basic_correctness.py

+1 -1

@@ -34,7 +34,7 @@ class TestSetting:

34 model_args=["--max-model-len", "2048"],

35 pp_size=2,

36 tp_size=2,

37 - attn_backend="FLASHINFER",

37 + attn_backend="FLASH_ATTN",

38 method="generate",

39 fullgraph=True,

40 ),

tests/core/block/e2e/test_correctness_sliding_window.py

+2 -6

@@ -32,7 +32,7 @@ BLOCK_SIZE = 16

32 @pytest.mark.parametrize("test_llm_kwargs", [{}])

33 @pytest.mark.parametrize("batch_size", [5])

34 @pytest.mark.parametrize("seed", [1])

35 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"])

35 +@pytest.mark.parametrize("backend", ["FLASH_ATTN", "XFORMERS"])

36 def test_sliding_window_retrieval(baseline_llm_generator, test_llm_generator,

37 batch_size, seed, backend, monkeypatch):

38 """

@@ -43,8 +43,6 @@ def test_sliding_window_retrieval(baseline_llm_generator, test_llm_generator,

44 Additionally, we compare the results of the v1 and v2 managers.

45 """

46 - if backend == "FLASHINFER" and current_platform.is_rocm():

47 - pytest.skip("Flashinfer does not support ROCm/HIP.")

46 if backend == "XFORMERS" and current_platform.is_rocm():

47 pytest.skip("Xformers does not support ROCm/HIP.")

@@ -96,7 +94,7 @@ def test_sliding_window_retrieval(baseline_llm_generator, test_llm_generator,

94 @pytest.mark.parametrize("test_llm_kwargs", [{"enable_chunked_prefill": True}])

95 @pytest.mark.parametrize("batch_size", [5])

96 @pytest.mark.parametrize("seed", [1])

99 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER", "XFORMERS"])

97 +@pytest.mark.parametrize("backend", ["FLASH_ATTN", "XFORMERS"])

98 def test_sliding_window_chunked_prefill(test_llm_generator, batch_size, seed,

99 backend, monkeypatch):

100 """

@@ -107,8 +105,6 @@ def test_sliding_window_chunked_prefill(test_llm_generator, batch_size, seed,

105 The results with and without chunked prefill are not the same due to

106 numerical instabilities.

107 """

110 - if backend == "FLASHINFER" and current_platform.is_rocm():

111 - pytest.skip("Flashinfer does not support ROCm/HIP.")

108 if backend == "XFORMERS" and current_platform.is_rocm():

109 pytest.skip("Xformers does not support ROCm/HIP.")

110 override_backend_env_variable(monkeypatch, backend)

tests/distributed/test_pp_cudagraph.py

-1

@@ -17,7 +17,6 @@ if TYPE_CHECKING:

17 ])

18 @pytest.mark.parametrize("ATTN_BACKEND", [

19 "FLASH_ATTN",

20 - "FLASHINFER",

20 ])

21 @create_new_process_for_each_test()

22 def test_pp_cudagraph(

tests/kernels/attention/test_attention_selector.py

@@ -81,6 +81,9 @@ def test_env(

81 m.setenv(STR_BACKEND_ENV_VAR, name)

82 m.setenv("VLLM_MLA_DISABLE", "1" if use_mla else "0")

84 + if name == "FLASHINFER" and not use_v1:

85 + pytest.skip("FlashInfer backend is only available on V1 engine")

86 +

87 if device == "cpu":

88 if not use_v1:

89 pytest.skip("CPU backend only supports V1")

tests/models/quantization/test_fp8.py

+1 -4

@@ -32,7 +32,7 @@ from ..utils import check_logprobs_close

32 # Due to low-precision numerical divergence, we only test logprob of 4 tokens

33 @pytest.mark.parametrize("max_tokens", [4])

34 @pytest.mark.parametrize("enforce_eager", [True])

35 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "XFORMERS", "FLASHINFER"])

35 +@pytest.mark.parametrize("backend", ["FLASH_ATTN", "XFORMERS"])

36 # NOTE: Increasing this in this suite will fail CI because we currently cannot

37 # reset distributed env properly. Use a value > 1 just when you test.

38 @pytest.mark.parametrize("tensor_parallel_size", [1])

@@ -57,9 +57,6 @@ def test_models(

57 numerical sensitive kernels.

58 """

60 - if backend == "FLASHINFER" and current_platform.is_rocm():

61 - pytest.skip("Flashinfer does not support ROCm/HIP.")

62 -

60 if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():

61 pytest.skip(

62 f"{kv_cache_dtype} is currently not supported on ROCm/HIP.")

vllm/attention/backends/flashinfer.py deleted

-1098

@@ -1,1098 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import dataclasses

5 -from collections import defaultdict

6 -from contextlib import contextmanager

7 -from dataclasses import dataclass

8 -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Type

9 -

10 -from vllm.multimodal import MultiModalPlaceholderMap

11 -

12 -try:

13 - from flashinfer import BatchDecodeWithPagedKVCacheWrapper

14 - from flashinfer.decode import (CUDAGraphBatchDecodeWithPagedKVCacheWrapper,

15 - trtllm_batch_decode_with_kv_cache)

16 - from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper

17 -

18 - from vllm.vllm_flash_attn import flash_attn_varlen_func

19 - FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024

20 -except ImportError:

21 - # Avoid turning these types into variables during type checking

22 - if not TYPE_CHECKING:

23 - BatchDecodeWithPagedKVCacheWrapper = None

24 - CUDAGraphBatchDecodeWithPagedKVCacheWrapper = None

25 - BatchPrefillWithPagedKVCacheWrapper = None

26 - trtllm_batch_decode_with_kv_cache = None

27 - FLASHINFER_WORKSPACE_BUFFER_SIZE = 0

28 - raise ImportError("FlashInfer is not installed. Please install it from "

29 - "https://github.com/flashinfer-ai/flashinfer") from None

30 -

31 -import torch

32 -

33 -import vllm.envs as envs

34 -from vllm import _custom_ops as ops

35 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

36 - AttentionLayer,

37 - AttentionMetadata,

38 - AttentionMetadataBuilder,

39 - AttentionState, AttentionType)

40 -from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,

41 - compute_slot_mapping_start_idx,

42 - is_block_tables_empty)

43 -from vllm.attention.layer import Attention

44 -from vllm.attention.ops.paged_attn import PagedAttention

45 -from vllm.config import VllmConfig, get_layers_from_vllm_config

46 -from vllm.logger import init_logger

47 -from vllm.utils import (async_tensor_h2d, get_kv_cache_torch_dtype,

48 - make_tensor_with_pad)

49 -from vllm.utils.flashinfer import use_trtllm_attention

50 -

51 -logger = init_logger(__name__)

52 -

53 -if TYPE_CHECKING:

54 - from vllm.worker.model_runner import ModelInputForGPUBuilder

55 -

56 -

57 -class FlashInferBackend(AttentionBackend):

58 -

59 - @staticmethod

60 - def get_name() -> str:

61 - return "FLASHINFER"

62 -

63 - @staticmethod

64 - def get_impl_cls() -> Type["FlashInferImpl"]:

65 - return FlashInferImpl

66 -

67 - @staticmethod

68 - def get_metadata_cls() -> Type["AttentionMetadata"]:

69 - return FlashInferMetadata

70 -

71 - @staticmethod

72 - def get_builder_cls() -> Type["FlashInferMetadataBuilder"]:

73 - return FlashInferMetadataBuilder

74 -

75 - @staticmethod

76 - def get_state_cls() -> Type["FlashInferState"]:

77 - return FlashInferState

78 -

79 - @staticmethod

80 - def get_kv_cache_shape(

81 - num_blocks: int,

82 - block_size: int,

83 - num_kv_heads: int,

84 - head_size: int,

85 - ) -> Tuple[int, ...]:

86 - return (num_blocks, 2, block_size, num_kv_heads, head_size)

87 -

88 - @staticmethod

89 - def get_kv_cache_stride_order() -> Tuple[int, ...]:

90 - cache_layout = FlashInferState.get_kv_cache_layout()

91 - assert (cache_layout in ("NHD", "HND"))

92 - stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3,

93 - 2, 4)

94 - return stride_order

95 -

96 - @staticmethod

97 - def swap_blocks(

98 - src_kv_cache: torch.Tensor,

99 - dst_kv_cache: torch.Tensor,

100 - src_to_dst: torch.Tensor,

101 - ) -> None:

102 - PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)

103 -

104 - @staticmethod

105 - def copy_blocks(

106 - kv_caches: List[torch.Tensor],

107 - src_to_dists: torch.Tensor,

108 - ) -> None:

109 - PagedAttention.copy_blocks(kv_caches, src_to_dists)

110 -

111 - @staticmethod

112 - def get_supported_head_sizes() -> List[int]:

113 - return [64, 128, 256]

114 -

115 - @staticmethod

116 - def get_fp8_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:

117 - if kv_cache_dtype in ("fp8", "fp8_e4m3"):

118 - return torch.float8_e4m3fn

119 - elif kv_cache_dtype == "fp8_e5m2":

120 - return torch.float8_e5m2

121 - else:

122 - raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")

123 -

124 -

125 -@dataclass

126 -class PerLayerParameters:

127 - """

128 - Currently, FlashInfer backend only support models in which all layers share

129 - the same values for the following hyperparameters.

130 - """

131 -

132 - window_left: int

133 - logits_soft_cap: Optional[float]

134 - sm_scale: float

135 -

136 -

137 -def get_per_layer_parameters(

138 - vllm_config: VllmConfig) -> Dict[str, PerLayerParameters]:

139 - """

140 - Scan all attention layers and determine some hyperparameters

141 - to use during `plan`.

142 - """

143 -

144 - layers = get_layers_from_vllm_config(vllm_config, Attention)

145 - per_layer_params: Dict[str, PerLayerParameters] = {}

146 -

147 - for key, layer in layers.items():

148 - impl = layer.impl

149 - assert isinstance(impl, FlashInferImpl)

150 -

151 - # Infer hyperparameters from the attention layer

152 - window_size = impl.sliding_window

153 - window_left = window_size[0] if window_size is not None else -1

154 - logits_soft_cap = impl.logits_soft_cap

155 - sm_scale = impl.scale

156 -

157 - per_layer_params[key] = PerLayerParameters(window_left,

158 - logits_soft_cap, sm_scale)

159 -

160 - return per_layer_params

161 -

162 -

163 -def infer_global_hyperparameters(

164 - per_layer_params: Dict[str, PerLayerParameters]) -> PerLayerParameters:

165 - """

166 - Currently, FlashInfer backend only support models in which all layers share

167 - the same values for the following hyperparameters:

168 - - `window_left`

169 - - `logits_soft_cap`

170 - - `sm_scale`

171 -

172 - So this function asserts that all layers share the same values for these

173 - hyperparameters and returns the global values.

174 - """

175 -

176 - assert len(per_layer_params) > 0, "No attention layers found in the model."

177 -

178 - param_sets = list(per_layer_params.values())

179 - global_params = param_sets[0]

180 - for params in param_sets:

181 - assert params == global_params, (

182 - "FlashInfer backend currently only supports models in which all "

183 - "layers share the same values for the following hyperparameters: "

184 - "`window_left`, `logits_soft_cap`, `sm_scale`.")

185 -

186 - return global_params

187 -

188 -

189 -class FlashInferState(AttentionState):

190 -

191 - def __init__(self, runner):

192 - self.runner = runner

193 - self._is_graph_capturing = False

194 - self._workspace_buffer = None

195 - self._decode_wrapper = None

196 - self._prefill_wrapper = None

197 -

198 - # Global hyperparameters shared by all attention layers

199 - self.global_hyperparameters: Optional[PerLayerParameters] = None

200 -

201 - self.vllm_config = self.runner.vllm_config

202 - self._kv_cache_layout = None

203 -

204 - def _get_workspace_buffer(self):

205 - if self._workspace_buffer is None:

206 - self._workspace_buffer = torch.zeros(

207 - FLASHINFER_WORKSPACE_BUFFER_SIZE,

208 - dtype=torch.uint8,

209 - device=self.runner.device)

210 - return self._workspace_buffer

211 -

212 - @staticmethod

213 - def get_kv_cache_layout():

214 - from vllm.v1.attention.backends.utils import _KV_CACHE_LAYOUT_OVERRIDE

215 - if _KV_CACHE_LAYOUT_OVERRIDE is not None:

216 - logger.info_once("Using KV cache layout %s",

217 - _KV_CACHE_LAYOUT_OVERRIDE)

218 - return _KV_CACHE_LAYOUT_OVERRIDE

219 - cache_layout = envs.VLLM_KV_CACHE_LAYOUT

220 - if cache_layout is None:

221 - logger.info_once("Using default KV cache layout NHD")

222 - return "NHD"

223 - logger.info_once("Using KV cache layout %s", cache_layout)

224 - return cache_layout

225 -

226 - def _get_prefill_wrapper(self):

227 - if self._prefill_wrapper is None:

228 - self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(

229 - self._get_workspace_buffer(), self.get_kv_cache_layout())

230 - return self._prefill_wrapper

231 -

232 - def _get_decode_wrapper(self):

233 - if self._decode_wrapper is None:

234 - num_qo_heads = (self.runner.model_config.get_num_attention_heads(

235 - self.runner.parallel_config))

236 - num_kv_heads = self.runner.model_config.get_num_kv_heads(

237 - self.runner.parallel_config)

238 - use_tensor_cores = envs.VLLM_FLASHINFER_FORCE_TENSOR_CORES or (

239 - num_qo_heads // num_kv_heads > 4)

240 - self._decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(

241 - self._get_workspace_buffer(),

242 - self.get_kv_cache_layout(),

243 - use_tensor_cores=use_tensor_cores)

244 - return self._decode_wrapper

245 -

246 - @contextmanager

247 - def graph_capture(self, max_batch_size: int):

248 - self._is_graph_capturing = True

249 - self._graph_decode_wrapper = None

250 - self._graph_slot_mapping = torch.full((max_batch_size, ),

251 - PAD_SLOT_ID,

252 - dtype=torch.long,

253 - device=self.runner.device)

254 - self._graph_seq_lens = torch.ones(max_batch_size,

255 - dtype=torch.int32,

256 - device=self.runner.device)

257 - self._graph_block_tables = torch.from_numpy(

258 - self.runner.graph_block_tables).to(device=self.runner.device)

259 - self._graph_decode_workspace_buffer = self._get_workspace_buffer()

260 - self._graph_indices_buffer = torch.empty(

261 - max_batch_size * self.runner.cache_config.num_gpu_blocks,

262 - dtype=torch.int32,

263 - device=self.runner.device)

264 - self._graph_indptr_buffer = torch.empty(max_batch_size + 1,

265 - dtype=torch.int32,

266 - device=self.runner.device)

267 - self._graph_last_page_len_buffer = torch.empty(

268 - max_batch_size, dtype=torch.int32, device=self.runner.device)

269 - yield

270 - self._is_graph_capturing = False

271 - del self._graph_slot_mapping

272 - del self._graph_seq_lens

273 - del self._graph_block_tables

274 - del self._graph_decode_workspace_buffer

275 - del self._graph_indices_buffer

276 - del self._graph_indptr_buffer

277 - del self._graph_last_page_len_buffer

278 - del self._graph_decode_wrapper

279 -

280 - def graph_clone(self, batch_size: int):

281 - assert self._is_graph_capturing

282 - state = self.__class__(self.runner)

283 - state._workspace_buffer = self._graph_decode_workspace_buffer

284 - state._decode_wrapper = self._graph_decode_wrapper

285 - state._prefill_wrapper = self._get_prefill_wrapper()

286 - return state

287 -

288 - def graph_capture_get_metadata_for_batch(

289 - self, batch_size: int, is_encoder_decoder_model: bool = False):

290 - assert self._is_graph_capturing

291 - _indptr_buffer = self._graph_indptr_buffer[:batch_size + 1]

292 - _last_page_len_buffer = self._graph_last_page_len_buffer[:batch_size]

293 -

294 - num_qo_heads = (self.runner.model_config.get_num_attention_heads(

295 - self.runner.parallel_config))

296 - num_kv_heads = self.runner.model_config.get_num_kv_heads(

297 - self.runner.parallel_config)

298 - use_tensor_cores = envs.VLLM_FLASHINFER_FORCE_TENSOR_CORES or (

299 - num_qo_heads // num_kv_heads > 4)

300 - self._graph_decode_wrapper = \

301 - CUDAGraphBatchDecodeWithPagedKVCacheWrapper(

302 - self._graph_decode_workspace_buffer, _indptr_buffer,

303 - self._graph_indices_buffer, _last_page_len_buffer,

304 - self.get_kv_cache_layout(),

305 - use_tensor_cores)

306 - if self.runner.kv_cache_dtype.startswith("fp8"):

307 - kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(

308 - self.runner.kv_cache_dtype)

309 - else:

310 - kv_cache_dtype = get_kv_cache_torch_dtype(

311 - self.runner.kv_cache_dtype, self.runner.model_config.dtype)

312 -

313 - paged_kv_indptr_tensor_host = torch.arange(0,

314 - batch_size + 1,

315 - dtype=torch.int32)

316 - paged_kv_indices_tensor_host = torch.arange(0,

317 - batch_size,

318 - dtype=torch.int32)

319 - paged_kv_last_page_len_tensor_host = torch.full((batch_size, ),

320 - self.runner.block_size,

321 - dtype=torch.int32)

322 - query_start_loc_host = torch.arange(0,

323 - batch_size + 1,

324 - dtype=torch.int32)

325 -

326 - global_params = infer_global_hyperparameters(

327 - get_per_layer_parameters(self.vllm_config))

328 -

329 - attn_metadata = self.runner.attn_backend.make_metadata(

330 - num_prefills=0,

331 - slot_mapping=self._graph_slot_mapping[:batch_size],

332 - multi_modal_placeholder_index_maps=None,

333 - enable_kv_scales_calculation=False,

334 - num_prefill_tokens=0,

335 - num_decode_tokens=batch_size,

336 - max_prefill_seq_len=0,

337 - max_decode_seq_len=0,

338 - seq_lens_tensor=self._graph_seq_lens,

339 - block_tables=self._graph_block_tables,

340 - paged_kv_indptr=paged_kv_indptr_tensor_host,

341 - paged_kv_indices=paged_kv_indices_tensor_host,

342 - paged_kv_last_page_len=paged_kv_last_page_len_tensor_host,

343 - num_qo_heads=num_qo_heads,

344 - num_kv_heads=num_kv_heads,

345 - head_dim=self.runner.model_config.get_head_size(),

346 - page_size=self.runner.block_size,

347 - seq_start_loc=None,

348 - query_start_loc=query_start_loc_host,

349 - device=self.runner.device,

350 - data_type=kv_cache_dtype,

351 - q_data_type=self.runner.model_config.dtype,

352 - use_cuda_graph=True,

353 - decode_wrapper=self._graph_decode_wrapper,

354 - prefill_wrapper=None,

355 - **dataclasses.asdict(global_params),

356 - )

357 - attn_metadata.begin_forward()

358 - return attn_metadata

359 -

360 - def get_graph_input_buffers(self,

361 - attn_metadata,

362 - is_encoder_decoder_model: bool = False):

363 - return {

364 - "block_tables": attn_metadata.block_tables,

365 - "seq_lens_tensor": attn_metadata.seq_lens_tensor,

366 - "slot_mapping": attn_metadata.slot_mapping,

367 - }

368 -

369 - def prepare_graph_input_buffers(self,

370 - input_buffers,

371 - attn_metadata,

372 - is_encoder_decoder_model: bool = False):

373 - # FlashInfer-specific logic: copy additional tensors

374 - num_total_blocks = attn_metadata.decode_metadata.seq_lens_tensor.shape[

375 - 0]

376 - input_buffers["seq_lens_tensor"][:num_total_blocks].copy_(

377 - attn_metadata.seq_lens_tensor, non_blocking=True)

378 - input_buffers["block_tables"][:num_total_blocks].copy_(

379 - attn_metadata.block_tables, non_blocking=True)

380 -

381 - def begin_forward(self, model_input):

382 - assert not self._is_graph_capturing

383 - state = self

384 - use_cuda_graph = model_input.attn_metadata.use_cuda_graph

385 - is_decode = model_input.attn_metadata.num_prefills == 0

386 - # In case of multistep chunked-prefill, there might be prefill requests

387 - # scheduled while CUDA graph mode is enabled. We don't run graph in that

388 - # case.

389 - if use_cuda_graph and is_decode:

390 - if model_input.inputs_embeds is None:

391 - batch_size = model_input.input_tokens.shape[0]

392 - state = (

393 - self.runner.graph_runners[model_input.virtual_engine][(

394 - batch_size, False)].attn_state)

395 - else:

396 - batch_size = model_input.inputs_embeds.shape[0]

397 - state = (

398 - self.runner.graph_runners[model_input.virtual_engine][(

399 - batch_size, True)].attn_state)

400 -

401 - model_input.attn_metadata.prefill_wrapper = state._get_prefill_wrapper(

402 - )

403 - model_input.attn_metadata.decode_wrapper = state._get_decode_wrapper()

404 - model_input.attn_metadata.begin_forward()

405 -

406 -

407 -@dataclass

408 -class FlashInferMetadata(AttentionMetadata):

409 - # Maximum sequence length among prefill batch. 0 if there are decoding

410 - # requests only.

411 - max_prefill_seq_len: int

412 - max_decode_seq_len: int

413 -

414 - # Number of query tokens for each request in the batch.

415 - # Currently, we require that all requests have the same number of query

416 - # tokens during the decoding phase. When speculavie decoding is enabled,

417 - # decode_query_len might be greater than 1. In all other cases, it is 1.

418 - decode_query_len: Optional[int] = 1

419 -

420 - use_cuda_graph: bool = True

421 -

422 - prefill_wrapper: Optional[BatchPrefillWithPagedKVCacheWrapper] = None

423 - decode_wrapper: Optional[BatchDecodeWithPagedKVCacheWrapper] = None

424 -

425 - # Metadata for the prefill stage

426 - seq_start_loc: Optional[torch.Tensor] = None

427 - query_start_loc: Optional[torch.Tensor] = None

428 - block_tables: Optional[torch.Tensor] = None

429 -

430 - # used for GPU operations

431 - seq_lens_tensor: Optional[torch.Tensor] = None

432 - block_table_bound: Optional[torch.Tensor] = None

433 -

434 - # An example for paged_kv_indices, paged_kv_indptr:

435 - # request 1, page indices [0, 5, 8]

436 - # request 2, page indices [1, 6, 7]

437 - # request 3, page indices [3, 4]

438 - # paged_kv_indices is a concatenation of page indices of all requests:

439 - # [0, 5, 8, 1, 6, 7, 3, 4]

440 - # paged_kv_indptr is used to index into paged_kv_indices:

441 - # [0, 3, 6, 8]

442 - # The indptr of the paged kv cache, shape: [batch_size + 1]

443 - paged_kv_indptr: Optional[torch.Tensor] = None

444 - # The page indices of the paged kv cache

445 - paged_kv_indices: Optional[torch.Tensor] = None

446 - # The number of entries in the last page of each request in

447 - # the paged kv cache, shape: [batch_size]

448 - paged_kv_last_page_len: Optional[torch.Tensor] = None

449 - # The number of query/output heads

450 - num_qo_heads: Optional[int] = None

451 - # The number of key/value heads

452 - num_kv_heads: Optional[int] = None

453 - # The dimension of the attention heads

454 - head_dim: Optional[int] = None

455 - # Block size of vllm

456 - page_size: Optional[int] = None

457 - # The data type of the paged kv cache

458 - data_type: torch.dtype = None

459 - # The data type of the query

460 - q_data_type: torch.dtype = None

461 - # FlashInfer 0.2 encourages passing host tensors

462 - device: torch.device = torch.device("cpu")

463 - is_profile_run: bool = False

464 -

465 - # The FlashInfer backend currently supports only models in which all layers

466 - # share the same following hyperparameters:

467 -

468 - # The left (inclusive) window size for the attention window, when

469 - # set to `-1`, the window size will be set to the full length of

470 - # the sequence. Defaults to `-1`.

471 - window_left: int = -1

472 - # The attention logits soft capping value (used in Gemini, Grok and

473 - # Gemma-2, etc.), if not provided, will be set to `0`. If greater

474 - # than 0, the logits will be capped according to formula:

475 - # $$\texttt{logits\_soft\_cap} \times

476 - # \mathrm{tanh}(x / \texttt{logits\_soft\_cap})$$,

477 - # where $x$ is the input logits.

478 - logits_soft_cap: Optional[float] = None

479 - # The scale used in softmax, if not provided, will be set to

480 - # `1.0 / sqrt(head_dim)`.

481 - sm_scale: Optional[float] = None

482 -

483 - def __post_init__(self):

484 - # Refer to

485 - # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157

486 - supported_head_sizes = FlashInferBackend.get_supported_head_sizes()

487 - if self.head_dim is not None and self.head_dim \

488 - not in supported_head_sizes:

489 - raise ValueError(

490 - f"Only {supported_head_sizes} are supported for head_dim,",

491 - f" received {self.head_dim}.")

492 -

493 - def begin_forward(self):

494 - if self.num_prefill_tokens > 0:

495 - if self.paged_kv_indices is None:

496 - return

497 -

498 - assert self.prefill_wrapper is not None

499 - assert self.query_start_loc is not None

500 - assert self.paged_kv_indices is not None

501 - assert self.paged_kv_indptr is not None

502 - assert self.paged_kv_last_page_len is not None

503 - assert self.block_table_bound is not None

504 - assert self.seq_lens_tensor is not None

505 - self.query_start_loc = self.query_start_loc[:self.num_prefills + 1]

506 - batch_size = self.query_start_loc.shape[0] - 1

507 - assert batch_size >= 0

508 - # We will use flash attention for profiling to

509 - # determine the number of blocks. Therefore,

510 - # we don't need to prepare the input for flashinfer for profile run.

511 - if not self.is_profile_run:

512 - self.paged_kv_indptr = self.paged_kv_indptr.to(self.device)

513 - self.paged_kv_last_page_len = self.paged_kv_last_page_len.to(

514 - self.device)

515 - self.block_table_bound = self.block_table_bound.to(self.device)

516 - self.seq_lens_tensor = self.seq_lens_tensor.to(self.device)

517 - self.paged_kv_indices = self.paged_kv_indices.to(self.device)

518 - self.prefill_wrapper.plan(

519 - self.query_start_loc,

520 - self.paged_kv_indptr[:self.num_prefills + 1],

521 - self.paged_kv_indices,

522 - self.paged_kv_last_page_len[:self.num_prefills],

523 - self.num_qo_heads,

524 - self.num_kv_heads,

525 - self.head_dim,

526 - self.page_size,

527 - causal=True,

528 - sm_scale=self.sm_scale,

529 - window_left=self.window_left,

530 - logits_soft_cap=self.logits_soft_cap,

531 - q_data_type=self.q_data_type,

532 - kv_data_type=self.data_type)

533 - if self.num_decode_tokens > 0:

534 - assert self.paged_kv_indices is not None

535 - assert self.paged_kv_indptr is not None

536 - assert self.paged_kv_last_page_len is not None

537 - self.paged_kv_indices = self.paged_kv_indices.to(self.device)

538 - self.paged_kv_indptr = self.paged_kv_indptr.to(self.device)

539 - self.paged_kv_last_page_len = self.paged_kv_last_page_len.to(

540 - self.device)

541 - # handle model warmup path

542 - if self.block_table_bound is not None:

543 - self.block_table_bound = self.block_table_bound.to(self.device)

544 - if self.seq_lens_tensor is not None:

545 - self.seq_lens_tensor = self.seq_lens_tensor.to(self.device)

546 -

547 - assert self.decode_wrapper is not None

548 - self.decode_wrapper.plan(

549 - self.paged_kv_indptr[self.num_prefills:],

550 - self.paged_kv_indices,

551 - self.paged_kv_last_page_len[self.num_prefills:],

552 - self.num_qo_heads,

553 - self.num_kv_heads,

554 - self.head_dim,

555 - self.page_size,

556 - # Disable flashinfer's pos encoding and use vllm's rope.

557 - pos_encoding_mode="NONE",

558 - window_left=self.window_left,

559 - logits_soft_cap=self.logits_soft_cap,

560 - sm_scale=self.sm_scale,

561 - # kv-cache data type.

562 - kv_data_type=self.data_type,

563 - # query data type.

564 - q_data_type=self.q_data_type)

565 -

566 - def asdict_zerocopy(self,

567 - skip_fields: Optional[Set[str]] = None

568 - ) -> Dict[str, Any]:

569 - if skip_fields is None:

570 - skip_fields = set()

571 - # We need to skip the prefill/decode_wrapper field since it cannot be

572 - # broadcasted with nccl when TP is enabled.

573 - skip_fields.add('prefill_wrapper')

574 - skip_fields.add('decode_wrapper')

575 - return super().asdict_zerocopy(skip_fields)

576 -

577 - @property

578 - def prefill_metadata(self) -> Optional["FlashInferMetadata"]:

579 - if self.num_prefills == 0:

580 - return None

581 - return self

582 -

583 - @property

584 - def decode_metadata(self) -> Optional["FlashInferMetadata"]:

585 - if self.num_decode_tokens == 0:

586 - return None

587 - return self

588 -

589 -

590 -class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

591 -

592 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):

593 -

594 - self.input_builder = input_builder

595 - self.runner = input_builder.runner

596 -

597 - self.sliding_window = input_builder.sliding_window

598 - self.block_size = input_builder.block_size

599 -

600 - # Global hyperparameters shared by all attention layers

601 - self.global_hyperparameters: Optional[PerLayerParameters] = None

602 -

603 - self.vllm_config = self.runner.vllm_config

604 -

605 - def prepare(self):

606 - self.slot_mapping: List[int] = []

607 - self.prefill_seq_lens: List[int] = []

608 - self.context_lens: List[int] = []

609 - self.block_tables: List[List[int]] = []

610 - self.curr_seq_lens: List[int] = []

611 - self.multimodal_placeholder_maps: Dict[

612 - str,

613 - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)

614 - self.num_prefills = 0

615 - self.num_prefill_tokens = 0

616 - self.num_decode_tokens = 0

617 -

618 - # Please follow https://docs.flashinfer.ai/tutorials/kv_layout.html#page-layout

619 - # for the precise definition of the following fields.

620 - # An example:

621 - # request 1, page indices [0, 5, 8]

622 - # request 2, page indices [1, 6, 7]

623 - # request 3, page indices [3, 4]

624 - # paged_kv_indices is a concatenation of page indices of all requests:

625 - # [0, 5, 8, 1, 6, 7, 3, 4]

626 - # paged_kv_indptr is used to index into paged_kv_indices:

627 - # [0, 3, 6, 8]

628 - self.paged_kv_indices: List[int] = []

629 - # 0 at the beginning of paged_kv_indptr indicates the start of the

630 - # first request’s page indices in the paged_kv_indices list.

631 - self.paged_kv_indptr: List[int] = [0]

632 - # paged_kv_last_page_len is the length of the last page of each request

633 - self.paged_kv_last_page_len: List[int] = []

634 - self.total_blocks = 0

635 - self.is_profile_run: bool = False

636 -

637 - if self.global_hyperparameters is None:

638 - # Infer global hyperparameters, since currently we only support

639 - # models in which all layers share the same values for the

640 - # following hyperparameters:

641 - # - `window_left`

642 - # - `logits_soft_cap`

643 - # - `sm_scale`

644 - inferred_params = infer_global_hyperparameters(

645 - get_per_layer_parameters(self.vllm_config))

646 - self.global_hyperparameters = inferred_params

647 - self.window_left = inferred_params.window_left

648 - self.logits_soft_cap = inferred_params.logits_soft_cap

649 - self.sm_scale = inferred_params.sm_scale

650 -

651 - def _add_seq_group(

652 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

653 - chunked_prefill_enabled: bool):

654 - """Add a sequence group to the metadata. Specifically update/append

655 - 1. context length.

656 - 2. block table.

657 - 3. slot mapping.

658 - """

659 - is_prompt = inter_data.is_prompt

660 - block_tables = inter_data.block_tables

661 - computed_block_nums = inter_data.computed_block_nums

662 -

663 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

664 - curr_sliding_window_block) in zip(

665 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

666 - inter_data.orig_seq_lens, inter_data.seq_lens,

667 - inter_data.query_lens, inter_data.context_lens,

668 - inter_data.curr_sliding_window_blocks):

669 - self.context_lens.append(context_len)

670 - if is_prompt:

671 - mm_maps = inter_data.multi_modal_placeholder_maps

672 - if mm_maps:

673 - for modality, placeholders in mm_maps.items():

674 - self.multimodal_placeholder_maps[modality].extend(

675 - placeholders)

676 - self.num_prefills += 1

677 - self.num_prefill_tokens += token_len

678 - self.prefill_seq_lens.append(seq_len)

679 - else:

680 - assert query_len == 1, (

681 - "seq_len: {}, context_len: {}, query_len: {}".format(

682 - seq_len, context_len, query_len))

683 - self.num_decode_tokens += query_len

684 - self.curr_seq_lens.append(curr_seq_len)

685 -

686 - # Compute block table.

687 - # TODO(sang): Combine chunked prefill and prefix caching by

688 - # only allowing multiple of block_size chunk size.

689 - # NOTE: This only works for oooooooxxx style attention.

690 - block_table = []

691 - if inter_data.prefix_cache_hit:

692 - block_table = computed_block_nums

693 - elif ((chunked_prefill_enabled or not is_prompt)

694 - and block_tables is not None):

695 - block_table = block_tables[seq_id][-curr_sliding_window_block:]

696 - self.block_tables.append(block_table)

697 -

698 - is_profile_run = is_block_tables_empty(block_tables)

699 -

700 - # Compute slot mapping.

701 - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,

702 - context_len,

703 - self.sliding_window)

704 - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,

705 - seq_len, context_len, start_idx,

706 - self.block_size, inter_data.block_tables)

707 -

708 - # It is not necessary to add paged_kv_indices, paged_kv_indptr,

709 - # and paged_kv_last_page_len for profile run because we will

710 - # create dummy inputs.

711 - if is_profile_run:

712 - self.is_profile_run = is_profile_run

713 - return

714 -

715 - block_table = block_tables[seq_id]

716 - self._update_paged_kv_tensors(block_table, seq_len)

717 -

718 - def _update_paged_kv_tensors(self, block_table: List[int], seq_len: int):

719 - # Get the number of valid blocks based on sequence length.

720 - # If seq_len = 16, block_size = 16,

721 - # block_table_bound is 1 with 1 valid block.

722 - # If seq_len = 15, block_size = 16,

723 - # block_table_bound is 0 + 1 with 1 valid block.

724 - self.total_blocks += len(block_table)

725 - block_table_bound = seq_len // self.block_size + 1 \

726 - if seq_len % self.block_size != 0 \

727 - else seq_len // self.block_size

728 - self.paged_kv_indices.extend(block_table[:block_table_bound])

729 - self.paged_kv_indptr.append(self.paged_kv_indptr[-1] +

730 - block_table_bound)

731 -

732 - last_page_len = seq_len % self.block_size

733 - if last_page_len == 0:

734 - last_page_len = self.block_size

735 - self.paged_kv_last_page_len.append(last_page_len)

736 -

737 - def build(self, seq_lens: List[int], query_lens: List[int],

738 - cuda_graph_pad_size: int, batch_size: int):

739 - """Build attention metadata with on-device tensors.

740 -

741 - Args:

742 - seq_lens: The maybe padded sequence lengths of the input sequences.

743 - query_lens: The query lengths of the input sequences.

744 - cuda_graph_pad_size: The padding size for cuda graph.

745 - -1 if cuda graph is not used.

746 - batch_size: The maybe padded batch size.

747 - """

748 - for inter_data in self.input_builder.inter_data_list:

749 - self._add_seq_group(inter_data,

750 - self.input_builder.chunked_prefill_enabled)

751 -

752 - device = self.runner.device

753 - use_captured_graph = cuda_graph_pad_size != -1

754 -

755 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0)

756 - max_decode_seq_len = max(self.curr_seq_lens, default=0)

757 - num_decode_tokens = self.num_decode_tokens

758 - decode_query_len = max(query_lens[self.num_prefills:], default=1)

759 -

760 - if use_captured_graph:

761 - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)

762 - self.block_tables.extend([] * cuda_graph_pad_size)

763 - num_decode_tokens = batch_size - self.num_prefill_tokens

764 -

765 - # The shape of graph_block_tables is

766 - # [max batch size, max context len // block size].

767 - input_block_tables = self.runner.graph_block_tables[:batch_size]

768 - max_blocks = input_block_tables.shape[1]

769 - for i, block_table in enumerate(self.block_tables):

770 - if block_table:

771 - num_blocks = len(block_table)

772 - if num_blocks <= max_blocks:

773 - input_block_tables[i, :num_blocks] = block_table

774 - else:

775 - # It may be possible to have more blocks allocated due

776 - # to lookahead slots of multi-step, however, they are

777 - # not used anyway, so can be safely ignored.

778 - input_block_tables[

779 - i, :max_blocks] = block_table[:max_blocks]

780 -

781 - block_tables = torch.from_numpy(input_block_tables).to(

782 - device, non_blocking=True)

783 -

784 - last_paged_kv_indptr = self.paged_kv_indptr[-1]

785 - self.paged_kv_indptr.extend([last_paged_kv_indptr] *

786 - cuda_graph_pad_size)

787 - self.paged_kv_last_page_len.extend([0] * cuda_graph_pad_size)

788 - else:

789 - block_tables = make_tensor_with_pad(

790 - self.block_tables,

791 - pad=0,

792 - dtype=torch.int,

793 - device=device,

794 - )

795 -

796 - assert device is not None

797 - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,

798 - self.runner.pin_memory)

799 - query_lens_tensor = async_tensor_h2d(query_lens, torch.long, device,

800 - self.runner.pin_memory)

801 - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,

802 - device, self.runner.pin_memory)

803 - query_start_loc = torch.zeros(query_lens_tensor.shape[0] + 1,

804 - dtype=torch.int32,

805 - device=device)

806 - seq_start_loc = torch.zeros(seq_lens_tensor.shape[0] + 1,

807 - dtype=torch.int32,

808 - device=device)

809 - placeholder_index_maps = {

810 - modality: placeholder_map.index_map()

811 - for modality, placeholder_map in

812 - self.multimodal_placeholder_maps.items()

813 - }

814 - torch.cumsum(seq_lens_tensor,

815 - dim=0,

816 - dtype=seq_start_loc.dtype,

817 - out=seq_start_loc[1:])

818 - torch.cumsum(query_lens_tensor,

819 - dim=0,

820 - dtype=query_start_loc.dtype,

821 - out=query_start_loc[1:])

822 -

823 - if len(self.paged_kv_indptr) > 0:

824 - # extend to the maximum number of blocks as returned by the

825 - # scheduler

826 - self.paged_kv_indices.extend(

827 - [0] * (self.total_blocks - len(self.paged_kv_indices)))

828 - paged_kv_indices_tensor = torch.tensor(self.paged_kv_indices,

829 - device="cpu",

830 - dtype=torch.int)

831 - paged_kv_indptr_tensor = torch.tensor(self.paged_kv_indptr,

832 - device="cpu",

833 - dtype=torch.int)

834 - paged_kv_last_page_len_tensor = torch.tensor(

835 - self.paged_kv_last_page_len, device="cpu", dtype=torch.int)

836 - block_table_bound_tensor = torch.zeros(len(self.paged_kv_indptr) -

837 - 1,

838 - device="cpu",

839 - dtype=torch.int)

840 - else:

841 - paged_kv_indices_tensor = None

842 - paged_kv_indptr_tensor = None

843 - paged_kv_last_page_len_tensor = None

844 - block_table_bound_tensor = None

845 -

846 - if self.runner.kv_cache_dtype.startswith("fp8"):

847 - kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(

848 - self.runner.kv_cache_dtype)

849 - else:

850 - kv_cache_dtype = get_kv_cache_torch_dtype(

851 - self.runner.kv_cache_dtype, self.runner.model_config.dtype)

852 -

853 - return FlashInferMetadata(

854 - decode_query_len=decode_query_len,

855 - num_prefills=self.num_prefills,

856 - slot_mapping=slot_mapping_tensor,

857 - multi_modal_placeholder_index_maps=placeholder_index_maps,

858 - enable_kv_scales_calculation=False,

859 - num_prefill_tokens=self.num_prefill_tokens,

860 - num_decode_tokens=num_decode_tokens,

861 - max_prefill_seq_len=max_prefill_seq_len,

862 - max_decode_seq_len=max_decode_seq_len,

863 - block_tables=block_tables,

864 - paged_kv_indptr=paged_kv_indptr_tensor,

865 - paged_kv_indices=paged_kv_indices_tensor,

866 - paged_kv_last_page_len=paged_kv_last_page_len_tensor,

867 - block_table_bound=block_table_bound_tensor,

868 - seq_lens_tensor=seq_lens_tensor,

869 - num_qo_heads=self.runner.model_config.get_num_attention_heads(

870 - self.runner.parallel_config),

871 - num_kv_heads=self.runner.model_config.get_num_kv_heads(

872 - self.runner.parallel_config),

873 - head_dim=self.runner.model_config.get_head_size(),

874 - page_size=self.block_size,

875 - seq_start_loc=seq_start_loc,

876 - query_start_loc=query_start_loc,

877 - device=device,

878 - data_type=kv_cache_dtype,

879 - q_data_type=self.runner.model_config.dtype,

880 - use_cuda_graph=use_captured_graph,

881 - is_profile_run=self.is_profile_run,

882 - window_left=self.window_left,

883 - logits_soft_cap=self.logits_soft_cap,

884 - sm_scale=self.sm_scale,

885 - )

886 -

887 -

888 -class FlashInferImpl(AttentionImpl):

889 -

890 - def __init__(

891 - self,

892 - num_heads: int,

893 - head_size: int,

894 - scale: float,

895 - num_kv_heads: int,

896 - alibi_slopes: Optional[List[float]],

897 - sliding_window: Optional[int],

898 - kv_cache_dtype: str,

899 - logits_soft_cap: Optional[float] = None,

900 - attn_type: str = AttentionType.DECODER,

901 - kv_sharing_target_layer_name: Optional[str] = None,

902 - use_irope: bool = False,

903 - ) -> None:

904 - if kv_sharing_target_layer_name is not None:

905 - raise NotImplementedError("KV sharing is not supported in V0 "

906 - "FLASHINFER backend.")

907 - if use_irope:

908 - logger.warning_once(

909 - "Using irope in FlashInfer is not supported yet, it will fall"

910 - " back to global attention for long context.")

911 - self.num_heads = num_heads

912 - self.head_size = head_size

913 - self.scale = float(scale)

914 - self.num_kv_heads = num_kv_heads

915 - if alibi_slopes is not None:

916 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

917 - self.alibi_slopes = alibi_slopes

918 - self.sliding_window = ((sliding_window - 1,

919 - 0) if sliding_window is not None else (-1, -1))

920 - self.kv_cache_dtype = kv_cache_dtype

921 - self.logits_soft_cap = logits_soft_cap

922 -

923 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

924 -

925 - if attn_type != AttentionType.DECODER:

926 - raise NotImplementedError("Encoder self-attention and "

927 - "encoder/decoder cross-attention "

928 - "are not implemented for "

929 - "FlashInferImpl")

930 -

931 - def forward(

932 - self,

933 - layer: AttentionLayer,

934 - query: torch.Tensor,

935 - key: torch.Tensor,

936 - value: torch.Tensor,

937 - kv_cache: torch.Tensor,

938 - attn_metadata: FlashInferMetadata,

939 - output: Optional[torch.Tensor] = None,

940 - output_scale: Optional[torch.Tensor] = None,

941 - ) -> torch.Tensor:

942 -

943 - if output_scale is not None:

944 - raise NotImplementedError(

945 - "fused output quantization is not yet supported"

946 - " for FlashInferImpl")

947 -

948 - # TODO: directly write to output tensor

949 - num_heads: int = self.num_heads

950 - head_size: int = self.head_size

951 - num_kv_heads: int = self.num_kv_heads

952 - kv_cache_dtype: str = self.kv_cache_dtype

953 - softmax_scale: float = self.scale

954 - window_size = self.sliding_window

955 - alibi_slopes = self.alibi_slopes

956 - logits_soft_cap = self.logits_soft_cap

957 -

958 - num_tokens, hidden_size = query.shape

959 - query = query.view(-1, num_heads, head_size)

960 - key = key.view(-1, num_kv_heads, head_size)

961 - value = value.view(-1, num_kv_heads, head_size)

962 -

963 - if kv_cache.numel() > 0:

964 - # Use the same reshape and cache kernel as flash attention.

965 - ops.reshape_and_cache_flash(

966 - key,

967 - value,

968 - kv_cache[:, 0],

969 - kv_cache[:, 1],

970 - attn_metadata.slot_mapping.flatten(),

971 - kv_cache_dtype,

972 - layer._k_scale,

973 - layer._v_scale,

974 - )

975 - # The FlashInfer api requires data to be in fp8_e4m3 or fp8_e5m2

976 - # to process the cache when the kv_cache_dtype is fp8

977 - if kv_cache_dtype.startswith("fp8"):

978 - torch_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(

979 - kv_cache_dtype)

980 - kv_cache = kv_cache.view(torch_dtype)

981 -

982 - num_prefill_tokens = attn_metadata.num_prefill_tokens

983 - num_decode_tokens = attn_metadata.num_decode_tokens

984 - assert key.shape[0] == num_prefill_tokens + num_decode_tokens, \

985 - f"key : {key.shape} : #prefill tokens {num_prefill_tokens} : #decode tokens {num_decode_tokens}" # noqa

986 - assert value.shape[0] == num_prefill_tokens + num_decode_tokens, \

987 - f"value : {value.shape} : #prefill toks {num_prefill_tokens} : #decode toks {num_decode_tokens}" # noqa

988 - query = query.contiguous(

989 - ) # Flashinfer requires query to be contiguous

990 - # Query for decode. KV is not needed because it is already cached.

991 - # QKV for prefill.

992 - decode_query = query[num_prefill_tokens:]

993 - query = query[:num_prefill_tokens]

994 -

995 - key = key[:num_prefill_tokens]

996 - value = value[:num_prefill_tokens]

997 -

998 - assert query.shape[0] == num_prefill_tokens

999 - assert decode_query.shape[0] == num_decode_tokens

1000 -

1001 - window_left = window_size[0] if window_size is not None else -1

1002 -

1003 - prefill_output: Optional[torch.Tensor] = None

1004 - if num_decode_tokens > 0:

1005 - decode_output = torch.empty(decode_query.shape,

1006 - dtype=decode_query.dtype,

1007 - device=decode_query.device)

1008 - else:

1009 - decode_output = None

1010 - stride_order = FlashInferBackend.get_kv_cache_stride_order()

1011 - if prefill_meta := attn_metadata.prefill_metadata:

1012 - # We will use flash attention for prefill

1013 - # when kv_cache is not provided.

1014 - # This happens when vllm runs the profiling to

1015 - # determine the number of blocks.

1016 - if kv_cache.numel() == 0:

1017 - prefill_output = flash_attn_varlen_func(

1018 - q=query,

1019 - k=key,

1020 - v=value,

1021 - cu_seqlens_q=prefill_meta.seq_start_loc,

1022 - cu_seqlens_k=prefill_meta.seq_start_loc,

1023 - max_seqlen_q=prefill_meta.max_prefill_seq_len,

1024 - max_seqlen_k=prefill_meta.max_prefill_seq_len,

1025 - softmax_scale=softmax_scale,

1026 - causal=True,

1027 - window_size=window_size,

1028 - alibi_slopes=alibi_slopes,

1029 - )

1030 - else:

1031 - assert prefill_meta is not None

1032 - assert prefill_meta.prefill_wrapper is not None

1033 -

1034 - assert prefill_meta.prefill_wrapper._causal

1035 - assert prefill_meta.prefill_wrapper._window_left == window_left

1036 - assert prefill_meta.prefill_wrapper._logits_soft_cap == (

1037 - logits_soft_cap or 0.0)

1038 - assert prefill_meta.prefill_wrapper._sm_scale == softmax_scale

1039 -

1040 - prefill_output = prefill_meta.prefill_wrapper.run(

1041 - query,

1042 - kv_cache.permute(*stride_order),

1043 - k_scale=layer._k_scale_float,

1044 - v_scale=layer._v_scale_float,

1045 - )

1046 - if decode_meta := attn_metadata.decode_metadata:

1047 - assert decode_meta is not None

1048 - assert decode_meta.decode_wrapper is not None

1049 -

1050 - assert decode_meta.decode_wrapper._window_left == window_left

1051 - assert decode_meta.decode_wrapper._logits_soft_cap == (

1052 - logits_soft_cap or 0.0)

1053 - assert decode_meta.decode_wrapper._sm_scale == softmax_scale

1054 - # TODO: @pavanimajety Remove this once the switch happens

1055 - # inside flashinfer.

1056 - if not use_trtllm_attention(

1057 - num_decode_tokens, attn_metadata.max_decode_seq_len,

1058 - kv_cache_dtype, attn_metadata.num_qo_heads,

1059 - attn_metadata.num_kv_heads, attn_metadata.head_dim):

1060 - decode_meta.decode_wrapper.run(

1061 - decode_query,

1062 - kv_cache.permute(*stride_order),

1063 - k_scale=layer._k_scale_float,

1064 - v_scale=layer._v_scale_float,

1065 - out=decode_output,

1066 - )

1067 - else:

1068 - workspace_buffer = (

1069 - decode_meta.decode_wrapper._float_workspace_buffer)

1070 - assert FlashInferState.get_kv_cache_layout() == "HND"

1071 - trtllm_batch_decode_with_kv_cache(

1072 - query=decode_query,

1073 - kv_cache=kv_cache.permute(*stride_order),

1074 - workspace_buffer=workspace_buffer,

1075 - block_tables=attn_metadata.block_tables,

1076 - seq_lens=decode_meta.seq_lens_tensor,

1077 - max_seq_len=attn_metadata.max_decode_seq_len,

1078 - bmm1_scale=layer._k_scale_float * softmax_scale,

1079 - bmm2_scale=layer._v_scale_float,

1080 - out=decode_output,

1081 - )

1082 -

1083 - if prefill_output is None and decode_output is not None:

1084 - # Decode only batch.

1085 - output, num_tokens = decode_output, num_decode_tokens

1086 - elif decode_output is None and prefill_output is not None:

1087 - # Prefill only batch.

1088 - output, num_tokens = prefill_output, num_prefill_tokens

1089 - else:

1090 - # Chunked prefill batch does not work with speculative decoding in

1091 - # FlashInfer backend, so the query length for decode should be 1.

1092 - assert prefill_output is not None

1093 - assert decode_output is not None

1094 - assert decode_meta is not None

1095 - assert decode_meta.decode_query_len == 1

1096 - decode_output = decode_output.squeeze(1)

1097 - output = torch.cat([prefill_output, decode_output], dim=0)

1098 - return output.view(num_tokens, hidden_size)

vllm/platforms/cuda.py

+1 -15

@@ -350,17 +350,7 @@ class CudaPlatformBase(Platform):

350 return FLEX_ATTENTION_V1

351

352 # Backends for V0 engine

353 - if selected_backend == _Backend.FLASHINFER:

354 - logger.info("Using FlashInfer backend.")

355 - if cls.has_device_capability(100):

356 - from vllm.v1.attention.backends.utils import (

357 - set_kv_cache_layout)

358 - logger.info_once(

359 - "Using HND KV cache layout on V1 engine by default for "

360 - "Blackwell (SM 10.0) GPUs.")

361 - set_kv_cache_layout("HND")

362 - return "vllm.attention.backends.flashinfer.FlashInferBackend"

363 - elif selected_backend == _Backend.XFORMERS:

353 + if selected_backend == _Backend.XFORMERS:

354 logger.info("Using XFormers backend.")

355 return "vllm.attention.backends.xformers.XFormersBackend"

356 elif selected_backend == _Backend.DUAL_CHUNK_FLASH_ATTN:

@@ -416,10 +406,6 @@ class CudaPlatformBase(Platform):

406 if (fp8_kv_cache and not flash_attn_supports_fp8()):

407 logger.info(

408 "Cannot use FlashAttention backend for FP8 KV cache.")

419 - logger.warning(

420 - "Please use FlashInfer backend with FP8 KV Cache for "

421 - "better performance by setting environment variable "

422 - "VLLM_ATTENTION_BACKEND=FLASHINFER")

409 target_backend = _Backend.XFORMERS

410 except ImportError:

411 logger.info(