[Misc] Add EPLB debug logging for balance diagnostics

Add comprehensive debug logging to the EPLB system to help diagnose expert load balancing issues in wideEP deployments: - Per-step balance breakdown: worst/best layer indices, min/max rank token counts for the worst layer, replica distribution stats - Pre-rearrange diagnostics: window utilization, load distribution across logical experts, top-5 hottest experts - Post-rearrange diagnostics: number of changed slots, replica count stats, predicted post-rearrange balancedness (simulates expected balance with the new mapping applied to current load data) - Warning when window_size > step_interval (stale data risk) Signed-off-by: Travis Shears <travis@neuralmagic.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
[BugFix] Fix EPLB balancedness metric using wrong dimension
2026-03-06 22:20:28 -05:00 · 2026-03-06 22:18:49 -05:00
68 changed files with 351 additions and 1377 deletions
@@ -83,6 +83,7 @@ We test the throughput by using `vllm bench serve` with request rate = inf to co
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3-8B",
            "tensor_parallel_size": 1,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -10,6 +10,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
            "tensor_parallel_size": 1,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy",
            "max-model-len": 2048,
@@ -36,6 +37,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
            "tensor_parallel_size": 4,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy",
            "max-model-len": 2048,
@@ -62,6 +64,7 @@
        "server_parameters": {
            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
            "tensor_parallel_size": 2,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy",
            "max-model-len": 2048,
@@ -88,6 +91,7 @@
        "server_parameters": {
            "model": "deepseek-ai/DeepSeek-R1",
            "tensor_parallel_size": 8,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy",
            "max-model-len": 2048,
@@ -5,6 +5,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
            "tensor_parallel_size": 1,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -22,6 +23,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
            "tensor_parallel_size": 4,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -39,6 +41,7 @@
        "server_parameters": {
            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
            "tensor_parallel_size": 2,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -56,6 +59,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct", 
            "tensor_parallel_size": 4,
+            "swap_space": 16,
            "speculative_config": {
                "model": "turboderp/Qwama-0.5B-Instruct",
                "num_speculative_tokens": 4,
@@ -2801,7 +2801,7 @@ steps:
  - vllm/v1/attention/selector.py
  - vllm/platforms/cuda.py
  commands:
-    - rocm-smi
+    rocm-smi
    - python3 examples/offline_inference/basic/chat.py
    # Attention
    # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
@@ -3283,7 +3283,7 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010

- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355)
+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200/MI355)
  mirror_hardwares: [amdexperimental, amdproduction, amdmi355]
  agent_pool: mi355_2
  timeout_in_minutes: 60
@@ -3305,7 +3305,7 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040

- label: Attention Benchmarks Smoke Test (B200-MI355)
+- label: Attention Benchmarks Smoke Test (B200/MI355)
  device: b200
  mirror_hardwares: [amdexperimental, amdmi355]
  agent_pool: mi355_2
@@ -145,6 +145,7 @@ def create_minimal_vllm_config(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=False,
    )
@@ -141,6 +141,7 @@ def _create_vllm_config(
    cache_config = CacheConfig(
        block_size=config.block_size,
        cache_dtype="auto",
+        swap_space=0,
    )
    cache_config.num_gpu_blocks = max_num_blocks
    cache_config.num_cpu_blocks = 0
@@ -507,10 +507,10 @@ longer relevant in v1:
 - `vllm:num_requests_swapped`
 - `vllm:cpu_cache_usage_perc`

-In this mode, when a request was preempted (e.g. to make room in KV
-cache to complete other requests), kv cache blocks were swapped out to
-CPU memory. The `--swap-space` flag has been removed as this feature
-is no longer used in V1.
+In this mode, when a request is preempted (e.g. to make room in KV
+cache to complete other requests), we swap kv cache blocks out to CPU
+memory. This is also known as "KV cache offloading" and is configured
+with `--swap-space` and `--preemption-mode`.

 Historically, [vLLM has long supported beam search](https://github.com/vllm-project/vllm/issues/6226). The
 SequenceGroup encapsulated the idea of N Sequences which
@@ -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. | | |
@@ -17,7 +17,7 @@ llm = Vllm(
    model="microsoft/Orca-2-7b",
    tensor_parallel_size=4,
    max_new_tokens=100,
-    vllm_kwargs={"gpu_memory_utilization": 0.5},
+    vllm_kwargs={"swap_space": 1, "gpu_memory_utilization": 0.5},
 )
 ```

@@ -794,6 +794,7 @@ class VllmRunner:
        tensor_parallel_size: int = 1,
        block_size: int = 16 if not torch.xpu.is_available() else 64,
        enable_chunked_prefill: bool | None = False,
+        swap_space: int = 4,
        enforce_eager: bool | None = False,
        # Set this to avoid hanging issue
        default_torch_num_threads: int | None = None,
@@ -830,6 +831,7 @@ class VllmRunner:
                trust_remote_code=trust_remote_code,
                dtype=dtype,
                seed=seed,
+                swap_space=swap_space,
                enforce_eager=enforce_eager,
                disable_log_stats=disable_log_stats,
                tensor_parallel_size=tensor_parallel_size,
@@ -22,7 +22,7 @@ prompts = [

 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

-# set different `gpu_memory_utilization` for different ranks,
+# set different `gpu_memory_utilization` and `swap_space` for different ranks,
 # to test if all ranks agree on the same kv cache configuration.
 llm = LLM(
    model="facebook/opt-125m",
@@ -30,6 +30,7 @@ llm = LLM(
    pipeline_parallel_size=int(os.getenv("PP_SIZE", 1)),
    distributed_executor_backend="external_launcher",
    gpu_memory_utilization=random.uniform(0.7, 0.9),
+    swap_space=random.randint(1, 4),
    seed=0,
 )

@@ -28,7 +28,7 @@ if dp_size > 1:

 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

-# set different `gpu_memory_utilization` for different ranks,
+# set different `gpu_memory_utilization` and `swap_space` for different ranks,
 # to test if all ranks agree on the same kv cache configuration.
 llm = LLM(
    model="microsoft/Phi-mini-MoE-instruct",
@@ -37,6 +37,7 @@ llm = LLM(
    enable_expert_parallel=int(os.getenv("ENABLE_EP", "0")) == 1,
    distributed_executor_backend="external_launcher",
    gpu_memory_utilization=random.uniform(0.7, 0.9),
+    swap_space=random.randint(1, 4),
    seed=0,
 )

@@ -13,7 +13,7 @@ import websockets

 from vllm.assets.audio import AudioAsset

-from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
+from ...utils import RemoteOpenAIServer
 from .conftest import add_attention_backend

 MISTRAL_FORMAT_ARGS = [
@@ -23,7 +23,7 @@ MISTRAL_FORMAT_ARGS = [
    "mistral",
    "--load_format",
    "mistral",
-] + ROCM_EXTRA_ARGS
+]

 MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"

@@ -77,9 +77,7 @@ async def test_multi_chunk_streaming(

    add_attention_backend(server_args, rocm_aiter_fa_attention)

-    with RemoteOpenAIServer(
-        model_name, server_args, env_dict=ROCM_ENV_OVERRIDES
-    ) as remote_server:
+    with RemoteOpenAIServer(model_name, server_args) as remote_server:
        ws_url = _get_websocket_url(remote_server)
        async with websockets.connect(ws_url) as ws:
            # Receive session.created
@@ -180,9 +178,7 @@ async def test_empty_commit_does_not_crash_engine(

    add_attention_backend(server_args, rocm_aiter_fa_attention)

-    with RemoteOpenAIServer(
-        model_name, server_args, env_dict=ROCM_ENV_OVERRIDES
-    ) as remote_server:
+    with RemoteOpenAIServer(model_name, server_args) as remote_server:
        ws_url = _get_websocket_url(remote_server)

        # --- First connection: empty commit (no audio appended) ----------
@@ -35,8 +35,6 @@ def server():
        "--trust-remote-code",
        "--limit-mm-per-prompt",
        json.dumps({"video": MAXIMUM_VIDEOS}),
-        "--media-io-kwargs",
-        json.dumps({"video": {"num_frames": 32}}),
    ]

    # ROCm: Increase timeouts to handle potential network delays and slower
@@ -129,73 +127,6 @@ async def test_single_chat_session_video(
    assert message.content is not None and len(message.content) >= 0


-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("video_url", [TEST_VIDEO_URLS[0]])
-async def test_request_media_io_kwargs_override_uses_fewer_video_frames(
-    client: openai.AsyncOpenAI, model_name: str, video_url: str
-):
-    messages = dummy_messages_from_video_url(video_url)
-
-    default_resp = await client.chat.completions.create(
-        model=model_name,
-        messages=messages,
-        max_completion_tokens=1,
-        temperature=0.0,
-    )
-    override_resp = await client.chat.completions.create(
-        model=model_name,
-        messages=messages,
-        max_completion_tokens=1,
-        temperature=0.0,
-        extra_body={
-            "media_io_kwargs": {
-                "video": {
-                    "num_frames": 4,
-                }
-            }
-        },
-    )
-
-    assert default_resp.usage is not None
-    assert override_resp.usage is not None
-    assert override_resp.usage.prompt_tokens < default_resp.usage.prompt_tokens
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-@pytest.mark.parametrize("video_url", [TEST_VIDEO_URLS[0]])
-async def test_invalid_num_frames_request_recoverable(
-    client: openai.AsyncOpenAI, model_name: str, video_url: str
-):
-    messages = dummy_messages_from_video_url(video_url)
-
-    with pytest.raises((openai.BadRequestError, openai.APIStatusError)):
-        await client.chat.completions.create(
-            model=model_name,
-            messages=messages,
-            max_completion_tokens=1,
-            temperature=0.0,
-            extra_body={
-                "media_io_kwargs": {
-                    "video": {
-                        "num_frames": "invalid",
-                    }
-                }
-            },
-        )
-
-    # Server should still handle subsequent requests after the failed one.
-    recovery_resp = await client.chat.completions.create(
-        model=model_name,
-        messages=messages,
-        max_completion_tokens=1,
-        temperature=0.0,
-    )
-    recovery_msg = recovery_resp.choices[0].message
-    assert recovery_msg.content is not None and len(recovery_msg.content) >= 0
-
-
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
@@ -127,39 +127,6 @@ def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
    assert output.usage.prompt_tokens == 767


-@pytest.mark.parametrize("model_name", [MODEL_NAME])
-def test_chat_image_with_media_io_kwargs(server: RemoteOpenAIServer, model_name: str):
-    rgba_image_url = (
-        "https://vllm-public-assets.s3.us-west-2.amazonaws.com"
-        "/vision_model_images/RGBA_comp.png"
-    )
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": "Represent the user's input."},
-                {"type": "image_url", "image_url": {"url": rgba_image_url}},
-            ],
-        }
-    ]
-
-    response = requests.post(
-        server.url_for("v1/embeddings"),
-        json={
-            "model": model_name,
-            "messages": messages,
-            "media_io_kwargs": {
-                "image": {"rgba_background_color": [0, 0, 0]},
-            },
-        },
-    )
-    response.raise_for_status()
-
-    output = EmbeddingResponse.model_validate(response.json())
-    assert len(output.data) == 1
-    assert len(output.data[0].embedding) == 3072
-
-
 def get_hf_prompt_tokens(model_name, content, image_url):
    processor = AutoProcessor.from_pretrained(
        model_name, trust_remote_code=True, num_crops=4
@@ -64,6 +64,7 @@ def test_worker_apply_lora(qwen3_lora_files):
        device_config=DeviceConfig("cuda"),
        cache_config=CacheConfig(
            block_size=16,
+            swap_space=0,
            cache_dtype="auto",
        ),
        lora_config=LoRAConfig(
@@ -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,
@@ -182,6 +182,7 @@ def create_vllm_config(
    cache_config = CacheConfig(
        block_size=block_size,
        cache_dtype="auto",
+        swap_space=0,
    )
    # Set cache blocks for testing
    #   (these may be set during initialization normally)
@@ -1776,6 +1776,7 @@ def create_scheduler_with_priority(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=enable_prefix_caching,
    )
@@ -3725,6 +3726,7 @@ def _create_encoder_decoder_scheduler(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=False,
    )
@@ -94,6 +94,7 @@ def create_scheduler(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=enable_prefix_caching,
    )
@@ -506,6 +506,7 @@ def test_encoder_instance_zero_kv_cache(
    cache_config = CacheConfig(
        block_size=16,
        gpu_memory_utilization=gpu_memory_utilization,
+        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=enable_prefix_caching,
    )
@@ -206,6 +206,7 @@ def create_vllm_config(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=True,
    )
@@ -118,6 +118,7 @@ def create_vllm_config(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype=cache_dtype,
        enable_prefix_caching=True,
    )
@@ -96,6 +96,7 @@ def get_vllm_config():
    cache_config = CacheConfig(
        block_size=BLOCK_SIZE,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
    )
    parallel_config = ParallelConfig()
@@ -808,6 +809,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
    cache_config = CacheConfig(
        block_size=BLOCK_SIZE,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
    )
    parallel_config = ParallelConfig()
@@ -1240,6 +1242,7 @@ def test_cudagraph_sizes_capped_for_mamba_cache():
    cache_config = CacheConfig(
        block_size=BLOCK_SIZE,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
    )
    parallel_config = ParallelConfig()
@@ -38,7 +38,6 @@ from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.inputs import TextPrompt, TokensPrompt
 from vllm.lora.request import LoRARequest
 from vllm.outputs import RequestOutput
-from vllm.platforms import current_platform
 from vllm.sampling_params import BeamSearchParams
 from vllm.tokenizers import TokenizerLike, get_tokenizer
 from vllm.utils.async_utils import merge_async_iterators
@@ -257,21 +256,17 @@ def run_hf(
    max_batch_size: int,
    trust_remote_code: bool,
    disable_detokenize: bool = False,
-    dtype: torch.dtype | None = torch.float16,
-    enable_torch_compile: bool = False,
 ) -> float:
    assert isinstance(tokenizer, PreTrainedTokenizerBase), (
        "the hf backend only supports HF tokenizers"
    )
    llm = AutoModelForCausalLM.from_pretrained(
-        model, dtype=dtype, trust_remote_code=trust_remote_code
+        model, dtype=torch.float16, trust_remote_code=trust_remote_code
    )
    if llm.config.model_type == "llama":
        # To enable padding in the HF backend.
        tokenizer.pad_token = tokenizer.eos_token
-    llm = llm.to(current_platform.device_type)
-    if enable_torch_compile:
-        llm = torch.compile(llm)
+    llm = llm.cuda()

    pbar = tqdm(total=len(requests))
    start = time.perf_counter()
@@ -300,7 +295,7 @@ def run_hf(
        # Generate the sequences.
        input_ids = tokenizer(batch, return_tensors="pt", padding=True).input_ids
        llm_outputs = llm.generate(
-            input_ids=input_ids.to(current_platform.device_type),
+            input_ids=input_ids.cuda(),
            do_sample=True,
            num_return_sequences=n,
            temperature=1.0,
@@ -738,12 +733,6 @@ def add_cli_args(parser: argparse.ArgumentParser):
        default=None,
        help="Maximum batch size for HF backend.",
    )
-    parser.add_argument(
-        "--hf-enable-torch-compile",
-        action="store_true",
-        default=False,
-        help="Enable Torch compile for HF backend.",
-    )
    parser.add_argument(
        "--output-json",
        type=str,
@@ -895,8 +884,6 @@ def main(args: argparse.Namespace):
            args.hf_max_batch_size,
            args.trust_remote_code,
            args.disable_detokenize,
-            dtype=args.dtype,
-            enable_torch_compile=args.hf_enable_torch_compile,
        )
    elif args.backend == "vllm-chat":
        elapsed_time, request_outputs = run_vllm_chat(
@@ -1,13 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import math
 from dataclasses import field
-from typing import Literal
+from typing import TYPE_CHECKING, Any, Literal

 from pydantic import Field, SkipValidation, field_validator

 from vllm.config.utils import config
 from vllm.logger import init_logger
+from vllm.utils.mem_constants import GiB_bytes
+from vllm.utils.mem_utils import format_gib, get_cpu_memory
+
+if TYPE_CHECKING:
+    from vllm.config.parallel import ParallelConfig
+else:
+    ParallelConfig = Any

 logger = init_logger(__name__)

@@ -45,6 +53,8 @@ class CacheConfig:
    not matter if you have another vLLM instance running on the same GPU. For
    example, if you have two vLLM instances running on the same GPU, you can
    set the GPU memory utilization to 0.5 for each instance."""
+    swap_space: float = Field(default=4, ge=0)
+    """Size of the CPU swap space per GPU (in GiB)."""
    cache_dtype: CacheDType = "auto"
    """Data type for kv cache storage. If "auto", will use model data type.
    CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports
@@ -163,6 +173,7 @@ class CacheConfig:
        ignored_factors = {
            # Runtime/derived knobs that don't affect compiled graph shape
            "gpu_memory_utilization",
+            "swap_space",
            "is_attention_free",
            "num_gpu_blocks_override",
            "enable_prefix_caching",
@@ -197,3 +208,24 @@ class CacheConfig:
                "scaling factor."
            )
        return cache_dtype
+
+    def verify_with_parallel_config(
+        self,
+        parallel_config: ParallelConfig,
+    ) -> None:
+        swap_space_bytes = math.ceil(self.swap_space * GiB_bytes)
+        total_cpu_memory = get_cpu_memory()
+        # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel
+        # group are in the same node. However, the GPUs may span multiple nodes.
+        num_gpus_per_node = parallel_config.tensor_parallel_size
+        cpu_memory_usage = swap_space_bytes * num_gpus_per_node
+
+        msg = (
+            f"{format_gib(cpu_memory_usage)} GiB out of the "
+            f"{format_gib(total_cpu_memory)} GiB total CPU memory "
+            "is allocated for the swap space."
+        )
+        if cpu_memory_usage > 0.7 * total_cpu_memory:
+            raise ValueError("Too large swap space. " + msg)
+        elif cpu_memory_usage > 0.4 * total_cpu_memory:
+            logger.warning("Possibly too large swap space. %s", msg)
@@ -674,6 +674,8 @@ class VllmConfig:

            self.parallel_config.is_moe_model = self.model_config.is_moe

+        self.cache_config.verify_with_parallel_config(self.parallel_config)
+
        if self.lora_config is not None:
            self.lora_config.verify_with_model_config(self.model_config)

@@ -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:
    """
@@ -701,6 +602,71 @@ class EplbState:
                        - self.expert_rearrangement_step,
                    )

+                    # Per-layer breakdown: worst/best layers,
+                    # per-rank token counts for the worst layer
+                    worst_layer = int(per_layer_balance.argmin().item())
+                    best_layer = int(per_layer_balance.argmax().item())
+                    worst_balance = float(
+                        per_layer_balance[worst_layer].item()
+                    )
+                    best_balance = float(
+                        per_layer_balance[best_layer].item()
+                    )
+
+                    worst_layer_ranks = num_tokens_per_rank[worst_layer]
+                    worst_min_rank = int(
+                        worst_layer_ranks.argmin().item()
+                    )
+                    worst_max_rank = int(
+                        worst_layer_ranks.argmax().item()
+                    )
+
+                    logger.info(
+                        "EPLB balance breakdown: "
+                        "worst_layer=%d (balance=%.4f, "
+                        "min_rank=%d[%.0f], max_rank=%d[%.0f]), "
+                        "best_layer=%d (balance=%.4f), "
+                        "num_layers=%d",
+                        worst_layer,
+                        worst_balance,
+                        worst_min_rank,
+                        float(
+                            worst_layer_ranks[worst_min_rank].item()
+                        ),
+                        worst_max_rank,
+                        float(
+                            worst_layer_ranks[worst_max_rank].item()
+                        ),
+                        best_layer,
+                        best_balance,
+                        num_tokens_per_rank.shape[0],
+                    )
+
+                    # Log replica distribution for debug
+                    replica_count = (
+                        eplb_model_state.logical_replica_count
+                    )
+                    if (
+                        replica_count is not None
+                        and replica_count.numel() > 0
+                    ):
+                        rc_float = replica_count.float()
+                        logger.debug(
+                            "EPLB replica stats (layer avg): "
+                            "min=%.1f, max=%.1f, mean=%.2f, "
+                            "num_with_replicas=%d/%d",
+                            float(rc_float.min().item()),
+                            float(rc_float.max().item()),
+                            float(rc_float.mean().item()),
+                            int(
+                                (rc_float > 1)
+                                .any(dim=0)
+                                .sum()
+                                .item()
+                            ),
+                            replica_count.shape[-1],
+                        )
+
        # Update the expert load sliding window
        if not is_dummy:
            for eplb_model_state in self.model_states.values():
@@ -777,6 +743,40 @@ class EplbState:
            )

        # Map the physical expert load to global logical experts
+        if is_main_rank:
+            # Log window utilization diagnostics
+            nonzero_slots = sum(
+                int(
+                    (ms.expert_load_window.sum(dim=(1, 2)) > 0)
+                    .sum()
+                    .item()
+                )
+                for ms in self.model_states.values()
+            )
+            logger.info(
+                "EPLB window state: window_step=%d/%d, "
+                "rearrangement_step=%d/%d, "
+                "nonzero_window_slots=%d/%d",
+                self.expert_load_window_step,
+                self.expert_load_window_size,
+                self.expert_rearrangement_step,
+                self.expert_rearrangement_step_interval,
+                nonzero_slots,
+                self.expert_load_window_size,
+            )
+            if (
+                self.expert_load_window_size
+                > self.expert_rearrangement_step_interval
+            ):
+                logger.warning(
+                    "EPLB: window_size (%d) > step_interval (%d). "
+                    "Stale window entries from before the last "
+                    "rearrangement will be converted with the current "
+                    "physical->logical mapping, which may be incorrect. "
+                    "Consider setting window_size <= step_interval.",
+                    self.expert_load_window_size,
+                    self.expert_rearrangement_step_interval,
+                )
        global_expert_load_windows = []
        for eplb_model_state in self.model_states.values():
            expert_load_window = eplb_model_state.expert_load_window[
@@ -839,6 +839,40 @@ class EplbState:
        for eplb_model_state, global_expert_load_window in zip(
            self.model_states.values(), global_expert_load_windows
        ):
+            if is_main_rank:
+                # Log load statistics the algorithm will use
+                load = global_expert_load_window.float()
+                load_per_layer = load.sum(dim=-1)
+                logger.info(
+                    "EPLB rearrange input: "
+                    "num_replicas=%d, num_groups=%d, "
+                    "num_nodes=%d, num_gpus=%d, "
+                    "total_load_per_layer: "
+                    "min=%.0f, max=%.0f, mean=%.0f",
+                    num_replicas,
+                    num_groups,
+                    num_nodes,
+                    num_gpus,
+                    float(load_per_layer.min().item()),
+                    float(load_per_layer.max().item()),
+                    float(load_per_layer.mean().item()),
+                )
+                # Top-5 hottest experts (averaged across layers)
+                avg_load = load.mean(dim=0)
+                top5_vals, top5_ids = avg_load.topk(
+                    min(5, avg_load.shape[0])
+                )
+                logger.info(
+                    "EPLB top-5 hottest logical experts "
+                    "(avg across layers): %s",
+                    ", ".join(
+                        f"e{int(eid)}={float(val):.0f}"
+                        for eid, val in zip(
+                            top5_ids.tolist(), top5_vals.tolist()
+                        )
+                    ),
+                )
+
            if not self.is_async or is_profile:
                # Get new expert mappings for the model
                (
@@ -854,6 +888,66 @@ class EplbState:
                    eplb_model_state.physical_to_logical_map,
                )

+                if is_main_rank and not is_profile:
+                    # Log what the algorithm decided
+                    old_p2l = eplb_model_state.physical_to_logical_map
+                    new_p2l = new_physical_to_logical_map.to(
+                        old_p2l.device
+                    )
+                    changed_slots = int(
+                        (old_p2l != new_p2l).sum().item()
+                    )
+                    total_slots = old_p2l.numel()
+                    rc = new_logical_replica_count.float()
+                    logger.info(
+                        "EPLB rearrange result: "
+                        "changed_slots=%d/%d (%.1f%%), "
+                        "replica_count: "
+                        "min=%.0f, max=%.0f, mean=%.2f",
+                        changed_slots,
+                        total_slots,
+                        100.0 * changed_slots / max(total_slots, 1),
+                        float(rc.min().item()),
+                        float(rc.max().item()),
+                        float(rc.mean().item()),
+                    )
+
+                    # Simulate new per-rank load to preview
+                    # balancedness
+                    new_rc = new_logical_replica_count.to(
+                        load.device
+                    ).float()
+                    per_expert_load = load / new_rc.clamp(min=1)
+                    phys_load = per_expert_load.gather(
+                        dim=-1,
+                        index=new_p2l.to(load.device).long(),
+                    )
+                    per_rank_load = phys_load.reshape(
+                        phys_load.shape[0], num_gpus, -1
+                    ).sum(dim=-1)
+                    avg_rl = per_rank_load.mean(dim=-1)
+                    max_rl = per_rank_load.max(dim=-1).values
+                    predicted_balance = torch.where(
+                        max_rl > 0,
+                        avg_rl / max_rl,
+                        torch.ones_like(max_rl),
+                    )
+                    logger.info(
+                        "EPLB predicted post-rearrange "
+                        "balancedness: mean=%.4f, "
+                        "min=%.4f (layer %d), max=%.4f",
+                        float(
+                            predicted_balance.mean().item()
+                        ),
+                        float(predicted_balance.min().item()),
+                        int(
+                            predicted_balance.argmin().item()
+                        ),
+                        float(
+                            predicted_balance.max().item()
+                        ),
+                    )
+
                # Update expert weights
                rearrange_expert_weights_inplace(
                    eplb_model_state.physical_to_logical_map,
@@ -865,9 +959,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]
@@ -447,6 +447,7 @@ class EngineArgs:
    )
    disable_sliding_window: bool = ModelConfig.disable_sliding_window
    disable_cascade_attn: bool = ModelConfig.disable_cascade_attn
+    swap_space: float = CacheConfig.swap_space
    offload_backend: str = OffloadConfig.offload_backend
    cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb
    cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params")
@@ -960,6 +961,7 @@ class EngineArgs:
        cache_group.add_argument(
            "--kv-cache-memory-bytes", **cache_kwargs["kv_cache_memory_bytes"]
        )
+        cache_group.add_argument("--swap-space", **cache_kwargs["swap_space"])
        cache_group.add_argument("--kv-cache-dtype", **cache_kwargs["cache_dtype"])
        cache_group.add_argument(
            "--num-gpu-blocks-override", **cache_kwargs["num_gpu_blocks_override"]
@@ -1524,6 +1526,7 @@ class EngineArgs:
            block_size=self.block_size,
            gpu_memory_utilization=self.gpu_memory_utilization,
            kv_cache_memory_bytes=self.kv_cache_memory_bytes,
+            swap_space=self.swap_space,
            cache_dtype=resolved_cache_dtype,  # type: ignore[arg-type]
            is_attention_free=model_config.is_attention_free,
            num_gpu_blocks_override=self.num_gpu_blocks_override,
@@ -462,15 +462,10 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
    maximum per prompt.
    """

-    def __init__(
-        self,
-        model_config: ModelConfig,
-        media_io_kwargs: dict[str, dict[str, Any]] | None = None,
-    ):
+    def __init__(self, model_config: ModelConfig):
        super().__init__()

        self._model_config = model_config
-        self._media_io_kwargs = media_io_kwargs

        self._items_by_modality = defaultdict[str, list[_T]](list)
        # Track original modality for each vision_chunk item (image or video)
@@ -492,14 +487,6 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
        model_cls = get_model_cls(self.model_config)
        return cast(type[SupportsMultiModal], model_cls)

-    @property
-    def media_io_kwargs(self) -> dict[str, dict[str, Any]] | None:
-        return self._media_io_kwargs or (
-            self._model_config.multimodal_config.media_io_kwargs
-            if self._model_config.multimodal_config
-            else None
-        )
-
    @property
    def allowed_local_media_path(self):
        return self._model_config.allowed_local_media_path
@@ -782,10 +769,12 @@ class MultiModalContentParser(BaseMultiModalContentParser):
        super().__init__()

        self._tracker = tracker
+        multimodal_config = self._tracker.model_config.multimodal_config
+        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)

        self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
-            media_io_kwargs=tracker.media_io_kwargs,
+            media_io_kwargs=media_io_kwargs,
            allowed_local_media_path=tracker.allowed_local_media_path,
            allowed_media_domains=tracker.allowed_media_domains,
        )
@@ -892,9 +881,11 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):
        super().__init__()

        self._tracker = tracker
+        multimodal_config = self._tracker.model_config.multimodal_config
+        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)
        self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
-            media_io_kwargs=tracker.media_io_kwargs,
+            media_io_kwargs=media_io_kwargs,
            allowed_local_media_path=tracker.allowed_local_media_path,
            allowed_media_domains=tracker.allowed_media_domains,
        )
@@ -1539,14 +1530,13 @@ def parse_chat_messages(
    messages: list[ChatCompletionMessageParam],
    model_config: ModelConfig,
    content_format: ChatTemplateContentFormat,
-    media_io_kwargs: dict[str, dict[str, Any]] | None = None,
 ) -> tuple[
    list[ConversationMessage],
    MultiModalDataDict | None,
    MultiModalUUIDDict | None,
 ]:
    conversation: list[ConversationMessage] = []
-    mm_tracker = MultiModalItemTracker(model_config, media_io_kwargs=media_io_kwargs)
+    mm_tracker = MultiModalItemTracker(model_config)

    for msg in messages:
        sub_messages = _parse_chat_message_content(
@@ -1573,16 +1563,13 @@ async def parse_chat_messages_async(
    messages: list[ChatCompletionMessageParam],
    model_config: ModelConfig,
    content_format: ChatTemplateContentFormat,
-    media_io_kwargs: dict[str, dict[str, Any]] | None = None,
 ) -> tuple[
    list[ConversationMessage],
    MultiModalDataDict | None,
    MultiModalUUIDDict | None,
 ]:
    conversation: list[ConversationMessage] = []
-    mm_tracker = AsyncMultiModalItemTracker(
-        model_config, media_io_kwargs=media_io_kwargs
-    )
+    mm_tracker = AsyncMultiModalItemTracker(model_config)

    for msg in messages:
        sub_messages = _parse_chat_message_content(
@@ -164,6 +164,12 @@ class LLM:
            compared with using gpu_memory_utilization. Note that
            kv_cache_memory_bytes (when not-None) ignores
            gpu_memory_utilization
+        swap_space: The size (GiB) of CPU memory per GPU to use as swap space.
+            This can be used for temporarily storing the states of the requests
+            when their `best_of` sampling parameters are larger than 1. If all
+            requests will have `best_of=1`, you can safely set this to 0.
+            Noting that `best_of` is only supported in V0. Otherwise, too small
+            values may cause out-of-memory (OOM) errors.
        cpu_offload_gb: The size (GiB) of CPU memory to use for offloading
            the model weights. This virtually increases the GPU memory space
            you can use to hold the model weights, at the cost of CPU-GPU data
@@ -234,6 +240,7 @@ class LLM:
        chat_template: Path | str | None = None,
        seed: int = 0,
        gpu_memory_utilization: float = 0.9,
+        swap_space: float = 4,
        cpu_offload_gb: float = 0,
        offload_group_size: int = 0,
        offload_num_in_group: int = 1,
@@ -258,17 +265,6 @@ class LLM:
    ) -> None:
        """LLM constructor."""

-        if "swap_space" in kwargs:
-            kwargs.pop("swap_space")
-            import warnings
-
-            warnings.warn(
-                "The 'swap_space' parameter is deprecated and ignored. "
-                "It will be removed in a future version.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-
        if "disable_log_stats" not in kwargs:
            kwargs["disable_log_stats"] = True

@@ -357,6 +353,7 @@ class LLM:
            seed=seed,
            gpu_memory_utilization=gpu_memory_utilization,
            kv_cache_memory_bytes=kv_cache_memory_bytes,
+            swap_space=swap_space,
            cpu_offload_gb=cpu_offload_gb,
            offload_group_size=offload_group_size,
            offload_num_in_group=offload_num_in_group,
@@ -268,13 +268,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
            "Will be accessible by the chat template."
        ),
    )
-    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
-        default=None,
-        description=(
-            "Additional kwargs to pass to the media IO connectors, "
-            "keyed by modality. Merged with engine-level media_io_kwargs."
-        ),
-    )
    mm_processor_kwargs: dict[str, Any] | None = Field(
        default=None,
        description=("Additional kwargs to pass to the HF processor."),
@@ -373,7 +366,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
                    reasoning_effort=self.reasoning_effort,
                ),
            ),
-            media_io_kwargs=self.media_io_kwargs,
        )

    def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
@@ -900,15 +900,10 @@ class OpenAIServing:
            ),
        )

-        mm_config = self.model_config.multimodal_config
-
        tok_params = request.build_tok_params(self.model_config)
        chat_params = request.build_chat_params(
            default_template, default_template_content_format
-        ).with_defaults(
-            default_template_kwargs,
-            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
-        )
+        ).with_defaults(default_template_kwargs)

        (conversation,), (engine_prompt,) = await renderer.render_chat_async(
            [messages],
@@ -197,13 +197,6 @@ class ResponsesRequest(OpenAIBaseModel):
            "through out the inference process and return in response."
        ),
    )
-    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
-        default=None,
-        description=(
-            "Additional kwargs to pass to the media IO connectors, "
-            "keyed by modality. Merged with engine-level media_io_kwargs."
-        ),
-    )
    mm_processor_kwargs: dict[str, Any] | None = Field(
        default=None,
        description=("Additional kwargs to pass to the HF processor."),
@@ -283,7 +276,6 @@ class ResponsesRequest(OpenAIBaseModel):
                    reasoning_effort=None if reasoning is None else reasoning.effort,
                ),
            ),
-            media_io_kwargs=self.media_io_kwargs,
        )

    def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
@@ -123,15 +123,10 @@ class PoolingIOProcessor:
            ),
        )

-        mm_config = self.model_config.multimodal_config
-
        tok_params = request.build_tok_params(self.model_config)
        chat_params = request.build_chat_params(
            default_template, default_template_content_format
-        ).with_defaults(
-            default_template_kwargs,
-            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
-        )
+        ).with_defaults(default_template_kwargs)

        (conversation,), (engine_prompt,) = renderer.render_chat(
            [messages],
@@ -124,13 +124,6 @@ class ChatRequestMixin(OpenAIBaseModel):
            "Will be accessible by the chat template."
        ),
    )
-    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
-        default=None,
-        description=(
-            "Additional kwargs to pass to the media IO connectors, "
-            "keyed by modality. Merged with engine-level media_io_kwargs."
-        ),
-    )
    # --8<-- [end:chat-extra-params]

    @model_validator(mode="before")
@@ -158,7 +151,6 @@ class ChatRequestMixin(OpenAIBaseModel):
                    continue_final_message=self.continue_final_message,
                ),
            ),
-            media_io_kwargs=self.media_io_kwargs,
        )


@@ -100,13 +100,6 @@ class TokenizeChatRequest(OpenAIBaseModel):
            "Will be accessible by the chat template."
        ),
    )
-    media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
-        default=None,
-        description=(
-            "Additional kwargs to pass to the media IO connectors, "
-            "keyed by modality. Merged with engine-level media_io_kwargs."
-        ),
-    )
    mm_processor_kwargs: dict[str, Any] | None = Field(
        default=None,
        description="Additional kwargs to pass to the HF processor.",
@@ -141,7 +134,6 @@ class TokenizeChatRequest(OpenAIBaseModel):
                    continue_final_message=self.continue_final_message,
                ),
            ),
-            media_io_kwargs=self.media_io_kwargs,
        )

    def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
@@ -905,10 +905,6 @@ def unified_mla_kv_cache_update(
    the data dependency between them to ensure torch.compile preserves ordering.
    """
    forward_context = get_forward_context()
-    if forward_context.attn_metadata is None:
-        # Dummy/profile forwards should not update live KV cache pages.
-        return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)
-
    attn_layer = forward_context.no_compile_layers[layer_name]
    kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]

@@ -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:
@@ -756,7 +756,7 @@ direct_register_custom_op(
 )


-class DeepSeekV2FusedQkvAProjLinear(MergedColumnParallelLinear):
+class DeepSeekV2FusedQkvAProj(MergedColumnParallelLinear):
    def __init__(
        self,
        input_size: int,
@@ -848,7 +848,7 @@ class DeepseekV2MLAAttention(nn.Module):
        self.max_position_embeddings = max_position_embeddings

        if self.q_lora_rank is not None:
-            self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
+            self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProj(
                self.hidden_size,
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                quant_config=quant_config,
@@ -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))
@@ -83,17 +83,11 @@ def extract_audio_from_video_bytes(


 class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]):
-    """Configuration values can be user-provided either by --media-io-kwargs or
-    by the runtime API field "media_io_kwargs". Ensure proper validation and
-    error handling.
-    """
-
    def __init__(self, **kwargs) -> None:
        super().__init__()

        # `kwargs` contains custom arguments from
-        # --media-io-kwargs for this modality, merged with
-        # per-request runtime media_io_kwargs via merge_kwargs().
+        # --media-io-kwargs for this modality.
        # They can be passed to the underlying
        # media loaders (e.g. custom implementations)
        # for flexible control.
@@ -128,11 +122,6 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]):


 class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):
-    """Configuration values can be user-provided either by --media-io-kwargs or
-    by the runtime API field "media_io_kwargs". Ensure proper validation and
-    error handling.
-    """
-
    def __init__(self) -> None:
        super().__init__()

@@ -44,28 +44,6 @@ class MediaWithBytes(Generic[_T]):


 class MediaIO(ABC, Generic[_T]):
-    """Configuration values can be user-provided either by --media-io-kwargs or
-    by the runtime API field "media_io_kwargs". Ensure proper validation and
-    error handling.
-    """
-
-    @classmethod
-    def merge_kwargs(
-        cls,
-        default_kwargs: dict[str, Any] | None,
-        runtime_kwargs: dict[str, Any] | None,
-    ) -> dict[str, Any]:
-        """Merge config-level kwargs and request-level kwargs.
-
-        By default this performs a shallow merge where runtime kwargs override
-        keys in default kwargs. Subclasses may override to apply modality-
-        specific behavior.
-        """
-        merged = dict(default_kwargs or {})
-        if runtime_kwargs:
-            merged.update(runtime_kwargs)
-        return merged
-
    @abstractmethod
    def load_bytes(self, data: bytes) -> _T:
        raise NotImplementedError
@@ -32,43 +32,9 @@ atexit.register(global_thread_pool.shutdown)

 MEDIA_CONNECTOR_REGISTRY = ExtensionManager()

-MODALITY_IO_MAP: dict[str, type[MediaIO]] = {
-    "audio": AudioMediaIO,
-    "image": ImageMediaIO,
-    "video": VideoMediaIO,
-}
-
-
-def merge_media_io_kwargs(
-    defaults: dict[str, dict[str, Any]] | None,
-    overrides: dict[str, dict[str, Any]] | None,
-) -> dict[str, dict[str, Any]] | None:
-    """Merge config-level and per-request media_io_kwargs per modality.
-
-    Each modality key is merged using the corresponding MediaIO subclass's
-    ``merge_kwargs``, which may apply modality-specific logic (e.g.
-    VideoMediaIO clears cross-dependent fps/num_frames fields).
-    """
-    if not defaults and not overrides:
-        return None
-    all_keys = set(defaults or {}) | set(overrides or {})
-    merged = {}
-    for key in all_keys:
-        io_cls = MODALITY_IO_MAP.get(key, MediaIO)
-        merged[key] = io_cls.merge_kwargs(
-            (defaults or {}).get(key),
-            (overrides or {}).get(key),
-        )
-    return merged or None
-

@MEDIA_CONNECTOR_REGISTRY.register("http")
 class MediaConnector:
-    """Configuration values can be user-provided either by --media-io-kwargs or
-    by the runtime API field "media_io_kwargs". Ensure proper validation and
-    error handling.
-    """
-
    def __init__(
        self,
        media_io_kwargs: dict[str, dict[str, Any]] | None = None,
@@ -15,18 +15,12 @@ from .base import MediaIO, MediaWithBytes


 class ImageMediaIO(MediaIO[Image.Image]):
-    """Configuration values can be user-provided either by --media-io-kwargs or
-    by the runtime API field "media_io_kwargs". Ensure proper validation and
-    error handling.
-    """
-
    def __init__(self, image_mode: str = "RGB", **kwargs) -> None:
        super().__init__()

        self.image_mode = image_mode
        # `kwargs` contains custom arguments from
-        # --media-io-kwargs for this modality, merged with
-        # per-request runtime media_io_kwargs via merge_kwargs().
+        # --media-io-kwargs for this modality.
        # They can be passed to the underlying
        # media loaders (e.g. custom implementations)
        # for flexible control.
@@ -94,13 +88,6 @@ class ImageMediaIO(MediaIO[Image.Image]):


 class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
-    """Image embedding MediaIO implementation.
-
-    Configuration values can be user-provided either by --media-io-kwargs or
-    by the runtime API field "media_io_kwargs". Ensure proper validation and
-    error handling.
-    """
-
    def __init__(self) -> None:
        super().__init__()

@@ -17,28 +17,6 @@ from .image import ImageMediaIO


 class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]):
-    """Configuration values can be user-provided either by --media-io-kwargs or
-    by the runtime API field "media_io_kwargs". Ensure proper validation and
-    error handling.
-    """
-
-    @classmethod
-    def merge_kwargs(
-        cls,
-        default_kwargs: dict[str, Any] | None,
-        runtime_kwargs: dict[str, Any] | None,
-    ) -> dict[str, Any]:
-        merged = super().merge_kwargs(default_kwargs, runtime_kwargs)
-        # fps and num_frames interact with each other, so if either is
-        # overridden at request time, wipe the other from defaults to
-        # avoid unintuitive cross-field interactions.
-        if runtime_kwargs:
-            if "num_frames" in runtime_kwargs and "fps" not in runtime_kwargs:
-                merged.pop("fps", None)
-            elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs:
-                merged.pop("num_frames", None)
-        return merged
-
    def __init__(
        self,
        image_io: ImageMediaIO,
@@ -50,8 +28,7 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]):
        self.image_io = image_io
        self.num_frames = num_frames
        # `kwargs` contains custom arguments from
-        # --media-io-kwargs for this modality, merged with
-        # per-request runtime media_io_kwargs via merge_kwargs().
+        # --media-io-kwargs for this modality.
        # They can be passed to the underlying
        # media loaders (e.g. custom implementations)
        # for flexible control.
@@ -22,6 +22,13 @@ _PARSERS_TO_REGISTER = {
    ),
 }

+# Register lazy parsers
+ParserManager.register_lazy_module(
+    name="minimax_m2",
+    module_path="vllm.parser.minimax_m2_parser",
+    class_name="MiniMaxM2Parser",
+)
+

 def register_lazy_parsers():
    for name, (file_name, class_name) in _PARSERS_TO_REGISTER.items():
@@ -49,7 +49,6 @@ class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):
            messages,
            self.model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = tokenizer.apply_chat_template(
@@ -76,7 +75,6 @@ class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):
            messages,
            self.model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = tokenizer.apply_chat_template(
@@ -49,7 +49,6 @@ class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):
            messages,
            self.model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = tokenizer.apply_chat_template(
@@ -76,7 +75,6 @@ class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):
            messages,
            self.model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = tokenizer.apply_chat_template(
@@ -635,7 +635,6 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
                tokenizer=tokenizer,
                model_config=model_config,
            ),
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = safe_apply_chat_template(
@@ -690,7 +689,6 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
                tokenizer=tokenizer,
                model_config=model_config,
            ),
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = safe_apply_chat_template(
@@ -90,7 +90,6 @@ class MistralRenderer(BaseRenderer[MistralTokenizer]):
            messages,
            self.model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = safe_apply_chat_template(
@@ -117,7 +116,6 @@ class MistralRenderer(BaseRenderer[MistralTokenizer]):
            messages,
            self.model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt_raw = await self._apply_chat_template_async(
@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, TypeVar
 from vllm.exceptions import VLLMValidationError
 from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
 from vllm.logger import init_logger
-from vllm.multimodal.media.connector import merge_media_io_kwargs
 from vllm.tokenizers import TokenizerLike
 from vllm.utils.import_utils import LazyLoader

@@ -53,15 +52,8 @@ class ChatParams:
    chat_template_kwargs: dict[str, Any] = field(default_factory=dict)
    """The kwargs to pass to the chat template."""

-    media_io_kwargs: dict[str, dict[str, Any]] | None = None
-    """Per-modality kwargs for media I/O (loading/decoding images, videos, etc.)."""
-
-    def with_defaults(
-        self,
-        default_chat_template_kwargs: dict[str, Any] | None = None,
-        default_media_io_kwargs: dict[str, dict[str, Any]] | None = None,
-    ):
-        if not default_chat_template_kwargs and not default_media_io_kwargs:
+    def with_defaults(self, default_chat_template_kwargs: dict[str, Any] | None):
+        if not default_chat_template_kwargs:
            return self

        return ChatParams(
@@ -71,10 +63,6 @@ class ChatParams:
                default_chat_template_kwargs,
                self.chat_template_kwargs,
            ),
-            media_io_kwargs=merge_media_io_kwargs(
-                default_media_io_kwargs,
-                self.media_io_kwargs,
-            ),
        )

    def get_apply_chat_template_kwargs(self) -> dict[str, Any]:
@@ -43,7 +43,6 @@ class TerratorchRenderer(BaseRenderer):
            messages,
            model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt = parse_dec_only_prompt([1])  # Dummy token IDs
@@ -65,7 +64,6 @@ class TerratorchRenderer(BaseRenderer):
            messages,
            model_config,
            content_format="string",
-            media_io_kwargs=params.media_io_kwargs,
        )

        prompt = parse_dec_only_prompt([1])  # Dummy token IDs
@@ -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):
@@ -24,10 +24,7 @@ from transformers.utils import CONFIG_NAME as HF_CONFIG_NAME
 from vllm import envs
 from vllm.logger import init_logger
 from vllm.transformers_utils.repo_utils import is_mistral_model_repo
-from vllm.transformers_utils.utils import (
-    parse_safetensors_file_metadata,
-    without_trust_remote_code,
-)
+from vllm.transformers_utils.utils import parse_safetensors_file_metadata

 from .config_parser_base import ConfigParserBase
 from .gguf_utils import (
@@ -143,12 +140,11 @@ class HFConfigParser(ConfigParserBase):
        **kwargs,
    ) -> tuple[dict, PretrainedConfig]:
        kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE
-        trust_remote_code |= kwargs.get("trust_remote_code", False)
-        kwargs = without_trust_remote_code(kwargs)
        config_dict, _ = PretrainedConfig.get_config_dict(
            model,
            revision=revision,
            code_revision=code_revision,
+            trust_remote_code=trust_remote_code,
            **kwargs,
        )
        # Use custom model class if it's in our registry
@@ -229,7 +225,7 @@ class MistralConfigParser(ConfigParserBase):
                model,
                revision=revision,
                code_revision=code_revision,
-                **without_trust_remote_code(kwargs),
+                **kwargs,
            )
        except OSError:  # Not found
            hf_config_dict = {}
@@ -525,7 +521,8 @@ def maybe_override_with_speculators(
    config_dict, _ = PretrainedConfig.get_config_dict(
        model if gguf_model_repo is None else gguf_model_repo,
        revision=revision,
-        **without_trust_remote_code(kwargs),
+        trust_remote_code=trust_remote_code,
+        **kwargs,
    )
    speculators_config = config_dict.get("speculators_config")

@@ -5,8 +5,6 @@ import os

 from transformers import AutoConfig, DeepseekV2Config, PretrainedConfig

-from vllm.transformers_utils.utils import without_trust_remote_code
-

 class EAGLEConfig(PretrainedConfig):
    model_type = "eagle"
@@ -81,7 +79,7 @@ class EAGLEConfig(PretrainedConfig):
        **kwargs,
    ) -> "EAGLEConfig":
        config_dict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
+            pretrained_model_name_or_path, **kwargs
        )
        return cls.from_dict(config_dict, **kwargs)

@@ -7,8 +7,6 @@ import os

 from transformers import PretrainedConfig

-from vllm.transformers_utils.utils import without_trust_remote_code
-

 class ExtractHiddenStatesConfig(PretrainedConfig):
    model_type = "extract_hidden_states"
@@ -44,7 +42,7 @@ class ExtractHiddenStatesConfig(PretrainedConfig):
        **kwargs,
    ) -> "ExtractHiddenStatesConfig":
        config_dict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
+            pretrained_model_name_or_path, **kwargs
        )
        return cls.from_dict(config_dict, **kwargs)

@@ -5,8 +5,6 @@ import os

 from transformers import PretrainedConfig

-from vllm.transformers_utils.utils import without_trust_remote_code
-

 class MedusaConfig(PretrainedConfig):
    model_type = "medusa"
@@ -44,7 +42,7 @@ class MedusaConfig(PretrainedConfig):
        **kwargs,
    ) -> "MedusaConfig":
        config_dict, kwargs = cls.get_config_dict(
-            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
+            pretrained_model_name_or_path, **kwargs
        )
        for k in list(config_dict.keys()):
            if "num" in k:
@@ -11,8 +11,6 @@ from vllm.transformers_utils.configs.speculators.algos import (

 __all__ = ["SpeculatorsConfig"]

-from vllm.transformers_utils.utils import without_trust_remote_code
-

 class SpeculatorsConfig(PretrainedConfig):
    model_type = "speculators"
@@ -24,9 +22,7 @@ class SpeculatorsConfig(PretrainedConfig):
        **kwargs,
    ) -> "SpeculatorsConfig":
        """Load speculators Eagle config and convert to vLLM format."""
-        config_dict, _ = cls.get_config_dict(
-            pretrained_model_name_or_path, **without_trust_remote_code(kwargs)
-        )
+        config_dict, _ = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)

        vllm_config = cls.extract_transformers_pre_trained_config(config_dict)
        return cls(**vllm_config)
@@ -27,13 +27,6 @@ def is_cloud_storage(model_or_path: str) -> bool:
    return is_s3(model_or_path) or is_gcs(model_or_path)


-def without_trust_remote_code(kwargs: dict[str, Any]) -> dict[str, Any]:
-    """Return kwargs without trust_remote_code without modifying original dict."""
-    if "trust_remote_code" not in kwargs:
-        return kwargs
-    return {k: v for k, v in kwargs.items() if k != "trust_remote_code"}
-
-
 def modelscope_list_repo_files(
    repo_id: str,
    revision: str | None = None,
@@ -1152,10 +1152,11 @@ class AiterFlashAttentionImpl(AttentionImpl):
                decode_max_query_len = attn_metadata.decode_metadata.max_query_len

                # Use unified_attention for speculative decoding (multi-token)
-                if decode_max_query_len > 1:
+                # or when sliding window is enabled
+                if self.sliding_window[0] != -1 or decode_max_query_len > 1:
                    assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
-                        "Shuffle KV cache layout is not supported with "
-                        "speculative decoding (multi-token decode)."
+                        "Shuffle KV cache layout is not supported with sliding "
+                        "window or speculative decoding (multi-token decode)."
                    )
                    from aiter.ops.triton.unified_attention import (
                        unified_attention,
@@ -567,7 +567,10 @@ class MPClient(EngineCoreClient):
                )

                with launch_core_engines(
-                    vllm_config, executor_class, log_stats, addresses
+                    vllm_config,
+                    executor_class,
+                    log_stats,
+                    addresses,
                ) as (engine_manager, coordinator, addresses):
                    self.resources.coordinator = coordinator
                    self.resources.engine_manager = engine_manager
@@ -635,10 +638,10 @@ class MPClient(EngineCoreClient):

    def shutdown(self, timeout: float | None = None) -> None:
        """Shutdown engine manager under timeout and clean up resources."""
-        if self._finalizer.detach() is not None:
-            if self.resources.engine_manager is not None:
-                self.resources.engine_manager.shutdown(timeout=timeout)
-            self.resources()
+        self._finalizer.detach()
+        if self.resources.engine_manager is not None:
+            self.resources.engine_manager.shutdown(timeout=timeout)
+        self.resources()

    def _format_exception(self, e: Exception) -> Exception:
        """If errored, use EngineDeadError so root cause is clear."""
@@ -682,7 +685,7 @@ class MPClient(EngineCoreClient):
            sentinels = [proc.sentinel for proc in engine_processes]
            died = multiprocessing.connection.wait(sentinels)
            _self = self_ref()
-            if not _self or not _self._finalizer.alive or _self.resources.engine_dead:
+            if not _self or _self.resources.engine_dead:
                return
            _self.resources.engine_dead = True
            proc_name = next(
@@ -1,9 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-from collections.abc import Callable
-from typing import Any
-
 import numpy as np
 import torch

@@ -20,14 +17,9 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner


@torch.inference_mode()
-def warmup_kernels(
-    model_runner: GPUModelRunner,
-    worker_execute_model: Callable[[SchedulerOutput], Any],
-    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
-) -> None:
+def warmup_kernels(model_runner: GPUModelRunner) -> None:
    """Run two execute_model + sample_tokens iterations to JIT compile
-    triton kernels. We must call the provided worker's execute_model for
-    pipeline parallel coordination.
+    triton kernels.

    The first iteration simulates a prefill with requests of 2 prompt
    tokens each. The second iteration simulates a decode step with all
@@ -91,7 +83,7 @@ def warmup_kernels(

    # Disable KV connector for warmup run.
    model_runner.kv_connector.set_disabled(True)
-    worker_execute_model(prefill_output)
+    model_runner.execute_model(prefill_output)

    if not model_runner.is_pooling_model:
        # Warm up sampler and perform a decode step for non-pooling models.
@@ -109,7 +101,7 @@ def warmup_kernels(
                structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask
            )

-        worker_sample_tokens(grammar_output)
+        model_runner.sample_tokens(grammar_output)

        # Step 2: Decode all requests with 1 token each.
        cached_req_data = CachedRequestData.make_empty()
@@ -128,12 +120,12 @@ def warmup_kernels(
        decode_output.total_num_scheduled_tokens = num_reqs
        decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups

-        worker_execute_model(decode_output)
-        worker_sample_tokens(None)
+        model_runner.execute_model(decode_output)
+        model_runner.sample_tokens(None)

    # Clean up - process finish_req_ids.
    cleanup_output = SchedulerOutput.make_empty()
    cleanup_output.finished_req_ids = set(req_ids)
-    worker_execute_model(cleanup_output)
+    model_runner.execute_model(cleanup_output)
    model_runner.kv_connector.set_disabled(False)
    torch.accelerator.synchronize()
@@ -584,7 +584,7 @@ class Worker(WorkerBase):

        if self.use_v2_model_runner:
            # V2: Run full execute_model + sample_tokens to JIT compile triton kernels.
-            warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)
+            warmup_kernels(self.model_runner)
        elif get_pp_group().is_last_rank:
            # V1: Warm up sampler and preallocate memory buffer for logits and other
            # sampling related tensors of max possible shape to avoid memory