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
Remove chunked_prefill_enabled flag in V1 MLA (#23183)
Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com>
Matthew Bonanni committed
Aug 20, 2025 at 17:43 UTC
a4fbb32fab3d2f91b3672bf581565378aaa18d6c
1 file changed
+23
-27
vllm/v1/attention/backends/mla/common.py
+23
-27
@@ -416,7 +416,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
416
self.model_config = vllm_config.model_config
417
cache_config = vllm_config.cache_config
418
parallel_config = vllm_config.parallel_config
419
- self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
419
self.num_heads = self.model_config.get_num_attention_heads(
420
parallel_config)
421
self.mla_dims = get_mla_dims(self.model_config)
@@ -426,30 +425,28 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
425
if self.aot_schedule:
426
self.page_size = self.kv_cache_spec.block_size
427
429
- if self.chunked_prefill_enabled:
430
- self.chunked_prefill_workspace_size = min(
431
- # Max sure there is enough for 8 full length request or at least
432
- # 4 pages of cache per request
433
- max(
434
- 8 * self.model_config.max_model_len, 4 *
435
- scheduler_config.max_num_seqs * cache_config.block_size),
436
- # For long-context models try not to over-allocate limiting
437
- # kv-cache space, limiting it to 64k tokens,
438
- # which would result in the workspace being:
439
- # 2*(576)*(64*1024) = 144mb
440
- # (assuming 576 MLA head dim, and fp16)
441
- # which would result in up-projected context being
442
- # 2*(192*128)*(64*1024) = 3gb
443
- # (assuming 192 QK head dim, 128 heads, and fp16)
444
- 128 * 1024)
445
- assert self.chunked_prefill_workspace_size >= \
446
- scheduler_config.max_num_seqs * cache_config.block_size
447
- self.chunked_prefill_workspace = torch.empty(
448
- (self.chunked_prefill_workspace_size,
449
- self.model_config.get_head_size()),
450
- dtype=self.model_config.dtype,
451
- device=device,
452
- )
428
+ self.chunked_prefill_workspace_size = min(
429
+ # Max sure there is enough for 8 full length request or at least
430
+ # 4 pages of cache per request
431
+ max(8 * self.model_config.max_model_len,
432
+ 4 * scheduler_config.max_num_seqs * cache_config.block_size),
433
+ # For long-context models try not to over-allocate limiting
434
+ # kv-cache space, limiting it to 64k tokens,
435
+ # which would result in the workspace being:
436
+ # 2*(576)*(64*1024) = 144mb
437
+ # (assuming 576 MLA head dim, and fp16)
438
+ # which would result in up-projected context being
439
+ # 2*(192*128)*(64*1024) = 3gb
440
+ # (assuming 192 QK head dim, 128 heads, and fp16)
441
+ 128 * 1024)
442
+ assert self.chunked_prefill_workspace_size >= \
443
+ scheduler_config.max_num_seqs * cache_config.block_size
444
+ self.chunked_prefill_workspace = torch.empty(
445
+ (self.chunked_prefill_workspace_size,
446
+ self.model_config.get_head_size()),
447
+ dtype=self.model_config.dtype,
448
+ device=device,
449
+ )
450
451
self._use_cudnn_prefill = use_cudnn_prefill()
452
self._use_fi_prefill = use_flashinfer_prefill()
@@ -620,8 +617,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
617
reqs_start:] - query_start_loc[reqs_start]
618
619
chunked_context_metadata = None
623
- if self.chunked_prefill_enabled and num_prefills > 0 \
624
- and max_context_len_cpu > 0:
620
+ if max_context_len_cpu > 0:
621
# NOTE: it is recommend you read the `Chunked Prefill` section
622
# in the comment at the top of the file before trying to
623
# understand the following code