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, 22:15 UTC
[ROCm] Disable chunked prefill/prefix caching when running MLA on non-cuda platforms (#13844)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Feb 25, 2025 at 22:56 UTC
1d35662e6dc199431bfe4004cc84d66fd9b297b1
2 files changed
+44
-12
vllm/attention/backends/mla/common.py
+30
-12
@@ -232,6 +232,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
232
from vllm.model_executor.layers.rotary_embedding import (
233
DeepseekScalingRotaryEmbedding, RotaryEmbedding)
234
from vllm.multimodal import MultiModalPlaceholderMap
235
+from vllm.platforms import current_platform
236
from vllm.utils import async_tensor_h2d, cdiv, make_tensor_with_pad, round_down
237
238
try:
@@ -1371,18 +1372,35 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1372
v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
1373
value=0)
1374
1374
- output = self.flash_attn_varlen_func(
1375
- q=q,
1376
- k=k,
1377
- v=v_padded,
1378
- cu_seqlens_q=prefill_metadata.query_start_loc,
1379
- cu_seqlens_k=prefill_metadata.query_start_loc,
1380
- max_seqlen_q=prefill_metadata.max_prefill_seq_len,
1381
- max_seqlen_k=prefill_metadata.max_prefill_seq_len,
1382
- softmax_scale=self.scale,
1383
- causal=True,
1384
- return_softmax_lse=has_context,
1385
- )
1375
+ if has_context:
1376
+ if not current_platform.is_cuda():
1377
+ raise NotImplementedError(
1378
+ "Chunked Prefill for MLA is not currently supported on"
1379
+ "non-cuda platforms")
1380
+ output = self.flash_attn_varlen_func(
1381
+ q=q,
1382
+ k=k,
1383
+ v=v_padded,
1384
+ cu_seqlens_q=prefill_metadata.query_start_loc,
1385
+ cu_seqlens_k=prefill_metadata.query_start_loc,
1386
+ max_seqlen_q=prefill_metadata.max_prefill_seq_len,
1387
+ max_seqlen_k=prefill_metadata.max_prefill_seq_len,
1388
+ softmax_scale=self.scale,
1389
+ causal=True,
1390
+ return_softmax_lse=True,
1391
+ )
1392
+ else:
1393
+ output = self.flash_attn_varlen_func(
1394
+ q=q,
1395
+ k=k,
1396
+ v=v_padded,
1397
+ cu_seqlens_q=prefill_metadata.query_start_loc,
1398
+ cu_seqlens_k=prefill_metadata.query_start_loc,
1399
+ max_seqlen_q=prefill_metadata.max_prefill_seq_len,
1400
+ max_seqlen_k=prefill_metadata.max_prefill_seq_len,
1401
+ softmax_scale=self.scale,
1402
+ causal=True,
1403
+ )
1404
1405
if has_context:
1406
suffix_output, suffix_lse = output
vllm/config.py
+14
@@ -3422,6 +3422,20 @@ class VllmConfig:
3422
"Disabling `torch.compile`.")
3423
self.compilation_config.level = CompilationLevel.NO_COMPILATION
3424
3425
+ if self.model_config and self.model_config.use_mla and \
3426
+ not current_platform.is_cuda():
3427
+ logger.info(
3428
+ "MLA is enabled on a non-cuda platform; forcing chunked "
3429
+ "prefill and prefix caching to be disabled.")
3430
+ self.scheduler_config.enable_chunked_prefill = False
3431
+ self.scheduler_config.chunked_prefill_enabled = False
3432
+ self.scheduler_config.max_num_batched_tokens = max(
3433
+ self.scheduler_config.max_model_len,
3434
+ _DEFAULT_MAX_NUM_BATCHED_TOKENS)
3435
+
3436
+ if self.cache_config is not None:
3437
+ self.cache_config.enable_prefix_caching = False
3438
+
3439
current_platform.check_and_update_config(self)
3440
3441
if not self.instance_id: