[wip] support PD seperate mode for fault-tolerance

.buildkite/release-pipeline.yaml

@@ -978,14 +978,22 @@ steps:
       - step: build-cpu-release-image-arm64
         allow_failure: true
 
-  - label: "Publish release images to DockerHub"
+  - label: "Publish {{matrix}} release images to DockerHub"
     depends_on:
       - block-publish-release-images
     key: publish-release-images-dockerhub
     agents:
       queue: small_cpu_queue_release
     commands:
-      - "bash .buildkite/scripts/publish-release-images.sh"
+      - "bash .buildkite/scripts/publish-release-images.sh {{matrix}}"
+    matrix:
+      - "cuda-13-0"
+      - "cuda-12-9"
+      - "cuda-13-0-ubuntu-24-04"
+      - "cuda-12-9-ubuntu-24-04"
+      - "rocm"
+      - "xpu"
+      - "cpu"
     plugins:
       - docker-login#v3.0.0:
           username: vllmbot

cmake/external_projects/deepgemm.cmake

@@ -30,7 +30,7 @@ else()
   # Keep in sync with tools/install_deepgemm.sh
   set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git")
   # TODO: switch to nv_dev branch after it support situ
-  set(_DEEPGEMM_UPSTREAM_TAG "f5a76426fa084087169693fd0cd815223576d6e9")
+  set(_DEEPGEMM_UPSTREAM_TAG "e21c821f39a2056d68067a466c64ddc942200106")
 
   set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
   if(NOT _deepgemm_fc_root)

csrc/cpu/sgl-kernels/fla.cpp

@@ -101,7 +101,7 @@ struct l2norm_kernel {
     using bVec = at::vec::Vectorized<scalar_t>;
     using fVec = at::vec::Vectorized<float>;
     constexpr int bVecSize = bVec::size();
-    constexpr float scale = 1.f / std::sqrt(static_cast<float>(D));
+    const float scale = 1.f / std::sqrt(static_cast<float>(D));
 
     fVec sum_fvec0(0.f), sum_fvec1(0.f);
     int d = 0;
@@ -155,7 +155,7 @@ struct l2norm_kernel<at::BFloat16, D, has_scale> {
     __m512bh va[COLS];
     __m512 vrscale;
 
-    constexpr float scale = 1.f / std::sqrt(D);
+    const float scale = 1.f / std::sqrt(D);
     __m512 vscale = _mm512_set1_ps(scale);
 
     // step 1: load input and do reduce with avx512-bf16

docker/Dockerfile

@@ -811,7 +811,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.15.post1
+ARG FLASHINFER_VERSION=0.6.16.post3
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -949,7 +949,8 @@ COPY requirements/test/cuda.txt requirements/test/cuda.txt
 COPY requirements/dev.txt requirements/dev.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt
-RUN --mount=type=cache,target=/opt/uv/cache \
+# Source-built wheels in this cache must match the Ubuntu 22.04 ABI.
+RUN --mount=type=cache,id=uv-v027-ubuntu2204,target=/opt/uv/cache \
     CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
     if [ "$CUDA_MAJOR" -ge 12 ]; then \
         if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
@@ -971,6 +972,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \
             && uv pip install --system -r requirements/dev.txt \
             --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
         fi \
+        && if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \
+            python3 -c "from arctic_inference.suffix_decoding import _C"; \
+        fi; \
     fi
 
 # install development dependencies (for testing)

docker/versions.json

@@ -71,7 +71,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.15.post1"
+      "default": "0.6.16.post3"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/design/moe_kernel_features.md

@@ -32,7 +32,7 @@ th {
 
 | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
 | ------- | ------------------ | ------------ | ------------- | ----- | --------------------- | --------- |
-| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.runner.MoERunner] |
+| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [`MoERunner`][vllm.model_executor.layers.fused_moe.runner.moe_runner.MoERunner] |
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
 | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |

docs/models/supported_models.md

@@ -364,7 +364,7 @@ th {
 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
 | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. | ✅︎ | ✅︎ |
 | `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ |
-| `ExaoneMoEForCausalLM` | K-EXAONE | `LGAI-EXAONE/K-EXAONE-236B-A23B`, etc. | | |
+| `ExaoneMoeForCausalLM` | K-EXAONE, K-EXAONE-2 | `LGAI-EXAONE/K-EXAONE-236B-A23B`, `LGAI-EXAONE/K-EXAONE-2.0-750B-A37B`, etc. | | |
 | `Exaone4ForCausalLM` | EXAONE-4 | `LGAI-EXAONE/EXAONE-4.0-32B`, etc. | ✅︎ | ✅︎ |
 | `Fairseq2LlamaForCausalLM` | Llama (fairseq2 format) | `mgleize/fairseq2-dummy-Llama-3.2-1B`, etc. | ✅︎ | ✅︎ |
 | `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ |

requirements/cuda.txt

@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
 # flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
 # install_requires so the published wheel does not carry an unresolvable pin
 --extra-index-url https://flashinfer.ai/whl/
-flashinfer-python==0.6.15.post1
-flashinfer-cubin==0.6.15.post1
+flashinfer-python==0.6.16.post3
+flashinfer-cubin==0.6.16.post3
 apache-tvm-ffi==0.1.11
 tilelang==0.1.12
 nvidia-cudnn-frontend>=1.19.1
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
 
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
 nvidia-cutlass-dsl[cu13]==4.6.0
-quack-kernels>=0.6.1  # Required for CUTLASS DSL 4.6 by MSA
+quack-kernels==0.6.1  # Required for CUTLASS DSL 4.6 by MSA
 
 # Tokenspeed_MLA for faster mla with spec decode
 tokenspeed-mla==0.1.8; platform_system == "Linux"

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml

@@ -13,5 +13,6 @@ server_args: >-
   --enable-expert-parallel
   --kv-cache-dtype fp8
 env:
+  PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
   VLLM_LOGGING_LEVEL: "DEBUG"
   VLLM_USE_V2_MODEL_RUNNER: "1"

tests/kernels/mamba/test_ssu_dispatch.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from unittest.mock import Mock
+
 import pytest
 import torch
 
-from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     FlashInferSSUBackend,
     TritonSSUBackend,
@@ -69,6 +71,48 @@ def test_flashinfer_backend_init():
     assert backend.name == "flashinfer"
 
 
+@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
+@pytest.mark.parametrize(
+    ("algorithm", "expected"),
+    [
+        (None, "auto"),
+        ("auto", "auto"),
+        ("simple", "simple"),
+        ("vertical", "vertical"),
+        ("horizontal", "horizontal"),
+    ],
+)
+def test_flashinfer_forwards_ssu_algorithm(
+    algorithm: MambaSSUAlgorithm | None,
+    expected: MambaSSUAlgorithm,
+    monkeypatch,
+):
+    import flashinfer.mamba
+
+    kernel = Mock()
+    monkeypatch.setattr(flashinfer.mamba, "selective_state_update", kernel)
+    backend = FlashInferSSUBackend(
+        MambaConfig(
+            backend=MambaBackendEnum.FLASHINFER,
+            ssu_algorithm=algorithm,
+        )
+    )
+
+    tensor = torch.empty(1)
+    backend(
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+    )
+
+    assert kernel.call_args.kwargs["algorithm"] == expected
+
+
 def test_uninitialized_backend_raises():
     import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
 

tests/models/kimi_k3/test_sequence_parallel.py

@@ -10,6 +10,7 @@
 from torch import nn
 
 from vllm.config import ParallelConfig
+from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.models.common.ops import sequence_parallel as sp_ops
 from vllm.models.kimi_k3.nvidia import model as kimi_model
 from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp
@@ -265,6 +266,84 @@ def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch):
     torch.testing.assert_close(final_norm.call_args.args[0], expected_hidden_states)
 
 
+@pytest.mark.parametrize(
+    ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"),
+    [
+        (True, True, True, 8, True),
+        (False, True, True, 8, False),  # opt-in only
+        (True, False, True, 8, False),  # replication only exists under SP
+        (True, True, False, 8, False),  # FusedMoE path owns the reduction
+        (True, True, True, 1, False),  # nothing to shard
+        (True, True, True, 5, False),  # 6144 % 5 -- would fail divide()
+    ],
+)
+def test_shard_sequence_parallel_mlp_gating(
+    monkeypatch,
+    enabled: bool,
+    use_sequence_parallel: bool,
+    eligible: bool,
+    tp_size: int,
+    expected: bool,
+):
+    monkeypatch.setattr(kimi_model.envs, "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", enabled)
+    monkeypatch.setattr(
+        kimi_model, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+
+    assert (
+        kimi_model.shard_sequence_parallel_mlp(
+            hidden_size=7168,
+            intermediate_size=6144,
+            use_sequence_parallel=use_sequence_parallel,
+            eligible=eligible,
+        )
+        is expected
+    )
+
+
+def test_sharded_sequence_parallel_mlp_matches_replicated(default_vllm_config):
+    """Sharded SP MLP must reproduce the replicated result for every token.
+
+    Each rank owns a *disjoint* token shard, so a weight shard alone cannot
+    finish a rank's own tokens: the ranks must gather the full token set,
+    compute partial sums over their intermediate shard, and reduce-scatter.
+    Splicing per-rank feature slices together instead silently mixes different
+    tokens and produces plausible-looking garbage.
+    """
+    tp_size, hidden, intermediate, tokens_per_rank = 4, 16, 12, 3
+    torch.manual_seed(0)
+    num_tokens = tp_size * tokens_per_rank
+    x = torch.randn(num_tokens, hidden)
+    gate_weight = torch.randn(intermediate, hidden)
+    up_weight = torch.randn(intermediate, hidden)
+    down_weight = torch.randn(hidden, intermediate)
+    act_fn = SiluAndMul()
+
+    replicated = act_fn(x @ torch.cat([gate_weight, up_weight]).T) @ down_weight.T
+
+    shard = intermediate // tp_size
+    # Every rank all-gathers the full token set, then computes its partial.
+    partials = [
+        act_fn(
+            x
+            @ torch.cat(
+                [
+                    gate_weight[r * shard : (r + 1) * shard],
+                    up_weight[r * shard : (r + 1) * shard],
+                ]
+            ).T
+        )
+        @ down_weight[:, r * shard : (r + 1) * shard].T
+        for r in range(tp_size)
+    ]
+    reduced = torch.stack(partials).sum(0)
+    # Reduce-scatter: rank r keeps only its own token shard.
+    for r in range(tp_size):
+        mine = reduced[r * tokens_per_rank : (r + 1) * tokens_per_rank]
+        expected = replicated[r * tokens_per_rank : (r + 1) * tokens_per_rank]
+        torch.testing.assert_close(mine, expected, atol=1e-5, rtol=1e-5)
+
+
 def test_sp_all_gather_uses_custom_kernel(monkeypatch):
     hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)
     expected = torch.cat([hidden_states, hidden_states])

tests/models/registry.py

@@ -267,7 +267,7 @@ def check_available_online(
         "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct", trust_remote_code=True
     ),
     "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"),
-    "ExaoneMoEForCausalLM": _HfExamplesInfo(
+    "ExaoneMoeForCausalLM": _HfExamplesInfo(
         "LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.1.0"
     ),
     "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"),

tests/quantization/test_auto_round.py

@@ -496,7 +496,7 @@ class DummyKernel:
     monkeypatch.setattr(
         "vllm.model_executor.layers.quantization.inc.schemes."
         "inc_mxfp4_linear.init_mxfp4_linear_kernel",
-        lambda: DummyKernel(),
+        lambda **kwargs: DummyKernel(),
     )
 
     from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import (  # noqa: E501
@@ -614,7 +614,7 @@ def process_weights_after_loading(self, layer) -> None:
     monkeypatch.setattr(
         "vllm.model_executor.layers.quantization.inc.schemes."
         "inc_mxfp4_linear.init_mxfp4_linear_kernel",
-        lambda: DummyKernel(),
+        lambda **kwargs: DummyKernel(),
     )
     monkeypatch.setattr(
         "vllm.model_executor.parameter.get_tensor_model_parallel_rank",

tests/quantization/test_modelopt.py

@@ -22,6 +22,7 @@
 from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.modelopt import (
     ModelOptFp8Config,
+    ModelOptFp8LinearMethod,
     ModelOptMixedPrecisionConfig,
     ModelOptMxFp8Config,
     ModelOptNvFp4Config,
@@ -114,6 +115,28 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head():
     assert isinstance(method, ModelOptNvFp4LinearMethod)
 
 
+def test_modelopt_fp8_updates_weight_dims_after_transpose():
+    layer = torch.nn.Module()
+    layer.register_parameter(
+        "weight", torch.nn.Parameter(torch.empty(3, 2), requires_grad=False)
+    )
+    layer.register_parameter(
+        "weight_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False)
+    )
+    layer.register_parameter(
+        "input_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False)
+    )
+
+    method = ModelOptFp8LinearMethod.__new__(ModelOptFp8LinearMethod)
+    method.fp8_linear = Mock()
+    method.process_weights_after_loading(layer)
+
+    assert layer.weight.shape == (2, 3)
+    assert layer.weight.input_dim == 0
+    assert layer.weight.output_dim == 1
+    method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer)
+
+
 def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():
     config = ModelOptNvFp4Config(
         is_checkpoint_nvfp4_serialized=True,