fix ep_enabled

Dockerfile.ref

@@ -18,7 +18,4 @@ COPY ./vllm/core/block_manager_v2.py /workspace/qwen3_6_scripts/vendor_overrides
 COPY ./vllm/sampling_params.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py
 COPY ./vllm/model_executor/sampling_metadata.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py
 COPY ./vllm/model_executor/layers/sampler.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
-COPY ./vllm/spec_decode/spec_decode_worker.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/spec_decode/spec_decode_worker.py
-COPY ./vllm/spec_decode/multi_step_worker.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/spec_decode/multi_step_worker.py
-COPY ./vllm/spec_decode/metrics.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/spec_decode/metrics.py
-RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh
+RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh

qwen3_6_scripts/ep_fused_moe_patch.py

@@ -39,43 +39,30 @@ from vllm.logger import init_logger
 logger = init_logger(__name__)
 
 
-def _get_ep_config():
-    """Read EP configuration from ParallelConfig.
+def _is_ep_requested():
+    """Check if EP is requested via environment variable.
 
-    Returns (ep_enabled, ep_size, ep_rank, ep_group).
-    Safe to call during model init — returns disabled if ParallelConfig
-    doesn't have EP fields (backward compat with unpatched vllm).
+    This is safe to call at any time (no dist dependency).
     """
-    try:
-        from vllm.config import ParallelConfig as _PC
-        # ParallelConfig is a per-worker singleton-ish object.
-        # We check if the fields exist (set by our config.py patch).
-        # During weight loading, the global TP group is already initialized.
-        from vllm.distributed import (
-            get_tensor_model_parallel_world_size,
-            get_tensor_model_parallel_rank,
-        )
-        # EP config is on the parallel_config instance that was passed
-        # through the worker. We can't easily get it here without plumbing,
-        # so we use env vars (set in arg_utils.py) as the source of truth.
-        import os
-        ep_enabled = bool(int(os.environ.get("VLLM_ENABLE_EXPERT_PARALLEL", "0")))
-        if not ep_enabled:
-            return False, 1, 0, None
+    import os
+    return bool(int(os.environ.get("VLLM_ENABLE_EXPERT_PARALLEL", "0")))
 
-        # When EP is enabled, ep_size = world_size (all ranks participate).
-        # ep_rank = global rank (since EP group = all ranks).
-        if dist.is_initialized():
-            ep_size = dist.get_world_size()
-            ep_rank = dist.get_rank()
-        else:
-            ep_size = 1
-            ep_rank = 0
-            ep_enabled = False
 
-        return ep_enabled, ep_size, ep_rank, None
-    except Exception:
-        return False, 1, 0, None
+def _get_ep_runtime_config():
+    """Get EP size and rank at runtime when dist is initialized.
+
+    Must only be called from within worker processes (during __init__,
+    weight_loader, or forward) where torch.distributed is guaranteed
+    to be initialized.
+
+    Returns (ep_size, ep_rank).
+    """
+    if dist.is_initialized():
+        return dist.get_world_size(), dist.get_rank()
+    # Fallback — should not happen during normal model loading
+    logger.warning("[PR #2269] dist not initialized when querying EP config, "
+                   "falling back to ep_size=1")
+    return 1, 0
 
 
 def patch_fused_moe_for_ep():
@@ -83,15 +70,20 @@ def patch_fused_moe_for_ep():
 
     Call this during model loading (e.g. in qwen3_5.py module init or
     patch_ops.sh) BEFORE any FusedMoE layers are constructed.
+
+    Only checks the env var at patch time. Actual ep_size/ep_rank are
+    queried at runtime (inside __init__, weight_loader, forward) when
+    torch.distributed is guaranteed to be initialized.
     """
     from vllm.model_executor.layers.fused_moe.layer import FusedMoE
 
-    ep_enabled, ep_size, ep_rank, _ = _get_ep_config()
-    if not ep_enabled or ep_size <= 1:
-        logger.info("[PR #2269] EP not enabled or ep_size<=1, skipping FusedMoE EP patch")
+    if not _is_ep_requested():
+        logger.info("[PR #2269] EP not enabled (VLLM_ENABLE_EXPERT_PARALLEL!=1), "
+                    "skipping FusedMoE EP patch")
         return
 
-    num_experts_per_rank = None  # set per-layer in __init__
+    logger.info("[PR #2269] EP requested via env var, installing FusedMoE monkey-patches "
+                "(ep_size/ep_rank will be resolved at runtime)")
 
     _orig_init = FusedMoE.__init__
 
@@ -102,6 +94,9 @@ def patch_fused_moe_for_ep():
                  prefix="", custom_routing_function=None):
         """EP-aware __init__: allocate only local experts, full intermediate."""
 
+        # Query ep_size/ep_rank NOW — dist is initialized in worker processes
+        ep_size, ep_rank = _get_ep_runtime_config()
+
         # Store global expert info before modifying
         self._ep_enabled = True
         self._ep_size = ep_size
@@ -154,7 +149,19 @@ def patch_fused_moe_for_ep():
 
     def _ep_weight_loader(self, param, loaded_weight, weight_name,
                           shard_id, expert_id):
-        """EP-aware weight_loader: skip experts not owned by this rank."""
+        """EP-aware weight_loader: skip experts not owned by this rank.
+
+        Critical fix: the original weight_loader calls
+        get_tensor_model_parallel_rank() internally to compute tp_rank,
+        then uses it to narrow() the loaded weight for TP sharding.
+        Under EP we pass tp_size=1 (no TP split on MoE), so weights are
+        full-sized, but tp_rank still returns the real rank (0-3).
+        rank >= 1 causes narrow(dim, 512*rank, 512) to exceed dim size 512.
+
+        Solution: temporarily monkey-patch get_tensor_model_parallel_rank
+        to return 0 during EP weight loading, so narrow() always starts
+        at offset 0 and loads the full (unsplit) weight.
+        """
         if not getattr(self, '_ep_enabled', False):
             return _orig_weight_loader(self, param, loaded_weight,
                                        weight_name, shard_id, expert_id)
@@ -169,10 +176,16 @@ def patch_fused_moe_for_ep():
         # Remap global expert_id to local index
         local_expert_id = expert_id - start
 
-        # Call original weight_loader with local expert_id.
-        # The param tensor is sized for local experts only.
-        _orig_weight_loader(self, param, loaded_weight, weight_name,
-                            shard_id, local_expert_id)
+        # Monkey-patch tp rank to 0 during load so narrow() doesn't OOB.
+        # The weight is already full-sized (tp_size=1), no split needed.
+        import vllm.model_executor.layers.fused_moe.layer as _fused_moe_mod
+        _real_get_tp_rank = _fused_moe_mod.get_tensor_model_parallel_rank
+        _fused_moe_mod.get_tensor_model_parallel_rank = lambda: 0
+        try:
+            _orig_weight_loader(self, param, loaded_weight, weight_name,
+                                shard_id, local_expert_id)
+        finally:
+            _fused_moe_mod.get_tensor_model_parallel_rank = _real_get_tp_rank
 
     _orig_forward = FusedMoE.forward
 
@@ -330,6 +343,5 @@ def patch_fused_moe_for_ep():
     FusedMoE.forward = _ep_forward
 
     logger.info(
-        "[PR #2269] FusedMoE patched for EP: ep_size=%d, ep_rank=%d. "
-        "Each card holds %d/%d-th of experts with full intermediate_size.",
-        ep_size, ep_rank, 1, ep_size)
+        "[PR #2269] FusedMoE patched for EP. "
+        "ep_size/ep_rank will be resolved per-layer at runtime.")

qwen3_6_scripts/patch_ops.sh

@@ -3,6 +3,8 @@ echo "[build] trigger 202609011237"
 echo "[build] trigger 202609011636"
 echo "[build] trigger 202609011655"
 echo "[build] trigger 202609011710"
+echo "[build] trigger 202609031035"
+echo "[build] trigger 202609031049"
 #!/usr/bin/env bash
 # BI-V100 patch script for Qwen3.6-35B-A3B (Qwen3_5 MoE architecture)
 #
@@ -162,17 +164,6 @@ install_patch_file \
     "${VLLM_OVERRIDE_ROOT}/worker/worker.py" \
     "${VLLM_ROOT}/worker/worker.py"
 
-build_stage "installing PR #2269 DP/EP-aware speculative decoding overrides"
-install_patch_file \
-    "${VLLM_OVERRIDE_ROOT}/spec_decode/spec_decode_worker.py" \
-    "${VLLM_ROOT}/spec_decode/spec_decode_worker.py"
-install_patch_file \
-    "${VLLM_OVERRIDE_ROOT}/spec_decode/multi_step_worker.py" \
-    "${VLLM_ROOT}/spec_decode/multi_step_worker.py"
-install_patch_file \
-    "${VLLM_OVERRIDE_ROOT}/spec_decode/metrics.py" \
-    "${VLLM_ROOT}/spec_decode/metrics.py"
-
 build_stage "installing hash-pinned CoreX 3.2.3 extensions"
 bash ./install_prebuilt_corex.sh "${VLLM_ROOT}"
 

qwen3_6_scripts/vendor_overrides/vllm/engine/arg_utils.py

@@ -840,7 +840,14 @@ class EngineArgs:
         # Get the list of attributes of this dataclass.
         attrs = [attr.name for attr in dataclasses.fields(cls)]
         # Set the attributes from the parsed arguments.
-        engine_args = cls(**{attr: getattr(args, attr) for attr in attrs})
+        # Fields that exist in the dataclass but have no corresponding CLI
+        # argument (e.g. data_parallel_size, enable_expert_parallel) will
+        # fall back to their dataclass defaults (typically read from env vars).
+        engine_args = cls(**{
+            attr: getattr(args, attr)
+            for attr in attrs
+            if hasattr(args, attr)
+        })
         return engine_args
 
     def create_model_config(self) -> ModelConfig: