From 7c9e7b1b014146e1e88288d874d185a15332fb4b Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 10:52:28 +0200
Subject: [PATCH 01/25] test(glm5next): cover direct TP3 geometry

(cherry picked from commit 2fe594deae43ee4d2463e64d94265dedef0b25fe)
---
 tests/config/test_glm53_tp3_geometry.py | 377 ++++++++++++++++++++++++
 1 file changed, 377 insertions(+)
 create mode 100644 tests/config/test_glm53_tp3_geometry.py

diff --git a/tests/config/test_glm53_tp3_geometry.py b/tests/config/test_glm53_tp3_geometry.py
new file mode 100644
index 000000000000..ae6feaf4ece6
--- /dev/null
+++ b/tests/config/test_glm53_tp3_geometry.py
@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from copy import deepcopy
+from types import SimpleNamespace
+from typing import Any, cast
+
+import pytest
+
+from vllm.config import ParallelConfig
+from vllm.config.speculative import SpeculativeConfig
+from vllm.transformers_utils.configs.glm53_tp3 import (
+    apply_glm53_tp3_draft_geometry,
+    apply_glm53_tp3_target_geometry,
+)
+
+
+class FakeGlm53ModelConfig:
+    """Device-free stand-in for the released GLM-5.3 checkpoint config."""
+
+    def __init__(
+        self,
+        *,
+        architecture: str = "Glm5NextForConditionalGeneration",
+        model_type: str = "glm5_next",
+        mm_encoder_tp_mode: str = "weights",
+    ) -> None:
+        self.hf_text_config = SimpleNamespace(
+            model_type=model_type,
+            architectures=[architecture],
+            hidden_size=4096,
+            num_attention_heads=64,
+            num_key_value_heads=64,
+            linear_num_heads=64,
+            linear_attn_config={"num_heads": 64},
+            moe_intermediate_size=2048,
+            n_routed_experts=288,
+            n_shared_experts=1,
+            vocab_size=154880,
+        )
+        self.hf_config = SimpleNamespace(
+            model_type=model_type,
+            architectures=[architecture],
+            text_config=self.hf_text_config,
+            vision_config=SimpleNamespace(
+                hidden_size=1024,
+                num_heads=16,
+                intermediate_size=4096,
+                projection_intermediate_size=10240,
+            ),
+        )
+        self.multimodal_config = SimpleNamespace(
+            mm_encoder_tp_mode=mm_encoder_tp_mode
+        )
+        self.model_arch_config = self.get_model_arch_config()
+
+    def get_model_arch_config(self) -> SimpleNamespace:
+        return SimpleNamespace(
+            total_num_attention_heads=self.hf_text_config.num_attention_heads,
+            total_num_kv_heads=self.hf_text_config.num_key_value_heads,
+            vocab_size=self.hf_text_config.vocab_size,
+        )
+
+
+class FakeDFlashModelConfig:
+    """Device-free stand-in for the released dense DFlash checkpoint config."""
+
+    def __init__(self) -> None:
+        self.hf_text_config = SimpleNamespace(
+            model_type="qwen3",
+            architectures=["DFlash2DraftModel"],
+            num_attention_heads=32,
+            num_key_value_heads=8,
+            vocab_size=154880,
+        )
+        self.hf_config = self.hf_text_config
+        self.model_arch_config = self.get_model_arch_config()
+
+    def get_model_arch_config(self) -> SimpleNamespace:
+        return SimpleNamespace(
+            total_num_attention_heads=self.hf_text_config.num_attention_heads,
+            total_num_kv_heads=self.hf_text_config.num_key_value_heads,
+            vocab_size=self.hf_text_config.vocab_size,
+        )
+
+
+@pytest.fixture
+def glm53_model_config() -> FakeGlm53ModelConfig:
+    return FakeGlm53ModelConfig()
+
+
+@pytest.fixture
+def tp3_ep_parallel_config() -> ParallelConfig:
+    return ParallelConfig(tensor_parallel_size=3, enable_expert_parallel=True)
+
+
+def _snapshot(model_config: object) -> dict[str, Any]:
+    return deepcopy(vars(model_config))
+
+
+def test_glm53_tp3_target_geometry_uses_parallel_config_and_preserves_logical_axes(
+    glm53_model_config: FakeGlm53ModelConfig,
+    tp3_ep_parallel_config: ParallelConfig,
+) -> None:
+    applied = apply_glm53_tp3_target_geometry(
+        cast(Any, glm53_model_config), tp3_ep_parallel_config
+    )
+
+    assert applied
+    text_config = glm53_model_config.hf_text_config
+    assert (
+        text_config.num_attention_heads,
+        text_config.num_key_value_heads,
+        text_config.linear_num_heads,
+        text_config.linear_attn_config["num_heads"],
+    ) == (72, 72, 66, 66)
+    assert (
+        text_config.original_num_attention_heads,
+        text_config.original_num_key_value_heads,
+        text_config.original_linear_num_heads,
+    ) == (64, 64, 64)
+
+    # Routed experts keep the checkpoint width under EP. Only the replicated
+    # shared expert gets a physical TP3 storage width.
+    assert tp3_ep_parallel_config.enable_expert_parallel
+    assert text_config.moe_intermediate_size == 2048
+    assert text_config.n_routed_experts == 288
+    assert text_config.n_shared_experts == 1
+    assert text_config.glm53_tp3_shared_expert_intermediate_size == 2112
+
+    # The tokenizer/logits contract remains the checkpoint vocabulary while
+    # parameters are allocated using the explicitly recorded storage size.
+    assert text_config.vocab_size == 154880
+    assert text_config.glm53_tp3_vocab_padding_size == 192
+    assert text_config.glm53_tp3_vocab_storage_size == 154944
+    assert text_config.glm53_tp3_mtp_projection_size == 4098
+    assert text_config.glm53_tp3_padding is True
+
+    wrapper_config = glm53_model_config.hf_config
+    assert wrapper_config.glm53_tp3_shared_expert_intermediate_size == 2112
+    assert wrapper_config.glm53_tp3_vocab_storage_size == 154944
+    assert wrapper_config.glm53_tp3_padding is True
+    assert glm53_model_config.model_arch_config.total_num_attention_heads == 72
+    assert glm53_model_config.model_arch_config.total_num_kv_heads == 72
+    assert glm53_model_config.model_arch_config.vocab_size == 154880
+
+
+def test_glm53_tp3_weights_mode_records_exact_vision_storage_geometry(
+    glm53_model_config: FakeGlm53ModelConfig,
+    tp3_ep_parallel_config: ParallelConfig,
+) -> None:
+    assert apply_glm53_tp3_target_geometry(
+        cast(Any, glm53_model_config), tp3_ep_parallel_config
+    )
+
+    vision_config = glm53_model_config.hf_config.vision_config
+    assert (
+        vision_config.num_heads,
+        vision_config.original_num_heads,
+    ) == (18, 16)
+    assert vision_config.hidden_size == 1024
+    assert vision_config.glm53_tp3_attention_projection_size == 1152
+    assert (
+        vision_config.intermediate_size,
+        vision_config.original_intermediate_size,
+    ) == (4098, 4096)
+    assert (
+        vision_config.projection_intermediate_size,
+        vision_config.original_projection_intermediate_size,
+    ) == (10242, 10240)
+    assert vision_config.glm53_tp3_padding is True
+
+
+def test_glm53_tp3_target_geometry_is_idempotent(
+    glm53_model_config: FakeGlm53ModelConfig,
+    tp3_ep_parallel_config: ParallelConfig,
+) -> None:
+    assert apply_glm53_tp3_target_geometry(
+        cast(Any, glm53_model_config), tp3_ep_parallel_config
+    )
+    once = _snapshot(glm53_model_config)
+
+    assert apply_glm53_tp3_target_geometry(
+        cast(Any, glm53_model_config), tp3_ep_parallel_config
+    )
+
+    assert _snapshot(glm53_model_config) == once
+
+
+@pytest.mark.parametrize(
+    ("attribute", "invalid_value"),
+    [
+        ("num_attention_heads", 63),
+        ("num_key_value_heads", 63),
+        ("linear_num_heads", 63),
+        ("moe_intermediate_size", 2047),
+        ("hidden_size", 4095),
+        ("vocab_size", 154879),
+    ],
+)
+def test_glm53_tp3_invalid_target_checkpoint_geometry_fails_closed(
+    glm53_model_config: FakeGlm53ModelConfig,
+    tp3_ep_parallel_config: ParallelConfig,
+    attribute: str,
+    invalid_value: int,
+) -> None:
+    setattr(glm53_model_config.hf_text_config, attribute, invalid_value)
+    before = _snapshot(glm53_model_config)
+
+    with pytest.raises(ValueError, match=rf"expected {attribute}="):
+        apply_glm53_tp3_target_geometry(
+            cast(Any, glm53_model_config), tp3_ep_parallel_config
+        )
+
+    assert _snapshot(glm53_model_config) == before
+
+
+def test_glm53_tp4_parallel_config_is_an_exact_attribute_noop(
+    glm53_model_config: FakeGlm53ModelConfig,
+) -> None:
+    parallel_config = ParallelConfig(tensor_parallel_size=4)
+    before = _snapshot(glm53_model_config)
+
+    assert not apply_glm53_tp3_target_geometry(
+        cast(Any, glm53_model_config), parallel_config
+    )
+
+    assert parallel_config.tensor_parallel_size == 4
+    assert _snapshot(glm53_model_config) == before
+    assert not hasattr(glm53_model_config.hf_text_config, "glm53_tp3_padding")
+
+
+def test_unrelated_tp3_model_is_an_exact_noop() -> None:
+    model_config = FakeGlm53ModelConfig(
+        architecture="LlamaForCausalLM", model_type="llama"
+    )
+    parallel_config = ParallelConfig(tensor_parallel_size=3)
+    before = _snapshot(model_config)
+
+    assert not apply_glm53_tp3_target_geometry(
+        cast(Any, model_config), parallel_config
+    )
+
+    assert _snapshot(model_config) == before
+    assert not hasattr(model_config.hf_text_config, "glm53_tp3_padding")
+
+
+def test_glm53_tp3_mtp_draft_preserves_expert_parallel_topology() -> None:
+    target_model_config = FakeGlm53ModelConfig()
+    draft_model_config = FakeGlm53ModelConfig(
+        architecture="Glm5NextMTPModel", model_type="glm5_next_mtp"
+    )
+    target_parallel_config = ParallelConfig(
+        tensor_parallel_size=3,
+        data_parallel_size=2,
+        data_parallel_size_local=2,
+        enable_expert_parallel=True,
+    )
+    draft_parallel_config = ParallelConfig(tensor_parallel_size=3)
+    speculative_config = SimpleNamespace(
+        method="mtp",
+        target_model_config=target_model_config,
+        target_parallel_config=target_parallel_config,
+        draft_model_config=draft_model_config,
+        draft_parallel_config=draft_parallel_config,
+    )
+
+    SpeculativeConfig._apply_glm53_tp3_draft_geometry(
+        cast(Any, speculative_config)
+    )
+
+    assert draft_parallel_config.tensor_parallel_size == 3
+    assert draft_parallel_config.data_parallel_size == 2
+    assert draft_parallel_config.data_parallel_size_local == 2
+    assert draft_parallel_config.enable_expert_parallel
+    draft_text_config = draft_model_config.hf_text_config
+    assert draft_text_config.num_attention_heads == 72
+    assert draft_text_config.linear_num_heads == 66
+    assert draft_text_config.moe_intermediate_size == 2048
+    assert draft_text_config.glm53_tp3_shared_expert_intermediate_size == 2112
+
+
+def test_glm53_tp3_dflash_drops_ep_and_couples_heads_with_vocab_storage() -> None:
+    target_model_config = FakeGlm53ModelConfig()
+    draft_model_config = FakeDFlashModelConfig()
+    target_parallel_config = ParallelConfig(
+        tensor_parallel_size=3,
+        data_parallel_size=2,
+        data_parallel_size_local=2,
+        enable_expert_parallel=True,
+    )
+    draft_parallel_config = ParallelConfig(
+        tensor_parallel_size=3, enable_expert_parallel=True
+    )
+    speculative_config = SimpleNamespace(
+        method="dflash",
+        target_model_config=target_model_config,
+        target_parallel_config=target_parallel_config,
+        draft_model_config=draft_model_config,
+        draft_parallel_config=draft_parallel_config,
+    )
+
+    SpeculativeConfig._apply_glm53_tp3_draft_geometry(
+        cast(Any, speculative_config)
+    )
+
+    assert draft_parallel_config.tensor_parallel_size == 3
+    assert draft_parallel_config.data_parallel_size == 2
+    assert draft_parallel_config.data_parallel_size_local == 2
+    assert not draft_parallel_config.enable_expert_parallel
+    draft_config = draft_model_config.hf_text_config
+    assert (
+        draft_config.num_attention_heads,
+        draft_config.num_key_value_heads,
+    ) == (36, 9)
+    assert (
+        draft_config.original_num_attention_heads,
+        draft_config.original_num_key_value_heads,
+    ) == (32, 8)
+    assert draft_config.vocab_size == 154880
+    assert draft_config.original_vocab_size == 154880
+    assert draft_config.draft_vocab_size == 154880
+    assert draft_config.glm53_tp3_vocab_padding_size == 192
+    assert draft_config.glm53_tp3_vocab_storage_size == 154944
+    assert draft_model_config.model_arch_config.total_num_attention_heads == 36
+    assert draft_model_config.model_arch_config.total_num_kv_heads == 9
+    assert draft_model_config.model_arch_config.vocab_size == 154880
+
+
+@pytest.mark.parametrize(
+    ("attribute", "invalid_value"),
+    [
+        ("num_attention_heads", 31),
+        ("num_key_value_heads", 7),
+        ("vocab_size", 154879),
+    ],
+)
+def test_glm53_tp3_invalid_dflash_checkpoint_geometry_fails_closed(
+    attribute: str,
+    invalid_value: int,
+) -> None:
+    target_model_config = FakeGlm53ModelConfig()
+    draft_model_config = FakeDFlashModelConfig()
+    target_parallel_config = ParallelConfig(tensor_parallel_size=3)
+    draft_parallel_config = ParallelConfig(tensor_parallel_size=3)
+    setattr(draft_model_config.hf_text_config, attribute, invalid_value)
+    before = _snapshot(draft_model_config)
+
+    with pytest.raises(ValueError, match=rf"expected {attribute}="):
+        apply_glm53_tp3_draft_geometry(
+            cast(Any, target_model_config),
+            target_parallel_config,
+            cast(Any, draft_model_config),
+            draft_parallel_config,
+        )
+
+    assert _snapshot(draft_model_config) == before
+
+
+def test_glm53_dflash_tp4_parallel_configs_are_an_exact_noop() -> None:
+    target_model_config = FakeGlm53ModelConfig()
+    draft_model_config = FakeDFlashModelConfig()
+    target_parallel_config = ParallelConfig(tensor_parallel_size=4)
+    draft_parallel_config = ParallelConfig(tensor_parallel_size=4)
+    before = _snapshot(draft_model_config)
+
+    assert not apply_glm53_tp3_draft_geometry(
+        cast(Any, target_model_config),
+        target_parallel_config,
+        cast(Any, draft_model_config),
+        draft_parallel_config,
+    )
+
+    assert target_parallel_config.tensor_parallel_size == 4
+    assert draft_parallel_config.tensor_parallel_size == 4
+    assert _snapshot(draft_model_config) == before
+    assert not hasattr(draft_model_config.hf_text_config, "glm53_tp3_padding")

From 357fbb8f53ebf6fc137739c76b0f3a9f0b23313a Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:15:20 +0200
Subject: [PATCH 02/25] feat(tp): load explicit padded tensor shards locally

(cherry picked from commit 2fac3e59ce139883c050b3565e13ebde340c8996)
---
 vllm/model_executor/layers/linear.py | 254 +++++++++++++++++++++++----
 vllm/model_executor/parameter.py     | 126 +++++++++----
 2 files changed, 320 insertions(+), 60 deletions(-)

diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index b195fdd550ef..20d48b820ea2 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -39,12 +39,42 @@
     PackedvLLMParameter,
     PerTensorScaleParameter,
     RowvLLMParameter,
+    load_tensor_parallel_weight,
 )
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
+def _validate_loaded_axis_size(
+    name: str,
+    loaded_size: int,
+    physical_size: int,
+) -> None:
+    if loaded_size <= 0:
+        raise ValueError(f"{name} must be positive, got {loaded_size}")
+    if loaded_size > physical_size:
+        raise ValueError(
+            f"{name}={loaded_size} exceeds physical size {physical_size}"
+        )
+
+
+def _validate_packed_loaded_size(
+    param: Parameter,
+    dim: int,
+    loaded_size: int,
+    name: str,
+) -> None:
+    """Reject a padded packed layout before any destination writes occur."""
+    if getattr(param, "packed_dim", None) != dim:
+        return
+    packed_factor = getattr(param, "packed_factor", 1)
+    if loaded_size % packed_factor:
+        raise ValueError(
+            f"{name}={loaded_size} is not divisible by packed_factor="
+            f"{packed_factor}"
+        )
+
 WEIGHT_LOADER_V2_SUPPORTED = [
     "UnquantizedLinearMethod",
     "CompressedTensorsLinearMethod",
@@ -431,6 +461,9 @@ class ColumnParallelLinear(LinearBase):
             shard per rank (see ``DCPGroupColumnParallelLinear``).
         tp_size: Override the tensor-parallel world size used for sharding.
             Defaults to the global TP world size.
+        loaded_output_size: Output dimension represented by the checkpoint.
+            Defaults to ``output_size``. A smaller explicit value enables
+            destination-local zero padding for the physical TP shard tail.
     """
 
     # --8<-- [end:column_parallel_linear]
@@ -450,6 +483,7 @@ def __init__(
         disable_tp: bool = False,
         tp_rank: int | None = None,
         tp_size: int | None = None,
+        loaded_output_size: int | None = None,
     ):
         # Divide the weight matrix along the last dimension.
         if disable_tp:
@@ -463,6 +497,16 @@ def __init__(
                 if tp_size is not None
                 else get_tensor_model_parallel_world_size()
             )
+        self.loaded_output_size = (
+            output_size if loaded_output_size is None else loaded_output_size
+        )
+        _validate_loaded_axis_size(
+            "loaded_output_size", self.loaded_output_size, output_size
+        )
+        self._allow_loaded_output_padding = (
+            loaded_output_size is not None
+            and self.loaded_output_size != output_size
+        )
         self.input_size_per_partition = input_size
         self.output_size_per_partition = divide(output_size, self.tp_size)
         self.output_partition_sizes = [self.output_size_per_partition]
@@ -552,9 +596,22 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
 
         param_data = param.data
         if output_dim is not None and not is_sharded_weight:
+            _validate_packed_loaded_size(
+                param,
+                output_dim,
+                self.loaded_output_size,
+                "loaded_output_size",
+            )
             shard_size = param_data.shape[output_dim]
             start_idx = self.tp_rank * shard_size
-            loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
+            load_tensor_parallel_weight(
+                param_data,
+                loaded_weight,
+                output_dim,
+                start_idx,
+                allow_padding=self._allow_loaded_output_padding,
+            )
+            return
 
         # Special case for loading scales off disk, which often do not
         # have a shape (such as in the case of AutoFP8).
@@ -570,7 +627,10 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
         if len(loaded_weight.shape) == 0:
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
-        param.load_column_parallel_weight(loaded_weight=loaded_weight)
+        param.load_column_parallel_weight(
+            loaded_weight=loaded_weight,
+            allow_padding=self._allow_loaded_output_padding,
+        )
 
     def forward(
         self,
@@ -666,6 +726,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, all weights matrix won't be sharded, this layer
                     will be treated as a "Replicated" MergedLinear.
+        loaded_output_sizes: Output dimensions represented by the checkpoint.
+            Defaults to ``output_sizes``. Smaller explicit values enable
+            destination-local zero padding for physical TP shard tails.
     """
 
     def __init__(
@@ -681,8 +744,30 @@ def __init__(
         *,
         return_bias: bool = True,
         disable_tp: bool = False,
+        loaded_output_sizes: list[int] | None = None,
     ):
         self.output_sizes = output_sizes
+        self.loaded_output_sizes = (
+            output_sizes if loaded_output_sizes is None else loaded_output_sizes
+        )
+        if len(self.loaded_output_sizes) != len(self.output_sizes):
+            raise ValueError(
+                "loaded_output_sizes must have the same length as output_sizes"
+            )
+        for shard_id, (loaded_size, physical_size) in enumerate(
+            zip(self.loaded_output_sizes, self.output_sizes)
+        ):
+            _validate_loaded_axis_size(
+                f"loaded_output_sizes[{shard_id}]",
+                loaded_size,
+                physical_size,
+            )
+        self._allow_loaded_output_padding = [
+            loaded_output_sizes is not None and loaded_size != physical_size
+            for loaded_size, physical_size in zip(
+                self.loaded_output_sizes, self.output_sizes
+            )
+        ]
         self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
         self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0
 
@@ -753,15 +838,16 @@ def weight_loader(
                 param_data.copy_(loaded_weight)
                 return
 
-            output_sizes = (
-                self.output_sizes[loaded_shard_id[0] : loaded_shard_id[-1] + 1]
+            shard_ids = (
+                list(loaded_shard_id)
                 if loaded_shard_id is not None
-                else self.output_sizes
+                else list(range(len(self.loaded_output_sizes)))
             )
+            output_sizes = [self.loaded_output_sizes[idx] for idx in shard_ids]
             current_shard_offset = 0
             shard_offsets: list[tuple[int, int, int]] = []
-            for i, output_size in enumerate(output_sizes):
-                shard_offsets.append((i, current_shard_offset, output_size))
+            for shard_id, output_size in zip(shard_ids, output_sizes):
+                shard_offsets.append((shard_id, current_shard_offset, output_size))
                 current_shard_offset += output_size
             packed_dim = getattr(param, "packed_dim", None)
             for shard_id, shard_offset, shard_size in shard_offsets:
@@ -776,6 +862,12 @@ def weight_loader(
                     )
 
                 if packed_dim == output_dim:
+                    _validate_packed_loaded_size(
+                        param,
+                        output_dim,
+                        self.loaded_output_sizes[shard_id],
+                        f"loaded_output_sizes[{shard_id}]",
+                    )
                     shard_size = shard_size // param.packed_factor
                     shard_offset = shard_offset // param.packed_factor
                     # Special case for Marlin.
@@ -818,7 +910,16 @@ def weight_loader(
             param_data = param_data.narrow(output_dim, shard_offset, shard_size)
             start_idx = self.tp_rank * shard_size
             if not is_sharded_weight:
-                loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
+                load_tensor_parallel_weight(
+                    param_data,
+                    loaded_weight,
+                    output_dim,
+                    start_idx,
+                    allow_padding=self._allow_loaded_output_padding[
+                        loaded_shard_id
+                    ],
+                )
+                return
         # Special case for per-tensor scales in fused case.
         elif needs_scalar_to_array:
             param_data, loaded_weight = adjust_scalar_to_fused_array(
@@ -842,6 +943,7 @@ def _load_fused_module_from_checkpoint(
         param: BasevLLMParameter,
         loaded_weight: torch.Tensor,
         output_sizes: list[int] | None = None,
+        shard_ids: list[int] | None = None,
     ):
         """
         Handle special case for models where MLP layers are already
@@ -855,9 +957,12 @@ def _load_fused_module_from_checkpoint(
 
         current_shard_offset = 0
         shard_offsets: list[tuple[int, int, int]] = []
-        output_sizes = output_sizes or self.output_sizes
-        for i, output_size in enumerate(output_sizes):
-            shard_offsets.append((i, current_shard_offset, output_size))
+        output_sizes = output_sizes or self.loaded_output_sizes
+        shard_ids = shard_ids or list(range(len(output_sizes)))
+        if len(shard_ids) != len(output_sizes):
+            raise ValueError("shard_ids and output_sizes must have the same length")
+        for shard_id, output_size in zip(shard_ids, output_sizes):
+            shard_offsets.append((shard_id, current_shard_offset, output_size))
             current_shard_offset += output_size
 
         for shard_id, shard_offset, shard_size in shard_offsets:
@@ -905,20 +1010,24 @@ def weight_loader_v2(
             elif type(param) in (RowvLLMParameter, BasevLLMParameter):
                 param.load_merged_column_weight(loaded_weight=loaded_weight)
                 return
-            output_sizes = (
-                [self.output_sizes[idx] for idx in loaded_shard_id]
-                if loaded_shard_id
-                else None
+            shard_ids = (
+                list(loaded_shard_id)
+                if loaded_shard_id is not None
+                else list(range(len(self.loaded_output_sizes)))
             )
+            output_sizes = [self.loaded_output_sizes[idx] for idx in shard_ids]
             if isinstance(param, BlockQuantScaleParameter):
                 weight_block_size = getattr(self, "weight_block_size", None)
                 output_sizes = [
                     adjust_block_scale_shard(weight_block_size, size, 0)[0]
-                    for size in (output_sizes or self.output_sizes)
+                    for size in output_sizes
                 ]
             # TODO: @dsikka - move to parameter.py
             self._load_fused_module_from_checkpoint(
-                param, loaded_weight, output_sizes=output_sizes
+                param,
+                loaded_weight,
+                output_sizes=output_sizes,
+                shard_ids=shard_ids,
             )
             return
 
@@ -940,6 +1049,7 @@ def weight_loader_v2(
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
+            allow_padding=self._allow_loaded_output_padding[loaded_shard_id],
         )
 
     def load_weights(
@@ -994,6 +1104,10 @@ class QKVParallelLinear(ColumnParallelLinear):
                         (e.g. model.layers.0.qkv_proj)
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, weights matrix won't be sharded through tp rank.
+        loaded_total_num_heads: Query-head count represented by the checkpoint.
+            Defaults to ``total_num_heads``.
+        loaded_total_num_kv_heads: KV-head count represented by the checkpoint.
+            Defaults to ``total_num_kv_heads``.
     """
 
     def __init__(
@@ -1011,6 +1125,8 @@ def __init__(
         return_bias: bool = True,
         disable_tp: bool = False,
         v_head_size: int | None = None,
+        loaded_total_num_heads: int | None = None,
+        loaded_total_num_kv_heads: int | None = None,
     ):
         self.hidden_size = hidden_size
         self.head_size = head_size
@@ -1019,6 +1135,34 @@ def __init__(
         if total_num_kv_heads is None:
             total_num_kv_heads = total_num_heads
         self.total_num_kv_heads = total_num_kv_heads
+        self.loaded_total_num_heads = (
+            total_num_heads
+            if loaded_total_num_heads is None
+            else loaded_total_num_heads
+        )
+        self.loaded_total_num_kv_heads = (
+            total_num_kv_heads
+            if loaded_total_num_kv_heads is None
+            else loaded_total_num_kv_heads
+        )
+        _validate_loaded_axis_size(
+            "loaded_total_num_heads",
+            self.loaded_total_num_heads,
+            self.total_num_heads,
+        )
+        _validate_loaded_axis_size(
+            "loaded_total_num_kv_heads",
+            self.loaded_total_num_kv_heads,
+            self.total_num_kv_heads,
+        )
+        self._allow_loaded_qkv_padding = {
+            "q": loaded_total_num_heads is not None
+            and self.loaded_total_num_heads != self.total_num_heads,
+            "k": loaded_total_num_kv_heads is not None
+            and self.loaded_total_num_kv_heads != self.total_num_kv_heads,
+            "v": loaded_total_num_kv_heads is not None
+            and self.loaded_total_num_kv_heads != self.total_num_kv_heads,
+        }
         # Divide the weight matrix along the last dimension.
         tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
         self.num_heads = divide(self.total_num_heads, tp_size)
@@ -1089,16 +1233,17 @@ def _load_fused_module_from_checkpoint(
         """
         shard_offsets = [
             # (shard_id, shard_offset, shard_size)
-            ("q", 0, self.total_num_heads * self.head_size),
+            ("q", 0, self.loaded_total_num_heads * self.head_size),
             (
                 "k",
-                self.total_num_heads * self.head_size,
-                self.total_num_kv_heads * self.head_size,
+                self.loaded_total_num_heads * self.head_size,
+                self.loaded_total_num_kv_heads * self.head_size,
             ),
             (
                 "v",
-                (self.total_num_heads + self.total_num_kv_heads) * self.head_size,
-                self.total_num_kv_heads * self.v_head_size,
+                (self.loaded_total_num_heads + self.loaded_total_num_kv_heads)
+                * self.head_size,
+                self.loaded_total_num_kv_heads * self.v_head_size,
             ),
         ]
 
@@ -1166,6 +1311,7 @@ def weight_loader_v2(
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
+            allow_padding=self._allow_loaded_qkv_padding[loaded_shard_id],
         )
 
     def weight_loader(
@@ -1196,16 +1342,17 @@ def weight_loader(
                 return
             shard_offsets = [
                 # (shard_id, shard_offset, shard_size)
-                ("q", 0, self.total_num_heads * self.head_size),
+                ("q", 0, self.loaded_total_num_heads * self.head_size),
                 (
                     "k",
-                    self.total_num_heads * self.head_size,
-                    self.total_num_kv_heads * self.head_size,
+                    self.loaded_total_num_heads * self.head_size,
+                    self.loaded_total_num_kv_heads * self.head_size,
                 ),
                 (
                     "v",
-                    (self.total_num_heads + self.total_num_kv_heads) * self.head_size,
-                    self.total_num_kv_heads * self.v_head_size,
+                    (self.loaded_total_num_heads + self.loaded_total_num_kv_heads)
+                    * self.head_size,
+                    self.loaded_total_num_kv_heads * self.v_head_size,
                 ),
             ]
             packed_dim = getattr(param, "packed_dim", None)
@@ -1221,6 +1368,14 @@ def weight_loader(
                     )
 
                 if packed_dim == output_dim:
+                    loaded_size = {
+                        "q": self.loaded_total_num_heads * self.head_size,
+                        "k": self.loaded_total_num_kv_heads * self.head_size,
+                        "v": self.loaded_total_num_kv_heads * self.v_head_size,
+                    }[shard_id]
+                    _validate_packed_loaded_size(
+                        param, output_dim, loaded_size, f"loaded {shard_id} size"
+                    )
                     shard_size = round(shard_size // param.packed_factor)
                     shard_offset = round(shard_offset // param.packed_factor)
 
@@ -1277,7 +1432,16 @@ def weight_loader(
             start_idx = shard_rank * shard_size
 
             if not is_sharded_weight:
-                loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
+                load_tensor_parallel_weight(
+                    param_data,
+                    loaded_weight,
+                    output_dim,
+                    start_idx,
+                    allow_padding=self._allow_loaded_qkv_padding[
+                        loaded_shard_id
+                    ],
+                )
+                return
 
         # Special case for per-tensor scales in fused case.
         elif needs_scalar_to_array:
@@ -1538,6 +1702,9 @@ class RowParallelLinear(LinearBase):
                         (e.g. model.layers.0.down_proj)
         return_bias: If true, return bias together with outputs in forward pass.
         disable_tp: If true, weights matrix won't be sharded through tp rank.
+        loaded_input_size: Input dimension represented by the checkpoint.
+            Defaults to ``input_size``. A smaller explicit value enables
+            destination-local zero padding for the physical TP shard tail.
     """
 
     # --8<-- [end:row_parallel_linear]
@@ -1556,10 +1723,21 @@ def __init__(
         *,
         return_bias: bool = True,
         disable_tp: bool = False,
+        loaded_input_size: int | None = None,
     ):
         # Divide the weight matrix along the first dimension.
         self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0
         self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
+        self.loaded_input_size = (
+            input_size if loaded_input_size is None else loaded_input_size
+        )
+        _validate_loaded_axis_size(
+            "loaded_input_size", self.loaded_input_size, input_size
+        )
+        self._allow_loaded_input_padding = (
+            loaded_input_size is not None
+            and self.loaded_input_size != input_size
+        )
         self.input_size_per_partition = divide(input_size, self.tp_size)
         self.output_size_per_partition = output_size
         self.output_partition_sizes = [output_size]
@@ -1617,9 +1795,22 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
 
         param_data = param.data
         if input_dim is not None and not is_sharded_weight:
+            _validate_packed_loaded_size(
+                param,
+                input_dim,
+                self.loaded_input_size,
+                "loaded_input_size",
+            )
             shard_size = param_data.shape[input_dim]
             start_idx = self.tp_rank * shard_size
-            loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size)
+            load_tensor_parallel_weight(
+                param_data,
+                loaded_weight,
+                input_dim,
+                start_idx,
+                allow_padding=self._allow_loaded_input_padding,
+            )
+            return
 
         # Special case for loading scales off disk, which often do not
         # have a shape (such as in the case of AutoFP8).
@@ -1636,7 +1827,10 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
 
-        param.load_row_parallel_weight(loaded_weight=loaded_weight)
+        param.load_row_parallel_weight(
+            loaded_weight=loaded_weight,
+            allow_padding=self._allow_loaded_input_padding,
+        )
 
     def forward(
         self,
diff --git a/vllm/model_executor/parameter.py b/vllm/model_executor/parameter.py
index 7f96ceda09db..93df17a3b760 100644
--- a/vllm/model_executor/parameter.py
+++ b/vllm/model_executor/parameter.py
@@ -28,6 +28,57 @@
 
 logger = init_logger(__name__)
 
+def load_tensor_parallel_weight(
+    param_data: torch.Tensor,
+    loaded_weight: torch.Tensor,
+    dim: int,
+    start_idx: int,
+    *,
+    allow_padding: bool = False,
+) -> None:
+    """Copy one TP shard, optionally zero-filling its unavailable tail.
+
+    Padding is opt-in so ordinary loaders retain strict ``Tensor.narrow``
+    failures.  The padded path writes directly into the rank-local destination
+    rather than allocating a checkpoint-sized or padded source tensor.
+    """
+    shard_size = param_data.shape[dim]
+    if not allow_padding:
+        loaded_shard = loaded_weight.narrow(dim, start_idx, shard_size)
+        assert param_data.shape == loaded_shard.shape
+        param_data.copy_(loaded_shard)
+        return
+
+    if start_idx < 0:
+        raise ValueError(f"TP shard start must be non-negative, got {start_idx}")
+    if loaded_weight.ndim != param_data.ndim:
+        raise ValueError(
+            "Padded TP weight rank mismatch: "
+            f"destination {tuple(param_data.shape)}, "
+            f"checkpoint {tuple(loaded_weight.shape)}"
+        )
+    dim = dim if dim >= 0 else loaded_weight.ndim + dim
+    if dim < 0 or dim >= loaded_weight.ndim:
+        raise IndexError(
+            f"TP shard dimension {dim} is invalid for rank {loaded_weight.ndim}"
+        )
+    for axis, (destination_size, loaded_size) in enumerate(
+        zip(param_data.shape, loaded_weight.shape)
+    ):
+        if axis != dim and destination_size != loaded_size:
+            raise ValueError(
+                "Padded TP weight shape mismatch outside the sharded axis: "
+                f"destination {tuple(param_data.shape)}, "
+                f"checkpoint {tuple(loaded_weight.shape)}"
+            )
+
+    available = max(0, loaded_weight.shape[dim] - start_idx)
+    copy_size = min(shard_size, available)
+    param_data.zero_()
+    if copy_size:
+        loaded_shard = loaded_weight.narrow(dim, start_idx, copy_size)
+        param_data.narrow(dim, 0, copy_size).copy_(loaded_shard)
+
 
 class BasevLLMParameter(Parameter):
     """
@@ -145,17 +196,24 @@ def __init__(self, output_dim: int, **kwargs):
     def output_dim(self):
         return self._output_dim
 
-    def load_column_parallel_weight(self, loaded_weight: torch.Tensor):
-        shard_size = self.data.shape[self.output_dim]
-        loaded_weight = loaded_weight.narrow(
-            self.output_dim, self.tp_rank * shard_size, shard_size
+    def load_column_parallel_weight(
+        self,
+        loaded_weight: torch.Tensor,
+        *,
+        allow_padding: bool = False,
+    ):
+        load_tensor_parallel_weight(
+            self.data,
+            loaded_weight,
+            self.output_dim,
+            self.tp_rank * self.data.shape[self.output_dim],
+            allow_padding=allow_padding,
         )
-        assert self.data.shape == loaded_weight.shape
-        self.data.copy_(loaded_weight)
 
     def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):
         shard_offset: int = kwargs["shard_offset"]
         shard_size: int = kwargs["shard_size"]
+        allow_padding: bool = kwargs.get("allow_padding", False)
 
         # TODO: move these to PackedColumnParameter and PackedvLLMParameter
         if (
@@ -166,20 +224,23 @@ def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):
                 shard_offset=shard_offset, shard_size=shard_size
             )
 
-        param_data = self.data
-
-        param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
-        loaded_weight = loaded_weight.narrow(
-            self.output_dim, self.tp_rank * shard_size, shard_size
+        param_data = self.data.narrow(
+            self.output_dim, shard_offset, shard_size
+        )
+        load_tensor_parallel_weight(
+            param_data,
+            loaded_weight,
+            self.output_dim,
+            self.tp_rank * shard_size,
+            allow_padding=allow_padding,
         )
-        assert param_data.shape == loaded_weight.shape
-        param_data.copy_(loaded_weight)
 
     def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):
         shard_offset: int = kwargs["shard_offset"]
         shard_size: int = kwargs["shard_size"]
         shard_id: str = kwargs["shard_id"]
         num_heads: int = kwargs["num_heads"]
+        allow_padding: bool = kwargs.get("allow_padding", False)
 
         # TODO: move these to PackedColumnParameter and PackedvLLMParameter
         if (
@@ -190,15 +251,17 @@ def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):
                 shard_offset=shard_offset, shard_size=shard_size
             )
 
-        param_data = self.data
         shard_id_int = self.tp_rank if shard_id == "q" else self.tp_rank // num_heads
-        param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
-        loaded_weight = loaded_weight.narrow(
-            self.output_dim, shard_id_int * shard_size, shard_size
+        param_data = self.data.narrow(
+            self.output_dim, shard_offset, shard_size
+        )
+        load_tensor_parallel_weight(
+            param_data,
+            loaded_weight,
+            self.output_dim,
+            shard_id_int * shard_size,
+            allow_padding=allow_padding,
         )
-
-        assert param_data.shape == loaded_weight.shape
-        param_data.copy_(loaded_weight)
 
 
 class RowvLLMParameter(BasevLLMParameter):
@@ -217,18 +280,21 @@ def __init__(self, input_dim: int, **kwargs):
     def input_dim(self):
         return self._input_dim
 
-    def load_row_parallel_weight(self, loaded_weight: torch.Tensor):
-        shard_size = self.data.shape[self.input_dim]
-        loaded_weight = loaded_weight.narrow(
-            self.input_dim, self.tp_rank * shard_size, shard_size
+    def load_row_parallel_weight(
+        self,
+        loaded_weight: torch.Tensor,
+        *,
+        allow_padding: bool = False,
+    ):
+        param_data = self.data
+        load_tensor_parallel_weight(
+            param_data,
+            loaded_weight,
+            self.input_dim,
+            self.tp_rank * param_data.shape[self.input_dim],
+            allow_padding=allow_padding,
         )
 
-        if len(loaded_weight.shape) == 0:
-            loaded_weight = loaded_weight.reshape(1)
-
-        assert self.data.shape == loaded_weight.shape
-        self.data.copy_(loaded_weight)
-
 
 class ModelWeightParameter(_ColumnvLLMParameter, RowvLLMParameter):
     """

From 24af676e90f8211b082c174d9fccefae127d29ca Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:03:55 +0200
Subject: [PATCH 03/25] feat(dflash): support GLM-5.3 TP3 physical geometry

(cherry picked from commit 890384d6aa391cf4bd6f612ced9515fbcd479a5d)
---
 vllm/model_executor/models/qwen3_dflash.py  | 87 +++++++++++++++++++--
 vllm/model_executor/models/qwen3_dflash2.py |  5 +-
 2 files changed, 83 insertions(+), 9 deletions(-)

diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py
index 1147ba9aaa72..105824aec12d 100644
--- a/vllm/model_executor/models/qwen3_dflash.py
+++ b/vllm/model_executor/models/qwen3_dflash.py
@@ -101,6 +101,41 @@ def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None:
             return style
     return None
 
+def _get_dflash_draft_vocab_size(config: Qwen3Config) -> int:
+    """Return the logical draft vocabulary across old and current config APIs."""
+    draft_vocab_size = getattr(config, "draft_vocab_size", None)
+    if draft_vocab_size is None:
+        draft_vocab_size = getattr(config, "original_vocab_size", None)
+    if draft_vocab_size is None:
+        draft_vocab_size = config.vocab_size
+    return int(draft_vocab_size)
+
+
+def _get_glm53_tp3_head_geometry(
+    config: Qwen3Config,
+) -> tuple[int, int] | None:
+    if not getattr(config, "glm53_tp3_padding", False):
+        return None
+
+    physical = (int(config.num_attention_heads), int(config.num_key_value_heads))
+    logical = (
+        int(config.original_num_attention_heads),
+        int(config.original_num_key_value_heads),
+    )
+    if (*physical, *logical) != (36, 9, 32, 8):
+        raise ValueError(
+            "GLM-5.3 DFlash TP3 requires physical/logical Q/KV heads "
+            f"36/9 from 32/8, got {physical[0]}/{physical[1]} from "
+            f"{logical[0]}/{logical[1]}."
+        )
+    return logical
+
+
+def _get_glm53_tp3_vocab_kwargs(config: Qwen3Config) -> dict[str, int]:
+    if not getattr(config, "glm53_tp3_padding", False):
+        return {}
+    return {"padding_size": int(config.glm53_tp3_vocab_padding_size)}
+
 
 def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
     spec_config = vllm_config.speculative_config
@@ -218,6 +253,7 @@ def __init__(
         hidden_size: int,
         num_heads: int,
         num_kv_heads: int,
+        config: Qwen3Config,
         rope_parameters: dict,
         max_position: int = 4096 * 32,
         head_dim: int | None = None,
@@ -245,11 +281,23 @@ def __init__(
         else:
             assert tp_size % self.total_num_kv_heads == 0
         self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
-        self.head_dim = head_dim or hidden_size // self.total_num_heads
+        head_geometry = _get_glm53_tp3_head_geometry(config)
+        loaded_num_heads = (
+            head_geometry[0] if head_geometry is not None else self.total_num_heads
+        )
+        self.head_dim = head_dim or hidden_size // loaded_num_heads
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
 
+        qkv_loader_kwargs = (
+            {
+                "loaded_total_num_heads": head_geometry[0],
+                "loaded_total_num_kv_heads": head_geometry[1],
+            }
+            if head_geometry is not None
+            else {}
+        )
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
             self.head_dim,
@@ -258,6 +306,12 @@ def __init__(
             bias=attention_bias,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
+            **qkv_loader_kwargs,
+        )
+        o_proj_loader_kwargs = (
+            {"loaded_input_size": head_geometry[0] * self.head_dim}
+            if head_geometry is not None
+            else {}
         )
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
@@ -265,6 +319,7 @@ def __init__(
             bias=attention_bias,  # DFlash has o_proj bias when using attention bias
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
+            **o_proj_loader_kwargs,
         )
 
         self.rotary_emb = get_rope(
@@ -362,6 +417,7 @@ def __init__(
 
         self.self_attn = DFlashQwen3Attention(
             hidden_size=self.hidden_size,
+            config=config,
             num_heads=config.num_attention_heads,
             max_position=config.max_position_embeddings,
             num_kv_heads=config.num_key_value_heads,
@@ -460,6 +516,7 @@ def __init__(
             self.config.vocab_size,
             self.config.hidden_size,
             prefix=maybe_prefix(prefix, "embed_tokens"),
+            **_get_glm53_tp3_vocab_kwargs(self.config),
         )
 
         # Masked query slots are fed to the draft as `mask_token_id`. Most DFlash
@@ -800,11 +857,27 @@ def _preprocess(
         for name, loaded_weight in weights:
             if "attention_sink_bias" in name:
                 # Sink bias is per-head; shard it across TP ranks like the
-                # attention heads themselves.
-                heads_per_rank = loaded_weight.shape[0] // tp_size
-                loaded_weight = loaded_weight.narrow(
-                    0, tp_rank * heads_per_rank, heads_per_rank
+                # attention heads themselves. GLM-5.3 TP3 writes the short
+                # rank-2 source shard into a zeroed physical destination.
+                heads_per_rank = self.config.num_attention_heads // tp_size
+                source_start = tp_rank * heads_per_rank
+                source_size = min(
+                    heads_per_rank,
+                    max(loaded_weight.shape[0] - source_start, 0),
                 )
+                if source_size == heads_per_rank:
+                    loaded_weight = loaded_weight.narrow(
+                        0, source_start, heads_per_rank
+                    )
+                else:
+                    local_weight = loaded_weight.new_zeros(
+                        heads_per_rank, *loaded_weight.shape[1:]
+                    )
+                    if source_size:
+                        local_weight[:source_size].copy_(
+                            loaded_weight.narrow(0, source_start, source_size)
+                        )
+                    loaded_weight = local_weight
             yield name, loaded_weight
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -821,8 +894,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)
         self.draft_model_config = vllm_config.speculative_config.draft_model_config
         self.config = self.draft_model_config.hf_config
-        if getattr(self.config, "draft_vocab_size", None) is None:
-            self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
+        self.config.draft_vocab_size = _get_dflash_draft_vocab_size(self.config)
         target_layer_num = vllm_config.model_config.get_num_layers(
             vllm_config.parallel_config
         )
@@ -837,6 +909,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             self.config.draft_vocab_size,
             self.config.hidden_size,
             prefix=maybe_prefix(prefix, "lm_head"),
+            **_get_glm53_tp3_vocab_kwargs(self.config),
         )
         self.logits_processor = LogitsProcessor(
             self.config.draft_vocab_size, scale=logit_scale
diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py
index be49d65f33c2..1163589a1de3 100644
--- a/vllm/model_executor/models/qwen3_dflash2.py
+++ b/vllm/model_executor/models/qwen3_dflash2.py
@@ -18,6 +18,7 @@
     DFlashQwen3DecoderLayer,
     DFlashQwen3ForCausalLM,
     DFlashQwen3Model,
+    _get_dflash_draft_vocab_size,
 )
 from .utils import maybe_prefix
 
@@ -398,7 +399,7 @@ def __init__(
         with set_model_tag("dflash2_candidate_selector"):
             self.candidate_selector = CandidateSelector(
                 hidden_size=self.config.hidden_size,
-                vocab_size=self.config.vocab_size,
+                vocab_size=_get_dflash_draft_vocab_size(self.config),
                 rank=int(draft_config["selector_rank"]),
                 top_k=int(draft_config["selector_top_k"]),
                 params_dtype=vllm_config.model_config.dtype,
@@ -418,7 +419,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         draft_config = self.config.dflash_config
         softcap = float(draft_config.get("final_logit_softcapping") or 0.0)
         self.candidate_logits_processor = LogitsProcessor(
-            vllm_config.model_config.get_vocab_size(),
+            _get_dflash_draft_vocab_size(self.config),
             scale=float(draft_config.get("output_multiplier", 1.0)),
             soft_cap=softcap if softcap > 0 else None,
         )

From 8bff0d5e9696def0fa5b5eac8e87de881627b68b Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:11:38 +0200
Subject: [PATCH 04/25] test(dflash): cover GLM-5.3 TP3 geometry

(cherry picked from commit a80e350644e69a9180a7a7b005fd9ac1ec8d63f2)
---
 tests/models/test_glm53_tp3_dflash.py | 190 ++++++++++++++++++++++++++
 1 file changed, 190 insertions(+)
 create mode 100644 tests/models/test_glm53_tp3_dflash.py

diff --git a/tests/models/test_glm53_tp3_dflash.py b/tests/models/test_glm53_tp3_dflash.py
new file mode 100644
index 000000000000..8d67fbbd992a
--- /dev/null
+++ b/tests/models/test_glm53_tp3_dflash.py
@@ -0,0 +1,190 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm.config.speculative import SpeculativeConfig
+from vllm.model_executor.layers import vocab_parallel_embedding
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.models import qwen3_dflash
+from vllm.model_executor.models.qwen3_dflash import (
+    DFlashQwen3Model,
+    _get_dflash_draft_vocab_size,
+    _get_glm53_tp3_head_geometry,
+    _get_glm53_tp3_vocab_kwargs,
+)
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+
+
+def _tp3_dflash_config(**overrides):
+    values = {
+        "glm53_tp3_padding": True,
+        "glm53_tp3_vocab_padding_size": 192,
+        "glm53_tp3_vocab_storage_size": 154944,
+        "num_attention_heads": 36,
+        "num_key_value_heads": 9,
+        "original_num_attention_heads": 32,
+        "original_num_key_value_heads": 8,
+        "original_vocab_size": 154880,
+        "draft_vocab_size": 154880,
+        "vocab_size": 154880,
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_dflash_tp3_geometry_and_vocab_storage(monkeypatch) -> None:
+    config = _tp3_dflash_config()
+
+    assert _get_glm53_tp3_head_geometry(config) == (32, 8)
+    assert _get_dflash_draft_vocab_size(config) == 154880
+    vocab_kwargs = _get_glm53_tp3_vocab_kwargs(config)
+    assert vocab_kwargs == {"padding_size": 192}
+
+    monkeypatch.setattr(
+        vocab_parallel_embedding,
+        "get_tensor_model_parallel_world_size",
+        lambda: 3,
+    )
+    monkeypatch.setattr(
+        vocab_parallel_embedding,
+        "get_tensor_model_parallel_rank",
+        lambda: 2,
+    )
+    embedding = VocabParallelEmbedding(
+        _get_dflash_draft_vocab_size(config),
+        8,
+        **vocab_kwargs,
+    )
+    assert embedding.num_embeddings == 154880
+    assert embedding.num_embeddings_padded == 154944
+    assert embedding.num_embeddings_per_partition == 51648
+    assert embedding.weight.shape == (51648, 8)
+
+
+def test_dflash_tp3_sink_bias_pads_only_rank_local_tail(monkeypatch) -> None:
+    monkeypatch.setattr(qwen3_dflash, "get_tensor_model_parallel_world_size", lambda: 3)
+    monkeypatch.setattr(qwen3_dflash, "get_tensor_model_parallel_rank", lambda: 2)
+    model = DFlashQwen3Model.__new__(DFlashQwen3Model)
+    torch.nn.Module.__init__(model)
+    model.config = _tp3_dflash_config()
+    loaded = torch.arange(32, dtype=torch.float32)
+
+    [(name, local)] = model._preprocess([("attention_sink_bias", loaded)])
+
+    assert name == "attention_sink_bias"
+    assert local.shape == (12,)
+    torch.testing.assert_close(local[:8], loaded[24:32])
+    torch.testing.assert_close(local[8:], torch.zeros(4))
+    assert local.untyped_storage().data_ptr() != loaded.untyped_storage().data_ptr()
+
+
+def test_dflash_tp4_paths_are_exact_noops(monkeypatch) -> None:
+    config = SimpleNamespace(
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        vocab_size=154880,
+    )
+    original_fields = vars(config).copy()
+    assert _get_glm53_tp3_head_geometry(config) is None
+    assert _get_glm53_tp3_vocab_kwargs(config) == {}
+    assert _get_dflash_draft_vocab_size(config) == 154880
+    assert vars(config) == original_fields
+
+    monkeypatch.setattr(
+        vocab_parallel_embedding,
+        "get_tensor_model_parallel_world_size",
+        lambda: 4,
+    )
+    monkeypatch.setattr(
+        vocab_parallel_embedding,
+        "get_tensor_model_parallel_rank",
+        lambda: 2,
+    )
+    embedding = VocabParallelEmbedding(config.vocab_size, 8)
+    assert embedding.num_embeddings == 154880
+    assert embedding.num_embeddings_padded == 154880
+    assert embedding.num_embeddings_per_partition == 38720
+
+    monkeypatch.setattr(qwen3_dflash, "get_tensor_model_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(qwen3_dflash, "get_tensor_model_parallel_rank", lambda: 2)
+    model = DFlashQwen3Model.__new__(DFlashQwen3Model)
+    torch.nn.Module.__init__(model)
+    model.config = config
+    loaded = torch.arange(32, dtype=torch.float32)
+    [(_, local)] = model._preprocess([("attention_sink_bias", loaded)])
+    torch.testing.assert_close(local, loaded[16:24])
+    assert local.untyped_storage().data_ptr() == loaded.untyped_storage().data_ptr()
+
+
+def test_dense_dflash_tp3_drops_target_expert_parallel(monkeypatch) -> None:
+    from vllm.transformers_utils.configs import glm53_tp3
+
+    monkeypatch.setattr(glm53_tp3, "is_glm53_config", lambda _: True)
+    applied = []
+    monkeypatch.setattr(
+        glm53_tp3,
+        "apply_glm53_tp3_draft_geometry",
+        lambda *args: applied.append(args),
+    )
+    placement = {
+        "prefill_context_parallel_size": 1,
+        "data_parallel_size": 2,
+        "data_parallel_size_local": 2,
+        "data_parallel_rank": 1,
+        "data_parallel_rank_local": 1,
+        "data_parallel_master_ip": "127.0.0.1",
+        "data_parallel_rpc_port": 1234,
+        "data_parallel_master_port": 4321,
+        "data_parallel_backend": "mp",
+        "data_parallel_external_lb": False,
+        "data_parallel_hybrid_lb": False,
+    }
+    target_parallel = SimpleNamespace(
+        tensor_parallel_size=3,
+        enable_expert_parallel=True,
+        **placement,
+    )
+    draft_parallel = SimpleNamespace(
+        tensor_parallel_size=3,
+        enable_expert_parallel=True,
+    )
+    spec = object.__new__(SpeculativeConfig)
+    object.__setattr__(spec, "method", "dflash")
+    object.__setattr__(spec, "target_model_config", object())
+    object.__setattr__(spec, "target_parallel_config", target_parallel)
+    object.__setattr__(spec, "draft_model_config", object())
+    object.__setattr__(spec, "draft_parallel_config", draft_parallel)
+
+    spec._apply_glm53_tp3_draft_geometry()
+
+    assert draft_parallel.enable_expert_parallel is False
+    for name, value in placement.items():
+        assert getattr(draft_parallel, name) == value
+    assert len(applied) == 1
+
+
+def test_dflash7_uses_eight_target_kda_state_columns(monkeypatch) -> None:
+    def fake_base_init(self, vllm_config, device) -> None:
+        self.max_num_tokens = 16
+        self.hidden_size = 4
+        self.dtype = torch.float32
+        self.num_speculative_steps = 7
+        self.max_num_reqs = 2
+        self.draft_model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(
+                num_hidden_layers=0,
+                dflash_config={"mask_token_id": 0},
+            )
+        )
+
+    monkeypatch.setattr(DraftModelSpeculator, "__init__", fake_base_init)
+    speculator = DFlashSpeculator(SimpleNamespace(), torch.device("cpu"))
+
+    # The target KDA keeps one initial-state column plus one per draft token.
+    assert speculator.num_speculative_steps == 7
+    assert speculator.num_query_per_req == 8
+    assert speculator.sample_col.shape == (14,)

From 217c52980d5cd60d99461b7821bb305e0aa78f85 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:04:38 +0200
Subject: [PATCH 05/25] feat(glm5next): implement direct TP3 vision geometry

(cherry picked from commit 3b30789214e4a550d47ebc55b8015b3b2b5a6a24)
---
 tests/models/test_glm5next_vision_tp3.py  | 275 ++++++++++++++++++++++
 vllm/models/glm5next/nvidia/multimodal.py |  91 ++++++-
 2 files changed, 361 insertions(+), 5 deletions(-)
 create mode 100644 tests/models/test_glm5next_vision_tp3.py

diff --git a/tests/models/test_glm5next_vision_tp3.py b/tests/models/test_glm5next_vision_tp3.py
new file mode 100644
index 000000000000..65ab8250bc90
--- /dev/null
+++ b/tests/models/test_glm5next_vision_tp3.py
@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor import parameter
+from vllm.model_executor.layers import linear
+from vllm.models.glm5next.nvidia import multimodal as glm5next_multimodal
+
+
+@pytest.fixture
+def tp3_linear_state(monkeypatch):
+    monkeypatch.setattr(
+        glm5next_multimodal, "get_tensor_model_parallel_world_size", lambda: 3
+    )
+    monkeypatch.setattr(
+        glm5next_multimodal.parallel_state,
+        "get_tensor_model_parallel_rank",
+        lambda: 2,
+    )
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 3)
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 2)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 3)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 2)
+
+
+def test_glm5next_vision_tp3_attention_shards_and_zeros_local_tail(
+    monkeypatch, tp3_linear_state, default_vllm_config
+) -> None:
+    class FakeEncoderAttention(torch.nn.Module):
+        def __init__(self, **kwargs) -> None:
+            super().__init__()
+            self.kwargs = kwargs
+
+    monkeypatch.setattr(glm5next_multimodal, "is_vit_use_data_parallel", lambda: False)
+    monkeypatch.setattr(glm5next_multimodal, "MMEncoderAttention", FakeEncoderAttention)
+
+    attention = glm5next_multimodal.Glm5NextVisionAttention(
+        embed_dim=8,
+        num_heads=18,
+        projection_size=1152,
+        loaded_num_heads=16,
+        loaded_projection_size=1024,
+    )
+
+    assert attention.head_dim == 64
+    assert attention.num_attention_heads_per_partition == 6
+    assert attention.q_norm.weight.shape == (64,)
+    assert attention.qkv.weight.shape == (1152, 8)
+    assert attention.proj.weight.shape == (8, 384)
+
+    q = torch.arange(1024 * 8, dtype=torch.float32).view(1024, 8)
+    k = q + 10000
+    v = q + 20000
+    attention.qkv.weight.weight_loader(attention.qkv.weight, q, "q")
+    attention.qkv.weight.weight_loader(attention.qkv.weight, k, "k")
+    attention.qkv.weight.weight_loader(attention.qkv.weight, v, "v")
+
+    for offset, checkpoint in zip((0, 384, 768), (q, k, v)):
+        torch.testing.assert_close(
+            attention.qkv.weight[offset : offset + 256], checkpoint[768:1024]
+        )
+        torch.testing.assert_close(
+            attention.qkv.weight[offset + 256 : offset + 384],
+            torch.zeros(128, 8),
+        )
+
+    proj = torch.arange(8 * 1024, dtype=torch.float32).view(8, 1024)
+    attention.proj.weight.weight_loader(attention.proj.weight, proj)
+    torch.testing.assert_close(attention.proj.weight[:, :256], proj[:, 768:1024])
+    torch.testing.assert_close(attention.proj.weight[:, 256:], torch.zeros(8, 128))
+
+
+def test_glm5next_vision_tp3_mlp_shards_and_zeros_local_tail(
+    monkeypatch, tp3_linear_state, default_vllm_config
+) -> None:
+    monkeypatch.setattr(glm5next_multimodal, "is_vit_use_data_parallel", lambda: False)
+    mlp = glm5next_multimodal.Glm5NextVisionMLP(
+        in_features=8,
+        hidden_features=4098,
+        loaded_hidden_features=4096,
+        swiglu_limit=10.0,
+    )
+
+    assert mlp.gate_up_proj.weight.shape == (2732, 8)
+    assert mlp.down_proj.weight.shape == (8, 1366)
+
+    gate_up = torch.arange(8192 * 8, dtype=torch.float32).view(8192, 8)
+    mlp.gate_up_proj.weight.weight_loader(mlp.gate_up_proj.weight, gate_up)
+    for local_offset, checkpoint_offset in ((0, 0), (1366, 4096)):
+        torch.testing.assert_close(
+            mlp.gate_up_proj.weight[local_offset : local_offset + 1364],
+            gate_up[checkpoint_offset + 2732 : checkpoint_offset + 4096],
+        )
+        torch.testing.assert_close(
+            mlp.gate_up_proj.weight[local_offset + 1364 : local_offset + 1366],
+            torch.zeros(2, 8),
+        )
+
+    down = torch.arange(8 * 4096, dtype=torch.float32).view(8, 4096)
+    mlp.down_proj.weight.weight_loader(mlp.down_proj.weight, down)
+    torch.testing.assert_close(mlp.down_proj.weight[:, :1364], down[:, 2732:4096])
+    torch.testing.assert_close(mlp.down_proj.weight[:, 1364:], torch.zeros(8, 2))
+
+
+def test_glm5next_vision_tp3_merger_shards_only_divisible_weights(
+    monkeypatch, tp3_linear_state, default_vllm_config
+) -> None:
+    class FakeProjection(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+            self.disable_tp = kwargs["disable_tp"]
+
+    monkeypatch.setattr(glm5next_multimodal, "is_vit_use_data_parallel", lambda: False)
+    monkeypatch.setattr(glm5next_multimodal, "ColumnParallelLinear", FakeProjection)
+
+    merger = glm5next_multimodal.Glm5NextPatchMerger(
+        d_model=4,
+        context_dim=10242,
+        loaded_context_dim=10240,
+        swiglu_limit=10.0,
+    )
+
+    assert merger.proj.disable_tp
+    assert merger.gate_up_proj.weight.shape == (6828, 4)
+    assert merger.down_proj.weight.shape == (4, 3414)
+
+    gate = torch.arange(10240 * 4, dtype=torch.float32).view(10240, 4)
+    merger.gate_up_proj.weight.weight_loader(
+        merger.gate_up_proj.weight, gate, loaded_shard_id=0
+    )
+    torch.testing.assert_close(merger.gate_up_proj.weight[:3412], gate[6828:10240])
+    torch.testing.assert_close(merger.gate_up_proj.weight[3412:3414], torch.zeros(2, 4))
+
+    up = gate + 100000
+    merger.gate_up_proj.weight.weight_loader(
+        merger.gate_up_proj.weight, up, loaded_shard_id=1
+    )
+    torch.testing.assert_close(merger.gate_up_proj.weight[3414:6826], up[6828:10240])
+    torch.testing.assert_close(merger.gate_up_proj.weight[6826:6828], torch.zeros(2, 4))
+
+    down = torch.arange(4 * 10240, dtype=torch.float32).view(4, 10240)
+    merger.down_proj.weight.weight_loader(merger.down_proj.weight, down)
+    torch.testing.assert_close(merger.down_proj.weight[:, :3412], down[:, 6828:10240])
+    torch.testing.assert_close(merger.down_proj.weight[:, 3412:], torch.zeros(4, 2))
+
+
+def _record_vision_geometry(monkeypatch, vision_config, *, data_parallel: bool, tp: int):
+    recorded = SimpleNamespace(block=None, merger=None, rope=None)
+
+    class FakeModule(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+            self.proj = SimpleNamespace(weight=torch.empty(0))
+
+    class FakeBlock(torch.nn.Module):
+        def __init__(self, **kwargs) -> None:
+            super().__init__()
+            recorded.block = kwargs
+
+    class FakeMerger(torch.nn.Module):
+        def __init__(self, **kwargs) -> None:
+            super().__init__()
+            recorded.merger = kwargs
+
+    def fake_rope(**kwargs):
+        recorded.rope = kwargs
+        return object()
+
+    monkeypatch.setattr(
+        glm5next_multimodal, "is_vit_use_data_parallel", lambda: data_parallel
+    )
+    monkeypatch.setattr(
+        glm5next_multimodal, "get_tensor_model_parallel_world_size", lambda: tp
+    )
+    monkeypatch.setattr(glm5next_multimodal, "Glm5NextVisionPatchEmbed", FakeModule)
+    monkeypatch.setattr(glm5next_multimodal, "Glm5NextVisionBlock", FakeBlock)
+    monkeypatch.setattr(glm5next_multimodal, "Glm5NextPatchMerger", FakeMerger)
+    monkeypatch.setattr(glm5next_multimodal, "Conv2dLayer", FakeModule)
+    monkeypatch.setattr(glm5next_multimodal, "RMSNorm", FakeModule)
+    monkeypatch.setattr(glm5next_multimodal, "get_rope", fake_rope)
+    monkeypatch.setattr(
+        glm5next_multimodal, "get_vit_attn_backend", lambda **kwargs: object()
+    )
+
+    transformer = glm5next_multimodal.Glm5NextVisionTransformer(
+        SimpleNamespace(swiglu_limit=10.0), vision_config
+    )
+    return transformer, recorded
+
+
+def test_glm5next_vision_tp3_consumes_direct_physical_geometry(monkeypatch) -> None:
+    vision_config = SimpleNamespace(
+        patch_size=14,
+        temporal_patch_size=2,
+        in_channels=3,
+        depth=1,
+        hidden_size=1024,
+        num_heads=18,
+        original_num_heads=16,
+        intermediate_size=4098,
+        original_intermediate_size=4096,
+        spatial_merge_size=2,
+        out_hidden_size=4096,
+        projection_intermediate_size=10242,
+        original_projection_intermediate_size=10240,
+        glm53_tp3_attention_projection_size=1152,
+        glm53_tp3_padding=True,
+        rms_norm_eps=1e-6,
+        swiglu_limit=10.0,
+    )
+    before = vars(vision_config).copy()
+
+    transformer, recorded = _record_vision_geometry(
+        monkeypatch, vision_config, data_parallel=False, tp=3
+    )
+
+    assert vars(vision_config) == before
+    assert transformer.tp_size == 3
+    assert transformer.num_heads == 18
+    assert transformer.attention_projection_size == 1152
+    assert recorded.rope["head_size"] == 64
+    assert recorded.block["num_heads"] == 18
+    assert recorded.block["loaded_num_heads"] == 16
+    assert recorded.block["projection_size"] == 1152
+    assert recorded.block["loaded_projection_size"] == 1024
+    assert recorded.block["mlp_hidden_dim"] == 4098
+    assert recorded.block["loaded_mlp_hidden_dim"] == 4096
+    assert recorded.merger["context_dim"] == 10242
+    assert recorded.merger["loaded_context_dim"] == 10240
+
+
+@pytest.mark.parametrize(
+    ("data_parallel", "tp", "expected_tp"),
+    [(True, 3, 1), (False, 4, 4)],
+)
+def test_glm5next_vision_unpadded_modes_are_exact_geometry_noops(
+    monkeypatch, data_parallel: bool, tp: int, expected_tp: int
+) -> None:
+    vision_config = SimpleNamespace(
+        patch_size=14,
+        temporal_patch_size=2,
+        in_channels=3,
+        depth=1,
+        hidden_size=1024,
+        num_heads=16,
+        intermediate_size=4096,
+        spatial_merge_size=2,
+        out_hidden_size=4096,
+        projection_intermediate_size=10240,
+        rms_norm_eps=1e-6,
+        swiglu_limit=10.0,
+    )
+    before = vars(vision_config).copy()
+
+    transformer, recorded = _record_vision_geometry(
+        monkeypatch, vision_config, data_parallel=data_parallel, tp=tp
+    )
+
+    assert vars(vision_config) == before
+    assert transformer.tp_size == expected_tp
+    assert transformer.num_heads == 16
+    assert transformer.attention_projection_size == 1024
+    assert recorded.rope["head_size"] == 64
+    assert recorded.block["num_heads"] == 16
+    assert recorded.block["loaded_num_heads"] is None
+    assert recorded.block["projection_size"] == 1024
+    assert recorded.block["loaded_projection_size"] is None
+    assert recorded.block["mlp_hidden_dim"] == 4096
+    assert recorded.block["loaded_mlp_hidden_dim"] is None
+    assert recorded.merger["context_dim"] == 10240
+    assert recorded.merger["loaded_context_dim"] is None
diff --git a/vllm/models/glm5next/nvidia/multimodal.py b/vllm/models/glm5next/nvidia/multimodal.py
index 951fef21d876..6b6114bd34f3 100644
--- a/vllm/models/glm5next/nvidia/multimodal.py
+++ b/vllm/models/glm5next/nvidia/multimodal.py
@@ -80,6 +80,7 @@ def __init__(
         in_features: int,
         hidden_features: int,
         swiglu_limit: float,
+        loaded_hidden_features: int | None = None,
         bias: bool = True,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
@@ -93,6 +94,11 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.gate_up_proj",
             disable_tp=use_data_parallel,
+            **(
+                {"loaded_output_sizes": [loaded_hidden_features] * 2}
+                if loaded_hidden_features is not None
+                else {}
+            ),
         )
         self.down_proj = RowParallelLinear(
             hidden_features,
@@ -101,6 +107,11 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
             disable_tp=use_data_parallel,
+            **(
+                {"loaded_input_size": loaded_hidden_features}
+                if loaded_hidden_features is not None
+                else {}
+            ),
         )
         # GLM-5.3-Flash clamps the vision SwiGLU gate/up unlike GLM-OCR/GLM-4V.
         self.act_fn = SiluAndMulWithClamp(swiglu_limit=swiglu_limit)
@@ -118,6 +129,8 @@ def __init__(
         embed_dim: int,
         num_heads: int,
         projection_size: int,
+        loaded_num_heads: int | None = None,
+        loaded_projection_size: int | None = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
@@ -132,11 +145,21 @@ def __init__(
         self.hidden_size_per_attention_head = dist_utils.divide(
             projection_size, num_heads
         )
+        if loaded_projection_size is not None:
+            checkpoint_head_dim = dist_utils.divide(
+                loaded_projection_size,
+                loaded_num_heads if loaded_num_heads is not None else num_heads,
+            )
+            if checkpoint_head_dim != self.hidden_size_per_attention_head:
+                raise ValueError(
+                    "Runtime and checkpoint vision attention head dimensions differ: "
+                    f"{self.hidden_size_per_attention_head} != {checkpoint_head_dim}"
+                )
         self.num_attention_heads_per_partition = dist_utils.divide(
             num_heads, self.tp_size
         )
 
-        self.head_dim = embed_dim // num_heads
+        self.head_dim = self.hidden_size_per_attention_head
 
         # q/k norm eps hard-coded 1e-5 — distinct from block/post norm eps.
         self.q_norm = RMSNorm(self.head_dim, eps=1e-5)
@@ -147,6 +170,14 @@ def __init__(
             head_size=self.hidden_size_per_attention_head,
             total_num_heads=num_heads,
             total_num_kv_heads=num_heads,
+            **(
+                {
+                    "loaded_total_num_heads": loaded_num_heads,
+                    "loaded_total_num_kv_heads": loaded_num_heads,
+                }
+                if loaded_num_heads is not None
+                else {}
+            ),
             bias=True,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj" if quant_config else f"{prefix}.qkv",
@@ -159,6 +190,11 @@ def __init__(
             prefix=f"{prefix}.proj",
             bias=True,
             disable_tp=use_data_parallel,
+            **(
+                {"loaded_input_size": loaded_projection_size}
+                if loaded_projection_size is not None
+                else {}
+            ),
         )
 
         self.attn = MMEncoderAttention(
@@ -235,6 +271,10 @@ def __init__(
         dim: int,
         num_heads: int,
         mlp_hidden_dim: int,
+        projection_size: int,
+        loaded_num_heads: int | None,
+        loaded_mlp_hidden_dim: int | None,
+        loaded_projection_size: int | None,
         swiglu_limit: float,
         norm_layer: partial[nn.Module] | None = None,
         quant_config: QuantizationConfig | None = None,
@@ -248,13 +288,16 @@ def __init__(
         self.attn = Glm5NextVisionAttention(
             embed_dim=dim,
             num_heads=num_heads,
-            projection_size=dim,
+            projection_size=projection_size,
+            loaded_num_heads=loaded_num_heads,
+            loaded_projection_size=loaded_projection_size,
             quant_config=quant_config,
             prefix=f"{prefix}.attn",
         )
         self.mlp = Glm5NextVisionMLP(
             dim,
             mlp_hidden_dim,
+            loaded_hidden_features=loaded_mlp_hidden_dim,
             swiglu_limit=swiglu_limit,
             bias=True,
             quant_config=quant_config,
@@ -287,13 +330,17 @@ def __init__(
         d_model: int,
         context_dim: int,
         swiglu_limit: float,
+        loaded_context_dim: int | None = None,
         quant_config: QuantizationConfig | None = None,
         bias: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
         use_data_parallel = is_vit_use_data_parallel()
+        tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
         self.hidden_size = d_model
+        # The 4096-wide projection cannot be divided across TP3. Keep only this
+        # comparatively small projection replicated; the padded merger MLP is sharded.
         self.proj = ColumnParallelLinear(
             self.hidden_size,
             self.hidden_size,
@@ -301,7 +348,7 @@ def __init__(
             gather_output=True,
             quant_config=quant_config,
             prefix=f"{prefix}.proj",
-            disable_tp=use_data_parallel,
+            disable_tp=use_data_parallel or self.hidden_size % tp_size != 0,
         )
         self.post_projection_norm = nn.LayerNorm(self.hidden_size)
         self.gate_up_proj = MergedColumnParallelLinear(
@@ -311,6 +358,11 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.gate_up_proj",
             disable_tp=use_data_parallel,
+            **(
+                {"loaded_output_sizes": [loaded_context_dim] * 2}
+                if loaded_context_dim is not None
+                else {}
+            ),
         )
         self.down_proj = RowParallelLinear(
             context_dim,
@@ -319,6 +371,11 @@ def __init__(
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
             disable_tp=use_data_parallel,
+            **(
+                {"loaded_input_size": loaded_context_dim}
+                if loaded_context_dim is not None
+                else {}
+            ),
         )
         # GLM-5.3-Flash also clamps the merger SwiGLU.
         self.act_fn = SiluAndMulWithClamp(swiglu_limit=swiglu_limit)
@@ -366,6 +423,25 @@ def __init__(
         self.hidden_size = vision_config.hidden_size
         self.num_heads = vision_config.num_heads
 
+        # The config pass writes these fields only for weights-mode TP3.
+        # Without its marker, keep the original TP/data-parallel construction.
+        padded_tp3 = bool(getattr(vision_config, "glm53_tp3_padding", False))
+        if padded_tp3:
+            projection_size = vision_config.glm53_tp3_attention_projection_size
+            loaded_num_heads = vision_config.original_num_heads
+            loaded_projection_size = self.hidden_size
+            loaded_intermediate_size = vision_config.original_intermediate_size
+            loaded_projection_intermediate_size = (
+                vision_config.original_projection_intermediate_size
+            )
+        else:
+            projection_size = self.hidden_size
+            loaded_num_heads = None
+            loaded_projection_size = None
+            loaded_intermediate_size = None
+            loaded_projection_intermediate_size = None
+        self.attention_projection_size = projection_size
+
         self.patch_size = vision_config.patch_size
         self.spatial_merge_size = vision_config.spatial_merge_size
         self.out_hidden_size = vision_config.out_hidden_size
@@ -387,7 +463,7 @@ def __init__(
         )
 
         norm_layer = partial(RMSNorm, eps=norm_eps)
-        head_dim = self.hidden_size // self.num_heads
+        head_dim = projection_size // self.num_heads
         self.rotary_pos_emb = get_rope(
             head_size=head_dim,
             max_position=8192,
@@ -399,7 +475,11 @@ def __init__(
                 Glm5NextVisionBlock(
                     dim=self.hidden_size,
                     num_heads=self.num_heads,
+                    projection_size=projection_size,
+                    loaded_num_heads=loaded_num_heads,
+                    loaded_projection_size=loaded_projection_size,
                     mlp_hidden_dim=vision_config.intermediate_size,
+                    loaded_mlp_hidden_dim=loaded_intermediate_size,
                     swiglu_limit=swiglu_limit,
                     norm_layer=norm_layer,
                     quant_config=quant_config,
@@ -412,6 +492,7 @@ def __init__(
         self.merger = Glm5NextPatchMerger(
             d_model=vision_config.out_hidden_size,
             context_dim=vision_config.projection_intermediate_size,
+            loaded_context_dim=loaded_projection_intermediate_size,
             swiglu_limit=swiglu_limit,
             quant_config=quant_config,
             bias=False,
@@ -550,7 +631,7 @@ def prepare_encoder_metadata(
         metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(
             self.attn_backend,
             cu_seqlens,
-            self.hidden_size,
+            self.attention_projection_size,
             self.tp_size,
             device,
         )

From 300eeba7dfdfca4e488048a3118c91c312a664aa Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:20:45 +0200
Subject: [PATCH 06/25] feat(glm53): prove and load direct TP3 runtime

(cherry picked from commit 70e736a53b63b1f9284c7ab411293ba83a960f5d)
---
 tests/models/test_glm5next_model.py           |  42 ++++-
 tests/v1/core/test_contiguous_kv_packing.py   |  31 ++++
 tests/v1/worker/test_gpu_model_runner.py      |  75 ++++++++
 .../layers/mamba/gdn/kimi_gdn_linear_attn.py  | 161 ++++++++++++++++--
 vllm/v1/core/kv_cache_utils.py                |   5 +-
 vllm/v1/worker/gpu/model_runner.py            |   3 +
 vllm/v1/worker/gpu_model_runner.py            |   3 +
 vllm/v1/worker/utils.py                       |  77 +++++++++
 8 files changed, 375 insertions(+), 22 deletions(-)

diff --git a/tests/models/test_glm5next_model.py b/tests/models/test_glm5next_model.py
index b6d3f89f1b27..800e7bbced34 100644
--- a/tests/models/test_glm5next_model.py
+++ b/tests/models/test_glm5next_model.py
@@ -117,6 +117,40 @@ def test_glm5next_checkpoint_weight_name_remapping(
 ) -> None:
     assert _remap_glm5next_weight_name(checkpoint_name) == parameter_name
 
+def test_glm5next_kda_a_log_loader_pads_tp3_tail(monkeypatch) -> None:
+    monkeypatch.setattr(
+        kimi_gdn_linear_attn, "get_tensor_model_parallel_rank", lambda: 2
+    )
+    param = torch.nn.Parameter(torch.full((22,), -1.0))
+    loaded_weight = torch.arange(64, dtype=torch.float32)
+
+    kimi_gdn_linear_attn.a_log_weight_loader(0, logical_size=64)(
+        param, loaded_weight
+    )
+
+    torch.testing.assert_close(param[:20], loaded_weight[44:64])
+    torch.testing.assert_close(param[20:], torch.zeros(2))
+
+
+def test_glm5next_kda_conv_loader_pads_tp3_tail() -> None:
+    param = torch.nn.Parameter(torch.full((132, 1, 3), -1.0))
+    loaded_weight = torch.arange(128, dtype=torch.float32).view(128, 1, 1)
+    loaded_weight = loaded_weight.expand(-1, 1, 3)
+
+    loader = kimi_gdn_linear_attn._make_fused_conv1d_weight_loader(
+        [132, 132, 132],
+        tp_size=3,
+        tp_rank=2,
+        loaded_dims=[128, 128, 128],
+    )
+    loader(param, loaded_weight, loaded_shard_id=1)
+
+    torch.testing.assert_close(param[44:84], loaded_weight[88:128])
+    torch.testing.assert_close(param[84:88], torch.zeros(4, 1, 3))
+    torch.testing.assert_close(param[:44], torch.full((44, 1, 3), -1.0))
+    torch.testing.assert_close(param[88:], torch.full((44, 1, 3), -1.0))
+
+
 
 def test_glm5next_mixed_precision_resolves_fused_attention_projections() -> None:
     quant_config = ModelOptMixedPrecisionConfig.__new__(ModelOptMixedPrecisionConfig)
@@ -1186,7 +1220,7 @@ def plan(caps):
 
 
 def test_b12x_kda_binds_live_invocations_and_shares_metadata(monkeypatch) -> None:
-    calls: dict[str, list] = {"bind": [], "run": []}
+    calls: dict[str, list] = {"bind": [], "retain": [], "run": []}
 
     class FakeApi:
         @staticmethod
@@ -1205,6 +1239,11 @@ def run_kda(binding, **kwargs):
         "get_forward_context",
         lambda: forward_context,
     )
+    monkeypatch.setattr(
+        kimi_gdn_linear_attn,
+        "retain_cuda_graph_capture_resource",
+        calls["retain"].append,
+    )
 
     plan = SimpleNamespace(caps=SimpleNamespace(max_state_slots=32))
     api = FakeApi()
@@ -1256,6 +1295,7 @@ def make_layer():
 
     assert len(calls["bind"]) == 2
     assert len(calls["run"]) == 2
+    assert calls["retain"] == calls["bind"]
     for binding, output in zip(calls["bind"], outputs):
         assert binding.mixed_qkv is mixed_qkv
         assert binding.raw_g is raw_g
diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py
index 00477cd317e5..1949f11ae233 100644
--- a/tests/v1/core/test_contiguous_kv_packing.py
+++ b/tests/v1/core/test_contiguous_kv_packing.py
@@ -113,6 +113,37 @@ def test_glm53_split_env_does_not_pad_other_mla_models(self, monkeypatch):
 
         assert _get_kv_cache_bytes_per_block(groups) == expected
 
+    def test_glm5_next_stride_is_c4_page_aligned_without_split_env(
+        self, monkeypatch
+    ):
+        monkeypatch.delenv("VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE", raising=False)
+        block_size = 3328
+        target = MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=528,
+            dtype=torch.uint8,
+            model_version="glm5_next",
+            page_tail_bytes_per_token=33,
+        )
+        larger_unaligned = MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=562,
+            dtype=torch.uint8,
+        )
+        groups = [
+            KVCacheGroupSpec(["target"], target),
+            KVCacheGroupSpec(["draft"], larger_unaligned),
+        ]
+        raw_bytes = larger_unaligned.page_size_bytes
+        c4_page_bytes = 64 * 132
+
+        assert raw_bytes % c4_page_bytes
+        assert _get_kv_cache_bytes_per_block(groups) == (
+            (raw_bytes + c4_page_bytes - 1) // c4_page_bytes * c4_page_bytes
+        )
+
     def test_layers_within_a_group_are_dense(self):
         groups, _, _ = _mixed_page_groups()
         pages = _pages(groups)
diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py
index 28503fefc241..47cca7c258d5 100644
--- a/tests/v1/worker/test_gpu_model_runner.py
+++ b/tests/v1/worker/test_gpu_model_runner.py
@@ -11,6 +11,7 @@
 import torch
 
 import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module
+import vllm.v1.worker.utils as worker_utils
 from vllm.config import (
     AttentionConfig,
     CacheConfig,
@@ -79,6 +80,80 @@ def _restore_default_dtype():
     torch.set_default_dtype(old)
 
 
+@pytest.mark.parametrize("collective_enabled", (True, False))
+def test_glm53_r17_tp3_runtime_proof_is_observed_and_fail_closed(
+    monkeypatch: pytest.MonkeyPatch,
+    collective_enabled: bool,
+) -> None:
+    class Glm5NextLinearAttention(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self._b12x_kda_api = object()
+            self.kda_prefill_backend = "flashkda"
+
+    messages = []
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            multimodal_config=SimpleNamespace(mm_encoder_tp_mode="weights"),
+        ),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=3,
+            enable_expert_parallel=True,
+        ),
+    )
+    model = torch.nn.Sequential(Glm5NextLinearAttention())
+    b12x_ar = object.__new__(worker_utils.B12xPcieAllReduce)
+    b12x_ar.disabled = not collective_enabled
+    b12x_ar.world_size = 3
+    b12x_ar._runtime = object()
+    b12x_ar.allreduce_max_bytes = 65536
+    monkeypatch.setenv("GLM53_R17_REQUIRE_RUNTIME_PROOF", "1")
+    monkeypatch.setattr(worker_utils, "is_glm53_config", lambda _: True)
+    monkeypatch.setattr(
+        worker_utils,
+        "get_ep_group",
+        lambda: SimpleNamespace(world_size=3),
+    )
+    monkeypatch.setattr(
+        worker_utils,
+        "get_tp_group",
+        lambda: SimpleNamespace(
+            device_communicator=SimpleNamespace(b12x_ar_comm=b12x_ar)
+        ),
+    )
+    monkeypatch.setattr(
+        worker_utils.logger,
+        "info_once",
+        lambda message, payload, **kwargs: messages.append(message % payload),
+    )
+
+    if not collective_enabled:
+        with pytest.raises(RuntimeError, match="runtime proof failed"):
+            worker_utils.log_glm53_r17_tp3_runtime_proof(vllm_config, model)
+        return
+
+    worker_utils.log_glm53_r17_tp3_runtime_proof(vllm_config, model)
+    assert messages == [
+        "GLM53_R17_TP3_RUNTIME_PROOF "
+        '{"collective_backend":"b12x_pcie_oneshot","expert_parallel_size":3,'
+        '"kda_decode_backend":"b12x","kda_prefill_backend":"flashkda",'
+        '"mm_encoder_tp_mode":"weights"}'
+    ]
+
+
+def test_glm53_r17_runtime_proof_does_not_change_tp4(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("GLM53_R17_REQUIRE_RUNTIME_PROOF", "1")
+    monkeypatch.setattr(worker_utils, "is_glm53_config", pytest.fail)
+    vllm_config = SimpleNamespace(
+        model_config=pytest.fail,
+        parallel_config=SimpleNamespace(tensor_parallel_size=4),
+    )
+
+    worker_utils.log_glm53_r17_tp3_runtime_proof(vllm_config, pytest.fail)
+
+
 def initialize_kv_cache(runner: GPUModelRunner):
     """
     Only perform necessary steps in GPUModelRunner.initialize_kv_cache()
diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
index 0cf33f772949..fa5da17d099f 100644
--- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
@@ -16,10 +16,6 @@
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
-from vllm.model_executor.model_loader.weight_utils import (
-    default_weight_loader,
-    sharded_weight_loader,
-)
 from vllm.model_executor.parameter import BasevLLMParameter
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
@@ -35,7 +31,10 @@
 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.kv_cache_interface import MambaSpec
-from vllm.v1.worker.workspace import current_workspace_manager
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    retain_cuda_graph_capture_resource,
+)
 
 from ...linear import (
     ColumnParallelLinear,
@@ -234,16 +233,61 @@ def resolve_kda_prefill_backend(
     return "triton"
 
 
-def a_log_weight_loader(
+def _load_rank_local_tail(
+    destination: torch.Tensor,
+    loaded_weight: torch.Tensor,
+    dim: int,
+    start_idx: int,
+    shard_size: int,
+    logical_size: int | None,
+) -> None:
+    """Load one shard, zero-filling only a proven logical checkpoint tail."""
+    if logical_size is None or loaded_weight.shape[dim] != logical_size:
+        destination.copy_(loaded_weight.narrow(dim, start_idx, shard_size))
+        return
+
+    available = max(0, min(shard_size, logical_size - start_idx))
+    if available == shard_size:
+        destination.copy_(loaded_weight.narrow(dim, start_idx, shard_size))
+        return
+
+    if destination.shape[dim] != shard_size:
+        raise ValueError(
+            "KDA TP3 tail destination has the wrong shard size: "
+            f"expected {shard_size}, got {destination.shape[dim]}."
+        )
+    destination.zero_()
+    if available:
+        destination.narrow(dim, 0, available).copy_(
+            loaded_weight.narrow(dim, start_idx, available)
+        )
+
+
+def _rank_local_tail_weight_loader(
     shard_axis: int,
+    logical_size: int | None = None,
 ) -> Callable[[torch.Tensor, torch.Tensor], None]:
-    """Load KDA A_log stored as either old 4D or current 1D weights."""
-
     def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
-        tp_rank = get_tensor_model_parallel_rank()
         shard_size = param.data.shape[shard_axis]
-        start_idx = tp_rank * shard_size
+        _load_rank_local_tail(
+            param.data,
+            loaded_weight,
+            shard_axis,
+            get_tensor_model_parallel_rank() * shard_size,
+            shard_size,
+            logical_size,
+        )
+
+    return loader
 
+
+def a_log_weight_loader(
+    shard_axis: int,
+    logical_size: int | None = None,
+) -> Callable[[torch.Tensor, torch.Tensor], None]:
+    rank_local_loader = _rank_local_tail_weight_loader(shard_axis, logical_size)
+
+    def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
         if loaded_weight.dim() == 4:
             assert loaded_weight.shape[:2] == (1, 1), (
                 f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}"
@@ -253,8 +297,7 @@ def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
             )
             loaded_weight = loaded_weight.view(loaded_weight.shape[2])
 
-        loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
-        return default_weight_loader(param, loaded_weight)
+        rank_local_loader(param, loaded_weight)
 
     return loader
 
@@ -263,8 +306,12 @@ def _make_fused_conv1d_weight_loader(
     dims: list[int],
     tp_size: int,
     tp_rank: int,
+    loaded_dims: list[int] | None = None,
 ) -> Callable[..., None]:
     sharded_dims = [dim // tp_size for dim in dims]
+    loaded_dims = dims if loaded_dims is None else loaded_dims
+    if len(loaded_dims) != len(dims):
+        raise ValueError("loaded_dims must match the fused convolution shard count")
 
     def weight_loader(
         param: torch.Tensor,
@@ -276,8 +323,15 @@ def weight_loader(
         shard_size = sharded_dims[loaded_shard_id]
         source_start = tp_rank * shard_size
         target_start = sum(sharded_dims[:loaded_shard_id])
-        loaded_shard = loaded_weight[source_start : source_start + shard_size]
-        param.data[target_start : target_start + shard_size].copy_(loaded_shard)
+        destination = param.data.narrow(0, target_start, shard_size)
+        _load_rank_local_tail(
+            destination,
+            loaded_weight,
+            0,
+            source_start,
+            shard_size,
+            loaded_dims[loaded_shard_id],
+        )
 
     return weight_loader
 
@@ -459,10 +513,27 @@ def __init__(
         assert kda_config is not None, "linear_attn_config must be set"
         self.head_dim = kda_config["head_dim"]
         self.num_heads = kda_config["num_heads"]
+        self._glm53_tp3_padding = bool(getattr(config, "glm53_tp3_padding", False))
+        self.logical_num_heads = int(
+            getattr(config, "original_linear_num_heads", self.num_heads)
+        )
+        if self._glm53_tp3_padding and (
+            self.tp_size != 3
+            or self.logical_num_heads != 64
+            or self.num_heads != 66
+        ):
+            raise ValueError(
+                "GLM-5.3 KDA TP3 padding requires logical64, physical66, "
+                f"and TP3; got logical{self.logical_num_heads}, "
+                f"physical{self.num_heads}, TP{self.tp_size}."
+            )
+        if not self._glm53_tp3_padding:
+            self.logical_num_heads = self.num_heads
         assert self.num_heads % self.tp_size == 0
         self.local_num_heads = divide(self.num_heads, self.tp_size)
 
         self.projection_size = self.head_dim * self.num_heads
+        self.logical_projection_size = self.head_dim * self.logical_num_heads
         self.local_projection_size = divide(self.projection_size, self.tp_size)
         self.conv_size = kda_config["short_conv_kernel_size"]
         self.use_full_rank_gate = kda_config.get("use_full_rank_gate", False)
@@ -476,17 +547,27 @@ def __init__(
                 self.head_dim,
                 self.num_heads,
             ]
+            loaded_in_proj_output_sizes = [self.logical_projection_size] * 4 + [
+                self.head_dim,
+                self.logical_num_heads,
+            ]
             local_output_size = (
                 4 * self.local_projection_size + self.head_dim + self.local_num_heads
             )
             self.in_proj_padding = -local_output_size % 16
             if self.in_proj_padding:
-                in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
+                padding_size = self.in_proj_padding * self.tp_size
+                in_proj_output_sizes.append(padding_size)
+                loaded_in_proj_output_sizes.append(padding_size)
         else:
             in_proj_output_sizes = [self.projection_size] * 3 + [
                 self.num_heads,
                 self.head_dim,
             ]
+            loaded_in_proj_output_sizes = [self.logical_projection_size] * 3 + [
+                self.logical_num_heads,
+                self.head_dim,
+            ]
             self.in_proj_padding = 0
         self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
             self.hidden_size,
@@ -496,6 +577,11 @@ def __init__(
             bias=False,
             quant_config=self.quant_config,
             prefix=f"{prefix}.in_proj_qkvgfab",
+            **(
+                {"loaded_output_sizes": loaded_in_proj_output_sizes}
+                if self._glm53_tp3_padding
+                else {}
+            ),
         )
         if self.in_proj_padding:
             self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
@@ -506,12 +592,29 @@ def __init__(
             bias=False,
             quant_config=self.quant_config,
             prefix=f"{prefix}.f_b_proj",
+            **(
+                {"loaded_output_size": self.logical_projection_size}
+                if self._glm53_tp3_padding
+                else {}
+            ),
         )
         self.dt_bias = nn.Parameter(
             torch.empty(self.local_projection_size, dtype=torch.float32)
         )
 
-        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
+        set_weight_attrs(
+            self.dt_bias,
+            {
+                "weight_loader": _rank_local_tail_weight_loader(
+                    0,
+                    (
+                        self.logical_projection_size
+                        if self._glm53_tp3_padding
+                        else None
+                    ),
+                )
+            },
+        )
 
         # One packed parameter and cache let decode run a single conv update.
         # Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
@@ -531,6 +634,11 @@ def __init__(
                     [self.projection_size] * 3,
                     self.tp_size,
                     self.tp_rank,
+                    (
+                        [self.logical_projection_size] * 3
+                        if self._glm53_tp3_padding
+                        else None
+                    ),
                 )
             },
         )
@@ -538,7 +646,15 @@ def __init__(
         self.A_log = nn.Parameter(
             torch.empty(self.local_num_heads, dtype=torch.float32)
         )
-        set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
+        set_weight_attrs(
+            self.A_log,
+            {
+                "weight_loader": a_log_weight_loader(
+                    0,
+                    self.logical_num_heads if self._glm53_tp3_padding else None,
+                )
+            },
+        )
 
         self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
         if self.gate_lower_bound is not None:
@@ -601,6 +717,11 @@ def __init__(
                 bias=False,
                 quant_config=self.quant_config,
                 prefix=f"{prefix}.g_b_proj",
+                **(
+                    {"loaded_output_size": self.logical_projection_size}
+                    if self._glm53_tp3_padding
+                    else {}
+                ),
             )
         self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
         self._b12x_kda_api: Any | None = None
@@ -615,6 +736,11 @@ def __init__(
             bias=False,
             quant_config=self.quant_config,
             prefix=f"{prefix}.o_proj",
+            **(
+                {"loaded_input_size": self.logical_projection_size}
+                if self._glm53_tp3_padding
+                else {}
+            ),
         )
 
         compilation_config = vllm_config.compilation_config
@@ -1063,6 +1189,7 @@ def _run_b12x_kda_decode_post_conv(
             num_tokens=num_tokens_tensor,
             output=output,
         )
+        retain_cuda_graph_capture_resource(binding)
         api.run_kda(
             binding,
             lower_bound=self.gate_lower_bound,
diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py
index 85df272c04bb..d4c507fba89f 100644
--- a/vllm/v1/core/kv_cache_utils.py
+++ b/vllm/v1/core/kv_cache_utils.py
@@ -1346,10 +1346,7 @@ def _get_kv_cache_bytes_per_block(
         for group in kv_cache_groups
         for layer_name in group.layer_names
     )
-    if (
-        os.getenv("VLLM_GLM53_SPLIT_TARGET_BLOCK_SIZE") is not None
-        and contains_glm5_next_mla
-    ):
+    if contains_glm5_next_mla:
         # GLM-5.3 stores two 64-row by 132-byte FP8 C4 index pages in the
         # target MLA page tail when the target block contains 512 tokens.
         # The block-outermost pool stride must preserve the index-page unit
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index 4c5e09ed6799..e1d066228bd5 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -160,6 +160,7 @@
 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
 from vllm.v1.worker.utils import (
     KVBlockZeroer,
+    log_glm53_r17_tp3_runtime_proof,
     copy_kv_cache_blocks_inplace,
     get_uniform_decode_token_count,
 )
@@ -411,6 +412,8 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:
             format_gib(m.consumed_memory),
             time_after_load - time_before_load,
         )
+        log_glm53_r17_tp3_runtime_proof(self.vllm_config, self.model)
+
 
         # Initialize the components that require the model.
         self.model_state = init_model_state(
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index e80af9842dbe..2f23942058b1 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -235,6 +235,7 @@
 from vllm.v1.worker.utils import (
     EncoderTimingStats,
     is_residual_scattered_for_sp,
+    log_glm53_r17_tp3_runtime_proof,
     raise_if_nan_logits,
 )
 from vllm.v1.worker.workspace import lock_workspace
@@ -5520,6 +5521,8 @@ def load_model(self, load_dummy_weights: bool = False) -> None:
             format_gib(self.model_memory_usage),
             time_after_load - time_before_load,
         )
+        log_glm53_r17_tp3_runtime_proof(self.vllm_config, self.model)
+
 
         mm_config = self.model_config.multimodal_config
         self.is_multimodal_pruning_enabled = (
diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py
index 7ed3e5bcd97a..f3a2a103471a 100644
--- a/vllm/v1/worker/utils.py
+++ b/vllm/v1/worker/utils.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
 import math
+import os
 from collections import defaultdict
 from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass, field
@@ -11,11 +13,16 @@
 import torch
 
 from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed.device_communicators.b12x_pcie_all_reduce import (
+    B12xPcieAllReduce,
+)
+from vllm.distributed.parallel_state import get_ep_group, get_tp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.models.interfaces import MultiModalEmbeddings
 from vllm.model_executor.models.utils import extract_layer_index
 from vllm.platforms import current_platform
+from vllm.transformers_utils.configs.glm53_tp3 import is_glm53_config
 from vllm.triton_utils import tl, triton
 from vllm.utils.mem_utils import MemorySnapshot, format_gib
 from vllm.utils.torch_utils import async_tensor_h2d
@@ -41,6 +48,76 @@
 logger = init_logger(__name__)
 
 
+def log_glm53_r17_tp3_runtime_proof(vllm_config: VllmConfig, model: Any) -> None:
+    """Emit the required resolved runtime receipt for GLM-5.3 R17 TP3."""
+    parallel_config = vllm_config.parallel_config
+    model_config = vllm_config.model_config
+    if (
+        os.environ.get("GLM53_R17_REQUIRE_RUNTIME_PROOF") != "1"
+        or parallel_config.tensor_parallel_size != 3
+        or not is_glm53_config(model_config)
+    ):
+        return
+
+    kda_layers = [
+        module
+        for module in model.modules()
+        if type(module).__name__ == "Glm5NextLinearAttention"
+    ]
+    kda_prefill_backends = {
+        getattr(layer, "kda_prefill_backend", None) for layer in kda_layers
+    }
+    device_communicator = get_tp_group().device_communicator
+    b12x_ar = getattr(device_communicator, "b12x_ar_comm", None)
+    collective_backend = (
+        "b12x_pcie_oneshot"
+        if isinstance(b12x_ar, B12xPcieAllReduce)
+        and not b12x_ar.disabled
+        and b12x_ar.world_size == 3
+        and b12x_ar._runtime is not None
+        and b12x_ar.allreduce_max_bytes > 0
+        else "fallback"
+    )
+    mm_config = model_config.multimodal_config
+    proof = {
+        "collective_backend": collective_backend,
+        "expert_parallel_size": (
+            get_ep_group().world_size if parallel_config.enable_expert_parallel else 1
+        ),
+        "kda_decode_backend": (
+            "b12x"
+            if kda_layers
+            and all(
+                getattr(layer, "_b12x_kda_api", None) is not None
+                for layer in kda_layers
+            )
+            else "fallback"
+        ),
+        "kda_prefill_backend": (
+            next(iter(kda_prefill_backends))
+            if len(kda_prefill_backends) == 1
+            and kda_prefill_backends <= {"flashkda", "triton"}
+            else "fallback"
+        ),
+        "mm_encoder_tp_mode": (
+            getattr(mm_config, "mm_encoder_tp_mode", None)
+            if mm_config is not None
+            else None
+        ),
+    }
+    expected = {
+        "collective_backend": "b12x_pcie_oneshot",
+        "expert_parallel_size": 3,
+        "kda_decode_backend": "b12x",
+        "kda_prefill_backend": "flashkda",
+        "mm_encoder_tp_mode": "weights",
+    }
+    payload = json.dumps(proof, sort_keys=True, separators=(",", ":"))
+    if proof != expected:
+        raise RuntimeError(f"GLM-5.3 R17 TP3 runtime proof failed: {payload}")
+    logger.info_once("GLM53_R17_TP3_RUNTIME_PROOF %s", payload, scope="global")
+
+
 def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None:
     if not any(num_nans_in_logits.values()):
         return

From 39054db882ae936f3c5cfa153f6aa0386d36db47 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:23:33 +0200
Subject: [PATCH 07/25] test(glm5next): keep TP3 vision loader checks CPU-only

(cherry picked from commit 391f880f64aee8a1c0508fb0b6ae421c7a57c424)
---
 tests/models/test_glm5next_vision_tp3.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/models/test_glm5next_vision_tp3.py b/tests/models/test_glm5next_vision_tp3.py
index 65ab8250bc90..294535f6035e 100644
--- a/tests/models/test_glm5next_vision_tp3.py
+++ b/tests/models/test_glm5next_vision_tp3.py
@@ -28,7 +28,7 @@ def tp3_linear_state(monkeypatch):
 
 
 def test_glm5next_vision_tp3_attention_shards_and_zeros_local_tail(
-    monkeypatch, tp3_linear_state, default_vllm_config
+    monkeypatch, tp3_linear_state
 ) -> None:
     class FakeEncoderAttention(torch.nn.Module):
         def __init__(self, **kwargs) -> None:
@@ -75,7 +75,7 @@ def __init__(self, **kwargs) -> None:
 
 
 def test_glm5next_vision_tp3_mlp_shards_and_zeros_local_tail(
-    monkeypatch, tp3_linear_state, default_vllm_config
+    monkeypatch, tp3_linear_state
 ) -> None:
     monkeypatch.setattr(glm5next_multimodal, "is_vit_use_data_parallel", lambda: False)
     mlp = glm5next_multimodal.Glm5NextVisionMLP(
@@ -107,7 +107,7 @@ def test_glm5next_vision_tp3_mlp_shards_and_zeros_local_tail(
 
 
 def test_glm5next_vision_tp3_merger_shards_only_divisible_weights(
-    monkeypatch, tp3_linear_state, default_vllm_config
+    monkeypatch, tp3_linear_state
 ) -> None:
     class FakeProjection(torch.nn.Module):
         def __init__(self, *args, **kwargs) -> None:

From d1738321a6159d2d869d5890c3f5c8e8be12870a Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:25:26 +0200
Subject: [PATCH 08/25] feat(glm5next): consume direct TP3 model geometry

(cherry picked from commit fa9bbdf3cae151eb74eb116879d2f761ee2726c8)
---
 tests/models/test_glm53_tp3_model.py     | 353 +++++++++++++++++++++++
 vllm/model_executor/layers/linear.py     |   6 +-
 vllm/models/glm5next/nvidia/attention.py |  20 ++
 vllm/models/glm5next/nvidia/model.py     |  41 ++-
 vllm/models/glm5next/nvidia/mtp.py       |  56 +++-
 5 files changed, 467 insertions(+), 9 deletions(-)
 create mode 100644 tests/models/test_glm53_tp3_model.py

diff --git a/tests/models/test_glm53_tp3_model.py b/tests/models/test_glm53_tp3_model.py
new file mode 100644
index 000000000000..91effebf5923
--- /dev/null
+++ b/tests/models/test_glm53_tp3_model.py
@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor.layers import linear
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor import parameter
+from vllm.models.glm5next.nvidia import attention as glm_attention
+from vllm.models.glm5next.nvidia import model as glm_model
+from vllm.models.glm5next.nvidia import mtp as glm_mtp
+from vllm.models.glm5next.nvidia.mtp import Glm5NextMultiTokenPredictorLayer
+
+
+def _set_tp3_rank2(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 3)
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 2)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 3)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 2)
+
+
+def test_explicit_loaded_sizes_zero_rank_local_destination_tails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _set_tp3_rank2(monkeypatch)
+
+    column = ColumnParallelLinear(
+        1, 6, bias=False, loaded_output_size=5
+    )
+    column.weight.weight_loader(
+        column.weight, torch.arange(1, 6, dtype=column.weight.dtype).unsqueeze(1)
+    )
+    torch.testing.assert_close(
+        column.weight[:, 0], column.weight.new_tensor([5, 0])
+    )
+
+    merged = MergedColumnParallelLinear(
+        1, [6, 6], bias=False, loaded_output_sizes=[5, 5]
+    )
+    merged.weight.weight_loader(
+        merged.weight,
+        torch.arange(1, 6, dtype=merged.weight.dtype).unsqueeze(1),
+        0,
+    )
+    torch.testing.assert_close(
+        merged.weight[:2, 0], merged.weight.new_tensor([5, 0])
+    )
+
+    qkv = QKVParallelLinear(
+        hidden_size=1,
+        head_size=1,
+        total_num_heads=6,
+        total_num_kv_heads=3,
+        loaded_total_num_heads=4,
+        loaded_total_num_kv_heads=2,
+        bias=False,
+    )
+    for shard_id, size in (("q", 4), ("k", 2), ("v", 2)):
+        qkv.weight.weight_loader(
+            qkv.weight,
+            torch.ones((size, 1), dtype=qkv.weight.dtype),
+            shard_id,
+        )
+    torch.testing.assert_close(qkv.weight[:, 0], qkv.weight.new_zeros(4))
+
+    row = RowParallelLinear(
+        6, 1, bias=False, loaded_input_size=5
+    )
+    row.weight.weight_loader(
+        row.weight, torch.arange(1, 6, dtype=row.weight.dtype).unsqueeze(0)
+    )
+    torch.testing.assert_close(
+        row.weight[0], row.weight.new_tensor([5, 0])
+    )
+
+
+def test_loaded_sizes_reject_invalid_physical_layout(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _set_tp3_rank2(monkeypatch)
+    with pytest.raises(ValueError, match="exceeds physical size"):
+        ColumnParallelLinear(1, 6, bias=False, loaded_output_size=7)
+    with pytest.raises(ValueError, match="same length"):
+        MergedColumnParallelLinear(
+            1, [6, 6], bias=False, loaded_output_sizes=[5]
+        )
+    with pytest.raises(ValueError, match="exceeds physical size"):
+        RowParallelLinear(6, 1, bias=False, loaded_input_size=7)
+
+
+@pytest.mark.parametrize("tp3", [False, True])
+def test_mla_projection_loaded_sizes_are_tp3_only(
+    monkeypatch: pytest.MonkeyPatch,
+    tp3: bool,
+) -> None:
+    captured: dict[str, dict] = {}
+
+    class FakeLinear(torch.nn.Module):
+        def __init__(self, *args, prefix: str, **kwargs) -> None:
+            super().__init__()
+            captured[prefix] = kwargs
+
+    class FakeModule(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+
+    monkeypatch.setattr(glm_attention, "ColumnParallelLinear", FakeLinear)
+    monkeypatch.setattr(glm_attention, "RowParallelLinear", FakeLinear)
+    monkeypatch.setattr(glm_attention, "DeepSeekV2FusedQkvAProjLinear", FakeLinear)
+    monkeypatch.setattr(glm_attention, "RMSNorm", FakeModule)
+    monkeypatch.setattr(
+        glm_attention, "MultiHeadLatentAttentionWrapper", FakeModule
+    )
+    monkeypatch.setattr(
+        glm_attention,
+        "get_tensor_model_parallel_world_size",
+        lambda: 3 if tp3 else 4,
+    )
+
+    config = SimpleNamespace(
+        rms_norm_eps=1e-5,
+        rope_parameters=None,
+        index_topk=None,
+        glm53_tp3_padding=tp3,
+    )
+    if tp3:
+        config.original_num_attention_heads = 64
+
+    glm_attention.Glm5NextMLAAttention(
+        vllm_config=SimpleNamespace(),
+        config=config,
+        hidden_size=4096,
+        num_heads=72 if tp3 else 64,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        v_head_dim=128,
+        q_lora_rank=1536,
+        kv_lora_rank=512,
+        skip_rope=True,
+        prefix="attn",
+    )
+
+    loaded_key = "loaded_output_size"
+    if tp3:
+        assert captured["attn.q_b_proj"][loaded_key] == 64 * 192
+        assert captured["attn.kv_b_proj"][loaded_key] == 64 * 256
+        assert captured["attn.o_proj"]["loaded_input_size"] == 64 * 128
+    else:
+        assert loaded_key not in captured["attn.q_b_proj"]
+        assert loaded_key not in captured["attn.kv_b_proj"]
+        assert "loaded_input_size" not in captured["attn.o_proj"]
+
+
+def test_shared_expert_uses_physical_tp3_width_and_logical_load_width(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = {}
+
+    class FakeMerged(torch.nn.Module):
+        def __init__(self, input_size, output_sizes, **kwargs) -> None:
+            super().__init__()
+            calls["gate_up"] = (input_size, output_sizes, kwargs)
+
+    class FakeRow(torch.nn.Module):
+        def __init__(self, input_size, output_size, **kwargs) -> None:
+            super().__init__()
+            calls["down"] = (input_size, output_size, kwargs)
+
+    monkeypatch.setattr(glm_model, "MergedColumnParallelLinear", FakeMerged)
+    monkeypatch.setattr(glm_model, "RowParallelLinear", FakeRow)
+
+    glm_model.Glm5NextMLP(
+        hidden_size=4096,
+        intermediate_size=2112,
+        loaded_intermediate_size=2048,
+        hidden_act="silu",
+    )
+
+    assert calls["gate_up"][1] == [2112, 2112]
+    assert calls["gate_up"][2]["loaded_output_sizes"] == [2048, 2048]
+    assert calls["down"][0] == 2112
+    assert calls["down"][2]["loaded_input_size"] == 2048
+
+
+@pytest.mark.parametrize("tp3", [False, True])
+def test_mtp_tp3_objects_are_created_only_for_active_geometry(
+    monkeypatch: pytest.MonkeyPatch,
+    tp3: bool,
+) -> None:
+    calls = {}
+
+    class FakeNorm(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+
+    class FakeColumn(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+            calls["projection"] = (args, kwargs)
+
+    class FakeSharedHead(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+            calls["standard_shared_head"] = True
+
+    class FakeParallelLMHead(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+            calls["padded_head"] = (args, kwargs)
+
+    class FakeDecoder(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+
+    monkeypatch.setattr(glm_mtp, "RMSNorm", FakeNorm)
+    monkeypatch.setattr(glm_mtp, "ColumnParallelLinear", FakeColumn)
+    monkeypatch.setattr(glm_mtp, "SharedHead", FakeSharedHead)
+    monkeypatch.setattr(glm_mtp, "ParallelLMHead", FakeParallelLMHead)
+    monkeypatch.setattr(glm_mtp, "Glm5NextDecoderLayer", FakeDecoder)
+    monkeypatch.setattr(
+        glm_mtp, "current_platform", SimpleNamespace(device_type="cpu")
+    )
+
+    config = SimpleNamespace(
+        hidden_size=4096,
+        rms_norm_eps=1e-5,
+        index_topk=4,
+        index_kpool=1,
+        vocab_size=154880,
+        glm53_tp3_padding=tp3,
+    )
+    if tp3:
+        config.glm53_tp3_mtp_projection_size = 4098
+        config.glm53_tp3_vocab_padding_size = 192
+    vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(
+            draft_model_config=SimpleNamespace(hf_config=config)
+        ),
+        quant_config=None,
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+    )
+
+    layer = glm_mtp.Glm5NextMultiTokenPredictorLayer(
+        vllm_config, "model.layers.45"
+    )
+
+    if tp3:
+        assert isinstance(layer.eh_proj, FakeColumn)
+        assert calls["projection"][0] == (8192, 4098)
+        assert calls["projection"][1]["loaded_output_size"] == 4096
+        assert calls["projection"][1]["gather_output"]
+        assert calls["padded_head"][1]["padding_size"] == 192
+        assert "standard_shared_head" not in calls
+    else:
+        assert type(layer.eh_proj) is torch.nn.Linear
+        assert layer.eh_proj.in_features == 8192
+        assert layer.eh_proj.out_features == 4096
+        assert calls["standard_shared_head"]
+        assert "projection" not in calls
+        assert "padded_head" not in calls
+
+
+@pytest.mark.parametrize("tp3", [False, True])
+def test_target_vocab_storage_padding_is_tp3_only(
+    monkeypatch: pytest.MonkeyPatch,
+    tp3: bool,
+) -> None:
+    calls = {}
+
+    class FakeModel(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+
+    class FakeHead(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+            calls["head"] = (args, kwargs)
+
+    class FakeLogits:
+        def __init__(self, vocab_size, **kwargs) -> None:
+            calls["logits_vocab_size"] = vocab_size
+
+    monkeypatch.setattr(glm_model, "Glm5NextModel", FakeModel)
+    monkeypatch.setattr(glm_model, "ParallelLMHead", FakeHead)
+    monkeypatch.setattr(glm_model, "LogitsProcessor", FakeLogits)
+    monkeypatch.setattr(
+        glm_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_last_rank=True),
+    )
+
+    config = SimpleNamespace(vocab_size=154880, hidden_size=4096)
+    if tp3:
+        config.glm53_tp3_padding = True
+        config.glm53_tp3_vocab_padding_size = 192
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=config),
+        quant_config=None,
+    )
+
+    glm_model.Glm5NextForCausalLM(vllm_config=vllm_config)
+
+    assert calls["head"][0] == (154880, 4096)
+    assert calls["logits_vocab_size"] == 154880
+    if tp3:
+        assert calls["head"][1]["padding_size"] == 192
+    else:
+        assert "padding_size" not in calls["head"][1]
+
+
+def test_mtp_projection_narrows_padded_output_contiguously() -> None:
+    class PaddedProjection(torch.nn.Module):
+        def forward(self, value: torch.Tensor) -> torch.Tensor:
+            output = value.new_zeros((*value.shape[:-1], 6))
+            output[..., :4] = value[..., :4]
+            return output
+
+    class RecordingBlock(torch.nn.Module):
+        def forward(self, *, hidden_states: torch.Tensor, **kwargs):
+            assert hidden_states.shape == (2, 4)
+            assert hidden_states.is_contiguous()
+            return hidden_states, None, None, None
+
+    class SharedHead(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.norm = lambda hidden_states, residual: (hidden_states, None)
+
+    layer = object.__new__(Glm5NextMultiTokenPredictorLayer)
+    torch.nn.Module.__init__(layer)
+    layer.enorm = torch.nn.Identity()
+    layer.hnorm = torch.nn.Identity()
+    layer.eh_proj = PaddedProjection()
+    layer.mtp_block = RecordingBlock()
+    layer.shared_head = SharedHead()
+
+    hidden_states, recycled = layer(
+        input_ids=torch.zeros(2, dtype=torch.long),
+        positions=torch.arange(2),
+        previous_hidden_states=torch.ones(2, 4),
+        inputs_embeds=torch.ones(2, 4),
+    )
+
+    assert hidden_states.is_contiguous()
+    torch.testing.assert_close(recycled, hidden_states)
diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index 20d48b820ea2..6a9df1e58f1d 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -762,7 +762,7 @@ def __init__(
                 loaded_size,
                 physical_size,
             )
-        self._allow_loaded_output_padding = [
+        self._allow_loaded_output_shard_padding = [
             loaded_output_sizes is not None and loaded_size != physical_size
             for loaded_size, physical_size in zip(
                 self.loaded_output_sizes, self.output_sizes
@@ -915,7 +915,7 @@ def weight_loader(
                     loaded_weight,
                     output_dim,
                     start_idx,
-                    allow_padding=self._allow_loaded_output_padding[
+                    allow_padding=self._allow_loaded_output_shard_padding[
                         loaded_shard_id
                     ],
                 )
@@ -1049,7 +1049,7 @@ def weight_loader_v2(
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
-            allow_padding=self._allow_loaded_output_padding[loaded_shard_id],
+            allow_padding=self._allow_loaded_output_shard_padding[loaded_shard_id],
         )
 
     def load_weights(
diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/nvidia/attention.py
index f3a093a14ff5..33c45d149b67 100644
--- a/vllm/models/glm5next/nvidia/attention.py
+++ b/vllm/models/glm5next/nvidia/attention.py
@@ -81,6 +81,22 @@ def __init__(
         self.scaling = self.qk_head_dim**-0.5
         self.max_position_embeddings = max_position_embeddings
         proj_input_size = input_size if input_size is not None else hidden_size
+        column_head_kwargs = {}
+        row_head_kwargs = {}
+        if getattr(config, "glm53_tp3_padding", False):
+            logical_num_heads = config.original_num_attention_heads
+            column_head_kwargs["loaded_output_size"] = (
+                logical_num_heads * self.qk_head_dim
+            )
+            row_head_kwargs["loaded_input_size"] = (
+                logical_num_heads * self.v_head_dim
+            )
+        kv_column_head_kwargs = {}
+        if getattr(config, "glm53_tp3_padding", False):
+            kv_column_head_kwargs["loaded_output_size"] = logical_num_heads * (
+                qk_nope_head_dim + v_head_dim
+            )
+
 
         if q_lora_rank is not None:
             self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
@@ -96,6 +112,7 @@ def __init__(
                 bias=False,
                 quant_config=quant_config,
                 prefix=f"{prefix}.q_b_proj",
+                **column_head_kwargs,
             )
         else:
             self.kv_a_proj_with_mqa = ReplicatedLinear(
@@ -111,6 +128,7 @@ def __init__(
                 bias=False,
                 quant_config=quant_config,
                 prefix=f"{prefix}.q_proj",
+                **column_head_kwargs,
             )
 
         self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps)
@@ -120,6 +138,7 @@ def __init__(
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.kv_b_proj",
+            **kv_column_head_kwargs,
         )
         self.o_proj = RowParallelLinear(
             num_heads * v_head_dim,
@@ -127,6 +146,7 @@ def __init__(
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
+            **row_head_kwargs,
         )
 
         if not skip_rope:
diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/nvidia/model.py
index c78d910c0b25..075d540d72c5 100644
--- a/vllm/models/glm5next/nvidia/model.py
+++ b/vllm/models/glm5next/nvidia/model.py
@@ -153,8 +153,25 @@ def __init__(
         is_sequence_parallel=False,
         prefix: str = "",
         swiglu_limit: float | None = None,
+        loaded_intermediate_size: int | None = None,
     ) -> None:
         super().__init__()
+        if loaded_intermediate_size is None:
+            loaded_intermediate_size = intermediate_size
+        if not 0 < loaded_intermediate_size <= intermediate_size:
+            raise ValueError(
+                "loaded_intermediate_size must be positive and no greater than "
+                f"intermediate_size, got {loaded_intermediate_size} and "
+                f"{intermediate_size}"
+            )
+        gate_up_kwargs = {}
+        down_kwargs = {}
+        if loaded_intermediate_size != intermediate_size:
+            gate_up_kwargs["loaded_output_sizes"] = [
+                loaded_intermediate_size
+            ] * 2
+            down_kwargs["loaded_input_size"] = loaded_intermediate_size
+
 
         # If is_sequence_parallel, the input and output tensors are sharded
         # across the ranks within the tp_group. In this case the weights are
@@ -167,6 +184,7 @@ def __init__(
             quant_config=quant_config,
             disable_tp=is_sequence_parallel,
             prefix=f"{prefix}.gate_up_proj",
+            **gate_up_kwargs,
         )
         self.down_proj = RowParallelLinear(
             intermediate_size,
@@ -176,6 +194,7 @@ def __init__(
             reduce_results=reduce_results,
             disable_tp=is_sequence_parallel,
             prefix=f"{prefix}.down_proj",
+            **down_kwargs,
         )
         if hidden_act != "silu":
             raise ValueError(
@@ -256,7 +275,14 @@ def __init__(
         if config.n_shared_experts is None:
             self.shared_experts = None
         else:
-            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+            checkpoint_intermediate_size = (
+                config.moe_intermediate_size * config.n_shared_experts
+            )
+            intermediate_size = getattr(
+                config,
+                "glm53_tp3_shared_expert_intermediate_size",
+                checkpoint_intermediate_size,
+            )
 
             self.shared_experts = Glm5NextMLP(
                 hidden_size=config.hidden_size,
@@ -267,6 +293,7 @@ def __init__(
                 reduce_results=False,
                 prefix=f"{prefix}.shared_experts",
                 swiglu_limit=swiglu_limit,
+                loaded_intermediate_size=checkpoint_intermediate_size,
             )
 
         self.experts = FusedMoEFactory(
@@ -928,10 +955,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             pool_topk_indices_buffer = None
 
         if get_pp_group().is_first_rank:
+            vocab_kwargs = {}
+            if getattr(config, "glm53_tp3_padding", False):
+                vocab_kwargs["padding_size"] = (
+                    config.glm53_tp3_vocab_padding_size
+                )
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
                 config.hidden_size,
                 prefix=f"{prefix}.embed_tokens",
+                **vocab_kwargs,
             )
         else:
             self.embed_tokens = PPMissingLayer()
@@ -1285,11 +1318,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
         if get_pp_group().is_last_rank:
+            vocab_kwargs = {}
+            if getattr(self.config, "glm53_tp3_padding", False):
+                vocab_kwargs["padding_size"] = (
+                    self.config.glm53_tp3_vocab_padding_size
+                )
             self.lm_head = ParallelLMHead(
                 self.config.vocab_size,
                 self.config.hidden_size,
                 quant_config=quant_config,
                 prefix=maybe_prefix(prefix, "lm_head"),
+                **vocab_kwargs,
             )
         else:
             self.lm_head = PPMissingLayer()
diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/nvidia/mtp.py
index 90bdb5b64214..0208ac2a7216 100644
--- a/vllm/models/glm5next/nvidia/mtp.py
+++ b/vllm/models/glm5next/nvidia/mtp.py
@@ -11,8 +11,10 @@
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import (
@@ -36,6 +38,21 @@
 from .pooled_indexer import Glm5NextPooledIndexer
 
 
+class _Glm53TP3SharedHead(nn.Module):
+    """MTP shared head with TP3-divisible physical vocabulary storage."""
+
+    def __init__(self, config, prefix: str, quant_config=None) -> None:
+        super().__init__()
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            padding_size=config.glm53_tp3_vocab_padding_size,
+            prefix=maybe_prefix(prefix, "head"),
+        )
+
+
 class Glm5NextMultiTokenPredictorLayer(nn.Module):
     def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
         super().__init__()
@@ -46,8 +63,24 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
 
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
-
+        self.eh_proj_output_size = getattr(
+            config, "glm53_tp3_mtp_projection_size", config.hidden_size
+        )
+        if getattr(config, "glm53_tp3_padding", False):
+            self.eh_proj = ColumnParallelLinear(
+                config.hidden_size * 2,
+                self.eh_proj_output_size,
+                gather_output=True,
+                bias=False,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.eh_proj",
+                loaded_output_size=config.hidden_size,
+            )
+        else:
+            self.eh_proj = nn.Linear(
+                config.hidden_size * 2, config.hidden_size, bias=False
+            )
         topk_tokens = config.index_topk
         kpool = getattr(config, "index_kpool", 1) or 1
         buffer_width = topk_tokens + (kpool - 1 if kpool > 1 else 0)
@@ -63,9 +96,14 @@ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
             dtype=torch.int32,
             device=current_platform.device_type,
         )
-        self.shared_head = SharedHead(
-            config=config, prefix=prefix, quant_config=quant_config
-        )
+        if getattr(config, "glm53_tp3_padding", False):
+            self.shared_head = _Glm53TP3SharedHead(
+                config=config, prefix=prefix, quant_config=quant_config
+            )
+        else:
+            self.shared_head = SharedHead(
+                config=config, prefix=prefix, quant_config=quant_config
+            )
         # MTP layers sit past the base model's hidden layers; parse the index
         # from the prefix (e.g. "...layers.32") so the decoder builds an MLA
         # (DSA) layer rather than KDA for the MTP path.
@@ -95,6 +133,10 @@ def forward(
             dim=-1,
         )
         hidden_states = self.eh_proj(eh_input)
+        if hidden_states.shape[-1] != previous_hidden_states.shape[-1]:
+            hidden_states = hidden_states[
+                ..., : previous_hidden_states.shape[-1]
+            ].contiguous()
         # Fuse the residual add and final RMSNorm. Glm5NextMoE already performs
         # its all-reduce, so no collective is needed here. The post-norm result
         # feeds both draft logits and the next recycled hidden state.
@@ -125,10 +167,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                 )
             }
         )
+        vocab_kwargs = {}
+        if getattr(config, "glm53_tp3_padding", False):
+            vocab_kwargs["padding_size"] = config.glm53_tp3_vocab_padding_size
         self.embed_tokens = VocabParallelEmbedding(
             config.vocab_size,
             config.hidden_size,
             prefix=maybe_prefix(prefix, "embed_tokens"),
+            **vocab_kwargs,
         )
         # Plain list for the per-propose lookup: ModuleDict[str(...)] builds a
         # string and hashes it on every draft step.

From 602fa1161d8c65191f4250eb8b23172d21050055 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:26:20 +0200
Subject: [PATCH 09/25] test(glm53): cover every TP3 runtime proof field

(cherry picked from commit 6b55f2308b83dc3fa212e8d8bdffb382497afe07)
---
 tests/v1/worker/test_gpu_model_runner.py | 29 +++++++++++++++++-------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py
index 47cca7c258d5..f9127754a27b 100644
--- a/tests/v1/worker/test_gpu_model_runner.py
+++ b/tests/v1/worker/test_gpu_model_runner.py
@@ -80,21 +80,32 @@ def _restore_default_dtype():
     torch.set_default_dtype(old)
 
 
-@pytest.mark.parametrize("collective_enabled", (True, False))
+@pytest.mark.parametrize(
+    "broken_contract",
+    (None, "collective", "expert_parallel", "kda_decode", "kda_prefill", "vision"),
+)
 def test_glm53_r17_tp3_runtime_proof_is_observed_and_fail_closed(
     monkeypatch: pytest.MonkeyPatch,
-    collective_enabled: bool,
+    broken_contract: str | None,
 ) -> None:
     class Glm5NextLinearAttention(torch.nn.Module):
         def __init__(self) -> None:
             super().__init__()
-            self._b12x_kda_api = object()
-            self.kda_prefill_backend = "flashkda"
+            self._b12x_kda_api = (
+                None if broken_contract == "kda_decode" else object()
+            )
+            self.kda_prefill_backend = (
+                "triton" if broken_contract == "kda_prefill" else "flashkda"
+            )
 
     messages = []
     vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(
-            multimodal_config=SimpleNamespace(mm_encoder_tp_mode="weights"),
+            multimodal_config=SimpleNamespace(
+                mm_encoder_tp_mode=(
+                    "data" if broken_contract == "vision" else "weights"
+                )
+            ),
         ),
         parallel_config=SimpleNamespace(
             tensor_parallel_size=3,
@@ -103,7 +114,7 @@ def __init__(self) -> None:
     )
     model = torch.nn.Sequential(Glm5NextLinearAttention())
     b12x_ar = object.__new__(worker_utils.B12xPcieAllReduce)
-    b12x_ar.disabled = not collective_enabled
+    b12x_ar.disabled = broken_contract == "collective"
     b12x_ar.world_size = 3
     b12x_ar._runtime = object()
     b12x_ar.allreduce_max_bytes = 65536
@@ -112,7 +123,9 @@ def __init__(self) -> None:
     monkeypatch.setattr(
         worker_utils,
         "get_ep_group",
-        lambda: SimpleNamespace(world_size=3),
+        lambda: SimpleNamespace(
+            world_size=2 if broken_contract == "expert_parallel" else 3
+        ),
     )
     monkeypatch.setattr(
         worker_utils,
@@ -127,7 +140,7 @@ def __init__(self) -> None:
         lambda message, payload, **kwargs: messages.append(message % payload),
     )
 
-    if not collective_enabled:
+    if broken_contract is not None:
         with pytest.raises(RuntimeError, match="runtime proof failed"):
             worker_utils.log_glm53_r17_tp3_runtime_proof(vllm_config, model)
         return

From 7c3e50c74959d9e58b7796898f317ba843e3c4de Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:27:51 +0200
Subject: [PATCH 10/25] test(glm5next): isolate vision loaders from global
 config

(cherry picked from commit 4d97ca2cfaa5f6a59b1b8533f2adfadf3a192d89)
---
 tests/models/test_glm5next_vision_tp3.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tests/models/test_glm5next_vision_tp3.py b/tests/models/test_glm5next_vision_tp3.py
index 294535f6035e..849d2e07f264 100644
--- a/tests/models/test_glm5next_vision_tp3.py
+++ b/tests/models/test_glm5next_vision_tp3.py
@@ -13,6 +13,15 @@
 
 @pytest.fixture
 def tp3_linear_state(monkeypatch):
+    compilation_config = SimpleNamespace(
+        custom_ops=["none"],
+        enabled_custom_ops=set(),
+        disabled_custom_ops=set(),
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.custom_op.get_cached_compilation_config",
+        lambda: compilation_config,
+    )
     monkeypatch.setattr(
         glm5next_multimodal, "get_tensor_model_parallel_world_size", lambda: 3
     )

From fde3aee9706cd35e4934bc7637718cd005339e71 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:28:22 +0200
Subject: [PATCH 11/25] test(dflash): assert padded tails stay inert

(cherry picked from commit 90f984b65821872e0a26a034f8739c88166992cb)
---
 tests/models/test_glm53_tp3_dflash.py | 135 ++++++++++++++++++++++++--
 1 file changed, 129 insertions(+), 6 deletions(-)

diff --git a/tests/models/test_glm53_tp3_dflash.py b/tests/models/test_glm53_tp3_dflash.py
index 8d67fbbd992a..3b402d63c704 100644
--- a/tests/models/test_glm53_tp3_dflash.py
+++ b/tests/models/test_glm53_tp3_dflash.py
@@ -5,11 +5,15 @@
 
 import torch
 
+from tests.utils import ensure_current_vllm_config
+
 from vllm.config.speculative import SpeculativeConfig
 from vllm.model_executor.layers import vocab_parallel_embedding
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.models import qwen3_dflash
 from vllm.model_executor.models.qwen3_dflash import (
+    DFlashQwen3Attention,
     DFlashQwen3Model,
     _get_dflash_draft_vocab_size,
     _get_glm53_tp3_head_geometry,
@@ -35,6 +39,60 @@ def _tp3_dflash_config(**overrides):
     values.update(overrides)
     return SimpleNamespace(**values)
 
+def _capture_attention_projection_kwargs(monkeypatch, config, tp_size):
+    calls = {}
+
+    class FakeProjection(torch.nn.Module):
+        def __init__(self, kind, args, kwargs):
+            super().__init__()
+            calls[kind] = (args, kwargs)
+
+    monkeypatch.setattr(
+        qwen3_dflash,
+        "QKVParallelLinear",
+        lambda *args, **kwargs: FakeProjection("qkv", args, kwargs),
+    )
+    monkeypatch.setattr(
+        qwen3_dflash,
+        "RowParallelLinear",
+        lambda *args, **kwargs: FakeProjection("o", args, kwargs),
+    )
+    monkeypatch.setattr(
+        qwen3_dflash,
+        "DFlashAttention",
+        lambda *args, **kwargs: FakeProjection("attention", args, kwargs),
+    )
+    monkeypatch.setattr(qwen3_dflash, "get_rope", lambda *args, **kwargs: object())
+    monkeypatch.setattr(
+        qwen3_dflash, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    DFlashQwen3Attention(
+        hidden_size=64,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=config.num_key_value_heads,
+        config=config,
+        rope_parameters={},
+        head_dim=2,
+        prefix="model.layers.0.self_attn",
+    )
+    return calls
+
+
+def test_dflash_tp3_wires_logical_checkpoint_projection_sizes(monkeypatch) -> None:
+    calls = _capture_attention_projection_kwargs(
+        monkeypatch, _tp3_dflash_config(), tp_size=3
+    )
+    qkv_args, qkv_kwargs = calls["qkv"]
+    o_args, o_kwargs = calls["o"]
+
+    assert qkv_args[:4] == (64, 2, 36, 9)
+    assert qkv_kwargs["loaded_total_num_heads"] == 32
+    assert qkv_kwargs["loaded_total_num_kv_heads"] == 8
+    assert qkv_kwargs["prefix"] == "model.layers.0.self_attn.qkv_proj"
+    assert o_args[:2] == (72, 64)
+    assert o_kwargs["loaded_input_size"] == 64
+    assert o_kwargs["prefix"] == "model.layers.0.self_attn.o_proj"
+
 
 def test_dflash_tp3_geometry_and_vocab_storage(monkeypatch) -> None:
     config = _tp3_dflash_config()
@@ -54,11 +112,12 @@ def test_dflash_tp3_geometry_and_vocab_storage(monkeypatch) -> None:
         "get_tensor_model_parallel_rank",
         lambda: 2,
     )
-    embedding = VocabParallelEmbedding(
-        _get_dflash_draft_vocab_size(config),
-        8,
-        **vocab_kwargs,
-    )
+    with ensure_current_vllm_config():
+        embedding = VocabParallelEmbedding(
+            _get_dflash_draft_vocab_size(config),
+            8,
+            **vocab_kwargs,
+        )
     assert embedding.num_embeddings == 154880
     assert embedding.num_embeddings_padded == 154944
     assert embedding.num_embeddings_per_partition == 51648
@@ -82,6 +141,60 @@ def test_dflash_tp3_sink_bias_pads_only_rank_local_tail(monkeypatch) -> None:
     assert local.untyped_storage().data_ptr() != loaded.untyped_storage().data_ptr()
 
 
+def test_zero_padded_qkv_o_projection_matches_logical_reference() -> None:
+    generator = torch.Generator().manual_seed(17)
+    hidden_states = torch.randn(3, 64, generator=generator)
+    logical_q_weight = torch.randn(64, 64, generator=generator)
+    logical_kv_weight = torch.randn(16, 64, generator=generator)
+
+    logical_q = torch.nn.functional.linear(hidden_states, logical_q_weight)
+    logical_k = torch.nn.functional.linear(hidden_states, logical_kv_weight)
+    physical_q = torch.zeros(3, 72)
+    physical_k = torch.zeros(3, 18)
+    physical_q[:, :64] = logical_q
+    physical_k[:, :16] = logical_k
+
+    torch.testing.assert_close(physical_q[:, 64:], torch.zeros(3, 8))
+    torch.testing.assert_close(physical_k[:, 16:], torch.zeros(3, 2))
+
+    logical_o_weight = torch.randn(64, 64, generator=generator)
+    physical_o_weight = torch.zeros(64, 72)
+    physical_o_weight[:, :64] = logical_o_weight
+    logical_output = torch.nn.functional.linear(logical_q, logical_o_weight)
+    padded_output = torch.nn.functional.linear(physical_q, physical_o_weight)
+    torch.testing.assert_close(padded_output, logical_output)
+
+
+def test_dflash_logits_and_selector_exclude_physical_vocab_tail() -> None:
+    logical_vocab_size = _get_dflash_draft_vocab_size(_tp3_dflash_config())
+    storage_logits = torch.zeros(1, 154944)
+    storage_logits[:, logical_vocab_size - 1] = 2
+    storage_logits[:, logical_vocab_size:] = 100
+
+    processor = LogitsProcessor.__new__(LogitsProcessor)
+    torch.nn.Module.__init__(processor)
+    processor.org_vocab_size = logical_vocab_size
+    processor.scale = 1.0
+    processor.soft_cap = None
+    processor._apply_head = lambda *args, **kwargs: storage_logits.clone()
+    lm_head = SimpleNamespace(
+        tp_size=1,
+        shard_indices=SimpleNamespace(
+            num_org_vocab_padding=154944 - logical_vocab_size,
+            org_vocab_start_index=0,
+        ),
+    )
+
+    logits = processor._get_logits(torch.empty(1, 1), lm_head, None)
+    assert logits is not None
+    assert logits.shape == (1, logical_vocab_size)
+    token_ids, values = processor.get_top_k_tokens(
+        lm_head, torch.empty(1, 1), k=1
+    )
+    assert token_ids.item() == logical_vocab_size - 1
+    assert values.item() == 2
+
+
 def test_dflash_tp4_paths_are_exact_noops(monkeypatch) -> None:
     config = SimpleNamespace(
         num_attention_heads=32,
@@ -94,6 +207,15 @@ def test_dflash_tp4_paths_are_exact_noops(monkeypatch) -> None:
     assert _get_dflash_draft_vocab_size(config) == 154880
     assert vars(config) == original_fields
 
+    calls = _capture_attention_projection_kwargs(monkeypatch, config, tp_size=4)
+    qkv_args, qkv_kwargs = calls["qkv"]
+    o_args, o_kwargs = calls["o"]
+    assert qkv_args[:4] == (64, 2, 32, 8)
+    assert "loaded_total_num_heads" not in qkv_kwargs
+    assert "loaded_total_num_kv_heads" not in qkv_kwargs
+    assert o_args[:2] == (64, 64)
+    assert "loaded_input_size" not in o_kwargs
+
     monkeypatch.setattr(
         vocab_parallel_embedding,
         "get_tensor_model_parallel_world_size",
@@ -104,7 +226,8 @@ def test_dflash_tp4_paths_are_exact_noops(monkeypatch) -> None:
         "get_tensor_model_parallel_rank",
         lambda: 2,
     )
-    embedding = VocabParallelEmbedding(config.vocab_size, 8)
+    with ensure_current_vllm_config():
+        embedding = VocabParallelEmbedding(config.vocab_size, 8)
     assert embedding.num_embeddings == 154880
     assert embedding.num_embeddings_padded == 154880
     assert embedding.num_embeddings_per_partition == 38720

From ec6ebc24cd2c942c1e3549b0709f8ea8f0b5fb68 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:28:52 +0200
Subject: [PATCH 12/25] test(glm5next): isolate TP3 shared expert geometry

(cherry picked from commit d7c8f680eac6b45ad9ace6a1cd78eb005de6e825)
---
 tests/models/test_glm53_tp3_model.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/models/test_glm53_tp3_model.py b/tests/models/test_glm53_tp3_model.py
index 91effebf5923..a2edec1d89ce 100644
--- a/tests/models/test_glm53_tp3_model.py
+++ b/tests/models/test_glm53_tp3_model.py
@@ -174,8 +174,14 @@ def __init__(self, input_size, output_size, **kwargs) -> None:
             super().__init__()
             calls["down"] = (input_size, output_size, kwargs)
 
+    class FakeActivation(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+
+
     monkeypatch.setattr(glm_model, "MergedColumnParallelLinear", FakeMerged)
     monkeypatch.setattr(glm_model, "RowParallelLinear", FakeRow)
+    monkeypatch.setattr(glm_model, "SiluAndMul", FakeActivation)
 
     glm_model.Glm5NextMLP(
         hidden_size=4096,

From 2dbc947444aa2c08c5093b32bfab8a725d083ff6 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:32:31 +0200
Subject: [PATCH 13/25] test(glm5next): disable compilation in vision loader
 tests

(cherry picked from commit 70c9b4d4836fbf66f95d83f19d6a19cb61172571)
---
 tests/models/test_glm5next_vision_tp3.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/models/test_glm5next_vision_tp3.py b/tests/models/test_glm5next_vision_tp3.py
index 849d2e07f264..392a516ed829 100644
--- a/tests/models/test_glm5next_vision_tp3.py
+++ b/tests/models/test_glm5next_vision_tp3.py
@@ -6,6 +6,7 @@
 import pytest
 import torch
 
+from vllm.config.compilation import CompilationMode
 from vllm.model_executor import parameter
 from vllm.model_executor.layers import linear
 from vllm.models.glm5next.nvidia import multimodal as glm5next_multimodal
@@ -17,6 +18,7 @@ def tp3_linear_state(monkeypatch):
         custom_ops=["none"],
         enabled_custom_ops=set(),
         disabled_custom_ops=set(),
+        mode=CompilationMode.NONE,
     )
     monkeypatch.setattr(
         "vllm.model_executor.custom_op.get_cached_compilation_config",

From c36c34a6160037234ba916bf925da424f7ab0ee0 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:35:04 +0200
Subject: [PATCH 14/25] fix(tp): preserve strict unpadded loader calls

(cherry picked from commit 93a2d0fd5bf4193f22b82c2d57aa8e62feacb5a6)
---
 vllm/model_executor/layers/linear.py | 56 ++++++++++++++++------------
 vllm/model_executor/parameter.py     | 24 +++++++++++-
 2 files changed, 55 insertions(+), 25 deletions(-)

diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index 6a9df1e58f1d..0280ca30152b 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -627,10 +627,13 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
         if len(loaded_weight.shape) == 0:
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
-        param.load_column_parallel_weight(
-            loaded_weight=loaded_weight,
-            allow_padding=self._allow_loaded_output_padding,
-        )
+        if self._allow_loaded_output_padding:
+            param.load_column_parallel_weight(
+                loaded_weight=loaded_weight,
+                allow_padding=True,
+            )
+        else:
+            param.load_column_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(
         self,
@@ -1044,13 +1047,15 @@ def weight_loader_v2(
                 weight_block_size, shard_size, shard_offset
             )
 
-        param.load_merged_column_weight(
-            loaded_weight=loaded_weight,
-            shard_id=loaded_shard_id,
-            shard_offset=shard_offset,
-            shard_size=shard_size,
-            allow_padding=self._allow_loaded_output_shard_padding[loaded_shard_id],
-        )
+        load_kwargs = {
+            "loaded_weight": loaded_weight,
+            "shard_id": loaded_shard_id,
+            "shard_offset": shard_offset,
+            "shard_size": shard_size,
+        }
+        if self._allow_loaded_output_shard_padding[loaded_shard_id]:
+            load_kwargs["allow_padding"] = True
+        param.load_merged_column_weight(**load_kwargs)
 
     def load_weights(
         self, weights: Iterable[tuple[str, torch.Tensor]]
@@ -1305,14 +1310,16 @@ def weight_loader_v2(
                 weight_block_size, shard_size, shard_offset
             )
 
-        param.load_qkv_weight(
-            loaded_weight=loaded_weight,
-            num_heads=self.num_kv_head_replicas,
-            shard_id=loaded_shard_id,
-            shard_offset=shard_offset,
-            shard_size=shard_size,
-            allow_padding=self._allow_loaded_qkv_padding[loaded_shard_id],
-        )
+        load_kwargs = {
+            "loaded_weight": loaded_weight,
+            "num_heads": self.num_kv_head_replicas,
+            "shard_id": loaded_shard_id,
+            "shard_offset": shard_offset,
+            "shard_size": shard_size,
+        }
+        if self._allow_loaded_qkv_padding[loaded_shard_id]:
+            load_kwargs["allow_padding"] = True
+        param.load_qkv_weight(**load_kwargs)
 
     def weight_loader(
         self,
@@ -1827,10 +1834,13 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
 
-        param.load_row_parallel_weight(
-            loaded_weight=loaded_weight,
-            allow_padding=self._allow_loaded_input_padding,
-        )
+        if self._allow_loaded_input_padding:
+            param.load_row_parallel_weight(
+                loaded_weight=loaded_weight,
+                allow_padding=True,
+            )
+        else:
+            param.load_row_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(
         self,
diff --git a/vllm/model_executor/parameter.py b/vllm/model_executor/parameter.py
index 93df17a3b760..bdb2d4b714b4 100644
--- a/vllm/model_executor/parameter.py
+++ b/vllm/model_executor/parameter.py
@@ -147,10 +147,30 @@ def _assert_and_load(self, loaded_weight: torch.Tensor):
         )
         self.data.copy_(loaded_weight)
 
-    def load_column_parallel_weight(self, loaded_weight: torch.Tensor):
+    def load_column_parallel_weight(
+        self,
+        loaded_weight: torch.Tensor,
+        *,
+        allow_padding: bool = False,
+    ):
+        if allow_padding and self.data.shape != loaded_weight.shape:
+            raise ValueError(
+                "Padded column-parallel loading requires a parameter with an "
+                "output sharding dimension"
+            )
         self._assert_and_load(loaded_weight)
 
-    def load_row_parallel_weight(self, loaded_weight: torch.Tensor):
+    def load_row_parallel_weight(
+        self,
+        loaded_weight: torch.Tensor,
+        *,
+        allow_padding: bool = False,
+    ):
+        if allow_padding and self.data.shape != loaded_weight.shape:
+            raise ValueError(
+                "Padded row-parallel loading requires a parameter with an "
+                "input sharding dimension"
+            )
         self._assert_and_load(loaded_weight)
 
     def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):

From 7d5ec25fc3abe8425795a7efa4f71f4aec3bb9e0 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 11:37:41 +0200
Subject: [PATCH 15/25] test(dflash): isolate CPU config context

(cherry picked from commit f7a2d2cf386370da6fd9df8340687512dc82d085)
---
 tests/models/test_glm53_tp3_dflash.py | 40 ++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/tests/models/test_glm53_tp3_dflash.py b/tests/models/test_glm53_tp3_dflash.py
index 3b402d63c704..9f2036faaeea 100644
--- a/tests/models/test_glm53_tp3_dflash.py
+++ b/tests/models/test_glm53_tp3_dflash.py
@@ -5,9 +5,9 @@
 
 import torch
 
-from tests.utils import ensure_current_vllm_config
-
+from vllm.config.compilation import CompilationMode
 from vllm.config.speculative import SpeculativeConfig
+from vllm.model_executor import custom_op
 from vllm.model_executor.layers import vocab_parallel_embedding
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
@@ -23,6 +23,20 @@
 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
 
 
+def _disable_custom_ops(monkeypatch) -> None:
+    compilation_config = SimpleNamespace(
+        mode=CompilationMode.NONE,
+        custom_ops=["none"],
+        enabled_custom_ops=set(),
+        disabled_custom_ops=set(),
+    )
+    monkeypatch.setattr(
+        custom_op,
+        "get_cached_compilation_config",
+        lambda: compilation_config,
+    )
+
+
 def _tp3_dflash_config(**overrides):
     values = {
         "glm53_tp3_padding": True,
@@ -34,12 +48,16 @@ def _tp3_dflash_config(**overrides):
         "original_num_key_value_heads": 8,
         "original_vocab_size": 154880,
         "draft_vocab_size": 154880,
+        "hidden_size": 4096,
+        "intermediate_size": 12288,
         "vocab_size": 154880,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
 
+
 def _capture_attention_projection_kwargs(monkeypatch, config, tp_size):
+    _disable_custom_ops(monkeypatch)
     calls = {}
 
     class FakeProjection(torch.nn.Module):
@@ -95,12 +113,16 @@ def test_dflash_tp3_wires_logical_checkpoint_projection_sizes(monkeypatch) -> No
 
 
 def test_dflash_tp3_geometry_and_vocab_storage(monkeypatch) -> None:
+    _disable_custom_ops(monkeypatch)
     config = _tp3_dflash_config()
 
     assert _get_glm53_tp3_head_geometry(config) == (32, 8)
     assert _get_dflash_draft_vocab_size(config) == 154880
     vocab_kwargs = _get_glm53_tp3_vocab_kwargs(config)
     assert vocab_kwargs == {"padding_size": 192}
+    assert config.intermediate_size % 3 == 0
+    assert config.intermediate_size // 3 == 4096
+    assert config.intermediate_size % 4 == 0
 
     monkeypatch.setattr(
         vocab_parallel_embedding,
@@ -112,12 +134,11 @@ def test_dflash_tp3_geometry_and_vocab_storage(monkeypatch) -> None:
         "get_tensor_model_parallel_rank",
         lambda: 2,
     )
-    with ensure_current_vllm_config():
-        embedding = VocabParallelEmbedding(
-            _get_dflash_draft_vocab_size(config),
-            8,
-            **vocab_kwargs,
-        )
+    embedding = VocabParallelEmbedding(
+        _get_dflash_draft_vocab_size(config),
+        8,
+        **vocab_kwargs,
+    )
     assert embedding.num_embeddings == 154880
     assert embedding.num_embeddings_padded == 154944
     assert embedding.num_embeddings_per_partition == 51648
@@ -226,8 +247,7 @@ def test_dflash_tp4_paths_are_exact_noops(monkeypatch) -> None:
         "get_tensor_model_parallel_rank",
         lambda: 2,
     )
-    with ensure_current_vllm_config():
-        embedding = VocabParallelEmbedding(config.vocab_size, 8)
+    embedding = VocabParallelEmbedding(config.vocab_size, 8)
     assert embedding.num_embeddings == 154880
     assert embedding.num_embeddings_padded == 154880
     assert embedding.num_embeddings_per_partition == 38720

From ff5b7fb21bde23d361cafc4c0cd91fd025e77349 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 12:26:04 +0200
Subject: [PATCH 16/25] Fix GLM TP3 draft config isolation

(cherry picked from commit 396576d29f13f1891ccf8982b05b7a2ff69c98ff)
---
 tests/config/test_glm53_tp3_geometry.py      |  50 +++++
 tests/models/test_glm53_tp3_dflash.py        |  41 +++-
 vllm/transformers_utils/configs/glm53_tp3.py | 200 +++++++++++++++++++
 vllm/v1/spec_decode/dflash.py                |   9 +
 4 files changed, 299 insertions(+), 1 deletion(-)
 create mode 100644 vllm/transformers_utils/configs/glm53_tp3.py

diff --git a/tests/config/test_glm53_tp3_geometry.py b/tests/config/test_glm53_tp3_geometry.py
index ae6feaf4ece6..d53dac532577 100644
--- a/tests/config/test_glm53_tp3_geometry.py
+++ b/tests/config/test_glm53_tp3_geometry.py
@@ -215,6 +215,56 @@ def test_glm53_tp3_invalid_target_checkpoint_geometry_fails_closed(
     assert _snapshot(glm53_model_config) == before
 
 
+@pytest.mark.parametrize(
+    ("attribute", "invalid_value"),
+    [
+        ("num_heads", 15),
+        ("hidden_size", 1023),
+        ("intermediate_size", 4095),
+        ("projection_intermediate_size", 10239),
+    ],
+)
+def test_glm53_tp3_invalid_vision_geometry_is_transactional(
+    glm53_model_config: FakeGlm53ModelConfig,
+    tp3_ep_parallel_config: ParallelConfig,
+    attribute: str,
+    invalid_value: int,
+) -> None:
+    vision_config = glm53_model_config.hf_config.vision_config
+    setattr(vision_config, attribute, invalid_value)
+    before = _snapshot(glm53_model_config)
+
+    with pytest.raises(ValueError, match=rf"expected {attribute}="):
+        apply_glm53_tp3_target_geometry(
+            cast(Any, glm53_model_config), tp3_ep_parallel_config
+        )
+
+    assert _snapshot(glm53_model_config) == before
+    assert not hasattr(glm53_model_config.hf_text_config, "glm53_tp3_padding")
+    assert not hasattr(glm53_model_config.hf_config, "glm53_tp3_padding")
+    assert not hasattr(vision_config, "glm53_tp3_padding")
+
+
+def test_glm53_tp3_invalid_wrapper_geometry_is_transactional(
+    glm53_model_config: FakeGlm53ModelConfig,
+    tp3_ep_parallel_config: ParallelConfig,
+) -> None:
+    glm53_model_config.hf_config.hidden_size = 4095
+    before = _snapshot(glm53_model_config)
+
+    with pytest.raises(ValueError, match=r"expected hidden_size=4096"):
+        apply_glm53_tp3_target_geometry(
+            cast(Any, glm53_model_config), tp3_ep_parallel_config
+        )
+
+    assert _snapshot(glm53_model_config) == before
+    assert not hasattr(glm53_model_config.hf_text_config, "glm53_tp3_padding")
+    assert not hasattr(glm53_model_config.hf_config, "glm53_tp3_padding")
+    assert not hasattr(
+        glm53_model_config.hf_config.vision_config, "glm53_tp3_padding"
+    )
+
+
 def test_glm53_tp4_parallel_config_is_an_exact_attribute_noop(
     glm53_model_config: FakeGlm53ModelConfig,
 ) -> None:
diff --git a/tests/models/test_glm53_tp3_dflash.py b/tests/models/test_glm53_tp3_dflash.py
index 9f2036faaeea..a85b0aba5059 100644
--- a/tests/models/test_glm53_tp3_dflash.py
+++ b/tests/models/test_glm53_tp3_dflash.py
@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
 from types import SimpleNamespace
+from typing import Any
 
 import torch
 
+from vllm.config import AttentionConfig, ParallelConfig
 from vllm.config.compilation import CompilationMode
 from vllm.config.speculative import SpeculativeConfig
 from vllm.model_executor import custom_op
@@ -19,6 +22,8 @@
     _get_glm53_tp3_head_geometry,
     _get_glm53_tp3_vocab_kwargs,
 )
+from vllm.v1.spec_decode.dflash import DFlashProposer
+from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
 
@@ -289,9 +294,10 @@ def test_dense_dflash_tp3_drops_target_expert_parallel(monkeypatch) -> None:
     target_parallel = SimpleNamespace(
         tensor_parallel_size=3,
         enable_expert_parallel=True,
+        rank=2,
         **placement,
     )
-    draft_parallel = SimpleNamespace(
+    draft_parallel = ParallelConfig(
         tensor_parallel_size=3,
         enable_expert_parallel=True,
     )
@@ -309,6 +315,39 @@ def test_dense_dflash_tp3_drops_target_expert_parallel(monkeypatch) -> None:
         assert getattr(draft_parallel, name) == value
     assert len(applied) == 1
 
+    @dataclass
+    class DraftVllmConfig:
+        model_config: Any
+        parallel_config: Any
+        attention_config: AttentionConfig
+
+    base = DraftVllmConfig(
+        model_config=SimpleNamespace(
+            model_arch_config=SimpleNamespace(is_mm_prefix_lm=False)
+        ),
+        parallel_config=target_parallel,
+        attention_config=AttentionConfig(),
+    )
+    proposer = object.__new__(DFlashProposer)
+    proposer.speculative_config = spec
+    proposer.vllm_config = base
+    proposer.dflash_causal = True
+    monkeypatch.setattr(
+        SpecDecodeBaseProposer,
+        "_create_draft_vllm_config",
+        lambda _: base,
+    )
+
+    draft_vllm_config = proposer._create_draft_vllm_config()
+
+    assert draft_vllm_config.parallel_config is not draft_parallel
+    assert draft_vllm_config.parallel_config.tensor_parallel_size == 3
+    assert not draft_vllm_config.parallel_config.enable_expert_parallel
+    assert draft_vllm_config.parallel_config.rank == target_parallel.rank
+    assert draft_parallel.rank == 0
+    assert target_parallel.tensor_parallel_size == 3
+    assert target_parallel.enable_expert_parallel
+
 
 def test_dflash7_uses_eight_target_kda_state_columns(monkeypatch) -> None:
     def fake_base_init(self, vllm_config, device) -> None:
diff --git a/vllm/transformers_utils/configs/glm53_tp3.py b/vllm/transformers_utils/configs/glm53_tp3.py
new file mode 100644
index 000000000000..9c163badb188
--- /dev/null
+++ b/vllm/transformers_utils/configs/glm53_tp3.py
@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Physical TP3 geometry for the GLM-5.3 target and its draft models.
+
+The checkpoint dimensions remain available as ``original_*`` attributes.  The
+model implementations use the physical dimensions to allocate TP-sharded
+parameters and the original dimensions while loading and producing logits.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from vllm.config.model import ModelConfig
+    from vllm.config.parallel import ParallelConfig
+
+_GLM53_ARCHITECTURES = {
+    "Glm5NextForCausalLM",
+    "Glm5NextForConditionalGeneration",
+    "Glm5NextMTPModel",
+}
+_DFLASH_ARCHITECTURES = {"DFlashDraftModel", "DFlash2DraftModel"}
+
+
+def _iter_hf_configs(model_config: ModelConfig):
+    seen: set[int] = set()
+    hf_config = model_config.hf_config
+    for config in (
+        hf_config,
+        getattr(model_config, "hf_text_config", None),
+        getattr(hf_config, "text_config", None),
+    ):
+        if config is not None and id(config) not in seen:
+            seen.add(id(config))
+            yield config
+
+
+def _has_architecture(model_config: ModelConfig, names: set[str]) -> bool:
+    for config in _iter_hf_configs(model_config):
+        if names.intersection(getattr(config, "architectures", None) or ()):
+            return True
+    return False
+
+
+def is_glm53_config(model_config: ModelConfig | None) -> bool:
+    if model_config is None:
+        return False
+    if _has_architecture(model_config, _GLM53_ARCHITECTURES):
+        return True
+    return any(
+        getattr(config, "model_type", None)
+        in {"glm5_next", "glm5_next_text", "glm5_next_mtp"}
+        for config in _iter_hf_configs(model_config)
+    )
+
+
+def _logical_value(config: Any, name: str) -> int:
+    original = getattr(config, f"original_{name}", None)
+    return int(original if original is not None else getattr(config, name))
+
+
+def _require_shape(config: Any, name: str, expected: int) -> None:
+    value = _logical_value(config, name)
+    if value != expected:
+        raise ValueError(
+            "GLM-5.3 TP3 padding only supports the released checkpoint "
+            f"geometry: expected {name}={expected}, got {value}."
+        )
+
+
+def apply_glm53_tp3_target_geometry(
+    model_config: ModelConfig | None,
+    parallel_config: ParallelConfig | None,
+) -> bool:
+    """Apply GLM-5.3's physical TP3 axes using the actual parallel config.
+
+    Returns whether the target is using the TP3 layout.  TP1/2/4 and unrelated
+    configurations are exact no-ops.
+    """
+    if (
+        model_config is None
+        or parallel_config is None
+        or parallel_config.tensor_parallel_size != 3
+        or not is_glm53_config(model_config)
+    ):
+        return False
+
+    text_config = model_config.hf_text_config
+    target_shapes = (
+        ("num_attention_heads", 64),
+        ("num_key_value_heads", 64),
+        ("linear_num_heads", 64),
+        ("moe_intermediate_size", 2048),
+        ("hidden_size", 4096),
+        ("vocab_size", 154880),
+    )
+    for name, expected in target_shapes:
+        _require_shape(text_config, name, expected)
+
+    for config in _iter_hf_configs(model_config):
+        if config is text_config:
+            continue
+        for name, expected in target_shapes:
+            if hasattr(config, name):
+                _require_shape(config, name, expected)
+
+    vision_config = getattr(model_config.hf_config, "vision_config", None)
+    multimodal_config = getattr(model_config, "multimodal_config", None)
+    uses_weights_mode_vision = (
+        vision_config is not None
+        and multimodal_config is not None
+        and multimodal_config.mm_encoder_tp_mode == "weights"
+    )
+    if uses_weights_mode_vision:
+        for name, expected in (
+            ("num_heads", 16),
+            ("hidden_size", 1024),
+            ("intermediate_size", 4096),
+            ("projection_intermediate_size", 10240),
+        ):
+            _require_shape(vision_config, name, expected)
+    for config in _iter_hf_configs(model_config):
+        if hasattr(config, "num_attention_heads"):
+            config.original_num_attention_heads = 64
+            config.num_attention_heads = 72
+        if hasattr(config, "num_key_value_heads"):
+            config.original_num_key_value_heads = 64
+            config.num_key_value_heads = 72
+        if hasattr(config, "linear_num_heads"):
+            config.original_linear_num_heads = 64
+            config.linear_num_heads = 66
+        linear_config = getattr(config, "linear_attn_config", None)
+        if isinstance(linear_config, dict) and "num_heads" in linear_config:
+            config.linear_attn_config = {**linear_config, "num_heads": 66}
+        config.glm53_tp3_padding = True
+        config.glm53_tp3_shared_expert_intermediate_size = 2112
+        config.glm53_tp3_mtp_projection_size = 4098
+        config.glm53_tp3_vocab_padding_size = 192
+        config.glm53_tp3_vocab_storage_size = 154944
+
+    if uses_weights_mode_vision:
+        vision_config.original_num_heads = 16
+        vision_config.num_heads = 18
+        vision_config.original_intermediate_size = 4096
+        vision_config.intermediate_size = 4098
+        vision_config.original_projection_intermediate_size = 10240
+        vision_config.projection_intermediate_size = 10242
+        vision_config.glm53_tp3_attention_projection_size = 1152
+        vision_config.glm53_tp3_padding = True
+
+    model_config.model_arch_config = model_config.get_model_arch_config()
+    return True
+
+
+def apply_glm53_tp3_draft_geometry(
+    target_model_config: ModelConfig | None,
+    target_parallel_config: ParallelConfig | None,
+    draft_model_config: ModelConfig | None,
+    draft_parallel_config: ParallelConfig | None,
+) -> bool:
+    """Apply the target's TP3 contract to an MTP or DFlash draft config."""
+    if (
+        target_model_config is None
+        or target_parallel_config is None
+        or draft_model_config is None
+        or draft_parallel_config is None
+        or target_parallel_config.tensor_parallel_size != 3
+        or draft_parallel_config.tensor_parallel_size != 3
+        or not is_glm53_config(target_model_config)
+    ):
+        return False
+
+    if is_glm53_config(draft_model_config):
+        return apply_glm53_tp3_target_geometry(
+            draft_model_config, draft_parallel_config
+        )
+    if not _has_architecture(draft_model_config, _DFLASH_ARCHITECTURES):
+        return False
+
+    text_config = draft_model_config.hf_text_config
+    _require_shape(text_config, "num_attention_heads", 32)
+    _require_shape(text_config, "num_key_value_heads", 8)
+    _require_shape(text_config, "vocab_size", 154880)
+    for config in _iter_hf_configs(draft_model_config):
+        if hasattr(config, "num_attention_heads"):
+            config.original_num_attention_heads = 32
+            config.num_attention_heads = 36
+        if hasattr(config, "num_key_value_heads"):
+            config.original_num_key_value_heads = 8
+            config.num_key_value_heads = 9
+        if hasattr(config, "vocab_size"):
+            config.original_vocab_size = 154880
+            config.draft_vocab_size = 154880
+        config.glm53_tp3_padding = True
+        config.glm53_tp3_vocab_padding_size = 192
+        config.glm53_tp3_vocab_storage_size = 154944
+
+    draft_model_config.model_arch_config = draft_model_config.get_model_arch_config()
+    return True
diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py
index e773ecf30596..354d0a302416 100644
--- a/vllm/v1/spec_decode/dflash.py
+++ b/vllm/v1/spec_decode/dflash.py
@@ -98,6 +98,15 @@ def load_model(self, target_model: torch.nn.Module) -> None:
     @override
     def _create_draft_vllm_config(self) -> VllmConfig:
         base = super()._create_draft_vllm_config()
+        draft_parallel_config = self.speculative_config.draft_parallel_config
+        assert draft_parallel_config is not None
+        base = replace(
+            base,
+            parallel_config=replace(
+                draft_parallel_config,
+                rank=self.vllm_config.parallel_config.rank,
+            ),
+        )
         # The draft model is text-only — clear the target's multimodal
         # flag so flash_attn is not rejected for mm_prefix support.
         arch = base.model_config.model_arch_config

From 91d35654e2d0dd88abac5bdf63f584c73fbd22e7 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 12:50:49 +0200
Subject: [PATCH 17/25] fix(tp): validate padded checkpoint layouts

(cherry picked from commit 1f7f257bf0294400bbe9f36d05498294e73d4a82)
---
 tests/models/test_glm53_tp3_model.py |  97 ++++++-
 vllm/model_executor/layers/linear.py | 374 ++++++++++++++++++++++++---
 vllm/model_executor/parameter.py     |  18 ++
 3 files changed, 450 insertions(+), 39 deletions(-)

diff --git a/tests/models/test_glm53_tp3_model.py b/tests/models/test_glm53_tp3_model.py
index a2edec1d89ce..3e1caac08d6a 100644
--- a/tests/models/test_glm53_tp3_model.py
+++ b/tests/models/test_glm53_tp3_model.py
@@ -6,6 +6,7 @@
 import pytest
 import torch
 
+from vllm.model_executor import parameter
 from vllm.model_executor.layers import linear
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -13,7 +14,10 @@
     QKVParallelLinear,
     RowParallelLinear,
 )
-from vllm.model_executor import parameter
+from vllm.model_executor.parameter import (
+    BlockQuantScaleParameter,
+    PackedvLLMParameter,
+)
 from vllm.models.glm5next.nvidia import attention as glm_attention
 from vllm.models.glm5next.nvidia import model as glm_model
 from vllm.models.glm5next.nvidia import mtp as glm_mtp
@@ -96,6 +100,97 @@ def test_loaded_sizes_reject_invalid_physical_layout(
         RowParallelLinear(6, 1, bias=False, loaded_input_size=7)
 
 
+def test_padded_loader_rejects_truncated_checkpoint_before_write(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _set_tp3_rank2(monkeypatch)
+    column = ColumnParallelLinear(1, 6, bias=False, loaded_output_size=5)
+    column.weight.data.fill_(17)
+
+    with pytest.raises(ValueError, match="expected 5, got 4"):
+        column.weight.weight_loader(
+            column.weight,
+            torch.ones((4, 1), dtype=column.weight.dtype),
+        )
+
+    torch.testing.assert_close(column.weight, column.weight.new_full((2, 1), 17))
+
+    merged = MergedColumnParallelLinear(
+        1, [6, 6], bias=False, loaded_output_sizes=[5, 5]
+    )
+    merged.weight.data.fill_(17)
+    with pytest.raises(ValueError, match="expected 5, got 4"):
+        merged.weight.weight_loader(
+            merged.weight,
+            torch.ones((4, 1), dtype=merged.weight.dtype),
+            0,
+        )
+    torch.testing.assert_close(merged.weight, merged.weight.new_full((4, 1), 17))
+
+    qkv = QKVParallelLinear(
+        hidden_size=1,
+        head_size=1,
+        total_num_heads=6,
+        total_num_kv_heads=3,
+        loaded_total_num_heads=4,
+        loaded_total_num_kv_heads=2,
+        bias=False,
+    )
+    qkv.weight.data.fill_(17)
+    with pytest.raises(ValueError, match="expected 4, got 3"):
+        qkv.weight.weight_loader(
+            qkv.weight,
+            torch.ones((3, 1), dtype=qkv.weight.dtype),
+            "q",
+        )
+    torch.testing.assert_close(qkv.weight, qkv.weight.new_full((4, 1), 17))
+
+    row = RowParallelLinear(6, 1, bias=False, loaded_input_size=5)
+    row.weight.data.fill_(17)
+    with pytest.raises(ValueError, match="expected 5, got 4"):
+        row.weight.weight_loader(
+            row.weight,
+            torch.ones((1, 4), dtype=row.weight.dtype),
+        )
+    torch.testing.assert_close(row.weight, row.weight.new_full((1, 2), 17))
+
+
+@pytest.mark.parametrize(
+    ("parameter_type", "error"),
+    [
+        (PackedvLLMParameter, "packed_factor=4"),
+        (BlockQuantScaleParameter, "quantization block size 4"),
+    ],
+)
+def test_padded_loader_rejects_unaligned_quantized_boundaries_before_write(
+    monkeypatch: pytest.MonkeyPatch,
+    parameter_type: type[parameter.BasevLLMParameter],
+    error: str,
+) -> None:
+    _set_tp3_rank2(monkeypatch)
+    column = ColumnParallelLinear(1, 6, bias=False, loaded_output_size=4)
+    column.weight_block_size = (4, 4)
+    kwargs = {
+        "data": torch.full((1, 1), 23.0),
+        "input_dim": 1,
+        "output_dim": 0,
+        "weight_loader": lambda *_args, **_kwargs: None,
+    }
+    if parameter_type is PackedvLLMParameter:
+        kwargs.update(packed_factor=4, packed_dim=0)
+    quantized_param = parameter_type(**kwargs)
+
+    with pytest.raises(ValueError, match=error):
+        column.weight_loader(
+            quantized_param,
+            torch.ones((1, 1), dtype=quantized_param.dtype),
+        )
+
+    torch.testing.assert_close(
+        quantized_param, quantized_param.new_full((1, 1), 23)
+    )
+
+
 @pytest.mark.parametrize("tp3", [False, True])
 def test_mla_projection_loaded_sizes_are_tp3_only(
     monkeypatch: pytest.MonkeyPatch,
diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index 0280ca30152b..daff56194f0c 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -65,7 +65,6 @@ def _validate_packed_loaded_size(
     loaded_size: int,
     name: str,
 ) -> None:
-    """Reject a padded packed layout before any destination writes occur."""
     if getattr(param, "packed_dim", None) != dim:
         return
     packed_factor = getattr(param, "packed_factor", 1)
@@ -75,6 +74,69 @@ def _validate_packed_loaded_size(
             f"{packed_factor}"
         )
 
+
+def _validate_padded_axis_layout(
+    param: Parameter,
+    dim: int,
+    loaded_size: int,
+    physical_size: int,
+    physical_shard_size: int,
+    physical_shard_offset: int,
+    name: str,
+    weight_block_size: tuple[int, ...] | None = None,
+) -> int:
+    """Validate a padded logical axis and return its checkpoint tensor size."""
+    if isinstance(param, BlockQuantScaleParameter):
+        if weight_block_size is None:
+            raise ValueError(f"{name} requires a weight block size")
+        if dim == getattr(param, "output_dim", None):
+            block_index = 0
+        elif dim == getattr(param, "input_dim", None):
+            block_index = 1
+        else:
+            raise ValueError(f"{name} does not identify a block-scaled axis")
+        if block_index >= len(weight_block_size):
+            raise ValueError(
+                f"{name} has no block size for parameter dimension {dim}"
+            )
+        block_size = weight_block_size[block_index]
+        shard_end = physical_shard_offset + physical_shard_size
+        boundaries = [
+            boundary
+            for boundary in (physical_shard_offset, shard_end)
+            if 0 < boundary < physical_size
+        ]
+        if physical_shard_size < physical_size:
+            boundaries.append(physical_shard_size)
+        if any(boundary % block_size for boundary in boundaries):
+            raise ValueError(
+                f"{name} physical TP boundaries are not aligned to "
+                f"quantization block size {block_size}"
+            )
+        return (loaded_size + block_size - 1) // block_size
+
+    if getattr(param, "packed_dim", None) == dim:
+        packed_factor = getattr(param, "packed_factor", 1)
+        shard_end = physical_shard_offset + physical_shard_size
+        for boundary_name, boundary in (
+            ("loaded size", loaded_size),
+            ("physical shard offset", physical_shard_offset),
+            ("physical shard end", shard_end),
+        ):
+            if boundary % packed_factor:
+                raise ValueError(
+                    f"{name} {boundary_name} {boundary} is not aligned to "
+                    f"packed_factor={packed_factor}"
+                )
+        if isinstance(param, (PackedColumnParameter, PackedvLLMParameter)):
+            packed_size, _ = param.adjust_shard_indexes_for_packing(
+                shard_size=loaded_size, shard_offset=0
+            )
+            return packed_size
+
+    return loaded_size
+
+
 WEIGHT_LOADER_V2_SUPPORTED = [
     "UnquantizedLinearMethod",
     "CompressedTensorsLinearMethod",
@@ -602,6 +664,18 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
                 self.loaded_output_size,
                 "loaded_output_size",
             )
+            expected_loaded_size = None
+            if self._allow_loaded_output_padding:
+                expected_loaded_size = _validate_padded_axis_layout(
+                    param,
+                    output_dim,
+                    self.loaded_output_size,
+                    self.output_size,
+                    self.output_size_per_partition,
+                    self.tp_rank * self.output_size_per_partition,
+                    "loaded_output_size",
+                    getattr(self, "weight_block_size", None),
+                )
             shard_size = param_data.shape[output_dim]
             start_idx = self.tp_rank * shard_size
             load_tensor_parallel_weight(
@@ -610,6 +684,7 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
                 output_dim,
                 start_idx,
                 allow_padding=self._allow_loaded_output_padding,
+                expected_loaded_size=expected_loaded_size,
             )
             return
 
@@ -628,9 +703,25 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
         if self._allow_loaded_output_padding:
+            output_dim = getattr(param, "output_dim", None)
+            if output_dim is None:
+                raise ValueError(
+                    "Padded column-parallel loading requires an output dimension"
+                )
+            expected_loaded_size = _validate_padded_axis_layout(
+                param,
+                output_dim,
+                self.loaded_output_size,
+                self.output_size,
+                self.output_size_per_partition,
+                self.tp_rank * self.output_size_per_partition,
+                "loaded_output_size",
+                getattr(self, "weight_block_size", None),
+            )
             param.load_column_parallel_weight(
                 loaded_weight=loaded_weight,
                 allow_padding=True,
+                expected_loaded_size=expected_loaded_size,
             )
         else:
             param.load_column_parallel_weight(loaded_weight=loaded_weight)
@@ -852,32 +943,59 @@ def weight_loader(
             for shard_id, output_size in zip(shard_ids, output_sizes):
                 shard_offsets.append((shard_id, current_shard_offset, output_size))
                 current_shard_offset += output_size
+            requires_padding = any(
+                self._allow_loaded_output_shard_padding[shard_id]
+                for shard_id in shard_ids
+            )
             packed_dim = getattr(param, "packed_dim", None)
+            transformed_shards: list[tuple[int, int, int]] = []
             for shard_id, shard_offset, shard_size in shard_offsets:
-                # Special case for Quantization.
-                # If quantized, we need to adjust the offset and size to account
-                # for the packing.
-                # Add check to adjust the size/offset for FP8 block scales
-                if isinstance(param, BlockQuantScaleParameter):
-                    weight_block_size = getattr(self, "weight_block_size", None)
-                    shard_size, shard_offset = adjust_block_scale_shard(
-                        weight_block_size, shard_size, shard_offset
+                _validate_packed_loaded_size(
+                    param,
+                    output_dim,
+                    self.loaded_output_sizes[shard_id],
+                    f"loaded_output_sizes[{shard_id}]",
+                )
+                if requires_padding:
+                    physical_shard_size = self.output_sizes[shard_id] // self.tp_size
+                    physical_shard_offset = (
+                        sum(self.output_sizes[:shard_id]) // self.tp_size
                     )
-
-                if packed_dim == output_dim:
-                    _validate_packed_loaded_size(
+                    _validate_padded_axis_layout(
                         param,
                         output_dim,
                         self.loaded_output_sizes[shard_id],
+                        sum(self.output_sizes) // self.tp_size,
+                        physical_shard_size,
+                        physical_shard_offset,
                         f"loaded_output_sizes[{shard_id}]",
+                        getattr(self, "weight_block_size", None),
                     )
+                if isinstance(param, BlockQuantScaleParameter):
+                    weight_block_size = getattr(self, "weight_block_size", None)
+                    shard_size, shard_offset = adjust_block_scale_shard(
+                        weight_block_size, shard_size, shard_offset
+                    )
+                if packed_dim == output_dim:
                     shard_size = shard_size // param.packed_factor
                     shard_offset = shard_offset // param.packed_factor
-                    # Special case for Marlin.
                     shard_size, shard_offset = adjust_marlin_shard(
                         param, shard_size, shard_offset
                     )
+                transformed_shards.append((shard_id, shard_offset, shard_size))
 
+            if requires_padding:
+                expected_loaded_size = sum(
+                    shard_size for _, _, shard_size in transformed_shards
+                )
+                if loaded_weight.shape[output_dim] != expected_loaded_size:
+                    raise ValueError(
+                        "Padded merged checkpoint axis size mismatch: "
+                        f"expected {expected_loaded_size}, "
+                        f"got {loaded_weight.shape[output_dim]}"
+                    )
+
+            for shard_id, shard_offset, shard_size in transformed_shards:
                 loaded_weight_shard = loaded_weight.narrow(
                     output_dim, shard_offset, shard_size
                 )
@@ -890,6 +1008,18 @@ def weight_loader(
             shard_size = self.output_sizes[loaded_shard_id]
             shard_offset //= self.tp_size
             shard_size //= self.tp_size
+            expected_loaded_size = None
+            if self._allow_loaded_output_shard_padding[loaded_shard_id]:
+                expected_loaded_size = _validate_padded_axis_layout(
+                    param,
+                    output_dim,
+                    self.loaded_output_sizes[loaded_shard_id],
+                    sum(self.output_sizes) // self.tp_size,
+                    shard_size,
+                    shard_offset,
+                    f"loaded_output_sizes[{loaded_shard_id}]",
+                    getattr(self, "weight_block_size", None),
+                )
 
             if isinstance(param, BlockQuantScaleParameter):
                 weight_block_size = getattr(self, "weight_block_size", None)
@@ -897,14 +1027,10 @@ def weight_loader(
                     weight_block_size, shard_size, shard_offset
                 )
 
-            # Special case for quantization.
-            # If quantized, we need to adjust the offset and size to account
-            # for the packing.
             packed_dim = getattr(param, "packed_dim", None)
             if packed_dim == output_dim:
                 shard_size = round(shard_size // param.packed_factor)
                 shard_offset = round(shard_offset // param.packed_factor)
-                # Special case for Marlin.
                 shard_size, shard_offset = adjust_marlin_shard(
                     param, shard_size, shard_offset
                 )
@@ -921,6 +1047,7 @@ def weight_loader(
                     allow_padding=self._allow_loaded_output_shard_padding[
                         loaded_shard_id
                     ],
+                    expected_loaded_size=expected_loaded_size,
                 )
                 return
         # Special case for per-tensor scales in fused case.
@@ -947,6 +1074,7 @@ def _load_fused_module_from_checkpoint(
         loaded_weight: torch.Tensor,
         output_sizes: list[int] | None = None,
         shard_ids: list[int] | None = None,
+        require_exact_axis: bool = False,
     ):
         """
         Handle special case for models where MLP layers are already
@@ -968,10 +1096,8 @@ def _load_fused_module_from_checkpoint(
             shard_offsets.append((shard_id, current_shard_offset, output_size))
             current_shard_offset += output_size
 
+        transformed_shards: list[tuple[int, int, int]] = []
         for shard_id, shard_offset, shard_size in shard_offsets:
-            # Special case for Quantization.
-            # If quantized, we need to adjust the offset and size to account
-            # for the packing.
             if (
                 isinstance(param, (PackedColumnParameter, PackedvLLMParameter))
                 and param.packed_dim == param.output_dim
@@ -979,7 +1105,20 @@ def _load_fused_module_from_checkpoint(
                 shard_size, shard_offset = param.adjust_shard_indexes_for_packing(
                     shard_size=shard_size, shard_offset=shard_offset
                 )
+            transformed_shards.append((shard_id, shard_offset, shard_size))
 
+        if require_exact_axis:
+            expected_loaded_size = sum(
+                shard_size for _, _, shard_size in transformed_shards
+            )
+            if loaded_weight.shape[param.output_dim] != expected_loaded_size:
+                raise ValueError(
+                    "Padded merged checkpoint axis size mismatch: "
+                    f"expected {expected_loaded_size}, "
+                    f"got {loaded_weight.shape[param.output_dim]}"
+                )
+
+        for shard_id, shard_offset, shard_size in transformed_shards:
             loaded_weight_shard = loaded_weight.narrow(
                 param.output_dim, shard_offset, shard_size
             )
@@ -1019,6 +1158,28 @@ def weight_loader_v2(
                 else list(range(len(self.loaded_output_sizes)))
             )
             output_sizes = [self.loaded_output_sizes[idx] for idx in shard_ids]
+            requires_padding = any(
+                self._allow_loaded_output_shard_padding[shard_id]
+                for shard_id in shard_ids
+            )
+            for shard_id in shard_ids:
+                _validate_packed_loaded_size(
+                    param,
+                    param.output_dim,
+                    self.loaded_output_sizes[shard_id],
+                    f"loaded_output_sizes[{shard_id}]",
+                )
+                if requires_padding:
+                    _validate_padded_axis_layout(
+                        param,
+                        param.output_dim,
+                        self.loaded_output_sizes[shard_id],
+                        sum(self.output_sizes) // self.tp_size,
+                        self.output_sizes[shard_id] // self.tp_size,
+                        sum(self.output_sizes[:shard_id]) // self.tp_size,
+                        f"loaded_output_sizes[{shard_id}]",
+                        getattr(self, "weight_block_size", None),
+                    )
             if isinstance(param, BlockQuantScaleParameter):
                 weight_block_size = getattr(self, "weight_block_size", None)
                 output_sizes = [
@@ -1031,6 +1192,7 @@ def weight_loader_v2(
                 loaded_weight,
                 output_sizes=output_sizes,
                 shard_ids=shard_ids,
+                require_exact_axis=requires_padding,
             )
             return
 
@@ -1040,6 +1202,18 @@ def weight_loader_v2(
         shard_size = self.output_sizes[loaded_shard_id]
         shard_offset //= self.tp_size
         shard_size //= self.tp_size
+        expected_loaded_size = None
+        if self._allow_loaded_output_shard_padding[loaded_shard_id]:
+            expected_loaded_size = _validate_padded_axis_layout(
+                param,
+                param.output_dim,
+                self.loaded_output_sizes[loaded_shard_id],
+                sum(self.output_sizes) // self.tp_size,
+                shard_size,
+                shard_offset,
+                f"loaded_output_sizes[{loaded_shard_id}]",
+                getattr(self, "weight_block_size", None),
+            )
 
         if isinstance(param, BlockQuantScaleParameter):
             weight_block_size = getattr(self, "weight_block_size", None)
@@ -1055,6 +1229,7 @@ def weight_loader_v2(
         }
         if self._allow_loaded_output_shard_padding[loaded_shard_id]:
             load_kwargs["allow_padding"] = True
+            load_kwargs["expected_loaded_size"] = expected_loaded_size
         param.load_merged_column_weight(**load_kwargs)
 
     def load_weights(
@@ -1224,6 +1399,34 @@ def _get_shard_size_mapping(self, loaded_shard_id: str):
         }
         return shard_size_mapping.get(loaded_shard_id)
 
+    def _get_loaded_shard_size_mapping(self, loaded_shard_id: str) -> int:
+        return {
+            "q": self.loaded_total_num_heads * self.head_size,
+            "k": self.loaded_total_num_kv_heads * self.head_size,
+            "v": self.loaded_total_num_kv_heads * self.v_head_size,
+        }[loaded_shard_id]
+
+    def _validate_padded_qkv_layout(
+        self,
+        param: Parameter,
+        output_dim: int,
+        loaded_shard_id: str,
+        shard_offset: int,
+        shard_size: int,
+    ) -> int:
+        physical_size = self._get_shard_offset_mapping("total")
+        assert physical_size is not None
+        return _validate_padded_axis_layout(
+            param,
+            output_dim,
+            self._get_loaded_shard_size_mapping(loaded_shard_id),
+            physical_size,
+            shard_size,
+            shard_offset,
+            f"loaded {loaded_shard_id} size",
+            getattr(self, "weight_block_size", None),
+        )
+
     def _load_fused_module_from_checkpoint(
         self, param: BasevLLMParameter, loaded_weight: torch.Tensor
     ):
@@ -1252,10 +1455,26 @@ def _load_fused_module_from_checkpoint(
             ),
         ]
 
+        requires_padding = any(self._allow_loaded_qkv_padding.values())
+        transformed_shards: list[tuple[str, int, int]] = []
         for shard_id, shard_offset, shard_size in shard_offsets:
-            # Special case for Quantization.
-            # If quantized, we need to adjust the offset and size to account
-            # for the packing.
+            local_shard_offset = self._get_shard_offset_mapping(shard_id)
+            local_shard_size = self._get_shard_size_mapping(shard_id)
+            assert local_shard_offset is not None and local_shard_size is not None
+            _validate_packed_loaded_size(
+                param,
+                param.output_dim,
+                self._get_loaded_shard_size_mapping(shard_id),
+                f"loaded {shard_id} size",
+            )
+            if requires_padding:
+                self._validate_padded_qkv_layout(
+                    param,
+                    param.output_dim,
+                    shard_id,
+                    local_shard_offset,
+                    local_shard_size,
+                )
             if isinstance(param, BlockQuantScaleParameter):
                 weight_block_size = getattr(self, "weight_block_size", None)
                 shard_size, shard_offset = adjust_block_scale_shard(
@@ -1268,7 +1487,20 @@ def _load_fused_module_from_checkpoint(
                 shard_size, shard_offset = param.adjust_shard_indexes_for_packing(
                     shard_size=shard_size, shard_offset=shard_offset
                 )
+            transformed_shards.append((shard_id, shard_offset, shard_size))
+
+        if requires_padding:
+            expected_loaded_size = sum(
+                shard_size for _, _, shard_size in transformed_shards
+            )
+            if loaded_weight.shape[param.output_dim] != expected_loaded_size:
+                raise ValueError(
+                    "Padded QKV checkpoint axis size mismatch: "
+                    f"expected {expected_loaded_size}, "
+                    f"got {loaded_weight.shape[param.output_dim]}"
+                )
 
+        for shard_id, shard_offset, shard_size in transformed_shards:
             loaded_weight_shard = loaded_weight.narrow(
                 param.output_dim, shard_offset, shard_size
             )
@@ -1303,6 +1535,15 @@ def weight_loader_v2(
         shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
         shard_size = self._get_shard_size_mapping(loaded_shard_id)
         assert shard_offset is not None and shard_size is not None
+        expected_loaded_size = None
+        if self._allow_loaded_qkv_padding[loaded_shard_id]:
+            expected_loaded_size = self._validate_padded_qkv_layout(
+                param,
+                param.output_dim,
+                loaded_shard_id,
+                shard_offset,
+                shard_size,
+            )
 
         if isinstance(param, BlockQuantScaleParameter):
             weight_block_size = getattr(self, "weight_block_size", None)
@@ -1319,6 +1560,7 @@ def weight_loader_v2(
         }
         if self._allow_loaded_qkv_padding[loaded_shard_id]:
             load_kwargs["allow_padding"] = True
+            load_kwargs["expected_loaded_size"] = expected_loaded_size
         param.load_qkv_weight(**load_kwargs)
 
     def weight_loader(
@@ -1363,34 +1605,51 @@ def weight_loader(
                 ),
             ]
             packed_dim = getattr(param, "packed_dim", None)
+            transformed_shards: list[tuple[str, int, int]] = []
+            requires_padding = any(self._allow_loaded_qkv_padding.values())
             for shard_id, shard_offset, shard_size in shard_offsets:
-                # Special case for Quantized Weights.
-                # If quantized, we need to adjust the offset and size to account
-                # for the packing.
-                # Add check to adjust the size/offset for FP8 block scales
+                local_shard_offset = self._get_shard_offset_mapping(shard_id)
+                local_shard_size = self._get_shard_size_mapping(shard_id)
+                assert local_shard_offset is not None and local_shard_size is not None
+                _validate_packed_loaded_size(
+                    param,
+                    output_dim,
+                    self._get_loaded_shard_size_mapping(shard_id),
+                    f"loaded {shard_id} size",
+                )
+                if requires_padding:
+                    self._validate_padded_qkv_layout(
+                        param,
+                        output_dim,
+                        shard_id,
+                        local_shard_offset,
+                        local_shard_size,
+                    )
                 if isinstance(param, BlockQuantScaleParameter):
                     weight_block_size = getattr(self, "weight_block_size", None)
                     shard_size, shard_offset = adjust_block_scale_shard(
                         weight_block_size, shard_size, shard_offset
                     )
-
                 if packed_dim == output_dim:
-                    loaded_size = {
-                        "q": self.loaded_total_num_heads * self.head_size,
-                        "k": self.loaded_total_num_kv_heads * self.head_size,
-                        "v": self.loaded_total_num_kv_heads * self.v_head_size,
-                    }[shard_id]
-                    _validate_packed_loaded_size(
-                        param, output_dim, loaded_size, f"loaded {shard_id} size"
-                    )
                     shard_size = round(shard_size // param.packed_factor)
                     shard_offset = round(shard_offset // param.packed_factor)
-
-                    # Special case for Marlin.
                     shard_size, shard_offset = adjust_marlin_shard(
                         param, shard_size, shard_offset
                     )
+                transformed_shards.append((shard_id, shard_offset, shard_size))
 
+            if requires_padding:
+                expected_loaded_size = sum(
+                    shard_size for _, _, shard_size in transformed_shards
+                )
+                if loaded_weight.shape[output_dim] != expected_loaded_size:
+                    raise ValueError(
+                        "Padded QKV checkpoint axis size mismatch: "
+                        f"expected {expected_loaded_size}, "
+                        f"got {loaded_weight.shape[output_dim]}"
+                    )
+
+            for shard_id, shard_offset, shard_size in transformed_shards:
                 loaded_weight_shard = loaded_weight.narrow(
                     output_dim, shard_offset, shard_size
                 )
@@ -1410,6 +1669,15 @@ def weight_loader(
             elif loaded_shard_id == "v":
                 shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size
                 shard_size = self.num_kv_heads * self.v_head_size
+            expected_loaded_size = None
+            if self._allow_loaded_qkv_padding[loaded_shard_id]:
+                expected_loaded_size = self._validate_padded_qkv_layout(
+                    param,
+                    output_dim,
+                    loaded_shard_id,
+                    shard_offset,
+                    shard_size,
+                )
 
             if isinstance(param, BlockQuantScaleParameter):
                 weight_block_size = getattr(self, "weight_block_size", None)
@@ -1447,6 +1715,7 @@ def weight_loader(
                     allow_padding=self._allow_loaded_qkv_padding[
                         loaded_shard_id
                     ],
+                    expected_loaded_size=expected_loaded_size,
                 )
                 return
 
@@ -1808,6 +2077,18 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
                 self.loaded_input_size,
                 "loaded_input_size",
             )
+            expected_loaded_size = None
+            if self._allow_loaded_input_padding:
+                expected_loaded_size = _validate_padded_axis_layout(
+                    param,
+                    input_dim,
+                    self.loaded_input_size,
+                    self.input_size,
+                    self.input_size_per_partition,
+                    self.tp_rank * self.input_size_per_partition,
+                    "loaded_input_size",
+                    getattr(self, "weight_block_size", None),
+                )
             shard_size = param_data.shape[input_dim]
             start_idx = self.tp_rank * shard_size
             load_tensor_parallel_weight(
@@ -1816,6 +2097,7 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
                 input_dim,
                 start_idx,
                 allow_padding=self._allow_loaded_input_padding,
+                expected_loaded_size=expected_loaded_size,
             )
             return
 
@@ -1835,9 +2117,25 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             loaded_weight = loaded_weight.reshape(1)
 
         if self._allow_loaded_input_padding:
+            input_dim = getattr(param, "input_dim", None)
+            if input_dim is None:
+                raise ValueError(
+                    "Padded row-parallel loading requires an input dimension"
+                )
+            expected_loaded_size = _validate_padded_axis_layout(
+                param,
+                input_dim,
+                self.loaded_input_size,
+                self.input_size,
+                self.input_size_per_partition,
+                self.tp_rank * self.input_size_per_partition,
+                "loaded_input_size",
+                getattr(self, "weight_block_size", None),
+            )
             param.load_row_parallel_weight(
                 loaded_weight=loaded_weight,
                 allow_padding=True,
+                expected_loaded_size=expected_loaded_size,
             )
         else:
             param.load_row_parallel_weight(loaded_weight=loaded_weight)
diff --git a/vllm/model_executor/parameter.py b/vllm/model_executor/parameter.py
index bdb2d4b714b4..541da213cbaf 100644
--- a/vllm/model_executor/parameter.py
+++ b/vllm/model_executor/parameter.py
@@ -35,6 +35,7 @@ def load_tensor_parallel_weight(
     start_idx: int,
     *,
     allow_padding: bool = False,
+    expected_loaded_size: int | None = None,
 ) -> None:
     """Copy one TP shard, optionally zero-filling its unavailable tail.
 
@@ -62,6 +63,15 @@ def load_tensor_parallel_weight(
         raise IndexError(
             f"TP shard dimension {dim} is invalid for rank {loaded_weight.ndim}"
         )
+    if expected_loaded_size is None:
+        raise ValueError(
+            "Padded TP loading requires the expected checkpoint axis size"
+        )
+    if loaded_weight.shape[dim] != expected_loaded_size:
+        raise ValueError(
+            "Padded TP checkpoint axis size mismatch: "
+            f"expected {expected_loaded_size}, got {loaded_weight.shape[dim]}"
+        )
     for axis, (destination_size, loaded_size) in enumerate(
         zip(param_data.shape, loaded_weight.shape)
     ):
@@ -152,6 +162,7 @@ def load_column_parallel_weight(
         loaded_weight: torch.Tensor,
         *,
         allow_padding: bool = False,
+        expected_loaded_size: int | None = None,
     ):
         if allow_padding and self.data.shape != loaded_weight.shape:
             raise ValueError(
@@ -165,6 +176,7 @@ def load_row_parallel_weight(
         loaded_weight: torch.Tensor,
         *,
         allow_padding: bool = False,
+        expected_loaded_size: int | None = None,
     ):
         if allow_padding and self.data.shape != loaded_weight.shape:
             raise ValueError(
@@ -221,6 +233,7 @@ def load_column_parallel_weight(
         loaded_weight: torch.Tensor,
         *,
         allow_padding: bool = False,
+        expected_loaded_size: int | None = None,
     ):
         load_tensor_parallel_weight(
             self.data,
@@ -228,6 +241,7 @@ def load_column_parallel_weight(
             self.output_dim,
             self.tp_rank * self.data.shape[self.output_dim],
             allow_padding=allow_padding,
+            expected_loaded_size=expected_loaded_size,
         )
 
     def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):
@@ -253,6 +267,7 @@ def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):
             self.output_dim,
             self.tp_rank * shard_size,
             allow_padding=allow_padding,
+            expected_loaded_size=kwargs.get("expected_loaded_size"),
         )
 
     def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):
@@ -281,6 +296,7 @@ def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):
             self.output_dim,
             shard_id_int * shard_size,
             allow_padding=allow_padding,
+            expected_loaded_size=kwargs.get("expected_loaded_size"),
         )
 
 
@@ -305,6 +321,7 @@ def load_row_parallel_weight(
         loaded_weight: torch.Tensor,
         *,
         allow_padding: bool = False,
+        expected_loaded_size: int | None = None,
     ):
         param_data = self.data
         load_tensor_parallel_weight(
@@ -313,6 +330,7 @@ def load_row_parallel_weight(
             self.input_dim,
             self.tp_rank * param_data.shape[self.input_dim],
             allow_padding=allow_padding,
+            expected_loaded_size=expected_loaded_size,
         )
 
 

From 2bc24bc07d3cede76eb5be070922602ad2825e1d Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 12:56:24 +0200
Subject: [PATCH 18/25] Copy DFlash parallel config safely

(cherry picked from commit 9dcde152f65381fb4d072b60d711a6004a890ea0)
---
 vllm/v1/spec_decode/dflash.py | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py
index 354d0a302416..7ab450c8ca2d 100644
--- a/vllm/v1/spec_decode/dflash.py
+++ b/vllm/v1/spec_decode/dflash.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from copy import copy
 from dataclasses import replace
 from typing import Any
 
@@ -98,15 +99,12 @@ def load_model(self, target_model: torch.nn.Module) -> None:
     @override
     def _create_draft_vllm_config(self) -> VllmConfig:
         base = super()._create_draft_vllm_config()
-        draft_parallel_config = self.speculative_config.draft_parallel_config
-        assert draft_parallel_config is not None
-        base = replace(
-            base,
-            parallel_config=replace(
-                draft_parallel_config,
-                rank=self.vllm_config.parallel_config.rank,
-            ),
+        draft_parallel_config = copy(
+            self.speculative_config.draft_parallel_config
         )
+        assert draft_parallel_config is not None
+        draft_parallel_config.rank = self.vllm_config.parallel_config.rank
+        base = replace(base, parallel_config=draft_parallel_config)
         # The draft model is text-only — clear the target's multimodal
         # flag so flash_attn is not rejected for mm_prefix support.
         arch = base.model_config.model_arch_config

From 5bdf4375e38afa6498da0fb4f462d5766f226056 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 13:40:47 +0200
Subject: [PATCH 19/25] fix(tp): load padded NVFP4 storage layouts

(cherry picked from commit 977efc978fb004051b45ee390e0305e054bd26f2)
---
 tests/models/test_glm53_tp3_model.py          | 69 +++++++++++++++++++
 vllm/model_executor/layers/linear.py          | 57 ++++++++++-----
 .../schemes/compressed_tensors_w4a4_nvfp4.py  |  2 +
 vllm/model_executor/parameter.py              | 14 +++-
 4 files changed, 122 insertions(+), 20 deletions(-)

diff --git a/tests/models/test_glm53_tp3_model.py b/tests/models/test_glm53_tp3_model.py
index 3e1caac08d6a..fe4e7a995f92 100644
--- a/tests/models/test_glm53_tp3_model.py
+++ b/tests/models/test_glm53_tp3_model.py
@@ -16,7 +16,10 @@
 )
 from vllm.model_executor.parameter import (
     BlockQuantScaleParameter,
+    GroupQuantScaleParameter,
+    ModelWeightParameter,
     PackedvLLMParameter,
+    PerTensorScaleParameter,
 )
 from vllm.models.glm5next.nvidia import attention as glm_attention
 from vllm.models.glm5next.nvidia import model as glm_model
@@ -191,6 +194,72 @@ def test_padded_loader_rejects_unaligned_quantized_boundaries_before_write(
     )
 
 
+def test_padded_v2_loader_preserves_unsharded_per_tensor_scale(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _set_tp3_rank2(monkeypatch)
+    row = RowParallelLinear(6, 1, bias=False, loaded_input_size=4)
+    scale = PerTensorScaleParameter(
+        data=torch.zeros(1, dtype=torch.float32),
+        weight_loader=lambda *_args, **_kwargs: None,
+    )
+
+    row.weight_loader_v2(scale, torch.tensor([3.0]))
+
+    torch.testing.assert_close(scale, torch.tensor([3.0]))
+
+
+@pytest.mark.parametrize(
+    ("parameter_type", "storage_factor", "physical_size", "loaded_size", "dtype"),
+    [
+        (ModelWeightParameter, 2, 6, 4, torch.uint8),
+        (GroupQuantScaleParameter, 16, 96, 64, torch.float8_e4m3fn),
+    ],
+)
+def test_padded_nvfp4_row_loader_converts_logical_to_storage_width(
+    monkeypatch: pytest.MonkeyPatch,
+    parameter_type: type[parameter.BasevLLMParameter],
+    storage_factor: int,
+    physical_size: int,
+    loaded_size: int,
+    dtype: torch.dtype,
+) -> None:
+    for module in (linear, parameter):
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_world_size", lambda: 3
+        )
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_rank", lambda: 1
+        )
+    row = RowParallelLinear(
+        physical_size,
+        1,
+        bias=False,
+        loaded_input_size=loaded_size,
+    )
+    local_storage_size = physical_size // 3 // storage_factor
+    loaded_storage_size = loaded_size // storage_factor
+    quantized_param = parameter_type(
+        data=torch.zeros((1, local_storage_size), dtype=dtype),
+        input_dim=1,
+        input_dim_storage_factor=storage_factor,
+        output_dim=0,
+        weight_loader=lambda *_args, **_kwargs: None,
+    )
+    loaded = torch.arange(
+        1,
+        loaded_storage_size + 1,
+        dtype=torch.uint8,
+    ).to(dtype)
+
+    row.weight_loader_v2(quantized_param, loaded.unsqueeze(0))
+
+    torch.testing.assert_close(
+        quantized_param,
+        loaded[local_storage_size : 2 * local_storage_size].unsqueeze(0),
+    )
+
+
 @pytest.mark.parametrize("tp3", [False, True])
 def test_mla_projection_loaded_sizes_are_tp3_only(
     monkeypatch: pytest.MonkeyPatch,
diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py
index daff56194f0c..b65c85b9d426 100644
--- a/vllm/model_executor/layers/linear.py
+++ b/vllm/model_executor/layers/linear.py
@@ -115,6 +115,25 @@ def _validate_padded_axis_layout(
             )
         return (loaded_size + block_size - 1) // block_size
 
+    storage_factor = (
+        getattr(param, "input_dim_storage_factor", 1)
+        if dim == getattr(param, "input_dim", None)
+        else 1
+    )
+    if storage_factor != 1:
+        shard_end = physical_shard_offset + physical_shard_size
+        for boundary_name, boundary in (
+            ("loaded size", loaded_size),
+            ("physical shard offset", physical_shard_offset),
+            ("physical shard end", shard_end),
+        ):
+            if boundary % storage_factor:
+                raise ValueError(
+                    f"{name} {boundary_name} {boundary} is not aligned to "
+                    f"input_dim_storage_factor={storage_factor}"
+                )
+        return loaded_size // storage_factor
+
     if getattr(param, "packed_dim", None) == dim:
         packed_factor = getattr(param, "packed_factor", 1)
         shard_end = physical_shard_offset + physical_shard_size
@@ -702,12 +721,8 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
         if len(loaded_weight.shape) == 0:
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
-        if self._allow_loaded_output_padding:
-            output_dim = getattr(param, "output_dim", None)
-            if output_dim is None:
-                raise ValueError(
-                    "Padded column-parallel loading requires an output dimension"
-                )
+        output_dim = getattr(param, "output_dim", None)
+        if self._allow_loaded_output_padding and output_dim is not None:
             expected_loaded_size = _validate_padded_axis_layout(
                 param,
                 output_dim,
@@ -1202,11 +1217,16 @@ def weight_loader_v2(
         shard_size = self.output_sizes[loaded_shard_id]
         shard_offset //= self.tp_size
         shard_size //= self.tp_size
+        output_dim = getattr(param, "output_dim", None)
+        has_padded_axis = (
+            self._allow_loaded_output_shard_padding[loaded_shard_id]
+            and output_dim is not None
+        )
         expected_loaded_size = None
-        if self._allow_loaded_output_shard_padding[loaded_shard_id]:
+        if has_padded_axis:
             expected_loaded_size = _validate_padded_axis_layout(
                 param,
-                param.output_dim,
+                output_dim,
                 self.loaded_output_sizes[loaded_shard_id],
                 sum(self.output_sizes) // self.tp_size,
                 shard_size,
@@ -1227,7 +1247,7 @@ def weight_loader_v2(
             "shard_offset": shard_offset,
             "shard_size": shard_size,
         }
-        if self._allow_loaded_output_shard_padding[loaded_shard_id]:
+        if has_padded_axis:
             load_kwargs["allow_padding"] = True
             load_kwargs["expected_loaded_size"] = expected_loaded_size
         param.load_merged_column_weight(**load_kwargs)
@@ -1535,11 +1555,16 @@ def weight_loader_v2(
         shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
         shard_size = self._get_shard_size_mapping(loaded_shard_id)
         assert shard_offset is not None and shard_size is not None
+        output_dim = getattr(param, "output_dim", None)
+        has_padded_axis = (
+            self._allow_loaded_qkv_padding[loaded_shard_id]
+            and output_dim is not None
+        )
         expected_loaded_size = None
-        if self._allow_loaded_qkv_padding[loaded_shard_id]:
+        if has_padded_axis:
             expected_loaded_size = self._validate_padded_qkv_layout(
                 param,
-                param.output_dim,
+                output_dim,
                 loaded_shard_id,
                 shard_offset,
                 shard_size,
@@ -1558,7 +1583,7 @@ def weight_loader_v2(
             "shard_offset": shard_offset,
             "shard_size": shard_size,
         }
-        if self._allow_loaded_qkv_padding[loaded_shard_id]:
+        if has_padded_axis:
             load_kwargs["allow_padding"] = True
             load_kwargs["expected_loaded_size"] = expected_loaded_size
         param.load_qkv_weight(**load_kwargs)
@@ -2116,12 +2141,8 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
 
-        if self._allow_loaded_input_padding:
-            input_dim = getattr(param, "input_dim", None)
-            if input_dim is None:
-                raise ValueError(
-                    "Padded row-parallel loading requires an input dimension"
-                )
+        input_dim = getattr(param, "input_dim", None)
+        if self._allow_loaded_input_padding and input_dim is not None:
             expected_loaded_size = _validate_padded_axis_layout(
                 param,
                 input_dim,
diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
index c737b057fcf0..d5b695a4e9aa 100644
--- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
@@ -57,6 +57,7 @@ def create_weights(
                 dtype=torch.uint8,
             ),
             input_dim=1,
+            input_dim_storage_factor=2,
             output_dim=0,
             weight_loader=weight_loader,
         )
@@ -77,6 +78,7 @@ def create_weights(
                 dtype=torch.float8_e4m3fn,
             ),
             input_dim=1,
+            input_dim_storage_factor=self.group_size,
             output_dim=0,
             weight_loader=weight_loader,
         )
diff --git a/vllm/model_executor/parameter.py b/vllm/model_executor/parameter.py
index 541da213cbaf..8f9bf59c96b2 100644
--- a/vllm/model_executor/parameter.py
+++ b/vllm/model_executor/parameter.py
@@ -305,17 +305,27 @@ class RowvLLMParameter(BasevLLMParameter):
     Parameter class defining weight_loading functionality
     (load_row_parallel_weight) for parameters being loaded
     into linear layers with row parallel functionality.
-    Requires an input_dim to be defined.
+    Requires an input dimension to be defined.
     """
 
-    def __init__(self, input_dim: int, **kwargs):
+    def __init__(
+        self, input_dim: int, input_dim_storage_factor: int = 1, **kwargs
+    ):
+        if input_dim_storage_factor < 1:
+            raise ValueError("input_dim_storage_factor must be positive")
         self._input_dim = input_dim
+        self._input_dim_storage_factor = input_dim_storage_factor
         super().__init__(**kwargs)
 
     @property
     def input_dim(self):
         return self._input_dim
 
+    @property
+    def input_dim_storage_factor(self):
+        """Logical input-axis elements represented by one storage element."""
+        return self._input_dim_storage_factor
+
     def load_row_parallel_weight(
         self,
         loaded_weight: torch.Tensor,

From 6e50fa06451be52b8def9a7d529c441dcba3f607 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 13:50:17 +0200
Subject: [PATCH 20/25] fix(tp): load ModelOpt padded NVFP4 weights

(cherry picked from commit 8bb285c9d26e5bf0e7e3241b13ed0d9eee6ccd43)
---
 tests/models/test_glm53_tp3_model.py          | 65 +++++++++++++++++++
 .../v1/attention/test_b12x_sparse_mla_api.py  |  2 +-
 .../layers/quantization/modelopt.py           |  4 ++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/tests/models/test_glm53_tp3_model.py b/tests/models/test_glm53_tp3_model.py
index fe4e7a995f92..e6d7aeb08cdb 100644
--- a/tests/models/test_glm53_tp3_model.py
+++ b/tests/models/test_glm53_tp3_model.py
@@ -8,6 +8,7 @@
 
 from vllm.model_executor import parameter
 from vllm.model_executor.layers import linear
+from vllm.model_executor.layers.quantization import modelopt
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -260,6 +261,70 @@ def test_padded_nvfp4_row_loader_converts_logical_to_storage_width(
     )
 
 
+
+@pytest.mark.parametrize(
+    "method_type",
+    (
+        modelopt.ModelOptNvFp4LinearMethod,
+        modelopt.ModelOptNvFp4W4A16LinearMethod,
+    ),
+)
+def test_modelopt_nvfp4_row_parameters_declare_storage_widths(
+    monkeypatch: pytest.MonkeyPatch,
+    method_type: type,
+) -> None:
+    for module in (linear, parameter):
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_world_size", lambda: 3
+        )
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_rank", lambda: 1
+        )
+    monkeypatch.setattr(
+        modelopt,
+        "init_nvfp4_linear_kernel",
+        lambda **_kwargs: SimpleNamespace(input_quant_key=lambda: None),
+    )
+    config = modelopt.ModelOptNvFp4Config(
+        quant_method=(
+            "NVFP4"
+            if method_type is modelopt.ModelOptNvFp4LinearMethod
+            else "W4A16_NVFP4"
+        ),
+        is_checkpoint_nvfp4_serialized=True,
+        group_size=16,
+    )
+    holder = torch.nn.Module()
+    method_type(config).create_weights(
+        holder,
+        input_size_per_partition=32,
+        output_partition_sizes=[1],
+        input_size=96,
+        output_size=1,
+        params_dtype=torch.bfloat16,
+        weight_loader=lambda *_args, **_kwargs: None,
+    )
+
+    assert holder.weight.input_dim_storage_factor == 2
+    assert holder.weight_scale.input_dim_storage_factor == 16
+
+    row = RowParallelLinear(96, 1, bias=False, loaded_input_size=64)
+    for param, storage_factor in (
+        (holder.weight, 2),
+        (holder.weight_scale, 16),
+    ):
+        loaded = torch.arange(
+            1,
+            64 // storage_factor + 1,
+            dtype=torch.uint8,
+        ).to(param.dtype)
+        row.weight_loader_v2(param, loaded.unsqueeze(0))
+        local_width = 32 // storage_factor
+        torch.testing.assert_close(
+            param,
+            loaded[local_width : 2 * local_width].unsqueeze(0),
+        )
+
 @pytest.mark.parametrize("tp3", [False, True])
 def test_mla_projection_loaded_sizes_are_tp3_only(
     monkeypatch: pytest.MonkeyPatch,
diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py
index d8e8a3d9ef08..73bc7c4d33bf 100644
--- a/tests/v1/attention/test_b12x_sparse_mla_api.py
+++ b/tests/v1/attention/test_b12x_sparse_mla_api.py
@@ -926,7 +926,7 @@ def test_b12x_glm_dsa_nvfp4_cache_writer_keeps_rope() -> None:
     impl._uses_glm_dsa_nvfp4_cache = True
     impl._concat_and_cache_nvfp4_mla_fp8_rope = lambda *args: calls.append(args)
     kv_c = torch.empty((3, 512), dtype=torch.bfloat16)
-    k_pe = torch.empty((3, 1, 64), dtype=torch.bfloat16)
+    k_pe = torch.zeros((3, 1, 64), dtype=torch.bfloat16)
     kv_cache = torch.empty((2, 64, 368), dtype=torch.uint8)
     slots = torch.tensor([0, 64, -1], dtype=torch.int64)
     scale = torch.ones((), dtype=torch.float32)
diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py
index 241a7c8adaf3..72f70b56a1fe 100644
--- a/vllm/model_executor/layers/quantization/modelopt.py
+++ b/vllm/model_executor/layers/quantization/modelopt.py
@@ -1164,6 +1164,7 @@ def create_weights(
                 dtype=torch.uint8,
             ),
             input_dim=1,
+            input_dim_storage_factor=2,
             output_dim=0,
             weight_loader=weight_loader,
         )
@@ -1191,6 +1192,7 @@ def create_weights(
                 dtype=weight_dtype,
             ),
             input_dim=1,
+            input_dim_storage_factor=self.quant_config.group_size,
             output_dim=0,
             weight_loader=weight_loader,
         )
@@ -1308,6 +1310,7 @@ def create_weights(
                 dtype=torch.uint8,
             ),
             input_dim=1,
+            input_dim_storage_factor=2,
             output_dim=0,
             weight_loader=weight_loader,
         )
@@ -1330,6 +1333,7 @@ def create_weights(
                 dtype=torch.float8_e4m3fn,
             ),
             input_dim=1,
+            input_dim_storage_factor=self.quant_config.group_size,
             output_dim=0,
             weight_loader=weight_loader,
         )

From b0040a01a55314a58740fd894116b8397a33f5e2 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 16:20:31 +0200
Subject: [PATCH 21/25] fix(modelopt): shard MXFP8 scales by storage width

(cherry picked from commit 90fec81284e69e982201285194f80f0b6b392b80)
---
 tests/models/test_glm53_tp3_model.py          | 35 +++++++++++++++++++
 .../layers/quantization/modelopt.py           |  1 +
 2 files changed, 36 insertions(+)

diff --git a/tests/models/test_glm53_tp3_model.py b/tests/models/test_glm53_tp3_model.py
index e6d7aeb08cdb..092b5a7bd9f3 100644
--- a/tests/models/test_glm53_tp3_model.py
+++ b/tests/models/test_glm53_tp3_model.py
@@ -325,6 +325,41 @@ def test_modelopt_nvfp4_row_parameters_declare_storage_widths(
             loaded[local_width : 2 * local_width].unsqueeze(0),
         )
 
+
+def test_modelopt_mxfp8_scale_loads_padded_tp3_storage_width(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    for module in (linear, parameter):
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_world_size", lambda: 3
+        )
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_rank", lambda: 1
+        )
+    monkeypatch.setattr(modelopt, "init_mxfp8_linear_kernel", lambda: object())
+    config = modelopt.ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+    )
+    holder = torch.nn.Module()
+    modelopt.ModelOptMxFp8LinearMethod(config).create_weights(
+        holder,
+        input_size_per_partition=4096,
+        output_partition_sizes=[1],
+        input_size=12288,
+        output_size=1,
+        params_dtype=torch.bfloat16,
+        weight_loader=lambda *_args, **_kwargs: None,
+    )
+
+    assert holder.weight_scale.input_dim_storage_factor == 32
+    loaded = torch.arange(1, 129, dtype=holder.weight_scale.dtype).unsqueeze(0)
+    row = RowParallelLinear(12288, 1, bias=False, loaded_input_size=4096)
+    row.weight_loader_v2(holder.weight_scale, loaded)
+    torch.testing.assert_close(holder.weight_scale, torch.zeros_like(holder.weight_scale))
+
+
 @pytest.mark.parametrize("tp3", [False, True])
 def test_mla_projection_loaded_sizes_are_tp3_only(
     monkeypatch: pytest.MonkeyPatch,
diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py
index 72f70b56a1fe..2b42454c3260 100644
--- a/vllm/model_executor/layers/quantization/modelopt.py
+++ b/vllm/model_executor/layers/quantization/modelopt.py
@@ -1836,6 +1836,7 @@ def create_weights(
                 dtype=MXFP8_SCALE_DTYPE,
             ),
             input_dim=1,
+            input_dim_storage_factor=MXFP8_BLOCK_SIZE,
             output_dim=0,
             weight_loader=weight_loader,
         )

From 531fc1cf902e2254ced467805a2afdd443626be3 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 18:16:00 +0200
Subject: [PATCH 22/25] fix(spec-decode): propagate engine DP identity to draft

(cherry picked from commit 4025af4404c01bdf72a892553d1ede1d1a6f77a5)
---
 tests/models/test_glm53_tp3_dflash.py | 30 ++++++++++++--
 vllm/config/speculative.py            | 57 +++++++++++++++++++++++++++
 vllm/v1/engine/core.py                | 23 ++++++++++-
 3 files changed, 106 insertions(+), 4 deletions(-)

diff --git a/tests/models/test_glm53_tp3_dflash.py b/tests/models/test_glm53_tp3_dflash.py
index a85b0aba5059..362e480ca87b 100644
--- a/tests/models/test_glm53_tp3_dflash.py
+++ b/tests/models/test_glm53_tp3_dflash.py
@@ -214,9 +214,7 @@ def test_dflash_logits_and_selector_exclude_physical_vocab_tail() -> None:
     logits = processor._get_logits(torch.empty(1, 1), lm_head, None)
     assert logits is not None
     assert logits.shape == (1, logical_vocab_size)
-    token_ids, values = processor.get_top_k_tokens(
-        lm_head, torch.empty(1, 1), k=1
-    )
+    token_ids, values = processor.get_top_k_tokens(lm_head, torch.empty(1, 1), k=1)
     assert token_ids.item() == logical_vocab_size - 1
     assert values.item() == 2
 
@@ -284,6 +282,7 @@ def test_dense_dflash_tp3_drops_target_expert_parallel(monkeypatch) -> None:
         "data_parallel_size_local": 2,
         "data_parallel_rank": 1,
         "data_parallel_rank_local": 1,
+        "data_parallel_index": 1,
         "data_parallel_master_ip": "127.0.0.1",
         "data_parallel_rpc_port": 1234,
         "data_parallel_master_port": 4321,
@@ -349,6 +348,31 @@ class DraftVllmConfig:
     assert target_parallel.enable_expert_parallel
 
 
+def test_engine_dp_identity_reaches_speculative_draft() -> None:
+    from vllm.v1.engine.core import EngineCoreProc
+
+    target = SimpleNamespace(
+        data_parallel_index=3,
+        data_parallel_rank=3,
+        data_parallel_rank_local=1,
+    )
+    draft = SimpleNamespace(
+        data_parallel_index=0,
+        data_parallel_rank=0,
+        data_parallel_rank_local=0,
+    )
+    vllm_config = SimpleNamespace(
+        parallel_config=target,
+        speculative_config=SimpleNamespace(draft_parallel_config=draft),
+    )
+
+    EngineCoreProc._sync_speculative_draft_dp_identity(vllm_config)
+
+    assert draft.data_parallel_index == 3
+    assert draft.data_parallel_rank == 3
+    assert draft.data_parallel_rank_local == 1
+
+
 def test_dflash7_uses_eight_target_kda_state_columns(monkeypatch) -> None:
     def fake_base_init(self, vllm_config, device) -> None:
         self.max_num_tokens = 16
diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
index e159da01a267..b8cb6f0e36eb 100644
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
@@ -1075,6 +1075,15 @@ def __post_init__(self):
         # will be detected automatically if possible. If the speculative method
         # can not be detected, it will be considered as the "draft_model" by
         # default.
+        # The target must expose its physical TP3 axes before a draft config is
+        # derived from it. This is an exact no-op for TP4 and other models.
+        from vllm.transformers_utils.configs.glm53_tp3 import (
+            apply_glm53_tp3_target_geometry,
+        )
+
+        apply_glm53_tp3_target_geometry(
+            self.target_model_config, self.target_parallel_config
+        )
 
         # infer method from user args
         if self.method is None and SpeculativeConfig._is_custom_proposer_path(
@@ -1537,6 +1546,7 @@ def __post_init__(self):
                 "adaptive_speculative_tokens_initial must not exceed "
                 "num_speculative_tokens."
             )
+        self._apply_glm53_tp3_draft_geometry()
 
         return self
 
@@ -1708,6 +1718,53 @@ def update_arch_(self):
         self.draft_model_config._model_info = model_info
         self.draft_model_config._architecture = arch
 
+    def _apply_glm53_tp3_draft_geometry(self) -> None:
+        from vllm.transformers_utils.configs.glm53_tp3 import (
+            apply_glm53_tp3_draft_geometry,
+            is_glm53_config,
+        )
+
+        if (
+            self.method not in ("mtp", "dflash")
+            or self.target_model_config is None
+            or self.target_parallel_config is None
+            or self.draft_model_config is None
+            or self.draft_parallel_config is None
+            or self.target_parallel_config.tensor_parallel_size != 3
+            or not is_glm53_config(self.target_model_config)
+        ):
+            return
+
+        # MTP participates in the target MoE's topology. DFlash is dense, so it
+        # keeps DP/PCP placement but must not join the routed-expert EP group.
+        target = self.target_parallel_config
+        draft = self.draft_parallel_config
+        for name in (
+            "prefill_context_parallel_size",
+            "data_parallel_size",
+            "data_parallel_size_local",
+            "data_parallel_rank",
+            "data_parallel_rank_local",
+            "data_parallel_index",
+            "data_parallel_master_ip",
+            "data_parallel_rpc_port",
+            "data_parallel_master_port",
+            "data_parallel_backend",
+            "data_parallel_external_lb",
+            "data_parallel_hybrid_lb",
+        ):
+            setattr(draft, name, getattr(target, name))
+        draft.enable_expert_parallel = (
+            target.enable_expert_parallel if self.method == "mtp" else False
+        )
+
+        apply_glm53_tp3_draft_geometry(
+            self.target_model_config,
+            target,
+            self.draft_model_config,
+            draft,
+        )
+
     @staticmethod
     def create_draft_parallel_config(
         target_parallel_config: ParallelConfig,
diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
index a69b89558a87..082e0b6c0bd0 100644
--- a/vllm/v1/engine/core.py
+++ b/vllm/v1/engine/core.py
@@ -1419,6 +1419,23 @@ def startup_handshake(
 
         return init_message.addresses
 
+    @staticmethod
+    def _sync_speculative_draft_dp_identity(vllm_config: VllmConfig) -> None:
+        speculative_config = vllm_config.speculative_config
+        if (
+            speculative_config is None
+            or speculative_config.draft_parallel_config is None
+        ):
+            return
+        target = vllm_config.parallel_config
+        draft = speculative_config.draft_parallel_config
+        for name in (
+            "data_parallel_index",
+            "data_parallel_rank",
+            "data_parallel_rank_local",
+        ):
+            setattr(draft, name, getattr(target, name))
+
     @staticmethod
     def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs):
         """Launch EngineCore busy loop in background process."""
@@ -1459,12 +1476,15 @@ def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs):
             if data_parallel and vllm_config.model_config.is_moe:
                 # Set data parallel rank for this engine process.
                 parallel_config.data_parallel_rank = dp_rank
-                engine_core = DPEngineCoreProc(*args, **kwargs)
             else:
                 # Non-MoE DP ranks are completely independent, so treat like DP=1.
                 # Note that parallel_config.data_parallel_index will still reflect
                 # the original DP rank.
                 parallel_config.reconfigure_for_independent_dp_rank()
+            EngineCoreProc._sync_speculative_draft_dp_identity(vllm_config)
+            if data_parallel and vllm_config.model_config.is_moe:
+                engine_core = DPEngineCoreProc(*args, **kwargs)
+            else:
                 engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)
 
             assert engine_core is not None
@@ -2558,6 +2578,7 @@ def __init__(
         self.addresses = addresses
         vllm_config.parallel_config.data_parallel_index = dp_rank
         vllm_config.parallel_config.data_parallel_rank_local = local_dp_rank
+        EngineCoreProc._sync_speculative_draft_dp_identity(vllm_config)
 
         self._set_nixl_side_channel_host()
 

From 02bad884049babfcb50596ead00ac34955184486 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Thu, 3 Sep 2026 21:26:47 +0200
Subject: [PATCH 23/25] test(glm53): clean TP3 regression coverage

(cherry picked from commit 059f85d25832d0542e774d2dccfb3a2c94def51e)
---
 tests/models/test_glm53_tp3_model.py     | 72 +++++++-----------------
 tests/models/test_glm5next_vision_tp3.py |  4 +-
 tests/v1/attention/test_dflash_attn.py   |  3 -
 3 files changed, 23 insertions(+), 56 deletions(-)

diff --git a/tests/models/test_glm53_tp3_model.py b/tests/models/test_glm53_tp3_model.py
index 092b5a7bd9f3..0bdadd7e1f0d 100644
--- a/tests/models/test_glm53_tp3_model.py
+++ b/tests/models/test_glm53_tp3_model.py
@@ -8,13 +8,13 @@
 
 from vllm.model_executor import parameter
 from vllm.model_executor.layers import linear
-from vllm.model_executor.layers.quantization import modelopt
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
     RowParallelLinear,
 )
+from vllm.model_executor.layers.quantization import modelopt
 from vllm.model_executor.parameter import (
     BlockQuantScaleParameter,
     GroupQuantScaleParameter,
@@ -40,15 +40,11 @@ def test_explicit_loaded_sizes_zero_rank_local_destination_tails(
 ) -> None:
     _set_tp3_rank2(monkeypatch)
 
-    column = ColumnParallelLinear(
-        1, 6, bias=False, loaded_output_size=5
-    )
+    column = ColumnParallelLinear(1, 6, bias=False, loaded_output_size=5)
     column.weight.weight_loader(
         column.weight, torch.arange(1, 6, dtype=column.weight.dtype).unsqueeze(1)
     )
-    torch.testing.assert_close(
-        column.weight[:, 0], column.weight.new_tensor([5, 0])
-    )
+    torch.testing.assert_close(column.weight[:, 0], column.weight.new_tensor([5, 0]))
 
     merged = MergedColumnParallelLinear(
         1, [6, 6], bias=False, loaded_output_sizes=[5, 5]
@@ -58,9 +54,7 @@ def test_explicit_loaded_sizes_zero_rank_local_destination_tails(
         torch.arange(1, 6, dtype=merged.weight.dtype).unsqueeze(1),
         0,
     )
-    torch.testing.assert_close(
-        merged.weight[:2, 0], merged.weight.new_tensor([5, 0])
-    )
+    torch.testing.assert_close(merged.weight[:2, 0], merged.weight.new_tensor([5, 0]))
 
     qkv = QKVParallelLinear(
         hidden_size=1,
@@ -79,15 +73,11 @@ def test_explicit_loaded_sizes_zero_rank_local_destination_tails(
         )
     torch.testing.assert_close(qkv.weight[:, 0], qkv.weight.new_zeros(4))
 
-    row = RowParallelLinear(
-        6, 1, bias=False, loaded_input_size=5
-    )
+    row = RowParallelLinear(6, 1, bias=False, loaded_input_size=5)
     row.weight.weight_loader(
         row.weight, torch.arange(1, 6, dtype=row.weight.dtype).unsqueeze(0)
     )
-    torch.testing.assert_close(
-        row.weight[0], row.weight.new_tensor([5, 0])
-    )
+    torch.testing.assert_close(row.weight[0], row.weight.new_tensor([5, 0]))
 
 
 def test_loaded_sizes_reject_invalid_physical_layout(
@@ -97,9 +87,7 @@ def test_loaded_sizes_reject_invalid_physical_layout(
     with pytest.raises(ValueError, match="exceeds physical size"):
         ColumnParallelLinear(1, 6, bias=False, loaded_output_size=7)
     with pytest.raises(ValueError, match="same length"):
-        MergedColumnParallelLinear(
-            1, [6, 6], bias=False, loaded_output_sizes=[5]
-        )
+        MergedColumnParallelLinear(1, [6, 6], bias=False, loaded_output_sizes=[5])
     with pytest.raises(ValueError, match="exceeds physical size"):
         RowParallelLinear(6, 1, bias=False, loaded_input_size=7)
 
@@ -190,9 +178,7 @@ def test_padded_loader_rejects_unaligned_quantized_boundaries_before_write(
             torch.ones((1, 1), dtype=quantized_param.dtype),
         )
 
-    torch.testing.assert_close(
-        quantized_param, quantized_param.new_full((1, 1), 23)
-    )
+    torch.testing.assert_close(quantized_param, quantized_param.new_full((1, 1), 23))
 
 
 def test_padded_v2_loader_preserves_unsharded_per_tensor_scale(
@@ -226,12 +212,8 @@ def test_padded_nvfp4_row_loader_converts_logical_to_storage_width(
     dtype: torch.dtype,
 ) -> None:
     for module in (linear, parameter):
-        monkeypatch.setattr(
-            module, "get_tensor_model_parallel_world_size", lambda: 3
-        )
-        monkeypatch.setattr(
-            module, "get_tensor_model_parallel_rank", lambda: 1
-        )
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
     row = RowParallelLinear(
         physical_size,
         1,
@@ -261,7 +243,6 @@ def test_padded_nvfp4_row_loader_converts_logical_to_storage_width(
     )
 
 
-
 @pytest.mark.parametrize(
     "method_type",
     (
@@ -274,12 +255,8 @@ def test_modelopt_nvfp4_row_parameters_declare_storage_widths(
     method_type: type,
 ) -> None:
     for module in (linear, parameter):
-        monkeypatch.setattr(
-            module, "get_tensor_model_parallel_world_size", lambda: 3
-        )
-        monkeypatch.setattr(
-            module, "get_tensor_model_parallel_rank", lambda: 1
-        )
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
     monkeypatch.setattr(
         modelopt,
         "init_nvfp4_linear_kernel",
@@ -330,12 +307,8 @@ def test_modelopt_mxfp8_scale_loads_padded_tp3_storage_width(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     for module in (linear, parameter):
-        monkeypatch.setattr(
-            module, "get_tensor_model_parallel_world_size", lambda: 3
-        )
-        monkeypatch.setattr(
-            module, "get_tensor_model_parallel_rank", lambda: 1
-        )
+        monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
+        monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
     monkeypatch.setattr(modelopt, "init_mxfp8_linear_kernel", lambda: object())
     config = modelopt.ModelOptMxFp8Config(
         is_checkpoint_mxfp8_serialized=True,
@@ -357,7 +330,9 @@ def test_modelopt_mxfp8_scale_loads_padded_tp3_storage_width(
     loaded = torch.arange(1, 129, dtype=holder.weight_scale.dtype).unsqueeze(0)
     row = RowParallelLinear(12288, 1, bias=False, loaded_input_size=4096)
     row.weight_loader_v2(holder.weight_scale, loaded)
-    torch.testing.assert_close(holder.weight_scale, torch.zeros_like(holder.weight_scale))
+    torch.testing.assert_close(
+        holder.weight_scale, torch.zeros_like(holder.weight_scale)
+    )
 
 
 @pytest.mark.parametrize("tp3", [False, True])
@@ -380,9 +355,7 @@ def __init__(self, *args, **kwargs) -> None:
     monkeypatch.setattr(glm_attention, "RowParallelLinear", FakeLinear)
     monkeypatch.setattr(glm_attention, "DeepSeekV2FusedQkvAProjLinear", FakeLinear)
     monkeypatch.setattr(glm_attention, "RMSNorm", FakeModule)
-    monkeypatch.setattr(
-        glm_attention, "MultiHeadLatentAttentionWrapper", FakeModule
-    )
+    monkeypatch.setattr(glm_attention, "MultiHeadLatentAttentionWrapper", FakeModule)
     monkeypatch.setattr(
         glm_attention,
         "get_tensor_model_parallel_world_size",
@@ -442,7 +415,6 @@ class FakeActivation(torch.nn.Module):
         def __init__(self, *args, **kwargs) -> None:
             super().__init__()
 
-
     monkeypatch.setattr(glm_model, "MergedColumnParallelLinear", FakeMerged)
     monkeypatch.setattr(glm_model, "RowParallelLinear", FakeRow)
     monkeypatch.setattr(glm_model, "SiluAndMul", FakeActivation)
@@ -495,9 +467,7 @@ def __init__(self, *args, **kwargs) -> None:
     monkeypatch.setattr(glm_mtp, "SharedHead", FakeSharedHead)
     monkeypatch.setattr(glm_mtp, "ParallelLMHead", FakeParallelLMHead)
     monkeypatch.setattr(glm_mtp, "Glm5NextDecoderLayer", FakeDecoder)
-    monkeypatch.setattr(
-        glm_mtp, "current_platform", SimpleNamespace(device_type="cpu")
-    )
+    monkeypatch.setattr(glm_mtp, "current_platform", SimpleNamespace(device_type="cpu"))
 
     config = SimpleNamespace(
         hidden_size=4096,
@@ -518,9 +488,7 @@ def __init__(self, *args, **kwargs) -> None:
         scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
     )
 
-    layer = glm_mtp.Glm5NextMultiTokenPredictorLayer(
-        vllm_config, "model.layers.45"
-    )
+    layer = glm_mtp.Glm5NextMultiTokenPredictorLayer(vllm_config, "model.layers.45")
 
     if tp3:
         assert isinstance(layer.eh_proj, FakeColumn)
diff --git a/tests/models/test_glm5next_vision_tp3.py b/tests/models/test_glm5next_vision_tp3.py
index 392a516ed829..8829d2ea81f5 100644
--- a/tests/models/test_glm5next_vision_tp3.py
+++ b/tests/models/test_glm5next_vision_tp3.py
@@ -159,7 +159,9 @@ def __init__(self, *args, **kwargs) -> None:
     torch.testing.assert_close(merger.down_proj.weight[:, 3412:], torch.zeros(4, 2))
 
 
-def _record_vision_geometry(monkeypatch, vision_config, *, data_parallel: bool, tp: int):
+def _record_vision_geometry(
+    monkeypatch, vision_config, *, data_parallel: bool, tp: int
+):
     recorded = SimpleNamespace(block=None, merger=None, rope=None)
 
     class FakeModule(torch.nn.Module):
diff --git a/tests/v1/attention/test_dflash_attn.py b/tests/v1/attention/test_dflash_attn.py
index a4633e3fecbb..b79724ba5308 100644
--- a/tests/v1/attention/test_dflash_attn.py
+++ b/tests/v1/attention/test_dflash_attn.py
@@ -2,8 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Split-KV draft attention against FlashAttention 2 at the DFlash draft shape."""
 
-import os
-
 import pytest
 import torch
 
@@ -140,4 +138,3 @@ def test_workspace_rejects_oversized_batch():
     op = dfa.DFlashDecodeAttention(device, HKV, max_batch=1, window=WINDOW)
     with pytest.raises(ValueError):
         op(q, k, v, block_table, seqused, cu, SCALE, torch.empty_like(q))
-    assert os.getenv("VLLM_GLM53_DFLASH_ATTN", "0") in ("0", "1")

From a3189ee5407699a8777ec386fba649c861902f8b Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Fri, 4 Sep 2026 14:29:47 +0200
Subject: [PATCH 24/25] fix(glm53): port amended TP3 readiness from #547 squash

- dflash proposer: TP3-gated draft parallel copy imported from glm53_tp3
- VllmConfig: materialize TP3 target geometry before generic validation
- cuda_communicator: RoCE fallback only when B12X all-reduce is absent/disabled
- glm53_tp3: EP requirement, full draft shape sweep, preserve draft_vocab_size
- speculative: draft TP3 check, DCP field propagation, world_size recompute
---
 vllm/config/speculative.py                    | 13 ++++++++++
 vllm/config/vllm.py                           |  8 ++++++
 .../device_communicators/cuda_communicator.py |  6 ++++-
 vllm/transformers_utils/configs/glm53_tp3.py  | 25 ++++++++++++++++---
 vllm/v1/spec_decode/dflash.py                 | 18 ++++++++-----
 5 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
index b8cb6f0e36eb..8965f3e502a5 100644
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
@@ -1731,6 +1731,7 @@ def _apply_glm53_tp3_draft_geometry(self) -> None:
             or self.draft_model_config is None
             or self.draft_parallel_config is None
             or self.target_parallel_config.tensor_parallel_size != 3
+            or self.draft_parallel_config.tensor_parallel_size != 3
             or not is_glm53_config(self.target_model_config)
         ):
             return
@@ -1741,6 +1742,11 @@ def _apply_glm53_tp3_draft_geometry(self) -> None:
         draft = self.draft_parallel_config
         for name in (
             "prefill_context_parallel_size",
+            "decode_context_parallel_size",
+            "dcp_kv_cache_interleave_size",
+            "dcp_comm_backend",
+            "dcp_q_replicate",
+            "cp_kv_cache_interleave_size",
             "data_parallel_size",
             "data_parallel_size_local",
             "data_parallel_rank",
@@ -1757,6 +1763,13 @@ def _apply_glm53_tp3_draft_geometry(self) -> None:
         draft.enable_expert_parallel = (
             target.enable_expert_parallel if self.method == "mtp" else False
         )
+        draft.world_size = (
+            draft.pipeline_parallel_size
+            * draft.tensor_parallel_size
+            * draft.prefill_context_parallel_size
+        )
+        if draft.distributed_executor_backend == "external_launcher":
+            draft.world_size *= draft.data_parallel_size
 
         apply_glm53_tp3_draft_geometry(
             self.target_model_config,
diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py
index f0cc6a625b07..56875556a169 100644
--- a/vllm/config/vllm.py
+++ b/vllm/config/vllm.py
@@ -1138,6 +1138,14 @@ def __post_init__(self):
         if self.performance_mode != "balanced":
             logger.info_once("Performance mode set to '%s'.", self.performance_mode)
 
+        # GLM-5.3's TP3 storage geometry must be materialized before generic
+        # parallel-shape validation. Derive it from ParallelConfig, never from
+        # process environment or checkpoint paths.
+        from vllm.transformers_utils.configs.glm53_tp3 import (
+            apply_glm53_tp3_target_geometry,
+        )
+
+        apply_glm53_tp3_target_geometry(self.model_config, self.parallel_config)
         self.try_verify_and_update_config()
 
         # Models may have supplied their own DCP defaults above; anything still
diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py
index 5161104e67b4..3a452e9ab8ac 100644
--- a/vllm/distributed/device_communicators/cuda_communicator.py
+++ b/vllm/distributed/device_communicators/cuda_communicator.py
@@ -122,7 +122,11 @@ def __init__(
                 device_group=self.device_group,
                 device=self.device,
             )
-        elif self.use_roce_allreduce and self.world_size > 1:
+        if (
+            self.use_roce_allreduce
+            and self.world_size > 1
+            and (self.b12x_ar_comm is None or self.b12x_ar_comm.disabled)
+        ):
             # RoCEnante: multi-node DGX Spark one-shot RDMA collectives
             # from b12x.comm.roce.
             from .b12x_roce_all_reduce import B12xRoceAllReduce
diff --git a/vllm/transformers_utils/configs/glm53_tp3.py b/vllm/transformers_utils/configs/glm53_tp3.py
index 9c163badb188..8b524e953ebf 100644
--- a/vllm/transformers_utils/configs/glm53_tp3.py
+++ b/vllm/transformers_utils/configs/glm53_tp3.py
@@ -86,6 +86,12 @@ def apply_glm53_tp3_target_geometry(
     ):
         return False
 
+    if not parallel_config.enable_expert_parallel:
+        raise ValueError(
+            "GLM-5.3 physical TP3 requires expert parallelism because the "
+            "released routed-expert width is not divisible by 3."
+        )
+
     text_config = model_config.hf_text_config
     target_shapes = (
         ("num_attention_heads", 64),
@@ -179,9 +185,19 @@ def apply_glm53_tp3_draft_geometry(
         return False
 
     text_config = draft_model_config.hf_text_config
-    _require_shape(text_config, "num_attention_heads", 32)
-    _require_shape(text_config, "num_key_value_heads", 8)
-    _require_shape(text_config, "vocab_size", 154880)
+    draft_shapes = (
+        ("num_attention_heads", 32),
+        ("num_key_value_heads", 8),
+        ("vocab_size", 154880),
+    )
+    for name, expected in draft_shapes:
+        _require_shape(text_config, name, expected)
+    for config in _iter_hf_configs(draft_model_config):
+        if config is text_config:
+            continue
+        for name, expected in draft_shapes:
+            if hasattr(config, name):
+                _require_shape(config, name, expected)
     for config in _iter_hf_configs(draft_model_config):
         if hasattr(config, "num_attention_heads"):
             config.original_num_attention_heads = 32
@@ -191,7 +207,8 @@ def apply_glm53_tp3_draft_geometry(
             config.num_key_value_heads = 9
         if hasattr(config, "vocab_size"):
             config.original_vocab_size = 154880
-            config.draft_vocab_size = 154880
+            if getattr(config, "draft_vocab_size", None) is None:
+                config.draft_vocab_size = 154880
         config.glm53_tp3_padding = True
         config.glm53_tp3_vocab_padding_size = 192
         config.glm53_tp3_vocab_storage_size = 154944
diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py
index 7ab450c8ca2d..a75c8cdd6121 100644
--- a/vllm/v1/spec_decode/dflash.py
+++ b/vllm/v1/spec_decode/dflash.py
@@ -10,6 +10,7 @@
 
 from vllm.config import VllmConfig
 from vllm.forward_context import set_forward_context
+from vllm.transformers_utils.configs.glm53_tp3 import is_glm53_config
 from vllm.logger import init_logger
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
@@ -99,12 +100,17 @@ def load_model(self, target_model: torch.nn.Module) -> None:
     @override
     def _create_draft_vllm_config(self) -> VllmConfig:
         base = super()._create_draft_vllm_config()
-        draft_parallel_config = copy(
-            self.speculative_config.draft_parallel_config
-        )
-        assert draft_parallel_config is not None
-        draft_parallel_config.rank = self.vllm_config.parallel_config.rank
-        base = replace(base, parallel_config=draft_parallel_config)
+        target_parallel_config = self.vllm_config.parallel_config
+        if (
+            target_parallel_config.tensor_parallel_size == 3
+            and is_glm53_config(self.speculative_config.target_model_config)
+        ):
+            draft_parallel_config = copy(
+                self.speculative_config.draft_parallel_config
+            )
+            assert draft_parallel_config is not None
+            draft_parallel_config.rank = target_parallel_config.rank
+            base = replace(base, parallel_config=draft_parallel_config)
         # The draft model is text-only — clear the target's multimodal
         # flag so flash_attn is not rejected for mm_prefix support.
         arch = base.model_config.model_arch_config

From e96b18dbb8c19230591e79e0ed056b12947b2ea1 Mon Sep 17 00:00:00 2001
From: Gerben Meijer <infernix@gmail.com>
Date: Fri, 4 Sep 2026 15:07:19 +0200
Subject: [PATCH 25/25] test(glm53): adopt #547 amended TP3 fixtures

---
 tests/config/test_glm53_tp3_geometry.py | 97 +++++++++++++++++++++++++
 tests/models/test_glm53_tp3_dflash.py   | 22 +++++-
 2 files changed, 117 insertions(+), 2 deletions(-)

diff --git a/tests/config/test_glm53_tp3_geometry.py b/tests/config/test_glm53_tp3_geometry.py
index d53dac532577..1fd05e9544c1 100644
--- a/tests/config/test_glm53_tp3_geometry.py
+++ b/tests/config/test_glm53_tp3_geometry.py
@@ -98,6 +98,21 @@ def _snapshot(model_config: object) -> dict[str, Any]:
     return deepcopy(vars(model_config))
 
 
+def test_glm53_tp3_target_geometry_requires_expert_parallel(
+    glm53_model_config: FakeGlm53ModelConfig,
+) -> None:
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=3,
+        enable_expert_parallel=False,
+    )
+
+    with pytest.raises(ValueError, match="requires expert parallelism"):
+        apply_glm53_tp3_target_geometry(
+            cast(Any, glm53_model_config),
+            parallel_config,
+        )
+
+
 def test_glm53_tp3_target_geometry_uses_parallel_config_and_preserves_logical_axes(
     glm53_model_config: FakeGlm53ModelConfig,
     tp3_ep_parallel_config: ParallelConfig,
@@ -330,6 +345,36 @@ def test_glm53_tp3_mtp_draft_preserves_expert_parallel_topology() -> None:
     assert draft_text_config.glm53_tp3_shared_expert_intermediate_size == 2112
 
 
+def test_glm53_tp3_draft_tp1_is_an_exact_noop() -> None:
+    target_model_config = FakeGlm53ModelConfig()
+    draft_model_config = FakeDFlashModelConfig()
+    target_parallel_config = ParallelConfig(
+        tensor_parallel_size=3,
+        data_parallel_size=2,
+        data_parallel_size_local=2,
+        enable_expert_parallel=True,
+    )
+    draft_parallel_config = ParallelConfig(tensor_parallel_size=1)
+    before = _snapshot(draft_model_config)
+    speculative_config = SimpleNamespace(
+        method="dflash",
+        target_model_config=target_model_config,
+        target_parallel_config=target_parallel_config,
+        draft_model_config=draft_model_config,
+        draft_parallel_config=draft_parallel_config,
+    )
+
+    SpeculativeConfig._apply_glm53_tp3_draft_geometry(
+        cast(Any, speculative_config)
+    )
+
+    assert draft_parallel_config.tensor_parallel_size == 1
+    assert draft_parallel_config.data_parallel_size == 1
+    assert draft_parallel_config.data_parallel_size_local == 1
+    assert not draft_parallel_config.enable_expert_parallel
+    assert _snapshot(draft_model_config) == before
+
+
 def test_glm53_tp3_dflash_drops_ep_and_couples_heads_with_vocab_storage() -> None:
     target_model_config = FakeGlm53ModelConfig()
     draft_model_config = FakeDFlashModelConfig()
@@ -377,6 +422,58 @@ def test_glm53_tp3_dflash_drops_ep_and_couples_heads_with_vocab_storage() -> Non
     assert draft_model_config.model_arch_config.vocab_size == 154880
 
 
+def test_glm53_tp3_dflash_preserves_reduced_draft_vocabulary() -> None:
+    target_model_config = FakeGlm53ModelConfig()
+    draft_model_config = FakeDFlashModelConfig()
+    draft_model_config.hf_text_config.draft_vocab_size = 32000
+
+    assert apply_glm53_tp3_draft_geometry(
+        cast(Any, target_model_config),
+        ParallelConfig(tensor_parallel_size=3),
+        cast(Any, draft_model_config),
+        ParallelConfig(tensor_parallel_size=3),
+    )
+
+    assert draft_model_config.hf_text_config.draft_vocab_size == 32000
+
+
+@pytest.mark.parametrize(
+    ("attribute", "invalid_value"),
+    [
+        ("num_attention_heads", 31),
+        ("num_key_value_heads", 7),
+        ("vocab_size", 154879),
+    ],
+)
+def test_glm53_tp3_invalid_nested_dflash_geometry_is_transactional(
+    attribute: str,
+    invalid_value: int,
+) -> None:
+    target_model_config = FakeGlm53ModelConfig()
+    draft_model_config = FakeDFlashModelConfig()
+    text_config = draft_model_config.hf_text_config
+    draft_model_config.hf_config = SimpleNamespace(
+        model_type="qwen3",
+        architectures=["DFlash2DraftModel"],
+        text_config=text_config,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        vocab_size=154880,
+    )
+    setattr(draft_model_config.hf_config, attribute, invalid_value)
+    before = _snapshot(draft_model_config)
+
+    with pytest.raises(ValueError, match=rf"expected {attribute}="):
+        apply_glm53_tp3_draft_geometry(
+            cast(Any, target_model_config),
+            ParallelConfig(tensor_parallel_size=3),
+            cast(Any, draft_model_config),
+            ParallelConfig(tensor_parallel_size=3),
+        )
+
+    assert _snapshot(draft_model_config) == before
+
+
 @pytest.mark.parametrize(
     ("attribute", "invalid_value"),
     [
diff --git a/tests/models/test_glm53_tp3_dflash.py b/tests/models/test_glm53_tp3_dflash.py
index 362e480ca87b..b391819793c8 100644
--- a/tests/models/test_glm53_tp3_dflash.py
+++ b/tests/models/test_glm53_tp3_dflash.py
@@ -277,7 +277,12 @@ def test_dense_dflash_tp3_drops_target_expert_parallel(monkeypatch) -> None:
         lambda *args: applied.append(args),
     )
     placement = {
-        "prefill_context_parallel_size": 1,
+        "prefill_context_parallel_size": 2,
+        "decode_context_parallel_size": 2,
+        "dcp_kv_cache_interleave_size": 4,
+        "dcp_comm_backend": "a2a",
+        "dcp_q_replicate": True,
+        "cp_kv_cache_interleave_size": 4,
         "data_parallel_size": 2,
         "data_parallel_size_local": 2,
         "data_parallel_rank": 1,
@@ -302,14 +307,23 @@ def test_dense_dflash_tp3_drops_target_expert_parallel(monkeypatch) -> None:
     )
     spec = object.__new__(SpeculativeConfig)
     object.__setattr__(spec, "method", "dflash")
-    object.__setattr__(spec, "target_model_config", object())
     object.__setattr__(spec, "target_parallel_config", target_parallel)
+    target_hf_config = SimpleNamespace(architectures=["Glm5NextForCausalLM"])
+    object.__setattr__(
+        spec,
+        "target_model_config",
+        SimpleNamespace(
+            hf_config=target_hf_config,
+            hf_text_config=target_hf_config,
+        ),
+    )
     object.__setattr__(spec, "draft_model_config", object())
     object.__setattr__(spec, "draft_parallel_config", draft_parallel)
 
     spec._apply_glm53_tp3_draft_geometry()
 
     assert draft_parallel.enable_expert_parallel is False
+    assert draft_parallel.world_size == 6
     for name, value in placement.items():
         assert getattr(draft_parallel, name) == value
     assert len(applied) == 1
@@ -347,6 +361,10 @@ class DraftVllmConfig:
     assert target_parallel.tensor_parallel_size == 3
     assert target_parallel.enable_expert_parallel
 
+    target_parallel.tensor_parallel_size = 4
+    tp4_vllm_config = proposer._create_draft_vllm_config()
+    assert tp4_vllm_config.parallel_config is target_parallel
+
 
 def test_engine_dp_identity_reaches_speculative_draft() -> None:
     from vllm.v1.engine.core import EngineCoreProc