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 15, 21:00 UTC
[BugFix] Fix prefix caching V0 MLA (#14255)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Co-authored-by: Ying Zhong <zhongyingmatrix@gmail.com>
Lucas Wilkinson committed
Mar 5, 2025 at 20:07 UTC
4dacaa4a8346b1ecd8bf24ec43e0fe5af8d91048
1 file changed
+24
-20
vllm/attention/backends/mla/common.py
+24
-20
@@ -313,9 +313,10 @@ class MLACommonState(AttentionState, Generic[T]):
313
cache_config = runner.cache_config
314
315
self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
316
+ self.enable_prefix_caching = cache_config.enable_prefix_caching
317
317
- if self.chunked_prefill_enabled:
318
- self.chunked_prefill_workspace_size = min(
318
+ if self.chunked_prefill_enabled or self.enable_prefix_caching:
319
+ self.context_chunk_workspace_size = min(
320
# Max sure there is enough for 8 full length request or at least
321
# 4 pages of cache per request
322
max(
@@ -330,7 +331,7 @@ class MLACommonState(AttentionState, Generic[T]):
331
# 2*(192*128)*(64*1024) = 3gb
332
# (assuming 192 QK head dim, 128 heads, and fp16)
333
128 * 1024)
333
- assert self.chunked_prefill_workspace_size >= \
334
+ assert self.context_chunk_workspace_size >= \
335
scheduler_config.max_num_seqs * cache_config.block_size
336
337
@contextmanager
@@ -430,23 +431,23 @@ class MLACommonState(AttentionState, Generic[T]):
431
"TritonMLAState does not support encoder/decoder yet")
432
433
def begin_forward(self, model_input):
433
- if self.chunked_prefill_enabled:
434
- if not hasattr(self, "chunked_prefill_workspace"):
434
+ if self.chunked_prefill_enabled or self.enable_prefix_caching:
435
+ if not hasattr(self, "context_chunk_workspace"):
436
# not self.runner.device does not return the correct device
437
# for this process, (init_device sets the correct device but
438
# only on the Worker). The only way Ive figured out to get the
439
# correct device is to allocate the workspace on the first call
440
# to begin_forward and use the device of the input tokens
441
assert model_input.input_tokens is not None
441
- self.chunked_prefill_workspace = torch.empty(
442
- (self.chunked_prefill_workspace_size,
442
+ self.context_chunk_workspace = torch.empty(
443
+ (self.context_chunk_workspace_size,
444
self.model_config.get_head_size()),
445
dtype=self.model_config.dtype,
446
device=model_input.input_tokens.device,
447
)
448
448
- model_input.attn_metadata.chunked_prefill_workspace = \
449
- self.chunked_prefill_workspace
449
+ model_input.attn_metadata.context_chunk_workspace = \
450
+ self.context_chunk_workspace
451
452
453
@dataclass
@@ -537,7 +538,7 @@ class MLACommonMetadata(AttentionMetadata):
538
context_chunk_seq_tot: Optional[List[int]] = None
539
context_chunk_max_seq_lens: Optional[List[int]] = None
540
# Set by MLAAttentionState in `begin_forward` so it doesn't get broadcasted
540
- chunked_prefill_workspace: Optional[torch.Tensor] = None
541
+ context_chunk_workspace: Optional[torch.Tensor] = None
542
543
def __post_init__(self):
544
supported_head_sizes = MLACommonBackend.get_supported_head_sizes()
@@ -747,11 +748,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
748
self.block_size = input_builder.block_size
749
self.chunked_prefill_enabled = \
750
self.runner.scheduler_config.chunked_prefill_enabled
751
+ self.enable_prefix_caching = \
752
+ self.runner.cache_config.enable_prefix_caching
753
751
- if self.chunked_prefill_enabled:
754
+ if self.chunked_prefill_enabled or self.enable_prefix_caching:
755
attn_state = self.input_builder.runner.attn_state
753
- self.chunked_prefill_workspace_size = \
754
- attn_state.chunked_prefill_workspace_size
756
+ self.context_chunk_workspace_size = \
757
+ attn_state.context_chunk_workspace_size
758
self.page_size = self.runner.block_size
759
760
def prepare(self):
@@ -920,7 +923,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
923
context_chunk_seq_tot = None
924
context_chunk_max_seq_lens = None
925
923
- if self.chunked_prefill_enabled and self.num_prefills > 0 \
926
+ if (self.chunked_prefill_enabled or self.enable_prefix_caching) \
927
+ and self.num_prefills > 0 \
928
and context_lens_tensor is not None \
929
and context_lens_tensor[:self.num_prefills].max() > 0:
930
@@ -936,7 +940,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
940
# algorithm here and allocate more workspace to prefills with
941
# longer context lengths
942
max_context_chunk = \
939
- self.chunked_prefill_workspace_size // num_prefills_with_context
943
+ self.context_chunk_workspace_size // num_prefills_with_context
944
945
# align max_context_chunk to page_size by rounding down,
946
# currently the `gather_cache` kernel cannot handle
@@ -965,7 +969,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
969
chunk_seq_lens.max(dim=1).values.tolist()
970
context_chunk_seq_tot = chunk_seq_lens.sum(dim=1).tolist()
971
assert max(context_chunk_seq_tot) <= \
968
- self.chunked_prefill_workspace_size
972
+ self.context_chunk_workspace_size
973
974
return self.runner.attn_backend.make_metadata(
975
# Required by ModelRunner
@@ -1288,8 +1292,8 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1292
# Fetch from attn_metadata directly, since it late bound by
1293
# MLAAttentionState, grabbing it directly `attn_metadata` can avoid
1294
# any weirdness around prefill_metadata caching
1291
- assert attn_metadata.chunked_prefill_workspace is not None
1292
- workspace = attn_metadata.chunked_prefill_workspace
1295
+ assert attn_metadata.context_chunk_workspace is not None
1296
+ workspace = attn_metadata.context_chunk_workspace
1297
1298
for i in range(iters):
1299
toks = prefill_metadata.context_chunk_seq_tot[i]
@@ -1502,12 +1506,12 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1506
"output is not yet supported for MLAImplBase")
1507
1508
if attn_metadata.is_profile_run and \
1505
- attn_metadata.chunked_prefill_workspace is not None:
1509
+ attn_metadata.context_chunk_workspace is not None:
1510
# During the profile run try to simulate to worse case output size
1511
# for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`
1512
# since this can be large
1513
_ = torch.empty(
1510
- (attn_metadata.chunked_prefill_workspace.shape[0],
1514
+ (attn_metadata.context_chunk_workspace.shape[0],
1515
self.num_heads, self.qk_nope_head_dim + self.v_head_dim),
1516
device=k_c_normed.device,
1517
dtype=k_c_normed.dtype,