Compare commits

..
Author SHA1 Message Date
esmeetu 4b6ff605fc update frontend
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-03-08 04:23:28 +00:00
esmeetu c5eacd261d Merge remote-tracking branch 'origin/main' into vllm-dashboard 2026-03-07 15:08:32 +00:00
esmeetu 549b8242d6 refactor
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-03-03 09:21:45 +00:00
esmeetu 3fde6cc382 Merge branch 'main' into vllm-dashboard
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-03-03 08:31:16 +00:00
esmeetu f0b536640c Merge branch 'main' into vllm-dashboard
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-03-01 12:28:39 +08:00
Roy WangandGitHub 9e173a5cdc Merge branch 'main' into vllm-dashboard 2026-02-03 13:16:14 +08:00
Roy WangandGitHub 09a977e897 Merge branch 'main' into vllm-dashboard 2026-01-29 16:51:43 +08:00
esmeetu 1b2b0043d8 add more metrics
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 21:15:37 +08:00
esmeetu a87a2676f2 add multi-modal support test
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 19:34:19 +08:00
esmeetu f54775736a add params
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 19:33:22 +08:00
esmeetu 934e6b2a46 update style & add filter
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 19:09:47 +08:00
esmeetu e9966f1c52 remove engine stats
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 18:53:46 +08:00
esmeetu 12e25a1d3f fix null stats
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 18:41:41 +08:00
esmeetu 3bebb3d705 refactor webui & add Engine Stats
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 18:34:22 +08:00
esmeetu 36547719b0 support healthcheck
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-15 17:56:36 +08:00
esmeetu d4d3e90c3e Add --enable-dashboard for vLLM Web UI
Signed-off-by: esmeetu <jasonailu87@gmail.com>
2026-01-13 18:49:10 +08:00
17 changed files with 3842 additions and 994 deletions
-2
View File
@@ -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. | | |
+292
View File
@@ -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
-12
View File
@@ -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,
+11 -117
View File
@@ -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]
+5
View File
@@ -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.
+13 -6
View File
@@ -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,
+6
View File
@@ -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
+7 -11
View File
@@ -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,
@@ -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:
-2
View File
@@ -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"),
-786
View File
@@ -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))
+1 -1
View File
@@ -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):