| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MoonEP (https://github.com/MoonshotAI/MoonEP) prepare/finalize. |
| 4 | |
| 5 | BF16 correctness-first proof of concept on top of vLLM's modular kernel |
| 6 | interface. |
| 7 | |
| 8 | MoonEP differs from DeepEP-style backends in two ways that shape this |
| 9 | integration: |
| 10 | |
| 11 | - ``dispatch`` returns tokens already grouped by expert *row* (a fixed |
| 12 | ``[NvS, H]`` layout with ``NvS = S x K`` real slots plus padding) together |
| 13 | with a ``cu_seqlens[E+B]`` segment table and an opaque ``plan``. There is |
| 14 | no per-token topk id tensor after dispatch; the expert compute must be a |
| 15 | grouped GEMM over ``cu_seqlens`` segments. |
| 16 | - Rows ``[E, E+B)`` of the weight/segment space are dynamic redundant-expert |
| 17 | prefetch slots. ``plan.experts_to_copy`` names the source expert of each |
| 18 | slot and ``Buffer.prefetch_weight`` must run between dispatch and expert |
| 19 | compute. |
| 20 | |
| 21 | PoC limitations: |
| 22 | - BF16 / unquantized only, eager only. |
| 23 | - Expert weights are replicated in global-expert order on every rank |
| 24 | (memory-heavy). Production Kimi-K3 serving requires sharded |
| 25 | symmetric-memory expert ownership, where rows ``[0, E)`` physically alias |
| 26 | each home rank's parameter memory. |
| 27 | - Route weights are applied inside the expert compute and MoonEP's |
| 28 | ``combine`` performs the K-sum, so ``finalize`` requires |
| 29 | ``TopKWeightAndReduceNoOP``. |
| 30 | """ |
| 31 | |
| 32 | from typing import Any, NamedTuple |
| 33 | |
| 34 | import torch |
| 35 | import torch.nn.functional as F |
| 36 | |
| 37 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 38 | from vllm.forward_context import get_forward_context |
| 39 | from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig |
| 40 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 41 | TopKWeightAndReduceNoOP, |
| 42 | ) |
| 43 | |
| 44 | MOONEP_DEFAULT_NUM_PREFETCH_SLOTS = 4 |
| 45 | MOONEP_DEFAULT_TOKEN_PADDING = 128 |
| 46 | MOONEP_DEFAULT_NUM_SMS = 32 |
| 47 | # MoonEP's weight-prefetch kernel tiles hidden and intermediate dims by 128. |
| 48 | MOONEP_WEIGHT_TILE = 128 |
| 49 | |
| 50 | |
| 51 | class MoonEPExpertWeightLayout(NamedTuple): |
| 52 | """Contiguous BF16 expert weights in MoonEP ``[E+B, ...]`` layout. |
| 53 | |
| 54 | Rows ``[0, E)`` hold expert weights in global expert order; rows |
| 55 | ``[E, E+B)`` are mutable prefetch slots filled by |
| 56 | ``Buffer.prefetch_weight``. |
| 57 | |
| 58 | The prefetch slots never need re-zeroing between calls: the planner |
| 59 | marks unused slots with ``-1`` in ``plan.experts_to_copy`` (skipped by |
| 60 | ``prefetch_weight``) and gives them an empty ``cu_seqlens`` segment, so |
| 61 | stale slot contents are never read; used slots are fully overwritten. |
| 62 | """ |
| 63 | |
| 64 | full_gate_weight: torch.Tensor |
| 65 | full_up_weight: torch.Tensor |
| 66 | full_down_weight: torch.Tensor |
| 67 | num_prefetch_slots: int |
| 68 | |
| 69 | |
| 70 | def make_moonep_weight_layout( |
| 71 | w13_weight: torch.Tensor, |
| 72 | w2_weight: torch.Tensor, |
| 73 | num_prefetch_slots: int, |
| 74 | ) -> MoonEPExpertWeightLayout: |
| 75 | """Build the replicated ``[E+B, ...]`` PoC weight layout. |
| 76 | |
| 77 | ``w13_weight`` must be ``[E, 2I, H]`` (gate rows first) and ``w2_weight`` |
| 78 | ``[E, H, I]``, both BF16 in global expert order. |
| 79 | """ |
| 80 | if num_prefetch_slots <= 0: |
| 81 | raise ValueError( |
| 82 | f"num_prefetch_slots must be positive, got {num_prefetch_slots}" |
| 83 | ) |
| 84 | if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16: |
| 85 | raise NotImplementedError("MoonEP PoC supports BF16 weights only.") |
| 86 | num_experts, two_i, hidden_size = w13_weight.shape |
| 87 | intermediate_size = two_i // 2 |
| 88 | if tuple(w2_weight.shape) != (num_experts, hidden_size, intermediate_size): |
| 89 | raise ValueError( |
| 90 | f"w2_weight shape {tuple(w2_weight.shape)} does not match " |
| 91 | f"w13_weight shape {tuple(w13_weight.shape)}" |
| 92 | ) |
| 93 | if hidden_size % MOONEP_WEIGHT_TILE or intermediate_size % MOONEP_WEIGHT_TILE: |
| 94 | raise ValueError( |
| 95 | "MoonEP weight prefetch requires hidden_size and intermediate_size " |
| 96 | f"to be multiples of {MOONEP_WEIGHT_TILE}; got H={hidden_size}, " |
| 97 | f"I={intermediate_size}" |
| 98 | ) |
| 99 | |
| 100 | full_gate_weight = torch.empty( |
| 101 | num_experts + num_prefetch_slots, |
| 102 | intermediate_size, |
| 103 | hidden_size, |
| 104 | dtype=torch.bfloat16, |
| 105 | device=w13_weight.device, |
| 106 | ) |
| 107 | full_up_weight = torch.empty_like(full_gate_weight) |
| 108 | full_down_weight = torch.empty( |
| 109 | num_experts + num_prefetch_slots, |
| 110 | hidden_size, |
| 111 | intermediate_size, |
| 112 | dtype=torch.bfloat16, |
| 113 | device=w2_weight.device, |
| 114 | ) |
| 115 | |
| 116 | full_gate_weight[:num_experts].copy_(w13_weight[:, :intermediate_size, :]) |
| 117 | full_up_weight[:num_experts].copy_(w13_weight[:, intermediate_size:, :]) |
| 118 | full_down_weight[:num_experts].copy_(w2_weight) |
| 119 | full_gate_weight[num_experts:].zero_() |
| 120 | full_up_weight[num_experts:].zero_() |
| 121 | full_down_weight[num_experts:].zero_() |
| 122 | |
| 123 | return MoonEPExpertWeightLayout( |
| 124 | full_gate_weight=full_gate_weight.contiguous(), |
| 125 | full_up_weight=full_up_weight.contiguous(), |
| 126 | full_down_weight=full_down_weight.contiguous(), |
| 127 | num_prefetch_slots=num_prefetch_slots, |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | def gather_moonep_weight_layout( |
| 132 | w13_local: torch.Tensor, |
| 133 | w2_local: torch.Tensor, |
| 134 | num_global_experts: int, |
| 135 | num_prefetch_slots: int, |
| 136 | ) -> MoonEPExpertWeightLayout: |
| 137 | """Build the replicated ``[E+B, ...]`` layout from this rank's local experts. |
| 138 | |
| 139 | PoC bridge: each EP rank loads only its own experts (linear placement), |
| 140 | so all-gather them once at load time into global expert order on every |
| 141 | rank. Production MoonEP instead maps rows ``[0, E)`` onto each home |
| 142 | rank's parameter memory via symmetric memory (RFC #52095 item 5). |
| 143 | """ |
| 144 | from vllm.distributed import get_ep_group |
| 145 | |
| 146 | ep_group = get_ep_group() |
| 147 | ep_size = ep_group.world_size |
| 148 | if ep_size == 1: |
| 149 | return make_moonep_weight_layout(w13_local, w2_local, num_prefetch_slots) |
| 150 | if w13_local.size(0) * ep_size != num_global_experts: |
| 151 | raise NotImplementedError( |
| 152 | "MoonEP PoC requires num_experts to be evenly divisible across EP " |
| 153 | f"ranks: {num_global_experts} experts, {ep_size} ranks" |
| 154 | ) |
| 155 | w13_global = ep_group.all_gather(w13_local.contiguous(), dim=0) |
| 156 | w2_global = ep_group.all_gather(w2_local.contiguous(), dim=0) |
| 157 | return make_moonep_weight_layout(w13_global, w2_global, num_prefetch_slots) |
| 158 | |
| 159 | |
| 160 | MOONEP_MIN_DISPATCH_TOKENS = 128 |
| 161 | |
| 162 | |
| 163 | class MoonEPBufferPool: |
| 164 | """Lazily creates ``moonep.Buffer`` instances at power-of-two capacities. |
| 165 | |
| 166 | ``Buffer`` fixes its token capacity ``S`` at construction, so a single |
| 167 | full-capacity buffer would pad every dispatch to |
| 168 | ``max_num_batched_tokens``. The pool instead serves the smallest |
| 169 | power-of-two capacity that fits the step's token count, so decode-sized |
| 170 | batches dispatch a few hundred slots rather than the full capacity. |
| 171 | |
| 172 | Selection must be identical on every EP rank (dispatch is collective and |
| 173 | ``Buffer`` construction is itself a collective): callers derive the token |
| 174 | count from ``dp_metadata`` (the max across DP ranks) so all ranks create |
| 175 | and pick the same buffer at the same step. |
| 176 | """ |
| 177 | |
| 178 | def __init__(self, buffer_kwargs: dict[str, Any], max_tokens_per_rank: int): |
| 179 | self._kwargs = {k: v for k, v in buffer_kwargs.items() if k != "S"} |
| 180 | self.max_tokens_per_rank = max_tokens_per_rank |
| 181 | self._buffers: dict[int, Any] = {} |
| 182 | |
| 183 | def capacity_for(self, num_tokens: int) -> int: |
| 184 | if num_tokens >= self.max_tokens_per_rank: |
| 185 | return self.max_tokens_per_rank |
| 186 | cap = 1 << max(num_tokens - 1, 0).bit_length() |
| 187 | return min(max(cap, MOONEP_MIN_DISPATCH_TOKENS), self.max_tokens_per_rank) |
| 188 | |
| 189 | def get(self, num_tokens: int) -> tuple[int, Any]: |
| 190 | """Return ``(capacity, buffer)`` for the given step token count.""" |
| 191 | from moonep import Buffer # type: ignore[import-not-found] |
| 192 | |
| 193 | capacity = self.capacity_for(num_tokens) |
| 194 | buffer = self._buffers.get(capacity) |
| 195 | if buffer is None: |
| 196 | buffer = Buffer(S=capacity, **self._kwargs) |
| 197 | self._buffers[capacity] = buffer |
| 198 | return capacity, buffer |
| 199 | |
| 200 | def destroy(self) -> None: |
| 201 | for buffer in self._buffers.values(): |
| 202 | buffer.destroy() |
| 203 | self._buffers.clear() |
| 204 | |
| 205 | |
| 206 | class MoonEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): |
| 207 | """Prepare/Finalize using MoonEP balanced dispatch/combine. |
| 208 | |
| 209 | ``prepare`` pads the batch to the selected buffer's static token |
| 210 | capacity, dispatches, runs ``prefetch_weight`` for the planned redundant |
| 211 | experts, and stashes the ``plan`` for ``finalize`` (the same pattern |
| 212 | DeepEP-HT uses for its handle). Downstream expert compute must consume |
| 213 | the expert-grouped ``[NvS, H]`` layout via ``cu_seqlens``. |
| 214 | """ |
| 215 | |
| 216 | def __init__( |
| 217 | self, |
| 218 | buffer_pool: "MoonEPBufferPool", |
| 219 | max_tokens_per_rank: int, |
| 220 | num_dispatchers: int, |
| 221 | num_global_experts: int, |
| 222 | weight_layout: MoonEPExpertWeightLayout | None = None, |
| 223 | ): |
| 224 | super().__init__() |
| 225 | self.buffer_pool = buffer_pool |
| 226 | self.max_tokens_per_rank = max_tokens_per_rank |
| 227 | self.num_dispatchers_ = num_dispatchers |
| 228 | self.num_global_experts = num_global_experts |
| 229 | self.weight_layout = weight_layout |
| 230 | self._fused_experts: Any = None |
| 231 | |
| 232 | # dispatch state consumed by finalize (and the expert runner) |
| 233 | self._plan: Any = None |
| 234 | self._cu_seqlens: torch.Tensor | None = None |
| 235 | self._num_tokens: int = 0 |
| 236 | self._buffer: Any = None # the pool buffer used by the current step |
| 237 | |
| 238 | def post_init_setup(self, fused_experts: mk.FusedMoEExperts) -> None: |
| 239 | # The [E+B] weight layout is attached to the experts by their |
| 240 | # process_weights_after_loading hook (after this runs), so keep the |
| 241 | # reference and resolve the layout lazily in prepare(). |
| 242 | self._fused_experts = fused_experts |
| 243 | |
| 244 | def _resolve_weight_layout(self) -> MoonEPExpertWeightLayout: |
| 245 | if self.weight_layout is None and self._fused_experts is not None: |
| 246 | layout = getattr(self._fused_experts, "weight_layout", None) |
| 247 | if isinstance(layout, MoonEPExpertWeightLayout): |
| 248 | self.weight_layout = layout |
| 249 | # Redundant experts' weights must be in rows [E, E+B) before the |
| 250 | # expert compute reads them; skipping prefetch silently corrupts |
| 251 | # output. |
| 252 | assert self.weight_layout is not None, ( |
| 253 | "MoonEPPrepareAndFinalize: weight layout not available (the " |
| 254 | "experts' process_weights_after_loading has not run)" |
| 255 | ) |
| 256 | return self.weight_layout |
| 257 | |
| 258 | @property |
| 259 | def num_dispatched_slots(self) -> int: |
| 260 | """``NvS`` of the buffer selected by the current step's prepare().""" |
| 261 | assert self._buffer is not None, "prepare() has not been called" |
| 262 | return int(self._buffer._ctx["NvS"]) |
| 263 | |
| 264 | def num_dispatchers(self) -> int: |
| 265 | return self.num_dispatchers_ |
| 266 | |
| 267 | def output_is_reduced(self) -> bool: |
| 268 | # combine returns the fully weighted+reduced token-major output |
| 269 | return True |
| 270 | |
| 271 | @property |
| 272 | def activation_format(self) -> mk.FusedMoEActivationFormat: |
| 273 | return mk.FusedMoEActivationFormat.Standard |
| 274 | |
| 275 | def max_num_tokens_per_rank(self) -> int | None: |
| 276 | return self.max_tokens_per_rank |
| 277 | |
| 278 | def topk_indices_dtype(self) -> torch.dtype | None: |
| 279 | return torch.int32 |
| 280 | |
| 281 | def supports_async(self) -> bool: |
| 282 | return False |
| 283 | |
| 284 | @property |
| 285 | def cu_seqlens(self) -> torch.Tensor: |
| 286 | assert self._cu_seqlens is not None, "prepare() has not been called" |
| 287 | return self._cu_seqlens |
| 288 | |
| 289 | @property |
| 290 | def plan(self) -> Any: |
| 291 | assert self._plan is not None, "prepare() has not been called" |
| 292 | return self._plan |
| 293 | |
| 294 | def _pad_to_capacity( |
| 295 | self, |
| 296 | a1: torch.Tensor, |
| 297 | topk_ids: torch.Tensor, |
| 298 | topk_weights: torch.Tensor, |
| 299 | capacity: int, |
| 300 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: |
| 301 | num_tokens = a1.shape[0] |
| 302 | if num_tokens > capacity: |
| 303 | raise ValueError( |
| 304 | f"MoonEP static buffer capacity exceeded: {num_tokens} tokens " |
| 305 | f"> S={capacity}." |
| 306 | ) |
| 307 | a1 = a1.contiguous() |
| 308 | topk_ids = topk_ids.to(dtype=torch.int32) |
| 309 | topk_weights = topk_weights.to(dtype=torch.float32) |
| 310 | # MoonEP's planner has no sentinel for invalid expert ids, so map |
| 311 | # negative ids (padded/invalid slots from the engine) to expert 0 |
| 312 | # with zero route weight — the same convention as the capacity |
| 313 | # padding rows below. Wasted compute, never wrong output. |
| 314 | invalid = topk_ids < 0 |
| 315 | topk_ids = torch.where(invalid, torch.zeros_like(topk_ids), topk_ids) |
| 316 | topk_weights = torch.where( |
| 317 | invalid, torch.zeros_like(topk_weights), topk_weights |
| 318 | ).contiguous() |
| 319 | topk_ids = topk_ids.contiguous() |
| 320 | if num_tokens == capacity: |
| 321 | return a1, topk_ids, topk_weights, num_tokens |
| 322 | pad = capacity - num_tokens |
| 323 | return ( |
| 324 | F.pad(a1, (0, 0, 0, pad)), |
| 325 | F.pad(topk_ids, (0, 0, 0, pad)), |
| 326 | F.pad(topk_weights, (0, 0, 0, pad)), |
| 327 | num_tokens, |
| 328 | ) |
| 329 | |
| 330 | def prepare( |
| 331 | self, |
| 332 | a1: torch.Tensor, |
| 333 | topk_weights: torch.Tensor, |
| 334 | topk_ids: torch.Tensor, |
| 335 | num_experts: int, |
| 336 | expert_map: torch.Tensor | None, |
| 337 | apply_router_weight_on_input: bool, |
| 338 | quant_config: FusedMoEQuantConfig, |
| 339 | defer_input_quant: bool = False, |
| 340 | ) -> mk.PrepareResultType: |
| 341 | if a1.dtype != torch.bfloat16: |
| 342 | raise NotImplementedError("MoonEP PoC supports BF16 hidden states only.") |
| 343 | if quant_config.quant_dtype is not None: |
| 344 | raise NotImplementedError("MoonEP PoC does not support quantized dispatch.") |
| 345 | assert num_experts == self.num_global_experts |
| 346 | assert self._plan is None, ( |
| 347 | "MoonEPPrepareAndFinalize.prepare() called again before finalize()" |
| 348 | ) |
| 349 | if apply_router_weight_on_input: |
| 350 | assert topk_ids.size(1) == 1, ( |
| 351 | "apply_router_weight_on_input requires top-1 routing" |
| 352 | ) |
| 353 | a1 = a1 * topk_weights.to(a1.dtype) |
| 354 | |
| 355 | # Buffer choice must be rank-symmetric: use the max token count |
| 356 | # across DP ranks (every rank sees the same dp_metadata), falling |
| 357 | # back to the local count outside a forward context (tests). |
| 358 | num_tokens = a1.shape[0] |
| 359 | try: |
| 360 | dp_meta = get_forward_context().dp_metadata |
| 361 | except AssertionError: |
| 362 | dp_meta = None |
| 363 | if dp_meta is not None: |
| 364 | num_tokens = int(dp_meta.num_tokens_across_dp_cpu.max()) |
| 365 | capacity, self._buffer = self.buffer_pool.get(num_tokens) |
| 366 | |
| 367 | a1, topk_ids, topk_weights, self._num_tokens = self._pad_to_capacity( |
| 368 | a1, topk_ids, topk_weights, capacity |
| 369 | ) |
| 370 | tokens_per_expert = torch.bincount( |
| 371 | topk_ids.reshape(-1).to(dtype=torch.int64), |
| 372 | minlength=num_experts, |
| 373 | ).to(dtype=torch.int32) |
| 374 | |
| 375 | hidden_nvsh, route_weights_nvs, cu_seqlens, plan = self._buffer.dispatch( |
| 376 | a1, |
| 377 | topk_weights, |
| 378 | topk_ids, |
| 379 | tokens_per_expert, |
| 380 | ) |
| 381 | self._plan = plan |
| 382 | self._cu_seqlens = cu_seqlens |
| 383 | |
| 384 | weight_layout = self._resolve_weight_layout() |
| 385 | self._buffer.prefetch_weight( |
| 386 | plan=plan, |
| 387 | full_gate_weight=weight_layout.full_gate_weight, |
| 388 | full_up_weight=weight_layout.full_up_weight, |
| 389 | full_down_weight=weight_layout.full_down_weight, |
| 390 | ) |
| 391 | |
| 392 | # Segment sizes per [E+B] weight row. NOTE: this is per *row*, not |
| 393 | # per local expert, and NvS rows are expert-grouped rather than |
| 394 | # token-major — only MoonEP-aware expert implementations can consume |
| 395 | # this activation layout. |
| 396 | expert_num_tokens = torch.diff(cu_seqlens, prepend=cu_seqlens.new_zeros(1)) |
| 397 | expert_tokens_meta = mk.ExpertTokensMetadata( |
| 398 | expert_num_tokens=expert_num_tokens, |
| 399 | expert_num_tokens_cpu=None, |
| 400 | ) |
| 401 | |
| 402 | # MoonEP has no post-dispatch per-token topk ids/weights; route |
| 403 | # weights come back in NvS order. |
| 404 | return hidden_nvsh, None, expert_tokens_meta, None, route_weights_nvs |
| 405 | |
| 406 | def finalize( |
| 407 | self, |
| 408 | output: torch.Tensor, |
| 409 | fused_expert_output: torch.Tensor, |
| 410 | topk_weights: torch.Tensor, |
| 411 | topk_ids: torch.Tensor, |
| 412 | apply_router_weight_on_input: bool, |
| 413 | weight_and_reduce_impl: mk.TopKWeightAndReduce, |
| 414 | ) -> None: |
| 415 | # Route weights are applied in the expert compute and MoonEP's |
| 416 | # combine performs the K-sum, so any weight/reduce work in finalize |
| 417 | # means MoonEP was paired with the wrong kind of experts. |
| 418 | assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP), ( |
| 419 | "MoonEP requires TopKWeightAndReduceNoOP, got " |
| 420 | f"{type(weight_and_reduce_impl).__name__}" |
| 421 | ) |
| 422 | combined, _, _ = self._buffer.combine( |
| 423 | plan=self.plan, |
| 424 | hidden_nvsh=fused_expert_output, |
| 425 | route_weights_nvs=None, |
| 426 | ) |
| 427 | output.copy_(combined[: self._num_tokens]) |
| 428 | self._plan = None |
| 429 | self._cu_seqlens = None |