iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][DCP] Default MLA DCP verify to the round-robin asm decode where it applies (#59965)

Browse Source 
Signed-off-by: hyukjlee <hyukjlee@amd.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Hyukjoon Lee
2026-10-07 03:10:54 +00:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 8a26869bca
commit 2342343746
3 changed files with 152 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_rocm_aiter_mla_dcp_cprr.py
+132 -2
View File
@@ -11,7 +11,11 @@ routes built; it is marked and skipped when unavailable.
from __future__ import annotations
import sys
from types import SimpleNamespace
import pytest
import torch
rocm_aiter_mla = pytest.importorskip(
"vllm.v1.attention.backends.mla.rocm_aiter_mla",
@@ -26,9 +30,15 @@ NATIVE = rocm_aiter_mla._NATIVE_CPRR_HEADS
MIN_QLEN = rocm_aiter_mla._MIN_CPRR_QLEN
def test_dcp_verify_env_defaults_to_segmented(monkeypatch):
def test_dcp_verify_env_defaults_to_auto(monkeypatch):
monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False)
assert rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "segmented"
assert rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "auto"
@pytest.mark.parametrize("route", ["auto", "asm", "segmented"])
def test_dcp_verify_env_accepts_routes(monkeypatch, route):
monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", route)
assert route == rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY
def test_dcp_verify_env_rejects_unknown_route(monkeypatch):
@@ -116,6 +126,126 @@ def test_route_needs_gfx950(monkeypatch):
)
# --------------------------------------------------------------------------
# AiterMLAMetadataBuilder -- the route the env value resolves to
# --------------------------------------------------------------------------
DCP = 8
def _build(monkeypatch, route, gathered_heads, segmented_supported=True):
"""Run the real builder __init__ on CPU at TP8/DCP8 with speculative decoding.
Stubs only what needs a GPU, AITER, or a DCP group, so the route decision
under test is the one the builder actually makes.
"""
def init_common_builder(self, *args, **kwargs):
self.num_heads = gathered_heads // DCP
self.dcp_world_size = DCP
self.reorder_batch_threshold = 5
monkeypatch.setattr(
rocm_aiter_mla.MLACommonMetadataBuilder, "__init__", init_common_builder
)
monkeypatch.setattr(
rocm_aiter_mla,
"_segmented_dcp_verify_supported",
lambda *args: segmented_supported,
)
monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False)
monkeypatch.setattr(
rocm_aiter_mla, "get_dcp_group", lambda: SimpleNamespace(rank_in_group=0)
)
monkeypatch.setattr(
torch.cuda,
"get_device_properties",
lambda device: SimpleNamespace(multi_processor_count=256),
)
monkeypatch.setitem(
sys.modules,
"aiter",
SimpleNamespace(
dtypes=SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16"),
get_mla_metadata_info_v1=lambda *args, **kwargs: tuple(
(1, torch.int32) for _ in range(6)
),
),
)
if route is None:
monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False)
else:
monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", route)
config = SimpleNamespace(
speculative_config=SimpleNamespace(num_speculative_tokens=4),
parallel_config=SimpleNamespace(
decode_context_parallel_size=DCP, cp_kv_cache_interleave_size=1
),
model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16),
scheduler_config=SimpleNamespace(max_num_seqs=2),
cache_config=SimpleNamespace(cache_dtype="fp8_e4m3", num_gpu_blocks=None),
compilation_config=SimpleNamespace(
cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False)
),
)
return rocm_aiter_mla.AiterMLAMetadataBuilder(
kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16),
layer_names=["layer.0"],
vllm_config=config,
device=torch.device("cpu"),
)
@pytest.mark.parametrize("route", [None, "auto", "AUTO"])
def test_auto_selects_cprr_for_k3_gathered_heads(on_gfx950, monkeypatch, route):
"""K3 at TP8/DCP8 gathers 96 heads; auto must take cprr padded to 128."""
builder = _build(monkeypatch, route, gathered_heads=96)
assert builder._asm_dcp_verify
assert builder._asm_dcp_verify_heads == 128
assert builder._num_attention_heads == 128
def test_auto_falls_back_to_segmented_above_largest_native(on_gfx950, monkeypatch):
builder = _build(monkeypatch, "auto", gathered_heads=256)
assert not builder._asm_dcp_verify
assert builder._supports_segmented_dcp_verify
def test_auto_without_any_route_names_the_missing_segmented_build(
on_gfx950, monkeypatch
):
with pytest.raises(ValueError) as exc:
_build(monkeypatch, "auto", gathered_heads=256, segmented_supported=False)
assert "VLLM_ROCM_AITER_MLA_DCP_VERIFY=auto" in str(exc.value)
assert "lacks segmented MLA decode" in str(exc.value)
@pytest.mark.parametrize("route", ["asm", "ASM"])
def test_explicit_asm_is_case_insensitive(on_gfx950, monkeypatch, route):
"""env_with_choices returns the raw string; ASM used to mean segmented."""
assert _build(monkeypatch, route, gathered_heads=96)._asm_dcp_verify
def test_explicit_asm_refuses_unservable_heads(on_gfx950, monkeypatch):
with pytest.raises(ValueError) as exc:
_build(monkeypatch, "ASM", gathered_heads=256)
assert "VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm" in str(exc.value)
assert "Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented" in str(exc.value)
@pytest.mark.parametrize("route", ["segmented", "SEGMENTED"])
def test_explicit_segmented_never_selects_cprr(on_gfx950, monkeypatch, route):
assert not _build(monkeypatch, route, gathered_heads=96)._asm_dcp_verify
def test_auto_off_gfx950_stays_segmented(monkeypatch):
import vllm.platforms.rocm as rocm
monkeypatch.setattr(rocm, "on_gfx950", lambda: False)
assert not _build(monkeypatch, "auto", gathered_heads=96)._asm_dcp_verify
# --------------------------------------------------------------------------
# qlen floor
# --------------------------------------------------------------------------
vllm/envs.py
+6 -5
View File
@@ -146,7 +146,7 @@ if TYPE_CHECKING:
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
VLLM_ROCM_USE_AITER_MLA: bool = True
VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto"
VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["asm", "segmented"] = "segmented"
VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["auto", "asm", "segmented"] = "auto"
VLLM_ROCM_USE_AITER_MHA: bool = True
VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: bool = False
@@ -1320,12 +1320,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1")
),
# Kernel for causal multi-token (spec-decode) verify steps under decode
# context parallelism on gfx950: "asm" uses AITER's round-robin ASM
# decode, "segmented" the Triton segmented MLA path.
# context parallelism: "asm" uses AITER's round-robin ASM decode (gfx950),
# "segmented" the Triton segmented MLA path. "auto" (default) takes "asm"
# wherever it can serve the shape and "segmented" otherwise.
"VLLM_ROCM_AITER_MLA_DCP_VERIFY": env_with_choices(
"VLLM_ROCM_AITER_MLA_DCP_VERIFY",
"segmented",
["asm", "segmented"],
"auto",
["auto", "asm", "segmented"],
case_sensitive=False,
),
# Small-head (<16) AITER MLA decode kernel selection. Small head counts
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+14 -7
View File
@@ -601,13 +601,14 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
parallel_config.decode_context_parallel_size,
parallel_config.cp_kv_cache_interleave_size,
)
verify_route = envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY.lower()
asm_dcp_verify_config = (
_asm_dcp_verify_configured(
parallel_config.decode_context_parallel_size,
parallel_config.cp_kv_cache_interleave_size,
multi_token_decode=vllm_config.speculative_config is not None,
)
and envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "asm"
and verify_route != "segmented"
)
super().__init__(
kv_cache_spec,
@@ -627,14 +628,20 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
self._asm_dcp_verify_heads = _asm_dcp_verify_heads(
self.num_heads * self.dcp_world_size
)
if not self._asm_dcp_verify_heads:
if self._asm_dcp_verify_heads:
self._asm_dcp_verify = True
elif verify_route == "asm" or not supports_segmented_dcp_verify:
raise ValueError(
"VLLM_ROCM_AITER_MLA_DCP_VERIFY=asm, but the round-robin asm "
f"decode has no kernel for {self.num_heads * self.dcp_world_size} "
f"DCP-gathered heads (native counts: {_NATIVE_CPRR_HEADS}). "
"Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented."
f"VLLM_ROCM_AITER_MLA_DCP_VERIFY={verify_route}, but the "
"round-robin asm decode has no kernel for "
f"{self.num_heads * self.dcp_world_size} DCP-gathered heads "
f"(native counts: {_NATIVE_CPRR_HEADS}). "
+ (
"Set VLLM_ROCM_AITER_MLA_DCP_VERIFY=segmented."
if supports_segmented_dcp_verify
else "This AITER build also lacks segmented MLA decode."
)
)
self._asm_dcp_verify = True
self._supports_segmented_dcp_verify = supports_segmented_dcp_verify
self._mla_max_split_per_batch = 0
if self._asm_dcp_verify: