diff --git a/README.md b/README.md index 3cd830dfce53..3bd447f5d22a 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ The plugin leverages the following key components: | vLLM Version | PyTorch Version | Engine | Status | |--------------|-----------------|---------|--------------| -| v0.24.0 | 2.9.x | V1 only | ✅ Supported | +| v0.24.0 | 2.11.x | V1 only | ✅ Supported | > **Note**: This plugin uses vLLM's V1 engine architecture (the V0 engine is not supported). Within the V1 engine, vLLM v0.24.0 auto-selects its **Model Runner V2** for certain architectures (e.g. Qwen3, DeepSeek-V2, Llama) and the V1 model runner for others; both are supported on MUSA. Set `VLLM_USE_V2_MODEL_RUNNER=1` or `0` to force one. @@ -79,7 +79,7 @@ install below. Most of the MUSA wheels are not published on public PyPI, so installing `requirements/musa.txt` with pip's default index fails with -`No matching distribution found for torch==2.9.1.post1+musa5.2.0s5000`. +`No matching distribution found for torch==2.11.0.post1+musa5.2.0`. The MUSA wheels must be **resolved** with the Moore Threads index as the sole `--index-url`, which is why the install below starts with a pass that installs diff --git a/README_CN.md b/README_CN.md index c6c36aad4cf2..b49b94d3b1d5 100644 --- a/README_CN.md +++ b/README_CN.md @@ -45,7 +45,7 @@ | vLLM 版本 | PyTorch 版本 | 引擎 | 状态 | |-----------|--------------|---------|--------------| -| v0.24.0 | 2.9.x | 仅 V1 | ✅ 已支持 | +| v0.24.0 | 2.11.x | 仅 V1 | ✅ 已支持 | > **注意**：本插件使用 vLLM 的 V1 引擎架构（不支持 V0 引擎）。在 V1 引擎内部，vLLM v0.24.0 会为部分架构（如 Qwen3、DeepSeek-V2、Llama）自动选用 **Model Runner V2**，其余架构使用 V1 model runner；两者在 MUSA 上均受支持。设置 `VLLM_USE_V2_MODEL_RUNNER=1` 或 `0` 可强制指定其一。 @@ -76,7 +76,7 @@ MUSA SDK 的主机上，请使用下面的源码安装。 大部分 MUSA wheel 并未发布到公共 PyPI，因此用 pip 默认索引安装 `requirements/musa.txt` 会失败并报 -`No matching distribution found for torch==2.9.1.post1+musa5.2.0s5000`。 +`No matching distribution found for torch==2.11.0.post1+musa5.2.0`。 MUSA wheel 必须在摩尔线程索引作为**唯一** `--index-url` 的前提下解析，因此下面的 安装步骤以“只装 MUSA wheel”的一步开始。 diff --git a/docker/README.md b/docker/README.md index 78f6471a043e..5082c5263d35 100644 --- a/docker/README.md +++ b/docker/README.md @@ -47,7 +47,7 @@ bash docker/build_image.sh With the defaults this produces: ``` -vllm-musa:ubuntu22.04_py3.10_musa_runtime_5.2_pytorch_release_2.9.1.post1_musa5.2.0s5000 +vllm-musa:ubuntu22.04_py3.10_musa_runtime_5.2_pytorch_release_2.11.0.post1_musa5.2.0 ``` The image accepts the model and engine arguments directly: @@ -180,7 +180,7 @@ Two build-time details make this work on such a host: ```bash docker run --rm \ --entrypoint python \ - vllm-musa:ubuntu22.04_py3.10_musa_runtime_5.2_pytorch_release_2.9.1.post1_musa5.2.0s5000 \ + vllm-musa:ubuntu22.04_py3.10_musa_runtime_5.2_pytorch_release_2.11.0.post1_musa5.2.0 \ -c "import torch, torch_musa; print('musa available:', torch.musa.is_available())" ``` @@ -218,12 +218,12 @@ for the validated runtime flags. 9. **vllm-openai** — the default serving target, with `vllm serve` as its entrypoint. -The Triton `3.2.0` pin is intentional for the torch `2.9.1` MUSA stack. Torch -Inductor reads `KernelMetadata.cluster_dims`; the MUSA backend in Triton `3.6.0` -does not expose that field, while the `3.2.0` `mtgpu` backend does. The explicit -vLLM runtime dependency pass also restores the `fastapi[standard]` extras and -`pycountry` that are skipped when the vendored requirements are installed with -`--no-deps` to protect the MUSA torch pins. +The Triton `3.2.0` pin is intentional for the MUSA 5.2 stack and is maintained +independently of the PyTorch release. Validate image imports, Inductor, and +compiled serving before changing it. The explicit vLLM runtime dependency pass +also restores the `fastapi[standard]` extras and `pycountry` that are skipped +when the vendored requirements are installed with `--no-deps` to protect the +MUSA torch pins. For the reasoning behind the pip-index split and the runtime shims, see the comments in `docker/musa.Dockerfile`. diff --git a/docker/musa.Dockerfile b/docker/musa.Dockerfile index 2c411cbbd570..156c7eccf0b3 100644 --- a/docker/musa.Dockerfile +++ b/docker/musa.Dockerfile @@ -68,7 +68,7 @@ RUN apt-get update && \ python -m pip install --upgrade pip && \ rm -rf /var/lib/apt/lists/* -# The torch 2.9.x MUSA wheel links MKL with .so.2 sonames. Ubuntu 22.04 apt +# The MUSA PyTorch wheel links MKL with .so.2 sonames. Ubuntu 22.04 apt # ships the same logical MKL components without that suffix. RUN ln -sf /usr/lib/x86_64-linux-gnu/libmkl_intel_lp64.so \ /usr/lib/x86_64-linux-gnu/libmkl_intel_lp64.so.2 && \ @@ -302,20 +302,31 @@ RUN printf '%s\n' \ '' \ 'def requirement_prefix(dist_name):' \ ' pattern = re.compile(rf"^{re.escape(dist_name)}==(.+)$")' \ - ' for line in Path("requirements/musa_private.txt").read_text().splitlines():' \ - ' match = pattern.match(line.strip())' \ - ' if match:' \ - ' return match.group(1).split("*", 1)[0]' \ - ' raise RuntimeError(f"missing {dist_name} pin in requirements/musa_private.txt")' \ + ' requirement_files = (' \ + ' "requirements/musa_private.txt",' \ + ' "requirements/common.txt",' \ + ' )' \ + ' for requirement_file in requirement_files:' \ + ' for line in Path(requirement_file).read_text().splitlines():' \ + ' match = pattern.match(line.strip())' \ + ' if match:' \ + ' return match.group(1).split("*", 1)[0]' \ + ' raise RuntimeError(f"missing {dist_name} pin in {requirement_files}")' \ + '' \ + 'exact_version_dists = frozenset({"torchada", "torch", "torch_musa", "torchvision", "torchaudio", "deep_ep"})' \ '' \ 'expected = (' \ + ' ("torchada", "torchada", requirement_prefix("torchada")),' \ ' ("numpy", "numpy", "1.26."),' \ ' ("torch", "torch", requirement_prefix("torch")),' \ ' ("torch_musa", "torch_musa", requirement_prefix("torch_musa")),' \ + ' ("torchvision", "torchvision", requirement_prefix("torchvision")),' \ + ' ("torchaudio", "torchaudio", requirement_prefix("torchaudio")),' \ ' ("mate", "mate", ""),' \ ' ("flash_attn_3", "flash_attn_3", ""),' \ ' ("flash_mla", "flash_mla", ""),' \ ' ("deep-gemm", "deep_gemm", ""),' \ + ' ("deep_ep", "deep_ep", requirement_prefix("deep_ep")),' \ ' ("tilelang_musa", "tilelang", ""),' \ ' ("triton", "triton", requirement_prefix("triton")),' \ ' ("uvloop", "uvloop", ""),' \ @@ -328,11 +339,13 @@ RUN printf '%s\n' \ 'for dist_name, module_name, prefix in expected:' \ ' importlib.import_module(module_name)' \ ' installed = version(dist_name)' \ - ' if prefix and not installed.startswith(prefix):' \ + ' if dist_name in exact_version_dists and installed != prefix:' \ + ' raise RuntimeError(f"{dist_name} expected exactly {prefix}, got {installed}")' \ + ' if dist_name not in exact_version_dists and prefix and not installed.startswith(prefix):' \ ' raise RuntimeError(f"{dist_name} expected {prefix}, got {installed}")' \ ' print(f"PASS import {module_name} version={installed}")' \ '' \ - 'for module_name in ("torchada", "vllm", "vllm_musa"):' \ + 'for module_name in ("vllm", "vllm_musa"):' \ ' module = importlib.import_module(module_name)' \ ' print("PASS import %s version=%s" % (module_name, getattr(module, "__version__", "unknown")))' \ > /tmp/vllm_musa_import_check.py && \ diff --git a/requirements/common.txt b/requirements/common.txt index 502b20836007..535a332005fb 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -1,6 +1,7 @@ # Runtime Python dependencies that are portable across MUSA images. -# 0.1.76 ports nested torch/cuda.h includes to torch/musa.h during builds. -torchada>=0.1.76 +# 0.1.79 contains the torch 2.11-compatible streams path used by compiled +# vLLM-MUSA. +torchada==0.1.79 mthreads-ml-py>=2.4.1 numpy==1.26 openai>=2.24.0 diff --git a/requirements/musa_private.txt b/requirements/musa_private.txt index f1d1ea56c106..53afab4aec1a 100644 --- a/requirements/musa_private.txt +++ b/requirements/musa_private.txt @@ -12,19 +12,19 @@ # projects on public PyPI, so they must resolve from the Moore Threads index # only. Versions below track the MUSA release_5.2.0 Python Wheel guide. -torch==2.9.1.post1+musa5.2.0 -torch_musa==2.9.1.post1+musa5.2.0 -torchvision==0.24.1.post1+musa5.2.0 -torchaudio==2.9.1+musa5.2.0 +torch==2.11.0.post1+musa5.2.0 +torch_musa==2.11.0.post1+musa5.2.0 +torchvision==0.26.0.post1+musa5.2.0 +torchaudio==2.11.0+musa5.2.0 mate==0.2.4 flash_attn_3==0.2.4 flash_mla==0.2.4 deep-gemm==0.2.4 sageattention==0.2.4 -deep_ep==1.1.0+musa5.2.0torch2.9.1.post1s5000 +deep_ep==1.1.0+musa5.2.0torch2.11.0.post1 tilelang_musa==0.1.8+musa.3 -# torch 2.9.1 Inductor reads KernelMetadata.cluster_dims. The MUSA backend in -# Triton 3.6.0 does not expose it; the 3.2.0 mtgpu backend does. +# The MUSA 5.2.0 Triton backend remains pinned independently of the PyTorch +# release; verify this pin in the 2.11 image before changing it. triton==3.2.0 apache-tvm-ffi==0.1.9.post3+musa.1 torch-c-dlpack-ext==0.1.5 diff --git a/requirements/test.txt b/requirements/test.txt index 4e897b5671f3..e262858c8715 100644 --- a/requirements/test.txt +++ b/requirements/test.txt @@ -1,2 +1,3 @@ # Test runner included in the Docker image. pytest +pytest-asyncio diff --git a/tests/test_build_ccache.py b/tests/test_build_ccache.py index 15d06f67ebb8..97b9ac8dc020 100644 --- a/tests/test_build_ccache.py +++ b/tests/test_build_ccache.py @@ -6,8 +6,41 @@ import subprocess from pathlib import Path +import pytest + from build_utils.ccache import configure_compiler_cache +_CONFIGURE_ENV_KEYS = ( + "PATH", + "CXX", + "PYTORCH_MCC", + "CCACHE_DIR", + "CCACHE_BASEDIR", + "CCACHE_COMPILERCHECK", + "CCACHE_NOHASHDIR", + "CCACHE_SLOPPINESS", + "CCACHE_MAXSIZE", + "VLLM_MUSA_REAL_CCACHE", + "VLLM_MUSA_REAL_MCC", + "VLLM_MUSA_CCACHE_MUSA_COMPILER", + "VLLM_MUSA_CCACHE_SOURCE_DIR", +) + + +@pytest.fixture(autouse=True) +def _restore_compiler_cache_environment(): + before = { + name: os.environ[name] + for name in _CONFIGURE_ENV_KEYS + if name in os.environ + } + yield + for name in _CONFIGURE_ENV_KEYS: + if name in before: + os.environ[name] = before[name] + else: + os.environ.pop(name, None) + def _write_executable(path: Path, content: str) -> None: path.write_text(content, encoding="utf-8") diff --git a/tests/test_deepseek_v4_swiglu_fp8_fusion_source.py b/tests/test_deepseek_v4_swiglu_fp8_fusion_source.py index 9593ab7950fb..5641bb4c882d 100644 --- a/tests/test_deepseek_v4_swiglu_fp8_fusion_source.py +++ b/tests/test_deepseek_v4_swiglu_fp8_fusion_source.py @@ -3,11 +3,14 @@ from pathlib import Path +import pytest ROOT = Path(__file__).resolve().parents[1] def test_native_kernel_keeps_cuda_spelling_for_torchada() -> None: + if not (ROOT / ".git").exists(): + pytest.skip("pre-build source contract requires a Git worktree") kernel = ( ROOT / "csrc" diff --git a/tests/test_inplace_porting_contract.py b/tests/test_inplace_porting_contract.py index 077d9fd471e7..4a2776e00880 100644 --- a/tests/test_inplace_porting_contract.py +++ b/tests/test_inplace_porting_contract.py @@ -6,10 +6,82 @@ ROOT = Path(__file__).resolve().parent.parent +SUPPORTED_MUSA_STACKS = { + "torch==2.9.1.post1+musa5.2.0": { + "private": ( + "torch_musa==2.9.1.post1+musa5.2.0", + "torchvision==0.24.1.post1+musa5.2.0", + "torchaudio==2.9.1+musa5.2.0", + "deep_ep==1.1.0+musa5.2.0torch2.9.1.post1s5000", + ), + "torchada": "torchada==0.1.77", + }, + "torch==2.11.0.post1+musa5.2.0": { + "private": ( + "torch_musa==2.11.0.post1+musa5.2.0", + "torchvision==0.26.0.post1+musa5.2.0", + "torchaudio==2.11.0+musa5.2.0", + "deep_ep==1.1.0+musa5.2.0torch2.11.0.post1", + ), + "torchada": "torchada==0.1.79", + }, +} + + +def _select_musa_stack(private_requirements): + torch_pins = sorted( + line for line in private_requirements if line.startswith("torch==") + ) + assert len(torch_pins) == 1, f"expected one torch pin, got {torch_pins}" + torch_pin = torch_pins[0] + assert torch_pin in SUPPORTED_MUSA_STACKS, f"unsupported MUSA stack: {torch_pin}" + return SUPPORTED_MUSA_STACKS[torch_pin] + + +def _declared_musa_stack(): + private_requirements = set( + (ROOT / "requirements" / "musa_private.txt").read_text().splitlines() + ) + common_requirements = set( + (ROOT / "requirements" / "common.txt").read_text().splitlines() + ) + return ( + private_requirements, + common_requirements, + _select_musa_stack(private_requirements), + ) + + +def test_supported_musa_stack_contract_cases_are_explicit(): + assert set(SUPPORTED_MUSA_STACKS) == { + "torch==2.9.1.post1+musa5.2.0", + "torch==2.11.0.post1+musa5.2.0", + } + assert ( + SUPPORTED_MUSA_STACKS["torch==2.9.1.post1+musa5.2.0"]["torchada"] + == "torchada==0.1.77" + ) + assert ( + SUPPORTED_MUSA_STACKS["torch==2.11.0.post1+musa5.2.0"]["torchada"] + == "torchada==0.1.79" + ) + assert ( + "torchvision==0.24.1.post1+musa5.2.0" + in SUPPORTED_MUSA_STACKS["torch==2.9.1.post1+musa5.2.0"]["private"] + ) + assert ( + "torchvision==0.26.0.post1+musa5.2.0" + in SUPPORTED_MUSA_STACKS["torch==2.11.0.post1+musa5.2.0"]["private"] + ) + for torch_pin, expected in SUPPORTED_MUSA_STACKS.items(): + assert _select_musa_stack({torch_pin}) is expected + def test_torchada_floor_is_consistent(): assert 'dynamic = ["dependencies"]' in (ROOT / "pyproject.toml").read_text() - assert "torchada>=0.1.76" in (ROOT / "requirements" / "common.txt").read_text() + private_requirements, common_requirements, expected = _declared_musa_stack() + assert set(expected["private"]).issubset(private_requirements) + assert expected["torchada"] in common_requirements def test_musa_image_runtime_dependency_contract(): @@ -22,8 +94,41 @@ def test_musa_image_runtime_dependency_contract(): dockerfile = (ROOT / "docker" / "musa.Dockerfile").read_text() assert "triton==3.2.0" in private_requirements + _, _, expected = _declared_musa_stack() + assert set(expected["private"]).issubset(private_requirements) assert "fastapi[standard]" in runtime_requirements assert "pycountry" in runtime_requirements + assert '"requirements/common.txt"' in dockerfile + exact_import_gates = ( + ("torchada", "torchada"), + ("torch", "torch"), + ("torch_musa", "torch_musa"), + ("torchvision", "torchvision"), + ("torchaudio", "torchaudio"), + ("deep_ep", "deep_ep"), + ) + for dist_name, module_name in exact_import_gates: + gate = f'("{dist_name}", "{module_name}", requirement_prefix("{dist_name}"))' + assert gate in dockerfile + + exact_version_gate = ( + 'exact_version_dists = frozenset({"torchada", "torch", "torch_musa", ' + '"torchvision", "torchaudio", "deep_ep"})' + ) + assert exact_version_gate in dockerfile + torchada_gate = '("torchada", "torchada", requirement_prefix("torchada"))' + torch_gate = '("torch", "torch", requirement_prefix("torch"))' + assert torchada_gate in dockerfile + assert dockerfile.index(torchada_gate) < dockerfile.index(torch_gate) + import_statement = "importlib.import_module(module_name)" + version_statement = "installed = version(dist_name)" + assert import_statement in dockerfile + assert dockerfile.index(import_statement) < dockerfile.index(version_statement) + assert "if dist_name in exact_version_dists and installed != prefix:" in dockerfile + assert ( + "if dist_name not in exact_version_dists and prefix " + "and not installed.startswith(prefix):" + ) in dockerfile assert '("triton", "triton", requirement_prefix("triton"))' in dockerfile assert '("uvloop", "uvloop", "")' in dockerfile assert '("pycountry", "pycountry", "")' in dockerfile @@ -275,7 +380,10 @@ def test_no_legacy_mirror_contract_remains(): "quantization_musa", "per-file _musa", ) - paths = [ROOT / ".gitignore", ROOT / "third_party" / "PINS"] + paths = [ROOT / "third_party" / "PINS"] + gitignore = ROOT / ".gitignore" + if gitignore.exists(): + paths.append(gitignore) paths.extend(sorted((ROOT / "vllm_musa" / "patches" / "series").glob("*.patch"))) offenders = [] diff --git a/tests/test_musa_sync.py b/tests/test_musa_sync.py index b3bddf612107..94607bfb1ce8 100644 --- a/tests/test_musa_sync.py +++ b/tests/test_musa_sync.py @@ -11,6 +11,7 @@ import shutil import subprocess import sys +from collections import Counter from pathlib import Path import pytest @@ -58,9 +59,9 @@ def test_report(ms, capsys): rc = ms.main(["report"]) out = capsys.readouterr().out assert rc == 0 - assert "125 divergences" in out - assert "'1': 67" in out and "'2': 19" in out and "'3': 1" in out - assert "'4a': 2" in out and "'5': 28" in out and "'6': 8" in out + assert "144 divergences" in out + assert "'1': 86" in out and "'2': 20" in out and "'3': 1" in out + assert "'4a': 2" in out and "'5': 27" in out and "'6': 8" in out def test_report_doc(ms, capsys): @@ -102,12 +103,26 @@ def test_series_readme_count_matches_directory(): assert f"Currently **{patch_count} patches**" in readme -def test_series_is_contiguous_and_uses_canonical_headers_and_authors(): +def test_series_uses_documented_prefixes_and_canonical_metadata(): series_dir = ROOT / "vllm_musa" / "patches" / "series" patches = sorted(series_dir.glob("*.patch")) - assert [p.name.split("-", 1)[0] for p in patches] == [ - f"{i:04d}" for i in range(1, len(patches) + 1) - ] + prefixes = [p.name.split("-", 1)[0] for p in patches] + assert all(len(prefix) == 4 and prefix.isdigit() for prefix in prefixes) + assert [int(prefix) for prefix in prefixes] == sorted( + int(prefix) for prefix in prefixes + ) + duplicate_prefixes = { + prefix for prefix, count in Counter(prefixes).items() if count > 1 + } + assert duplicate_prefixes == { + "0088", + "0089", + "0090", + "0094", + "0099", + "0100", + "0101", + } headers = [p.read_bytes().splitlines()[:2] for p in patches] zero_commit_header = ( b"From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001" diff --git a/tests/test_patches.py b/tests/test_patches.py index 0c41502b7e47..7e442f7feb52 100644 --- a/tests/test_patches.py +++ b/tests/test_patches.py @@ -789,6 +789,7 @@ def _make_vllm_config( else None, compilation_config=SimpleNamespace( custom_ops=[], + inductor_compile_config={}, mode=CompilationMode.NONE if is_deepseek_v4 else "NONE", cudagraph_mode=cudagraph_mode, max_cudagraph_capture_size=max_cudagraph_capture_size, @@ -811,6 +812,57 @@ def test_qwen3_moe_does_not_cap_default_cudagraph_capture_size(self): assert vllm_config.compilation_config.max_cudagraph_capture_size is None assert vllm_config.compilation_config.custom_ops == ["all"] + @pytest.mark.parametrize( + ("architectures", "torch211_or_newer", "expected"), + [ + (["Qwen3VLForConditionalGeneration"], True, 1), + (["Qwen3VLForConditionalGeneration"], False, None), + (["Qwen3ForCausalLM"], True, None), + ], + ) + def test_qwen3vl_inductor_tiling_default( + self, monkeypatch, architectures, torch211_or_newer, expected + ): + import vllm_musa.platform as musa_platform + from vllm_musa.platform import MUSAPlatformBase + + vllm_config = self._make_vllm_config(architectures=architectures) + monkeypatch.setattr( + musa_platform, + "_is_torch_211_or_newer", + lambda: torch211_or_newer, + ) + + MUSAPlatformBase.apply_config_platform_defaults(vllm_config) + + assert ( + vllm_config.compilation_config.inductor_compile_config.get( + "triton.max_tiles" + ) + == expected + ) + + def test_qwen3vl_preserves_user_tiling_override(self, monkeypatch): + import vllm_musa.platform as musa_platform + from vllm_musa.platform import MUSAPlatformBase + + vllm_config = self._make_vllm_config( + architectures=["Qwen3VLForConditionalGeneration"] + ) + vllm_config.compilation_config.inductor_compile_config["triton.max_tiles"] = 2 + monkeypatch.setattr( + musa_platform, + "_is_torch_211_or_newer", + lambda: True, + ) + + MUSAPlatformBase.apply_config_platform_defaults(vllm_config) + + assert ( + vllm_config.compilation_config.inductor_compile_config["triton.max_tiles"] + == 2 + ) + def test_qwen3_moe_fp8_preserves_user_cudagraph_capture_size(self): from vllm_musa.platform import MUSAPlatformBase diff --git a/tests/test_qwen3_qk_rope_kv_presplit.py b/tests/test_qwen3_qk_rope_kv_presplit.py index 176d2f3e2f62..9645e9efb787 100644 --- a/tests/test_qwen3_qk_rope_kv_presplit.py +++ b/tests/test_qwen3_qk_rope_kv_presplit.py @@ -89,7 +89,12 @@ def _meta(node: fx.Node, shape: tuple[int, ...], dtype=torch.bfloat16) -> fx.Nod return node -def _make_graph(*, q_heads: int = 16, extra_q_user: bool = False): +def _make_graph( + *, + q_heads: int = 16, + extra_q_user: bool = False, + hoisted_layer_name: bool = False, +): kv_heads = 8 q_width = q_heads * 128 kv_width = kv_heads * 128 @@ -100,6 +105,9 @@ def _make_graph(*, q_heads: int = 16, extra_q_user: bool = False): positions = _meta(graph.placeholder("positions"), (1,), torch.int64) cos = _meta(graph.placeholder("cos"), (32, 128)) output = _meta(graph.placeholder("output"), (1, q_heads, 128)) + layer_name: str | fx.Node = "model.layers.0.self_attn" + if hoisted_layer_name: + layer_name = graph.placeholder("layer_name") split = graph.call_method( "split", args=(qkv, [q_width, kv_width, kv_width]), @@ -152,7 +160,7 @@ def _make_graph(*, q_heads: int = 16, extra_q_user: bool = False): graph.call_function(torch.ops.aten.clone.default, args=(q_flat,)) kv_update = graph.call_function( torch.ops.vllm.unified_kv_cache_update, - kwargs={"key": key, "value": value, "layer_name": "model.layers.0.self_attn"}, + kwargs={"key": key, "value": value, "layer_name": layer_name}, ) attention = graph.call_function( torch.ops.vllm.unified_attention_with_output, @@ -161,7 +169,7 @@ def _make_graph(*, q_heads: int = 16, extra_q_user: bool = False): "key": key, "value": value, "output": output, - "layer_name": "model.layers.0.self_attn", + "layer_name": layer_name, "scale": 0.08838834764831845, "output_scale": 1.0, "kv_cache_dummy_dep": kv_update, @@ -179,6 +187,14 @@ def test_plan_matches_qwen3_06b_and_8b_shapes(): assert len(candidates) == 1 +def test_plan_matches_torch211_hoisted_layer_name(): + graph_module = _make_graph(hoisted_layer_name=True) + candidates = plan_qwen3_qk_rope_kv_presplit(graph_module, 1) + assert candidates is not None + assert len(candidates) == 1 + assert isinstance(candidates[0].layer_name, fx.Node) + + def test_plan_is_atomic_on_extra_q_user(): graph_module = _make_graph(extra_q_user=True) before = str(graph_module.graph) @@ -199,7 +215,8 @@ def test_temporary_schema_is_order_isolated(): assert _dispatch_has_schema(name) is registered_before -def test_apply_replaces_norm_rope_and_cache(monkeypatch): +@pytest.mark.parametrize("hoisted_layer_name", [False, True]) +def test_apply_replaces_norm_rope_and_cache(monkeypatch, hoisted_layer_name): module_name = "vllm_musa.kernels.qwen3_qk_rope_kv" if module_name not in sys.modules: kernels_package = sys.modules.get("vllm_musa.kernels") @@ -219,7 +236,12 @@ def test_apply_replaces_norm_rope_and_cache(monkeypatch): ) with _temporary_schema(_FUSED_OP_NAME, _FUSED_OP_SCHEMA): - graph_module = _make_graph() + graph_module = _make_graph(hoisted_layer_name=hoisted_layer_name) + layer_name = ( + next(node for node in graph_module.graph.nodes if node.name == "layer_name") + if hoisted_layer_name + else None + ) candidates = plan_qwen3_qk_rope_kv_presplit(graph_module, 1) assert candidates is not None assert apply_qwen3_qk_rope_kv_presplit(graph_module, candidates) == 1 @@ -247,3 +269,5 @@ def test_apply_replaces_norm_rope_and_cache(monkeypatch): assert len(attention_nodes) == 1 assert len(fused_nodes) == 1 assert fused_nodes[0] in attention_nodes[0].all_input_nodes + if layer_name is not None: + assert fused_nodes[0].kwargs["layer_name"] is layer_name diff --git a/tests/test_qwen_optimization_contract.py b/tests/test_qwen_optimization_contract.py index c128c09f5181..c2b46a49b5c4 100644 --- a/tests/test_qwen_optimization_contract.py +++ b/tests/test_qwen_optimization_contract.py @@ -405,6 +405,7 @@ def test_static_contract_rejects_one_field_mismatches() -> None: config.speculative_config = object() contract = resolve_optimization_contract(config) assert not contract.prefers(OptimizationFeature.QWEN3_QK_ROPE_KV_PRESPLIT) + assert not contract.prefers(OptimizationFeature.QWEN_LEGACY_SAMPLING) config.speculative_config = None config.cache_config.block_size = 16 diff --git a/tests/test_rms_deepgemm_fusion.py b/tests/test_rms_deepgemm_fusion.py index 66a41cac8386..9dfc4fe57ed5 100644 --- a/tests/test_rms_deepgemm_fusion.py +++ b/tests/test_rms_deepgemm_fusion.py @@ -12,12 +12,23 @@ pytest.importorskip("torchada") torch = pytest.importorskip("torch") pytest.importorskip("torch_musa") +auto_functionalized = pytest.importorskip( + "torch._higher_order_ops.auto_functionalize" +).auto_functionalized HIDDEN = 4096 GROUP_SIZE = 128 EPSILON = 1e-6 +def _logical_graph_targets(graph) -> set[object]: + targets = {node.target for node in graph.nodes} + for node in graph.nodes: + if node.target is auto_functionalized and node.args: + targets.add(node.args[0]) + return targets + + @pytest.fixture(scope="module", autouse=True) def _musa_device() -> Iterator[None]: if not hasattr(torch, "musa") or not torch.musa.is_available(): @@ -273,10 +284,10 @@ def __init__(self, fusion=None): self.inductor_config["post_grad_custom_post_pass"] = self.post_pass def post_pass(self, graph): - self.targets_before = {node.target for node in graph.nodes} + self.targets_before = _logical_graph_targets(graph) if self.fusion is not None: self.fusion(graph) - self.targets_after = {node.target for node in graph.nodes} + self.targets_after = _logical_graph_targets(graph) def __call__(self, graph_module, example_inputs): return compile_fx( diff --git a/tests/test_silu_deepgemm_fusion.py b/tests/test_silu_deepgemm_fusion.py index bc3d554fd5b7..44c885605fab 100644 --- a/tests/test_silu_deepgemm_fusion.py +++ b/tests/test_silu_deepgemm_fusion.py @@ -13,6 +13,17 @@ torch = pytest.importorskip("torch") F = pytest.importorskip("torch.nn.functional") pytest.importorskip("torch_musa") +auto_functionalized = pytest.importorskip( + "torch._higher_order_ops.auto_functionalize" +).auto_functionalized + + +def _logical_graph_targets(graph) -> set[object]: + targets = {node.target for node in graph.nodes} + for node in graph.nodes: + if node.target is auto_functionalized and node.args: + targets.add(node.args[0]) + return targets @pytest.fixture(scope="module", autouse=True) @@ -246,10 +257,10 @@ def __init__(self, fusion=None): self.inductor_config["post_grad_custom_post_pass"] = self.post_pass def post_pass(self, graph): - self.targets_before = {node.target for node in graph.nodes} + self.targets_before = _logical_graph_targets(graph) if self.fusion is not None: self.fusion(graph) - self.targets_after = {node.target for node in graph.nodes} + self.targets_after = _logical_graph_targets(graph) def __call__(self, graph_module, example_inputs): return compile_fx( diff --git a/tests/test_silu_deepgemm_fusion_source.py b/tests/test_silu_deepgemm_fusion_source.py index 5229142ceaa1..9969d3c63dbe 100644 --- a/tests/test_silu_deepgemm_fusion_source.py +++ b/tests/test_silu_deepgemm_fusion_source.py @@ -31,10 +31,16 @@ def test_fusion_matches_the_production_native_swiglu_deepgemm_pair(): ).read_text() assert "MatcherSiluAndMul(enabled=False)" in source - assert "torch.ops.vllm.musa_deepgemm_fp8_op(" in source + assert "from torch._higher_order_ops.auto_functionalize import" in source + assert "return auto_functionalized(" in source + assert "torch.ops.vllm.musa_deepgemm_fp8_op.default," in source + assert "group_size=128," in source + assert "use_deep_gemm_e8m0=False," in source + assert "output=None," in source + assert ")[0]" in source assert "torch.ops.vllm.musa_silu_deepgemm_fp8_op(" in source - assert source.count("\n 128,") == 2 - assert source.count("\n False,") == 2 + assert source.count("\n 128,") == 1 + assert source.count("\n False,") == 1 def test_pass_manager_keeps_the_experimental_scope_narrow(): diff --git a/tests/test_uniform_decode_logits_indices_source.py b/tests/test_uniform_decode_logits_indices_source.py index b73d1641cc4d..453e64ca5188 100644 --- a/tests/test_uniform_decode_logits_indices_source.py +++ b/tests/test_uniform_decode_logits_indices_source.py @@ -11,14 +11,16 @@ / "series" / "0090-perf-musa-unify-Qwen-runtime-fast-paths.patch" ) +CONTRACT = ROOT / "vllm_musa" / "optimization_contract" / "qwen.py" def test_uniform_decode_patch_reuses_uploaded_request_indices() -> None: source = PATCH.read_text() + contract = CONTRACT.read_text() assert "current_platform.is_musa()" in source - assert "not self.is_pooling_model" in source - assert "self.speculative_config is None" in source + assert "not execution.has_speculative_config" in contract + assert "not execution.is_pooling_model" in contract for architecture in ( "Qwen2ForCausalLM", "Qwen2MoeForCausalLM", @@ -27,7 +29,7 @@ def test_uniform_decode_patch_reuses_uploaded_request_indices() -> None: "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", ): - assert architecture in source + assert architecture in contract assert "VLLM_MUSA_QWEN_UNIFORM_DECODE_LOGITS_INDICES" not in source assert "max_num_scheduled_tokens == 1" in source assert "num_tokens_unpadded == num_reqs" in source diff --git a/vllm_musa/compilation/passes/rms_deepgemm_fusion.py b/vllm_musa/compilation/passes/rms_deepgemm_fusion.py index fa8b4254c980..400b35ee3485 100644 --- a/vllm_musa/compilation/passes/rms_deepgemm_fusion.py +++ b/vllm_musa/compilation/passes/rms_deepgemm_fusion.py @@ -3,6 +3,7 @@ import torch from torch import fx +from torch._higher_order_ops.auto_functionalize import auto_functionalized from vllm.compilation.passes.vllm_inductor_pass import ( VllmFusionPatternMatcherPass, VllmPatternReplacement, @@ -39,13 +40,15 @@ def _pattern( normalized, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( input, residual, norm_weight, 1e-6, None ) - output = torch.ops.vllm.musa_deepgemm_fp8_op( - normalized, - weight, - weight_scale, - 128, - False, - ) + output = auto_functionalized( + torch.ops.vllm.musa_deepgemm_fp8_op.default, + input=normalized, + weight=weight, + weight_scale=weight_scale, + group_size=128, + use_deep_gemm_e8m0=False, + output=None, + )[0] return output, residual_out return _pattern diff --git a/vllm_musa/compilation/passes/silu_deepgemm_fusion.py b/vllm_musa/compilation/passes/silu_deepgemm_fusion.py index 357156ab133a..e6889178b72e 100644 --- a/vllm_musa/compilation/passes/silu_deepgemm_fusion.py +++ b/vllm_musa/compilation/passes/silu_deepgemm_fusion.py @@ -3,6 +3,7 @@ import torch from torch import fx +from torch._higher_order_ops.auto_functionalize import auto_functionalized from vllm.compilation.passes.fusion.matcher_utils import MatcherSiluAndMul from vllm.compilation.passes.vllm_inductor_pass import ( VllmFusionPatternMatcherPass, @@ -39,13 +40,15 @@ def _pattern( weight_scale: torch.Tensor, ) -> torch.Tensor: activated = self.silu_and_mul_matcher(input) - return torch.ops.vllm.musa_deepgemm_fp8_op( - activated, - weight, - weight_scale, - 128, - False, - ) + return auto_functionalized( + torch.ops.vllm.musa_deepgemm_fp8_op.default, + input=activated, + weight=weight, + weight_scale=weight_scale, + group_size=128, + use_deep_gemm_e8m0=False, + output=None, + )[0] return _pattern diff --git a/vllm_musa/compilation/qwen3_qk_rope_kv_presplit.py b/vllm_musa/compilation/qwen3_qk_rope_kv_presplit.py index 6b3c8b243ed3..fd7116d2baf3 100644 --- a/vllm_musa/compilation/qwen3_qk_rope_kv_presplit.py +++ b/vllm_musa/compilation/qwen3_qk_rope_kv_presplit.py @@ -58,7 +58,7 @@ class Qwen3QKRoPEKVCandidate: k_weight: fx.Node positions: fx.Node cos_sin_cache: fx.Node - layer_name: str + layer_name: str | fx.Node def _is_call(node: fx.Node, qualified_name: str) -> bool: @@ -186,7 +186,7 @@ def plan_qwen3_qk_rope_kv_presplit( candidates: list[Qwen3QKRoPEKVCandidate] = [] seen_rope: set[fx.Node] = set() - seen_layer_names: set[str] = set() + seen_layer_names: set[str | fx.Node] = set() for kv_update in all_kv: if len(kv_update.users) != 1: return None @@ -197,14 +197,22 @@ def plan_qwen3_qk_rope_kv_presplit( key = _arg(attention, 1, "key") value = _arg(attention, 2, "value") layer_name = _arg(attention, 4, "layer_name") + kv_layer_name = _arg(kv_update, 2, "layer_name") dummy = _arg(attention, 7, "kv_cache_dummy_dep", None) + # torch >= 2.11 hoists vLLM's opaque LayerName into an FX placeholder. + # Preserve that node as the fused custom op already accepts LayerNameType. + same_layer_name = ( + kv_layer_name is layer_name + if isinstance(layer_name, fx.Node) + else kv_layer_name == layer_name + ) if ( not all(isinstance(node, fx.Node) for node in (query, key, value)) or dummy is not kv_update - or not isinstance(layer_name, str) + or not isinstance(layer_name, (str, fx.Node)) or _arg(kv_update, 0, "key") is not key or _arg(kv_update, 1, "value") is not value - or _arg(kv_update, 2, "layer_name") != layer_name + or not same_layer_name or layer_name in seen_layer_names ): return None diff --git a/vllm_musa/kernels/gated_qkv.py b/vllm_musa/kernels/gated_qkv.py index 3d3c3fba657c..99cb068893ad 100644 --- a/vllm_musa/kernels/gated_qkv.py +++ b/vllm_musa/kernels/gated_qkv.py @@ -35,8 +35,8 @@ def _gated_qk_norm_rope_token_kernel( k_out_ptr, position_stride, cache_stride, - epsilon, - weight_offset, + epsilon: tl.constexpr, + weight_offset: tl.constexpr, NUM_Q_HEADS: tl.constexpr, HEAD_DIM: tl.constexpr, ROTARY_DIM: tl.constexpr, diff --git a/vllm_musa/patches/series/0003-MUSA-vllm.compilation.passes.pass_manager.patch b/vllm_musa/patches/series/0003-MUSA-vllm.compilation.passes.pass_manager.patch index 20a92a1a70be..05e6703788d3 100644 --- a/vllm_musa/patches/series/0003-MUSA-vllm.compilation.passes.pass_manager.patch +++ b/vllm_musa/patches/series/0003-MUSA-vllm.compilation.passes.pass_manager.patch @@ -1,5 +1,5 @@ From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 -From: Elon Gao +From: musa Date: Wed, 3 Jun 2026 11:47:35 +0000 Subject: [PATCH] MUSA: vllm.compilation.passes.pass_manager diff --git a/vllm_musa/patches/series/0099-perf-reuse-common-gdn-full-decode-metadata.patch b/vllm_musa/patches/series/0099-perf-reuse-common-gdn-full-decode-metadata.patch index f0dd323b84d8..32ca13af6abd 100644 --- a/vllm_musa/patches/series/0099-perf-reuse-common-gdn-full-decode-metadata.patch +++ b/vllm_musa/patches/series/0099-perf-reuse-common-gdn-full-decode-metadata.patch @@ -1,5 +1,5 @@ -From 01ab4ea625085fda5ab5b6d0f82cbc3d0461f920 Mon Sep 17 00:00:00 2001 -From: codex +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa Date: Mon, 3 Aug 2026 06:09:42 +0000 Subject: [PATCH] perf: reuse common GDN full decode metadata @@ -110,4 +110,3 @@ index 0a034c131..1dfab4fa6 100644 num_prefills=num_prefills, -- 2.34.1 - diff --git a/vllm_musa/patches/series/0100-perf-musa-sharded-qwen-gumbel.patch b/vllm_musa/patches/series/0100-perf-musa-sharded-qwen-gumbel.patch index aab9b929dfdf..da570a4ab81d 100644 --- a/vllm_musa/patches/series/0100-perf-musa-sharded-qwen-gumbel.patch +++ b/vllm_musa/patches/series/0100-perf-musa-sharded-qwen-gumbel.patch @@ -1,5 +1,5 @@ -From e21ebcc656c5973792ccc6cfef0fa7b1205b9986 Mon Sep 17 00:00:00 2001 -From: codex +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa Date: Fri, 31 Jul 2026 14:28:18 +0800 Subject: [PATCH] perf: add gated sharded Qwen Gumbel logits path diff --git a/vllm_musa/patches/series/0100-perf-skip-ssm-slot-mapping-in-uniform-decode.patch b/vllm_musa/patches/series/0100-perf-skip-ssm-slot-mapping-in-uniform-decode.patch index 26faf939a002..a2e91a4654dd 100644 --- a/vllm_musa/patches/series/0100-perf-skip-ssm-slot-mapping-in-uniform-decode.patch +++ b/vllm_musa/patches/series/0100-perf-skip-ssm-slot-mapping-in-uniform-decode.patch @@ -1,5 +1,5 @@ -From 0000000000000000000000000000000000000100 Mon Sep 17 00:00:00 2001 -From: codex +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa Date: Mon, 3 Aug 2026 07:15:00 +0000 Subject: [PATCH] perf: skip SSM slot mapping in uniform decode @@ -69,4 +69,3 @@ diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runne ) -> None: -- 2.34.1 - diff --git a/vllm_musa/patches/series/0101-perf-specialize-qwen36-bf16-w1-c25.patch b/vllm_musa/patches/series/0101-perf-specialize-qwen36-bf16-w1-c25.patch index 669202b70929..af22e9baa9e6 100644 --- a/vllm_musa/patches/series/0101-perf-specialize-qwen36-bf16-w1-c25.patch +++ b/vllm_musa/patches/series/0101-perf-specialize-qwen36-bf16-w1-c25.patch @@ -1,5 +1,5 @@ -From 0000000000000000000000000000000000000101 Mon Sep 17 00:00:00 2001 -From: codex +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa Date: Tue, 4 Aug 2026 06:30:00 +0800 Subject: [PATCH] perf: specialize Qwen3.6 BF16 W1 tile at C25 diff --git a/vllm_musa/platform.py b/vllm_musa/platform.py index 73102b02f774..c4d227865a25 100644 --- a/vllm_musa/platform.py +++ b/vllm_musa/platform.py @@ -9,6 +9,8 @@ from functools import cache, wraps from typing import TYPE_CHECKING, Any, TypeVar +from packaging.version import Version + # isort: off import torchada # noqa: F401 import torch @@ -39,6 +41,13 @@ logger = init_logger(__name__) +_QWEN3_VL_ARCHITECTURES = {"Qwen3VLForConditionalGeneration"} + + +def _is_torch_211_or_newer() -> bool: + return Version(torch.__version__.split("+", 1)[0]) >= Version("2.11") + + _P = ParamSpec("_P") _R = TypeVar("_R") @@ -382,6 +391,30 @@ def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: if all(s not in compilation_config.custom_ops for s in ("all", "none")): compilation_config.custom_ops.append("all") + # torch 2.11's Inductor tiling heuristic turns Qwen3-VL's decode-time + # clone/index-select/split kernel into a slower 2D grid on MUSA. The + # torch 2.9 1D grid is restored by limiting this model's pointwise + # kernels to one tile dimension. Preserve an explicit user override. + architectures = set( + getattr( + getattr(vllm_config.model_config, "hf_config", None), + "architectures", + None, + ) + or () + ) + inductor_config = compilation_config.inductor_compile_config + if ( + _is_torch_211_or_newer() + and architectures & _QWEN3_VL_ARCHITECTURES + and "triton.max_tiles" not in inductor_config + ): + inductor_config["triton.max_tiles"] = 1 + logger.info( + "Using triton.max_tiles=1 for Qwen3-VL on torch >= 2.11 " + "to preserve the MUSA decode pointwise layout" + ) + @classmethod def check_and_update_config(cls, vllm_config: "VllmConfig") -> None: # when dflash spec-decode is active, coerce the draft-loop