forked from Karylab-cklius/vllm
Compare commits
16
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4b6ff605fc | ||
|
|
c5eacd261d | ||
|
|
549b8242d6 | ||
|
|
3fde6cc382 | ||
|
|
f0b536640c | ||
|
|
9e173a5cdc | ||
|
|
09a977e897 | ||
|
|
1b2b0043d8 | ||
|
|
a87a2676f2 | ||
|
|
f54775736a | ||
|
|
934e6b2a46 | ||
|
|
e9966f1c52 | ||
|
|
12e25a1d3f | ||
|
|
3bebb3d705 | ||
|
|
36547719b0 | ||
|
|
d4d3e90c3e |
@@ -469,8 +469,6 @@ th {
|
||||
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
|
||||
| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
|
||||
| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
|
||||
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for vLLM dashboard functionality."""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from vllm.entrypoints.serve.dashboard.api_router import (
|
||||
_get_dashboard_html,
|
||||
_get_explicit_cli_args,
|
||||
_get_vllm_env_vars,
|
||||
attach_router,
|
||||
router,
|
||||
)
|
||||
from vllm.version import __version__ as VLLM_VERSION
|
||||
|
||||
|
||||
class TestDashboardHelpers:
|
||||
"""Test helper functions in dashboard api_router."""
|
||||
|
||||
def test_get_vllm_env_vars(self):
|
||||
"""Test that _get_vllm_env_vars returns VLLM environment variables."""
|
||||
vllm_envs, explicit_envs = _get_vllm_env_vars()
|
||||
|
||||
# Should return a dict of VLLM_* variables
|
||||
assert isinstance(vllm_envs, dict)
|
||||
assert isinstance(explicit_envs, set)
|
||||
|
||||
# All keys should start with VLLM_
|
||||
for key in vllm_envs:
|
||||
assert key.startswith("VLLM_"), f"Key {key} doesn't start with VLLM_"
|
||||
|
||||
# Should not include keys with "KEY" (secrets)
|
||||
for key in vllm_envs:
|
||||
assert "KEY" not in key, f"Key {key} contains KEY (potential secret)"
|
||||
|
||||
def test_get_vllm_env_vars_explicit_detection(self):
|
||||
"""Test that explicitly set env vars are detected."""
|
||||
test_var = "VLLM_TEST_DASHBOARD_VAR"
|
||||
try:
|
||||
os.environ[test_var] = "test_value"
|
||||
# Need to reload envs module to pick up the new var
|
||||
# For this test, we just verify the mechanism works
|
||||
_, explicit_envs = _get_vllm_env_vars()
|
||||
# The test var won't be in vllm.envs, but the mechanism is tested
|
||||
assert isinstance(explicit_envs, set)
|
||||
finally:
|
||||
os.environ.pop(test_var, None)
|
||||
|
||||
def test_get_explicit_cli_args_none(self):
|
||||
"""Test _get_explicit_cli_args with None args."""
|
||||
result = _get_explicit_cli_args(None)
|
||||
assert result == set()
|
||||
|
||||
def test_get_explicit_cli_args_with_args(self):
|
||||
"""Test _get_explicit_cli_args with mock args."""
|
||||
mock_args = MagicMock()
|
||||
mock_args.model = "test-model"
|
||||
mock_args.served_model_name = "test-served-model"
|
||||
mock_args.dtype = "float16" # Non-default value
|
||||
|
||||
result = _get_explicit_cli_args(mock_args)
|
||||
|
||||
# model should always be explicit if set
|
||||
assert "model" in result
|
||||
assert "served_model_name" in result
|
||||
|
||||
def test_get_dashboard_html(self):
|
||||
"""Test that dashboard HTML is loaded correctly."""
|
||||
html = _get_dashboard_html()
|
||||
|
||||
assert isinstance(html, str)
|
||||
assert len(html) > 0
|
||||
assert "<!DOCTYPE html>" in html
|
||||
assert "vLLM Dashboard" in html
|
||||
assert "</html>" in html
|
||||
|
||||
|
||||
class TestDashboardRouter:
|
||||
"""Test dashboard router endpoints."""
|
||||
|
||||
@pytest.fixture
|
||||
def app(self):
|
||||
"""Create a test FastAPI app with dashboard router."""
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
|
||||
# Mock app state
|
||||
app.state.openai_serving_models = None
|
||||
app.state.vllm_config = None
|
||||
app.state.args = None
|
||||
app.state.engine_client = None
|
||||
|
||||
return app
|
||||
|
||||
@pytest.fixture
|
||||
def client(self, app):
|
||||
"""Create a test client."""
|
||||
return TestClient(app)
|
||||
|
||||
def test_dashboard_index(self, client):
|
||||
"""Test GET /dashboard returns HTML page."""
|
||||
response = client.get("/dashboard")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "text/html" in response.headers["content-type"]
|
||||
assert "vLLM Dashboard" in response.text
|
||||
|
||||
def test_dashboard_api_info(self, client):
|
||||
"""Test GET /dashboard/api/info returns server info."""
|
||||
response = client.get("/dashboard/api/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert "version" in data
|
||||
assert data["version"] == VLLM_VERSION
|
||||
assert "status" in data
|
||||
assert data["status"] == "running"
|
||||
|
||||
def test_dashboard_api_info_with_models(self, app):
|
||||
"""Test /dashboard/api/info includes model info when available."""
|
||||
# Mock serving models
|
||||
mock_model = MagicMock()
|
||||
mock_model.id = "test-model"
|
||||
mock_model.root = "test-model-root"
|
||||
|
||||
mock_models_response = MagicMock()
|
||||
mock_models_response.data = [mock_model]
|
||||
|
||||
mock_serving_models = AsyncMock()
|
||||
mock_serving_models.show_available_models = AsyncMock(
|
||||
return_value=mock_models_response
|
||||
)
|
||||
|
||||
app.state.openai_serving_models = mock_serving_models
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard/api/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert "models" in data
|
||||
assert len(data["models"]) == 1
|
||||
assert data["models"][0]["id"] == "test-model"
|
||||
assert data["models"][0]["root"] == "test-model-root"
|
||||
|
||||
def test_dashboard_api_info_health_check(self, app):
|
||||
"""Test /dashboard/api/info uses engine health check."""
|
||||
# Mock healthy engine client
|
||||
mock_engine_client = AsyncMock()
|
||||
mock_engine_client.check_health = AsyncMock(return_value=None)
|
||||
mock_engine_client.is_sleeping = AsyncMock(return_value=False)
|
||||
mock_engine_client.is_paused = AsyncMock(return_value=False)
|
||||
|
||||
app.state.engine_client = mock_engine_client
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard/api/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "running"
|
||||
assert data["is_sleeping"] is False
|
||||
assert data["is_paused"] is False
|
||||
|
||||
def test_dashboard_api_info_unhealthy_engine(self, app):
|
||||
"""Test /dashboard/api/info reports unhealthy when engine is dead."""
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
|
||||
# Mock unhealthy engine client
|
||||
mock_engine_client = AsyncMock()
|
||||
mock_engine_client.check_health = AsyncMock(side_effect=EngineDeadError())
|
||||
mock_engine_client.is_sleeping = AsyncMock(return_value=False)
|
||||
mock_engine_client.is_paused = AsyncMock(return_value=False)
|
||||
|
||||
app.state.engine_client = mock_engine_client
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard/api/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "unhealthy"
|
||||
|
||||
def test_dashboard_api_info_sleeping_engine(self, app):
|
||||
"""Test /dashboard/api/info reports is_sleeping when engine is sleeping."""
|
||||
# Mock sleeping engine client
|
||||
mock_engine_client = AsyncMock()
|
||||
mock_engine_client.check_health = AsyncMock(return_value=None)
|
||||
mock_engine_client.is_sleeping = AsyncMock(return_value=True)
|
||||
mock_engine_client.is_paused = AsyncMock(return_value=False)
|
||||
|
||||
app.state.engine_client = mock_engine_client
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard/api/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "running"
|
||||
assert data["is_sleeping"] is True
|
||||
|
||||
def test_dashboard_api_info_server_load(self, app):
|
||||
"""Test /dashboard/api/info includes server_load when available."""
|
||||
app.state.server_load_metrics = 5
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard/api/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["server_load"] == 5
|
||||
|
||||
def test_dashboard_api_info_paused_engine(self, app):
|
||||
"""Test /dashboard/api/info reports is_paused when engine is paused."""
|
||||
# Mock paused engine client
|
||||
mock_engine_client = AsyncMock()
|
||||
mock_engine_client.check_health = AsyncMock(return_value=None)
|
||||
mock_engine_client.is_sleeping = AsyncMock(return_value=False)
|
||||
mock_engine_client.is_paused = AsyncMock(return_value=True)
|
||||
|
||||
app.state.engine_client = mock_engine_client
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard/api/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "running"
|
||||
assert data["is_paused"] is True
|
||||
|
||||
def test_dashboard_api_metrics(self, client):
|
||||
"""Test GET /dashboard/api/metrics returns metrics."""
|
||||
response = client.get("/dashboard/api/metrics")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Should return a dict (may be empty if no metrics available)
|
||||
assert isinstance(data, dict)
|
||||
|
||||
|
||||
class TestAttachRouter:
|
||||
"""Test attach_router function."""
|
||||
|
||||
def test_attach_router_disabled(self):
|
||||
"""Test router is not attached when dashboard is disabled."""
|
||||
app = FastAPI()
|
||||
mock_args = MagicMock()
|
||||
mock_args.enable_dashboard = False
|
||||
app.state.args = mock_args
|
||||
|
||||
attach_router(app)
|
||||
|
||||
# Dashboard routes should not be available
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard")
|
||||
assert response.status_code == 404
|
||||
|
||||
def test_attach_router_enabled(self):
|
||||
"""Test router is attached when dashboard is enabled."""
|
||||
app = FastAPI()
|
||||
mock_args = MagicMock()
|
||||
mock_args.enable_dashboard = True
|
||||
app.state.args = mock_args
|
||||
app.state.openai_serving_models = None
|
||||
app.state.vllm_config = None
|
||||
|
||||
attach_router(app)
|
||||
|
||||
# Dashboard routes should be available
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard")
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_attach_router_no_args(self):
|
||||
"""Test router is not attached when args is None."""
|
||||
app = FastAPI()
|
||||
app.state.args = None
|
||||
|
||||
attach_router(app)
|
||||
|
||||
# Dashboard routes should not be available
|
||||
client = TestClient(app)
|
||||
response = client.get("/dashboard")
|
||||
assert response.status_code == 404
|
||||
@@ -480,18 +480,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
min_transformers_version="4.56.3",
|
||||
),
|
||||
"RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),
|
||||
"SarvamMoEForCausalLM": _HfExamplesInfo(
|
||||
"sarvamai/sarvam-30b",
|
||||
trust_remote_code=True,
|
||||
max_model_len=4096,
|
||||
is_available_online=True,
|
||||
),
|
||||
"SarvamMLAForCausalLM": _HfExamplesInfo(
|
||||
"sarvamai/sarvam-105b",
|
||||
trust_remote_code=True,
|
||||
max_model_len=4096,
|
||||
is_available_online=True,
|
||||
),
|
||||
"SeedOssForCausalLM": _HfExamplesInfo(
|
||||
"ByteDance-Seed/Seed-OSS-36B-Instruct",
|
||||
trust_remote_code=True,
|
||||
|
||||
@@ -56,105 +56,6 @@ from .rebalance_execute import (
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _verify_expert_weights_after_rearrange(
|
||||
model_state: "EplbModelState",
|
||||
ep_rank: int,
|
||||
) -> None:
|
||||
"""
|
||||
Post-rearrangement diagnostic: verify expert weight consistency.
|
||||
|
||||
Checks:
|
||||
1. g1_alphas[i] == a13_scale_val * w13_weight_scale_2[i] for all experts
|
||||
2. g2_alphas[i] == a2_scale_val * w2_weight_scale_2[i] for all experts
|
||||
3. Per-weight checksums for tracking corruption across rearrangements
|
||||
"""
|
||||
torch.cuda.synchronize()
|
||||
model = model_state.model
|
||||
g1_broken = 0
|
||||
g2_broken = 0
|
||||
g1_max_diff_all = 0.0
|
||||
g2_max_diff_all = 0.0
|
||||
for layer_idx, layer in enumerate(model.moe_layers):
|
||||
g1 = getattr(layer, "g1_alphas", None)
|
||||
g2 = getattr(layer, "g2_alphas", None)
|
||||
s2_13 = getattr(layer, "w13_weight_scale_2", None)
|
||||
s2_2 = getattr(layer, "w2_weight_scale_2", None)
|
||||
a13 = getattr(layer, "w13_input_scale", None)
|
||||
a2 = getattr(layer, "w2_input_scale", None)
|
||||
|
||||
# Invariant checks
|
||||
if g1 is not None and s2_13 is not None and a13 is not None:
|
||||
a13_val = a13.float().max().item()
|
||||
expected_g1 = a13_val * s2_13.float()
|
||||
diff = (g1.float() - expected_g1).abs()
|
||||
max_diff = diff.max().item()
|
||||
g1_max_diff_all = max(g1_max_diff_all, max_diff)
|
||||
if max_diff > 1e-6:
|
||||
g1_broken += 1
|
||||
bad = (diff > 1e-6).nonzero(as_tuple=True)[0]
|
||||
logger.error(
|
||||
"EPLB INVARIANT BROKEN rank %d layer %d: "
|
||||
"g1_alphas != a13_scale * w13_scale_2, "
|
||||
"max_diff=%.6e, broken_slots=%s "
|
||||
"(g1=%s, expected=%s)",
|
||||
ep_rank, layer_idx, max_diff,
|
||||
bad[:8].tolist(),
|
||||
g1.float()[bad[:4]].tolist(),
|
||||
expected_g1[bad[:4]].tolist(),
|
||||
)
|
||||
|
||||
if g2 is not None and s2_2 is not None and a2 is not None:
|
||||
a2_val = a2.float().max().item()
|
||||
expected_g2 = a2_val * s2_2.float()
|
||||
diff = (g2.float() - expected_g2).abs()
|
||||
max_diff = diff.max().item()
|
||||
g2_max_diff_all = max(g2_max_diff_all, max_diff)
|
||||
if max_diff > 1e-6:
|
||||
g2_broken += 1
|
||||
bad = (diff > 1e-6).nonzero(as_tuple=True)[0]
|
||||
logger.error(
|
||||
"EPLB INVARIANT BROKEN rank %d layer %d: "
|
||||
"g2_alphas != a2_scale * w2_scale_2, "
|
||||
"max_diff=%.6e, broken_slots=%s",
|
||||
ep_rank, layer_idx, max_diff,
|
||||
bad[:8].tolist(),
|
||||
)
|
||||
|
||||
# Per-weight checksums (rank 0 only, sample layers)
|
||||
if ep_rank == 0 and layer_idx % 20 == 0:
|
||||
checksums = []
|
||||
for name, param in layer.named_parameters():
|
||||
if name in {"w13_input_scale", "w2_input_scale",
|
||||
"e_score_correction_bias"}:
|
||||
continue
|
||||
if (name.startswith("_shared_experts.")
|
||||
or name.startswith("_gate.")):
|
||||
continue
|
||||
cs = param.float().abs().sum().item()
|
||||
checksums.append(f"{name}={cs:.4f}")
|
||||
logger.info(
|
||||
"EPLB checksums rank %d layer %d: %s",
|
||||
ep_rank, layer_idx, ", ".join(checksums),
|
||||
)
|
||||
|
||||
num_layers = model.num_moe_layers
|
||||
if g1_broken > 0 or g2_broken > 0:
|
||||
logger.error(
|
||||
"EPLB VERIFY rank %d: %d/%d layers g1 broken, "
|
||||
"%d/%d layers g2 broken (g1_max=%.2e, g2_max=%.2e)",
|
||||
ep_rank, g1_broken, num_layers,
|
||||
g2_broken, num_layers,
|
||||
g1_max_diff_all, g2_max_diff_all,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"EPLB VERIFY rank %d: all %d layers OK "
|
||||
"(g1_max=%.2e, g2_max=%.2e)",
|
||||
ep_rank, num_layers,
|
||||
g1_max_diff_all, g2_max_diff_all,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EplbStats:
|
||||
"""
|
||||
@@ -670,22 +571,18 @@ class EplbState:
|
||||
.float()
|
||||
)
|
||||
|
||||
# Compute per-layer balancedness ratio:
|
||||
# for each layer: (mean across ranks) / (max across ranks)
|
||||
# then average across layers.
|
||||
# dim=-1 is the rank dimension.
|
||||
avg_per_layer = num_tokens_per_rank.mean(dim=-1)
|
||||
max_per_layer = num_tokens_per_rank.max(dim=-1).values
|
||||
per_layer_balance = torch.where(
|
||||
max_per_layer > 0,
|
||||
avg_per_layer / max_per_layer,
|
||||
torch.ones_like(max_per_layer),
|
||||
)
|
||||
balancedness = float(per_layer_balance.mean().item())
|
||||
# Compute balancedness ratio:
|
||||
# for each layer:
|
||||
# (mean load across ranks) / (max load across ranks)
|
||||
avg_tokens_tensor = num_tokens_per_rank.mean(dim=0).sum(dim=0)
|
||||
max_tokens_tensor = num_tokens_per_rank.max(dim=0).values.sum(dim=0)
|
||||
|
||||
# Summary stats for logging
|
||||
avg_tokens = float(avg_per_layer.sum().item())
|
||||
max_tokens = float(max_per_layer.sum().item())
|
||||
# Just to make type checker happy
|
||||
tokens_tensors: list[float] = torch.stack(
|
||||
[avg_tokens_tensor, max_tokens_tensor]
|
||||
).tolist()
|
||||
avg_tokens, max_tokens = tokens_tensors
|
||||
balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
|
||||
|
||||
if ep_group.rank() == 0:
|
||||
logger.info(
|
||||
@@ -865,9 +762,6 @@ class EplbState:
|
||||
)
|
||||
|
||||
if not is_profile:
|
||||
_verify_expert_weights_after_rearrange(
|
||||
eplb_model_state, ep_rank
|
||||
)
|
||||
if (
|
||||
eplb_model_state.physical_to_logical_map.shape[1]
|
||||
!= new_physical_to_logical_map.shape[1]
|
||||
|
||||
@@ -281,6 +281,11 @@ class FrontendArgs(BaseFrontendArgs):
|
||||
Enable offline FastAPI documentation for air-gapped environments.
|
||||
Uses vendored static assets bundled with vLLM.
|
||||
"""
|
||||
enable_dashboard: bool = False
|
||||
"""
|
||||
Enable the vLLM web dashboard at /dashboard for monitoring and testing.
|
||||
Provides server info, metrics display, and a chat interface for testing.
|
||||
"""
|
||||
use_gpu_for_pooling_score: bool = False
|
||||
"""If set, run pooling score MaxSim on GPU in the API server process.
|
||||
Can significantly improve late-interaction scoring performance.
|
||||
|
||||
@@ -129,8 +129,8 @@ def get_uvicorn_log_config(args: Namespace) -> dict | None:
|
||||
|
||||
Priority:
|
||||
1. If log_config_file is specified, use it
|
||||
2. If disable_access_log_for_endpoints is specified, create a config with
|
||||
the access log filter
|
||||
2. If disable_access_log_for_endpoints is specified or dashboard is enabled,
|
||||
create a config with the access log filter
|
||||
3. Otherwise, return None (use uvicorn defaults)
|
||||
"""
|
||||
# First, try to load from file if specified
|
||||
@@ -138,16 +138,23 @@ def get_uvicorn_log_config(args: Namespace) -> dict | None:
|
||||
if log_config is not None:
|
||||
return log_config
|
||||
|
||||
# If endpoints to filter are specified, create a config with the filter
|
||||
if args.disable_access_log_for_endpoints:
|
||||
from vllm.logging_utils import create_uvicorn_log_config
|
||||
excluded_paths: list[str] = []
|
||||
|
||||
# Parse comma-separated string into list
|
||||
# User-specified endpoints to filter
|
||||
if args.disable_access_log_for_endpoints:
|
||||
excluded_paths = [
|
||||
p.strip()
|
||||
for p in args.disable_access_log_for_endpoints.split(",")
|
||||
if p.strip()
|
||||
]
|
||||
|
||||
# Dashboard polling endpoints create noise every 5s; suppress them
|
||||
if getattr(args, "enable_dashboard", False):
|
||||
excluded_paths.extend(["/dashboard/api/info", "/dashboard/api/metrics", "/dashboard"])
|
||||
|
||||
if excluded_paths:
|
||||
from vllm.logging_utils import create_uvicorn_log_config
|
||||
|
||||
return create_uvicorn_log_config(
|
||||
excluded_paths=excluded_paths,
|
||||
log_level=args.uvicorn_log_level,
|
||||
|
||||
@@ -52,6 +52,12 @@ def register_vllm_serve_api_routers(app: FastAPI):
|
||||
|
||||
attach_tokenize_router(app)
|
||||
|
||||
from vllm.entrypoints.serve.dashboard.api_router import (
|
||||
attach_router as attach_dashboard_router,
|
||||
)
|
||||
|
||||
attach_dashboard_router(app)
|
||||
|
||||
from .instrumentator import register_instrumentator_api_routers
|
||||
|
||||
register_instrumentator_api_routers(app)
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,458 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Dashboard router for vLLM web UI.
|
||||
|
||||
Provides endpoints for:
|
||||
- /dashboard - Main dashboard HTML page
|
||||
- /dashboard/api/info - Server information JSON (config, env, load, status)
|
||||
- /dashboard/api/metrics - Metrics JSON
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import pathlib
|
||||
|
||||
import pydantic
|
||||
from fastapi import APIRouter, FastAPI, Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.version import __version__ as VLLM_VERSION
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
PydanticVllmConfig = pydantic.TypeAdapter(VllmConfig)
|
||||
|
||||
|
||||
def _get_vllm_env_vars() -> tuple[dict, set]:
|
||||
"""Get all VLLM environment variables and which ones are explicitly set."""
|
||||
import os
|
||||
|
||||
from vllm.config.utils import normalize_value
|
||||
|
||||
vllm_envs = {}
|
||||
explicit_envs = set()
|
||||
|
||||
for key in dir(envs):
|
||||
if key.startswith("VLLM_") and "KEY" not in key:
|
||||
value = getattr(envs, key, None)
|
||||
if value is not None:
|
||||
value = normalize_value(value)
|
||||
vllm_envs[key] = value
|
||||
# Check if explicitly set in environment
|
||||
if key in os.environ:
|
||||
explicit_envs.add(key)
|
||||
|
||||
return vllm_envs, explicit_envs
|
||||
|
||||
|
||||
def _get_explicit_cli_args(args) -> tuple[set, dict]:
|
||||
"""Get CLI args that were explicitly set by user (non-default values).
|
||||
|
||||
Returns a tuple of (set of explicit arg names, dict of arg name -> value).
|
||||
"""
|
||||
from vllm.config.utils import normalize_value
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.entrypoints.openai.cli_args import FrontendArgs
|
||||
|
||||
explicit_args: set[str] = set()
|
||||
non_default_args: dict = {}
|
||||
if args is None:
|
||||
return explicit_args, non_default_args
|
||||
|
||||
# Get default values from dataclasses
|
||||
try:
|
||||
frontend_defaults = FrontendArgs()
|
||||
engine_defaults = AsyncEngineArgs(model="")
|
||||
|
||||
for key, default_val in vars(frontend_defaults).items():
|
||||
if hasattr(args, key):
|
||||
current_val = getattr(args, key)
|
||||
if current_val != default_val:
|
||||
explicit_args.add(key)
|
||||
non_default_args[key] = normalize_value(current_val)
|
||||
|
||||
for key, default_val in vars(engine_defaults).items():
|
||||
if hasattr(args, key) and key != "model":
|
||||
current_val = getattr(args, key)
|
||||
if current_val != default_val:
|
||||
explicit_args.add(key)
|
||||
non_default_args[key] = normalize_value(current_val)
|
||||
|
||||
# model is always explicit if set
|
||||
if hasattr(args, "model") and args.model:
|
||||
explicit_args.add("model")
|
||||
non_default_args["model"] = args.model
|
||||
if hasattr(args, "served_model_name") and args.served_model_name:
|
||||
explicit_args.add("served_model_name")
|
||||
non_default_args["served_model_name"] = normalize_value(
|
||||
args.served_model_name
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.debug("Failed to get explicit CLI args: %s", e)
|
||||
|
||||
return explicit_args, non_default_args
|
||||
|
||||
|
||||
_resolved_attn_backend: str | None = None
|
||||
|
||||
|
||||
def _resolve_attn_backend_name(vllm_config, configured_backend) -> str | None:
|
||||
"""Best-effort resolution of the actual attention backend.
|
||||
|
||||
When the user leaves the backend on auto (None), we call the platform's
|
||||
selection logic with the model's head_size / dtype so the dashboard can
|
||||
show the concrete backend (e.g. FLASH_ATTN) instead of "auto".
|
||||
Returns the resolved backend name string or None on failure.
|
||||
"""
|
||||
global _resolved_attn_backend
|
||||
if _resolved_attn_backend is not None:
|
||||
return _resolved_attn_backend
|
||||
|
||||
try:
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.selector import AttentionSelectorConfig
|
||||
|
||||
mc = vllm_config.model_config
|
||||
cc = vllm_config.cache_config
|
||||
|
||||
selector_cfg = AttentionSelectorConfig(
|
||||
head_size=mc.get_head_size(),
|
||||
dtype=mc.dtype,
|
||||
kv_cache_dtype=cc.cache_dtype if cc.cache_dtype != "auto" else None,
|
||||
block_size=getattr(cc, "block_size", None),
|
||||
use_mla=mc.use_mla,
|
||||
)
|
||||
cls_path = current_platform.get_attn_backend_cls(
|
||||
configured_backend,
|
||||
attn_selector_config=selector_cfg,
|
||||
)
|
||||
if not cls_path:
|
||||
return None
|
||||
for member in AttentionBackendEnum:
|
||||
if member.value and member.value == cls_path:
|
||||
_resolved_attn_backend = member.name
|
||||
return _resolved_attn_backend
|
||||
_resolved_attn_backend = cls_path.rsplit(".", 1)[-1].replace(
|
||||
"Backend", ""
|
||||
)
|
||||
return _resolved_attn_backend
|
||||
except Exception as e:
|
||||
logger.debug("Failed to resolve attention backend: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def _get_startup_info(vllm_config) -> dict:
|
||||
"""Derive startup info from vllm_config available in the API server.
|
||||
|
||||
Exposes values that are computed properties or enum-typed in the config
|
||||
and would otherwise be missing or opaque in the JSON-serialized output.
|
||||
Also computes KV cache capacity and max concurrency from cache_config
|
||||
(num_gpu_blocks is set via the engine READY handshake).
|
||||
"""
|
||||
info: dict = {}
|
||||
if vllm_config is None:
|
||||
return info
|
||||
|
||||
# Architecture is a @property, not serialized by Pydantic
|
||||
try:
|
||||
mc = vllm_config.model_config
|
||||
info["architecture"] = mc.architecture
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Attention backend — resolve the actual backend even when config is auto
|
||||
try:
|
||||
ac = vllm_config.attention_config
|
||||
configured = getattr(ac, "backend", None)
|
||||
resolved_name = _resolve_attn_backend_name(vllm_config, configured)
|
||||
if resolved_name:
|
||||
info["attention_backend"] = resolved_name
|
||||
elif configured is not None:
|
||||
info["attention_backend"] = configured.name
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Optimization level
|
||||
try:
|
||||
opt = vllm_config.optimization_level
|
||||
if opt is not None:
|
||||
info["optimization_level"] = f"O{int(opt)}"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Compilation config: mode, cudagraph, inductor partition, fusion passes
|
||||
try:
|
||||
cpc = vllm_config.compilation_config
|
||||
mode = getattr(cpc, "mode", None)
|
||||
if mode is not None:
|
||||
info["compilation_mode"] = mode.name
|
||||
cgm = getattr(cpc, "cudagraph_mode", None)
|
||||
if cgm is not None:
|
||||
info["cudagraph_mode"] = cgm.name
|
||||
igp = getattr(cpc, "use_inductor_graph_partition", None)
|
||||
if igp is not None:
|
||||
info["inductor_graph_partition"] = igp
|
||||
|
||||
pc = getattr(cpc, "pass_config", None)
|
||||
if pc is not None:
|
||||
fusions = {}
|
||||
for attr in (
|
||||
"fuse_norm_quant",
|
||||
"fuse_act_quant",
|
||||
"fuse_attn_quant",
|
||||
"enable_sp",
|
||||
"fuse_gemm_comms",
|
||||
"fuse_allreduce_rms",
|
||||
"enable_qk_norm_rope_fusion",
|
||||
):
|
||||
val = getattr(pc, attr, None)
|
||||
if val is True:
|
||||
fusions[attr] = True
|
||||
if fusions:
|
||||
info["pass_config_fusions"] = fusions
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Kernel config: flashinfer autotune, MoE backend
|
||||
try:
|
||||
kc = vllm_config.kernel_config
|
||||
fi = getattr(kc, "enable_flashinfer_autotune", None)
|
||||
if fi is not None:
|
||||
info["flashinfer_autotune"] = fi
|
||||
moe = getattr(kc, "moe_backend", None)
|
||||
if moe is not None:
|
||||
info["moe_backend"] = moe
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# KV cache capacity and max concurrency
|
||||
try:
|
||||
cc = vllm_config.cache_config
|
||||
mc = vllm_config.model_config
|
||||
|
||||
num_gpu_blocks = getattr(cc, "num_gpu_blocks", None)
|
||||
block_size = getattr(cc, "block_size", None)
|
||||
if num_gpu_blocks and block_size:
|
||||
kv_cache_tokens = num_gpu_blocks * block_size
|
||||
info["kv_cache_tokens"] = kv_cache_tokens
|
||||
|
||||
max_model_len = getattr(mc, "max_model_len", None)
|
||||
if max_model_len and max_model_len > 0:
|
||||
info["max_concurrency"] = round(
|
||||
kv_cache_tokens / max_model_len, 2
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return info
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _get_system_env_info_cached() -> dict:
|
||||
"""Get system environment info (GPU, CUDA, torch, OS, etc.).
|
||||
|
||||
Cached since this info never changes during the lifetime of the server.
|
||||
"""
|
||||
from vllm.collect_env import get_env_info
|
||||
|
||||
return get_env_info()._asdict()
|
||||
|
||||
|
||||
def _get_dashboard_html() -> str:
|
||||
"""Load the dashboard HTML from static file."""
|
||||
static_dir = pathlib.Path(__file__).parent / "static"
|
||||
html_file = static_dir / "index.html"
|
||||
if html_file.exists():
|
||||
return html_file.read_text(encoding="utf-8")
|
||||
return "<html><body><h1>Dashboard HTML not found</h1></body></html>"
|
||||
|
||||
|
||||
@router.get("/dashboard", response_class=HTMLResponse)
|
||||
async def dashboard_index() -> HTMLResponse:
|
||||
"""Serve the main dashboard page."""
|
||||
html_content = _get_dashboard_html()
|
||||
return HTMLResponse(content=html_content)
|
||||
|
||||
|
||||
@router.get("/dashboard/api/info")
|
||||
async def dashboard_info(request: Request) -> JSONResponse:
|
||||
"""Get server information for dashboard display."""
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
|
||||
info: dict = {
|
||||
"version": VLLM_VERSION,
|
||||
"status": "running",
|
||||
}
|
||||
|
||||
# Check engine health
|
||||
try:
|
||||
engine_client = getattr(request.app.state, "engine_client", None)
|
||||
if engine_client is not None:
|
||||
await engine_client.check_health()
|
||||
info["status"] = "running"
|
||||
except EngineDeadError:
|
||||
info["status"] = "unhealthy"
|
||||
except Exception as e:
|
||||
logger.debug("Failed to check engine health: %s", e)
|
||||
|
||||
# Get model information
|
||||
try:
|
||||
serving_models = request.app.state.openai_serving_models
|
||||
if serving_models is not None:
|
||||
models_response = await serving_models.show_available_models()
|
||||
info["models"] = [
|
||||
{"id": model.id, "root": model.root} for model in models_response.data
|
||||
]
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get model info for dashboard: %s", e)
|
||||
info["models"] = []
|
||||
|
||||
# Get full engine config and derive startup info
|
||||
vllm_config = getattr(request.app.state, "vllm_config", None)
|
||||
try:
|
||||
if vllm_config is not None:
|
||||
info["vllm_config"] = PydanticVllmConfig.dump_python(
|
||||
vllm_config, mode="json", fallback=str
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get vllm_config for dashboard: %s", e)
|
||||
|
||||
try:
|
||||
startup_info = _get_startup_info(vllm_config)
|
||||
if startup_info:
|
||||
info["startup_info"] = startup_info
|
||||
except Exception as e:
|
||||
logger.debug("Failed to get startup_info for dashboard: %s", e)
|
||||
|
||||
# Get environment variables (with explicit markers)
|
||||
try:
|
||||
vllm_env, explicit_envs = _get_vllm_env_vars()
|
||||
info["vllm_env"] = vllm_env
|
||||
info["explicit_envs"] = list(explicit_envs)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get vllm_env for dashboard: %s", e)
|
||||
|
||||
# Get system environment info (GPU, CUDA, torch, OS, etc.)
|
||||
try:
|
||||
info["system_env"] = await asyncio.to_thread(
|
||||
_get_system_env_info_cached
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Failed to get system_env for dashboard: %s", e)
|
||||
|
||||
# Get explicit CLI args and non-default args dict
|
||||
try:
|
||||
args = getattr(request.app.state, "args", None)
|
||||
explicit_args, non_default_args = _get_explicit_cli_args(args)
|
||||
info["explicit_args"] = list(explicit_args)
|
||||
info["non_default_args"] = non_default_args
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get explicit_args for dashboard: %s", e)
|
||||
|
||||
# Get server load metrics
|
||||
try:
|
||||
server_load = getattr(request.app.state, "server_load_metrics", None)
|
||||
if server_load is not None:
|
||||
info["server_load"] = server_load
|
||||
except Exception as e:
|
||||
logger.debug("Failed to get server_load for dashboard: %s", e)
|
||||
|
||||
# Check if engine is sleeping (only available in dev mode)
|
||||
try:
|
||||
engine_client = getattr(request.app.state, "engine_client", None)
|
||||
if engine_client is not None:
|
||||
is_sleeping_method = getattr(engine_client, "is_sleeping", None)
|
||||
if is_sleeping_method is not None:
|
||||
info["is_sleeping"] = await is_sleeping_method()
|
||||
except Exception as e:
|
||||
logger.debug("Failed to check is_sleeping for dashboard: %s", e)
|
||||
|
||||
# Check if engine is paused (for RLHF workflows)
|
||||
try:
|
||||
engine_client = getattr(request.app.state, "engine_client", None)
|
||||
if engine_client is not None:
|
||||
is_paused_method = getattr(engine_client, "is_paused", None)
|
||||
if is_paused_method is not None:
|
||||
info["is_paused"] = await is_paused_method()
|
||||
except Exception as e:
|
||||
logger.debug("Failed to check is_paused for dashboard: %s", e)
|
||||
|
||||
return JSONResponse(content=info)
|
||||
|
||||
|
||||
@router.get("/dashboard/api/metrics")
|
||||
async def dashboard_metrics(request: Request) -> JSONResponse:
|
||||
"""Get metrics for dashboard display.
|
||||
|
||||
Returns both Prometheus metrics and internal engine stats that are only
|
||||
accessible in-process (not available via external /metrics endpoint).
|
||||
"""
|
||||
metrics: dict = {}
|
||||
|
||||
# Try to get metrics from prometheus registry
|
||||
try:
|
||||
from vllm.v1.metrics.reader import (
|
||||
Counter,
|
||||
Gauge,
|
||||
Histogram,
|
||||
get_metrics_snapshot,
|
||||
)
|
||||
|
||||
snapshot = get_metrics_snapshot()
|
||||
for metric in snapshot:
|
||||
name = metric.name
|
||||
# For labeled metrics, create unique keys to avoid overwriting
|
||||
# e.g. vllm:request_success with finished_reason label
|
||||
if metric.labels:
|
||||
# Filter out common labels (model_name, engine) for key
|
||||
key_labels = {
|
||||
k: v
|
||||
for k, v in metric.labels.items()
|
||||
if k not in ("model_name", "engine")
|
||||
}
|
||||
if key_labels:
|
||||
# Create key like "vllm:request_success:stop"
|
||||
label_suffix = ":".join(str(v) for v in key_labels.values())
|
||||
name = f"{metric.name}:{label_suffix}"
|
||||
if isinstance(metric, (Counter, Gauge)):
|
||||
metrics[name] = {
|
||||
"value": metric.value,
|
||||
"labels": metric.labels,
|
||||
}
|
||||
elif isinstance(metric, Histogram):
|
||||
metrics[name] = {
|
||||
"count": metric.count,
|
||||
"sum": metric.sum,
|
||||
"labels": metric.labels,
|
||||
}
|
||||
except ImportError:
|
||||
logger.debug("Metrics reader not available")
|
||||
except Exception as e:
|
||||
logger.warning("Failed to get metrics for dashboard: %s", e)
|
||||
|
||||
# Get server load if tracking is enabled
|
||||
try:
|
||||
server_load = getattr(request.app.state, "server_load_metrics", None)
|
||||
if server_load is not None:
|
||||
metrics["server_load"] = {"value": server_load, "labels": {}}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return JSONResponse(content=metrics)
|
||||
|
||||
|
||||
def attach_router(app: FastAPI) -> None:
|
||||
"""Attach dashboard router if enabled via args."""
|
||||
args = getattr(app.state, "args", None)
|
||||
if args is None or not getattr(args, "enable_dashboard", False):
|
||||
return
|
||||
|
||||
logger.info("Enabling vLLM dashboard at /dashboard")
|
||||
app.include_router(router)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1392,23 +1392,19 @@ class FusedMoE(CustomOp):
|
||||
weights = list(self.named_parameters())
|
||||
weights = [(name, _maybe_make_contiguous(name, p)) for name, p in weights]
|
||||
|
||||
# `w13_input_scale` and `w2_input_scale` are global per-tensor
|
||||
# activation scales shared across all experts (e.g. NVFP4).
|
||||
# They are broadcast views (stride 0) from .expand() and are
|
||||
# not actual expert weights, so exclude them from EPLB.
|
||||
NON_EXPERT_WEIGHTS = {
|
||||
"e_score_correction_bias",
|
||||
"w13_input_scale",
|
||||
"w2_input_scale",
|
||||
}
|
||||
|
||||
assert all(
|
||||
weight.is_contiguous()
|
||||
for name, weight in weights
|
||||
if not (name.startswith("_shared_experts.") or name.startswith("_gate."))
|
||||
and name not in NON_EXPERT_WEIGHTS
|
||||
)
|
||||
|
||||
# Filter out the non-expert weights.
|
||||
# `e_score_correction_bias` is a bias for each logical expert,
|
||||
# with shape (num_logical_experts,), not an expert weight.
|
||||
NON_EXPERT_WEIGHTS = {
|
||||
"e_score_correction_bias",
|
||||
}
|
||||
|
||||
return [
|
||||
weight.view(self.local_num_experts, -1)
|
||||
for name, weight in weights
|
||||
|
||||
@@ -365,8 +365,6 @@ def make_nvfp4_moe_quant_config(
|
||||
w2_scale_2: torch.Tensor,
|
||||
a13_scale: torch.Tensor,
|
||||
a2_scale: torch.Tensor,
|
||||
g1_alphas: torch.Tensor | None = None,
|
||||
g2_alphas: torch.Tensor | None = None,
|
||||
) -> FusedMoEQuantConfig:
|
||||
if backend == NvFp4MoeBackend.MARLIN:
|
||||
return nvfp4_w4a16_moe_quant_config(
|
||||
@@ -376,10 +374,8 @@ def make_nvfp4_moe_quant_config(
|
||||
w2_scale=w2_scale,
|
||||
)
|
||||
|
||||
if g1_alphas is None:
|
||||
g1_alphas = a13_scale * w13_scale_2
|
||||
if g2_alphas is None:
|
||||
g2_alphas = a2_scale * w2_scale_2
|
||||
g1_alphas = a13_scale * w13_scale_2
|
||||
g2_alphas = a2_scale * w2_scale_2
|
||||
return nvfp4_moe_quant_config(
|
||||
g1_alphas=g1_alphas,
|
||||
g2_alphas=g2_alphas,
|
||||
|
||||
+2
-22
@@ -554,23 +554,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
layer.w13_input_scale = a13_scale
|
||||
layer.w2_input_scale = a2_scale
|
||||
|
||||
# Pre-compute g1/g2 alphas as registered parameters so EPLB
|
||||
# rearranges them alongside expert weights (see modelopt.py).
|
||||
if self.nvfp4_backend not in (
|
||||
NvFp4MoeBackend.FLASHINFER_TRTLLM,
|
||||
NvFp4MoeBackend.MARLIN,
|
||||
):
|
||||
layer.g1_alphas = torch.nn.Parameter(
|
||||
a13_scale * w13_scale_2, requires_grad=False
|
||||
)
|
||||
layer.g2_alphas = torch.nn.Parameter(
|
||||
a2_scale * w2_scale_2, requires_grad=False
|
||||
)
|
||||
|
||||
# Setup modular kernel for TP case and naive DP/EP case.
|
||||
# In non-naive DP/EP case, we will create a ModularKernelMethod.
|
||||
# TODO(rob): unify these so FP8MoEMethod owns the ModularKernel
|
||||
# in both cases.
|
||||
# Setup modular kernel.
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_nvfp4_moe_kernel(
|
||||
@@ -591,7 +575,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
|
||||
result = make_nvfp4_moe_quant_config(
|
||||
return make_nvfp4_moe_quant_config(
|
||||
backend=self.nvfp4_backend,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
@@ -599,11 +583,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
|
||||
w2_scale_2=layer.w2_weight_scale_2,
|
||||
a13_scale=layer.w13_input_scale,
|
||||
a2_scale=layer.w2_input_scale,
|
||||
g1_alphas=getattr(layer, "g1_alphas", None),
|
||||
g2_alphas=getattr(layer, "g2_alphas", None),
|
||||
)
|
||||
assert result is not None
|
||||
return result
|
||||
|
||||
def apply_monolithic(
|
||||
self,
|
||||
|
||||
@@ -29,7 +29,6 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
|
||||
select_fp8_moe_backend,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
|
||||
NvFp4MoeBackend,
|
||||
convert_to_nvfp4_moe_kernel_format,
|
||||
is_global_sf_supported_for_nvfp4_backend,
|
||||
make_nvfp4_moe_kernel,
|
||||
@@ -1374,29 +1373,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
|
||||
replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
|
||||
replace_parameter(layer, "w2_input_scale", a2_scale)
|
||||
|
||||
# Pre-compute g1/g2 alphas as registered parameters so EPLB
|
||||
# rearranges them alongside expert weights. Without this, the
|
||||
# quant config caches g1_alphas = a_scale * w_scale_2 once at
|
||||
# init, and EPLB's in-place rearrangement of w_scale_2 leaves
|
||||
# the cached product stale, corrupting dequantization.
|
||||
#
|
||||
# Use direct Parameter assignment (not replace_parameter) because
|
||||
# g1_alphas/g2_alphas are not pre-registered in create_weights.
|
||||
if self.nvfp4_backend not in (
|
||||
NvFp4MoeBackend.FLASHINFER_TRTLLM,
|
||||
NvFp4MoeBackend.MARLIN,
|
||||
):
|
||||
layer.g1_alphas = torch.nn.Parameter(
|
||||
(a13_scale * w13_scale_2).contiguous(), requires_grad=False
|
||||
)
|
||||
layer.g2_alphas = torch.nn.Parameter(
|
||||
(a2_scale * w2_scale_2).contiguous(), requires_grad=False
|
||||
)
|
||||
|
||||
# Setup modular kernel for TP case and naive DP/EP case.
|
||||
# In non-naive DP/EP case, we will create a ModularKernelMethod.
|
||||
# TODO(rob): unify these so FP8MoEMethod owns the ModularKernel
|
||||
# in both cases.
|
||||
# Setup modular kernel.
|
||||
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
||||
assert self.experts_cls is not None
|
||||
self.moe_kernel = make_nvfp4_moe_kernel(
|
||||
@@ -1408,7 +1385,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
|
||||
)
|
||||
|
||||
def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
|
||||
result = make_nvfp4_moe_quant_config(
|
||||
return make_nvfp4_moe_quant_config(
|
||||
backend=self.nvfp4_backend,
|
||||
w13_scale=layer.w13_weight_scale,
|
||||
w2_scale=layer.w2_weight_scale,
|
||||
@@ -1416,11 +1393,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
|
||||
w2_scale_2=layer.w2_weight_scale_2,
|
||||
a13_scale=layer.w13_input_scale,
|
||||
a2_scale=layer.w2_input_scale,
|
||||
g1_alphas=getattr(layer, "g1_alphas", None),
|
||||
g2_alphas=getattr(layer, "g2_alphas", None),
|
||||
)
|
||||
assert result is not None
|
||||
return result
|
||||
|
||||
@property
|
||||
def supports_eplb(self) -> bool:
|
||||
|
||||
@@ -191,8 +191,6 @@ _TEXT_GENERATION_MODELS = {
|
||||
"Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
|
||||
"Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
|
||||
"RWForCausalLM": ("falcon", "FalconForCausalLM"),
|
||||
"SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
|
||||
"SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
|
||||
"SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
|
||||
"Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
|
||||
"Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
|
||||
|
||||
@@ -1,786 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Copyright 2026 Sarvam AI team. All rights reserved.
|
||||
#
|
||||
# This code is based on Llama, Deepseek, and Bailing MoE implementations
|
||||
# in this library. It has been modified from its original forms to
|
||||
# accommodate Sarvam's MoE architectures.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable, Iterator
|
||||
from itertools import islice
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
|
||||
from vllm.distributed import (
|
||||
get_pp_group,
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
ReplicatedLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from .bailing_moe import BailingMoeForCausalLM
|
||||
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
|
||||
from .utils import (
|
||||
AutoWeightsLoader,
|
||||
PPMissingLayer,
|
||||
is_pp_missing_parameter,
|
||||
make_empty_intermediate_tensors_factory,
|
||||
make_layers,
|
||||
maybe_prefix,
|
||||
)
|
||||
|
||||
|
||||
def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
|
||||
if scale <= 1:
|
||||
return 1.0
|
||||
return 0.1 * mscale * math.log(scale) + 1.0
|
||||
|
||||
|
||||
def _is_gate_expert_bias_name(name: str) -> bool:
|
||||
return name.endswith(".mlp.gate.e_score_correction_bias") or name.endswith(
|
||||
".gate.e_score_correction_bias"
|
||||
)
|
||||
|
||||
|
||||
def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor:
|
||||
if t.numel() == 0:
|
||||
return t
|
||||
return t - t.mean()
|
||||
|
||||
|
||||
def _normalized_weights(
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> Iterator[tuple[str, torch.Tensor]]:
|
||||
for name, w in weights:
|
||||
if _is_gate_expert_bias_name(name):
|
||||
yield name, _zero_mean_tensor(w)
|
||||
else:
|
||||
yield name, w
|
||||
|
||||
|
||||
class SarvamMLAAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
config,
|
||||
cache_config: CacheConfig | None = None,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.qk_nope_head_dim = config.qk_nope_head_dim
|
||||
self.qk_rope_head_dim = config.qk_rope_head_dim
|
||||
self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
|
||||
self.v_head_dim = config.v_head_dim
|
||||
|
||||
self.q_lora_rank = getattr(config, "q_lora_rank", None)
|
||||
self.kv_lora_rank = config.kv_lora_rank
|
||||
|
||||
self.total_num_heads = config.num_attention_heads
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
assert self.total_num_heads % tp_size == 0
|
||||
self.num_local_heads = self.total_num_heads // tp_size
|
||||
|
||||
self.scaling = self.qk_head_dim**-0.5
|
||||
self.max_position_embeddings = config.max_position_embeddings
|
||||
|
||||
if self.q_lora_rank is not None:
|
||||
self.q_a_proj = ReplicatedLinear(
|
||||
self.hidden_size,
|
||||
self.q_lora_rank,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.q_a_proj",
|
||||
)
|
||||
self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
|
||||
self.q_b_proj = ColumnParallelLinear(
|
||||
self.q_lora_rank,
|
||||
self.total_num_heads * self.qk_head_dim,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.q_b_proj",
|
||||
)
|
||||
self.q_proj = None # type: ignore
|
||||
else:
|
||||
self.q_proj = ColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
self.total_num_heads * self.qk_head_dim,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.q_proj",
|
||||
)
|
||||
self.q_a_proj = None # type: ignore
|
||||
self.q_a_layernorm = None # type: ignore
|
||||
self.q_b_proj = None # type: ignore
|
||||
|
||||
# KV latent (MQA-style) A-proj
|
||||
self.kv_a_proj_with_mqa = ReplicatedLinear(
|
||||
self.hidden_size,
|
||||
self.kv_lora_rank + self.qk_rope_head_dim,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.kv_a_proj_with_mqa",
|
||||
)
|
||||
self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
|
||||
|
||||
# KV B-proj produces per-head K_nope and V
|
||||
self.kv_b_proj = ColumnParallelLinear(
|
||||
self.kv_lora_rank,
|
||||
self.total_num_heads * (self.qk_nope_head_dim + self.v_head_dim),
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.kv_b_proj",
|
||||
)
|
||||
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.total_num_heads * self.v_head_dim,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
)
|
||||
|
||||
self.rotary_emb = get_rope(
|
||||
self.qk_rope_head_dim,
|
||||
# rotary_dim=self.qk_rope_head_dim,
|
||||
max_position=config.max_position_embeddings,
|
||||
rope_parameters=config.rope_parameters,
|
||||
is_neox_style=False,
|
||||
)
|
||||
|
||||
if config.rope_parameters.get("rope_type", None) == "deepseek_yarn":
|
||||
mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
|
||||
scaling_factor = config.rope_parameters["factor"]
|
||||
mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
|
||||
self.scaling = self.scaling * mscale * mscale
|
||||
|
||||
mla_modules = MLAModules(
|
||||
kv_a_layernorm=self.kv_a_layernorm,
|
||||
kv_b_proj=self.kv_b_proj,
|
||||
rotary_emb=self.rotary_emb,
|
||||
o_proj=self.o_proj,
|
||||
fused_qkv_a_proj=None,
|
||||
kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
|
||||
q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
|
||||
q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
|
||||
q_proj=self.q_proj if self.q_lora_rank is None else None,
|
||||
indexer=None,
|
||||
indexer_rotary_emb=None,
|
||||
is_sparse=False,
|
||||
topk_indices_buffer=None,
|
||||
)
|
||||
|
||||
self.mla_attn = MultiHeadLatentAttentionWrapper(
|
||||
self.hidden_size,
|
||||
self.num_local_heads,
|
||||
self.scaling,
|
||||
self.qk_nope_head_dim,
|
||||
self.qk_rope_head_dim,
|
||||
self.v_head_dim,
|
||||
self.q_lora_rank,
|
||||
self.kv_lora_rank,
|
||||
mla_modules,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return self.mla_attn(positions, hidden_states, llama_4_scaling=None)
|
||||
|
||||
|
||||
class SarvamMLAMLP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
intermediate_size: int,
|
||||
config,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
reduce_results: bool = True,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.gate_up_proj = MergedColumnParallelLinear(
|
||||
config.hidden_size,
|
||||
[intermediate_size] * 2,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.gate_up_proj",
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
intermediate_size,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
reduce_results=reduce_results,
|
||||
prefix=f"{prefix}.down_proj",
|
||||
)
|
||||
self.act_fn = SiluAndMul()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
gate_up, _ = self.gate_up_proj(x)
|
||||
x = self.act_fn(gate_up)
|
||||
x, _ = self.down_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class SarvamMLAMoE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config,
|
||||
parallel_config: ParallelConfig,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.config = config
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
self.tp_rank = get_tensor_model_parallel_rank()
|
||||
self.hidden_size = config.hidden_size
|
||||
|
||||
self.num_experts = config.num_experts
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 2.5)
|
||||
|
||||
self.n_group = getattr(config, "n_group", None)
|
||||
self.topk_group = getattr(config, "topk_group", None)
|
||||
self.use_grouped_topk = self.n_group is not None and self.topk_group is not None
|
||||
|
||||
self.norm_expert_prob = getattr(config, "norm_topk_prob", True)
|
||||
|
||||
router_dtype_cfg = getattr(config, "router_dtype", "fp32")
|
||||
if router_dtype_cfg is None:
|
||||
self.router_dtype = None
|
||||
elif router_dtype_cfg == "fp32":
|
||||
self.router_dtype = torch.float32
|
||||
else:
|
||||
self.router_dtype = torch.bfloat16
|
||||
|
||||
self.gate = nn.Linear(
|
||||
self.hidden_size,
|
||||
self.num_experts,
|
||||
bias=False,
|
||||
dtype=self.router_dtype,
|
||||
)
|
||||
|
||||
if getattr(config, "moe_router_enable_expert_bias", True):
|
||||
self.gate.e_score_correction_bias = nn.Parameter(
|
||||
torch.empty(
|
||||
(self.num_experts,),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
)
|
||||
else:
|
||||
self.gate.e_score_correction_bias = None
|
||||
|
||||
self.score_function = getattr(config, "score_function", "sigmoid")
|
||||
self.num_shared_experts = getattr(config, "num_shared_experts", 1)
|
||||
if self.num_shared_experts > 0:
|
||||
if hasattr(config, "moe_shared_expert_intermediate_size"):
|
||||
shared_int = config.moe_shared_expert_intermediate_size
|
||||
else:
|
||||
shared_int = config.moe_intermediate_size
|
||||
shared_int *= self.num_shared_experts
|
||||
self.shared_experts = SarvamMLAMLP(
|
||||
intermediate_size=shared_int,
|
||||
config=config,
|
||||
quant_config=quant_config,
|
||||
reduce_results=False,
|
||||
prefix=f"{prefix}.shared_experts",
|
||||
)
|
||||
else:
|
||||
self.shared_experts = None
|
||||
|
||||
self.experts = SharedFusedMoE(
|
||||
shared_experts=self.shared_experts,
|
||||
num_experts=self.num_experts,
|
||||
top_k=self.top_k,
|
||||
hidden_size=self.hidden_size,
|
||||
intermediate_size=config.moe_intermediate_size,
|
||||
reduce_results=False,
|
||||
renormalize=self.norm_expert_prob,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.experts",
|
||||
scoring_func=self.score_function,
|
||||
e_score_correction_bias=self.gate.e_score_correction_bias,
|
||||
num_expert_group=self.n_group,
|
||||
topk_group=self.topk_group,
|
||||
use_grouped_topk=self.use_grouped_topk,
|
||||
routed_scaling_factor=self.routed_scaling_factor,
|
||||
)
|
||||
|
||||
def maybe_get_fused_moe(self) -> SharedFusedMoE:
|
||||
return self.experts
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
num_tokens, hidden_dim = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
router_logits = self.gate(
|
||||
hidden_states.to(self.router_dtype)
|
||||
if self.router_dtype is not None
|
||||
else hidden_states
|
||||
)
|
||||
router_logits = router_logits.to(hidden_states.dtype)
|
||||
final_hidden = self.experts(
|
||||
hidden_states=hidden_states,
|
||||
router_logits=router_logits,
|
||||
)
|
||||
|
||||
if self.shared_experts is not None:
|
||||
shared_output, expert_output = final_hidden
|
||||
else:
|
||||
shared_output, expert_output = None, final_hidden
|
||||
|
||||
if shared_output is not None:
|
||||
expert_output = expert_output + shared_output
|
||||
|
||||
if self.tp_size > 1:
|
||||
expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(
|
||||
expert_output
|
||||
)
|
||||
|
||||
return expert_output.view(num_tokens, hidden_dim)
|
||||
|
||||
|
||||
class SarvamMLABlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
cache_config = vllm_config.cache_config
|
||||
quant_config = vllm_config.quant_config
|
||||
parallel_config = vllm_config.parallel_config
|
||||
layer_idx = int(prefix.split(".")[-1])
|
||||
hidden_size = config.hidden_size
|
||||
dense_intermediate = getattr(config, "intermediate_size", 16384)
|
||||
|
||||
self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
|
||||
self.self_attn = SarvamMLAAttention(
|
||||
vllm_config=vllm_config,
|
||||
config=config,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.self_attn",
|
||||
)
|
||||
self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
|
||||
use_moe = hasattr(config, "num_experts") and config.num_experts is not None
|
||||
first_k_dense = getattr(config, "first_k_dense_replace", 1)
|
||||
moe_layer_freq = getattr(config, "moe_layer_freq", 1)
|
||||
if use_moe:
|
||||
is_moe_layer = layer_idx >= first_k_dense and (
|
||||
(layer_idx - first_k_dense) % moe_layer_freq == 0
|
||||
)
|
||||
else:
|
||||
is_moe_layer = False
|
||||
|
||||
if is_moe_layer:
|
||||
self.mlp = SarvamMLAMoE(
|
||||
config=config,
|
||||
parallel_config=parallel_config,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
else:
|
||||
self.mlp = SarvamMLAMLP(
|
||||
intermediate_size=dense_intermediate,
|
||||
config=config,
|
||||
quant_config=quant_config,
|
||||
reduce_results=True,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
residual: torch.Tensor | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if residual is None:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
else:
|
||||
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||
|
||||
hidden_states = self.self_attn(
|
||||
positions=positions,
|
||||
hidden_states=hidden_states,
|
||||
)
|
||||
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
class SarvamMLAModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
|
||||
self.config = config
|
||||
self.vocab_size = config.vocab_size
|
||||
self.embed_dim = config.hidden_size
|
||||
self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
|
||||
if get_pp_group().is_first_rank or (
|
||||
self.tie_word_embeddings and get_pp_group().is_last_rank
|
||||
):
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
self.vocab_size,
|
||||
self.embed_dim,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.embed_tokens",
|
||||
)
|
||||
else:
|
||||
self.embed_tokens = PPMissingLayer()
|
||||
|
||||
self.embedding_dropout = torch.nn.Dropout(
|
||||
getattr(config, "embedding_dropout", 0.0)
|
||||
)
|
||||
self.start_layer, self.end_layer, self.layers = make_layers(
|
||||
config.num_hidden_layers,
|
||||
lambda prefix: SarvamMLABlock(
|
||||
vllm_config=vllm_config,
|
||||
prefix=prefix,
|
||||
),
|
||||
prefix=f"{prefix}.layers",
|
||||
)
|
||||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
||||
["hidden_states", "residual"], config.hidden_size
|
||||
)
|
||||
if get_pp_group().is_last_rank:
|
||||
self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)
|
||||
else:
|
||||
self.norm = PPMissingLayer()
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.embed_tokens(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
if get_pp_group().is_first_rank:
|
||||
if inputs_embeds is not None:
|
||||
hidden_states = inputs_embeds
|
||||
else:
|
||||
hidden_states = self.embed_input_ids(input_ids)
|
||||
hidden_states = self.embedding_dropout(hidden_states)
|
||||
residual = None
|
||||
else:
|
||||
assert intermediate_tensors is not None
|
||||
hidden_states = intermediate_tensors["hidden_states"]
|
||||
residual = intermediate_tensors["residual"]
|
||||
|
||||
for layer in islice(self.layers, self.start_layer, self.end_layer):
|
||||
hidden_states, residual = layer(
|
||||
hidden_states,
|
||||
positions,
|
||||
residual,
|
||||
)
|
||||
if not get_pp_group().is_last_rank:
|
||||
return IntermediateTensors(
|
||||
{"hidden_states": hidden_states, "residual": residual}
|
||||
)
|
||||
if residual is None:
|
||||
hidden_states = self.norm(hidden_states)
|
||||
else:
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
return hidden_states
|
||||
|
||||
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
|
||||
return SharedFusedMoE.make_expert_params_mapping(
|
||||
self,
|
||||
ckpt_gate_proj_name="gate_proj",
|
||||
ckpt_down_proj_name="down_proj",
|
||||
ckpt_up_proj_name="up_proj",
|
||||
num_experts=self.config.num_experts,
|
||||
)
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
"""Load weights with stacked gate+up and MoE expert remapping."""
|
||||
weights = _normalized_weights(weights)
|
||||
stacked_params_mapping = [
|
||||
("gate_up_proj", "gate_proj", 0),
|
||||
("gate_up_proj", "up_proj", 1),
|
||||
]
|
||||
|
||||
params_dict = dict(self.named_parameters(remove_duplicate=False))
|
||||
loaded_params: set[str] = set()
|
||||
expert_params_mapping = self.get_expert_mapping()
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
if "mlp.experts" in name:
|
||||
continue
|
||||
new_name = name.replace(weight_name, param_name)
|
||||
if new_name.endswith(".bias") and new_name not in params_dict:
|
||||
continue
|
||||
if new_name not in params_dict:
|
||||
continue
|
||||
if is_pp_missing_parameter(new_name, self):
|
||||
continue
|
||||
|
||||
param = params_dict[new_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
loaded_params.add(new_name)
|
||||
break
|
||||
else:
|
||||
mapped = False
|
||||
for (
|
||||
param_name,
|
||||
weight_name,
|
||||
expert_id,
|
||||
shard_id,
|
||||
) in expert_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
|
||||
new_name = name.replace(weight_name, param_name)
|
||||
if is_pp_missing_parameter(new_name, self):
|
||||
continue
|
||||
if new_name not in params_dict:
|
||||
continue
|
||||
|
||||
param = params_dict[new_name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", default_weight_loader
|
||||
)
|
||||
weight_loader(
|
||||
param,
|
||||
loaded_weight,
|
||||
name,
|
||||
shard_id=shard_id,
|
||||
expert_id=expert_id,
|
||||
)
|
||||
loaded_params.add(new_name)
|
||||
mapped = True
|
||||
break
|
||||
|
||||
if mapped:
|
||||
continue
|
||||
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
if name not in params_dict:
|
||||
continue
|
||||
if is_pp_missing_parameter(name, self):
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
|
||||
return loaded_params
|
||||
|
||||
|
||||
class SarvamMixtureOfExperts(MixtureOfExperts):
|
||||
def extract_moe_parameters(self, example_moe: SarvamMLAMoE | None) -> None:
|
||||
if example_moe is None:
|
||||
raise RuntimeError("No SarvamMLAMoE layer found in model.layers.")
|
||||
|
||||
self.num_logical_experts = example_moe.num_experts
|
||||
self.num_routed_experts = example_moe.num_experts # routed pool size
|
||||
self.num_shared_experts = getattr(example_moe.config, "num_shared_experts", 1)
|
||||
|
||||
self.num_physical_experts = self.num_logical_experts
|
||||
self.num_local_physical_experts = self.num_logical_experts
|
||||
self.num_redundant_experts = 0
|
||||
|
||||
def update_physical_experts_metadata(
|
||||
self,
|
||||
num_physical_experts: int,
|
||||
num_local_physical_experts: int,
|
||||
) -> None:
|
||||
self.num_physical_experts = num_physical_experts
|
||||
self.num_local_physical_experts = num_local_physical_experts
|
||||
self.num_redundant_experts = num_physical_experts - self.num_logical_experts
|
||||
|
||||
for moe in self.moe_mlp_layers:
|
||||
moe.n_physical_experts = num_physical_experts
|
||||
moe.n_local_physical_experts = num_local_physical_experts
|
||||
moe.n_redundant_experts = self.num_redundant_experts
|
||||
|
||||
fused = moe.experts
|
||||
if hasattr(fused, "n_local_physical_experts"):
|
||||
fused.n_local_physical_experts = num_local_physical_experts
|
||||
if hasattr(fused, "n_physical_experts"):
|
||||
fused.n_physical_experts = num_physical_experts
|
||||
if hasattr(fused, "n_redundant_experts"):
|
||||
fused.n_redundant_experts = self.num_redundant_experts
|
||||
if hasattr(fused, "update_expert_map"):
|
||||
fused.update_expert_map()
|
||||
|
||||
def set_eplb_state(self, eplb_state) -> None:
|
||||
self.eplb_state = eplb_state
|
||||
for moe in self.moe_layers:
|
||||
if hasattr(moe, "set_eplb_state"):
|
||||
moe.set_eplb_state(eplb_state)
|
||||
|
||||
|
||||
class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfExperts):
|
||||
packed_modules_mapping = {
|
||||
"q_proj": ["q_proj"],
|
||||
"q_a_proj": ["q_a_proj"],
|
||||
"q_b_proj": ["q_b_proj"],
|
||||
"kv_a_proj_with_mqa": ["kv_a_proj_with_mqa"],
|
||||
"kv_b_proj": ["kv_b_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
self.config = config
|
||||
self.quant_config = quant_config
|
||||
|
||||
self.model = SarvamMLAModel(
|
||||
vllm_config=vllm_config,
|
||||
prefix=maybe_prefix(prefix, "model"),
|
||||
)
|
||||
|
||||
self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False)
|
||||
if get_pp_group().is_last_rank:
|
||||
if self.tie_word_embeddings:
|
||||
self.lm_head = self.model.embed_tokens
|
||||
else:
|
||||
self.lm_head = ParallelLMHead(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "lm_head"),
|
||||
)
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
else:
|
||||
self.lm_head = PPMissingLayer()
|
||||
self.logits_processor = None # type: ignore
|
||||
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.model.make_empty_intermediate_tensors
|
||||
)
|
||||
|
||||
self.expert_weights = []
|
||||
self.num_moe_layers = 0
|
||||
|
||||
self.moe_layers = []
|
||||
self.moe_mlp_layers = []
|
||||
|
||||
example_moe = None
|
||||
for layer in self.model.layers:
|
||||
if isinstance(layer, PPMissingLayer):
|
||||
continue
|
||||
if isinstance(layer.mlp, SarvamMLAMoE):
|
||||
example_moe = layer.mlp
|
||||
self.moe_mlp_layers.append(layer.mlp)
|
||||
self.moe_layers.append(layer.mlp.experts)
|
||||
self.num_moe_layers += 1
|
||||
|
||||
self.extract_moe_parameters(example_moe)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.model.embed_input_ids(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
return self.model(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
)
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor | None:
|
||||
if not get_pp_group().is_last_rank:
|
||||
return None
|
||||
logits = self.logits_processor(self.lm_head, hidden_states)
|
||||
return logits
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
loader = AutoWeightsLoader(
|
||||
self,
|
||||
skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
|
||||
)
|
||||
return loader.load_weights(weights)
|
||||
|
||||
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
|
||||
return self.model.get_expert_mapping()
|
||||
|
||||
|
||||
class SarvamMoEForCausalLM(BailingMoeForCausalLM):
|
||||
"""Same as BailingMoeForCausalLM, but normalizes gate expert_bias pre-load."""
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
return super().load_weights(_normalized_weights(weights))
|
||||
@@ -68,7 +68,7 @@ class ToolParser:
|
||||
# tool_choice: "Forced Function" or "required" will override
|
||||
# structured output json settings to make tool calling work correctly
|
||||
request.structured_outputs = StructuredOutputsParams(
|
||||
json=json_schema_from_tool # type: ignore[call-arg]
|
||||
json=json_schema_from_tool
|
||||
)
|
||||
request.response_format = None
|
||||
if isinstance(request, ResponsesRequest):
|
||||
|
||||
Reference in New Issue
Block a user