274 lines
11 KiB
Python
274 lines
11 KiB
Python
"""EP-aware FusedMoE override for vllm 0.6.3.
|
||
|
||
Ported from tpu-inference upstream (xllm):
|
||
- 04077875 PR #2577: Enable correct DP attention for hybrid models
|
||
- df7f5b35 PR #2679: Optimize attn DP + MoE EP ReduceScatter
|
||
- 57987c2 PR #3435: Reduce shared-expert output over in-group TP under attn DP
|
||
|
||
Architecture:
|
||
4 GPUs, --tensor-parallel-size 4, VLLM_ENABLE_EXPERT_PARALLEL=1
|
||
→ TP group = EP group = WORLD = {rank0, rank1, rank2, rank3}
|
||
→ Routed experts: each rank holds 64/256 experts, tp_size=1 (full intermediate)
|
||
→ Shared expert: TP=4 sharded (uses default tp_size=world_size)
|
||
|
||
Design:
|
||
1. __init__: allocate only local experts (64 per rank), tp_size=1
|
||
2. weight_loader: skip non-local experts, remap expert_id to local
|
||
3. forward (in qwen3_5.py): EP mask + per-expert compute + reduction
|
||
|
||
Reduction strategy (from upstream PR #2679 / #3435):
|
||
- routed_out: partial across EP ranks (each rank computed local experts only)
|
||
- shared_out: partial across TP ranks (RowParallelLinear reduce_results=False)
|
||
- Since EP group == TP group == WORLD: single all-reduce on (routed + shared)
|
||
- When DP > 1: can optimize to reduce-scatter (see dp_forward_moe_wrapper)
|
||
|
||
EP expert dispatch (from upstream fused_moe_gmm.py):
|
||
- topk routing produces global expert ids on ALL ranks (identical)
|
||
- Each rank masks non-local experts (weight=0), remaps to local [0, E_local)
|
||
- Non-local (token,expert) pairs distributed via mod to avoid expert 0 overload
|
||
- Per-expert loop skips experts with count=0 → no wasted GEMM for empty experts
|
||
- Remaining waste: tokens routed to non-local experts still enter local experts
|
||
with weight=0. Upstream uses ragged_gather to skip these entirely.
|
||
TODO: implement token-level skip for further speedup.
|
||
"""
|
||
|
||
import os
|
||
import torch
|
||
import torch.distributed as dist
|
||
|
||
from vllm.logger import init_logger
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Environment / runtime helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _is_ep_requested():
|
||
return bool(int(os.environ.get("VLLM_ENABLE_EXPERT_PARALLEL", "0")))
|
||
|
||
|
||
def _get_ep_runtime_config():
|
||
"""Return (ep_size, ep_rank). EP group = WORLD group."""
|
||
if dist.is_initialized():
|
||
return dist.get_world_size(), dist.get_rank()
|
||
return 1, 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# EP-aware expert mask + remap (called from qwen3_5.py _pure_pytorch_experts)
|
||
# ---------------------------------------------------------------------------
|
||
# Ported from tpu-inference fused_moe_gmm.py:
|
||
# _process_tokens_locally() → valid_rows_mask + ragged_gather
|
||
# Our CUDA version: mask weights to 0, distribute non-local ids via mod.
|
||
|
||
def ep_mask_and_remap(topk_ids, topk_weights, start_expert_id,
|
||
num_experts_per_rank, ep_rank):
|
||
"""Mask non-local experts and remap global ids to local.
|
||
|
||
Args:
|
||
topk_ids: [T, top_k] global expert ids
|
||
topk_weights: [T, top_k] routing weights
|
||
start_expert_id: first global expert id on this rank
|
||
num_experts_per_rank: number of local experts (E_local)
|
||
ep_rank: this rank's EP index
|
||
|
||
Returns:
|
||
topk_ids: [T, top_k] local expert ids in [0, E_local)
|
||
topk_weights: [T, top_k] with non-local entries zeroed
|
||
local_mask: [T, top_k] bool, True for locally-routed pairs
|
||
"""
|
||
E_local = num_experts_per_rank
|
||
end_expert_id = start_expert_id + E_local
|
||
|
||
# Identify local (token, expert) pairs
|
||
local_mask = (topk_ids >= start_expert_id) & (topk_ids < end_expert_id)
|
||
|
||
# Zero non-local weights (these contribute nothing to output)
|
||
topk_weights = topk_weights * local_mask.to(topk_weights.dtype)
|
||
|
||
# Remap global → local ids
|
||
# Local experts: id - start → [0, E_local)
|
||
# Non-local: distribute via mod to avoid inflating expert 0 count
|
||
# (upstream uses ragged_gather to skip non-local tokens entirely;
|
||
# this mod trick achieves balanced distribution with weight=0)
|
||
local_ids = topk_ids - start_expert_id
|
||
non_local = ~local_mask
|
||
if non_local.any():
|
||
local_ids[non_local] = topk_ids[non_local] % E_local
|
||
|
||
return local_ids, topk_weights, local_mask
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# EP-aware reduction (called from qwen3_5.py forward())
|
||
# ---------------------------------------------------------------------------
|
||
# Ported from tpu-inference:
|
||
# custom_ops/fused_moe.py VllmMoERunner._maybe_reduce_shared_expert_output
|
||
# custom_ops/fused_moe.py VllmMoERunner._maybe_reduce_final_output
|
||
# fused_moe_gmm.py scatter_results / defer_all_reduce logic
|
||
|
||
def ep_reduce_output(routed_out, shared_out):
|
||
"""Reduce EP + TP partial outputs into final MoE output.
|
||
|
||
In our topology (TP group == EP group == WORLD):
|
||
routed_out: [T, H] partial — each rank only computed local experts
|
||
shared_out: [T, H] partial — RowParallelLinear TP shard (reduce_results=False)
|
||
|
||
IMPORTANT: use tensor_model_parallel_all_reduce (goes through ixformer
|
||
optimized path on BI-V100) instead of raw dist.all_reduce (raw NCCL,
|
||
40x slower on this hardware). This is valid because TP group == WORLD.
|
||
|
||
Returns: [T, H] fully-reduced output.
|
||
"""
|
||
from vllm.distributed import tensor_model_parallel_all_reduce
|
||
out = routed_out + shared_out
|
||
out = tensor_model_parallel_all_reduce(out)
|
||
return out
|
||
|
||
|
||
def ep_reduce_scatter_output(routed_out, shared_out, dp_size, dp_rank):
|
||
"""Reduce-scatter variant for DP > 1 (upstream PR #2679 optimization).
|
||
|
||
When DP attention is active, the MoE forward receives all-gathered tokens
|
||
from all DP ranks (T_global = T_local * dp_size). Instead of:
|
||
all_reduce(T_global) → slice(T_local) # 2 steps
|
||
we use:
|
||
reduce_scatter(T_global) → T_local # 1 step, half bandwidth
|
||
|
||
Equivalent to upstream:
|
||
jax.lax.psum_scatter(out, axis_name=scatter_axes,
|
||
scatter_dimension=0, tiled=True)
|
||
|
||
Returns: [T_local, H] reduced output for this DP rank only.
|
||
"""
|
||
out = routed_out + shared_out
|
||
if not dist.is_initialized() or dp_size <= 1:
|
||
return out
|
||
|
||
T_global, H = out.shape
|
||
assert T_global % dp_size == 0, (
|
||
f"reduce_scatter: T_global={T_global} not divisible by dp_size={dp_size}")
|
||
|
||
output = torch.empty(T_global // dp_size, H,
|
||
dtype=out.dtype, device=out.device)
|
||
dist.reduce_scatter_tensor(output, out, op=dist.ReduceOp.SUM)
|
||
return output
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Monkey-patch FusedMoE
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def patch_fused_moe_for_ep():
|
||
"""Monkey-patch FusedMoE to support Expert Parallelism.
|
||
|
||
Only patches __init__ and weight_loader. Forward is NOT patched —
|
||
qwen3_5.py forward() handles EP dispatch via:
|
||
ep_mask_and_remap() for expert routing
|
||
ep_reduce_output() for reduction
|
||
"""
|
||
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
|
||
|
||
if not _is_ep_requested():
|
||
logger.info("[EP] not enabled (VLLM_ENABLE_EXPERT_PARALLEL!=1), skipping")
|
||
return
|
||
|
||
logger.info("[EP] installing FusedMoE __init__ + weight_loader patches")
|
||
|
||
_orig_init = FusedMoE.__init__
|
||
|
||
def _ep_init(self, num_experts, top_k, hidden_size, intermediate_size,
|
||
params_dtype=None, reduce_results=False, renormalize=True,
|
||
use_grouped_topk=False, num_expert_group=None,
|
||
topk_group=None, quant_config=None, tp_size=None,
|
||
prefix="", custom_routing_function=None):
|
||
"""EP-aware __init__: allocate only local experts, full intermediate."""
|
||
|
||
ep_size, ep_rank = _get_ep_runtime_config()
|
||
|
||
num_experts_per_rank = num_experts // ep_size
|
||
start_expert_id = ep_rank * num_experts_per_rank
|
||
assert num_experts % ep_size == 0, (
|
||
f"num_experts={num_experts} not divisible by ep_size={ep_size}")
|
||
|
||
logger.info(
|
||
"[EP] FusedMoE init: global=%d ep_size=%d ep_rank=%d "
|
||
"local=%d (experts %d..%d) tp_size=1 (full intermediate)",
|
||
num_experts, ep_size, ep_rank, num_experts_per_rank,
|
||
start_expert_id, start_expert_id + num_experts_per_rank - 1)
|
||
|
||
# tp_size=1: each rank holds full intermediate for its local experts.
|
||
# Why: with EP+TP on same group, a single all-reduce sums across ranks.
|
||
# If weights were TP-sharded (I/tp), each expert would only have 1/tp
|
||
# of intermediate on one rank — all-reduce sums DIFFERENT experts (EP)
|
||
# but does NOT reconstruct TP shards of the SAME expert.
|
||
# With tp_size=1: each rank holds full intermediate, all-reduce sums
|
||
# ep_size ranks × (num_experts/ep_size) experts = all experts. Correct.
|
||
_orig_init(
|
||
self,
|
||
num_experts=num_experts_per_rank,
|
||
top_k=top_k,
|
||
hidden_size=hidden_size,
|
||
intermediate_size=intermediate_size,
|
||
params_dtype=params_dtype,
|
||
reduce_results=False, # we reduce in forward()
|
||
renormalize=renormalize,
|
||
use_grouped_topk=use_grouped_topk,
|
||
num_expert_group=num_expert_group,
|
||
topk_group=topk_group,
|
||
quant_config=quant_config,
|
||
tp_size=1,
|
||
prefix=prefix,
|
||
custom_routing_function=custom_routing_function,
|
||
)
|
||
|
||
# Override num_experts back to global for routing (topk over 256)
|
||
self.num_experts = num_experts
|
||
|
||
# EP attributes — set AFTER _orig_init (super().__init__ wipes __dict__)
|
||
self._ep_enabled = True
|
||
self._ep_size = ep_size
|
||
self._ep_rank = ep_rank
|
||
self._global_num_experts = num_experts
|
||
self._num_experts_per_rank = num_experts_per_rank
|
||
self._start_expert_id = start_expert_id
|
||
|
||
_orig_weight_loader = FusedMoE.weight_loader
|
||
|
||
def _ep_weight_loader(self, param, loaded_weight, weight_name,
|
||
shard_id, expert_id):
|
||
"""EP-aware weight_loader: skip non-local experts, remap id.
|
||
|
||
Since tp_size=1, we force tp_rank=0 during loading so narrow()
|
||
reads offset 0 (full-size weight, no TP split to index into).
|
||
"""
|
||
if not getattr(self, '_ep_enabled', False):
|
||
return _orig_weight_loader(self, param, loaded_weight,
|
||
weight_name, shard_id, expert_id)
|
||
|
||
start = self._start_expert_id
|
||
end = start + self._num_experts_per_rank
|
||
|
||
if expert_id < start or expert_id >= end:
|
||
return # skip non-local expert
|
||
|
||
local_expert_id = expert_id - start
|
||
|
||
# Mock tp_rank=0 in layer.py's LOCAL namespace.
|
||
# layer.py does `from vllm.distributed import get_tensor_model_parallel_rank`
|
||
# so we must patch the name in layer.py's module, not in vllm.distributed.
|
||
import vllm.model_executor.layers.fused_moe.layer as _layer_mod
|
||
_real = _layer_mod.get_tensor_model_parallel_rank
|
||
_layer_mod.get_tensor_model_parallel_rank = lambda: 0
|
||
try:
|
||
_orig_weight_loader(self, param, loaded_weight, weight_name,
|
||
shard_id, local_expert_id)
|
||
finally:
|
||
_layer_mod.get_tensor_model_parallel_rank = _real
|
||
|
||
FusedMoE.__init__ = _ep_init
|
||
FusedMoE.weight_loader = _ep_weight_loader
|
||
|
||
logger.info("[EP] FusedMoE patches installed (__init__, weight_loader)") |