forked from Karylab-cklius/vllm
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0decac0d96 | ||
|
|
fd56c57bde | ||
|
|
7285178622 | ||
|
|
27509c8dde | ||
|
|
b284862ea9 | ||
|
|
932dfd5276 | ||
|
|
682ffebfef | ||
|
|
1be7a57a18 |
Executable
+39
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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. | | ✅︎ |
|
||||
|
||||
@@ -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"],
|
||||
|
||||
@@ -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."
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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"),
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user