[MoonEP] BF16 PoC integration for Kimi-K3

.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.

.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

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`.

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/<context>/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…_

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 <number> --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 <number> --repo wirybeaver/sglang --body "..."`
+- **Apply / remove labels**: `gh issue edit <number> --repo wirybeaver/sglang --add-label "..."` / `--remove-label "..."`
+- **Close**: `gh issue close <number> --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 <number> --repo wirybeaver/sglang --comments` and `gh pr diff <number> --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 <number> --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 #<map>` at the top of the child body. Labels: `wayfinder:<type>` (`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/<child>/dependencies/blocked_by -F issue_id=<blocker-db-id>`, where `<blocker-db-id>` is the blocker's numeric **database id** (`gh api repos/wirybeaver/sglang/issues/<n> --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: #<n>, #<n>` 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 <n> --repo wirybeaver/sglang --add-assignee @me` — the session's first write.
+- **Resolve**: `gh issue comment <n> --repo wirybeaver/sglang --body "<answer>"`, then `gh issue close <n> --repo wirybeaver/sglang`, then append a context pointer (gist + link) to the map's Decisions-so-far.

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.