diff --git a/.scratch/moonep/integration_research.md b/.scratch/moonep/integration_research.md new file mode 100644 index 000000000000..e53b078dbe57 --- /dev/null +++ b/.scratch/moonep/integration_research.md @@ -0,0 +1,108 @@ +# MoonEP Integration Research Notes + +## Context + +- Upstream issue `sgl-project/sglang#32607` tracks Kimi-K3 work and includes an + open roadmap item to integrate MoonEP for better all-to-all performance. +- Local MoonEP checkout used for these notes: `/home/user/MoonEP` at commit + `0f385f0`. +- SGLang already has several MoE A2A backends selected by `--moe-a2a-backend`; + DeepEP-style backends are routed through `MaybeTboDeepEPDispatcher` and the + common dispatcher abstraction + (`python/sglang/srt/layers/moe/fused_moe_triton/layer.py:109-150`, + `python/sglang/srt/batch_overlap/two_batch_overlap.py:1074-1100`). + +## Primary-source findings + +### Existing SGLang DeepEP seam + +- The MoE A2A enum and helpers live in + `python/sglang/srt/layers/moe/utils.py`. DeepEP, Mooncake, MoRI, NIXL, PPLX, + FlashInfer, MegaMoE, and Ascend backends are represented there + (`python/sglang/srt/layers/moe/utils.py:28-92`). +- `python/sglang/srt/layers/moe/fused_moe_triton/layer.py` creates the + dispatcher. DeepEP-class backends are wrapped by + `MaybeTboDeepEPDispatcher`, which can allocate one or two inner dispatchers + depending on two-batch overlap + (`python/sglang/srt/layers/moe/fused_moe_triton/layer.py:109-150`, + `python/sglang/srt/batch_overlap/two_batch_overlap.py:1074-1100`). +- `python/sglang/srt/layers/moe/token_dispatcher/deepep.py` exposes two + dispatch output contracts: + - normal mode: dispatched hidden states, top-k ids/weights, and + `num_recv_tokens_per_expert`; + - low-latency mode: packed hidden states, top-k ids/weights, `masked_m`, and + `expected_m` + (`python/sglang/srt/layers/moe/token_dispatcher/deepep.py:96-154`, + `python/sglang/srt/layers/moe/token_dispatcher/deepep.py:496-846`). +- Current expert runners consume those DeepEP contracts through + `DeepEPNormalCombineInput` / `DeepEPLLCombineInput`; they do not consume a + MoonEP communication plan or MoonEP `cu_seqlens` + (`python/sglang/srt/layers/moe/ep_moe/layer.py:270-341`). + +### MoonEP contract + +- MoonEP's public API is `moonep.Buffer` in `/home/user/MoonEP/moonep/api.py`. + It is constructed with static shape parameters `S`, `H`, `K`, `E`, EP rank + count, and optional `B` prefetch slots + (`/home/user/MoonEP/moonep/api.py:431-499`). +- `Buffer.dispatch(...)` returns + `(hidden_nvsh, route_weights_nvs, cu_seqlens, plan)` and requires + `topk_experts_sk` as `int32` plus a local `tokens_per_expert` vector when + creating a fresh plan (`/home/user/MoonEP/moonep/api.py:685-814`). +- `Buffer.combine(...)` requires the saved `MoonEPCommPlan` and returns + token-major output plus optional gathered route weights + (`/home/user/MoonEP/moonep/api.py:881-1007`). +- `Buffer.prefetch_weight(...)` requires the saved plan plus three contiguous + BF16 full expert-weight tensors for gate/up/down projections with shape + `[E+B, H, H']` (`/home/user/MoonEP/moonep/api.py:816-879`). +- MoonEP's README states that framework integration requires one contiguous + symmetric-memory expert weight tensor per projection. This is a stricter + contract than SGLang's existing local-expert DeepEP runner path + (`/home/user/MoonEP/README.md:43-81`). + +## Design implication + +MoonEP should be a distinct `moonep` backend, not an alias for `deepep`. +Pretending it is DeepEP-compatible would be incorrect because SGLang's current +DeepEP runners expect `num_recv_tokens_per_expert` or `masked_m`, while MoonEP +produces `cu_seqlens` and a reusable `MoonEPCommPlan`. + +## First implementation slice + +The initial patch should: + +1. add `moonep` to the CLI/backend enum and EP-size resolution path; +2. add an explicit `MoonEPDispatcher` placeholder so backend selection is + centralized and future implementation has a clear seam; +3. fail fast with a clear runtime message rather than falling into DeepEP or + standard-dispatch assumptions. + +## Follow-up tickets + +1. **MoonEP dispatch output contract** — done in + `9865fc123 Add MoonEP dispatcher data contract` + - Add `DispatchOutputFormat.MOONEP` and `CombineInputFormat.MOONEP`. + - Carry `hidden_nvsh`, `route_weights_nvs`, `cu_seqlens`, and + `MoonEPCommPlan` through the MoE core. + +2. **MoonEP buffer lifecycle** — initial facade done in + `MoonEPBuffer`; runtime sizing policy still needs validation + - Add a process-wide MoonEP buffer facade similar to `DeepEPBuffer`. + - Key the buffer by static `S`, `H`, `K`, `E`, EP group, and `B`. + - Decide how SGLang should choose `S` for prefill/decode/cuda-graph capture. + +3. **MoonEP weight layout** + - Add contiguous symmetric-memory `[E+B, H, H']` weight buffers for each + routed expert projection. + - Map existing checkpoint-loaded local expert weights into rows `[0, E)`. + - Allocate/pool rows `[E, E+B)` for prefetch slots. + +4. **MoonEP expert runner** + - Implement a grouped expert GEMM that consumes MoonEP `cu_seqlens`. + - Wire `prefetch_weight(plan, ...)` between dispatch and expert GEMM. + - Start with BF16 inference before quantized Kimi-K3 paths. + +5. **Kimi-K3 integration and validation** + - Add a Kimi-K3 launch recipe once runtime execution is implemented. + - Validate on the multi-GPU/NVLink MoonEP test matrix; local unit tests can + only cover backend selection and fail-fast behavior. diff --git a/.scratch/moonep/tickets.md b/.scratch/moonep/tickets.md new file mode 100644 index 000000000000..ba91e9b6ff1c --- /dev/null +++ b/.scratch/moonep/tickets.md @@ -0,0 +1,31 @@ +# MoonEP integration tickets + +Issue tracker: `wirybeaver/sglang` +Map: #7 https://github.com/wirybeaver/sglang/issues/7 + +## Completed + +- #8 https://github.com/wirybeaver/sglang/issues/8 — Research DeepEP and MoonEP integration seam +- #9 https://github.com/wirybeaver/sglang/issues/9 — Recognize backend and fail fast +- #10 https://github.com/wirybeaver/sglang/issues/10 — Add dispatcher data contract +- #11 https://github.com/wirybeaver/sglang/issues/11 — Add buffer facade and env sizing knobs + +## Completed + +- #12 https://github.com/wirybeaver/sglang/issues/12 — Decide static token-capacity policy + +## Completed + +- #13 https://github.com/wirybeaver/sglang/issues/13 — Add contiguous symmetric expert weight layout + +## Completed + +- #14 https://github.com/wirybeaver/sglang/issues/14 — Add BF16 expert runner for cu_seqlens + +## Completed + +- #15 https://github.com/wirybeaver/sglang/issues/15 — Wire runtime dispatch/prefetch/compute/combine path + +## Completed + +- #16 https://github.com/wirybeaver/sglang/issues/16 — Validate BF16 MoonEP PoC and prepare upstream handoff notes diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000000..101dcb2bc49e --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,13 @@ +## Agent skills + +### Issue tracker + +Issues for this work are tracked in GitHub Issues on `wirybeaver/sglang`; completed PoC issues are later cloned/escalated upstream to `sgl-project/sglang`. See `docs/agents/issue-tracker.md`. + +### Triage labels + +Use the default triage labels: `needs-triage`, `needs-info`, `ready-for-agent`, `ready-for-human`, `wontfix`. See `docs/agents/triage-labels.md`. + +### Domain docs + +This repo uses a single-context layout: root `CONTEXT.md` plus `docs/adr/`. See `docs/agents/domain.md`. diff --git a/docs/agents/domain.md b/docs/agents/domain.md new file mode 100644 index 000000000000..fcf62bb16f7c --- /dev/null +++ b/docs/agents/domain.md @@ -0,0 +1,51 @@ +# Domain Docs + +How the engineering skills should consume this repo's domain documentation when exploring the codebase. + +## Before exploring, read these + +- **`CONTEXT.md`** at the repo root, or +- **`CONTEXT-MAP.md`** at the repo root if it exists — it points at one `CONTEXT.md` per context. Read each one relevant to the topic. +- **`docs/adr/`** — read ADRs that touch the area you're about to work in. In multi-context repos, also check `src//docs/adr/` for context-scoped decisions. + +If any of these files don't exist, **proceed silently**. Don't flag their absence; don't suggest creating them upfront. The `/domain-modeling` skill (reached via `/grill-with-docs` and `/improve-codebase-architecture`) creates them lazily when terms or decisions actually get resolved. + +## File structure + +Single-context repo (most repos): + +```text +/ +├── CONTEXT.md +├── docs/adr/ +│ ├── 0001-event-sourced-orders.md +│ └── 0002-postgres-for-write-model.md +└── src/ +``` + +Multi-context repo (presence of `CONTEXT-MAP.md` at the root): + +```text +/ +├── CONTEXT-MAP.md +├── docs/adr/ ← system-wide decisions +└── src/ + ├── ordering/ + │ ├── CONTEXT.md + │ └── docs/adr/ ← context-specific decisions + └── billing/ + ├── CONTEXT.md + └── docs/adr/ +``` + +## Use the glossary's vocabulary + +When your output names a domain concept (in an issue title, a refactor proposal, a hypothesis, a test name), use the term as defined in `CONTEXT.md`. Don't drift to synonyms the glossary explicitly avoids. + +If the concept you need isn't in the glossary yet, that's a signal — either you're inventing language the project doesn't use (reconsider) or there's a real gap (note it for `/domain-modeling`). + +## Flag ADR conflicts + +If your output contradicts an existing ADR, surface it explicitly rather than silently overriding: + +> _Contradicts ADR-0007 (event-sourced orders) — but worth reopening because…_ diff --git a/docs/agents/issue-tracker.md b/docs/agents/issue-tracker.md new file mode 100644 index 000000000000..077026cef439 --- /dev/null +++ b/docs/agents/issue-tracker.md @@ -0,0 +1,55 @@ +# Issue tracker: GitHub + +Issues and PRDs for this repo live as GitHub issues on `wirybeaver/sglang`. Use the `gh` CLI for all operations. + +This repo also has upstream `sgl-project/sglang`. For PoC work, create and update issues in `wirybeaver/sglang`. Once a PoC is complete, clone or summarize the issue upstream before creating PRs against upstream. + +Use `--repo wirybeaver/sglang` when needed to avoid accidentally writing to upstream. + +## Conventions + +- **Create an issue**: `gh issue create --repo wirybeaver/sglang --title "..." --body "..."`. Use a heredoc for multi-line bodies. +- **Read an issue**: `gh issue view --repo wirybeaver/sglang --comments`, filtering comments by `jq` and also fetching labels. +- **List issues**: `gh issue list --repo wirybeaver/sglang --state open --json number,title,body,labels,comments --jq '[.[] | {number, title, body, labels: [.labels[].name], comments: [.comments[].body]}]'` with appropriate `--label` and `--state` filters. +- **Comment on an issue**: `gh issue comment --repo wirybeaver/sglang --body "..."` +- **Apply / remove labels**: `gh issue edit --repo wirybeaver/sglang --add-label "..."` / `--remove-label "..."` +- **Close**: `gh issue close --repo wirybeaver/sglang --comment "..."` + +## Pull requests as a triage surface + +**PRs as a request surface: no.** _(Set to `yes` if this repo treats external PRs as feature requests; `/triage` reads this flag.)_ + +When set to `yes`, PRs run through the same labels and states as issues, using the `gh pr` equivalents: + +- **Read a PR**: `gh pr view --repo wirybeaver/sglang --comments` and `gh pr diff --repo wirybeaver/sglang` for the diff. +- **List external PRs for triage**: `gh pr list --repo wirybeaver/sglang --state open --json number,title,body,labels,author,authorAssociation,comments` then keep only `authorAssociation` of `CONTRIBUTOR`, `FIRST_TIME_CONTRIBUTOR`, or `NONE` (drop `OWNER`/`MEMBER`/`COLLABORATOR`). +- **Comment / label / close**: `gh pr comment`, `gh pr edit --add-label`/`--remove-label`, `gh pr close`, all with `--repo wirybeaver/sglang`. + +GitHub shares one number space across issues and PRs, so a bare `#42` may be either — resolve with `gh pr view 42 --repo wirybeaver/sglang` and fall back to `gh issue view 42 --repo wirybeaver/sglang`. + +## Upstream handoff + +When a PoC issue is complete and ready for upstream work: + +1. Summarize the accepted scope, decisions, implementation notes, and validation results from the `wirybeaver/sglang` issue. +2. Create or update the corresponding upstream issue in `sgl-project/sglang`. +3. Link the upstream issue and PRs back to the fork issue before closing the fork issue. + +## When a skill says "publish to the issue tracker" + +Create a GitHub issue in `wirybeaver/sglang`. + +## When a skill says "fetch the relevant ticket" + +Run `gh issue view --repo wirybeaver/sglang --comments`. + +## Wayfinding operations + +Used by `/wayfinder`. The **map** is a single issue with **child** issues as tickets, all in `wirybeaver/sglang` unless explicitly handing off upstream. + +- **Map**: a single issue labelled `wayfinder:map`, holding the Notes / Decisions-so-far / Fog body. `gh issue create --repo wirybeaver/sglang --label wayfinder:map`. +- **Child ticket**: an issue linked to the map as a GitHub sub-issue (`gh api` on the sub-issues endpoint). Where sub-issues aren't enabled, add the child to a task list in the map body and put `Part of #` at the top of the child body. Labels: `wayfinder:` (`research`/`prototype`/`grilling`/`task`). Once claimed, the ticket is assigned to the driving dev. +- **Blocking**: GitHub's **native issue dependencies** — the canonical, UI-visible representation. Add an edge with `gh api --method POST repos/wirybeaver/sglang/issues//dependencies/blocked_by -F issue_id=`, where `` is the blocker's numeric **database id** (`gh api repos/wirybeaver/sglang/issues/ --jq .id`, _not_ the `#number` or `node_id`). GitHub reports `issue_dependencies_summary.blocked_by` (open blockers only — the live gate). Where dependencies aren't available, fall back to a `Blocked by: #, #` line at the top of the child body. A ticket is unblocked when every blocker is closed. +- **Frontier query**: list the map's open children (`gh issue list --repo wirybeaver/sglang --state open`, scoped to the map's sub-issues / task list), drop any with an open blocker (`issue_dependencies_summary.blocked_by > 0`, or an open issue in the `Blocked by` line) or an assignee; first in map order wins. +- **Claim**: `gh issue edit --repo wirybeaver/sglang --add-assignee @me` — the session's first write. +- **Resolve**: `gh issue comment --repo wirybeaver/sglang --body ""`, then `gh issue close --repo wirybeaver/sglang`, then append a context pointer (gist + link) to the map's Decisions-so-far. diff --git a/docs/agents/triage-labels.md b/docs/agents/triage-labels.md new file mode 100644 index 000000000000..b716855d485f --- /dev/null +++ b/docs/agents/triage-labels.md @@ -0,0 +1,15 @@ +# Triage Labels + +The skills speak in terms of five canonical triage roles. This file maps those roles to the actual label strings used in this repo's issue tracker. + +| Label in mattpocock/skills | Label in our tracker | Meaning | +| -------------------------- | -------------------- | ---------------------------------------- | +| `needs-triage` | `needs-triage` | Maintainer needs to evaluate this issue | +| `needs-info` | `needs-info` | Waiting on reporter for more information | +| `ready-for-agent` | `ready-for-agent` | Fully specified, ready for an AFK agent | +| `ready-for-human` | `ready-for-human` | Requires human implementation | +| `wontfix` | `wontfix` | Will not be actioned | + +When a skill mentions a role (e.g. "apply the AFK-ready triage label"), use the corresponding label string from this table. + +Edit the right-hand column to match whatever vocabulary you actually use. diff --git a/docs/specs/moonep-integration.md b/docs/specs/moonep-integration.md new file mode 100644 index 000000000000..845f7ab25b64 --- /dev/null +++ b/docs/specs/moonep-integration.md @@ -0,0 +1,248 @@ +# MoonEP integration for Kimi-K3 PoC + +## Status + +- Workflow phase: BF16 PoC validation complete; upstream PR handoff next. +- Branch/worktree: `sg-moonep` in `/home/user/sg-moonep`. +- Local MoonEP reference: `/home/user/MoonEP` at commit `0f385f0`. +- Fork issue map: `wirybeaver/sglang#7`. +- Upstream roadmap context: `sgl-project/sglang#32607` includes the Kimi-K3 + MoonEP integration item. + +Completed implementation slices: + +1. `1396477f8 Recognize MoonEP MoE A2A backend` +2. `9865fc123 Add MoonEP dispatcher data contract` +3. `d47680245 Add MoonEP buffer facade` +4. `eed214860 Document MoonEP token capacity policy` +5. `ae30e8e9d Add MoonEP BF16 weight layout` +6. `8b84353f3 Add MoonEP BF16 expert runner` +7. `0dc3f745a Wire MoonEP BF16 runtime path` +8. `c11c7040a Add MoonEP BF16 validation script` + +Runtime dispatch/prefetch/compute/combine is now executable for the BF16 +correctness PoC. It remains experimental and guarded by documented limitations +rather than being production-ready for full Kimi-K3 serving. + +## Problem + +SGLang currently has DeepEP-style MoE all-to-all backends whose expert runners +consume either DeepEP normal outputs (`num_recv_tokens_per_expert`) or DeepEP +low-latency outputs (`masked_m`/`expected_m`). MoonEP has a different runtime +contract: dispatch returns a `MoonEPCommPlan`, `cu_seqlens`, and expert-grouped +`[NvS, H]` token storage, and its expert compute requires contiguous +symmetric-memory expert weights in `[E+B, H, H']` layout. + +Treating MoonEP as a DeepEP alias would silently route MoonEP outputs into the +wrong runner contract. The integration must therefore introduce a distinct +MoonEP backend path. + +## Goals + +- Add a production-shaped `--moe-a2a-backend moonep` path for Kimi-K3 PoC work. +- Preserve fail-fast behavior until every required runtime contract is wired. +- Reuse existing SGLang MoE seams where possible: backend enum, dispatcher + abstraction, `MaybeTboDeepEPDispatcher`, and runtime resource ownership. +- Keep implementation slices modular and independently reviewable. +- Start with BF16 inference support before Kimi-K3 quantized fast paths. + +## Non-goals for the PoC + +- Do not implement MoonEP training backward support initially. +- Do not make MoonEP look DeepEP-compatible by adapting `cu_seqlens` into + DeepEP's `num_recv_tokens_per_expert`/`masked_m` contracts. +- Do not add quantized Kimi-K3 kernels before the BF16 path is functional. +- Do not remove the runtime guard until dispatch, weight layout, expert compute, + combine, and validation are all in place. + +## Design decisions + +### MoonEP is a distinct A2A backend + +`MoeA2ABackend.MOONEP` should stay separate from `DEEPEP`, `MOONCAKE`, `NIXL`, +and `PPLX`. SGLang can instantiate a `MoonEPDispatcher` via the same high-level +MoE dispatcher factory, but the dispatcher output format is `MOONEP`, not a +DeepEP format. + +### Dispatch contract + +MoonEP dispatch output carries: + +- `hidden_states`: MoonEP-dispatched `[NvS, H]` tensor. +- `route_weights_nvs`: optional `[NvS]` route weights. +- `cu_seqlens`: `[E+B]` expert-row cumulative sequence lengths. +- `plan`: opaque MoonEP plan object; typed as `Any` in SGLang so importing the + dispatcher module does not require the optional `moonep` package. + +MoonEP combine input carries: + +- expert output `hidden_states` in `[NvS, H]` layout. +- `route_weights_nvs`, when route-weight gather is needed. +- the saved `plan`. + +### Buffer lifecycle + +MoonEP buffer allocation is static by shape. `MoonEPBuffer` caches process-wide +buffers in `get_resources().buffers["moonep_ep_state"]`, keyed by: + +- `S`: max dispatch tokens per rank, +- `H`: hidden size, +- `K`: router top-k, +- `E`: total routed experts in EP group, +- EP group identity and rank count, +- `B`: prefetch slots, +- token padding, +- communication SM count. + +Environment knobs: + +- `SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` defaults to `128`. +- `SGLANG_MOONEP_NUM_PREFETCH_SLOTS` defaults to `-1`, meaning `E / EP`. +- `SGLANG_MOONEP_TOKEN_PADDING` defaults to `128`. +- `SGLANG_MOONEP_NUM_SMS` defaults to `32`. + +Accepted static token-capacity policy for the PoC: + +- `SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` is the single static `S` + capacity for both prefill and decode on each process. +- The dispatcher pads runtime batches with `num_tokens <= S` up to `S` before + calling MoonEP. Padding rows use zero hidden states, zero route weights, and + expert `0` top-k ids so MoonEP's `topk_experts_sk` and `tokens_per_expert` + remain internally consistent. The combined output is sliced back to the + original token count. +- Runtime batches with `num_tokens > S` fail fast and instruct the operator to + increase `SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK`. +- CUDA graph capture is disabled for the initial PoC. MoonEP has static buffer + shapes, but the first executable SGLang path should validate correctness and + NVLink behavior in eager mode before making capture/replay guarantees. + +### Weight layout + +Each routed expert projection needs one contiguous symmetric-memory tensor with +shape `[E+B, H, H']`: + +- rows `[0, E)`: all experts in global expert-id order; +- rows `[E, E+B)`: local prefetch slots filled by `buffer.prefetch_weight`. + +The layout must preserve SGLang checkpoint loading semantics while exposing the +MoonEP contiguous row contract to the expert GEMM. Prefer a small adapter around +existing expert weight ownership before introducing broad abstractions. + +PoC weight-layout limitation: + +- The first SGLang integration slice stores global BF16 expert rows in the + layer's `w13_weight`/`w2_weight` tensors and derives cached contiguous + gate/up/down tensors with `B` prefetch slots. This makes the BF16 runtime path + easy to validate but is memory-heavy for production-scale Kimi-K3. +- A production Kimi-K3 implementation should replace this with true + symmetric-memory expert rows whose physical storage remains sharded by home + rank. + +### Expert runner + +The first runnable path should be BF16 inference: + +1. Dispatch hidden states and route weights with MoonEP. +2. Prefetch selected remote expert weights using the saved plan. +3. Run a MoonEP-compatible grouped expert GEMM over `[NvS, H]` and + `cu_seqlens[E+B]`. +4. Combine expert outputs with the saved plan. + +Quantized Kimi-K3 variants should come after this path is correct. + +The initial runner is a correctness-first Python segment loop over +`cu_seqlens`. It applies BF16 gate/up/down weights per segment, multiplies by +MoonEP-dispatched route weights, and returns `MoonEPCombineInput`. This is not +the final performance kernel; it exists to validate the dispatcher/data/weight +contracts before replacing the segment loop with a grouped GEMM. + +Runtime wiring now performs the full eager PoC sequence in SGLang: + +1. pad the runtime token batch to the static MoonEP `S` capacity; +2. compute `tokens_per_expert` for the padded top-k ids; +3. call `MoonEPBuffer.dispatch`; +4. derive active expert ids for each `cu_seqlens` segment; +5. call `MoonEPBuffer.prefetch_weight` with the cached BF16 layout; +6. run the BF16 segment runner; +7. call `MoonEPBuffer.combine`; and +8. slice the combined output back to the original token count. + +Known runtime limitations: + +- eager only; CUDA graph capture is disabled by server-args handling; +- BF16/unquantized only; +- no training backward support; +- no production-scale non-replicated weight mapping yet; +- dispatch and prefetch currently run synchronously for correctness-first + validation. + +## Phased tickets + +The GitHub issue tracker lives on `wirybeaver/sglang`; completed PoC issues are +later cloned or summarized upstream to `sgl-project/sglang` before upstream PRs. + +| Phase | Status | Deliverable | +| --- | --- | --- | +| 1 | Done | Research DeepEP and MoonEP integration seam (`wirybeaver/sglang#8`). | +| 2 | Done | Recognize `moonep` backend and fail fast (`wirybeaver/sglang#9`). | +| 3 | Done | Add MoonEP dispatch/combine data contract (`wirybeaver/sglang#10`). | +| 4 | Done | Add process-wide MoonEP buffer facade (`wirybeaver/sglang#11`). | +| 5 | Done | Decide static token-capacity policy (`wirybeaver/sglang#12`). | +| 6 | Done | Add MoonEP contiguous symmetric weight layout (`wirybeaver/sglang#13`). | +| 7 | Done | Add BF16 MoonEP expert runner consuming `cu_seqlens` (`wirybeaver/sglang#14`). | +| 8 | Done | Wire runtime dispatcher dispatch/prefetch/compute/combine (`wirybeaver/sglang#15`). | +| 9 | Done | Add BF16 PoC validation and upstream handoff notes (`wirybeaver/sglang#16`). | + +Issue links are tracked in `.scratch/moonep/tickets.md`. + +## Validation plan + +Local/unit validation: + +- Backend enum/CLI recognizes `moonep`. +- Server args fail fast while runtime support is incomplete. +- Dispatch/combine formats type-check against the dispatcher protocol. +- Buffer facade keying, env defaults, and destruction work with a fake `moonep` + module. +- `python3 -m compileall` and `git diff --check` pass. + +Runtime validation completed: + +- Added `scripts/moonep/validate_moonep_bf16_poc.py`, a torchrun correctness + harness for `MoonEPDispatcher.dispatch -> prefetch_weight ->` BF16 segment + runner `-> combine`. The script compares every rank's token-major output + against a local PyTorch reference using the same global BF16 expert weights + and top-k routes. +- Vast.ai 4x H100 SXM (`NV6` topology) was not sufficient for MoonEP: CUDA + reported `CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED=0`, and MoonEP failed + during multicast granularity setup with `operation not supported`. This + confirms that pairwise NVLink alone is not enough for this MoonEP path. +- Vast.ai 8x H100 SXM (`NV18` topology) passed the BF16 PoC validation. All + eight ranks reported `global_ok=true`, `max_abs_err=0.0`, `relative_err=0.0` + for `tokens=128`, `hidden_size=1024`, `intermediate_size=128`, `top_k=2`, + `num_experts=16`, and `num_prefetch_slots=2`. +- Validation artifacts are local under + `.scratch/moonep/remote-validation-20260802T050429Z/`, including the 4x H100 + failure, 8x H100 success, environment records, command logs, and SHA-256 + manifests. + +Docker/NCU validation notes: + +- `docker-ncu-benchmark` preflight was run against the Vast.ai H100 hosts. The + H100 offers with the required NVSwitch/multicast capability are container + instances, not VM/Docker hosts: Docker, `nvidia-ctk`, and + `nvidia-container-cli` were absent, so `docker run --gpus all` could not be + tested there. +- Host Nsight Compute was present, but `RmProfilingAdminOnly=1`; a direct NCU + hardware-counter smoke test failed with `ERR_NVGPUCTRPERM`. No performance + claims are made from NCU profiling. +- Available Vast.ai VM/Docker offers with multiple RTX 5090/4090 GPUs had + `bw_nvlink=0` and are useful only for build/Docker smoke, not for MoonEP + all-to-all correctness. The BF16 MoonEP validation should use an 8x H100/H200 + or B200-class NVSwitch/multicast host. + +Known local environment blockers at time of validation: + +- `orjson` is missing, so Python unit tests fail during SGLang import. +- `pytest` is not installed. +- `cargo` is not installed, so Rust `cargo check` cannot run. diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index f1290b9c62f0..e1e1f84e913b 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -2201,6 +2201,7 @@ def _cutlass_moe_env_override(view: Any) -> dict: { "megamoe", "deepep", + "moonep", "mooncake", "nixl", "ascend_fuseep", diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 08a05b3ae5fa..a3fe4beb0b0c 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -27,6 +27,7 @@ ) from sglang.srt.layers.moe.token_dispatcher import ( DeepEPDispatcher, + MoonEPDispatcher, MooncakeEPDispatcher, MoriEPDispatcher, NixlEPDispatcher, @@ -1079,6 +1080,10 @@ def __init__(self, **kwargs): self._inners = [ DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) ] + elif get_moe_a2a_backend().is_moonep(): + self._inners = [ + MoonEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) + ] elif get_moe_a2a_backend().is_mooncake(): self._inners = [ MooncakeEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers) @@ -1122,6 +1127,9 @@ def combine_a(self, **kwargs): def combine_b(self, **kwargs): return self._execute("combine_b", **kwargs) + def prefetch_weight(self, **kwargs): + return self._execute("prefetch_weight", **kwargs) + def register_deepep_dispatch_hook(self, hook): handle_list = [] for inner in self._inners: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5d21605de71b..4befe5b99514 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -748,6 +748,13 @@ class Envs: SGLANG_NIXL_EP_BF16_DISPATCH = EnvBool(False) SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # MoonEP + SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # -1 uses MoonEP's training-safe default B = E / EP. + SGLANG_MOONEP_NUM_PREFETCH_SLOTS = EnvInt(-1) + SGLANG_MOONEP_TOKEN_PADDING = EnvInt(128) + SGLANG_MOONEP_NUM_SMS = EnvInt(32) + # PPLX-EP (Perplexity pplx-kernels NVSHMEM all-to-all) SGLANG_PPLX_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index f6a445b428b0..28713060a982 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -25,6 +25,10 @@ DeepEPLLCombineInput, DeepEPNormalCombineInput, ) +from sglang.srt.layers.moe.token_dispatcher.moonep import ( + get_moonep_expert_weight_layout, + run_moonep_bf16_expert, +) from sglang.srt.layers.moe.topk import ( StandardTopKOutput, TopKOutput, @@ -291,6 +295,26 @@ def run_moe_core( output = self.forward_cutlass_w4afp8_masked(dispatch_output) else: assert False, "forward_deepgemm_masked is deprecated" + elif DispatchOutputChecker.format_is_moonep(dispatch_output): + weight_layout = get_moonep_expert_weight_layout( + self, + num_prefetch_slots=( + int(dispatch_output.expert_ids.numel()) - self.num_experts + ), + ) + self.dispatcher.prefetch_weight( + plan=dispatch_output.plan, + weight_layout=weight_layout, + ) + return run_moonep_bf16_expert( + dispatch_output, + weight_layout, + activation=self.moe_runner_config.activation, + ) + else: + raise NotImplementedError( + f"Unsupported DeepEPMoE dispatch format: {dispatch_output.format}" + ) combine_input_wrapper = ( DeepEPNormalCombineInput @@ -346,6 +370,7 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]): if ( get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_moonep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_pplx() diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 006f0db17b45..d41172aef15d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -121,6 +121,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: return StandardDispatcher(moe_runner_config) elif ( a2a_backend.is_deepep() + or a2a_backend.is_moonep() or a2a_backend.is_mooncake() or a2a_backend.is_mori() or a2a_backend.is_nixl() @@ -363,9 +364,24 @@ def __init__( f"quant_method={type(self.quant_method).__name__})." ) + moonep_global_weight_storage = get_moe_a2a_backend().is_moonep() + if moonep_global_weight_storage: + if quant_config is not None: + raise NotImplementedError( + "MoonEP PoC supports unquantized BF16 MoE weights only." + ) + if num_fused_shared_experts != 0: + raise NotImplementedError( + "MoonEP PoC does not support fused shared experts yet." + ) + self.quant_method.create_weights( layer=self, - num_experts=self.num_local_experts, + num_experts=( + self.num_experts + if moonep_global_weight_storage + else self.num_local_experts + ), hidden_size=hidden_size, intermediate_size_per_partition=self.intermediate_size_per_partition, params_dtype=params_dtype, @@ -377,6 +393,13 @@ def __init__( with_bias=with_bias, moe_intermediate_size=intermediate_size, ) + if moonep_global_weight_storage: + self.w13_weight._sglang_require_global_experts = True + self.w2_weight._sglang_require_global_experts = True + if getattr(self, "w13_weight_bias", None) is not None: + self.w13_weight_bias._sglang_require_global_experts = True + if getattr(self, "w2_weight_bias", None) is not None: + self.w2_weight_bias._sglang_require_global_experts = True self.quant_method.create_moe_runner(self, self.moe_runner_config) self.dispatcher = create_moe_dispatcher(self.moe_runner_config) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py index 7f2c0942f95c..9c65901e685b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py @@ -30,6 +30,17 @@ MooncakeDispatchOutput, MooncakeEPDispatcher, ) +from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPBufferKey, + MoonEPCombineInput, + MoonEPDispatcher, + MoonEPDispatchOutput, + MoonEPExpertWeightLayout, + get_moonep_expert_weight_layout, + get_moonep_num_prefetch_slots, + run_moonep_bf16_expert, +) from sglang.srt.layers.moe.token_dispatcher.moriep import ( MoriEPDispatcher, MoriEPLLCombineInput, @@ -67,6 +78,15 @@ "MooncakeCombineInput", "MooncakeDispatchOutput", "MooncakeEPDispatcher", + "MoonEPCombineInput", + "MoonEPBuffer", + "MoonEPBufferKey", + "MoonEPDispatcher", + "MoonEPDispatchOutput", + "MoonEPExpertWeightLayout", + "get_moonep_expert_weight_layout", + "get_moonep_num_prefetch_slots", + "run_moonep_bf16_expert", "MoriEPNormalDispatchOutput", "MoriEPNormalCombineInput", "MoriEPLLDispatchOutput", diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index 1ff2beb5bb33..cad75fba7f16 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -29,6 +29,8 @@ DeepEPNormalDispatchOutput, FlashinferCombineInput, FlashinferDispatchOutput, + MoonEPCombineInput, + MoonEPDispatchOutput, StandardCombineInput, StandardDispatchOutput, ) @@ -165,12 +167,19 @@ def format_is_flashinfer( ) -> TypeGuard[FlashinferDispatchOutput]: return dispatch_output.format.is_flashinfer() + @staticmethod + def format_is_moonep( + dispatch_output: DispatchOutput, + ) -> TypeGuard[MoonEPDispatchOutput]: + return dispatch_output.format.is_moonep() + class DispatchOutputFormat(Enum): STANDARD = "standard" DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" + MOONEP = "moonep" FLASHINFER = "flashinfer" ASCEND_TP = "ascend_tp" @@ -195,6 +204,9 @@ def is_deepep(self) -> bool: def is_flashinfer(self) -> bool: return self == DispatchOutputFormat.FLASHINFER + def is_moonep(self) -> bool: + return self == DispatchOutputFormat.MOONEP + @runtime_checkable class DispatchOutput(Protocol): @@ -249,11 +261,18 @@ def format_is_flashinfer( ) -> TypeGuard[FlashinferCombineInput]: return combine_input.format == CombineInputFormat.FLASHINFER + @staticmethod + def format_is_moonep( + combine_input: CombineInput, + ) -> TypeGuard[MoonEPCombineInput]: + return combine_input.format == CombineInputFormat.MOONEP + class CombineInputFormat(Enum): STANDARD = "standard" DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" + MOONEP = "moonep" FLASHINFER = "flashinfer" ASCEND_TP = "ascend_tp" diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moonep.py b/python/sglang/srt/layers/moe/token_dispatcher/moonep.py new file mode 100644 index 000000000000..55ab3be63268 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/moonep.py @@ -0,0 +1,708 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, NamedTuple, NoReturn, Optional + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from sglang.srt.environ import envs +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.moe.topk import TopKOutputChecker +from sglang.srt.layers.moe.utils import DeepEPMode + + +_MOONEP_UNSUPPORTED_MESSAGE = ( + "MoonEP MoE A2A is recognized by SGLang, but the runtime dispatcher is not " + "implemented yet. MoonEP is not a drop-in DeepEP-compatible backend: it " + "returns a MoonEPCommPlan/cu_seqlens and requires MoonEP-compatible " + "contiguous symmetric-memory expert weights plus a VM-group expert GEMM." +) + + +class MoonEPDispatchOutput(NamedTuple): + """MoonEP dispatch output. + + ``plan`` is intentionally typed as ``Any`` so this module can define the + SGLang-side contract without importing the optional ``moonep`` package at + module import time. + """ + + hidden_states: torch.Tensor + route_weights_nvs: Optional[torch.Tensor] + cu_seqlens: torch.Tensor + plan: Any + expert_ids: torch.Tensor + num_tokens: int + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.MOONEP + + +assert isinstance(MoonEPDispatchOutput, DispatchOutput) + + +class MoonEPCombineInput(NamedTuple): + """MoonEP combine input.""" + + hidden_states: torch.Tensor + route_weights_nvs: Optional[torch.Tensor] + plan: Any + num_tokens: int + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.MOONEP + + +assert isinstance(MoonEPCombineInput, CombineInput) + + +class MoonEPExpertWeightLayout(NamedTuple): + """Contiguous BF16 expert weights in MoonEP prefetch layout.""" + + full_gate_weight: torch.Tensor + full_up_weight: torch.Tensor + full_down_weight: torch.Tensor + num_prefetch_slots: int + + +@dataclass(frozen=True) +class MoonEPBufferKey: + """Static MoonEP buffer dimensions. + + MoonEP allocates its communication buffers from static shape parameters, + unlike DeepEP's normal-dispatch path. Keep the dimensions explicit so the + process-wide facade never reuses a buffer with incompatible token capacity, + model shape, EP topology, or prefetch-slot layout. + """ + + num_max_dispatch_tokens_per_rank: int + hidden_size: int + router_topk: int + num_experts: int + num_ep_ranks: int + group_id: int + num_prefetch_slots: int + token_padding: int + num_sms: int + + +class MoonEPBuffer: + """Process-wide facade for MoonEP communication buffers. + + The underlying ``moonep.Buffer`` owns NVLink/VMM allocations and is keyed by + MoonEP's static allocation dimensions. The state lives on + ``ctx.resources.buffers`` so tests can reset it with ``reset_context()`` and + future runtime code has one lifecycle hook per process. + """ + + @classmethod + def _state(cls): + from types import SimpleNamespace + + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + state = buffers.get("moonep_ep_state") + if state is None: + state = SimpleNamespace( + buffers={}, + active_key=None, + ) + buffers["moonep_ep_state"] = state + return state + + @staticmethod + def _require_positive_int(name: str, value: int) -> int: + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + return value + + @staticmethod + def _resolve_num_ep_ranks(group: dist.ProcessGroup) -> int: + try: + num_ep_ranks = dist.get_world_size(group=group) + except (AssertionError, RuntimeError, TypeError, ValueError): + group_size = getattr(group, "size", None) + if not callable(group_size): + raise + num_ep_ranks = group_size() + return MoonEPBuffer._require_positive_int("num_ep_ranks", int(num_ep_ranks)) + + @staticmethod + def _resolve_num_prefetch_slots( + num_prefetch_slots: int | None, + num_experts: int, + num_ep_ranks: int, + ) -> int: + if num_experts % num_ep_ranks != 0: + raise ValueError( + "MoonEP requires num_experts to be divisible by the EP group size: " + f"num_experts={num_experts}, num_ep_ranks={num_ep_ranks}" + ) + + if num_prefetch_slots is None: + num_prefetch_slots = envs.SGLANG_MOONEP_NUM_PREFETCH_SLOTS.get() + num_prefetch_slots = int(num_prefetch_slots) + if num_prefetch_slots <= 0: + return num_experts // num_ep_ranks + return MoonEPBuffer._require_positive_int( + "num_prefetch_slots", num_prefetch_slots + ) + + @classmethod + def build_key( + cls, + group: dist.ProcessGroup, + hidden_size: int, + router_topk: int, + num_experts: int, + num_max_dispatch_tokens_per_rank: int | None = None, + num_prefetch_slots: int | None = None, + token_padding: int | None = None, + num_sms: int | None = None, + ) -> MoonEPBufferKey: + if num_max_dispatch_tokens_per_rank is None: + num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + if token_padding is None: + token_padding = envs.SGLANG_MOONEP_TOKEN_PADDING.get() + if num_sms is None: + num_sms = envs.SGLANG_MOONEP_NUM_SMS.get() + + num_ep_ranks = cls._resolve_num_ep_ranks(group) + num_experts = cls._require_positive_int("num_experts", int(num_experts)) + num_prefetch_slots = cls._resolve_num_prefetch_slots( + num_prefetch_slots, + num_experts, + num_ep_ranks, + ) + + return MoonEPBufferKey( + num_max_dispatch_tokens_per_rank=cls._require_positive_int( + "num_max_dispatch_tokens_per_rank", + int(num_max_dispatch_tokens_per_rank), + ), + hidden_size=cls._require_positive_int("hidden_size", int(hidden_size)), + router_topk=cls._require_positive_int("router_topk", int(router_topk)), + num_experts=num_experts, + num_ep_ranks=num_ep_ranks, + group_id=id(group), + num_prefetch_slots=num_prefetch_slots, + token_padding=cls._require_positive_int( + "token_padding", int(token_padding) + ), + num_sms=cls._require_positive_int("num_sms", int(num_sms)), + ) + + @classmethod + def get_existing_buffer( + cls, + key: MoonEPBufferKey | None = None, + ): + """Return an already-created buffer, if any. + + Without a key this returns the most recently requested MoonEP buffer. + Future runtime code should pass an explicit key when multiple static + capacities are in play. + """ + + state = cls._state() + if key is None: + key = state.active_key + if key is None: + return None + return state.buffers.get(key) + + @classmethod + def get_moonep_buffer( + cls, + group: dist.ProcessGroup, + hidden_size: int, + router_topk: int, + num_experts: int, + num_max_dispatch_tokens_per_rank: int | None = None, + num_prefetch_slots: int | None = None, + token_padding: int | None = None, + num_sms: int | None = None, + ): + key = cls.build_key( + group=group, + hidden_size=hidden_size, + router_topk=router_topk, + num_experts=num_experts, + num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, + num_prefetch_slots=num_prefetch_slots, + token_padding=token_padding, + num_sms=num_sms, + ) + + state = cls._state() + buffer = state.buffers.get(key) + if buffer is not None: + state.active_key = key + return buffer + + try: + from moonep import Buffer + except ImportError as exc: + raise ImportError( + "MoonEP is not installed. Install MoonEP before running SGLang " + "with --moe-a2a-backend moonep." + ) from exc + + buffer = Buffer( + S=key.num_max_dispatch_tokens_per_rank, + H=key.hidden_size, + K=key.router_topk, + E=key.num_experts, + num_ep_ranks=key.num_ep_ranks, + num_sms=key.num_sms, + token_padding=key.token_padding, + B=key.num_prefetch_slots, + group=group, + ) + state.buffers[key] = buffer + state.active_key = key + return buffer + + @classmethod + def destroy_buffer(cls, key: MoonEPBufferKey | None = None) -> None: + state = cls._state() + if key is None: + key = state.active_key + if key is None: + return + + buffer = state.buffers.pop(key, None) + destroy = getattr(buffer, "destroy", None) + if callable(destroy): + destroy() + if state.active_key == key: + state.active_key = next(reversed(state.buffers), None) + + @classmethod + def destroy_all_buffers(cls) -> None: + state = cls._state() + for key in list(state.buffers): + cls.destroy_buffer(key) + state.active_key = None + + +def get_moonep_num_prefetch_slots(num_experts: int, num_ep_ranks: int) -> int: + return MoonEPBuffer._resolve_num_prefetch_slots( + num_prefetch_slots=None, + num_experts=num_experts, + num_ep_ranks=num_ep_ranks, + ) + + +def get_moonep_expert_weight_layout( + layer: torch.nn.Module, + num_prefetch_slots: int, +) -> MoonEPExpertWeightLayout: + """Return cached contiguous BF16 gate/up/down tensors for MoonEP. + + The first executable PoC path supports only unquantized BF16 weights stored + in global expert-id order. Rows ``[E, E+B)`` are mutable prefetch slots and + are intentionally preserved across calls so ``buffer.prefetch_weight`` can + fill them before the MoonEP expert runner consumes the layout. + """ + + if num_prefetch_slots <= 0: + raise ValueError( + f"num_prefetch_slots must be positive, got {num_prefetch_slots}" + ) + + quant_config = getattr(layer, "quant_config", None) + if quant_config is not None: + raise NotImplementedError( + "MoonEP PoC expert weight layout supports unquantized BF16 only." + ) + + if getattr(layer.moe_runner_config, "num_fused_shared_experts", 0) != 0: + raise NotImplementedError( + "MoonEP PoC does not support fused shared experts yet." + ) + if not getattr(layer.moe_runner_config, "is_gated", True): + raise NotImplementedError("MoonEP PoC requires gated w13 experts.") + if getattr(layer, "use_triton_kernels", False): + raise NotImplementedError( + "MoonEP PoC expects canonical [E, 2I, H] w13 layout, not " + "triton-kernels transposed weight layout." + ) + + w13_weight = layer.w13_weight + w2_weight = layer.w2_weight + if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16: + raise NotImplementedError( + "MoonEP PoC expert runner supports BF16 weights only." + ) + if not w13_weight.is_contiguous() or not w2_weight.is_contiguous(): + raise ValueError("MoonEP expert source weights must be contiguous.") + + num_experts = int(layer.num_experts) + intermediate_size = int(layer.intermediate_size_per_partition) + hidden_size = int(layer.hidden_size) + expected_w13_shape = (num_experts, 2 * intermediate_size, hidden_size) + expected_w2_shape = (num_experts, hidden_size, intermediate_size) + if tuple(w13_weight.shape) != expected_w13_shape: + raise ValueError( + "MoonEP PoC requires global w13_weight shape " + f"{expected_w13_shape}, got {tuple(w13_weight.shape)}." + ) + if tuple(w2_weight.shape) != expected_w2_shape: + raise ValueError( + "MoonEP PoC requires global w2_weight shape " + f"{expected_w2_shape}, got {tuple(w2_weight.shape)}." + ) + + cache_key = ( + num_prefetch_slots, + w13_weight.data_ptr(), + w2_weight.data_ptr(), + tuple(w13_weight.shape), + tuple(w2_weight.shape), + ) + cache = getattr(layer, "_moonep_weight_layout_cache", None) + if cache is not None and cache[0] == cache_key: + return cache[1] + + full_gate_weight = torch.empty( + num_experts + num_prefetch_slots, + intermediate_size, + hidden_size, + dtype=torch.bfloat16, + device=w13_weight.device, + ) + full_up_weight = torch.empty_like(full_gate_weight) + full_down_weight = torch.empty( + num_experts + num_prefetch_slots, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + device=w2_weight.device, + ) + + full_gate_weight[:num_experts].copy_(w13_weight[:, :intermediate_size, :]) + full_up_weight[:num_experts].copy_( + w13_weight[:, intermediate_size : 2 * intermediate_size, :] + ) + full_down_weight[:num_experts].copy_(w2_weight) + full_gate_weight[num_experts:].zero_() + full_up_weight[num_experts:].zero_() + full_down_weight[num_experts:].zero_() + + layout = MoonEPExpertWeightLayout( + full_gate_weight=full_gate_weight.contiguous(), + full_up_weight=full_up_weight.contiguous(), + full_down_weight=full_down_weight.contiguous(), + num_prefetch_slots=num_prefetch_slots, + ) + layer._moonep_weight_layout_cache = (cache_key, layout) + return layout + + +def run_moonep_bf16_expert( + dispatch_output: MoonEPDispatchOutput, + weight_layout: MoonEPExpertWeightLayout, + *, + activation: str = "silu", +) -> MoonEPCombineInput: + """Run a simple BF16 MoonEP expert core over ``cu_seqlens`` segments. + + This is a correctness-first PoC runner. It consumes MoonEP's already + expert-grouped ``[NvS, H]`` token layout, applies gate/up/down expert + weights for each non-empty `cu_seqlens` segment, multiplies each dispatched + row by its route weight, and returns a `MoonEPCombineInput` for + `MoonEPDispatcher.combine`. + """ + + if activation != "silu": + raise NotImplementedError( + f"MoonEP BF16 PoC runner only supports silu, got {activation!r}" + ) + + hidden_states = dispatch_output.hidden_states + route_weights_nvs = dispatch_output.route_weights_nvs + cu_seqlens = dispatch_output.cu_seqlens + expert_ids = dispatch_output.expert_ids + + if hidden_states.ndim != 2: + raise ValueError( + f"MoonEP hidden states must be [NvS, H], got {hidden_states.shape}" + ) + if cu_seqlens.ndim != 1: + raise ValueError(f"cu_seqlens must be 1D, got {cu_seqlens.shape}") + if expert_ids.shape != cu_seqlens.shape: + raise ValueError( + f"expert_ids shape {expert_ids.shape} must match cu_seqlens " + f"shape {cu_seqlens.shape}" + ) + if route_weights_nvs is not None and route_weights_nvs.ndim != 1: + raise ValueError( + f"route_weights_nvs must be 1D, got {route_weights_nvs.shape}" + ) + + output = torch.empty_like(hidden_states) + prev = 0 + for group_id in range(cu_seqlens.numel()): + cur = int(cu_seqlens[group_id].item()) + if cur < prev: + raise ValueError("MoonEP cu_seqlens must be non-decreasing") + if cur == prev: + continue + + expert_id = int(expert_ids[group_id].item()) + if expert_id < 0: + output[prev:cur].zero_() + prev = cur + continue + if expert_id >= weight_layout.full_gate_weight.shape[0]: + raise ValueError( + f"expert_id {expert_id} exceeds MoonEP weight rows " + f"{weight_layout.full_gate_weight.shape[0]}" + ) + + x = hidden_states[prev:cur] + gate = F.linear(x, weight_layout.full_gate_weight[expert_id]) + up = F.linear(x, weight_layout.full_up_weight[expert_id]) + activated = F.silu(gate) * up + y = F.linear(activated, weight_layout.full_down_weight[expert_id]) + if route_weights_nvs is not None: + y = y * route_weights_nvs[prev:cur].to(dtype=y.dtype).unsqueeze(-1) + output[prev:cur].copy_(y) + prev = cur + + if prev < hidden_states.shape[0]: + output[prev:].zero_() + + return MoonEPCombineInput( + hidden_states=output, + route_weights_nvs=route_weights_nvs, + plan=dispatch_output.plan, + num_tokens=dispatch_output.num_tokens, + ) + + +class MoonEPDispatcher(BaseDispatcher): + """MoonEP dispatcher for the initial BF16 inference PoC.""" + + def __init__( + self, + group: torch.distributed.ProcessGroup, + router_topk: int, + permute_fusion: bool = False, + num_experts: int | None = None, + num_local_experts: int | None = None, + hidden_size: int | None = None, + params_dtype: torch.dtype | None = None, + deepep_mode: DeepEPMode = DeepEPMode.AUTO, + async_finish: bool = False, + return_recv_hook: bool = False, + ): + super().__init__() + self.group = group + self.router_topk = router_topk + self.permute_fusion = permute_fusion + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.params_dtype = params_dtype + self.deepep_mode = deepep_mode + self.async_finish = async_finish + self.return_recv_hook = return_recv_hook + self.expert_mask_gpu = None + self.num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + self.num_prefetch_slots = None + + @staticmethod + def _raise_unimplemented() -> NoReturn: + raise NotImplementedError(_MOONEP_UNSUPPORTED_MESSAGE) + + def _get_buffer(self): + if self.hidden_size is None or self.num_experts is None: + raise ValueError( + "MoonEPDispatcher requires hidden_size and num_experts to " + "create a MoonEP buffer." + ) + return MoonEPBuffer.get_moonep_buffer( + group=self.group, + hidden_size=self.hidden_size, + router_topk=self.router_topk, + num_experts=self.num_experts, + num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank, + num_prefetch_slots=self.num_prefetch_slots, + ) + + def _get_rank(self) -> int: + try: + return dist.get_rank(group=self.group) + except (AssertionError, RuntimeError, TypeError, ValueError): + return 0 + + def _pad_to_capacity( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + num_tokens = int(hidden_states.shape[0]) + capacity = int(self.num_max_dispatch_tokens_per_rank) + if num_tokens > capacity: + raise ValueError( + "MoonEP runtime batch has more tokens than its static buffer " + f"capacity: num_tokens={num_tokens}, capacity={capacity}. " + "Increase SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK." + ) + + hidden_states = hidden_states.contiguous() + topk_ids = topk_ids.to(dtype=torch.int32).contiguous() + topk_weights = topk_weights.to(dtype=torch.float32).contiguous() + if num_tokens == capacity: + return hidden_states, topk_ids, topk_weights, num_tokens + + pad_tokens = capacity - num_tokens + hidden_pad = hidden_states.new_zeros(pad_tokens, hidden_states.shape[1]) + id_pad = topk_ids.new_zeros(pad_tokens, topk_ids.shape[1]) + weight_pad = topk_weights.new_zeros(pad_tokens, topk_weights.shape[1]) + return ( + torch.cat((hidden_states, hidden_pad), dim=0).contiguous(), + torch.cat((topk_ids, id_pad), dim=0).contiguous(), + torch.cat((topk_weights, weight_pad), dim=0).contiguous(), + num_tokens, + ) + + def _tokens_per_expert(self, topk_ids: torch.Tensor) -> torch.Tensor: + assert self.num_experts is not None + return torch.bincount( + topk_ids.reshape(-1).to(dtype=torch.int64), + minlength=self.num_experts, + ).to(dtype=torch.int32) + + def _expert_ids_from_plan( + self, + cu_seqlens: torch.Tensor, + plan: Any, + ) -> torch.Tensor: + assert self.num_experts is not None + num_groups = int(cu_seqlens.numel()) + expert_ids = torch.full_like(cu_seqlens, -1) + experts_to_copy = plan.experts_to_copy + if experts_to_copy.ndim == 2: + experts_to_copy = experts_to_copy[self._get_rank()] + + prev = 0 + for group_id in range(num_groups): + cur = int(cu_seqlens[group_id].item()) + if cur > prev: + if group_id < self.num_experts: + expert_ids[group_id] = group_id + else: + expert_ids[group_id] = experts_to_copy[group_id - self.num_experts] + prev = cur + return expert_ids + + def dispatch( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ) -> DispatchOutput: + if not TopKOutputChecker.format_is_standard(topk_output): + raise NotImplementedError( + "MoonEP PoC requires standard top-k output before dispatch." + ) + if hidden_states.dtype != torch.bfloat16: + raise NotImplementedError( + "MoonEP PoC dispatcher supports BF16 hidden states only." + ) + if self.num_experts is None: + raise ValueError("MoonEPDispatcher requires num_experts.") + + hidden_states, topk_ids, topk_weights, num_tokens = self._pad_to_capacity( + hidden_states, + topk_output.topk_ids, + topk_output.topk_weights, + ) + tokens_per_expert = self._tokens_per_expert(topk_ids) + buffer = self._get_buffer() + hidden_nvsh, route_weights_nvs, cu_seqlens, plan = buffer.dispatch( + hidden_states, + topk_weights, + topk_ids, + tokens_per_expert, + async_finish=False, + ) + return MoonEPDispatchOutput( + hidden_states=hidden_nvsh, + route_weights_nvs=route_weights_nvs, + cu_seqlens=cu_seqlens, + plan=plan, + expert_ids=self._expert_ids_from_plan(cu_seqlens, plan), + num_tokens=num_tokens, + ) + + def dispatch_a( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + ): + self._raise_unimplemented() + + def dispatch_b(self): + self._raise_unimplemented() + + def combine( + self, + combine_input: CombineInput, + ) -> torch.Tensor: + if combine_input.format != CombineInputFormat.MOONEP: + raise TypeError( + f"MoonEPDispatcher.combine expected MOONEP input, got " + f"{combine_input.format}" + ) + hidden_states, _route_weights_sk, _event = self._get_buffer().combine( + plan=combine_input.plan, + hidden_nvsh=combine_input.hidden_states, + route_weights_nvs=None, + async_finish=False, + ) + return hidden_states[: combine_input.num_tokens].contiguous() + + def combine_a( + self, + combine_input: CombineInput, + ): + self._raise_unimplemented() + + def combine_b(self): + self._raise_unimplemented() + + def prefetch_weight( + self, + plan: Any, + weight_layout: MoonEPExpertWeightLayout, + ) -> None: + self._get_buffer().prefetch_weight( + plan=plan, + async_finish=False, + full_gate_weight=weight_layout.full_gate_weight, + full_up_weight=weight_layout.full_up_weight, + full_down_weight=weight_layout.full_down_weight, + ) + + def register_deepep_dispatch_hook(self, hook): + self._raise_unimplemented() diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 9ecaec671404..ffa40cce6ea0 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -29,6 +29,7 @@ class MoeA2ABackend(Enum): NONE = "none" DEEPEP = "deepep" + MOONEP = "moonep" MOONCAKE = "mooncake" NIXL = "nixl" MORI = "mori" @@ -54,6 +55,9 @@ def is_none(self): def is_deepep(self): return self == MoeA2ABackend.DEEPEP + def is_moonep(self): + return self == MoeA2ABackend.MOONEP + def is_mooncake(self): return self == MoeA2ABackend.MOONCAKE @@ -388,9 +392,15 @@ def is_sbo_enabled() -> bool: def is_deepep_class_backend() -> bool: - """Check if the MoE backend is DeepEP-family (DeepEP, Mooncake, Mori, or PPLX).""" + """Check if the MoE backend is DeepEP-family.""" b = get_moe_a2a_backend() - return b.is_deepep() or b.is_mooncake() or b.is_mori() or b.is_pplx() + return ( + b.is_deepep() + or b.is_moonep() + or b.is_mooncake() + or b.is_mori() + or b.is_pplx() + ) def uses_per_rank_fused_shared_slots() -> bool: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c297e8e33ddb..6b76f6713603 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -267,6 +267,7 @@ MOE_A2A_BACKEND_CHOICES = [ "none", "deepep", + "moonep", "mooncake", "nixl", "mori", @@ -2249,6 +2250,7 @@ class ServerArgs: Literal[ "none", "deepep", + "moonep", "mooncake", "nixl", "mori", @@ -6622,6 +6624,15 @@ def _handle_a2a_moe(self): f"DeepEP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." ) + if a2a_backend == "moonep": + logger.warning( + "MoonEP MoE is enabled in experimental BF16 PoC mode. " + "Cuda graph is disabled while the eager MoonEP dispatch/" + "prefetch/compute/combine path is validated." + ) + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED + if a2a_backend == "mooncake": logger.warning( f"Mooncake MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." diff --git a/scripts/moonep/validate_moonep_bf16_poc.py b/scripts/moonep/validate_moonep_bf16_poc.py new file mode 100755 index 000000000000..35f054f851a1 --- /dev/null +++ b/scripts/moonep/validate_moonep_bf16_poc.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +"""Distributed BF16 MoonEP PoC validation. + +Run with torchrun on a single NVLink/NVSwitch node, for example: + + SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128 \ + torchrun --standalone --nproc-per-node=4 \ + scripts/moonep/validate_moonep_bf16_poc.py --tokens 128 --hidden-size 1024 + +The script validates SGLang's MoonEP BF16 path: +MoonEPDispatcher.dispatch -> MoonEPBuffer.prefetch_weight -> BF16 segment runner +-> MoonEPDispatcher.combine. It compares the final token-major output against a +rank-local PyTorch reference using the same top-k and expert weights. +""" + +from __future__ import annotations + +import argparse +from enum import IntEnum, auto +import json +import os +import sys +from pathlib import Path +from types import ModuleType +from typing import NamedTuple, Protocol, TypeGuard, runtime_checkable + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, default=128) + parser.add_argument("--hidden-size", type=int, default=1024) + parser.add_argument("--intermediate-size", type=int, default=128) + parser.add_argument("--top-k", type=int, default=2) + parser.add_argument("--experts-per-rank", type=int, default=2) + parser.add_argument("--prefetch-slots", type=int, default=-1) + parser.add_argument("--seed", type=int, default=20260802) + parser.add_argument("--atol", type=float, default=5e-2) + parser.add_argument("--rtol", type=float, default=5e-2) + return parser.parse_args() + + +def setup_dist() -> tuple[int, int, int]: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + return rank, world_size, local_rank + + +def make_topk(tokens: int, top_k: int, num_experts: int, device: torch.device): + # Deterministic but rank-local routing. Keep weights normalized so the + # reference magnitude remains bounded. + topk_ids = torch.randint( + 0, + num_experts, + (tokens, top_k), + device=device, + dtype=torch.int64, + ) + raw_weights = torch.rand(tokens, top_k, device=device, dtype=torch.float32) + topk_weights = raw_weights / raw_weights.sum(dim=-1, keepdim=True) + return topk_ids, topk_weights + + +class LightweightTopKOutputFormat(IntEnum): + STANDARD = auto() + + +@runtime_checkable +class LightweightTopKOutput(Protocol): + @property + def format(self) -> LightweightTopKOutputFormat: ... + + +class LightweightStandardTopKOutput(NamedTuple): + topk_weights: torch.Tensor + topk_ids: torch.Tensor + router_logits: torch.Tensor + + @property + def format(self) -> LightweightTopKOutputFormat: + return LightweightTopKOutputFormat.STANDARD + + +class LightweightTopKOutputChecker: + @staticmethod + def format_is_standard( + topk_output: LightweightTopKOutput, + ) -> TypeGuard[LightweightStandardTopKOutput]: + return isinstance(topk_output, LightweightStandardTopKOutput) + + +class LightweightDeepEPMode(IntEnum): + NORMAL = 1 + LOW_LATENCY = 2 + AUTO = 3 + + +def install_lightweight_sglang_imports() -> None: + """Install validation-only SGLang package stubs. + + Full SGLang environments should use the normal imports. Minimal remote + validation images often do not carry every frontend/model dependency that + SGLang's package ``__init__`` imports, though this script only needs the + MoonEP dispatcher, its base protocol, envs, and runtime buffer registry. + These stubs keep the package search path pointed at the checked-out source + tree while bypassing heavyweight parent ``__init__`` files. + """ + + repo_root = Path(__file__).resolve().parents[2] + sglang_root = repo_root / "python" / "sglang" + + for name in list(sys.modules): + if name == "sglang" or name.startswith("sglang."): + del sys.modules[name] + + def add_package(name: str, path: Path) -> None: + module = ModuleType(name) + module.__path__ = [str(path)] + sys.modules[name] = module + + add_package("sglang", sglang_root) + add_package("sglang.srt", sglang_root / "srt") + add_package("sglang.srt.layers", sglang_root / "srt" / "layers") + add_package("sglang.srt.layers.moe", sglang_root / "srt" / "layers" / "moe") + add_package( + "sglang.srt.layers.moe.token_dispatcher", + sglang_root / "srt" / "layers" / "moe" / "token_dispatcher", + ) + + topk_module = ModuleType("sglang.srt.layers.moe.topk") + topk_module.TopKOutputFormat = LightweightTopKOutputFormat + topk_module.TopKOutput = LightweightTopKOutput + topk_module.StandardTopKOutput = LightweightStandardTopKOutput + topk_module.TopKOutputChecker = LightweightTopKOutputChecker + sys.modules[topk_module.__name__] = topk_module + + utils_module = ModuleType("sglang.srt.layers.moe.utils") + utils_module.DeepEPMode = LightweightDeepEPMode + sys.modules[utils_module.__name__] = utils_module + + +def import_moonep_validation_symbols(): + try: + from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + ) + from sglang.srt.layers.moe.topk import StandardTopKOutput + + return ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + StandardTopKOutput, + ) + except ModuleNotFoundError as exc: + print( + "Normal SGLang import failed; retrying with lightweight validation " + f"imports: {exc}", + file=sys.stderr, + flush=True, + ) + install_lightweight_sglang_imports() + from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + ) + from sglang.srt.layers.moe.topk import StandardTopKOutput + + return ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + StandardTopKOutput, + ) + + +def expert_mlp(x, expert_id: int, gate, up, down): + return F.linear( + F.silu(F.linear(x, gate[expert_id])) * F.linear(x, up[expert_id]), + down[expert_id], + ) + + +def reference_output(hidden, topk_ids, topk_weights, gate, up, down): + out = torch.zeros_like(hidden) + tokens, top_k = topk_ids.shape + for token_idx in range(tokens): + x = hidden[token_idx : token_idx + 1] + acc = torch.zeros_like(x) + for k in range(top_k): + expert_id = int(topk_ids[token_idx, k].item()) + acc += expert_mlp(x, expert_id, gate, up, down) * topk_weights[ + token_idx, k + ].to(hidden.dtype) + out[token_idx] = acc[0] + return out + + +def main() -> None: + args = parse_args() + os.environ["SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = str(args.tokens) + if args.prefetch_slots > 0: + os.environ["SGLANG_MOONEP_NUM_PREFETCH_SLOTS"] = str(args.prefetch_slots) + + rank, world_size, local_rank = setup_dist() + device = torch.device(f"cuda:{local_rank}") + torch.manual_seed(args.seed + rank) + + ( + MoonEPBuffer, + MoonEPDispatcher, + MoonEPExpertWeightLayout, + run_moonep_bf16_expert, + StandardTopKOutput, + ) = import_moonep_validation_symbols() + + num_experts = world_size * args.experts_per_rank + hidden = torch.randn( + args.tokens, + args.hidden_size, + device=device, + dtype=torch.bfloat16, + ) + topk_ids, topk_weights = make_topk(args.tokens, args.top_k, num_experts, device) + topk_output = StandardTopKOutput( + topk_weights=topk_weights, + topk_ids=topk_ids, + router_logits=torch.empty(0, device=device), + ) + + dispatcher = MoonEPDispatcher( + group=dist.group.WORLD, + router_topk=args.top_k, + num_experts=num_experts, + num_local_experts=args.experts_per_rank, + hidden_size=args.hidden_size, + params_dtype=torch.bfloat16, + ) + + dispatch_output = dispatcher.dispatch(hidden, topk_output) + num_prefetch_slots = int(dispatch_output.expert_ids.numel()) - num_experts + + # Full global rows are deliberately replicated for this correctness PoC. + # The production path should replace this with true symmetric expert-row + # mappings whose physical storage is owned by each expert's home rank. + torch.manual_seed(args.seed) + gate = torch.randn( + num_experts + num_prefetch_slots, + args.intermediate_size, + args.hidden_size, + device=device, + dtype=torch.bfloat16, + ) / 8 + up = torch.randn_like(gate) / 8 + down = torch.randn( + num_experts + num_prefetch_slots, + args.hidden_size, + args.intermediate_size, + device=device, + dtype=torch.bfloat16, + ) / 8 + gate[num_experts:].zero_() + up[num_experts:].zero_() + down[num_experts:].zero_() + layout = MoonEPExpertWeightLayout( + gate.contiguous(), + up.contiguous(), + down.contiguous(), + num_prefetch_slots, + ) + + dispatcher.prefetch_weight(dispatch_output.plan, layout) + combine_input = run_moonep_bf16_expert(dispatch_output, layout) + output = dispatcher.combine(combine_input) + + expected = reference_output(hidden, topk_ids, topk_weights, gate, up, down) + max_abs_err = (output.float() - expected.float()).abs().max() + rel_err = max_abs_err / expected.float().abs().max().clamp_min(1e-6) + local_ok = bool( + torch.allclose(output.float(), expected.float(), atol=args.atol, rtol=args.rtol) + ) + ok_tensor = torch.tensor([1 if local_ok else 0], device=device, dtype=torch.int32) + dist.all_reduce(ok_tensor, op=dist.ReduceOp.MIN) + + result = { + "rank": rank, + "world_size": world_size, + "tokens": args.tokens, + "hidden_size": args.hidden_size, + "intermediate_size": args.intermediate_size, + "top_k": args.top_k, + "num_experts": num_experts, + "num_prefetch_slots": num_prefetch_slots, + "max_abs_err": float(max_abs_err.item()), + "relative_err": float(rel_err.item()), + "local_ok": local_ok, + "global_ok": bool(ok_tensor.item()), + } + print(json.dumps(result, sort_keys=True), flush=True) + dist.barrier(device_ids=[local_rank]) + MoonEPBuffer.destroy_all_buffers() + dist.destroy_process_group() + if rank == 0 and not bool(ok_tensor.item()): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/test/registered/unit/layers/moe/test_moonep_buffer.py b/test/registered/unit/layers/moe/test_moonep_buffer.py new file mode 100644 index 000000000000..dd06f5f9fbc9 --- /dev/null +++ b/test/registered/unit/layers/moe/test_moonep_buffer.py @@ -0,0 +1,306 @@ +import sys +import types +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.moe.token_dispatcher.moonep import ( + MoonEPBuffer, + MoonEPDispatchOutput, + MoonEPExpertWeightLayout, + get_moonep_expert_weight_layout, + run_moonep_bf16_expert, +) +from sglang.srt.runtime_context import reset_context +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _FakeMoonEPBuffer: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.destroy_calls = 0 + self.__class__.instances.append(self) + + def destroy(self): + self.destroy_calls += 1 + + +def _fake_moonep_module(): + module = types.ModuleType("moonep") + module.Buffer = _FakeMoonEPBuffer + return module + + +class TestMoonEPBuffer(unittest.TestCase): + def setUp(self): + reset_context() + _FakeMoonEPBuffer.instances.clear() + + def tearDown(self): + try: + MoonEPBuffer.destroy_all_buffers() + finally: + reset_context() + _FakeMoonEPBuffer.instances.clear() + + def test_lazily_constructs_and_reuses_buffer_for_static_key(self): + group = object() + + with ( + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=4, + ), + ): + buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=256, + num_prefetch_slots=16, + token_padding=64, + num_sms=20, + ) + same_buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=256, + num_prefetch_slots=16, + token_padding=64, + num_sms=20, + ) + larger_buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=512, + num_prefetch_slots=16, + token_padding=64, + num_sms=20, + ) + + self.assertIs(buffer, same_buffer) + self.assertIsNot(buffer, larger_buffer) + self.assertEqual(len(_FakeMoonEPBuffer.instances), 2) + self.assertEqual( + buffer.kwargs, + { + "S": 256, + "H": 1024, + "K": 8, + "E": 64, + "num_ep_ranks": 4, + "num_sms": 20, + "token_padding": 64, + "B": 16, + "group": group, + }, + ) + self.assertIs(MoonEPBuffer.get_existing_buffer(), larger_buffer) + + def test_resolves_env_defaults_and_training_safe_prefetch_slots(self): + group = object() + + with ( + envs.SGLANG_MOONEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(384), + envs.SGLANG_MOONEP_NUM_PREFETCH_SLOTS.override(-1), + envs.SGLANG_MOONEP_TOKEN_PADDING.override(32), + envs.SGLANG_MOONEP_NUM_SMS.override(18), + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=8, + ), + ): + buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=2048, + router_topk=6, + num_experts=128, + ) + + self.assertEqual(buffer.kwargs["S"], 384) + self.assertEqual(buffer.kwargs["token_padding"], 32) + self.assertEqual(buffer.kwargs["num_sms"], 18) + self.assertEqual(buffer.kwargs["B"], 16) + + def test_rejects_non_divisible_experts_before_allocating(self): + with ( + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=6, + ), + ): + with self.assertRaisesRegex(ValueError, "divisible"): + MoonEPBuffer.get_moonep_buffer( + group=object(), + hidden_size=1024, + router_topk=8, + num_experts=64, + ) + + self.assertEqual(_FakeMoonEPBuffer.instances, []) + + def test_destroy_all_buffers_releases_cached_buffers(self): + group = object() + + with ( + patch.dict(sys.modules, {"moonep": _fake_moonep_module()}), + patch( + "sglang.srt.layers.moe.token_dispatcher.moonep.dist.get_world_size", + return_value=4, + ), + ): + buffer = MoonEPBuffer.get_moonep_buffer( + group=group, + hidden_size=1024, + router_topk=8, + num_experts=64, + num_max_dispatch_tokens_per_rank=256, + ) + + MoonEPBuffer.destroy_all_buffers() + + self.assertEqual(buffer.destroy_calls, 1) + self.assertIsNone(MoonEPBuffer.get_existing_buffer()) + + +class TestMoonEPExpertWeightLayout(unittest.TestCase): + def _fake_layer(self): + num_experts, hidden_size, intermediate_size = 3, 4, 5 + w13_weight = torch.arange( + num_experts * 2 * intermediate_size * hidden_size, + dtype=torch.bfloat16, + ).reshape(num_experts, 2 * intermediate_size, hidden_size) + w2_weight = torch.arange( + num_experts * hidden_size * intermediate_size, + dtype=torch.bfloat16, + ).reshape(num_experts, hidden_size, intermediate_size) + + return SimpleNamespace( + quant_config=None, + moe_runner_config=SimpleNamespace( + num_fused_shared_experts=0, + is_gated=True, + ), + use_triton_kernels=False, + w13_weight=w13_weight, + w2_weight=w2_weight, + num_experts=num_experts, + intermediate_size_per_partition=intermediate_size, + hidden_size=hidden_size, + ) + + def test_layout_splits_gate_up_down_and_adds_prefetch_slots(self): + layer = self._fake_layer() + + layout = get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + + self.assertEqual(tuple(layout.full_gate_weight.shape), (5, 5, 4)) + self.assertEqual(tuple(layout.full_up_weight.shape), (5, 5, 4)) + self.assertEqual(tuple(layout.full_down_weight.shape), (5, 4, 5)) + torch.testing.assert_close( + layout.full_gate_weight[:3], + layer.w13_weight[:, :5, :], + ) + torch.testing.assert_close( + layout.full_up_weight[:3], + layer.w13_weight[:, 5:10, :], + ) + torch.testing.assert_close(layout.full_down_weight[:3], layer.w2_weight) + self.assertTrue(torch.all(layout.full_gate_weight[3:] == 0)) + self.assertTrue(torch.all(layout.full_up_weight[3:] == 0)) + self.assertTrue(torch.all(layout.full_down_weight[3:] == 0)) + + def test_layout_is_cached_for_same_weight_storage(self): + layer = self._fake_layer() + + first = get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + second = get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + + self.assertIs(first, second) + + def test_layout_rejects_local_expert_storage(self): + layer = self._fake_layer() + layer.w13_weight = layer.w13_weight[:2].contiguous() + + with self.assertRaisesRegex(ValueError, "global w13_weight"): + get_moonep_expert_weight_layout(layer, num_prefetch_slots=2) + + +class TestMoonEPBf16ExpertRunner(unittest.TestCase): + def test_segment_runner_applies_expert_weights_and_route_weights(self): + hidden_states = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + dtype=torch.float32, + ) + route_weights = torch.tensor([1.0, 0.5, 2.0], dtype=torch.float32) + cu_seqlens = torch.tensor([2, 3], dtype=torch.int32) + expert_ids = torch.tensor([0, 1], dtype=torch.int32) + gate = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.5, 0.0], [0.0, 0.5]], + ] + ) + up = torch.tensor( + [ + [[2.0, 0.0], [0.0, 2.0]], + [[1.5, 0.0], [0.0, 1.5]], + ] + ) + down = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[2.0, 0.0], [0.0, 2.0]], + ] + ) + layout = MoonEPExpertWeightLayout( + full_gate_weight=gate, + full_up_weight=up, + full_down_weight=down, + num_prefetch_slots=0, + ) + dispatch_output = MoonEPDispatchOutput( + hidden_states=hidden_states, + route_weights_nvs=route_weights, + cu_seqlens=cu_seqlens, + plan=object(), + expert_ids=expert_ids, + num_tokens=2, + ) + + combine_input = run_moonep_bf16_expert(dispatch_output, layout) + + expected = torch.empty_like(hidden_states) + for start, end, expert in [(0, 2, 0), (2, 3, 1)]: + x = hidden_states[start:end] + y = torch.nn.functional.linear( + torch.nn.functional.silu( + torch.nn.functional.linear(x, gate[expert]) + ) + * torch.nn.functional.linear(x, up[expert]), + down[expert], + ) + expected[start:end] = y * route_weights[start:end, None] + + torch.testing.assert_close(combine_input.hidden_states, expected) + self.assertIs(combine_input.plan, dispatch_output.plan) + self.assertEqual(combine_input.num_tokens, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 648ffe52c6c9..cfd5cc4d2967 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1209,6 +1209,23 @@ def test_waterfill_supports_deepep_low_latency_mode(self): self.assertTrue(server_args.enforce_shared_experts_fusion) +class TestMoonEPArgs(CustomTestCase): + def test_moonep_backend_is_recognized_and_disables_cuda_graph(self): + server_args = ServerArgs( + model_path="dummy", + moe_a2a_backend="moonep", + ) + + server_args._handle_a2a_moe() + + from sglang.srt.arg_groups.overrides import resolved_view + + self.assertEqual(resolved_view(server_args).moe_a2a_backend, "moonep") + self.assertEqual(resolved_view(server_args).ep_size, server_args.tp_size) + self.assertEqual(server_args.cuda_graph_config.decode.backend, Backend.DISABLED) + self.assertEqual(server_args.cuda_graph_config.prefill.backend, Backend.DISABLED) + + class TestPrefillOnlyDisableKvCache(unittest.TestCase): """Validation for --prefill-only-disable-kv-cache. diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index d3ac25300597..5e5de78d6cdf 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -428,6 +428,39 @@ def test_initialize_materializes_group(self): self.assertTrue(is_tbo_enabled()) self.assertEqual(get_flags().moe.deepep_config, "") + def test_initialize_materializes_moonep_backend(self): + from sglang.srt.layers.moe.utils import get_moe_a2a_backend + + self._init(moe_a2a_backend="moonep") + self.assertTrue(get_moe_a2a_backend().is_moonep()) + + def test_moonep_dispatch_contract_formats(self): + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInputChecker, + DispatchOutputChecker, + MoonEPCombineInput, + MoonEPDispatchOutput, + ) + + dispatch_output = MoonEPDispatchOutput( + hidden_states=object(), + route_weights_nvs=None, + cu_seqlens=object(), + plan=object(), + expert_ids=object(), + num_tokens=1, + ) + combine_input = MoonEPCombineInput( + hidden_states=object(), + route_weights_nvs=None, + plan=dispatch_output.plan, + num_tokens=1, + ) + + self.assertTrue(dispatch_output.format.is_moonep()) + self.assertTrue(DispatchOutputChecker.format_is_moonep(dispatch_output)) + self.assertTrue(CombineInputChecker.format_is_moonep(combine_input)) + def test_speculative_swap_and_restore(self): from sglang.srt.layers.moe.utils import ( get_moe_a2a_backend,