Compare commits

...
14 Commits
Author SHA1 Message Date
khluu 0decac0d96 fix: resolve CUTLASS fmin compatibility for DeepSeek-V4 init
Signed-off-by: khluu <khluu000@gmail.com>
2026-06-03 17:11:47 -07:00
Harry Mellorandkhluu fd56c57bde Fix OlmoHybridForCausalLM not initialising (#43846)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
(cherry picked from commit 19af4e6dd4)
2026-06-03 16:56:07 -07:00
Kevin H. Luu 7285178622 [Bugfix] Fix HyperCLOVAX CI failure after upstream removed remote code (#43860)
Signed-off-by: Kevin Luu <kevin@inferact.ai>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
(cherry picked from commit 61288b5458)
2026-06-03 16:55:00 -07:00
Alecandkhluu 27509c8dde [Bugfix][CI] Normalize NIXL connector CUDA wheel installs (#44266)
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
(cherry picked from commit 816cc73a9b)
2026-06-02 23:21:24 -07:00
Kevin H. Luu b284862ea9 [docker] Stop using extra-index-url for flashinfer-jit-cache (#44366)
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2026-06-02 19:02:03 -07:00
932dfd5276 [Feature] Add support for JetBrains' Mellum v2 code generation model (#43992)
Signed-off-by: Madeesh Kannan <madeeswaran.kannan@jetbrains.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
2026-06-02 19:01:56 -07:00
682ffebfef [CPU][Zen] Route W8A8 and W4A16 linear inference through zentorch on AMD Zen CPUs (#41813)
Signed-off-by: R <Ganesh.R@amd.com>
Signed-off-by: Harshal Adhav <harshal.adhav@amd.com>
Signed-off-by: Aakar Dwivedi <aadwived@amd.com>
Co-authored-by: R <Ganesh.R@amd.com>
Co-authored-by: Harshal Adhav <harshal.adhav@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
2026-06-02 19:01:49 -07:00
Vadim Gimpelsonandkhluu 1be7a57a18 [Bugfix] Exclude Ray DP from #42585's deferred port allocation (#43864)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
2026-06-02 19:01:42 -07:00
khluu 0b3ba88f16 Revert "[CPU] Experimentally enable Triton and MRV2 (#43225)"
This reverts commit 65b7a812a2.
2026-05-29 02:28:43 -07:00
799c3afa5d [BugFix] Fix hard-coded timeout for multi-API-server startup (#43768)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-05-28 00:11:54 -07:00
Thien Tranandkhluu 64e25235c7 [Bugfix] Pass routed_scaling_factor to FlashInfer TRTLLM BF16 MoE (#43769) 2026-05-28 00:11:49 -07:00
TJianandkhluu a147dd0115 [ROCm][DSV4] Enable Tilelang MHC replacing torch/triton mhc (#43679)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-05-28 00:11:43 -07:00
amitz-nvandkhluu 0759293512 [Bugfix][Kernel] TRTLLM NVFP4 MoE chunking (#43599)
Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com>
2026-05-28 00:11:38 -07:00
Benjamin Bartelsandkhluu a930f5a58d Fix RunAI streamer tensor buffer reuse during weight loading (#43464)
Signed-off-by: bbartels <benjamin@bartels.dev>
2026-05-28 00:11:32 -07:00
52 changed files with 1746 additions and 324 deletions
-14
View File
@@ -54,20 +54,6 @@ steps:
pytest -x -v -s tests/models/language/generation -m cpu_model
pytest -x -v -s tests/models/language/pooling -m cpu_model"
- label: CPU-ModelRunnerV2 Tests
depends_on: []
device: intel_cpu
no_plugin: true
soft_fail: true
source_file_dependencies:
- vllm/v1/worker/cpu/
- vllm/v1/worker/gpu/
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
- label: CPU-Quantization Model Tests
depends_on: []
device: intel_cpu
+39
View File
@@ -0,0 +1,39 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
set -euo pipefail
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
uv pip install --system -r "${REQUIREMENTS_FILE}"
NIXL_METADATA=$(python3 - <<'PY'
import importlib.metadata as metadata
import torch
cuda_version = torch.version.cuda
if cuda_version is None:
raise SystemExit("torch.version.cuda is not set")
print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
PY
)
read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
# matching this CI image so nixl_ep_cpp links against the available libcudart.
uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
python3 - <<'PY'
import importlib.metadata as metadata
for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
try:
version = metadata.version(package_name)
except metadata.PackageNotFoundError:
version = "not installed"
print(f"{package_name}: {version}")
PY
+9 -9
View File
@@ -11,7 +11,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
@@ -22,7 +22,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -34,7 +34,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -46,7 +46,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
@@ -58,7 +58,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
@@ -73,7 +73,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
@@ -87,7 +87,7 @@ steps:
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
@@ -102,5 +102,5 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+1 -1
View File
@@ -86,7 +86,7 @@ steps:
- tests/v1/metrics
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# split the test to avoid interference
- pytest -v -s -m 'not cpu_test' v1/core
+1 -1
View File
@@ -760,7 +760,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
ARG FLASHINFER_VERSION=0.6.11.post2
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
# ============================================================
# OPENAI API SERVER DEPENDENCIES
+6 -5
View File
@@ -27,14 +27,11 @@ WORKDIR /workspace
ARG PYTHON_VERSION=3.12
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
ARG max_jobs=32
ENV MAX_JOBS=${max_jobs}
# Install minimal dependencies and uv
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
--mount=type=cache,target=/var/lib/apt,sharing=locked \
apt-get update -y \
&& apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \
&& apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates \
gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
&& curl -LsSf https://astral.sh/uv/install.sh | sh
@@ -126,6 +123,9 @@ RUN --mount=type=cache,target=/root/.cargo/registry \
######################### BUILD IMAGE #########################
FROM base AS vllm-build
ARG max_jobs=32
ENV MAX_JOBS=${max_jobs}
ARG GIT_REPO_CHECK=0
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
ARG VLLM_CPU_X86=0
@@ -257,7 +257,8 @@ WORKDIR /vllm-workspace
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=cache,target=/root/.cache/ccache \
--mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
uv pip install dist/*.whl && \
uv pip install "vllm[audio]"
# Add labels to document build configuration
LABEL org.opencontainers.image.title="vLLM CPU"
+1
View File
@@ -438,6 +438,7 @@ th {
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
+1
View File
@@ -16,3 +16,4 @@ wheel
jinja2>=3.1.6
amdsmi==7.0.2
timm>=1.0.17
tilelang==0.1.10
+1
View File
@@ -22,3 +22,4 @@ timm>=1.0.17
# amd-quark: required for Quark quantization on ROCm
# To be consistent with test_quark.py
amd-quark>=0.8.99
tilelang==0.1.10
+1
View File
@@ -43,6 +43,7 @@ schemathesis>=3.39.15 # Required for openai schema test
# quantization
bitsandbytes==0.49.2
buildkite-test-collector==0.1.9
tilelang==0.1.10
genai_perf>=0.0.8
tritonclient>=2.51.0
+21 -2
View File
@@ -43,7 +43,9 @@ anyio==4.13.0
# starlette
# watchfiles
apache-tvm-ffi==0.1.10
# via xgrammar
# via
# tilelang
# xgrammar
arctic-inference==0.1.1
# via -r requirements/test/rocm.in
argcomplete==3.6.3
@@ -129,7 +131,9 @@ click==8.3.1
# typer
# uvicorn
cloudpickle==3.1.2
# via -r requirements/test/../common.txt
# via
# -r requirements/test/../common.txt
# tilelang
colorama==0.4.6
# via
# perceptron
@@ -511,6 +515,8 @@ mistral-common==1.11.2
# -c requirements/common.txt
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
ml-dtypes==0.5.4
# via tilelang
model-hosting-container-standards==0.1.14
# via
# -c requirements/common.txt
@@ -587,6 +593,7 @@ numpy==2.2.6
# lm-eval
# matplotlib
# mistral-common
# ml-dtypes
# mteb
# numba
# opencv-python-headless
@@ -610,6 +617,7 @@ numpy==2.2.6
# statsmodels
# tensorizer
# tifffile
# tilelang
# torchvision
# transformers
# tritonclient
@@ -811,6 +819,7 @@ psutil==7.2.2
# accelerate
# peft
# tensorizer
# tilelang
py==1.11.0
# via pytest-forked
py-cpuinfo==9.0.0
@@ -1192,6 +1201,10 @@ tiktoken==0.12.0
# gpt-oss
# lm-eval
# mistral-common
tilelang==0.1.10
# via
# -c requirements/rocm.txt
# -r requirements/test/rocm.in
timm==1.0.17
# via
# -c requirements/rocm.txt
@@ -1208,6 +1221,8 @@ tomli==2.4.0
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch-c-dlpack-ext==0.1.5
# via tilelang
tqdm==4.67.3
# via
# -r requirements/test/../common.txt
@@ -1225,6 +1240,7 @@ tqdm==4.67.3
# pqdm
# segmentation-models-pytorch
# sentence-transformers
# tilelang
# transformers
transformers==5.5.3
# via
@@ -1293,6 +1309,7 @@ typing-extensions==4.15.0
# sentence-transformers
# sqlalchemy
# starlette
# tilelang
# torch
# typeguard
# typing-inspection
@@ -1359,6 +1376,8 @@ yarl==1.23.0
# via
# aiohttp
# schemathesis
z3-solver==4.15.4.0
# via tilelang
zipp==3.23.0
# via importlib-metadata
+1 -8
View File
@@ -1165,9 +1165,7 @@ setup(
install_requires=get_requirements(),
extras_require={
# AMD Zen CPU optimizations via zentorch
"zen": [
"zentorch-weekly==5.2.1.dev20260408"
], # Zentorch has weekly releases. This pulls the known-good version.
"zen": ["zentorch==2.11.0.0"],
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
@@ -1195,11 +1193,6 @@ setup(
"opentelemetry-exporter-otlp>=1.26.0",
"opentelemetry-semantic-conventions-ai>=0.4.1",
],
"triton-cpu": [
"triton @ "
"git+https://github.com/triton-lang/triton-cpu.git@270e696d ; "
"platform_machine == 'x86_64'",
], # Remove after stable release
},
cmdclass=cmdclass,
package_data=package_data,
+5 -2
View File
@@ -1656,7 +1656,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
layer.routing_method_type = RoutingMethodType.Renormalize
layer.expert_map = None
layer.apply_router_weight_on_input = False
layer.routed_scaling_factor = None
layer.routed_scaling_factor = 2.446
layer.shared_experts = None
layer._expert_routing_tables = lambda: None
@@ -1678,7 +1678,10 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
# Compute torch baseline
w1_original = w1.clone()
w2_original = w2.clone()
baseline_output = torch_moe(a, w1_original, w2_original, router_logits, topk)
baseline_output = (
torch_moe(a, w1_original, w2_original, router_logits, topk)
* layer.routed_scaling_factor
)
close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85)
assert close.float().mean() > 0.925
+166 -2
View File
@@ -4,7 +4,12 @@ import pytest
import torch
import vllm.model_executor.kernels.mhc # noqa: F401
from vllm.model_executor.kernels.mhc.tilelang import (
_tilelang_hc_prenorm_gemm,
_torch_hc_prenorm_gemm,
)
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_tilelang
from vllm.utils.torch_utils import set_random_seed
DEVICE = current_platform.device_type
@@ -92,8 +97,128 @@ def hc_head_ref(
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason="CUDA required",
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@pytest.mark.parametrize("hc_mult", [4])
def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
torch.set_default_device(DEVICE)
set_random_seed(0)
residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
hc_mult2 = hc_mult * hc_mult
hc_mult3 = 2 * hc_mult + hc_mult2
fn = (
torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float)
* 1e-4
* (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1))
).flatten(1, 2)
hc_scale = torch.randn((3,), dtype=torch.float) * 0.1
hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1
hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6
sinkhorn_repeat = 20
hc_post_alpha = 1.0
ref = mhc_pre_ref(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_alpha,
sinkhorn_repeat,
)
out = torch.ops.vllm.mhc_pre_tilelang(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_alpha,
sinkhorn_repeat,
)
for actual, expected in zip(out, ref, strict=True):
torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
)
@pytest.mark.parametrize(
("num_tokens", "hidden_size"),
[
(1, 1280),
(512, 1280),
(2048, 1280),
(1, 4096),
(64, 4096),
(512, 4096),
(2048, 4096),
(1, 7168),
(64, 7168),
(512, 7168),
(2048, 7168),
],
)
def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
torch.set_default_device(DEVICE)
set_random_seed(0)
hc_mult = 4
hc_mult3 = 2 * hc_mult + hc_mult * hc_mult
x = torch.randn((num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16)
fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
out_ref = torch.empty((1, num_tokens, hc_mult3), dtype=torch.float32)
sqrsum_ref = torch.empty((1, num_tokens), dtype=torch.float32)
out = torch.empty_like(out_ref)
sqrsum = torch.empty_like(sqrsum_ref)
_torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref)
_tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)
torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)
torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4)
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@pytest.mark.parametrize("hc_mult", [4])
def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult):
torch.set_default_device(DEVICE)
set_random_seed(0)
x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix)
out = torch.ops.vllm.mhc_post_tilelang(
x,
residual,
post_layer_mix,
comb_res_mix,
)
torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -196,3 +321,42 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() and has_tilelang()),
reason="CUDA or ROCm and tilelang required",
)
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
@pytest.mark.parametrize("hc_mult", [4])
def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
torch.set_default_device(DEVICE)
set_random_seed(0)
residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
rms_eps = hc_eps = 1e-6
out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)
out.fill_(float("nan"))
result = torch.ops.vllm.hc_head_fused_kernel_tilelang(
residual,
fn,
hc_scale,
hc_base,
out,
hidden_size,
rms_eps,
hc_eps,
hc_mult,
)
assert result is None
assert not torch.isnan(out).any()
out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
@@ -6,6 +6,7 @@ import tempfile
import huggingface_hub.constants
import torch
from safetensors.torch import save_file
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf,
@@ -14,6 +15,27 @@ from vllm.model_executor.model_loader.weight_utils import (
)
def test_runai_safetensors_weights_iterator_clones_reused_buffers(
tmp_path, monkeypatch
):
monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0")
weights_file = tmp_path / "model.safetensors"
expected_tensors = {
"first": torch.tensor([1.0, 2.0]),
"second": torch.tensor([3.0, 4.0]),
}
save_file(expected_tensors, weights_file)
actual_tensors = dict(
runai_safetensors_weights_iterator([str(weights_file)], False)
)
assert actual_tensors.keys() == expected_tensors.keys()
assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr()
for name, expected_tensor in expected_tensors.items():
assert torch.equal(actual_tensors[name], expected_tensor)
def test_runai_model_loader():
with tempfile.TemporaryDirectory() as tmpdir:
huggingface_hub.constants.HF_HUB_OFFLINE = False
+2 -1
View File
@@ -335,7 +335,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"HYV3ForCausalLM": _HfExamplesInfo("tencent/Hy3-preview", trust_remote_code=True),
"HyperCLOVAXForCausalLM": _HfExamplesInfo(
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
trust_remote_code=True,
min_transformers_version="5.9.0",
),
"InternLMForCausalLM": _HfExamplesInfo(
"internlm/internlm-chat-7b", trust_remote_code=True
@@ -523,6 +523,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
"Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
"Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
"MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"),
"Qwen3NextForCausalLM": _HfExamplesInfo(
"Qwen/Qwen3-Next-80B-A3B-Instruct",
extras={"tiny-random": "tiny-random/qwen3-next-moe"},
@@ -2,6 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import socket
import time
import uuid
from pathlib import Path
from types import SimpleNamespace
@@ -9,9 +11,17 @@ from typing import Any
import pytest
import ray
import zmq
from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
from vllm.v1.engine.core import EngineCoreActorMixin
from vllm.v1.engine.utils import CoreEngineActorManager, EngineZmqAddresses
from vllm.v1.engine.utils import (
CoreEngineActorManager,
EngineZmqAddresses,
get_engine_zmq_addresses,
launch_core_engines,
)
from vllm.v1.utils import APIServerProcessManager
class _StubEngineCoreActor(EngineCoreActorMixin):
@@ -42,6 +52,48 @@ class _StubEngineCoreActor(EngineCoreActorMixin):
def get_nixl_side_channel_host(self) -> str | None:
return os.environ.get("VLLM_NIXL_SIDE_CHANNEL_HOST")
def get_addresses(self) -> tuple[list[str], list[str]]:
"""Return the addresses snapshot the actor was constructed with.
Used by the Ray-DP regression test to assert that no ``tcp://host:0``
placeholders were pickled into the actor at ``.remote()`` time.
"""
return list(self.addresses.inputs), list(self.addresses.outputs)
# Module-level stub worker for the Ray-DP regression test. Must be importable
# by ``multiprocessing.spawn`` (no closures, no nesting). Mirrors the worker
# in ``tests/entrypoints/test_api_server_process_manager.py``.
def _bind_and_report_worker(listen_address, sock, args, client_config):
"""Bind ROUTER/PULL with a kernel-assigned port, report the actual
endpoints back via ``actual_address_pipe``, then exit."""
ctx = zmq.Context()
try:
in_sock = make_zmq_socket(
ctx, client_config["input_address"], zmq.ROUTER, bind=True
)
out_sock = make_zmq_socket(
ctx, client_config["output_address"], zmq.PULL, bind=True
)
try:
pipe = client_config["actual_address_pipe"]
try:
pipe.send(
{
"input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(),
"output_address": out_sock.getsockopt(
zmq.LAST_ENDPOINT
).decode(),
}
)
finally:
pipe.close()
finally:
in_sock.close(linger=0)
out_sock.close(linger=0)
finally:
ctx.term()
class _DummyExecutor:
pass
@@ -134,3 +186,172 @@ def test_driver_nixl_side_channel_host_does_not_leak_to_engine_core_actor(
else:
for pg in created_placement_groups:
ray.util.remove_placement_group(pg)
@pytest.fixture
def ray_context_dp2():
"""Ray context sized for two stub actors (each PG needs ~1 CPU)."""
started_ray = False
if not ray.is_initialized():
project_root = str(Path(__file__).resolve().parents[3])
ray.init(
num_cpus=4,
runtime_env={"env_vars": {"PYTHONPATH": project_root}},
log_to_driver=False,
)
started_ray = True
yield
if started_ray:
ray.shutdown()
def _make_vllm_config_ray_dp_multinode() -> SimpleNamespace:
"""Minimal vllm_config that drives the Ray-DP multi-API-server path:
``data_parallel_size != data_parallel_size_local`` forces TCP placeholders
(multi-node fan-out), and ``data_parallel_backend="ray"`` routes
``launch_core_engines`` through the Ray branch.
"""
return SimpleNamespace(
parallel_config=SimpleNamespace(
data_parallel_size=2,
data_parallel_size_local=1,
data_parallel_rank=0,
data_parallel_rank_local=None,
data_parallel_master_ip="127.0.0.1",
data_parallel_backend="ray",
data_parallel_rpc_port=29550,
local_engines_only=False,
enable_elastic_ep=False,
world_size=1,
),
model_config=SimpleNamespace(multimodal_config=None, is_moe=False),
cache_config=SimpleNamespace(),
needs_dp_coordinator=False,
kv_transfer_config=None,
# ``_apply_dp_identity_suffix`` reads and rewrites this.
instance_id="vllm-ray-dp-regression-test",
)
@pytest.mark.timeout(120)
@pytest.mark.usefixtures("ray_context_dp2")
def test_ray_dp_addresses_resolved_before_actor_creation(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Regression guard for the Ray-DP + multi-API-server hang from PR #42585.
``launch_core_engines`` Ray branch pickles ``addresses`` into each engine
actor at ``.remote()`` time, and ``EngineCoreActorMixin._perform_handshakes``
is a no-op, so the actor uses that pickled snapshot for the rest of its
life. If ``run_multi_api_server`` allocates ``addresses`` as
``tcp://host:0`` placeholders (its default), the actors hold placeholders
forever and DEALER-connect to port 0 — ZMQ ``connect`` is async and does
not raise, so the failure mode is a deterministic hang.
The Ray-DP carve-out in ``run_multi_api_server`` forces
``defer_api_server_ports=False`` when ``data_parallel_backend == "ray"``
so addresses are pre-allocated in the driver and Ray pickles real ports
into each actor. This test mirrors that call-site logic and asserts the
actors hold real (non-placeholder) endpoints. If the carve-out is
removed without an alternative fix, the test fails.
"""
created_placement_groups: list[Any] = []
def create_dp_placement_groups(vllm_config: Any):
pg1 = _make_cpu_placement_group()
pg2 = _make_cpu_placement_group()
created_placement_groups.extend([pg1, pg2])
return [pg1, pg2], [0, 0]
monkeypatch.setattr("vllm.v1.engine.core.EngineCoreActor", _StubEngineCoreActor)
monkeypatch.setattr(
CoreEngineActorManager,
"create_dp_placement_groups",
staticmethod(create_dp_placement_groups),
)
vllm_config = _make_vllm_config_ray_dp_multinode()
# Mirror run_multi_api_server's address-allocation logic. The Ray DP
# carve-out forces pre-allocation so the addresses pickled into engine
# actors at .remote() time are real, not ``tcp://host:0``.
is_ray_dp = vllm_config.parallel_config.data_parallel_backend == "ray"
addresses = get_engine_zmq_addresses(
vllm_config,
num_api_servers=2,
defer_api_server_ports=not is_ray_dp,
)
sock = socket.socket()
engine_manager: CoreEngineActorManager | None = None
actor_snapshots: list[tuple[list[str], list[str]]] = []
api_server_manager: APIServerProcessManager | None = None
try:
# Ray actors are spawned here, pickling ``addresses`` into each one.
with launch_core_engines(
vllm_config,
executor_class=_DummyExecutor,
log_stats=False,
addresses=addresses,
num_api_servers=2,
) as (
engine_manager,
_coordinator,
_addresses_out,
_tensor_queue,
):
assert isinstance(engine_manager, CoreEngineActorManager)
# API-server children bind to the pre-allocated ports.
api_server_manager = APIServerProcessManager(
listen_address="tcp://127.0.0.1:0",
sock=sock,
args="test_args",
num_servers=2,
input_addresses=addresses.inputs,
output_addresses=addresses.outputs,
target_server_fn=_bind_and_report_worker,
)
# run_multi_api_server skips ``gather_actual_addresses`` for
# Ray DP (addresses are already real). Mirror that.
if not is_ray_dp:
actual_inputs, actual_outputs = (
api_server_manager.gather_actual_addresses(timeout=15.0)
)
addresses.inputs = actual_inputs
addresses.outputs = actual_outputs
# Snapshot what each Ray actor actually holds.
actors = (
engine_manager.local_engine_actors + engine_manager.remote_engine_actors
)
actor_snapshots = ray.get(
[actor.get_addresses.remote() for actor in actors]
)
finally:
if api_server_manager is not None:
api_server_manager.shutdown()
time.sleep(0.2)
sock.close()
if engine_manager is not None:
engine_manager.shutdown()
else:
for pg in created_placement_groups:
ray.util.remove_placement_group(pg)
# Every Ray actor must hold real, non-placeholder addresses.
assert actor_snapshots, "expected at least one Ray actor to be created"
for actor_inputs, actor_outputs in actor_snapshots:
for url in actor_inputs + actor_outputs:
scheme, _host, port = split_zmq_path(url)
assert scheme == "tcp", url
assert port and int(port) > 0, (
f"Ray actor was pickled with placeholder address {url!r}; "
"``run_multi_api_server`` must pre-allocate ports for the "
"Ray DP backend so the actors hold real endpoints by the "
"time they DEALER-connect. See PR #42585 / Ray-DP "
"multi-API-server regression."
)
+224 -40
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
from functools import cache
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
import torch
@@ -10,8 +10,9 @@ from vllm.platforms import current_platform
from vllm.utils.import_utils import has_tilelang
from vllm.utils.math_utils import cdiv
# tilelang is only available on CUDA platforms
if TYPE_CHECKING or current_platform.is_cuda():
# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so
# registering the Python wrapper modules does not require TileLang everywhere.
if TYPE_CHECKING or current_platform.is_cuda_alike():
if not has_tilelang():
raise ImportError(
"tilelang is required for mhc but is not installed. Install it with "
@@ -23,6 +24,8 @@ else:
tilelang = None # type: ignore[assignment]
T = None # type: ignore[assignment]
ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()
@cache
def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
@@ -37,12 +40,17 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
return split_k
pass_configs: dict[tilelang.PassConfigKey, Any] = {
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
}
if current_platform.is_cuda():
pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
},
pass_configs=pass_configs,
)
def mhc_pre_big_fuse_tilelang(
gemm_out_mul,
@@ -78,7 +86,8 @@ def mhc_pre_big_fuse_tilelang(
layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
with T.Kernel(num_tokens, threads=96) as i:
T.pdl_sync()
if ENABLE_PDL:
T.pdl_sync()
##################################################################
# _pre_norm_fn_fwd_norm
rms = T.alloc_fragment(1, T.float32)
@@ -174,18 +183,16 @@ def mhc_pre_big_fuse_tilelang(
ol[i1_h] += pre * xl[i_hc, i1_h]
T.copy(ol, layer_input[i, i0_h * hidden_block])
T.pdl_trigger()
if ENABLE_PDL:
T.pdl_trigger()
# Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
},
pass_configs=pass_configs,
)
def mhc_pre_big_fuse_with_norm_tilelang(
gemm_out_mul,
@@ -230,7 +237,8 @@ def mhc_pre_big_fuse_with_norm_tilelang(
T.clear(mixes)
rms[0] = 0
T.pdl_sync()
if ENABLE_PDL:
T.pdl_sync()
for i_split in T.serial(n_splits):
rms[0] += gemm_out_sqrsum[i_split, i]
@@ -341,15 +349,12 @@ def mhc_pre_big_fuse_with_norm_tilelang(
T.copy(ol, layer_input[i, i0_h * hidden_block])
T.pdl_trigger()
if ENABLE_PDL:
T.pdl_trigger()
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
},
pass_configs=pass_configs,
)
def mhc_fused_tilelang(
comb_mix,
@@ -390,8 +395,8 @@ def mhc_fused_tilelang(
with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
tid = T.get_thread_binding()
warp_id = T.get_warp_idx()
lane = T.get_lane_idx()
warp_id = tid // 32
lane = tid % 32
s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
s_post = T.alloc_shared((hc,), T.float32)
@@ -407,7 +412,8 @@ def mhc_fused_tilelang(
T.clear(sqr)
h_split_start = i_ks * h_per_split
T.pdl_sync()
if ENABLE_PDL:
T.pdl_sync()
T.copy(post_mix[i_n, 0], s_post)
T.copy(comb_mix[i_n, 0, 0], s_comb)
@@ -466,15 +472,12 @@ def mhc_fused_tilelang(
v2 += s_warp[w, tile_n]
rp_out[i_ks, i_n] = v2
T.pdl_trigger()
if ENABLE_PDL:
T.pdl_trigger()
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
},
pass_configs=pass_configs,
)
def mhc_post_tilelang(
a,
@@ -507,7 +510,8 @@ def mhc_post_tilelang(
a_local = T.alloc_fragment((hc, hc), T.float32)
c_local = T.alloc_fragment(hc, T.float32)
T.pdl_sync()
if ENABLE_PDL:
T.pdl_sync()
T.copy(a[i_n, 0, 0], a_local)
T.copy(c[i_n, 0], c_local)
@@ -523,15 +527,193 @@ def mhc_post_tilelang(
x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
T.copy(x_local, x[i_n, 0, i0_h * h_blk])
T.pdl_trigger()
if ENABLE_PDL:
T.pdl_trigger()
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
},
pass_configs=pass_configs,
)
def hc_prenorm_gemm_tilelang(
x,
fn,
out,
sqrsum,
hidden_size: int,
hc_mult: int = 4,
n_out: int = 24,
n_thr: int = 512,
tile_n: int = 12,
n_splits: int = 1,
) -> tilelang.JITKernel:
num_tokens = T.dynamic("num_tokens")
hc_hidden_size = hc_mult * hidden_size
k_per_split = hc_hidden_size // n_splits
k_iters = k_per_split // n_thr
n_tiles = T.ceildiv(n_out, tile_n)
x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type]
fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type]
out: T.Tensor((n_splits, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type]
sqrsum: T.Tensor((n_splits, num_tokens), T.float32) # type: ignore[no-redef, valid-type]
with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as (
i_n,
i_t,
i_s,
):
tid = T.get_thread_binding()
acc = T.alloc_local((tile_n,), T.float32)
sqr = T.alloc_local((1,), T.float32)
T.clear(acc)
T.clear(sqr)
if ENABLE_PDL:
T.pdl_sync()
for it in T.serial(k_iters):
i_k = i_s * k_per_split + it * n_thr + tid
x_val = x[i_n, i_k]
for i_o in T.unroll(tile_n):
out_idx = i_t * tile_n + i_o
if out_idx < n_out:
acc[i_o] += x_val * fn[out_idx, i_k]
if i_t == 0:
sqr[0] += x_val * x_val
for i_o in T.unroll(tile_n):
acc[i_o] = T.warp_reduce_sum(acc[i_o])
if i_t == 0:
sqr[0] = T.warp_reduce_sum(sqr[0])
lane = tid % 32
warp_id = tid // 32
num_warps = n_thr // 32
warp_acc = T.alloc_shared((num_warps, tile_n), T.float32)
warp_sqr = T.alloc_shared(num_warps, T.float32)
if lane == 0:
for i_o in T.unroll(tile_n):
warp_acc[warp_id, i_o] = acc[i_o]
if i_t == 0:
warp_sqr[warp_id] = sqr[0]
T.sync_threads()
if warp_id == 0:
if lane < tile_n:
reduced_acc = T.alloc_var(T.float32, init=0.0)
for i_w in T.unroll(num_warps):
reduced_acc += warp_acc[i_w, lane]
out_idx = i_t * tile_n + lane
if out_idx < n_out:
out[i_s, i_n, out_idx] = reduced_acc
if lane == 0 and i_t == 0:
reduced_sqr = T.alloc_var(T.float32, init=0.0)
for i_w in T.unroll(num_warps):
reduced_sqr += warp_sqr[i_w]
sqrsum[i_s, i_n] = reduced_sqr
if ENABLE_PDL:
T.pdl_trigger()
@tilelang.jit(
pass_configs=pass_configs,
)
def hc_prenorm_gemm_block_m_tilelang(
x,
fn,
out,
sqrsum,
hidden_size: int,
hc_mult: int = 4,
n_out: int = 24,
n_thr: int = 512,
tile_n: int = 12,
block_m: int = 2,
) -> tilelang.JITKernel:
num_tokens = T.dynamic("num_tokens")
hc_hidden_size = hc_mult * hidden_size
k_iters = hc_hidden_size // n_thr
n_tiles = T.ceildiv(n_out, tile_n)
m_tiles = T.ceildiv(num_tokens, block_m)
x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type]
fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type]
out: T.Tensor((1, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type]
sqrsum: T.Tensor((1, num_tokens), T.float32) # type: ignore[no-redef, valid-type]
with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t):
tid = T.get_thread_binding()
acc = T.alloc_local((block_m, tile_n), T.float32)
sqr = T.alloc_local((block_m,), T.float32)
T.clear(acc)
T.clear(sqr)
if ENABLE_PDL:
T.pdl_sync()
for it in T.serial(k_iters):
i_k = it * n_thr + tid
fn_val = T.alloc_local((tile_n,), T.float32)
for i_o in T.unroll(tile_n):
out_idx = i_t * tile_n + i_o
if out_idx < n_out:
fn_val[i_o] = fn[out_idx, i_k]
else:
fn_val[i_o] = 0.0
for i_m in T.unroll(block_m):
token_idx = i_mt * block_m + i_m
if token_idx < num_tokens:
x_val = x[token_idx, i_k]
for i_o in T.unroll(tile_n):
acc[i_m, i_o] += x_val * fn_val[i_o]
if i_t == 0:
sqr[i_m] += x_val * x_val
for i_m in T.unroll(block_m):
for i_o in T.unroll(tile_n):
acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o])
if i_t == 0:
sqr[i_m] = T.warp_reduce_sum(sqr[i_m])
lane = tid % 32
warp_id = tid // 32
num_warps = n_thr // 32
warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32)
warp_sqr = T.alloc_shared((num_warps, block_m), T.float32)
if lane == 0:
for i_m in T.unroll(block_m):
for i_o in T.unroll(tile_n):
warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o]
if i_t == 0:
warp_sqr[warp_id, i_m] = sqr[i_m]
T.sync_threads()
if warp_id == 0:
for i_m in T.unroll(block_m):
token_idx = i_mt * block_m + i_m
if token_idx < num_tokens:
if lane < tile_n:
reduced_acc = T.alloc_var(T.float32, init=0.0)
for i_w in T.unroll(num_warps):
reduced_acc += warp_acc[i_w, i_m, lane]
out_idx = i_t * tile_n + lane
if out_idx < n_out:
out[0, token_idx, out_idx] = reduced_acc
if lane == 0 and i_t == 0:
reduced_sqr = T.alloc_var(T.float32, init=0.0)
for i_w in T.unroll(num_warps):
reduced_sqr += warp_sqr[i_w, i_m]
sqrsum[0, token_idx] = reduced_sqr
if ENABLE_PDL:
T.pdl_trigger()
@tilelang.jit(
pass_configs=pass_configs,
)
def hc_head_fuse_tilelang(
residual,
@@ -566,7 +748,8 @@ def hc_head_fuse_tilelang(
out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]
with T.Kernel(num_tokens, threads=n_thr) as i:
T.pdl_sync()
if ENABLE_PDL:
T.pdl_sync()
# ------------------------------------------------------------------
# Pass 1 for each residual channel m_c and h_block:
@@ -624,4 +807,5 @@ def hc_head_fuse_tilelang(
T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
T.pdl_trigger()
if ENABLE_PDL:
T.pdl_trigger()
+16 -9
View File
@@ -308,13 +308,16 @@ def run_multi_api_server(args: argparse.Namespace):
from vllm.v1.engine.utils import get_engine_zmq_addresses
# Per-API-server ports are picked by the kernel at each child's bind()
# to avoid parent-probe vs child-bind TOCTOU; Rust front-end opts out
# because it has no port-report-back channel.
# Defer port allocation to the child's bind() to avoid TOCTOU, except
# for Rust front-end and Ray DP, which can't see the post-bind rebind
# (CLI-arg subprocess / pickled-into-actor snapshot respectively) and
# so pre-allocate driver-side -- reintroducing the original race only
# there.
is_ray_dp = parallel_config.data_parallel_backend == "ray"
addresses = get_engine_zmq_addresses(
vllm_config,
num_api_servers,
defer_api_server_ports=not rust_frontend_path,
defer_api_server_ports=not (rust_frontend_path or is_ray_dp),
)
with launch_core_engines(
@@ -348,11 +351,15 @@ def run_multi_api_server(args: argparse.Namespace):
tensor_queue=tensor_queue,
)
# Forward each child's bound endpoints to the engine handshake
# (runs on ``with`` exit).
actual_inputs, actual_outputs = api_server_manager.gather_actual_addresses()
addresses.inputs = actual_inputs
addresses.outputs = actual_outputs
if not is_ray_dp:
# Forward each child's bound endpoints to the engine handshake
# (runs on ``with`` exit). Skipped for Ray DP, where addresses
# are pre-allocated above and Ray actors already hold them.
actual_inputs, actual_outputs = (
api_server_manager.gather_actual_addresses()
)
addresses.inputs = actual_inputs
addresses.outputs = actual_outputs
# Wait for API servers.
try:
+10 -1
View File
@@ -58,6 +58,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
XPUW4A8IntLinearKernel,
XPUwNa16LinearKernel,
)
from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
ZentorchWNA16LinearKernel,
)
from vllm.model_executor.kernels.linear.mxfp4 import (
MxFp4LinearKernel,
MxFp4LinearLayerConfig,
@@ -157,6 +160,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import (
from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
XPUFP8ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
ZentorchInt8ScaledMMLinearKernel,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
from vllm.platforms import PlatformEnum, current_platform
@@ -254,7 +260,7 @@ def _filter_kernels_by_backend(
# in priority/performance order (when available)
_POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = {
PlatformEnum.CPU: [CPUInt8ScaledMMLinearKernel],
PlatformEnum.CPU: [ZentorchInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel],
PlatformEnum.CUDA: [
CutlassInt8ScaledMMLinearKernel,
TritonInt8ScaledMMLinearKernel,
@@ -348,6 +354,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
],
PlatformEnum.CPU: [
Dynamic4bitLinearKernel,
ZentorchWNA16LinearKernel,
CPUWNA16LinearKernel,
],
}
@@ -1018,6 +1025,8 @@ __all__ = [
"RowWiseTorchFP8ScaledMMLinearKernel",
"ROCmFP8ScaledMMLinearKernel",
"TritonInt8ScaledMMLinearKernel",
"ZentorchInt8ScaledMMLinearKernel",
"ZentorchWNA16LinearKernel",
"MPLinearKernel",
"MPLinearLayerConfig",
"AllSparkLinearKernel",
@@ -36,6 +36,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
XPUW4A8IntLinearKernel,
XPUwNa16LinearKernel,
)
from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
ZentorchWNA16LinearKernel,
)
__all__ = [
"MPLinearKernel",
@@ -51,4 +54,5 @@ __all__ = [
"TritonW4A16LinearKernel",
"XPUW4A8IntLinearKernel",
"XPUwNa16LinearKernel",
"ZentorchWNA16LinearKernel",
]
@@ -0,0 +1,211 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Zentorch W4A16 GPTQ weight-only-quantized linear kernel for AMD Zen CPUs.
Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed
``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector
falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``.
"""
import torch
from vllm.logger import init_logger
from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from .cpu import CPUWNA16LinearKernel
from .MPLinearKernel import MPLinearLayerConfig
logger = init_logger(__name__)
def _import_unpack_from_int32():
"""Import compressed-tensors' ``unpack_from_int32`` across versions."""
try:
from compressed_tensors.compressors.pack_quantized.helpers import (
unpack_from_int32,
)
except ImportError:
from compressed_tensors.compressors.quantized_compressors.pack_quantized import ( # type: ignore[import-not-found] # noqa: E501
unpack_from_int32,
)
return unpack_from_int32
class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
"""W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``."""
@classmethod
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
ok, reason = super().can_implement(c)
if not ok:
return ok, reason
if not current_platform.is_zen_cpu():
return False, "ZentorchWNA16 requires an AMD Zen CPU."
if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]):
return (
False,
"torch.ops.zentorch.{zentorch_woq_repack_weight, "
"zentorch_woq_linear} are not registered.",
)
if c.has_g_idx:
return False, "ZentorchWNA16 does not support activation re-ordering."
return True, None
def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool:
"""Eligibility predicate for the zentorch W4A16 GPTQ fast path.
Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()``
with ``layer`` untouched).
"""
if (
self.w_gidx_name is not None
and getattr(layer, self.w_gidx_name, None) is not None
) or (getattr(self.config, "has_g_idx", False)):
return False
weight_packed = getattr(layer, self.w_q_name, None)
weight_scale = getattr(layer, self.w_s_name, None)
if weight_packed is None or weight_scale is None:
return False
bits = self.config.weight_type.mantissa
pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
# 4-bit -> 8 values per int32;
if pack_factor != 8:
return False
# GPTQ-only. AWQ packs along the output dim instead.
in_dim = getattr(weight_packed, "input_dim", None)
pk_dim = getattr(weight_packed, "packed_dim", None)
if in_dim is None or pk_dim is None or in_dim != pk_dim:
return False
is_ct_format = in_dim == pk_dim == 1
if not is_ct_format:
return False
if weight_packed.dim() != 2 or weight_scale.dim() != 2:
return False
# 4-bit -> 8 values per int32; in_features must be divisible by num_groups.
in_features = weight_packed.shape[1] * 8
num_groups = weight_scale.shape[1]
return num_groups > 0 and in_features % num_groups == 0
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Repack CT GPTQ weights into the zentorch WOQ layout.
Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading``
via ``super()`` when the layer doesn't satisfy
``_zentorch_woq_eligible``.
On success, ``layer._zentorch_processed_weights`` is set to ``True``
"""
if getattr(layer, "_zentorch_processed_weights", False):
return
if not self._zentorch_woq_eligible(layer):
logger.info_once(
"[zen_cpu] ZentorchWNA16 fast path not eligible for this "
"layer (AWQ pack layout, g_idx, or non-int32 storage); "
"falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)."
)
super().process_weights_after_loading(layer)
return
if (not self.config.zero_points) and (self.w_zp_name is not None):
setattr(layer, self.w_zp_name, None)
if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
setattr(layer, self.w_gidx_name, None)
weight_q = getattr(layer, self.w_q_name)
weight_s = getattr(layer, self.w_s_name)
weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s
bits = self.config.weight_type.mantissa
pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1]
in_features = weight_packed.shape[1] * pack_factor
original_shape = torch.Size([out_features, in_features])
unpack_from_int32 = _import_unpack_from_int32()
repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
weight_unpacked = unpack_from_int32(
weight_packed,
bits,
original_shape,
packed_dim=weight_q.packed_dim,
)
zp_param = (
getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None
)
needs_unsigned_offset = self.config.weight_type == scalar_types.uint4
if needs_unsigned_offset:
weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15)
repacked = repack_op(weight_unpacked.to(torch.int8).contiguous())
if zp_param is None:
zp_tc = None
else:
zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param
zp = unpack_from_int32(
zp_tensor,
bits,
(out_features, num_groups),
packed_dim=zp_param.packed_dim,
)
if needs_unsigned_offset:
zp = (zp.to(torch.int32) + 8).clamp(0, 15)
zp_tc = zp.to(torch.int8).t().contiguous()
layer._zentorch_woq_packed = repacked.t()
layer._zentorch_woq_scale = weight_scale.t().contiguous()
layer._zentorch_woq_zero_point = zp_tc
for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name):
if param_name is None:
continue
param = getattr(layer, param_name, None)
if param is None:
continue
if hasattr(param, "data"):
param.data = torch.empty(0)
else:
setattr(layer, param_name, torch.empty(0))
layer._zentorch_kind = "compressed_tensors_w4a16_gptq"
layer._zentorch_processed_weights = True
logger.info_once(
"[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ "
"(weight_type=%s, has_zp=%s)",
self.config.weight_type,
zp_tc is not None,
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
if getattr(layer, "_zentorch_processed_weights", False):
return torch.ops.zentorch.zentorch_woq_linear.default(
x,
layer._zentorch_woq_packed,
layer._zentorch_woq_scale,
layer._zentorch_woq_zero_point,
bias,
)
return super().apply_weights(layer, x, bias)
__all__ = ["ZentorchWNA16LinearKernel"]
@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
from vllm.model_executor.kernels.linear.scaled_mm.triton import (
TritonInt8ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
ZentorchInt8ScaledMMLinearKernel,
)
__all__ = [
"FP8ScaledMMLinearKernel",
@@ -58,6 +61,7 @@ __all__ = [
"RowWiseTorchFP8ScaledMMLinearKernel",
"ROCmFP8ScaledMMLinearKernel",
"TritonInt8ScaledMMLinearKernel",
"ZentorchInt8ScaledMMLinearKernel",
"Fp8BlockScaledMMLinearKernel",
"CPUFp8BlockScaledMMKernel",
]
@@ -0,0 +1,98 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Zentorch dynamic-symmetric W8A8 int8 linear kernel for AMD Zen CPUs.
Selected by ``choose_scaled_mm_linear_kernel`` ahead of the generic
oneDNN-backed ``CPUInt8ScaledMMLinearKernel``. When ``is_supported`` or
``can_implement`` rejects a layer, the selector falls through to the next
kernel in ``_POSSIBLE_INT8_KERNELS[PlatformEnum.CPU]``.
"""
import torch
from vllm.logger import init_logger
from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
from vllm.model_executor.layers.quantization.utils import replace_parameter
from vllm.platforms import current_platform
from .ScaledMMLinearKernel import (
Int8ScaledMMLinearKernel,
Int8ScaledMMLinearLayerConfig,
)
logger = init_logger(__name__)
class ZentorchInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
if not current_platform.is_cpu():
return False, "requires CPU."
if not current_platform.is_zen_cpu():
return False, "requires AMD Zen CPU."
if not has_zentorch_op(["zentorch_dynamic_qlinear"]):
return (
False,
"torch.ops.zentorch.zentorch_dynamic_qlinear is not registered.",
)
return True, None
@classmethod
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
if c.is_static_input_scheme:
return False, "requires dynamic activation quantization."
if not c.input_symmetric:
return False, "requires symmetric activation quantization."
if not c.is_channelwise:
return False, "requires per-channel weight quantization."
return True, None
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
"""Prepare weights for ``zentorch_dynamic_qlinear``.
Keeps weight in [N, K] layout (int8, contiguous) and converts the
per-channel weight scale to bf16 with shape ``(N,)``.
"""
w_q_name, w_s_name, _, _, _ = self.layer_param_names
weight = getattr(layer, w_q_name)
n = weight.shape[0]
replace_parameter(
layer,
w_q_name,
torch.nn.Parameter(weight.data.contiguous(), requires_grad=False),
)
weight_scale = getattr(layer, w_s_name)
ws = weight_scale.data
if ws.dim() == 2 and ws.shape[-1] == 1:
ws = ws.squeeze(-1)
ws = ws.to(torch.bfloat16).contiguous()
assert ws.shape == (n,), (
f"[zen_cpu] expected weight scale shape ({n},), got {tuple(ws.shape)}"
)
replace_parameter(
layer,
w_s_name,
torch.nn.Parameter(ws, requires_grad=False),
)
logger.info_once(
"[zen_cpu] Using zentorch_dynamic_qlinear for W8A8 (dynamic-symmetric)"
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
w_q_name, w_s_name, _, _, _ = self.layer_param_names
return torch.ops.zentorch.zentorch_dynamic_qlinear(
x,
getattr(layer, w_q_name),
getattr(layer, w_s_name),
bias,
zentorch_op_name="zentorch::zentorch_dynamic_qlinear",
)
@@ -0,0 +1,23 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Gates zentorch CPU linear dispatch on platform/op availability."""
from __future__ import annotations
import torch
from vllm.platforms import current_platform
__all__ = ["has_zentorch_op"]
def has_zentorch_op(op_names: list[str]) -> bool:
"""Return ``True`` when running on Zen CPU with all named ops registered."""
if not op_names:
raise ValueError("has_zentorch_op requires at least one op name")
if not current_platform.is_zen_cpu():
return False
ns = getattr(torch.ops, "zentorch", None)
if ns is None:
return False
return all(hasattr(ns, op_name) for op_name in op_names)
+144 -26
View File
@@ -5,6 +5,88 @@ import torch
from vllm.utils.torch_utils import direct_register_custom_op
def _torch_hc_prenorm_gemm(
x: torch.Tensor,
fn: torch.Tensor,
out: torch.Tensor,
sqrsum: torch.Tensor,
) -> None:
assert out.shape[0] == 1
assert sqrsum.shape[0] == 1
x_float = x.float()
out[0].copy_(x_float @ fn.t())
sqrsum[0].copy_(x_float.square().sum(dim=-1))
def _tilelang_hc_prenorm_gemm(
x: torch.Tensor,
fn: torch.Tensor,
out: torch.Tensor,
sqrsum: torch.Tensor,
hidden_size: int,
hc_mult: int,
tile_n: int = 12,
n_thr: int = 512,
n_splits: int = 1,
) -> None:
from vllm._tilelang_ops import (
hc_prenorm_gemm_block_m_tilelang,
hc_prenorm_gemm_tilelang,
)
assert out.shape[0] == n_splits
assert sqrsum.shape[0] == n_splits
assert x.shape[1] == hc_mult * hidden_size
assert x.shape[1] % n_splits == 0
assert (x.shape[1] // n_splits) % n_thr == 0
use_default_config = tile_n == 12 and n_thr == 512
if n_splits == 1 and use_default_config and x.shape[0] >= 1024:
hc_prenorm_gemm_block_m_tilelang(
x,
fn,
out,
sqrsum,
hidden_size,
hc_mult,
fn.shape[0],
n_thr,
tile_n,
2,
)
return
if (
n_splits == 1
and use_default_config
and x.shape[0] < 128
and x.shape[1] % 1024 == 0
):
hc_prenorm_gemm_tilelang(
x,
fn,
out,
sqrsum,
hidden_size,
hc_mult,
fn.shape[0],
1024,
4,
n_splits,
)
return
hc_prenorm_gemm_tilelang(
x,
fn,
out,
sqrsum,
hidden_size,
hc_mult,
fn.shape[0],
n_thr,
tile_n,
n_splits,
)
def mhc_pre_tilelang(
residual: torch.Tensor,
fn: torch.Tensor,
@@ -80,10 +162,16 @@ def mhc_pre_tilelang(
residual_flat = residual.view(-1, hc_mult, hidden_size)
num_tokens = residual_flat.shape[0]
# these numbers are from deepgemm kernel impl
block_k = 64
block_m = 64
n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
from vllm.utils.deep_gemm import is_deep_gemm_supported
use_deep_gemm = is_deep_gemm_supported()
if use_deep_gemm:
# these numbers are from deepgemm kernel impl
block_k = 64
block_m = 64
n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
else:
n_splits = 1
post_mix = torch.empty(
num_tokens, hc_mult, dtype=torch.float32, device=residual.device
@@ -102,13 +190,24 @@ def mhc_pre_tilelang(
n_splits, num_tokens, dtype=torch.float32, device=residual.device
)
tf32_hc_prenorm_gemm(
residual_flat.view(num_tokens, hc_mult * hidden_size),
fn,
gemm_out_mul,
gemm_out_sqrsum,
n_splits,
)
residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size)
if use_deep_gemm:
tf32_hc_prenorm_gemm(
residual_2d,
fn,
gemm_out_mul,
gemm_out_sqrsum,
n_splits,
)
else:
_tilelang_hc_prenorm_gemm(
residual_2d,
fn,
gemm_out_mul,
gemm_out_sqrsum,
hidden_size,
hc_mult,
)
if norm_weight is None:
mhc_pre_big_fuse_tilelang(
@@ -304,16 +403,24 @@ def mhc_fused_post_pre_tilelang(
post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
fma_token_threshold = 16
if num_tokens <= fma_token_threshold:
from vllm.utils.deep_gemm import is_deep_gemm_supported
use_deep_gemm = is_deep_gemm_supported()
use_small_fma = num_tokens <= 16
if use_small_fma:
# TODO(gnovack): investigate autotuning these heuristics
tile_n = 2 if num_tokens < 8 else 3
n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
else:
# these number are from deepgemm kernel impl
block_k = 64
block_m = 64
n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
if use_deep_gemm:
# these number are from deepgemm kernel impl
block_k = 64
block_m = 64
n_splits = compute_num_split(
block_k, hc_hidden_size, cdiv(num_tokens, block_m)
)
else:
n_splits = 1
gemm_out_mul = torch.empty(
n_splits,
@@ -348,7 +455,7 @@ def mhc_fused_post_pre_tilelang(
device=residual.device,
)
if num_tokens <= fma_token_threshold:
if use_small_fma:
mhc_fused_tilelang(
comb_res_mix_flat,
residual_flat,
@@ -375,15 +482,26 @@ def mhc_fused_post_pre_tilelang(
residual.shape[-1],
)
from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size)
if use_deep_gemm:
from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
tf32_hc_prenorm_gemm(
residual_cur.view(num_tokens, hc_mult * hidden_size),
fn,
gemm_out_mul,
gemm_out_sqrsum,
n_splits,
)
tf32_hc_prenorm_gemm(
residual_cur_2d,
fn,
gemm_out_mul,
gemm_out_sqrsum,
n_splits,
)
else:
_tilelang_hc_prenorm_gemm(
residual_cur_2d,
fn,
gemm_out_mul,
gemm_out_sqrsum,
hidden_size,
hc_mult,
)
if norm_weight is None:
mhc_pre_big_fuse_tilelang(
@@ -99,9 +99,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
) -> bool:
return True
def supports_chunking(self) -> bool:
return False
def supports_expert_map(self) -> bool:
return False
@@ -140,5 +137,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
intermediate_size=self.intermediate_size_per_partition,
local_expert_offset=self.ep_rank * self.local_num_experts,
local_num_experts=self.local_num_experts,
routed_scaling_factor=routed_scaling_factor,
routing_method_type=self.routing_method_type,
)
@@ -88,9 +88,6 @@ class TrtLlmFp8ExpertsBase:
or moe_parallel_config.use_ag_rs_all2all_kernels
) and not moe_parallel_config.enable_eplb
def supports_chunking(self) -> bool:
return False
def supports_expert_map(self) -> bool:
return False
@@ -113,9 +113,6 @@ class TrtLlmMxfp4ExpertsBase:
def activation_format() -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
def supports_chunking(self) -> bool:
return False
def supports_expert_map(self) -> bool:
return False
@@ -157,8 +157,27 @@ class TrtLlmNvFp4ExpertsBase:
def activation_format() -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
def supports_chunking(self) -> bool:
return False
def _get_chunk_size(self) -> int:
MAX_GRID_Y = 65535
MAX_TILE_TOKENS_DIM = 128
def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int:
"""Calculates the max number of supported tokens, so the CUDA grid.Y limit
won't be reached.
Based on getMaxNumCtasInBatchDim function in flashinfer's TRTLLM MoE runner:
https://github.com/flashinfer-ai/flashinfer/blob/719ee23fd82cb220d51ad118ca60198718f6c9d1/include/flashinfer/trtllm/fused_moe/runner.h#L97
Which given numTokens, topK, numExperts, tileTokensDim calculates maxNumCtas
which is used as the CUDA grid.Y dimension, which we want to
be <= MAX_GRID_Y. Solving for numTokens gives the formula below.
"""
return (
num_experts + (MAX_GRID_Y - num_experts + 1) * MAX_TILE_TOKENS_DIM - 1
) // top_k
# Using 305k or more causes IMA error in the kernel, so limit to 300k.
return min(
300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts)
)
def supports_expert_map(self) -> bool:
return False
@@ -199,7 +218,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
return TopKWeightAndReduceNoOP()
def apply(
def _invoke_kernel(
self,
output: torch.Tensor,
hidden_states: torch.Tensor,
@@ -209,18 +228,10 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
topk_ids: torch.Tensor,
activation: MoEActivation,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
a2_scale: torch.Tensor | None,
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
a1q_scale: torch.Tensor,
):
import flashinfer
assert self._supports_activation(activation)
assert a1q_scale is not None
assert self.quant_config.w1_scale is not None
assert self.quant_config.w2_scale is not None
@@ -262,6 +273,57 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
output=output,
)
def apply(
self,
output: torch.Tensor,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
activation: MoEActivation,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
a2_scale: torch.Tensor | None,
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
):
assert self._supports_activation(activation)
assert a1q_scale is not None
M = hidden_states.shape[0]
chunk_size = self._get_chunk_size()
if chunk_size >= M:
self._invoke_kernel(
output,
hidden_states,
w1,
w2,
topk_weights,
topk_ids,
activation,
global_num_experts,
a1q_scale,
)
else:
for start in range(0, M, chunk_size):
end = min(start + chunk_size, M)
self._invoke_kernel(
output[start:end],
hidden_states[start:end],
w1,
w2,
topk_weights[start:end],
topk_ids[start:end],
activation,
global_num_experts,
a1q_scale[start:end],
)
class TrtLlmNvFp4ExpertsMonolithic(
TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic
+126 -20
View File
@@ -3,8 +3,12 @@
import torch
# this import will also register the custom ops
# import vllm.model_executor.kernels.mhc # noqa: F401
import vllm.model_executor.kernels.mhc as mhc_kernels
from vllm.model_executor.custom_op import CustomOp
from vllm.utils.import_utils import has_tilelang
HAS_TILELANG = has_tilelang()
# --8<-- [start:mhc_pre]
@@ -85,6 +89,52 @@ class MHCPreOp(CustomOp):
# sinkhorn_repeat,
# )
# else:
if HAS_TILELANG:
return torch.ops.vllm.mhc_pre_tilelang(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
n_splits,
norm_weight,
norm_eps,
)
else:
return self.forward_native(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
n_splits,
norm_weight,
norm_eps,
)
def forward_native(
self,
residual: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
n_splits: int = 1,
norm_weight: torch.Tensor | None = None,
norm_eps: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return mhc_kernels.mhc_pre_torch(
residual,
fn,
@@ -97,9 +147,6 @@ class MHCPreOp(CustomOp):
sinkhorn_repeat,
)
def forward_native(self, *args, **kwargs):
raise NotImplementedError("Native implementation of mhc_pre is not available")
# --8<-- [start:mhc_post]
@CustomOp.register("mhc_post")
@@ -147,6 +194,20 @@ class MHCPostOp(CustomOp):
# comb_res_mix,
# )
# else:
if HAS_TILELANG:
return torch.ops.vllm.mhc_post_tilelang(
x, residual, post_layer_mix, comb_res_mix
)
else:
return self.forward_native(x, residual, post_layer_mix, comb_res_mix)
def forward_native(
self,
x: torch.Tensor,
residual: torch.Tensor,
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
) -> torch.Tensor:
return mhc_kernels.mhc_post_torch(
x,
residual,
@@ -154,9 +215,6 @@ class MHCPostOp(CustomOp):
comb_res_mix,
)
def forward_native(self, *args, **kwargs):
raise NotImplementedError("Native implementation of mhc_post is not available")
# --8<-- [start:hc_head]
@CustomOp.register("hc_head")
@@ -220,17 +278,32 @@ class HCHeadOp(CustomOp):
out = torch.empty(
num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
)
torch.ops.vllm.hc_head_triton(
hs_flat,
hc_fn,
hc_scale,
hc_base,
out,
hidden_size,
rms_norm_eps,
hc_eps,
hc_mult,
)
if HAS_TILELANG:
torch.ops.vllm.hc_head_fused_kernel_tilelang(
hs_flat,
hc_fn,
hc_scale,
hc_base,
out,
hidden_size,
rms_norm_eps,
hc_eps,
hc_mult,
)
else:
torch.ops.vllm.hc_head_triton(
hs_flat,
hc_fn,
hc_scale,
hc_base,
out,
hidden_size,
rms_norm_eps,
hc_eps,
hc_mult,
)
return out.view(*outer_shape, hidden_size)
def forward_native(self, *args, **kwargs):
@@ -290,9 +363,42 @@ class MHCFusedPostPreOp(CustomOp):
norm_eps,
)
def forward_hip(self, *args, **kwargs):
raise NotImplementedError(
"Hip implementation of mhc_fused_post_pre is not available"
def forward_hip(
self,
x: torch.Tensor,
residual: torch.Tensor,
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_eps: float,
hc_pre_eps: float,
hc_sinkhorn_eps: float,
hc_post_mult_value: float,
sinkhorn_repeat: int,
n_splits: int = 1,
tile_n: int = 1,
norm_weight: torch.Tensor | None = None,
norm_eps: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
return torch.ops.vllm.mhc_fused_post_pre_tilelang(
x,
residual,
post_layer_mix,
comb_res_mix,
fn,
hc_scale,
hc_base,
rms_eps,
hc_pre_eps,
hc_sinkhorn_eps,
hc_post_mult_value,
sinkhorn_repeat,
n_splits,
tile_n,
norm_weight,
norm_eps,
)
def forward_native(self, *args, **kwargs):
@@ -1090,7 +1090,8 @@ def runai_safetensors_weights_iterator(
mininterval=2,
)
yield from tensor_iter
for name, tensor in tensor_iter:
yield name, tensor.clone()
def _init_fastsafetensors_loader(
+253
View File
@@ -0,0 +1,253 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
from torch import nn
from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from .qwen3_moe import (
Qwen3MoeAttention,
Qwen3MoeDecoderLayer,
Qwen3MoeForCausalLM,
Qwen3MoeMLP,
Qwen3MoeModel,
Qwen3MoeSparseMoeBlock,
)
from .utils import PPMissingLayer, extract_layer_index, maybe_prefix
class MellumAttention(Qwen3MoeAttention):
"""
Differences from `Qwen3MoeAttention`:
- Supports `per_layer_sliding_window` for `Attention`.
"""
def __init__(
self,
hidden_size: int,
num_heads: int,
num_kv_heads: int,
rope_parameters: dict[str, Any],
max_position_embeddings: int = 8192,
head_dim: int | None = None,
rms_norm_eps: float = 1e-06,
qkv_bias: bool = False,
cache_config: Any | None = None,
quant_config: Any | None = None,
prefix: str = "",
dual_chunk_attention_config: dict[str, Any] | None = None,
per_layer_sliding_window: int | None = None,
) -> None:
nn.Module.__init__(self)
self.hidden_size = hidden_size
tp_size = get_tensor_model_parallel_world_size()
self.total_num_heads = num_heads
assert self.total_num_heads % tp_size == 0
self.num_heads = self.total_num_heads // tp_size
self.total_num_kv_heads = num_kv_heads
if self.total_num_kv_heads >= tp_size:
assert self.total_num_kv_heads % tp_size == 0
else:
assert tp_size % self.total_num_kv_heads == 0
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
self.head_dim = head_dim or (hidden_size // self.total_num_heads)
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.scaling = self.head_dim**-0.5
self.max_position_embeddings = max_position_embeddings
self.dual_chunk_attention_config = dual_chunk_attention_config
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=qkv_bias,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
)
self.rotary_emb = get_rope(
self.head_dim,
max_position=max_position_embeddings,
rope_parameters=rope_parameters,
dual_chunk_attention_config=dual_chunk_attention_config,
)
self.attn = Attention(
self.num_heads,
self.head_dim,
self.scaling,
num_kv_heads=self.num_kv_heads,
cache_config=cache_config,
quant_config=quant_config,
per_layer_sliding_window=per_layer_sliding_window,
prefix=f"{prefix}.attn",
**(
{
"layer_idx": extract_layer_index(prefix),
"dual_chunk_attention_config": dual_chunk_attention_config,
}
if dual_chunk_attention_config
else {}
),
)
self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
class MellumDecoderLayer(Qwen3MoeDecoderLayer):
"""
Differences from `Qwen3MoeDecoderLayer`:
- Supports interleaved SWA and per-layer RoPE scaling.
"""
def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
nn.Module.__init__(self)
config = vllm_config.model_config.hf_text_config
cache_config = vllm_config.cache_config
quant_config = vllm_config.quant_config
self.hidden_size = config.hidden_size
max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
dual_chunk_attention_config = getattr(
config, "dual_chunk_attention_config", None
)
layer_idx = extract_layer_index(prefix)
layer_type = config.layer_types[layer_idx]
if layer_type == "sliding_attention":
sliding_window = getattr(config, "sliding_window", None)
else:
sliding_window = None
rope_parameters = config.rope_parameters[layer_type]
self.self_attn = MellumAttention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
rope_parameters=rope_parameters,
max_position_embeddings=max_position_embeddings,
rms_norm_eps=config.rms_norm_eps,
qkv_bias=getattr(config, "attention_bias", False),
head_dim=getattr(config, "head_dim", None),
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
dual_chunk_attention_config=dual_chunk_attention_config,
per_layer_sliding_window=sliding_window,
)
if config.mlp_layer_types[layer_idx] == "sparse":
self.mlp = Qwen3MoeSparseMoeBlock(
vllm_config=vllm_config, prefix=f"{prefix}.mlp"
)
else:
self.mlp = Qwen3MoeMLP(
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
hidden_act=config.hidden_act,
quant_config=quant_config,
prefix=f"{prefix}.mlp",
)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
@support_torch_compile
class MellumModel(Qwen3MoeModel):
"""
Differences from `Qwen3MoeModel`:
- Uses `MellumDecoderLayer`.
"""
def __init__(
self,
*,
vllm_config: VllmConfig,
prefix: str = "",
):
super().__init__(
vllm_config=vllm_config,
prefix=prefix,
decoder_layer_type=MellumDecoderLayer,
)
class MellumForCausalLM(Qwen3MoeForCausalLM):
"""
Differences from `Qwen3MoeForCausalLM`:
- Uses `MellumModel`.
"""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
nn.Module.__init__(self)
config = vllm_config.model_config.hf_text_config
quant_config = vllm_config.quant_config
self.config = config
self.quant_config = quant_config
if "dense" in getattr(config, "mlp_layer_types", []):
self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
self.model = MellumModel(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
self.lm_head = ParallelLMHead(
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=maybe_prefix(prefix, "lm_head"),
)
if self.config.tie_word_embeddings:
self.lm_head.weight = self.model.embed_tokens.weight
self.logits_processor = LogitsProcessor(config.vocab_size)
self.make_empty_intermediate_tensors = (
self.model.make_empty_intermediate_tensors
)
self.expert_weights = []
self.moe_layers = []
example_layer = None
for layer in self.model.layers:
if isinstance(layer, PPMissingLayer):
continue
assert isinstance(layer, Qwen3MoeDecoderLayer)
if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
example_layer = layer.mlp
self.moe_layers.append(layer.mlp.experts)
if example_layer is None:
raise RuntimeError("No MoE layer found in the model.layers.")
self.num_moe_layers = len(self.moe_layers)
self.num_expert_groups = 1
self.num_shared_experts = 0
self.num_logical_experts = example_layer.n_logical_experts
self.num_physical_experts = example_layer.n_physical_experts
self.num_local_physical_experts = example_layer.n_local_physical_experts
self.num_routed_experts = example_layer.n_routed_experts
self.num_redundant_experts = example_layer.n_redundant_experts
+1
View File
@@ -160,6 +160,7 @@ _TEXT_GENERATION_MODELS = {
"LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"),
"MambaForCausalLM": ("mamba", "MambaForCausalLM"),
"Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),
"MellumForCausalLM": ("mellum", "MellumForCausalLM"),
"MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
"MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
"MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
+11 -7
View File
@@ -54,6 +54,7 @@ from vllm.models.deepseek_v4.attention import (
)
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.utils.import_utils import has_tilelang
class DeepseekV4MLP(nn.Module):
@@ -473,6 +474,7 @@ class DeepseekV4DecoderLayer(nn.Module):
self.mhc_pre = MHCPreOp()
self.mhc_post = MHCPostOp()
self.mhc_fused_post_pre = MHCFusedPostPreOp()
self.has_tilelang = has_tilelang()
def hc_pre(
self,
@@ -503,7 +505,7 @@ class DeepseekV4DecoderLayer(nn.Module):
):
return self.mhc_post(x, residual, post, comb)
def _forward_cuda(
def _forward_fused_post_pre(
self,
x: torch.Tensor,
positions: torch.Tensor,
@@ -555,7 +557,7 @@ class DeepseekV4DecoderLayer(nn.Module):
x = self.ffn(x, input_ids)
return x, residual, post_mix, res_mix
def _forward_rocm(
def _forward_unfused_post_pre(
self,
x: torch.Tensor,
positions: torch.Tensor,
@@ -594,12 +596,13 @@ class DeepseekV4DecoderLayer(nn.Module):
) -> tuple[
torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
]:
if current_platform.is_rocm():
return self._forward_rocm(
if not self.has_tilelang:
return self._forward_unfused_post_pre(
x, positions, input_ids, post_mix, res_mix, residual
)
return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)
return self._forward_fused_post_pre(
x, positions, input_ids, post_mix, res_mix, residual
)
@support_torch_compile
@@ -682,6 +685,7 @@ class DeepseekV4Model(nn.Module):
requires_grad=False,
)
self.hc_head_op = HCHeadOp()
self.has_tilelang = has_tilelang()
# Pre-hc_head residual stream buffer for the MTP draft. Stable
# address (outside the cudagraph pool) so the copy_ in forward()
# refreshes it correctly across captured shapes.
@@ -748,7 +752,7 @@ class DeepseekV4Model(nn.Module):
res_mix,
residual,
)
if layer is not None and current_platform.is_cuda():
if layer is not None and self.has_tilelang:
hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
if not get_pp_group().is_last_rank:
+3 -1
View File
@@ -39,6 +39,7 @@ from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weigh
from vllm.model_executor.models.utils import maybe_prefix
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.utils.import_utils import has_tilelang
from .model import DeepseekV4DecoderLayer
@@ -118,6 +119,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
)
self.hc_head_op = HCHeadOp()
self.has_tilelang = has_tilelang()
def forward(
self,
@@ -144,7 +146,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
hidden_states, residual, post_mix, res_mix = self.mtp_block(
positions=positions, x=hidden_states, input_ids=None
)
if current_platform.is_cuda():
if self.has_tilelang:
hidden_states = self.mtp_block.hc_post(
hidden_states, residual, post_mix, res_mix
)
@@ -320,11 +320,11 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
bits = _recast_val(scale_raw, Uint32)
ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
y0 = cute.arch.fmin(
y0 = cutlass.min(
cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
Float32(self.fp8_max),
)
y1 = cute.arch.fmin(
y1 = cutlass.min(
cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
Float32(self.fp8_max),
)
@@ -978,11 +978,11 @@ class SparseAttnNormRopeStoreKernel:
bits = _recast_val(scale_raw, Uint32)
ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
y0 = cute.arch.fmin(
y0 = cutlass.min(
cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
Float32(self.fp8_max),
)
y1 = cute.arch.fmin(
y1 = cutlass.min(
cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
Float32(self.fp8_max),
)
+3 -20
View File
@@ -14,7 +14,6 @@ from vllm.logger import init_logger
from vllm.utils.cpu_resource_utils import (
DEVICE_CONTROL_ENV_VAR,
get_memory_node_info,
get_visible_memory_node,
)
from vllm.utils.mem_constants import GiB_bytes
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -136,13 +135,9 @@ class CpuPlatform(Platform):
scheduler_config.async_scheduling = False
parallel_config = vllm_config.parallel_config
if (
os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING", "1") == "1"
and parallel_config.distributed_executor_backend == "uni"
):
# OMP requires the MP executor to function correctly, UniProc
# is not supported as it is not possible to set the OMP
# environment correctly
# OMP requires the MP executor to function correctly, UniProc is not
# supported as it is not possible to set the OMP environment correctly
if parallel_config.distributed_executor_backend == "uni":
parallel_config.distributed_executor_backend = "mp"
if parallel_config.worker_cls == "auto":
@@ -486,15 +481,3 @@ class CpuPlatform(Platform):
slot_mapping,
isa,
)
@classmethod
def get_current_memory_usage(
cls, device: torch.types.Device | None = None
) -> float:
allowed_mem_node_list = get_visible_memory_node()
mem_status_list = [get_memory_node_info(i) for i in allowed_mem_node_list]
memory_usage = 0
for s in mem_status_list:
memory_usage += s.total_memory - s.available_memory
return memory_usage
+9
View File
@@ -592,6 +592,15 @@ class CudaPlatformBase(Platform):
default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
)
@classmethod
def is_arch_support_pdl(cls) -> bool:
try:
device = torch.cuda.current_device()
major, _ = torch.cuda.get_device_capability(device)
except Exception:
return False
return major >= 9
# NVML utils
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
+7
View File
@@ -1016,6 +1016,13 @@ class Platform:
# Native always used by default. Platforms can override this behavior.
return IrOpPriorityConfig.with_default(["native"])
@classmethod
def is_arch_support_pdl(cls) -> bool:
"""
Does the current platform support PDL (Programmatic Dependent Launch)?
"""
return False
class UnspecifiedPlatform(Platform):
_enum = PlatformEnum.UNSPECIFIED
+4 -2
View File
@@ -101,6 +101,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
fireredlid="FireRedLIDConfig",
funaudiochat="FunAudioChatConfig",
granite4_vision="Granite4VisionConfig",
hyperclovax="HyperCLOVAXConfig",
hyperclovax_vlm="HCXVisionConfig",
hunyuan_vl="HunYuanVLConfig",
hy_v3="HYV3Config",
@@ -114,6 +115,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
jais="JAISConfig",
mlp_speculator="MLPSpeculatorConfig",
medusa="MedusaConfig",
mellum="MellumConfig",
midashenglm="MiDashengLMConfig",
moondream3="Moondream3Config",
eagle="EAGLEConfig",
@@ -429,9 +431,9 @@ def patch_legacy_rope_type(rope_parameters: dict[str, Any] | None) -> None:
if "rope_type" not in rope_parameters and "type" in rope_parameters:
rope_parameters["rope_type"] = rope_parameters["type"]
logger.info("Replacing legacy 'type' key with 'rope_type'")
# Case 3: No rope_type field at all - cannot determine RoPE type, raise error
# Case 3: No rope_type field present - nothing to patch
if "rope_type" not in rope_parameters:
raise ValueError("rope_parameters should have a 'rope_type' key")
return
# Patch legacy rope_type values with warning
if rope_parameters["rope_type"] == "su":
rope_parameters["rope_type"] = "longrope"
@@ -49,6 +49,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
"LagunaConfig": "vllm.transformers_utils.configs.laguna",
"Lfm2MoeConfig": "vllm.transformers_utils.configs.lfm2_moe",
"MedusaConfig": "vllm.transformers_utils.configs.medusa",
"MellumConfig": "vllm.transformers_utils.configs.mellum",
"MiDashengLMConfig": "vllm.transformers_utils.configs.midashenglm",
"MLPSpeculatorConfig": "vllm.transformers_utils.configs.mlp_speculator",
"Moondream3Config": "vllm.transformers_utils.configs.moondream3",
@@ -117,6 +118,7 @@ __all__ = [
"LagunaConfig",
"Lfm2MoeConfig",
"MedusaConfig",
"MellumConfig",
"MiDashengLMConfig",
"MLPSpeculatorConfig",
"Moondream3Config",
@@ -0,0 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from transformers import Qwen3MoeConfig
class MellumConfig(Qwen3MoeConfig):
model_type = "mellum"
-12
View File
@@ -3,7 +3,6 @@
import os
import types
from importlib.metadata import version
from importlib.util import find_spec
from vllm.logger import init_logger
@@ -49,17 +48,6 @@ if HAS_TRITON:
len(active_drivers),
)
HAS_TRITON = False
# Check Triton CPU
if "cpu" in version("vllm"):
if "cpu" in backends:
HAS_TRITON = True
else:
logger.warning(
"Triton is installed, but doesn't include CPU backend. "
"Disabling Triton."
)
HAS_TRITON = False
except ImportError:
# This can occur if Triton is partially installed or triton.backends
# is missing.
+1
View File
@@ -430,6 +430,7 @@ def has_triton_kernels() -> bool:
return is_available
@cache
def has_tilelang() -> bool:
"""Whether the optional `tilelang` package is available."""
return _has_module("tilelang")
+1 -3
View File
@@ -50,10 +50,8 @@ def is_pin_memory_available() -> bool:
def is_uva_available() -> bool:
"""Check if Unified Virtual Addressing (UVA) is available."""
# UVA requires pinned memory.
from vllm.platforms import current_platform
# TODO: Add more requirements for UVA if needed.
return is_pin_memory_available() or current_platform.is_cpu()
return is_pin_memory_available()
@cache
+1 -1
View File
@@ -241,7 +241,7 @@ class APIServerProcessManager:
def gather_actual_addresses(
self,
timeout: float = 60.0,
timeout: float = envs.VLLM_ENGINE_READY_TIMEOUT_S,
) -> tuple[list[str], list[str]]:
"""Return (inputs, outputs) reported by each child, indexed by
``client_index``. Raises ``RuntimeError`` on timeout or premature
View File
-16
View File
@@ -1,16 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
import torch
from vllm.utils.platform_utils import is_uva_available
class UvaBuffer:
def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
if not is_uva_available():
raise RuntimeError("UVA is not available")
self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
self.np = self.cpu.numpy()
self.uva = self.cpu
-16
View File
@@ -1,16 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.logger import init_logger
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
logger = init_logger(__name__)
class CPUModelRunner(GPUModelRunner):
# TBD: Whether need to move this to Worker?
def warming_up_model(self) -> None:
logger.info("Warming up model for the compilation...")
# Only generate graph for the generic shape
self.profile_run()
logger.info("Warming up done.")
-62
View File
@@ -1,62 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# isort: skip_file
# ruff: noqa: E402
# mypy: disable-error-code="misc, assignment"
from typing import Any
# Patch torch APIs
import torch
def noop(*args: Any, **kwargs: Any) -> None:
pass
class _EventPlaceholder:
def __init__(self, *args, **kwargs) -> None:
self.record = noop
self.synchronize = noop
class _StreamPlaceholder:
def __init__(self, *args, **kwargs) -> None:
self.wait_stream = noop
def __enter__(self, *args, **kwargs):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
pass
torch.Event = _EventPlaceholder
torch.cuda.Event = _EventPlaceholder
torch.cuda.Stream = _StreamPlaceholder
torch.cuda.set_stream = noop
torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
torch.accelerator.synchronize = noop
torch.accelerator.empty_cache = noop
# Patch vLLM torch utils
import vllm.utils.torch_utils as torch_utils
def async_tensor_h2d(
data: list,
dtype: torch.dtype,
device: str | torch.device,
pin_memory: bool = False,
) -> torch.Tensor:
return torch.tensor(data, dtype=dtype, device="cpu")
torch_utils.async_tensor_h2d = async_tensor_h2d
# Patch model runner APIs
import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils
import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils
gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer
+3 -16
View File
@@ -1,9 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Must be imported firstly
import vllm.v1.worker.cpu.shm # noqa # isort: skip
import math
import os
import sys
@@ -105,8 +101,6 @@ class CPUWorker(Worker):
)
def init_device(self):
self.device = torch.device("cpu")
# Check whether critical libraries are loaded
def check_preloaded_libs(name: str):
ld_preload_list = os.environ.get("LD_PRELOAD", "")
@@ -147,16 +141,9 @@ class CPUWorker(Worker):
set_random_seed(self.model_config.seed)
# Construct the model runner
if self.use_v2_model_runner:
from vllm.v1.worker.cpu.model_runner import (
CPUModelRunner as CPUModelRunnerV2,
)
self.model_runner: CPUModelRunner = CPUModelRunnerV2( # type: ignore
self.vllm_config, self.device
)
else:
self.model_runner = CPUModelRunner(self.vllm_config, torch.device("cpu"))
self.model_runner: CPUModelRunner = CPUModelRunner(
self.vllm_config, torch.device("cpu")
)
def sleep(self, level: int = 1) -> None:
logger.warning("sleep mode is not supported on CPU, ignore it.")