e61bac87ee · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -10,8 +10,7 @@ import torch

10 from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,

11 BatchPrefillWithPagedKVCacheWrapper,

12 MultiLevelCascadeAttentionWrapper)

13 -from flashinfer.decode import (_get_range_buf, get_seq_lens,

14 - trtllm_batch_decode_with_kv_cache)

13 +from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache

14 from flashinfer.prefill import trtllm_batch_context_with_kv_cache

15

16 import vllm.envs as envs

@@ -142,19 +141,10 @@ class FlashInferMetadata:

141 # The number of entries in the last page of each request in

142 # the paged kv cache, shape: [batch_size] (CPU for plan)

143 paged_kv_last_page_len_cpu: torch.Tensor

145 - # The number of query/output heads

146 - num_qo_heads: int

147 - # The number of key/value heads

148 - num_kv_heads: int

149 - # The dimension of the attention heads

150 - head_dim: int

151 - # Block size of vllm

152 - page_size: int

153 - # The data type of the paged kv cache

154 - kv_data_type: torch.dtype

144 # The data type of the query

145 q_data_type: torch.dtype

146

147 + seq_lens_cpu: torch.Tensor

148 slot_mapping: torch.Tensor

149

150 # For flashinfer trtllm batch decode

@@ -185,10 +175,6 @@ class FlashInferMetadata:

175 qo_indptr_gpu: Optional[torch.Tensor] = None

176 paged_kv_indptr_gpu: Optional[torch.Tensor] = None

177

188 - def __post_init__(self):

189 - if self.head_dim is not None:

190 - FlashInferBackend.validate_head_size(self.head_dim)

191 -

178

179 class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

180 cudagraph_support: ClassVar[AttentionCGSupport] = \

@@ -201,13 +187,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

187 self.device = device

188 self.vllm_config = vllm_config

189 self.cache_config = vllm_config.cache_config

190 + self.model_config = vllm_config.model_config

191 self.kv_cache_spec = kv_cache_spec

192 self._workspace_buffer = None

193 self._prefill_wrapper = None # Wrapper for prefill/append

194 self._decode_wrapper = None # Wrapper for decode (general shape)

195

196 self.compilation_config = vllm_config.compilation_config

210 - max_num_pages_per_req = cdiv(vllm_config.model_config.max_model_len,

197 + max_num_pages_per_req = cdiv(self.model_config.max_model_len,

198 self.kv_cache_spec.block_size)

199 max_num_reqs = vllm_config.scheduler_config.max_num_seqs

200 max_num_pages = max_num_reqs * max_num_pages_per_req

@@ -221,6 +208,29 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

208 self._decode_cudagraph_max_bs = min(

209 max_num_reqs, self.compilation_config.max_capture_size)

210

211 + self.num_qo_heads = self.model_config.get_num_attention_heads(

212 + self.vllm_config.parallel_config)

213 + self.num_kv_heads = self.kv_cache_spec.num_kv_heads

214 + self.head_dim = self.kv_cache_spec.head_size

215 + FlashInferBackend.validate_head_size(self.head_dim)

216 + self.page_size = self.kv_cache_spec.block_size

217 +

218 + self.enable_fusion = (

219 + self.compilation_config.pass_config.enable_attn_fusion)

220 + self.q_data_type = self.model_config.dtype

221 + self.cache_dtype = self.cache_config.cache_dtype

222 + if self.cache_dtype.startswith("fp8"):

223 + self.kv_cache_dtype = (

224 + FlashInferBackend.get_fp8_dtype_for_flashinfer(

225 + self.cache_dtype))

226 + # Insert FP8 quant for query if FP8 kv cache and attn fusion enabled

227 + if self.enable_fusion:

228 + self.q_data_type = self.kv_cache_dtype

229 + else:

230 + self.kv_cache_dtype = self.kv_cache_spec.dtype

231 + self.use_tensor_cores = (envs.VLLM_FLASHINFER_FORCE_TENSOR_CORES or

232 + (self.num_qo_heads // self.num_kv_heads > 4))

233 +

234 self._cascade_wrapper = None # Wrapper for cascade attention

235

236 # Global hyperparameters shared by all attention layers

@@ -282,14 +292,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

292 decode_wrapper = self._decode_wrapper

293

294 if decode_wrapper is None:

285 - num_qo_heads = (

286 - self.vllm_config.model_config.get_num_attention_heads(

287 - self.vllm_config.parallel_config))

288 - num_kv_heads = self.vllm_config.model_config.get_num_kv_heads(

289 - self.vllm_config.parallel_config)

290 - use_tensor_cores = envs.VLLM_FLASHINFER_FORCE_TENSOR_CORES or (

291 - num_qo_heads // num_kv_heads > 4)

292 -

295 if use_cudagraph:

296 paged_kv_indptr = self.paged_kv_indptr[:batch_size + 1]

297 paged_kv_indices = self.paged_kv_indices

@@ -306,7 +308,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

308 paged_kv_indptr_buffer=paged_kv_indptr,

309 paged_kv_indices_buffer=paged_kv_indices,

310 paged_kv_last_page_len_buffer=paged_kv_last_page_len,

309 - use_tensor_cores=use_tensor_cores)

311 + use_tensor_cores=self.use_tensor_cores)

312

313 # save the decode wrapper

314 if use_cudagraph:

@@ -342,16 +344,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

344 attn_metadata.shared_kv_last_page_len_cpu,

345 attn_metadata.paged_kv_last_page_len_cpu

346 ],

345 - attn_metadata.num_qo_heads,

346 - attn_metadata.num_kv_heads,

347 - attn_metadata.head_dim,

348 - attn_metadata.page_size,

347 + self.num_qo_heads,

348 + self.num_kv_heads,

349 + self.head_dim,

350 + self.page_size,

351 causal=True,

352 sm_scale=self.global_hyperparameters.sm_scale,

353 window_left=self.global_hyperparameters.window_left,

354 logits_soft_cap=self.global_hyperparameters.logits_soft_cap,

353 - q_data_type=attn_metadata.q_data_type,

354 - kv_data_type=attn_metadata.kv_data_type,

355 + q_data_type=self.q_data_type,

356 + kv_data_type=self.kv_cache_dtype,

357 )

358 else:

359 # Regular attention (common case).

@@ -383,17 +385,17 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

385 attn_metadata.paged_kv_indices,

386 attn_metadata.

387 paged_kv_last_page_len_cpu[prefill_start:],

386 - attn_metadata.num_qo_heads,

387 - attn_metadata.num_kv_heads,

388 - attn_metadata.head_dim,

389 - attn_metadata.page_size,

388 + self.num_qo_heads,

389 + self.num_kv_heads,

390 + self.head_dim,

391 + self.page_size,

392 causal=True,

393 sm_scale=self.global_hyperparameters.sm_scale,

394 window_left=self.global_hyperparameters.window_left,

395 logits_soft_cap=self.global_hyperparameters.

396 logits_soft_cap,

395 - q_data_type=attn_metadata.q_data_type,

396 - kv_data_type=attn_metadata.kv_data_type,

397 + q_data_type=self.q_data_type,

398 + kv_data_type=self.kv_cache_dtype,

399 )

400 else:

401 attn_metadata.qo_indptr_gpu = qo_indptr_cpu.to(self.device)

@@ -435,18 +437,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

437 self.paged_kv_indptr_cpu[:num_input_tokens + 1],

438 attn_metadata.paged_kv_indices,

439 self.paged_kv_last_page_len_cpu[:num_input_tokens],

438 - attn_metadata.num_qo_heads,

439 - attn_metadata.num_kv_heads,

440 - attn_metadata.head_dim,

441 - attn_metadata.page_size,

440 + attn_metadata.seq_lens_cpu[:num_input_tokens],

441 + self.num_qo_heads,

442 + self.num_kv_heads,

443 + self.head_dim,

444 + self.page_size,

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

446 pos_encoding_mode="NONE",

447 sm_scale=self.global_hyperparameters.sm_scale,

448 window_left=self.global_hyperparameters.window_left,

449 logits_soft_cap=self.global_hyperparameters.

450 logits_soft_cap,

448 - q_data_type=attn_metadata.q_data_type,

449 - kv_data_type=attn_metadata.kv_data_type,

451 + q_data_type=self.q_data_type,

452 + kv_data_type=self.kv_cache_dtype,

453 )

454

455 def build(self,

@@ -458,9 +461,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

461 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens =\

462 split_decodes_and_prefills(common_attn_metadata)

463

461 - page_size = self.kv_cache_spec.block_size

464 + page_size = self.page_size

465 max_q_len = common_attn_metadata.max_query_len

463 - max_seq_len = common_attn_metadata.seq_lens_cpu.max()

466 + max_seq_len = common_attn_metadata.seq_lens_cpu.max().item()

467 seq_lens = common_attn_metadata.seq_lens

468 seq_lens_cpu = common_attn_metadata.seq_lens_cpu

469 block_table_tensor = common_attn_metadata.block_table_tensor

@@ -495,7 +498,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

498 shared_kv_page_indices_cpu = None

499 shared_kv_last_page_len_cpu = None

500

498 - max_num_blocks = block_table_bounds_cpu.max()

501 + max_num_blocks = block_table_bounds_cpu.max().item()

502 block_table_bounds = block_table_bounds_cpu.to(self.device,

503 non_blocking=True)

504 mask = (self.block_table_arange[:max_num_blocks].unsqueeze(0)

@@ -520,42 +523,23 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

523 paged_kv_last_page_len_cpu,

524 out=self.paged_kv_last_page_len_cpu[:num_reqs])

525

523 - cache_dtype = self.cache_config.cache_dtype

524 - if cache_dtype.startswith("fp8"):

525 - kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(

526 - cache_dtype)

527 - else:

528 - kv_cache_dtype = self.kv_cache_spec.dtype

529 -

530 - config = self.vllm_config

531 - num_qo_heads = config.model_config.get_num_attention_heads(

532 - config.parallel_config)

533 - num_kv_heads = self.kv_cache_spec.num_kv_heads

534 - head_dim = self.kv_cache_spec.head_size

535 -

526 # Check if any layer uses sinks (requires TRTLLM attention)

527 has_sinks = self.global_hyperparameters.has_sinks

528

539 - # Insert FP8 quant for query if FP8 kv cache and attn fusion enabled

540 - q_dtype = config.model_config.dtype

541 - enable_fusion = config.compilation_config.pass_config.enable_attn_fusion

542 - if cache_dtype.startswith("fp8") and enable_fusion:

543 - q_dtype = kv_cache_dtype

544 -

545 - prefill_use_trtllm = use_trtllm_attention(num_qo_heads,

546 - num_kv_heads,

529 + prefill_use_trtllm = use_trtllm_attention(self.num_qo_heads,

530 + self.num_kv_heads,

531 num_prefill_tokens,

532 max_seq_len,

549 - cache_dtype,

550 - q_dtype,

533 + self.cache_dtype,

534 + self.q_data_type,

535 is_prefill=True,

536 has_sinks=has_sinks)

553 - decode_use_trtllm = use_trtllm_attention(num_qo_heads,

554 - num_kv_heads,

537 + decode_use_trtllm = use_trtllm_attention(self.num_qo_heads,

538 + self.num_kv_heads,

539 num_decode_tokens,

540 max_seq_len,

557 - cache_dtype,

558 - q_dtype,

541 + self.cache_dtype,

542 + self.q_data_type,

543 is_prefill=False,

544 has_sinks=has_sinks)

545

@@ -566,12 +550,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

550 paged_kv_indices=paged_kv_indices,

551 paged_kv_last_page_len_cpu=self.

552 paged_kv_last_page_len_cpu[:num_reqs],

569 - num_qo_heads=num_qo_heads,

570 - num_kv_heads=num_kv_heads,

571 - head_dim=head_dim,

572 - page_size=page_size,

573 - kv_data_type=kv_cache_dtype,

574 - q_data_type=q_dtype,

553 + q_data_type=self.q_data_type,

554 + seq_lens_cpu=seq_lens_cpu,

555 slot_mapping=common_attn_metadata.slot_mapping,

556 max_q_len=max_q_len,

557 max_seq_len=max_seq_len,

@@ -910,6 +890,7 @@ def fast_plan_decode(

890 indptr_cpu: torch.Tensor,

891 indices: torch.Tensor,

892 last_page_len_cpu: torch.Tensor,

893 + seq_lens_cpu: torch.Tensor,

894 num_qo_heads: int,

895 num_kv_heads: int,

896 head_dim: int,

@@ -987,9 +968,6 @@ def fast_plan_decode(

968 kv_data_type = getattr(torch, kv_data_type) if isinstance(

969 kv_data_type, str) else kv_data_type

970

990 - if self.use_tensor_cores:

991 - qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")

992 -

971 if batch_size != self._fixed_batch_size:

972 raise ValueError(

973 "The batch size should be fixed in cudagraph mode, the runtime "

@@ -1006,12 +984,8 @@ def fast_plan_decode(

984 self._paged_kv_last_page_len_buf.copy_(last_page_len_cpu,

985 non_blocking=True)

986

1009 - indptr_host = indptr_cpu

1010 - last_page_len_host = last_page_len_cpu

1011 -

987 if self.use_tensor_cores:

1013 - kv_lens_arr_host = get_seq_lens(indptr_host, last_page_len_host,

1014 - page_size)

988 + qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")

989

990 try:

991 # Make sure we pass exactly 15 arguments for tensor core version

@@ -1020,8 +994,8 @@ def fast_plan_decode(

994 self._int_workspace_buffer,

995 self._pin_memory_int_workspace_buffer,

996 qo_indptr_host,

1023 - indptr_host,

1024 - kv_lens_arr_host,

997 + indptr_cpu,

998 + seq_lens_cpu,

999 batch_size, # total_num_rows

1000 batch_size,

1001 num_qo_heads,

@@ -1041,7 +1015,7 @@ def fast_plan_decode(

1015 self._float_workspace_buffer,

1016 self._int_workspace_buffer,

1017 self._pin_memory_int_workspace_buffer,

1044 - indptr_host,

1018 + indptr_cpu,

1019 batch_size,

1020 num_qo_heads,

1021 num_kv_heads,

[Misc] Minor refactoring for FlashInfer backend (#23147)