Skip to content

GitLab

Unverified Commit 82f3f30e authored by Pleaplusone's avatar Pleaplusone Committed by GitHub
Browse files

[ROCm][Perf] Enable `sparse_mla`'s cudagraph on ROCm platform (#35719) 


Signed-off-by: default avatarganyi <ygan@amd.com>
parent 9095cbbf
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 4 deletions
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
View file @ 82f3f30e
......@@ -151,7 +151,9 @@ class ROCMAiterMLASparseMetadata(AttentionMetadata):
class ROCMAiterMLASparseMetadataBuilder(
AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]
):
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
_cudagraph_support: ClassVar[AttentionCGSupport] = (
AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
)
def __init__(
self,
......
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
View file @ 82f3f30e
......@@ -327,9 +327,6 @@ def rocm_fp8_paged_mqa_logits(
aiter_paged_mqa_logits_module = None
if rocm_aiter_ops.is_enabled():
aiter_paged_mqa_logits_module = paged_mqa_logits_module()
# FIXME(ganyi): Temporarily disable the aiter path until nightly docker
# update aiter to the fix PR.
aiter_paged_mqa_logits_module = None
if aiter_paged_mqa_logits_module is not None:
deepgemm_fp8_paged_mqa_logits_stage1 = (
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment