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, 19:45 UTC
[ROCm] Enable chunked prefill/paged attention in MLA on ROCm (#14316)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Mar 12, 2025 at 08:51 UTC
d9f83d62068baa8ac4923cdb45f7945c7d6651a3
2 files changed
+4
-18
vllm/attention/backends/mla/common.py
+2
-16
@@ -1327,21 +1327,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1327
[0, q.shape[-1] - v.shape[-1]],
1328
value=0)
1329
1330
- if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN:
1331
- attn_output, attn_softmax_lse = self.triton_fa_func(
1332
- q,
1333
- k,
1334
- v_padded,
1335
- None,
1336
- prefill_metadata.query_start_loc,
1337
- prefill_metadata.context_chunk_cu_seq_lens[i],
1338
- prefill_metadata.max_query_len,
1339
- prefill_metadata.context_chunk_max_seq_lens[i],
1340
- False, # causal
1341
- self.scale,
1342
- None, # attn_mask is None unless applying ALiBi mask
1343
- )
1344
- elif is_vllm_fa:
1330
+ if is_vllm_fa:
1331
attn_output, attn_softmax_lse = self.flash_attn_varlen_func(
1332
q=q,
1333
k=k,
@@ -1416,7 +1402,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1402
v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
1403
value=0)
1404
1419
- if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN:
1405
+ if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN and not has_context:
1406
output = self.triton_fa_func(
1407
q,
1408
k,
vllm/config.py
+2
-2
@@ -3450,9 +3450,9 @@ class VllmConfig:
3450
self.compilation_config.level = CompilationLevel.NO_COMPILATION
3451
3452
if self.model_config and self.model_config.use_mla and \
3453
- not current_platform.is_cuda():
3453
+ not (current_platform.is_cuda() or current_platform.is_rocm()):
3454
logger.info(
3455
- "MLA is enabled on a non-cuda platform; forcing chunked "
3455
+ "MLA is enabled on a non-GPU platform; forcing chunked "
3456
"prefill and prefix caching to be disabled.")
3457
self.scheduler_config.enable_chunked_prefill = False
3458
self.scheduler_config.chunked_prefill_enabled = False