diff --git a/docs/index.md b/docs/index.md index d4c38c8ea..1940ae1a1 100644 --- a/docs/index.md +++ b/docs/index.md @@ -72,6 +72,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/glm52/ep4-gb300.md` | DP4/EP4 (`--moe-topo ep4`) on 4xGB300: 64 whole experts/rank (188 GiB weights/rank), second DeepEP shim instantiation, and a weight-only routed-expert chain (bf16×fp8 mma on the aligned recv layout — no act re-quant/relayout/remap) replacing the sm_90a DeepGEMM masked GEMMs. Bring-up green 2026-07-12 incl. layer-6 EP4 oracle 63/64 ×4 buckets; c=32 `559 tok/s` / TPOT p50 `39.05ms` untuned. Opt-in warm-cache pinned staging cuts full four-rank weight-load wall time `54.985→17.175s` and engine/HTTP-ready from `63.192→25.418s`; it stays off by default because a forced-cold network-filesystem A/B regressed `144.002→221.887s`. Needs `EP_DISABLE_GIN=1` on NIC-less nodes. | | `models/glm52/cross-node-scaling.md` | Control plane SHIPPED (`feat/glm52-rank-host`): framed-TCP hub-and-spoke — coordinator `--rank-hosts host:port=N`, dumb `--glm52-rank-host` process, FIFO `Response` frames, fail-stop + 60s destroy watchdog. On GB300 NVL72 one rack = one NVLink/IMEX domain, so the single-LSA DeepEP shim works cross-tray with no GIN un-baking: 2-tray EP8 solo p50 23.61 / p99 24.00 ms (≈ loopback), EP widths {4,8,16,32,64} each a constexpr shim instantiation. Pitfalls: containers need the IMEX channel device; teardown must shutdown() the socket. GIN scale-out sections remain the design for IB/RoCE beyond one rack; >4 nodes upgrades control to the preserved SMR design. | | `models/glm52/dspark-mtp.md` | DSpark speculative decoding (community `RedHatAI/GLM-5.2-speculator.dspark`, not native MTP): qwen3-arch 5-layer draft at hidden 6144 + rank-256 Markov head; verify span rides the decode buckets (span-4 default). M1 span-steps, M2 draft lane, M3 greedy round loop (sharegpt c1 1.52×), M4 sampled verify — non-greedy speculation via prefix-match over sampled tokens (c1 code 2.38× at temp 1; full rejection sampling probe-measured out at ≤ +1.5%). | +| `models/glm52/native-mtp-accuracy.md` | Native MTP must consume the target model's final-normalized hidden, matching official vLLM rather than the pre-final-norm residual. The fix raises matched c8 accepted length from `1.753` to `3.725` versus `3.786`; repeated matched c1 measures OpenInfer `7.749 ms` TPOT versus official vLLM `8.814 ms`. | | `models/glm52/paged-kv-prefix-cache.md` | Static per-slot KV partitions → per-rank `BlockPool` of 64-token content-hashed pages (Kimi #239 pattern, zero kernel changes): full-lifetime admission reservation, coordinator-shipped `Glm52StepKv` page rows, prefix caching on by default (suffix-only prefill), DSpark × prefix-cache mutually exclusive, launch-ahead lease breaks at page boundaries. Merged as #588: all jz-38 gates green, warm prefix TTFT 14.12 s → 0.84 s (16.9×) byte-identical, step bench flat vs the D5 anchor. | | `models/glm52/continuous-batching.md` | D2 + D2.5 execution record: multi-slot admission (8 requests/rank, least-loaded first) + {1,2,4,8} batch-bucket graphs (smallest bucket covering the fullest rank, per-bucket `Glm52BucketState`). Solo 22.4 ms/step; D2's c9 cliff killed (47.1 → 31.8 ms/step, 171 → 254 tok/s); poisson soaks clean; pinned slot-3/7 parity PASS. Known: buckets are distinct FP associations (bucket-crossing requests can greedy-diverge at near-ties); open anomaly: one-off silent request drop (#551). | | `models/glm52/oracle-harness.md` | Self-contained accuracy oracle: `tools/accuracy/glm52_oracle.py` (pinned transformers 5.12.1 official `glm_moe_dsa`, fp8-precision-emulated) emits hardcodable Rust probe constants; `oracle/mla.rs` replays the seeded input and asserts. MLA gate green on jz38 (64/64 probes, diff RMS 1.8e-5), negative controls red. No MB fixtures in git. | diff --git a/docs/models/glm52/native-mtp-accuracy.md b/docs/models/glm52/native-mtp-accuracy.md new file mode 100644 index 000000000..6666c8e8c --- /dev/null +++ b/docs/models/glm52/native-mtp-accuracy.md @@ -0,0 +1,445 @@ +# GLM5.2 native MTP accuracy and acceptance + +> **TL;DR:** The c8 native-MTP acceptance gap was a correctness bug at the target/draft boundary: +> OpenInfer passed the target's pre-final-norm residual to MTP, while official vLLM passes the +> model-returned final-normalized hidden. Using `scratch.final_normed` raises the matched c8 mean +> accepted length from `1.753` to `3.725` versus official vLLM's `3.786`, and reduces TPOT from +> `23.44 ms` to `11.31 ms`. A subsequent three-run matched c1 measures OpenInfer at `7.749 ms` +> versus official vLLM at `8.814 ms`; their proxy round costs are `30.04` and `30.72 ms`, so there +> is no remaining c1 TPOT deficit. On the selected 251-token target trajectory, mean accepted +> length changes from `1.000` to `5.795`. +> +> **Last touched:** 2026-07 + +## Preparation + +- **Read**: + - `docs/index.md` — routes GLM5.2 model work and accuracy methodology. + - `docs/models/glm52/dspark-mtp.md` — establishes accepted-length accounting, span verification, + and the rule that speculative performance must be explained as round cost divided by accepted + tokens. + - `docs/models/glm52/moe-tp8-low-latency.md` — records the existing TP8/EP8 topology and its + numerical and performance characteristics. + - `docs/models/glm52/oracle-harness.md` — requires an external truth implementation, seeded + reproducibility, and RMS/p99 rather than max-only float checks. + - `docs/playbooks/accuracy-parity-playbook.md` — prescribes exact token IDs, first-diff + localization, and production-path teacher forcing. + - `docs/playbooks/bench-vs-vllm.md` — defines matched hardware, model, client, sampling, and + prefix-cache controls for comparative serving measurements. +- **Relevant history**: + - DSpark reached useful accepted lengths only after validating the real online hidden-state and + verify path; a small standalone forward match was not considered sufficient. + - GLM5.2 bucket/topology changes can move near-tie decisions without breaking task accuracy, so + token mismatches and structural drift must be classified separately. +- **Plan**: + 1. Choose one deterministic prompt and capture the official vLLM greedy token IDs plus per-step + target raw hidden, final target logits, and native-MTP intermediate tensors. + 2. Teacher-force those exact token IDs through the OpenInfer production path and capture the same + checkpoints, without comparing states after the token streams diverge. + 3. Find the first divergent checkpoint in this order: target raw hidden, MTP prepare/first logits, + then MTP recycle hidden and layer-78 KV across draft steps 2–5. + 4. Classify the first difference as target topology/numerics, MTP forward semantics, or MTP + paged-KV/recycle state; add the narrowest reproducible accuracy gate before changing code. + 5. Re-run online accepted-length and c8 TPOT A/B on the same 8×H200 node. Treat an optimization as + a win only when measured accepted length improves without target-quality regression. +- **Risks / open questions**: + - Official vLLM runs TP8+EP8 while the current OpenInfer path is TP1/DP8+EP8. Target raw hidden can + differ from collective and accumulation order even when final task quality is healthy. + - The existing five-row oracle proves MTP front/layer-78 behavior on official vLLM inputs, but + does not cover a long online teacher-forced trajectory or the production paged-KV recycle path. + - Near-tie top logits must not be reported as a structural model bug without regret/margin data. + +### C8 accepted-length decomposition follow-up + +- **Question**: Which component accounts for the measured `3.786` versus `1.753` c8 accepted-length + gap: different target trajectories, MTP forward numerics, or state carried between draft steps? +- **Plan**: + 1. Reconstruct the seeded 64-request token-ID corpus used by the retained benchmark and persist + it with tokenizer, seed, and checksum metadata. Feed identical token IDs to both engines so + prompt generation is no longer a hidden variable. + 2. Run the engines sequentially on the same 8×H200 node and retain per-request target tokens, + draft tokens, accepted-prefix lengths, and rejection margins. Stratify requests by low, + median, and high acceptance before collecting tensor checkpoints. + 3. For each selected request, stop at the first differing target token or first differing draft + token. Compare the raw target hidden supplied to MTP before comparing MTP prepare output, + first-step logits, recycled normalized hidden, and layer-78 KV for draft steps 2–5. + 4. Replay official target hidden through the OpenInfer MTP oracle at those proposal-entry states. + This intervention separates target-model drift from the MTP forward without requiring the two + target implementations to remain on the same greedy path. + 5. Classify the aggregate gap by counterfactual: measure how many acceptance decisions recover + under a shared target trajectory and official-hidden replay. Only then change code, add the + narrowest production-path regression gate, and rerun the matched c8 benchmark. +- **Interpretation gates**: + - Target raw hidden differs first: report target topology/numerics and quantify acceptance under + official-hidden replay; do not attribute the gap to MTP forward. + - Target raw hidden agrees but the first draft logits differ materially: localize MTP prepare, + layer 78, final norm, and shared head in that order. + - First draft agrees but a later draft diverges: inspect recycled hidden and MTP paged-KV + positions at the first differing step. + - Use top-1 regret, target/draft margins, RMS, and p99. Exact equality is not required when the + selected token and acceptance decision are stable. +- **Operational risks**: + - Official vLLM and OpenInfer consume the same eight GPUs, so they must run sequentially. Compact + artifacts must be saved before switching engines, and the clean OpenInfer service must be + restored afterward. + - Dumping full hidden states for `64 × 256` tokens is excessive. Only first-difference states from + the stratified subset should be retained; aggregate runs keep tokens, counters, and margins. + - TP8+EP8 versus TP1/DP8+EP8 may prevent bit parity even on a teacher-forced path. Conclusions + must distinguish harmless numeric drift from changes to selected tokens or accepted prefixes. + +### Matched c1 performance follow-up + +- **Question**: Does corrected OpenInfer still trail official vLLM by about `2 ms` at concurrency + one, or did the earlier comparison mix OpenInfer c8 (`11.31 ms`) with vLLM c1 (`9.10 ms`)? +- **Plan**: + 1. Reuse the retained 64-request random corpus, greedy sampling, 256 output tokens, disabled + prefix cache, and identical client/version flags. Validate the live model id and benchmark + client flags before collecting results. + 2. Run OpenInfer and official vLLM sequentially on the same 8×H200 node at concurrency one, + saving benchmark JSON and speculative-acceptance counters for every measured run. + 3. Compare TPOT, accepted length, and `TPOT × accepted length`; only attribute a compute-side + difference when the matched c1 round costs disagree beyond run-to-run noise. + 4. Restore the corrected OpenInfer service and health-check it after the official-vLLM run. +- **Risks / open questions**: + - The official vLLM and OpenInfer topologies differ, so accepted length remains + content-dependent; round cost is required alongside TPOT. + - A single c1 run can be distorted by startup and host noise. Use warmup plus repeated measured + runs and report their spread rather than selecting the best number. + +## Execution Log + +### Native MTP forward and serving integration + +- Added the native layer-78 MTP decoder, its separate one-layer KV state, proposal/recycle loop, and + scheduler-wide collective modes for reset-only, context-only, and proposal rounds. +- Verified release build, library and CLI tests, an 8-rank end-to-end run, and an official-vLLM + fixture gate. +- Result: functional. The official-vLLM fixture comparison measured MLP RMS `8.62e-3` + (`p99 2.54e-2`), chained hidden RMS `1.87e-2` (`p99 5.86e-2`), exact top-1, top-8 overlap `8/8`, + and top-32 overlap `30/32`. + +### PR hardening + +- Replaced the public `dspark path + native_mtp bool` pair with one + `None / Dspark(path) / NativeMtp` launch type, so callers cannot enable two drafters. +- Bumped the rank-host wire version and added a `BuildModel` round-trip test covering the drafter + payload. Mixed coordinator/rank-host binaries now reject at the handshake. +- Replaced the seven-argument MTP round command with `Reset / Context / Propose` variants. The + runner/model boundary can no longer carry a context-only mode with proposal data. +- Removed the per-round global test lock. The production gate records its first proposal once and + copies the request-local acceptance counters when the slot is released. +- Extended that gate, without another model load, to reuse rank 0's released slot and issue eight + concurrent rank-pinned requests with different prompt and output lengths. The reused slot has + its own first-proposal and request-level acceptance assertions; the other ranks provide the + mixed-state collective liveness coverage. +- Removed unrelated GSM8K evaluator changes from this branch. The evaluator results below remain + evidence gathered during bring-up, not part of the native-MTP serving change. + +### EP8 MoE correctness prerequisite and scope + +- This PR also corrects the shared EP8 routed-expert pipeline used by every target MoE layer, not + only native MTP's layer 78. The previous path multiplied route weights into the SwiGLU + activation before FP8 quantization and folded the routed scaling factor into router weights. + Official vLLM quantizes the unweighted activation before W2, applies each route weight to W2's + BF16 output, combines experts, and then applies the `2.5` routed scaling factor before adding the + unscaled shared-expert output. +- The correction is a prerequisite for the MTP oracle because the target and layer-78 decoder + share `glm52_moe_ep8_routed_forward`; retaining the old target semantics only for layers 0–77 + would create two numerically different implementations of the same checkpoint MoE. +- The official-vLLM layer-78 golden gate exercises the corrected EP8 kernel and reports MLP RMS + `8.62e-3`, chained hidden RMS `1.87e-2`, exact top-1, top-8 overlap `8/8`, and top-32 overlap + `30/32`. The existing layer-6 EP8 oracle exercises the same kernel in a plain target layer and is + part of the 8×H200 validation set. A fresh run passes all four `g64/g32/g16/g8` groupings with + `62/64` probes inside tolerance for each grouping and the two permitted near-tie outliers, in + `27.84 s`. +- The post-change plain serving result is GSM8K `195/200`; native MTP is `196/200`. No retained + pre-change plain GSM8K run used the identical harness, so this document does not claim a + before/after task-accuracy improvement from the MoE correction. The direct evidence is the + official tensor oracle plus the post-change plain-path regression gate. + +### End-task accuracy + +- Ran the same GSM8K five-shot evaluation with native MTP disabled and enabled. +- Result: target quality is stable on the measured slice: plain `195/200` (`97.5%`), native MTP + `196/200` (`98.0%`). This rules out a broad target-quality collapse, not an MTP acceptance issue. + +### Serving performance and acceptance attribution + +- Matched official vLLM and OpenInfer on one 8×H200 node with greedy decoding and fixed output + length. An exploratory c1 terminal snapshot measured: + + | Engine | c1 TPOT | Mean accepted length | TPOT × accepted length | + | --- | ---: | ---: | ---: | + | official vLLM native MTP | `9.10 ms` | `3.40` | `30.9 ms` | + | OpenInfer native MTP | `15.03 ms` | `2.00` | `30.1 ms` | + +- OpenInfer plain c1 TPOT was `18.66 ms`; native MTP therefore helps, but less than vLLM. +- The OpenInfer accepted-length value is reproducible from eight retained request histograms + (`1009` speculative rounds), not an average of per-request averages. The standalone c1 client + result was not retained, however, so the table is mechanism evidence rather than a regression + baseline. +- The retained c8 artifacts for the same 64 random requests, plus the corrected replay, report: + + | Engine | Mean TPOT | Mean accepted length | + | --- | ---: | ---: | + | official vLLM native MTP | `16.90 ms` | `3.786` | + | OpenInfer before hidden-boundary fix | `23.44 ms` | `1.753` | + | OpenInfer after hidden-boundary fix | `11.31 ms` | `3.725` | + + The OpenInfer values are weighted from the final 64 per-request histograms: `9191` rounds before + the fix and `4346` rounds after it. +- Result: the corrected A/B confirms that shorter accepted prefixes caused most of the observed + TPOT gap. `TPOT × accepted length` is `41.10 ms` before and `42.11 ms` after the fix, consistent + with unchanged speculative-round cost. The post-fix OpenInfer/vLLM TPOT comparison still + includes engine and topology differences; it is not a claim that the MTP round itself is faster. + +#### Retained c8 measurement record + +- Hardware/model: one 8×H200 node, the same GLM5.2 FP8 checkpoint, prefix cache disabled. +- Client workload: random dataset, default seed `0`, nominal input length `128`, output length + `256`, `64` prompts, concurrency `8`, unlimited request rate, temperature `0`, ignore EOS. + Both runs completed `64/64`, with `8152` input and `16384` output tokens. +- Official vLLM provenance: commit `dcfebf93`, benchmark JSON SHA-256 + `af909e6a8a06562ff42bfa882decf1462ee0c63771be08b2229201f9b766e780`. + Raw counters are `4347` proposal rounds, `21735` drafted tokens, and `12109` accepted draft + tokens. Thus mean accepted length including the target bonus is + `1 + 12109 / 4347 = 3.785599`. +- OpenInfer provenance: commit `fd6bd6e0`, benchmark JSON SHA-256 + `66fd2d4ea38fa7ceb1429612e35f2f9fce6eed26f35546cfb2da1c0563619760`. + The aggregate accepted-draft histogram for indices `0..7` is + `[4564, 3113, 955, 415, 65, 79, 0, 0]`: `9191` rounds and `6923` accepted drafts. Thus mean + accepted length including the target bonus is `1 + 6923 / 9191 = 1.753237`. +- Corrected OpenInfer replay: benchmark JSON SHA-256 + `5e4c3d7219fee985008e6932bdb07c74708b5a39691480259db357e6ece2db55`; + histogram record SHA-256 + `6c95eb444cbf8296773be94b425d2b83899b47e718f2c42e13dab0a5995fb2ba`. + The aggregate histogram is `[795, 650, 469, 852, 201, 1379, 0, 0]`: `4346` rounds and `11843` + accepted drafts. Mean accepted length is `1 + 11843 / 4346 = 3.725035`. + +### C8 first-difference diagnosis and fix + +- Reconstructed the retained 64-request corpus from the benchmark version that produced the + original artifact. Its token-ID corpus SHA-256 is + `59516edc33d2a7a36b63628db7cd4eb0888f3aec7f9ac97b49039920743928d3`; it contains `8152` + prompt tokens and requests exactly `16384` output tokens. +- Per-round target traces disprove the earlier working hypothesis that target-trajectory drift + explains the full acceptance gap. Six requests share at least 32 target tokens across engines. + An exploratory capture found one 251-token shared trajectory: before the fix OpenInfer rejects + every first draft (`251` rounds, mean accepted length `1.000`), while official vLLM reached + approximately `5.78`. That full official per-request capture was not retained as a durable + artifact, so the approximate value is diagnostic rather than a regression baseline. +- At the first proposal of that request, the embedding is bit-exact but the target hidden passed + into MTP differs before any layer-78 or MTP-KV work: + + | Boundary | Before-fix cosine | After-fix cosine | + | --- | ---: | ---: | + | target hidden supplied to MTP | `0.8361` | `0.9818` | + | `eh_proj` output | `0.8320` | `0.9771` | + | layer-78 raw hidden | `0.7404` | `0.9877` | + | recycled normalized hidden | `0.7171` | `0.9861` | + +- The before-fix target-hidden norm is `281.01`, versus official vLLM's `78.63`. OpenInfer's + post-fix norm is `77.32`. This large discontinuity is not attributable to BF16 reduction order. +- Official vLLM registers `GlmMoeDsaForCausalLM` on its DeepSeek-V2-compatible path. That target + returns final-RMSNorm hidden states, and the MTP proposer consumes that model return directly. + OpenInfer instead selected `scratch.hidden`, the residual before final RMSNorm. The fix changes + the source to `scratch.final_normed`; token shifting, layer 78, MTP KV, and the verifier remain + unchanged. +- On the selected shared trajectory, the first draft changes from the incorrect token `98863` to + official vLLM's `98825`. The corrected OpenInfer run records 44 rounds with accepted-draft + histogram `[1, 1, 0, 0, 0, 42, 0, 0]`, or mean accepted length `5.795`. +- The retained short official trace has SHA-256 + `02e97a2f4d23b573cf53b20611840dc098e5f694bfde9284f532bda2c972d999`. Its six captured + proposal records reject the first two drafts, then fully accept all five drafts in the next four + records. The ordered checksum-list SHA-256 for the retained official tensor trace is + `ee07cf72de367ffcb03d543bf8be4ea7ff9bb9bb7fa1d42a08cf1b25ba52461c`. + The OpenInfer tensor dumps used for the cosine table were transient, so those tensor metrics are + diagnostic rather than a standalone reproducible gate. +- The corrected selected-request response and acceptance log have SHA-256 + `8f5a6663d6ddb10c59458a1f6849fff468150265443398da63aa492d0b3ca3f3` and + `5807097280c2edf83a60b4f1841760ca60e7061dddad6916edb49a5835093d94`. +- Release validation passes `84` library tests with `19` GPU/model gates ignored, followed by the + two explicitly enabled official-vLLM MTP front and EP8 layer-78 golden gates. An ignored + production-path gate loads the full checkpoint on 8×H200, submits the selected 256-token request + through the real scheduler/executor, and asserts target trajectory, first draft `98825`, and mean + accepted length at least `5.0`; the original single-request form passes in `210.37 s`. PR + hardening extends the same model-loaded gate with a released-slot reuse request plus eight + concurrent rank-pinned requests of different lengths. The reused slot must again draft `98825` + first, complete at least 32 speculative rounds, and retain mean accepted length at least `5.0`; + the other requests exercise mixed-state collective liveness. The hardened gate passes in + `220.70 s`. The clean c8 replay completes `64/64` requests with the exact retained input/output + totals. + +### Matched c1 after the hidden-boundary fix + +- Reused the exact c8 random workload at concurrency one: seed `0`, `64` prompts, temperature `0`, + ignore EOS, prefix cache disabled, nominal input length `128`, and output length `256`. Every + measured run completed `64/64` requests with `8152` input and `16384` output tokens. One + readiness probe and two warmup requests preceded each measured run and are excluded from the + counters below. +- Both engines ran sequentially on the same 8×H200 node. OpenInfer used TP1/DP8+EP8; official vLLM + commit `dcfebf93` used TP8+EP8 from its upstream nightly image, with five native-MTP draft tokens. + Three measured runs per engine give: + + | Engine | Mean TPOT | Run-to-run σ | Mean accepted length | TPOT × accepted length | Mean TTFT | Output throughput | + | --- | ---: | ---: | ---: | ---: | ---: | ---: | + | OpenInfer native MTP | `7.749 ms` | `0.003 ms` | `3.8762` | `30.04 ms` | `432.50 ms` | `106.29 tok/s` | + | official vLLM native MTP | `8.814 ms` | `0.001 ms` | `3.4858` | `30.72 ms` | `232.57 ms` | `103.22 tok/s` | + +- Per-run OpenInfer TPOT is `[7.753, 7.746, 7.749] ms`; official vLLM is + `[8.815, 8.814, 8.813] ms`. OpenInfer is therefore `1.065 ms` (`12.1%`) lower on matched c1, + rather than approximately `2 ms` higher. OpenInfer output throughput is `3.0%` higher, while its + TTFT is about `200 ms` higher; the result is decode-specific, not an end-to-end latency win. +- OpenInfer reproduces the same aggregate accepted-draft histogram in all three runs: + `[815, 644, 443, 346, 212, 1723, 0, 0]`, or `4183` rounds and `12031` accepted drafts. + Official vLLM likewise reproduces `4714` rounds and `11718` accepted drafts in each run. +- The official acceptance counters are retained inside each enhanced benchmark JSON, not in a + separate log. This benchmark-client build fetches vLLM's cumulative `vllm:spec_decode*` + Prometheus counters after the readiness probe and warmups but immediately before the measured + requests, fetches them again after the measured requests, and writes their delta as + `spec_decode_num_drafts`, `spec_decode_draft_tokens`, `spec_decode_accepted_tokens`, + `spec_decode_acceptance_length`, and per-position rates. Thus `3.485787` is directly auditable as + `1 + 11718 / 4714` from every retained official JSON; readiness and warmup traffic are outside + the snapshot window. +- The proxy round costs differ by only `0.69 ms` (`2.2%`) and favor OpenInfer. Most of the TPOT + delta comes from OpenInfer accepting `0.390` more tokens per speculative round on this engine's + target trajectories. This does not imply identical target text or identical MTP numerics across + the two topologies. +- The benchmarked OpenInfer binary was built from `2dd5237e` (native-MTP commits `f35292c1`, + `fd6bd6e0`, `83389e75`, and `2dd5237e` over base `272c2f94`) in the + release profile. The runtime binary SHA-256 is + `1a803c63c64377a82d611a1713cbf769b6bee1440c14902561922974c3691f62`. The official-vLLM + image digest is `sha256:79460a12901891f5e74d7a6ee1259f8aad9aa3b405cc0753534ee4ff6124fd3b`. +- The ordered checksum-list SHA-256 for the three OpenInfer benchmark JSON files, three OpenInfer + acceptance logs, and three official-vLLM benchmark JSON files is + `a0dbc587ec16816dfae134a0f3f629745d8a6db5375bc09fb306fa4d82f165e9`. + The benchmark-client binary SHA-256 is + `953dab27d67e370645e8f78163bbd7e9cb423539f1290233988ab90ad617ed8c`; the corpus SHA-256 + remains `59516edc33d2a7a36b63628db7cd4eb0888f3aec7f9ac97b49039920743928d3`. +- The checkout-local benchmark binary terminated with `Illegal instruction` only after entering + the request path. No measured request reached the server in that attempt. The run switched to a + previously validated upstream-client build, repeated the exact dry-run token totals, passed a + one-request smoke test on both engines, and then collected the six retained runs above. +- After measurement, the official-vLLM container was stopped and the corrected OpenInfer service + was restored and health-checked. + +### Pre-fix fixed-prompt comparison + +- Used the deterministic prompt `The capital of France is` with greedy decoding. The official + trace was captured from vLLM commit `dcfebf93`; OpenInfer ran its normal target, scheduler, + long-lived MTP KV, and five-step proposal loop. +- The first proposal is exact: + + | Engine | Draft tokens | + | --- | --- | + | official vLLM | `[13, 576, 3283, 315, 12089]` | + | OpenInfer | `[13, 576, 3283, 315, 12089]` | + +- Repeating the OpenInfer request four times produced the same proposal and output, ruling out + request-to-request nondeterminism. +- The first four rounds that share the same target token trajectory accept: + + | Engine | Accepted drafts per round | Mean including target bonus | + | --- | --- | ---: | + | official vLLM | `[1, 1, 1, 2]` | `2.25` | + | OpenInfer | `[1, 1, 1, 1]` | `2.00` | + +- The durable proposal record before their round cadence diverges is: + + | Round | official vLLM | OpenInfer | + | ---: | --- | --- | + | 0 | `[13, 576, 3283, 315, 12089]` | `[13, 576, 3283, 315, 12089]` | + | 1 | `[504, 279, 6722, 315, 9621]` | `[504, 279, 3283, 315, 12089]` | + | 2 | `[311, 7148, 374, 12089, 11]` | `[311, 14915, 409, 93729, 273]` | + | 3 | `[374, 220, 16, 13, 20]` | `[374, 264, 3283, 315, 220]` | + +- The first two proposals agree through the acceptance-relevant prefix. Their first later top-1 + difference is `6722` versus `3283`; official logits are `18.25` versus `18.00`, and that + position is already behind a rejected draft, so it cannot change accepted length. +- The only acceptance-changing difference in the four comparable rounds is official token `220` + versus OpenInfer token `264`. In the official BF16 logits, `220` is `18.00` and `264` is + `17.25`. The target token stream remains identical through the compared interval; only the + drafter loses one accepted token. +- A suspected MTP-step raw-hidden recycle bug was rejected. Official vLLM traces show that each + next draft step receives the prior `shared_head`-normalized hidden bit-exactly; OpenInfer already + feeds the same normalized value. +- Result: this narrow prompt excluded a token shift, stale recycle hidden, and an obvious MTP + KV-position error, but did not validate the target-to-MTP input boundary. Its first proposal + happened to agree despite the wrong source tensor. The broader c8 first-difference trace above + supersedes the earlier topology-numerics attribution. + +#### Retained fixed-prompt record + +- Official side: vLLM commit `dcfebf93`, TP8+EP8, `max_tokens=8`; its five rank-0 proposal groups + contain 25 paired forward/logit records. The ordered checksum-list SHA-256 is + `cae3e354d39b21f16d207899d8addc3e12f3b14d1cdb4d4cfd893a0533f4778f`. +- OpenInfer side: commit `fd6bd6e0`, TP1/DP8+EP8, `max_tokens=20` so the short-tail policy exposes + at least four proposal rounds. The request uses temperature `0` and seed `1`. +- The prompt token IDs are `[6722, 315, 9621, 374, 12089]`. The shared target token trajectory + through the compared interval is + `[12089, 13, 31008, 504, 12089, 311, 54831, 374, 220, 101294]`. +- Official vLLM reports four verified rounds, five accepted drafts, and per-position acceptance + rates `[1.0, 0.25, 0.0, 0.0, 0.0]`; these counters and the proposal transitions reproduce + accepted drafts `[1,1,1,2]`. OpenInfer's four rounds are reproduced directly from the table and + the shared target trajectory as `[1,1,1,1]`. + +### Latest-main integration + +- Rebased the native-MTP series onto the TP4 prefill execution change. The semantic conflict was + larger than the textual conflict: the coordinator must retain `prefill_chunk_size: Option` + rather than reduce it to a boolean, while the speculative drafter remains a separate launch + choice. +- Updated the native-MTP KV arena and its byte ledger to use the same eight-slot pool geometry, + adopted the renamed paged-MQA indexer shape constructor, and kept remote rank hosts on the + decode-only build contract. +- Local validation passes the GLM5.2 release server check and `86` library tests, with `20` + hardware/model gates ignored. + +### How to use aggregate accepted length + +- The matched random benchmark uses identical input/output lengths, but official vLLM runs a + TP8+EP8 target while OpenInfer runs TP1/DP8+EP8. Their greedy generated texts diverge from the + first request, so each MTP instance is scored on a different output trajectory. +- Accepted length is strongly content-dependent. Comparing `3.40` to `2.00` therefore describes + the measured serving systems and suggests one contributor to their TPOT difference, but does + not isolate MTP forward correctness or quantify compute-side differences. +- Aggregate acceptance alone still cannot identify the faulty component. Here, per-request target + traces found long shared trajectories and tensor comparison found the first difference before + MTP forward. After that attribution, the matched c8 acceptance result is a useful regression + measurement. + +## Debrief + +- **Outcome**: Native MTP now consumes the same target-hidden boundary as official vLLM. Matched c8 + acceptance is `3.725` versus `3.786`, and the formerly pathological shared trajectory aligns. + Matched c1 TPOT is `7.749 ms` versus official vLLM's `8.814 ms`, with near-equal proxy round + cost. Task-level quality remains healthy on the measured GSM8K slice. +- **Pitfalls encountered**: + - A small offline layer-78 match cannot certify online accepted length because it bypasses the + production target hidden-state source and long-lived MTP KV. + - Calling both tensors “target hidden” concealed a material API boundary: the residual before + final RMSNorm and the model-returned hidden are not interchangeable MTP inputs. + - TPOT alone hid the mechanism; multiplying TPOT by accepted length motivated the acceptance + investigation, but is not a substitute for a retained per-round timing measurement. + - Reading the vLLM model return type in isolation suggested that raw hidden might be recycled. + The captured call boundary proved the opposite: vLLM returns normalized recycle hidden to the + proposer while retaining raw hidden only for logits. + - A top-1 mismatch after an earlier rejected draft is diagnostically useful but irrelevant to + accepted length. First-diff analysis must stop at the effective accepted prefix. + - A benchmark client's dry-run can succeed while its request path contains CPU instructions + unsupported by the current host. Require a real one-request smoke before a long measurement. +- **Lessons learned**: + - Native-MTP acceptance is itself an accuracy metric even when target-generated answers remain + correct. + - Compare c1 with c1. The earlier apparent `11.31` versus `9.10 ms` gap mixed OpenInfer c8 with + an exploratory official-vLLM c1 snapshot; matched repeated c1 reverses that conclusion. + - Cross-engine aggregate acceptance is meaningful only when both engines follow the same target + token trajectory; identical prompt lengths and sampling flags are insufficient. + - For distributed BF16/FP8 paths, classify top-1 differences with logit margin and whether the + position can affect acceptance before treating them as structural defects. +- **Follow-ups**: + - Retain the official-vLLM fixtures and fixed-prompt procedure as the accuracy gate for changes + to MTP prepare, recycle, sparse-index reuse, and KV positioning. + - Keep the reconstructed c8 corpus as a serving-level regression measurement. Use the selected + production-path gate for the target/MTP handoff so a future change from `final_normed` back to + the pre-norm residual fails before benchmark interpretation. diff --git a/openinfer-glm52/src/config.rs b/openinfer-glm52/src/config.rs index e22842436..840e0160d 100644 --- a/openinfer-glm52/src/config.rs +++ b/openinfer-glm52/src/config.rs @@ -9,6 +9,8 @@ use serde_json::Value; pub(crate) const GLM52_HIDDEN: usize = 6144; pub(crate) const GLM52_VOCAB: usize = 154_880; pub(crate) const GLM52_LAYERS: usize = 78; +/// Checkpoint layer containing the native multi-token prediction decoder. +pub(crate) const GLM52_MTP_LAYER: usize = GLM52_LAYERS; pub(crate) const GLM52_DENSE_LAYERS: usize = 3; /// The checkpoint's `max_position_embeddings` — `probe_config_json` pins the /// config to exactly this, so it doubles as the architecture ceiling any @@ -41,7 +43,7 @@ pub(crate) const GLM52_EXPERT_INTERMEDIATE: usize = 2048; pub(crate) const GLM52_ROUTED_EXPERTS: usize = 256; pub(crate) const GLM52_TOPK: usize = 8; const GLM52_SHARED_EXPERTS: usize = 1; -const GLM52_ROUTED_SCALING_FACTOR: f64 = 2.5; +pub(crate) const GLM52_ROUTED_SCALING_FACTOR: f64 = 2.5; const GLM52_RMS_NORM_EPS: f64 = 1.0e-5; /// The f32 the GPU norm kernels consume (every RMSNorm in the model shares /// the one checkpoint eps that `probe_config_json` validates). diff --git a/openinfer-glm52/src/fp8.rs b/openinfer-glm52/src/fp8.rs index 3c3a09c6e..5de125581 100644 --- a/openinfer-glm52/src/fp8.rs +++ b/openinfer-glm52/src/fp8.rs @@ -482,6 +482,16 @@ impl Glm52MlpScratch { .alloc_zeros::(rows * GLM52_GEMV_MMA_SCRATCH_FLOATS_PER_ROW)?, }) } + + #[cfg(test)] + pub(crate) fn gate_up(&self) -> &CudaSlice { + &self.gate_up + } + + #[cfg(test)] + pub(crate) fn silu_out(&self) -> &CudaSlice { + &self.silu_out + } } /// A plain fp8 SwiGLU MLP over the scratch's `rows` tokens into a diff --git a/openinfer-glm52/src/layer.rs b/openinfer-glm52/src/layer.rs index b01ed2f23..899191fe2 100644 --- a/openinfer-glm52/src/layer.rs +++ b/openinfer-glm52/src/layer.rs @@ -91,6 +91,16 @@ pub(crate) struct Glm52LayerCaches { pub(crate) index_k_cache: Option>, } +/// Sparse-index policy for one decoder-layer forward. Target layers compute +/// or inherit indices according to their checkpoint role. Native MTP computes +/// layer 78's indices on its first pass, then reuses the selected rows for +/// the remaining four proposal iterations (`index_share_for_mtp_iteration`). +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum Glm52LayerIndexMode { + Normal, + Reuse, +} + /// Everything one decode step shares across layers: the token position, the two /// rotary tables (MLA interleaved; indexer half-split — different conventions, /// same `[32]` cos/sin extent), and the paging plumbing common to every layer's @@ -168,6 +178,7 @@ pub(crate) fn glm52_layer_attention_half( parity: usize, first_layer: bool, tp_ar: Option<(&mut crate::moe_tp::Glm52MoeTpState, usize)>, + index_mode: Glm52LayerIndexMode, ) -> Result<()> { // Attention-TP: a head-sharded layer (8 of 64 heads) produces an o_proj // PARTIAL that must cross the AR brick before the residual add; holding @@ -195,8 +206,15 @@ pub(crate) fn glm52_layer_attention_half( let tokens = step.mla_sched.batch(); glm52_mla_front_q_into(ctx, &w.mla, &s.layer.normed, &mut s.mla_front)?; let mut topk_ready = None; - match &w.indexer { - Glm52LayerIndexer::Full(indexer) => { + match (&w.indexer, index_mode) { + (Glm52LayerIndexer::Full(_), Glm52LayerIndexMode::Reuse) => { + ensure!( + caches.index_k_cache.is_some(), + "GLM5.2 reused full-indexer layer is missing its index-K cache" + ); + *carry_ready = true; + } + (Glm52LayerIndexer::Full(indexer), Glm52LayerIndexMode::Normal) => { let index_k_cache = caches .index_k_cache .as_mut() @@ -227,12 +245,15 @@ pub(crate) fn glm52_layer_attention_half( } *carry_ready = true; } - Glm52LayerIndexer::Shared => { + (Glm52LayerIndexer::Shared, Glm52LayerIndexMode::Normal) => { ensure!( caches.index_k_cache.is_none(), "GLM5.2 shared-indexer layer unexpectedly owns an index-K cache" ); } + (Glm52LayerIndexer::Shared, Glm52LayerIndexMode::Reuse) => { + anyhow::bail!("GLM5.2 cannot request explicit top-k reuse on a shared-indexer layer") + } } ensure!( *carry_ready, @@ -396,7 +417,19 @@ pub(crate) fn glm52_decoder_layer_forward( tokens, s.layer.normed.data_mut(), )?; - glm52_layer_attention_half(ctx, None, w, caches, step, s, carry_ready, 0, true, None)?; + glm52_layer_attention_half( + ctx, + None, + w, + caches, + step, + s, + carry_ready, + 0, + true, + None, + Glm52LayerIndexMode::Normal, + )?; match &w.mlp { Glm52LayerMlp::Dense(dense) => glm52_dense_mlp_forward_into( ctx, diff --git a/openinfer-glm52/src/lib.rs b/openinfer-glm52/src/lib.rs index 300b89f55..0f92d2a9f 100644 --- a/openinfer-glm52/src/lib.rs +++ b/openinfer-glm52/src/lib.rs @@ -25,6 +25,7 @@ mod moe_decode; mod moe_ep8; mod moe_ep_wo; mod moe_tp; +mod mtp; #[cfg(test)] mod oracle; mod prefill_tp; @@ -75,6 +76,37 @@ use crate::model::glm52_pool_blocks; pub const GLM52_PREFILL_CHUNK_ALIGN: usize = GLM52_MODEL_LEN_ALIGN; pub const GLM52_DEFAULT_PREFILL_CHUNK_SIZE: usize = 16_384; +/// Optional speculative decoder used by the GLM5.2 engine. +#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +pub enum Glm52Drafter { + None, + /// External DSpark checkpoint. + Dspark(PathBuf), + /// Checkpoint-native layer-78 multi-token prediction decoder. + NativeMtp, +} + +impl Glm52Drafter { + fn enabled(&self) -> bool { + !matches!(self, Self::None) + } + + fn is_dspark(&self) -> bool { + matches!(self, Self::Dspark(_)) + } + + fn is_mtp(&self) -> bool { + matches!(self, Self::NativeMtp) + } + + fn dspark_path(&self) -> Option<&Path> { + match self { + Self::Dspark(path) => Some(path), + Self::None | Self::NativeMtp => None, + } + } +} + /// TP4 prefill-only configuration. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct Glm52PrefillOnlyOptions { @@ -87,12 +119,10 @@ pub struct Glm52PrefillOnlyOptions { pub struct Glm52LaunchOptions { pub tp_size: usize, pub dp_size: usize, - /// DSpark drafter checkpoint dir (`RedHatAI/GLM-5.2-speculator.dspark`). - /// Enables speculative decoding for greedy AND sampled requests (the - /// verify span prefix-matches per-row sampled tokens — lossless): verify - /// spans ride the decode buckets, accepted tokens commit in batches, - /// per-request accept stats are logged on release. - pub dspark_draft_model_path: Option, + /// Optional speculative decoder. DSpark enables lossless speculative + /// sampling from an external checkpoint; native MTP uses the checkpoint's + /// layer-78 decoder and currently requires single-node EP8. + pub drafter: Glm52Drafter, /// Per-request context cap (`prompt + max_tokens - 1 <= max_model_len`). /// `None` sizes it from the post-weight-load free VRAM (fleet minimum); /// an explicit value is still validated against that budget so an @@ -102,14 +132,15 @@ pub struct Glm52LaunchOptions { pub prefill_only: Option, /// vLLM-style kill switch: disable prefix matching outright (every /// prefill recomputes the full prompt). Prefix caching is also forced - /// off while the DSpark drafter is on — the draft lane needs the - /// aux-hidden captures a skipped prefix never produces. + /// off while a speculative decoder is on: DSpark needs aux-hidden + /// captures for every prefix row, while native MTP needs target hidden + /// states and uninterrupted MTP KV continuity. pub no_prefix_cache: bool, /// `Some` adds the pegaflow host tier under the prefix cache: sealed KV /// blocks flow to one shared pinned pool on request release, and a /// prompt whose prefix fell out of HBM restores from it at admission. - /// Requires the prefix cache (rejected at launch alongside the DSpark - /// drafter or `no_prefix_cache`). + /// Requires the prefix cache (rejected at launch alongside any + /// speculative decoder or `no_prefix_cache`). pub kv_offload: Option, /// Launch-time MoE sharding topology. `Ep8` (default) is the /// high-throughput configuration: 32 whole experts per rank, DeepEP @@ -413,7 +444,7 @@ pub fn launch(model_path: &Path, options: Glm52LaunchOptions) -> Result Result Result Result Result Result { let pool_slots = if prefill_only { @@ -619,8 +660,10 @@ fn glm52_cap_bytes( model::GLM52_MAX_BATCH_PER_RANK }; Ok(glm52_arena_bytes(max_model_len, pool_slots, prefill_only)? - + if dspark_enabled { + + if drafter.is_dspark() { crate::dspark::glm52_dspark_arena_bytes(max_model_len) + } else if drafter.is_mtp() { + crate::mtp::glm52_mtp_arena_bytes(max_model_len)? } else { 0 }) @@ -648,12 +691,12 @@ fn glm52_prefill_scratch_reservation( fn derive_max_model_len( requested: Option, min_free_vram_bytes: usize, - dspark_enabled: bool, + drafter: &Glm52Drafter, prefill_scratch_bytes: usize, prefill_only: bool, ) -> Result { let reserve_bytes = GLM52_VRAM_RESERVE_BYTES - + if dspark_enabled { + + if drafter.is_dspark() { GLM52_DSPARK_VRAM_RESERVE_BYTES } else { 0 @@ -677,7 +720,7 @@ fn derive_max_model_len( requested / GLM52_MODEL_LEN_ALIGN * GLM52_MODEL_LEN_ALIGN, requested.next_multiple_of(GLM52_MODEL_LEN_ALIGN), ); - let required = glm52_cap_bytes(requested, dspark_enabled, prefill_only)?; + let required = glm52_cap_bytes(requested, drafter, prefill_only)?; ensure!( required <= budget_bytes, "GLM5.2 --max-model-len {requested} needs {} of cache per rank but only {} \ @@ -694,8 +737,7 @@ fn derive_max_model_len( let (mut lo, mut hi) = (0, GLM52_MAX_CONTEXT / GLM52_MODEL_LEN_ALIGN); while lo < hi { let mid = (lo + hi).div_ceil(2); - if glm52_cap_bytes(mid * GLM52_MODEL_LEN_ALIGN, dspark_enabled, prefill_only)? - <= budget_bytes + if glm52_cap_bytes(mid * GLM52_MODEL_LEN_ALIGN, drafter, prefill_only)? <= budget_bytes { lo = mid; } else { @@ -715,7 +757,7 @@ fn derive_max_model_len( }; Ok(Glm52ContextBudget { max_model_len, - arena_bytes: glm52_cap_bytes(max_model_len, dspark_enabled, prefill_only)?, + arena_bytes: glm52_cap_bytes(max_model_len, drafter, prefill_only)?, reserve_bytes, budget_bytes, }) @@ -757,7 +799,7 @@ struct LoadedGlm52Runtime { fn start_engine( model_path: &Path, options: &Glm52LoadOptions, - dspark_path: Option<&Path>, + drafter: Glm52Drafter, requested_max_model_len: Option, prefill_only: Option, no_prefix_cache: bool, @@ -766,8 +808,7 @@ fn start_engine( weight_staging: bool, dump_graph_png: Option, ) -> Result { - let dspark_enabled = dspark_path.is_some(); - let startup = validate_startup(model_path, options, moe_topo)?; + let startup = validate_startup(model_path, options, moe_topo, drafter.is_mtp())?; let loaded = load_rank_weights_to_gpu(model_path, &startup, moe_topo, weight_staging)?; log::info!( "GLM5.2 load-weight startup complete: ranks={}, rank_plan_tensors={:?}, rank_gpu_tensors={:?}, rank_gpu_bytes={:?}", @@ -791,7 +832,7 @@ fn start_engine( let budget = derive_max_model_len( requested_max_model_len, min_free_vram_bytes.saturating_sub(qa_kva_twin_bytes), - dspark_enabled, + &drafter, glm52_prefill_scratch_reservation(prefill_only)?, prefill_only.is_some(), )?; @@ -818,8 +859,8 @@ fn start_engine( ByteSize(qa_kva_twin_bytes as u64), ByteSize(budget.arena_bytes as u64), model::GLM52_MAX_BATCH_PER_RANK, - if dspark_enabled { - " (dspark lane included)" + if drafter.enabled() { + " (draft lane included)" } else { "" }, @@ -849,7 +890,7 @@ fn start_engine( &loaded.workers, max_model_len, moe_topo, - dspark_enabled, + &drafter, prefill_only.map(|options| options.chunk_size), ) { Ok(rank_arenas) => rank_arenas, @@ -867,7 +908,7 @@ fn start_engine( if prefill_only.is_some() { preflight_prefill_kernels(&loaded.workers)?; } - if let Some(dspark_path) = dspark_path { + if let Some(dspark_path) = drafter.dspark_path() { load_dspark_drafters(&loaded.workers, dspark_path)?; } ensure_post_build_headroom(&loaded.workers)?; @@ -915,7 +956,7 @@ fn start_engine( submit_rx, loaded.workers, &eos_token_ids, - dspark_enabled, + drafter, prefill_only.map(|prefill| prefill.chunk_size), max_model_len, no_prefix_cache, @@ -1064,14 +1105,14 @@ fn build_rank_models( workers: &[Glm52Worker], max_model_len: usize, moe_topo: Glm52MoeTopo, - dspark_enabled: bool, + drafter: &Glm52Drafter, prefill_chunk_size: Option, ) -> Result>> { let build_started = Instant::now(); let responses = workers .iter() .map(|worker| { - worker.build_model_async(max_model_len, moe_topo, dspark_enabled, prefill_chunk_size) + worker.build_model_async(max_model_len, moe_topo, drafter.clone(), prefill_chunk_size) }) .collect::>>()?; let mut rank_arenas = Vec::with_capacity(responses.len()); @@ -1255,6 +1296,7 @@ fn validate_startup( model_path: &Path, options: &Glm52LoadOptions, moe_topo: Glm52MoeTopo, + native_mtp: bool, ) -> Result { let config_path = model_path.join("config.json"); let content = std::fs::read_to_string(&config_path) @@ -1296,7 +1338,7 @@ fn validate_startup( ); let manifest = Glm52WeightManifest::from_model_dir(model_path)?; - let rank_bundles = manifest.all_rank_load_bundles(moe_topo)?; + let rank_bundles = manifest.all_rank_load_bundles(moe_topo, native_mtp)?; let mut rank_tensor_counts = Vec::with_capacity(rank_bundles.len()); let mut rank_expert_ranges = Vec::with_capacity(rank_bundles.len()); for bundle in &rank_bundles { @@ -1428,58 +1470,103 @@ mod max_model_len_tests { /// Free VRAM that budgets exactly a `cap`-token context (exact ledger + /// reserve) — inverted through the same `glm52_cap_bytes` the derivation /// uses, so the tests exercise the policy, not a parallel formula. - fn free_for(cap: usize, dspark: bool, prefill_scratch_bytes: usize) -> usize { + fn free_for(cap: usize, drafter: &Glm52Drafter, prefill_scratch_bytes: usize) -> usize { let reserve = GLM52_VRAM_RESERVE_BYTES - + if dspark { + + if drafter.is_dspark() { GLM52_DSPARK_VRAM_RESERVE_BYTES } else { 0 } + prefill_scratch_bytes; - reserve + glm52_cap_bytes(cap, dspark, false).expect("cap bytes") + reserve + glm52_cap_bytes(cap, drafter, false).expect("cap bytes") } #[test] fn derived_cap_is_aligned_and_scales_with_free_vram() { - let cap = derive_max_model_len(None, free_for(10_048, false, 0), false, 0, false) - .expect("derive") - .max_model_len; + let cap = derive_max_model_len( + None, + free_for(10_048, &Glm52Drafter::None, 0), + &Glm52Drafter::None, + 0, + false, + ) + .expect("derive") + .max_model_len; assert_eq!(cap, 10_048, "exact budget for an aligned cap derives it"); assert!(cap.is_multiple_of(GLM52_MODEL_LEN_ALIGN)); - let larger = derive_max_model_len(None, free_for(50_048, false, 0), false, 0, false) - .expect("derive") - .max_model_len; + let larger = derive_max_model_len( + None, + free_for(50_048, &Glm52Drafter::None, 0), + &Glm52Drafter::None, + 0, + false, + ) + .expect("derive") + .max_model_len; assert!(larger > cap); } #[test] fn dspark_lane_shrinks_the_derived_cap() { - let free = free_for(50_048, false, 0); - let plain = derive_max_model_len(None, free, false, 0, false).expect("derive"); - let dspark = derive_max_model_len(None, free, true, 0, false).expect("derive"); + let free = free_for(50_048, &Glm52Drafter::None, 0); + let plain = + derive_max_model_len(None, free, &Glm52Drafter::None, 0, false).expect("derive"); + let dspark_drafter = Glm52Drafter::Dspark(PathBuf::from("draft")); + let dspark = derive_max_model_len(None, free, &dspark_drafter, 0, false).expect("derive"); assert!( dspark.max_model_len < plain.max_model_len, "dspark cap-scaled cost must shrink the cap" ); } + #[test] + fn native_mtp_lane_shrinks_the_derived_cap() { + let free = free_for(50_048, &Glm52Drafter::None, 0); + let plain = + derive_max_model_len(None, free, &Glm52Drafter::None, 0, false).expect("derive"); + let native_mtp = + derive_max_model_len(None, free, &Glm52Drafter::NativeMtp, 0, false).expect("derive"); + assert!( + native_mtp.max_model_len < plain.max_model_len, + "native MTP cap-scaled KV must shrink the target context cap" + ); + assert!( + glm52_cap_bytes(50_048, &Glm52Drafter::NativeMtp, false).expect("MTP cap bytes") + > glm52_cap_bytes(50_048, &Glm52Drafter::None, false).expect("plain cap bytes"), + "native MTP must be represented in the exact memory ledger" + ); + } + #[test] fn derived_cap_never_exceeds_the_checkpoint_ceiling() { - let budget = derive_max_model_len(None, usize::MAX / 2, false, 0, false).expect("derive"); + let budget = derive_max_model_len(None, usize::MAX / 2, &Glm52Drafter::None, 0, false) + .expect("derive"); assert_eq!(budget.max_model_len, GLM52_MAX_CONTEXT); } #[test] fn too_little_vram_fails_instead_of_serving_a_toy_cap() { - let err = derive_max_model_len(None, free_for(1024, false, 0), false, 0, false) - .expect_err("sub-minimum cap must fail"); + let err = derive_max_model_len( + None, + free_for(1024, &Glm52Drafter::None, 0), + &Glm52Drafter::None, + 0, + false, + ) + .expect_err("sub-minimum cap must fail"); assert!(err.to_string().contains("context cap"), "{err}"); } #[test] fn unaligned_requested_cap_is_rejected_with_the_nearest_valid_values() { - let err = derive_max_model_len(Some(5000), free_for(100_032, false, 0), false, 0, false) - .expect_err("unaligned cap must fail, not silently round"); + let err = derive_max_model_len( + Some(5000), + free_for(100_032, &Glm52Drafter::None, 0), + &Glm52Drafter::None, + 0, + false, + ) + .expect_err("unaligned cap must fail, not silently round"); let message = err.to_string(); assert!( message.contains("4992") && message.contains("5056"), @@ -1489,15 +1576,27 @@ mod max_model_len_tests { #[test] fn requested_cap_beyond_the_budget_fails_at_launch() { - let err = derive_max_model_len(Some(99_968), free_for(10_048, false, 0), false, 0, false) - .expect_err("over-budget cap must fail"); + let err = derive_max_model_len( + Some(99_968), + free_for(10_048, &Glm52Drafter::None, 0), + &Glm52Drafter::None, + 0, + false, + ) + .expect_err("over-budget cap must fail"); assert!(err.to_string().contains("--max-model-len"), "{err}"); } #[test] fn requested_cap_below_the_minimum_fails() { - derive_max_model_len(Some(1024), free_for(100_032, false, 0), false, 0, false) - .expect_err("sub-minimum cap must fail"); + derive_max_model_len( + Some(1024), + free_for(100_032, &Glm52Drafter::None, 0), + &Glm52Drafter::None, + 0, + false, + ) + .expect_err("sub-minimum cap must fail"); } #[test] @@ -1507,10 +1606,11 @@ mod max_model_len_tests { }; let scratch = glm52_prefill_scratch_reservation(Some(prefill)).expect("prefill reservation"); - let free = free_for(100_032, false, 0); - let decode = derive_max_model_len(None, free, false, 0, false).expect("decode budget"); - let prefill = - derive_max_model_len(None, free, false, scratch, true).expect("prefill budget"); + let free = free_for(100_032, &Glm52Drafter::None, 0); + let decode = + derive_max_model_len(None, free, &Glm52Drafter::None, 0, false).expect("decode budget"); + let prefill = derive_max_model_len(None, free, &Glm52Drafter::None, scratch, true) + .expect("prefill budget"); assert!( prefill.max_model_len > decode.max_model_len, "one shared prefill pool must fit a larger per-request cap than eight decode maxima" diff --git a/openinfer-glm52/src/model/mod.rs b/openinfer-glm52/src/model/mod.rs index 3bd89e1c2..38bc49676 100644 --- a/openinfer-glm52/src/model/mod.rs +++ b/openinfer-glm52/src/model/mod.rs @@ -77,8 +77,10 @@ use crate::weights::retype_owned; mod build; mod launch_ahead; +mod mtp; mod step_body; use launch_ahead::Glm52SpeculatedStep; +use mtp::Glm52NativeMtp; use step_body::run_step_body; /// The per-rank slot count and the largest decode bucket. A slot is a batch @@ -285,6 +287,7 @@ pub(crate) fn rope_tables(position: usize) -> (Vec, Vec) { pub(crate) struct Glm52RankModel { layers: Vec, caches: Vec, + mtp: Option, embed: DeviceMatrix, final_norm: DeviceVec, /// Full vocabulary head retained for DSpark and non-greedy sampling. @@ -527,7 +530,7 @@ impl Glm52RankModel { max_model_len: usize, moe_topo: crate::Glm52MoeTopo, attn_shard: Option, - dspark_enabled: bool, + drafter: &crate::Glm52Drafter, prefill_chunk_size: Option, ) -> Result { ensure!( @@ -596,6 +599,10 @@ impl Glm52RankModel { .transpose()?, }); } + let mtp = drafter + .is_mtp() + .then(|| Glm52NativeMtp::build(ctx, w, max_model_len)) + .transpose()?; let embed_raw = w.take_tensor("model.embed_tokens.weight")?; let lm_head_raw = w.take_tensor("lm_head.weight")?; @@ -703,7 +710,7 @@ impl Glm52RankModel { mqa_shape, mla_heads, mla_backend, - dspark_enabled, + drafter.is_dspark(), )?, graph: CudaGraphState::new(), block_table: bucket_table, @@ -790,6 +797,7 @@ impl Glm52RankModel { Ok(Self { layers, caches, + mtp, embed, final_norm, lm_head, @@ -862,6 +870,53 @@ impl Glm52RankModel { ) } + #[allow(clippy::too_many_arguments)] + pub(crate) fn mtp_propose( + &mut self, + ctx: &DeviceContext, + aux: &DeviceContext, + ep: &mut Glm52MoeEpState, + round: &crate::runner::Glm52MtpRound, + ) -> Result> { + let Some(source_bucket) = round.source_bucket() else { + self.mtp + .as_mut() + .context("GLM5.2 native MTP command reached a model without MTP weights")? + .reset_slots(round.resets())?; + return Ok(Vec::new()); + }; + let source_index = self + .buckets + .iter() + .position(|bucket| bucket.rows == source_bucket) + .with_context(|| { + format!( + "GLM5.2 MTP source bucket {source_bucket} is not in \ + {GLM52_DECODE_BUCKETS:?}" + ) + })?; + // Official vLLM feeds MTP the target model return, which is after + // final RMSNorm for GLM5.2. The pre-norm residual is not an + // interchangeable MTP input even when target top-1 is unchanged. + let target_final_normed = &self.buckets[source_index].scratch.final_normed; + let mtp = self + .mtp + .as_mut() + .context("GLM5.2 native MTP command reached a model without MTP weights")?; + mtp.reset_slots(round.resets())?; + mtp.propose( + ctx, + aux, + ep, + &self.embed, + &self.lm_head, + &self.cos_table, + &self.sin_table, + target_final_normed, + round, + ) + } + /// One lock-step step: feed `inputs[row]` = the `(token, position)` each /// forwarded row carries, return the next-token id per ROW (the fused /// greedy argmax, overwritten for the coordinator's `sampling` rows by a diff --git a/openinfer-glm52/src/model/mtp.rs b/openinfer-glm52/src/model/mtp.rs new file mode 100644 index 000000000..4d1be2243 --- /dev/null +++ b/openinfer-glm52/src/model/mtp.rs @@ -0,0 +1,607 @@ +//! Checkpoint-native GLM5.2 MTP serving lane. +//! +//! The target step keeps its final-normalized hidden rows resident. A draft round +//! packs only committed rows, shifts each sequence token one place left, and +//! runs checkpoint layer 78 once to synchronize MTP KV and produce draft 1. +//! Four single-token iterations then recycle the layer's shared-head-normalized +//! hidden. Rejected speculative KV is not copied back: the next committed +//! first pass overwrites it at the same positions. + +use anyhow::Context as _; +use anyhow::Result; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use half::bf16; +use openinfer_core::cuda_graph::CudaGraphState; +use openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_BYTES_PER_TOKEN; +use openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_PAGE_SIZE; +use openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_TOPK; +use openinfer_kernels::ops::Glm52FlashMlaSparseDecode; +use openinfer_kernels::ops::Glm52IndexerCacheLayout; +use openinfer_kernels::ops::argmax_bf16_split_into; +use openinfer_kernels::ops::embedding_rows_into; +use openinfer_kernels::ops::glm52_flashmla_sparse_decode_num_sm_parts; +use openinfer_kernels::ops::rms_norm_rows_into; +use openinfer_kernels::tensor::DeviceContext; +use openinfer_kernels::tensor::DeviceMatrix; + +use super::GLM52_DECODE_BUCKETS; +use super::GLM52_MAX_BATCH_PER_RANK; +use super::INDEX_CACHE_BLOCK; +use super::NUM_SMS; +use super::build; +use super::glm52_pool_blocks; +use super::glm52_table_width; +use super::step_body::glm52_moe_ep_layer; +use crate::bookend::glm52_embed_into; +use crate::bookend::glm52_lm_head_into; +use crate::config::GLM52_HIDDEN; +use crate::config::GLM52_INDEX_HEAD_DIM; +use crate::config::GLM52_MTP_LAYER; +use crate::config::GLM52_RMS_EPS; +use crate::config::GLM52_SM_SCALE; +use crate::config::GLM52_VOCAB; +use crate::indexer::Glm52IndexerScratch; +use crate::layer::Glm52DecodeStep; +use crate::layer::Glm52DecoderLayerWeights; +use crate::layer::Glm52LayerCaches; +use crate::layer::Glm52LayerIndexMode; +use crate::layer::Glm52LayerMlp; +use crate::layer::glm52_layer_attention_half; +use crate::layer::glm52_layer_finish; +use crate::mla_decode::Glm52MlaSchedMetadata; +use crate::mla_decode::glm52_select_mla_backend; +use crate::moe_ep_wo::Glm52MoeEpState; +use crate::mtp::GLM52_MTP_DRAFTS; +use crate::mtp::Glm52MtpBookendWeights; +use crate::mtp::Glm52MtpScratch; +use crate::mtp::glm52_mtp_prepare_into; +use crate::mtp::glm52_mtp_recycle_into; +use crate::rows::Rows; +use crate::runner::Glm52MtpRound; +use crate::scratch::Glm52DecodeScratch; +use crate::weights::Glm52RankGpuWeights; +use crate::weights::retype_owned; + +struct Glm52MtpBucket { + rows: usize, + sched: Glm52MlaSchedMetadata, + scratch: Glm52DecodeScratch, + bookend_scratch: Glm52MtpScratch, + embeds: Rows, + previous: Rows, + decoder_input: Rows, + block_table: CudaSlice, + compute_graph: CudaGraphState, + reuse_graph: CudaGraphState, +} + +pub(super) struct Glm52NativeMtp { + bookend: Glm52MtpBookendWeights, + layer: Glm52DecoderLayerWeights, + cache: Glm52LayerCaches, + buckets: [Glm52MtpBucket; GLM52_DECODE_BUCKETS.len()], + max_model_len: usize, + table_width: usize, + pages_per_slot: usize, + positions: CudaSlice, + cos: CudaSlice, + sin: CudaSlice, + token_ids: CudaSlice, + slot_mapping: CudaSlice, + seq_lens: CudaSlice, + shared_topk: CudaSlice, + committed_lens: [usize; GLM52_MAX_BATCH_PER_RANK], +} + +impl Glm52NativeMtp { + pub(super) fn build( + ctx: &DeviceContext, + weights: &mut Glm52RankGpuWeights, + max_model_len: usize, + ) -> Result { + let prefix = format!("model.layers.{GLM52_MTP_LAYER}"); + let enorm = build::take_bf16_vec( + ctx, + weights, + &format!("{prefix}.enorm.weight"), + GLM52_HIDDEN, + )?; + let hnorm = build::take_bf16_vec( + ctx, + weights, + &format!("{prefix}.hnorm.weight"), + GLM52_HIDDEN, + )?; + let eh_proj_raw = weights.take_tensor(&format!("{prefix}.eh_proj.weight"))?; + ensure!( + eh_proj_raw.len() == 2 * GLM52_HIDDEN * GLM52_HIDDEN * size_of::(), + "GLM5.2 MTP eh_proj byte length drifted" + ); + let eh_proj = DeviceMatrix { + data: retype_owned::(&ctx.stream, eh_proj_raw)?, + rows: GLM52_HIDDEN, + cols: 2 * GLM52_HIDDEN, + }; + let shared_norm = build::take_bf16_vec( + ctx, + weights, + &format!("{prefix}.shared_head.norm.weight"), + GLM52_HIDDEN, + )?; + let bookend = Glm52MtpBookendWeights::new(enorm, hnorm, eh_proj, shared_norm)?; + let layer = build::build_decoder_layer( + ctx, + weights, + GLM52_MTP_LAYER, + crate::Glm52MoeTopo::Ep8, + None, + )?; + + let num_blocks = glm52_pool_blocks(max_model_len, GLM52_MAX_BATCH_PER_RANK); + let table_width = glm52_table_width(max_model_len); + let index_layout = Glm52IndexerCacheLayout { + cache_blocks: num_blocks, + cache_block_size: INDEX_CACHE_BLOCK, + cache_block_stride_bytes: INDEX_CACHE_BLOCK * (GLM52_INDEX_HEAD_DIM + 4), + }; + let backend = glm52_select_mla_backend(crate::config::GLM52_HEADS)?; + let contract = Glm52FlashMlaSparseDecode { + batch_size: GLM52_MAX_BATCH_PER_RANK, + num_blocks, + topk: GLM52_FLASHMLA_SPARSE_TOPK, + num_sm_parts: glm52_flashmla_sparse_decode_num_sm_parts()?, + sm_scale: GLM52_SM_SCALE, + }; + let cache = Glm52LayerCaches { + mla_cache: ctx.stream.alloc_zeros::( + num_blocks + * GLM52_FLASHMLA_SPARSE_PAGE_SIZE + * GLM52_FLASHMLA_SPARSE_BYTES_PER_TOKEN, + )?, + index_k_cache: Some( + ctx.stream + .alloc_zeros::(index_layout.min_cache_bytes()?)?, + ), + }; + + let mut buckets = Vec::with_capacity(GLM52_DECODE_BUCKETS.len()); + for rows in GLM52_DECODE_BUCKETS { + let row_contract = Glm52FlashMlaSparseDecode { + batch_size: rows, + ..contract + }; + let mqa_shape = Glm52IndexerScratch::paged_mqa_shape( + rows, + index_layout, + table_width, + NUM_SMS, + max_model_len, + ); + buckets.push(Glm52MtpBucket { + rows, + sched: Glm52MlaSchedMetadata::new_for_backend( + ctx, + row_contract, + crate::config::GLM52_HEADS, + backend, + )?, + scratch: Glm52DecodeScratch::new_for_backend( + ctx, + &row_contract, + mqa_shape, + crate::config::GLM52_HEADS, + backend, + false, + )?, + bookend_scratch: Glm52MtpScratch::new(ctx, rows)?, + embeds: Rows::zeros(ctx, rows)?, + previous: Rows::zeros(ctx, rows)?, + decoder_input: Rows::zeros(ctx, rows)?, + block_table: ctx.stream.alloc_zeros::(rows * table_width)?, + compute_graph: CudaGraphState::new(), + reuse_graph: CudaGraphState::new(), + }); + } + Ok(Self { + bookend, + layer, + cache, + buckets: buckets + .try_into() + .map_err(|_| anyhow::anyhow!("GLM5.2 MTP bucket count drifted"))?, + max_model_len, + table_width, + pages_per_slot: (max_model_len + 1).div_ceil(GLM52_FLASHMLA_SPARSE_PAGE_SIZE), + positions: ctx.stream.alloc_zeros(GLM52_MAX_BATCH_PER_RANK)?, + cos: ctx + .stream + .alloc_zeros(GLM52_MAX_BATCH_PER_RANK * crate::config::GLM52_ROPE_HALF)?, + sin: ctx + .stream + .alloc_zeros(GLM52_MAX_BATCH_PER_RANK * crate::config::GLM52_ROPE_HALF)?, + token_ids: ctx.stream.alloc_zeros(GLM52_MAX_BATCH_PER_RANK)?, + slot_mapping: ctx.stream.alloc_zeros(GLM52_MAX_BATCH_PER_RANK)?, + seq_lens: ctx.stream.alloc_zeros(GLM52_MAX_BATCH_PER_RANK)?, + shared_topk: ctx + .stream + .alloc_zeros(GLM52_MAX_BATCH_PER_RANK * GLM52_FLASHMLA_SPARSE_TOPK)?, + committed_lens: [0; GLM52_MAX_BATCH_PER_RANK], + }) + } + + pub(super) fn reset_slots(&mut self, resets: &[usize]) -> Result<()> { + for &slot in resets { + ensure!( + slot < GLM52_MAX_BATCH_PER_RANK, + "GLM5.2 MTP reset slot {slot} is outside \ + 0..{GLM52_MAX_BATCH_PER_RANK}" + ); + self.committed_lens[slot] = 0; + } + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + pub(super) fn propose( + &mut self, + ctx: &DeviceContext, + aux: &DeviceContext, + ep: &mut Glm52MoeEpState, + embed: &DeviceMatrix, + lm_head: &DeviceMatrix, + cos_table: &DeviceMatrix, + sin_table: &DeviceMatrix, + target_final_normed: &Rows, + round: &Glm52MtpRound, + ) -> Result> { + let (context_bucket, appends, proposal) = match round { + Glm52MtpRound::Context { + context_bucket, + appends, + .. + } => (*context_bucket, appends.as_slice(), None), + Glm52MtpRound::Propose { + context_bucket, + draft_bucket, + appends, + proposal_slots, + .. + } => ( + *context_bucket, + appends.as_slice(), + Some((*draft_bucket, proposal_slots.as_slice())), + ), + Glm52MtpRound::Reset { .. } => { + unreachable!("reset-only MTP rounds return before target hidden is selected") + } + }; + let proposal_slots = proposal.map_or(&[][..], |(_, slots)| slots); + ensure!( + appends.len() <= context_bucket, + "GLM5.2 MTP context rows {} exceed collective bucket {context_bucket}", + appends.len(), + ); + ensure!( + proposal_slots.windows(2).all(|pair| pair[0] < pair[1]), + "GLM5.2 MTP proposal slots must be strictly ascending" + ); + let context_index = self.bucket_index(context_bucket)?; + for (packed, append) in appends.iter().enumerate() { + ensure!( + append.slot < GLM52_MAX_BATCH_PER_RANK + && append.target_row < target_final_normed.tokens(), + "GLM5.2 MTP append target row {} or slot {} is out of bounds \ + (target rows {}, slots {})", + append.target_row, + append.slot, + target_final_normed.tokens(), + GLM52_MAX_BATCH_PER_RANK, + ); + ensure!( + append.position == self.committed_lens[append.slot], + "GLM5.2 MTP slot {} first-pass position {} != committed {}", + append.slot, + append.position, + self.committed_lens[append.slot] + ); + let src = target_final_normed + .data() + .slice(append.target_row * GLM52_HIDDEN..(append.target_row + 1) * GLM52_HIDDEN); + let mut dst = self.buckets[context_index] + .previous + .data_mut() + .slice_mut(packed * GLM52_HIDDEN..(packed + 1) * GLM52_HIDDEN); + ctx.stream.memcpy_dtod(&src, &mut dst)?; + self.committed_lens[append.slot] += 1; + } + let context_inputs: Vec<(usize, u32, usize)> = appends + .iter() + .map(|append| (append.slot, append.input_token, append.position)) + .collect(); + self.forward( + ctx, + aux, + ep, + embed, + lm_head, + cos_table, + sin_table, + context_index, + &context_inputs, + Glm52LayerIndexMode::Normal, + )?; + let Some((draft_bucket, proposal_slots)) = proposal else { + return Ok(Vec::new()); + }; + ensure!( + proposal_slots.len() <= draft_bucket, + "GLM5.2 MTP proposal rows {} exceed collective bucket {draft_bucket}", + proposal_slots.len(), + ); + let mut last_rows = Vec::with_capacity(proposal_slots.len()); + for &slot in proposal_slots { + let row = appends + .iter() + .rposition(|append| append.slot == slot) + .with_context(|| format!("GLM5.2 MTP proposal slot {slot} has no append"))?; + last_rows.push(row); + } + let context_tokens = self.argmax_host(ctx, context_index)?; + let draft_index = self.bucket_index(draft_bucket)?; + for (packed, (&slot, &context_row)) in proposal_slots.iter().zip(&last_rows).enumerate() { + let src_topk = self.buckets[context_index].scratch.idx.global_slots.slice( + context_row * GLM52_FLASHMLA_SPARSE_TOPK + ..(context_row + 1) * GLM52_FLASHMLA_SPARSE_TOPK, + ); + let mut dst_topk = self.shared_topk.slice_mut( + packed * GLM52_FLASHMLA_SPARSE_TOPK..(packed + 1) * GLM52_FLASHMLA_SPARSE_TOPK, + ); + ctx.stream.memcpy_dtod(&src_topk, &mut dst_topk)?; + let src_hidden = self.buckets[context_index] + .scratch + .final_normed + .data() + .slice(context_row * GLM52_HIDDEN..(context_row + 1) * GLM52_HIDDEN); + let mut dst_hidden = self.buckets[draft_index] + .previous + .data_mut() + .slice_mut(packed * GLM52_HIDDEN..(packed + 1) * GLM52_HIDDEN); + ctx.stream.memcpy_dtod(&src_hidden, &mut dst_hidden)?; + ensure!( + self.committed_lens[slot] < self.max_model_len, + "GLM5.2 MTP slot {slot} exhausted its context cap" + ); + } + + let mut spans = vec![[0u32; GLM52_MTP_DRAFTS]; proposal_slots.len()]; + for (span, &row) in spans.iter_mut().zip(&last_rows) { + span[0] = context_tokens[row]; + } + for iteration in 1..GLM52_MTP_DRAFTS { + let inputs: Vec<(usize, u32, usize)> = proposal_slots + .iter() + .enumerate() + .map(|(row, &slot)| { + ( + slot, + spans[row][iteration - 1], + self.committed_lens[slot] + iteration - 1, + ) + }) + .collect(); + { + let topk_len = proposal_slots.len() * GLM52_FLASHMLA_SPARSE_TOPK; + if topk_len > 0 { + let src = self.shared_topk.slice(..topk_len); + let mut dst = self.buckets[draft_index] + .scratch + .idx + .global_slots + .slice_mut(..topk_len); + ctx.stream.memcpy_dtod(&src, &mut dst)?; + } + } + self.forward( + ctx, + aux, + ep, + embed, + lm_head, + cos_table, + sin_table, + draft_index, + &inputs, + Glm52LayerIndexMode::Reuse, + )?; + let tokens = self.argmax_host(ctx, draft_index)?; + for (row, span) in spans.iter_mut().enumerate() { + span[iteration] = tokens[row]; + let src = self.buckets[draft_index] + .scratch + .final_normed + .data() + .slice(row * GLM52_HIDDEN..(row + 1) * GLM52_HIDDEN); + let mut dst = self.buckets[draft_index] + .previous + .data_mut() + .slice_mut(row * GLM52_HIDDEN..(row + 1) * GLM52_HIDDEN); + ctx.stream.memcpy_dtod(&src, &mut dst)?; + } + } + Ok(spans) + } + + fn bucket_index(&self, rows: usize) -> Result { + self.buckets + .iter() + .position(|bucket| bucket.rows == rows) + .with_context(|| format!("GLM5.2 MTP bucket {rows} is not in {GLM52_DECODE_BUCKETS:?}")) + } + + #[allow(clippy::too_many_arguments)] + fn forward( + &mut self, + ctx: &DeviceContext, + aux: &DeviceContext, + ep: &mut Glm52MoeEpState, + embed: &DeviceMatrix, + lm_head: &DeviceMatrix, + cos_table: &DeviceMatrix, + sin_table: &DeviceMatrix, + bucket_index: usize, + inputs: &[(usize, u32, usize)], + index_mode: Glm52LayerIndexMode, + ) -> Result<()> { + let rows = self.buckets[bucket_index].rows; + let mut tokens = [0u32; GLM52_MAX_BATCH_PER_RANK]; + let mut positions = [0u32; GLM52_MAX_BATCH_PER_RANK]; + let mut seq_lens = [1i32; GLM52_MAX_BATCH_PER_RANK]; + let mut slot_mapping = [0i64; GLM52_MAX_BATCH_PER_RANK]; + let mut pages = vec![0i32; rows * self.table_width]; + for (row, &(slot, token, position)) in inputs.iter().enumerate() { + ensure!( + row < rows && slot < GLM52_MAX_BATCH_PER_RANK && position < self.max_model_len, + "GLM5.2 MTP input row {row}/{rows}, slot \ + {slot}/{GLM52_MAX_BATCH_PER_RANK}, or position \ + {position}/{} is out of bounds", + self.max_model_len, + ); + tokens[row] = token; + positions[row] = position as u32; + seq_lens[row] = (position + 1) as i32; + let page_offset = position / GLM52_FLASHMLA_SPARSE_PAGE_SIZE; + let page = 1 + slot * self.pages_per_slot + page_offset; + slot_mapping[row] = (page * GLM52_FLASHMLA_SPARSE_PAGE_SIZE + + position % GLM52_FLASHMLA_SPARSE_PAGE_SIZE) + as i64; + for logical_page in 0..=page_offset { + pages[row * self.table_width + logical_page] = + (1 + slot * self.pages_per_slot + logical_page) as i32; + } + } + ctx.stream.memcpy_htod(&tokens, &mut self.token_ids)?; + ctx.stream.memcpy_htod(&positions, &mut self.positions)?; + ctx.stream.memcpy_htod(&seq_lens, &mut self.seq_lens)?; + ctx.stream + .memcpy_htod(&slot_mapping, &mut self.slot_mapping)?; + embedding_rows_into(ctx, cos_table, &self.positions, rows, &mut self.cos)?; + embedding_rows_into(ctx, sin_table, &self.positions, rows, &mut self.sin)?; + ctx.stream + .memcpy_htod(&pages, &mut self.buckets[bucket_index].block_table)?; + + let bucket = &mut self.buckets[bucket_index]; + let Glm52MtpBucket { + sched, + scratch, + bookend_scratch, + embeds, + previous, + decoder_input, + block_table, + compute_graph, + reuse_graph, + .. + } = bucket; + let step = Glm52DecodeStep { + mla_cos: &self.cos, + mla_sin: &self.sin, + idx_cos: &self.cos, + idx_sin: &self.sin, + mla_sched: sched, + slot_mapping: &self.slot_mapping, + block_table, + seq_lens: &self.seq_lens, + }; + let graph = match index_mode { + Glm52LayerIndexMode::Normal => compute_graph, + Glm52LayerIndexMode::Reuse => reuse_graph, + }; + graph.run_or_capture(ctx, || { + glm52_embed_into(ctx, embed, &self.token_ids, embeds)?; + glm52_mtp_prepare_into( + ctx, + &self.bookend, + &self.positions, + embeds, + previous, + bookend_scratch, + decoder_input, + )?; + ctx.stream + .memcpy_dtod(decoder_input.data(), scratch.hidden.data_mut())?; + rms_norm_rows_into( + ctx, + scratch.hidden.data(), + &self.layer.input_ln, + GLM52_RMS_EPS, + GLM52_HIDDEN, + rows, + scratch.layer.normed.data_mut(), + )?; + let mut carry_ready = index_mode == Glm52LayerIndexMode::Reuse; + glm52_layer_attention_half( + ctx, + Some(aux), + &self.layer, + &mut self.cache, + &step, + scratch, + &mut carry_ready, + 0, + true, + None, + index_mode, + )?; + let Glm52LayerMlp::MoeEp8(moe) = &self.layer.mlp else { + anyhow::bail!("GLM5.2 MTP layer 78 is not EP MoE") + }; + glm52_moe_ep_layer( + ctx, + aux, + ep, + moe, + scratch, + rows, + crate::weights::GLM52_EP_RANKS * rows, + )?; + glm52_layer_finish(ctx, scratch, 0, false)?; + glm52_mtp_recycle_into( + ctx, + &self.bookend, + &scratch.hidden, + &mut scratch.final_normed, + )?; + glm52_lm_head_into(ctx, &scratch.final_normed, lm_head, &mut scratch.logits)?; + argmax_bf16_split_into( + ctx, + scratch.logits.data(), + rows, + GLM52_VOCAB, + &mut scratch.argmax_partial_values, + &mut scratch.argmax_partial_indices, + &mut scratch.argmax_values, + &mut scratch.argmax_indices, + ) + }) + } + + fn argmax_host(&self, ctx: &DeviceContext, bucket_index: usize) -> Result> { + let bucket = &self.buckets[bucket_index]; + let values = ctx.stream.clone_dtoh(&bucket.scratch.argmax_values)?; + let indices = ctx.stream.clone_dtoh(&bucket.scratch.argmax_indices)?; + values + .iter() + .zip(indices) + .enumerate() + .map(|(row, (value, index))| { + ensure!( + value.to_f32().is_finite() && index >= 0, + "GLM5.2 MTP row {row} produced invalid argmax value {} at index {index}", + value.to_f32(), + ); + u32::try_from(index).context("GLM5.2 MTP argmax does not fit u32") + }) + .collect() + } +} diff --git a/openinfer-glm52/src/model/step_body.rs b/openinfer-glm52/src/model/step_body.rs index b60d54514..affc26845 100644 --- a/openinfer-glm52/src/model/step_body.rs +++ b/openinfer-glm52/src/model/step_body.rs @@ -4,7 +4,7 @@ use anyhow::Context as _; use anyhow::Result; use anyhow::ensure; use cudarc::driver::CudaSlice; -use openinfer_kernels::ops::add_into; +use openinfer_kernels::ops::add_scaled_bf16_into; use openinfer_kernels::ops::argmax_bf16_split_into; use openinfer_kernels::ops::copy_hidden_rows_raw_into; use openinfer_kernels::ops::glm52_vocab_parallel_pack_launch; @@ -26,10 +26,12 @@ use crate::dense::glm52_dense_mlp_forward_into; use crate::layer::Glm52DecodeStep; use crate::layer::Glm52DecoderLayerWeights; use crate::layer::Glm52LayerCaches; +use crate::layer::Glm52LayerIndexMode; use crate::layer::Glm52LayerMlp; use crate::layer::glm52_layer_attention_half; use crate::layer::glm52_layer_finish; use crate::layer::glm52_layer_finish_fused; +use crate::moe_decode::run_ep_router_into; use crate::moe_decode::run_router_into; use crate::moe_ep_wo::Glm52MoeEpState; use crate::moe_ep8::Glm52MoeEp8LayerWeights; @@ -104,6 +106,7 @@ pub(super) fn run_step_body( parity, layer == 0, tp_ar, + Glm52LayerIndexMode::Normal, ) .with_context(|| format!("GLM5.2 layer {layer} attention half"))?; let mut tp_padded_mlp = false; @@ -266,7 +269,7 @@ pub(super) fn run_step_body( /// DeepEP dispatch/expert-GEMM/combine, joined by the closing add into /// `mlp_out`. The events recorded here during capture become graph edges; /// replay keeps the parallel branches. -fn glm52_moe_ep_layer( +pub(super) fn glm52_moe_ep_layer( ctx: &DeviceContext, aux: &DeviceContext, ep8: &mut Glm52MoeEpState, @@ -289,7 +292,7 @@ fn glm52_moe_ep_layer( )?; let shared_done = aux.stream.record_event(None)?; - run_router_into(ctx, &moe.router, s.layer.normed2.data(), &mut s.router)?; + run_ep_router_into(ctx, &moe.router, s.layer.normed2.data(), &mut s.router)?; let dispatched = ep8.routed_forward( ctx, &moe.bank, @@ -302,9 +305,10 @@ fn glm52_moe_ep_layer( ); // Join: the closing add consumes both branches. ctx.stream.wait(&shared_done)?; - add_into( + add_scaled_bf16_into( ctx, ep8.combined(), + crate::config::GLM52_ROUTED_SCALING_FACTOR as f32, s.layer.shared_out.data(), batch * GLM52_HIDDEN, s.layer.mlp_out.data_mut(), diff --git a/openinfer-glm52/src/moe_decode.rs b/openinfer-glm52/src/moe_decode.rs index 78d6f4fb7..cf8c435e0 100644 --- a/openinfer-glm52/src/moe_decode.rs +++ b/openinfer-glm52/src/moe_decode.rs @@ -254,9 +254,9 @@ impl Glm52MoeExpertBank { } } -/// Router output for one token: the top-8 GLOBAL expert ids and their -/// normalized, x2.5-scaled weights, both device-resident (never read back to -/// host). +/// Router output for one token: the top-8 GLOBAL expert ids and weights, both +/// device-resident (never read back to host). The caller chooses whether the +/// model's routed scale is folded into these weights. pub(crate) struct RoutedTopk { pub(crate) topk_idx: CudaSlice, pub(crate) topk_weight: CudaSlice, @@ -296,6 +296,31 @@ pub(crate) fn run_router_into( router: &Glm52MoeRouterWeights, normed_hidden: &CudaSlice, s: &mut Glm52RouterScratch, +) -> Result<()> { + run_router_into_with_config(ctx, router, normed_hidden, s, Glm52RouterConfig::glm52()) +} + +pub(crate) fn run_ep_router_into( + ctx: &DeviceContext, + router: &Glm52MoeRouterWeights, + normed_hidden: &CudaSlice, + s: &mut Glm52RouterScratch, +) -> Result<()> { + run_router_into_with_config( + ctx, + router, + normed_hidden, + s, + Glm52RouterConfig::glm52_unscaled(), + ) +} + +fn run_router_into_with_config( + ctx: &DeviceContext, + router: &Glm52MoeRouterWeights, + normed_hidden: &CudaSlice, + s: &mut Glm52RouterScratch, + config: Glm52RouterConfig, ) -> Result<()> { let mut router_out = Glm52RouterOutput { topk_weight: &mut s.route.topk_weight, @@ -303,7 +328,7 @@ pub(crate) fn run_router_into( }; glm52_router_noaux_tc_launch( ctx, - Glm52RouterConfig::glm52(), + config, Glm52RouterBatch { active_tokens: s.tokens, padded_tokens: s.tokens, @@ -363,15 +388,14 @@ pub(crate) fn run_router_rows_into( ) } -/// Allocating convenience over [`run_router_into`] for the oracle-gate/test -/// paths. +/// Allocating EP/unscaled router for oracle-gate/test paths. #[cfg(test)] -pub(crate) fn run_router( +pub(crate) fn run_ep_router( ctx: &DeviceContext, router: &Glm52MoeRouterWeights, normed_hidden: &CudaSlice, ) -> Result { let mut s = Glm52RouterScratch::new(ctx, 1)?; - run_router_into(ctx, router, normed_hidden, &mut s)?; + run_ep_router_into(ctx, router, normed_hidden, &mut s)?; Ok(s.route) } diff --git a/openinfer-glm52/src/moe_ep8.rs b/openinfer-glm52/src/moe_ep8.rs index 1ad21ab2f..2dcb5360c 100644 --- a/openinfer-glm52/src/moe_ep8.rs +++ b/openinfer-glm52/src/moe_ep8.rs @@ -10,8 +10,8 @@ //! dispatch(x bf16, global topk) # collective; recv = expert-major //! → metadata: psum i32 → offsets i64 + masked_m + row_map //! → re-quant recv rows → masked [32, 64, k] fp8 + mn-major scales -//! → DeepGEMM masked W13 (32 groups) → weighted SiLU·quant (recv weights) -//! → DeepGEMM masked W2 → remap masked→aligned slots +//! → DeepGEMM masked W13 (32 groups) → SiLU·quant +//! → DeepGEMM masked W2 → route-weighted remap to aligned slots //! → combine # collective; sums slots per token //! ``` //! @@ -51,7 +51,7 @@ use openinfer_kernels::ops::glm52_deepgemm_grouped_fp8_metadata_launch; use openinfer_kernels::ops::glm52_deepgemm_masked_grouped_fp8_launch; use openinfer_kernels::ops::glm52_deepgemm_masked_out_to_aligned_launch; use openinfer_kernels::ops::glm52_fp8_per_token_group_quant_bf16_masked_launch; -use openinfer_kernels::ops::glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_launch; +use openinfer_kernels::ops::glm52_silu_and_mul_per_token_group_quant_bf16_masked_launch; use openinfer_kernels::tensor::DeviceContext; use crate::moe_decode::EXPERTS; @@ -191,8 +191,9 @@ impl Glm52MoeEp8State { /// enter simultaneously per layer. A rank with tokens passes its /// post-attention normed hidden rows + router output (`[T, HIDDEN]` / /// `[T, 8]`) and the row count; on `Ok(true)` the routed output -/// `[T, HIDDEN]` (route weight and ×2.5 scaling already folded) is in -/// `state.combined()`. A token-less rank passes `None` and gets `Ok(false)`. +/// `[T, HIDDEN]` (normalized route weights folded, routed scale not applied) +/// is in `state.combined()`. A token-less rank passes `None` and gets +/// `Ok(false)`. /// The DP8 production path always passes `Some` (pad rows are dispatched like /// real ones); `None` survives for the EP8 layer oracle gate's /// single-dispatcher replay. @@ -314,10 +315,9 @@ pub(crate) fn glm52_moe_ep8_routed_forward( &mut state.w13_out_masked, )?; - // Weighted SwiGLU quant: silu(gate)*up*route_weight → fp8 W2 input. The - // gate|up rows are already masked (the W13 GEMM wrote them there); the - // per-slot weight is exactly what dispatch delivered per expanded row. - glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_launch( + // SwiGLU quant: vLLM quantizes the unweighted activation before W2 and + // applies the router weight to W2's BF16 output. + glm52_silu_and_mul_per_token_group_quant_bf16_masked_launch( ctx, Glm52MoeQuantShape { rows: bound_rows, @@ -327,7 +327,6 @@ pub(crate) fn glm52_moe_ep8_routed_forward( GLM52_DEEPGEMM_MASKED_GROUPS, GLM52_DEEPGEMM_MASKED_CAP, &state.w13_out_masked, - &state.recv_topk_weight, &mut state.w2_act_masked, &mut state.w2_act_scale_masked, &state.expert_offsets, @@ -347,13 +346,15 @@ pub(crate) fn glm52_moe_ep8_routed_forward( &mut state.expert_out_masked, )?; - // Masked GEMM output → the aligned recv slots decode_combine addresses. + // Masked GEMM output × router weight → the aligned recv slots + // decode_combine addresses. glm52_deepgemm_masked_out_to_aligned_launch( ctx, W2_N, &state.expert_out_masked, &state.masked_m, &state.expert_offsets, + &state.recv_topk_weight, &mut state.expert_out, )?; diff --git a/openinfer-glm52/src/mtp.rs b/openinfer-glm52/src/mtp.rs new file mode 100644 index 000000000..ab5d8bdc5 --- /dev/null +++ b/openinfer-glm52/src/mtp.rs @@ -0,0 +1,272 @@ +//! GLM5.2 MTP layer-78 accuracy-oracle bookends. +//! +//! The checkpoint's MTP decoder block is the same concrete decoder-layer +//! implementation as the target stack. This module owns only the math unique +//! to MTP: +//! +//! ```text +//! embed = where(position == 0, 0, embed) +//! decoder_input = eh_proj(cat(enorm(embed), hnorm(previous_hidden))) +//! raw_hidden = decoder_layer_78(decoder_input) +//! recycle_hidden = shared_head.norm(raw_hidden) +//! logits = lm_head(shared_head.norm(raw_hidden)) +//! ``` +//! +//! `raw_hidden` must remain available for target-head logits. The normalized +//! value is recycled into the next draft iteration; normalizing in place +//! would apply the shared norm twice on the logits path. +//! +//! Production serving owns residency and state in `model::mtp`; the oracle +//! tests call these same bookend operations directly. + +use anyhow::Context as _; +use anyhow::Result; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use half::bf16; +use openinfer_kernels::ops::copy_hidden_rows_raw_into; +use openinfer_kernels::ops::gemm_strided_batched_bf16; +use openinfer_kernels::ops::mask_position_zero_rows_into; +use openinfer_kernels::ops::rms_norm_rows_into; +use openinfer_kernels::tensor::DeviceContext; +use openinfer_kernels::tensor::DeviceMatrix; +use openinfer_kernels::tensor::DeviceVec; +use openinfer_kernels::tensor::HiddenStates; + +use crate::config::GLM52_HIDDEN; +use crate::config::GLM52_INDEX_HEAD_DIM; +use crate::config::GLM52_RMS_EPS; +use crate::model::GLM52_DECODE_BUCKETS; +use crate::model::GLM52_MAX_BATCH_PER_RANK; +use crate::model::GLM52_MODEL_LEN_ALIGN; +use crate::model::glm52_pool_blocks; +use crate::rows::Rows; + +const MTP_FUSED_INPUT: usize = 2 * GLM52_HIDDEN; +pub(crate) const GLM52_MTP_DRAFTS: usize = 5; + +/// Context-scaled device memory owned by the native MTP lane: one layer of +/// MLA + index-K cache and one set of per-bucket indexer logits/block tables. +/// Fixed-size weights and scratch are accounted by the post-build headroom +/// probe; this function is the exact monotone term used to derive the context +/// cap before those arenas are allocated. +pub(crate) fn glm52_mtp_arena_bytes(max_model_len: usize) -> Result { + let blocks = glm52_pool_blocks(max_model_len, GLM52_MAX_BATCH_PER_RANK); + let mla = blocks + .checked_mul(GLM52_MODEL_LEN_ALIGN) + .and_then(|v| v.checked_mul(openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_BYTES_PER_TOKEN)) + .context("GLM5.2 MTP MLA arena byte count overflow")?; + let index_k = blocks + .checked_mul(GLM52_MODEL_LEN_ALIGN) + .and_then(|v| v.checked_mul(GLM52_INDEX_HEAD_DIM + size_of::())) + .context("GLM5.2 MTP index-K arena byte count overflow")?; + let rows: usize = GLM52_DECODE_BUCKETS.iter().sum(); + let indexer_logits = rows + .checked_mul(max_model_len.next_multiple_of(256)) + .and_then(|v| v.checked_mul(size_of::() + size_of::())) + .context("GLM5.2 MTP indexer scratch byte count overflow")?; + let block_tables = rows + .checked_mul(max_model_len.div_ceil(GLM52_MODEL_LEN_ALIGN)) + .and_then(|v| v.checked_mul(size_of::())) + .context("GLM5.2 MTP block-table byte count overflow")?; + mla.checked_add(index_k) + .and_then(|v| v.checked_add(indexer_logits)) + .and_then(|v| v.checked_add(block_tables)) + .context("GLM5.2 MTP arena byte count overflow") +} + +/// The four BF16 weights around the ordinary layer-78 decoder block. +pub(crate) struct Glm52MtpBookendWeights { + enorm: DeviceVec, + hnorm: DeviceVec, + eh_proj: DeviceMatrix, + shared_norm: DeviceVec, +} + +impl Glm52MtpBookendWeights { + pub(crate) fn new( + enorm: DeviceVec, + hnorm: DeviceVec, + eh_proj: DeviceMatrix, + shared_norm: DeviceVec, + ) -> Result { + ensure!( + enorm.len == GLM52_HIDDEN, + "GLM5.2 MTP enorm must be [{GLM52_HIDDEN}], got [{}]", + enorm.len + ); + ensure!( + hnorm.len == GLM52_HIDDEN, + "GLM5.2 MTP hnorm must be [{GLM52_HIDDEN}], got [{}]", + hnorm.len + ); + ensure!( + eh_proj.rows == GLM52_HIDDEN && eh_proj.cols == MTP_FUSED_INPUT, + "GLM5.2 MTP eh_proj must be [{GLM52_HIDDEN}, {MTP_FUSED_INPUT}], got [{}, {}]", + eh_proj.rows, + eh_proj.cols + ); + ensure!( + shared_norm.len == GLM52_HIDDEN, + "GLM5.2 MTP shared norm must be [{GLM52_HIDDEN}], got [{}]", + shared_norm.len + ); + Ok(Self { + enorm, + hnorm, + eh_proj, + shared_norm, + }) + } + + #[cfg(test)] + pub(crate) fn from_host( + ctx: &DeviceContext, + enorm: &[u8], + hnorm: &[u8], + eh_proj: &[u8], + shared_norm: &[u8], + ) -> Result { + Self::new( + DeviceVec::from_safetensors(ctx, enorm)?, + DeviceVec::from_safetensors(ctx, hnorm)?, + DeviceMatrix::from_safetensors(ctx, eh_proj, GLM52_HIDDEN, MTP_FUSED_INPUT)?, + DeviceVec::from_safetensors(ctx, shared_norm)?, + ) + } +} + +/// Persistent MTP-only intermediates for one row bucket. +pub(crate) struct Glm52MtpScratch { + masked_embed: Rows, + normed_embed: Rows, + normed_previous: Rows, + fused_input: HiddenStates, +} + +impl Glm52MtpScratch { + pub(crate) fn new(ctx: &DeviceContext, tokens: usize) -> Result { + Ok(Self { + masked_embed: Rows::zeros(ctx, tokens)?, + normed_embed: Rows::zeros(ctx, tokens)?, + normed_previous: Rows::zeros(ctx, tokens)?, + fused_input: HiddenStates::zeros(ctx, MTP_FUSED_INPUT, tokens)?, + }) + } + + #[cfg(test)] + pub(crate) fn normed_embed(&self) -> &Rows { + &self.normed_embed + } + + #[cfg(test)] + pub(crate) fn normed_previous(&self) -> &Rows { + &self.normed_previous + } +} + +/// Build the ordinary layer-78 decoder input. One GEMM consumes the physical +/// concatenation so its accumulation and BF16 output boundary match vLLM's +/// `nn.Linear(torch.cat(...))`. +pub(crate) fn glm52_mtp_prepare_into( + ctx: &DeviceContext, + w: &Glm52MtpBookendWeights, + positions: &CudaSlice, + inputs_embeds: &Rows, + previous_hidden: &Rows, + s: &mut Glm52MtpScratch, + decoder_input: &mut Rows, +) -> Result<()> { + let tokens = inputs_embeds.tokens(); + ensure!( + previous_hidden.tokens() == tokens + && s.masked_embed.tokens() == tokens + && decoder_input.tokens() == tokens, + "GLM5.2 MTP row bucket mismatch" + ); + mask_position_zero_rows_into( + ctx, + inputs_embeds.data(), + positions, + GLM52_HIDDEN, + tokens, + s.masked_embed.data_mut(), + )?; + rms_norm_rows_into( + ctx, + s.masked_embed.data(), + &w.enorm, + GLM52_RMS_EPS, + GLM52_HIDDEN, + tokens, + s.normed_embed.data_mut(), + )?; + rms_norm_rows_into( + ctx, + previous_hidden.data(), + &w.hnorm, + GLM52_RMS_EPS, + GLM52_HIDDEN, + tokens, + s.normed_previous.data_mut(), + )?; + copy_hidden_rows_raw_into( + ctx, + s.normed_embed.data(), + GLM52_HIDDEN, + &mut s.fused_input.data, + MTP_FUSED_INPUT, + 0, + tokens, + )?; + copy_hidden_rows_raw_into( + ctx, + s.normed_previous.data(), + GLM52_HIDDEN, + &mut s.fused_input.data, + MTP_FUSED_INPUT, + GLM52_HIDDEN, + tokens, + )?; + gemm_strided_batched_bf16( + ctx, + true, + false, + GLM52_HIDDEN, + tokens, + MTP_FUSED_INPUT, + &w.eh_proj.data, + MTP_FUSED_INPUT, + 0, + &s.fused_input.data, + MTP_FUSED_INPUT, + 0, + decoder_input.data_mut(), + GLM52_HIDDEN, + 0, + 1, + ) +} + +/// Normalize layer 78's raw residual output for the next MTP iteration. +/// Callers retain `raw_hidden` unchanged for the shared target lm_head path. +pub(crate) fn glm52_mtp_recycle_into( + ctx: &DeviceContext, + w: &Glm52MtpBookendWeights, + raw_hidden: &Rows, + recycle_hidden: &mut Rows, +) -> Result<()> { + ensure!( + raw_hidden.tokens() == recycle_hidden.tokens(), + "GLM5.2 MTP recycle row bucket mismatch" + ); + rms_norm_rows_into( + ctx, + raw_hidden.data(), + &w.shared_norm, + GLM52_RMS_EPS, + GLM52_HIDDEN, + raw_hidden.tokens(), + recycle_hidden.data_mut(), + ) +} diff --git a/openinfer-glm52/src/oracle/layer_ep4.rs b/openinfer-glm52/src/oracle/layer_ep4.rs index edce187f5..a89a946d8 100644 --- a/openinfer-glm52/src/oracle/layer_ep4.rs +++ b/openinfer-glm52/src/oracle/layer_ep4.rs @@ -23,7 +23,7 @@ use openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_PAGE_SIZE; use openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_TOPK; use openinfer_kernels::ops::Glm52FlashMlaSparseDecode; use openinfer_kernels::ops::Glm52IndexerCacheLayout; -use openinfer_kernels::ops::add_into; +use openinfer_kernels::ops::add_scaled_bf16_into; use openinfer_kernels::ops::glm52_ep_deepep_unique_id; use openinfer_kernels::ops::glm52_flashmla_sparse_decode_num_sm_parts; use openinfer_kernels::tensor::DeviceContext; @@ -57,7 +57,7 @@ use crate::model::INDEX_CACHE_BLOCK; use crate::model::NUM_SMS; use crate::model::rope_tables; use crate::moe_decode::HIDDEN; -use crate::moe_decode::run_router; +use crate::moe_decode::run_ep_router; use crate::moe_ep_wo::Glm52MoeEpWoState; use crate::moe_ep_wo::glm52_moe_ep_wo_routed_forward; use crate::scratch::Glm52DecodeScratch; @@ -268,8 +268,9 @@ fn run_layer_prefill_ep4( 0, true, None, + crate::layer::Glm52LayerIndexMode::Normal, )?; - let route = run_router(ctx, &moe.router, scratch.layer.normed2.data())?; + let route = run_ep_router(ctx, &moe.router, scratch.layer.normed2.data())?; let dispatched = glm52_moe_ep_wo_routed_forward( ctx, ep4, @@ -284,9 +285,10 @@ fn run_layer_prefill_ep4( &mut scratch.shared_mlp, scratch.layer.shared_out.data_mut(), )?; - add_into( + add_scaled_bf16_into( ctx, ep4.combined(), + crate::config::GLM52_ROUTED_SCALING_FACTOR as f32, scratch.layer.shared_out.data(), HIDDEN, scratch.layer.mlp_out.data_mut(), diff --git a/openinfer-glm52/src/oracle/layer_ep8.rs b/openinfer-glm52/src/oracle/layer_ep8.rs index 356cbcf89..41519e49f 100644 --- a/openinfer-glm52/src/oracle/layer_ep8.rs +++ b/openinfer-glm52/src/oracle/layer_ep8.rs @@ -19,7 +19,7 @@ use openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_PAGE_SIZE; use openinfer_kernels::ops::GLM52_FLASHMLA_SPARSE_TOPK; use openinfer_kernels::ops::Glm52FlashMlaSparseDecode; use openinfer_kernels::ops::Glm52IndexerCacheLayout; -use openinfer_kernels::ops::add_into; +use openinfer_kernels::ops::add_scaled_bf16_into; use openinfer_kernels::ops::glm52_ep_deepep_unique_id; use openinfer_kernels::ops::glm52_flashmla_sparse_decode_num_sm_parts; use openinfer_kernels::tensor::DeviceContext; @@ -53,7 +53,7 @@ use crate::model::INDEX_CACHE_BLOCK; use crate::model::NUM_SMS; use crate::model::rope_tables; use crate::moe_decode::HIDDEN; -use crate::moe_decode::run_router; +use crate::moe_decode::run_ep_router; use crate::moe_ep8::Glm52MoeEp8State; use crate::moe_ep8::glm52_moe_ep8_routed_forward; use crate::scratch::Glm52DecodeScratch; @@ -127,7 +127,7 @@ fn layer_moe_ep8_oracle_gate() -> Result<()> { let mut ep8 = Glm52MoeEp8State::new(&ctx, &unique_id, EP_RANKS, 0)?; // Replay the layer once per global-token bucket, in the same order as the // expert threads' collective loops. - let outputs: Result>> = GLOBAL_TOKEN_BUCKETS + let outputs: Result> = GLOBAL_TOKEN_BUCKETS .into_iter() .map(|global_tokens| { run_layer_prefill_ep8( @@ -155,7 +155,7 @@ fn layer_moe_ep8_oracle_gate() -> Result<()> { for (outputs, global_tokens) in outputs?.iter().zip(GLOBAL_TOKEN_BUCKETS) { assert_layer_probes( &format!("layer6/moe/ep8/g{global_tokens}"), - outputs, + &outputs.hidden, MOE_ORACLE_LAYER_PROBES, MOE_ORACLE_LAYER_TOL, 4, @@ -167,14 +167,52 @@ fn layer_moe_ep8_oracle_gate() -> Result<()> { /// The EP8 variant of the gate's prefill-via-decode walk: same decode /// environment as `oracle::layer::run_layer_prefill`, with the MLP half /// driven through the collective. -fn run_layer_prefill_ep8( +pub(super) struct LayerEp8Outputs { + pub(super) hidden: Vec, + pub(super) post_attention: Vec, + pub(super) mlp: Vec, +} + +fn run_moe_ep8_half( + ctx: &DeviceContext, + moe: &crate::moe_ep8::Glm52MoeEp8LayerWeights, + ep8: &mut Glm52MoeEp8State, + scratch: &mut Glm52DecodeScratch, + global_tokens: usize, +) -> Result<()> { + let route = run_ep_router(ctx, &moe.router, scratch.layer.normed2.data())?; + let dispatched = glm52_moe_ep8_routed_forward( + ctx, + ep8, + &moe.bank, + Some((scratch.layer.normed2.data(), &route, 1)), + global_tokens, + )?; + ensure!(dispatched, "rank-0 EP8 MoE returned no combined output"); + moe.shared.forward_into( + ctx, + scratch.layer.normed2.data(), + &mut scratch.shared_mlp, + scratch.layer.shared_out.data_mut(), + )?; + add_scaled_bf16_into( + ctx, + ep8.combined(), + crate::config::GLM52_ROUTED_SCALING_FACTOR as f32, + scratch.layer.shared_out.data(), + HIDDEN, + scratch.layer.mlp_out.data_mut(), + ) +} + +pub(super) fn run_layer_prefill_ep8( ctx: &DeviceContext, w: &crate::layer::Glm52DecoderLayerWeights, ep8: &mut Glm52MoeEp8State, hidden_host: &[bf16], oracle_ctx: usize, global_tokens: usize, -) -> Result> { +) -> Result { let Glm52LayerMlp::MoeEp8(moe) = &w.mlp else { anyhow::bail!("ep8 gate requires the MoeEp8 layer weights"); }; @@ -221,7 +259,9 @@ fn run_layer_prefill_ep8( let mut scratch = Glm52DecodeScratch::new(ctx, &contract, mqa_shape, crate::config::GLM52_HEADS, false)?; - let mut outputs = Vec::with_capacity(oracle_ctx * HIDDEN); + let mut hidden_outputs = Vec::with_capacity(oracle_ctx * HIDDEN); + let mut post_attention_outputs = Vec::with_capacity(oracle_ctx * HIDDEN); + let mut mlp_outputs = Vec::with_capacity(oracle_ctx * HIDDEN); for position in 0..oracle_ctx { ctx.stream.memcpy_htod( &hidden_host[position * HIDDEN..(position + 1) * HIDDEN], @@ -268,8 +308,119 @@ fn run_layer_prefill_ep8( 0, true, None, + crate::layer::Glm52LayerIndexMode::Normal, + )?; + run_moe_ep8_half(ctx, moe, ep8, &mut scratch, global_tokens)?; + let post_attention_host = ctx.stream.clone_dtoh(scratch.layer.attn[0].data())?; + post_attention_outputs.extend(post_attention_host.iter().map(|v| v.to_f32())); + let mlp_host = ctx.stream.clone_dtoh(scratch.layer.mlp_out.data())?; + mlp_outputs.extend(mlp_host.iter().map(|v| v.to_f32())); + glm52_layer_finish(ctx, &mut scratch, 0, false)?; + let out_host = ctx.stream.clone_dtoh(scratch.hidden.data())?; + hidden_outputs.extend(out_host.iter().map(|v| v.to_f32())); + } + Ok(LayerEp8Outputs { + hidden: hidden_outputs, + post_attention: post_attention_outputs, + mlp: mlp_outputs, + }) +} + +pub(super) struct MoeEp8RowsOutputs { + pub(super) mlp: Vec, + pub(super) normed: Vec, + pub(super) topk_ids: Vec, + pub(super) topk_weights: Vec, + pub(super) routed: Vec, + pub(super) shared_gate_up: Vec, + pub(super) shared_silu: Vec, + pub(super) shared: Vec, +} + +pub(super) fn run_moe_ep8_rows( + ctx: &DeviceContext, + w: &crate::layer::Glm52DecoderLayerWeights, + ep8: &mut Glm52MoeEp8State, + post_attention_host: &[bf16], + reference_normed_host: &[bf16], + rows: usize, + global_tokens: usize, +) -> Result { + ensure!( + post_attention_host.len() == rows * HIDDEN, + "EP8 MoE oracle input has {} elements, expected {}", + post_attention_host.len(), + rows * HIDDEN + ); + ensure!( + reference_normed_host.len() == rows * HIDDEN, + "EP8 MoE oracle reference norm has {} elements, expected {}", + reference_normed_host.len(), + rows * HIDDEN + ); + let Glm52LayerMlp::MoeEp8(moe) = &w.mlp else { + anyhow::bail!("EP8 MoE oracle requires MoeEp8 layer weights"); + }; + let contract = Glm52FlashMlaSparseDecode { + batch_size: 1, + num_blocks: 1, + topk: GLM52_FLASHMLA_SPARSE_TOPK, + num_sm_parts: glm52_flashmla_sparse_decode_num_sm_parts()?, + sm_scale: GLM52_SM_SCALE, + }; + let index_cache_layout = Glm52IndexerCacheLayout { + cache_blocks: 1, + cache_block_size: INDEX_CACHE_BLOCK, + cache_block_stride_bytes: INDEX_CACHE_BLOCK * (GLM52_INDEX_HEAD_DIM + 4), + }; + let mqa_shape = + Glm52IndexerScratch::paged_mqa_shape(1, index_cache_layout, 1, NUM_SMS, rows.max(1)); + let mut scratch = + Glm52DecodeScratch::new(ctx, &contract, mqa_shape, crate::config::GLM52_HEADS, false)?; + let mut outputs = MoeEp8RowsOutputs { + mlp: Vec::with_capacity(post_attention_host.len()), + normed: Vec::with_capacity(post_attention_host.len()), + topk_ids: Vec::with_capacity(rows * crate::config::GLM52_TOPK), + topk_weights: Vec::with_capacity(rows * crate::config::GLM52_TOPK), + routed: Vec::with_capacity(post_attention_host.len()), + shared_gate_up: Vec::with_capacity( + rows * 2 * crate::moe_decode::GLM52_SHARED_EXPERT_INTERMEDIATE, + ), + shared_silu: Vec::with_capacity(rows * crate::moe_decode::GLM52_SHARED_EXPERT_INTERMEDIATE), + shared: Vec::with_capacity(post_attention_host.len()), + }; + for (row, reference_normed) in post_attention_host + .chunks_exact(HIDDEN) + .zip(reference_normed_host.chunks_exact(HIDDEN)) + { + ctx.stream.memcpy_htod(row, scratch.hidden.data_mut())?; + openinfer_kernels::ops::rms_norm_rows_into( + ctx, + scratch.hidden.data(), + &w.post_attn_ln, + crate::config::GLM52_RMS_EPS, + HIDDEN, + 1, + scratch.layer.normed2.data_mut(), )?; - let route = run_router(ctx, &moe.router, scratch.layer.normed2.data())?; + let normed = ctx.stream.clone_dtoh(scratch.layer.normed2.data())?; + outputs + .normed + .extend(normed.iter().map(|value| value.to_f32())); + // Feed the exact vLLM norm output into all downstream stages. This + // keeps the RMSNorm delta from contaminating router and expert-kernel + // diagnostics. + ctx.stream + .memcpy_htod(reference_normed, scratch.layer.normed2.data_mut())?; + let route = run_ep_router(ctx, &moe.router, scratch.layer.normed2.data())?; + let topk_ids = ctx.stream.clone_dtoh(&route.topk_idx)?; + outputs + .topk_ids + .extend_from_slice(&topk_ids[..crate::config::GLM52_TOPK]); + let topk_weights = ctx.stream.clone_dtoh(&route.topk_weight)?; + outputs + .topk_weights + .extend_from_slice(&topk_weights[..crate::config::GLM52_TOPK]); let dispatched = glm52_moe_ep8_routed_forward( ctx, ep8, @@ -278,22 +429,41 @@ fn run_layer_prefill_ep8( global_tokens, )?; ensure!(dispatched, "rank-0 EP8 MoE returned no combined output"); + let routed = ctx.stream.clone_dtoh(ep8.combined())?; + outputs.routed.extend(routed[..HIDDEN].iter().map(|value| { + bf16::from_f32(value.to_f32() * crate::config::GLM52_ROUTED_SCALING_FACTOR as f32) + .to_f32() + })); moe.shared.forward_into( ctx, scratch.layer.normed2.data(), &mut scratch.shared_mlp, scratch.layer.shared_out.data_mut(), )?; - add_into( + let shared_gate_up = ctx.stream.clone_dtoh(scratch.shared_mlp.gate_up())?; + outputs + .shared_gate_up + .extend(shared_gate_up.iter().map(|value| value.to_f32())); + let shared_silu = ctx.stream.clone_dtoh(scratch.shared_mlp.silu_out())?; + outputs + .shared_silu + .extend(shared_silu.iter().map(|value| value.to_f32())); + let shared = ctx.stream.clone_dtoh(scratch.layer.shared_out.data())?; + outputs + .shared + .extend(shared.iter().map(|value| value.to_f32())); + add_scaled_bf16_into( ctx, ep8.combined(), + crate::config::GLM52_ROUTED_SCALING_FACTOR as f32, scratch.layer.shared_out.data(), HIDDEN, scratch.layer.mlp_out.data_mut(), )?; - glm52_layer_finish(ctx, &mut scratch, 0, false)?; - let out_host = ctx.stream.clone_dtoh(scratch.hidden.data())?; - outputs.extend(out_host.iter().map(|v| v.to_f32())); + let output = ctx.stream.clone_dtoh(scratch.layer.mlp_out.data())?; + outputs + .mlp + .extend(output.iter().map(|value| value.to_f32())); } Ok(outputs) } diff --git a/openinfer-glm52/src/oracle/layer_tp8.rs b/openinfer-glm52/src/oracle/layer_tp8.rs index b7239345b..510bbfc0c 100644 --- a/openinfer-glm52/src/oracle/layer_tp8.rs +++ b/openinfer-glm52/src/oracle/layer_tp8.rs @@ -290,6 +290,7 @@ fn run_layer_prefill_tp8( 0, true, None, + crate::layer::Glm52LayerIndexMode::Normal, )?; // The production TP8 arm verbatim: router on the real path, then the // replicated kernel writes routed + shared into all 8 rows. diff --git a/openinfer-glm52/src/oracle/mod.rs b/openinfer-glm52/src/oracle/mod.rs index 964af6a46..8f438466c 100644 --- a/openinfer-glm52/src/oracle/mod.rs +++ b/openinfer-glm52/src/oracle/mod.rs @@ -24,5 +24,7 @@ mod layer_ep4; mod layer_ep8; mod layer_tp8; mod mla; +mod mtp; +mod mtp_production; mod sparse_mla_probe; mod tp8_ar; diff --git a/openinfer-glm52/src/oracle/mtp.rs b/openinfer-glm52/src/oracle/mtp.rs new file mode 100644 index 000000000..8c39d0719 --- /dev/null +++ b/openinfer-glm52/src/oracle/mtp.rs @@ -0,0 +1,680 @@ +//! Official-vLLM golden gates for the GLM5.2 MTP accuracy bring-up. +//! +//! The fixture is a five-row prompt forward captured from the official vLLM +//! nightly at commit `dcfebf93f4eccf30f71872283331eee757915daf`. It covers +//! the position-zero embedding mask, both input norms, the physical concat + +//! single BF16 `eh_proj` GEMM, layer 78, the shared-head recycle norm, and +//! sampled-row logits. +//! +//! Bookend operators are compared directly and should be bit-exact apart from +//! the cuBLAS GEMM tail. The full-layer reference runs vLLM's TP8 attention + +//! sequence-parallel MoE, while OpenInfer runs full attention + EP8 MoE. +//! Their reduction trees differ, so that gate bounds the hidden-state delta +//! and requires the draft top-1, top-8 set, and at least 30/32 top logits to +//! agree instead of pretending the two distributed topologies are bitwise +//! comparable. + +use std::path::Path; +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use anyhow::ensure; +use half::bf16; +use openinfer_kernels::ops::glm52_ep_deepep_unique_id; +use openinfer_kernels::tensor::DeviceContext; +use openinfer_kernels::tensor::DeviceMatrix; +use safetensors::Dtype; +use safetensors::SafeTensors; + +use super::layer::GateLayerMlp; +use super::layer::LayerTensors; +use super::layer::load_decoder_layer; +use super::layer::load_rank_expert_bank; +use super::layer::model_path; +use super::layer_ep8::run_layer_prefill_ep8; +use super::layer_ep8::run_moe_ep8_rows; +use crate::bookend::glm52_lm_head_into; +use crate::config::GLM52_HIDDEN; +use crate::config::GLM52_MTP_LAYER; +use crate::config::GLM52_VOCAB; +use crate::moe_ep8::Glm52MoeEp8State; +use crate::moe_ep8::glm52_moe_ep8_routed_forward; +use crate::mtp::Glm52MtpBookendWeights; +use crate::mtp::Glm52MtpScratch; +use crate::mtp::glm52_mtp_prepare_into; +use crate::mtp::glm52_mtp_recycle_into; +use crate::rows::Rows; +use crate::weights::Glm52WeightManifest; +use crate::weights::mmap_file; + +const ROWS: usize = 5; +const EP_RANKS: usize = 8; +const VLLM_COMMIT: &str = "dcfebf93f4eccf30f71872283331eee757915daf"; +const MODEL_CONFIG_SHA256: &str = + "d1539d36be7546a1d827fe9cf74c55874695652efb6a5aaa3e60cde1c76ba819"; +const MODEL_WEIGHT_INDEX_SHA256: &str = + "e0fe7f28c1f853d4824e4d796374e3dacf1fe470988773952c79b063768134bf"; +const GOLDEN: &[u8] = include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/glm52-mtp-front-vllm-dcfebf93.safetensors" +)); +const TP1_LAYER_GOLDEN: &[u8] = include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/glm52-mtp-layer78-vllm-tp1-dcfebf93.safetensors" +)); + +fn validate_fixture_metadata(bytes: &[u8], topology: &str) -> Result<()> { + let (_, header) = SafeTensors::read_metadata(bytes)?; + let metadata = header + .metadata() + .as_ref() + .context("MTP fixture has no provenance metadata")?; + for (key, expected) in [ + ("reference", "official-vllm"), + ("vllm_commit", VLLM_COMMIT), + ("topology", topology), + ("model", "GLM-5.2-FP8"), + ("model_config_sha256", MODEL_CONFIG_SHA256), + ("model_weight_index_sha256", MODEL_WEIGHT_INDEX_SHA256), + ] { + ensure!( + metadata.get(key).map(String::as_str) == Some(expected), + "MTP fixture metadata {key:?} is {:?}, expected {expected:?}", + metadata.get(key) + ); + } + Ok(()) +} + +fn bf16_tensor(tensors: &SafeTensors<'_>, name: &str, shape: &[usize]) -> Result> { + let view = tensors.tensor(name)?; + ensure!( + view.dtype() == Dtype::BF16 && view.shape() == shape, + "MTP golden {name} must be BF16 {shape:?}, got {:?} {:?}", + view.dtype(), + view.shape() + ); + Ok(view + .data() + .chunks_exact(2) + .map(|bytes| bf16::from_bits(u16::from_le_bytes([bytes[0], bytes[1]]))) + .collect()) +} + +fn positions(tensors: &SafeTensors<'_>) -> Result> { + i64_tensor(tensors, "positions", &[ROWS])? + .into_iter() + .map(|position| u32::try_from(position).context("MTP golden position is outside u32")) + .collect() +} + +fn i64_tensor(tensors: &SafeTensors<'_>, name: &str, shape: &[usize]) -> Result> { + let view = tensors.tensor(name)?; + ensure!( + view.dtype() == Dtype::I64 && view.shape() == shape, + "MTP golden {name} must be I64 {shape:?}, got {:?} {:?}", + view.dtype(), + view.shape() + ); + Ok(view + .data() + .chunks_exact(8) + .map(|bytes| i64::from_le_bytes(bytes.try_into().expect("eight-byte chunk"))) + .collect()) +} + +fn i32_tensor(tensors: &SafeTensors<'_>, name: &str, shape: &[usize]) -> Result> { + let view = tensors.tensor(name)?; + ensure!( + view.dtype() == Dtype::I32 && view.shape() == shape, + "MTP golden {name} must be I32 {shape:?}, got {:?} {:?}", + view.dtype(), + view.shape() + ); + Ok(view + .data() + .chunks_exact(4) + .map(|bytes| i32::from_le_bytes(bytes.try_into().expect("four-byte chunk"))) + .collect()) +} + +fn f32_tensor(tensors: &SafeTensors<'_>, name: &str, shape: &[usize]) -> Result> { + let view = tensors.tensor(name)?; + ensure!( + view.dtype() == Dtype::F32 && view.shape() == shape, + "MTP golden {name} must be F32 {shape:?}, got {:?} {:?}", + view.dtype(), + view.shape() + ); + Ok(view + .data() + .chunks_exact(4) + .map(|bytes| f32::from_le_bytes(bytes.try_into().expect("four-byte chunk"))) + .collect()) +} + +fn checkpoint_tensor( + manifest: &Glm52WeightManifest, + model_path: &Path, + name: &str, +) -> Result> { + let shard = manifest.shard_for(name)?; + let mmap = mmap_file(&model_path.join(shard))?; + let tensors = SafeTensors::deserialize(mmap.as_ref())?; + Ok(tensors.tensor(name)?.data().to_vec()) +} + +fn load_mtp_head_weights( + ctx: &DeviceContext, + manifest: &Glm52WeightManifest, + model_path: &Path, +) -> Result { + let prefix = "model.layers.78"; + Glm52MtpBookendWeights::from_host( + ctx, + &checkpoint_tensor(manifest, model_path, &format!("{prefix}.enorm.weight"))?, + &checkpoint_tensor(manifest, model_path, &format!("{prefix}.hnorm.weight"))?, + &checkpoint_tensor(manifest, model_path, &format!("{prefix}.eh_proj.weight"))?, + &checkpoint_tensor( + manifest, + model_path, + &format!("{prefix}.shared_head.norm.weight"), + )?, + ) +} + +fn upload_rows(ctx: &DeviceContext, host: &[bf16]) -> Result> { + ensure!( + host.len().is_multiple_of(C), + "MTP golden host tensor length {} is not divisible by {C}", + host.len() + ); + let mut rows = Rows::zeros(ctx, host.len() / C)?; + ctx.stream.memcpy_htod(host, rows.data_mut())?; + Ok(rows) +} + +fn assert_close( + ctx: &DeviceContext, + label: &str, + actual: &Rows, + expected: &[bf16], + rms_limit: f32, + p99_limit: f32, +) -> Result<()> { + let actual = ctx + .stream + .clone_dtoh(actual.data())? + .into_iter() + .map(bf16::to_f32) + .collect::>(); + assert_close_values(label, &actual, expected, rms_limit, p99_limit) +} + +fn assert_close_values( + label: &str, + actual: &[f32], + expected: &[bf16], + rms_limit: f32, + p99_limit: f32, +) -> Result<()> { + let expected = expected + .iter() + .map(|value| value.to_f32()) + .collect::>(); + assert_close_f32_values(label, actual, &expected, rms_limit, p99_limit) +} + +fn assert_close_f32_values( + label: &str, + actual: &[f32], + expected: &[f32], + rms_limit: f32, + p99_limit: f32, +) -> Result<()> { + ensure!( + actual.len() == expected.len(), + "{label}: actual length {} != expected {}", + actual.len(), + expected.len() + ); + let exact = actual + .iter() + .zip(expected) + .filter(|(a, b)| a.to_bits() == b.to_bits()) + .count(); + let mut diffs = actual + .iter() + .zip(expected) + .map(|(a, b)| (a - b).abs()) + .collect::>(); + diffs.sort_by(f32::total_cmp); + let rms = (diffs.iter().map(|diff| diff * diff).sum::() / diffs.len() as f32).sqrt(); + let p99 = diffs[(diffs.len() * 99 / 100).min(diffs.len() - 1)]; + let max = diffs[diffs.len() - 1]; + println!( + "{label}: exact={exact}/{} ({:.2}%) rms={rms:.6e} p99={p99:.6e} max={max:.6e}", + diffs.len(), + exact as f64 * 100.0 / diffs.len() as f64 + ); + ensure!( + rms <= rms_limit && p99 <= p99_limit, + "{label}: rms {rms:.6e} / p99 {p99:.6e} exceed limits \ + {rms_limit:.6e} / {p99_limit:.6e}" + ); + Ok(()) +} + +fn assert_vllm_topk( + actual_logits: &[bf16], + expected_ids: &[i64], + expected_values: &[bf16], +) -> Result<()> { + ensure!( + expected_ids.len() == expected_values.len() && !expected_ids.is_empty(), + "MTP logits golden ids/values shape mismatch" + ); + let mut ranked = (0..actual_logits.len()).collect::>(); + ranked.sort_unstable_by(|&left, &right| { + actual_logits[right] + .to_f32() + .total_cmp(&actual_logits[left].to_f32()) + .then_with(|| left.cmp(&right)) + }); + let actual_topk = &ranked[..expected_ids.len()]; + let expected_ids = expected_ids + .iter() + .map(|&id| usize::try_from(id).context("negative MTP golden token id")) + .collect::>>()?; + let overlap = actual_topk + .iter() + .filter(|id| expected_ids.contains(id)) + .count(); + let top8_overlap = actual_topk[..8] + .iter() + .filter(|id| expected_ids[..8].contains(id)) + .count(); + println!( + "mtp/logits: top1 actual={} expected={}, top8 overlap={top8_overlap}, \ + top{} overlap={overlap}, actual_top8={:?}", + actual_topk[0], + expected_ids[0], + expected_ids.len(), + &actual_topk[..8.min(actual_topk.len())], + ); + let actual_expected_values = expected_ids + .iter() + .map(|&id| actual_logits[id].to_f32()) + .collect::>(); + assert_close_values( + "mtp/logits/expected_ids", + &actual_expected_values, + expected_values, + 2.5e-1, + 5.0e-1, + )?; + ensure!( + actual_topk[0] == expected_ids[0], + "MTP draft top-1 differs from official vLLM" + ); + ensure!( + top8_overlap == 8, + "MTP draft top-8 set overlap with official vLLM is {top8_overlap}/8" + ); + ensure!( + overlap + 2 >= expected_ids.len(), + "MTP draft top-{} overlap with official vLLM is {overlap}/{}", + expected_ids.len(), + expected_ids.len() + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA + GLM-5.2-FP8 checkpoint"] +fn mtp_front_vllm_golden_gate() -> Result<()> { + validate_fixture_metadata(GOLDEN, "tp8-ep0")?; + let fixture = SafeTensors::deserialize(GOLDEN)?; + let model_path = model_path(); + let manifest = Glm52WeightManifest::from_model_dir(&model_path)?; + let ctx = DeviceContext::new()?; + let weights = load_mtp_head_weights(&ctx, &manifest, &model_path)?; + + let positions = ctx.stream.clone_htod(&positions(&fixture)?)?; + let inputs_embeds = upload_rows::( + &ctx, + &bf16_tensor(&fixture, "inputs_embeds_raw", &[ROWS, GLM52_HIDDEN])?, + )?; + let previous_hidden = upload_rows::( + &ctx, + &bf16_tensor(&fixture, "previous_hidden_raw", &[ROWS, GLM52_HIDDEN])?, + )?; + let mut scratch = Glm52MtpScratch::new(&ctx, ROWS)?; + let mut decoder_input = Rows::zeros(&ctx, ROWS)?; + glm52_mtp_prepare_into( + &ctx, + &weights, + &positions, + &inputs_embeds, + &previous_hidden, + &mut scratch, + &mut decoder_input, + )?; + + assert_close( + &ctx, + "mtp/enorm", + scratch.normed_embed(), + &bf16_tensor(&fixture, "inputs_embeds_norm", &[ROWS, GLM52_HIDDEN])?, + 1.0e-3, + 3.90625e-3, + )?; + assert_close( + &ctx, + "mtp/hnorm", + scratch.normed_previous(), + &bf16_tensor(&fixture, "previous_hidden_norm", &[ROWS, GLM52_HIDDEN])?, + 1.0e-3, + 3.90625e-3, + )?; + assert_close( + &ctx, + "mtp/eh_proj", + &decoder_input, + &bf16_tensor(&fixture, "eh_proj", &[ROWS, GLM52_HIDDEN])?, + 1.0e-3, + 3.90625e-3, + )?; + + let raw_hidden = upload_rows::( + &ctx, + &bf16_tensor(&fixture, "raw_hidden", &[ROWS, GLM52_HIDDEN])?, + )?; + let mut recycle_hidden = Rows::zeros(&ctx, ROWS)?; + glm52_mtp_recycle_into(&ctx, &weights, &raw_hidden, &mut recycle_hidden)?; + assert_close( + &ctx, + "mtp/shared_norm", + &recycle_hidden, + &bf16_tensor(&fixture, "recycle_hidden", &[ROWS, GLM52_HIDDEN])?, + 1.0e-3, + 3.90625e-3, + ) +} + +#[test] +#[ignore = "requires 8×H200 + GLM-5.2-FP8 checkpoint + NCCL >= 2.30.4 + DeepGEMM env"] +fn mtp_layer78_vllm_ep8_golden_gate() -> Result<()> { + validate_fixture_metadata(GOLDEN, "tp8-ep0")?; + validate_fixture_metadata(TP1_LAYER_GOLDEN, "tp1-ep0")?; + let fixture = SafeTensors::deserialize(GOLDEN)?; + let tp1_fixture = SafeTensors::deserialize(TP1_LAYER_GOLDEN)?; + let expected_post_attention = bf16_tensor(&fixture, "decoder_residual", &[ROWS, GLM52_HIDDEN])?; + let expected_tp8_mlp = bf16_tensor(&fixture, "decoder_hidden", &[ROWS, GLM52_HIDDEN])?; + let expected_tp1_mlp = bf16_tensor(&tp1_fixture, "decoder_hidden", &[ROWS, GLM52_HIDDEN])?; + let expected_tp1_normed = + bf16_tensor(&tp1_fixture, "post_attention_norm", &[ROWS, GLM52_HIDDEN])?; + let expected_tp1_topk_ids = i32_tensor(&tp1_fixture, "topk_ids", &[ROWS, 8])?; + let expected_tp1_topk_weights = f32_tensor(&tp1_fixture, "topk_weights", &[ROWS, 8])?; + let expected_tp1_routed = bf16_tensor(&tp1_fixture, "routed_hidden", &[ROWS, GLM52_HIDDEN])?; + let expected_tp1_shared_gate_up = bf16_tensor( + &tp1_fixture, + "shared_gate_up", + &[ + ROWS, + 2 * crate::moe_decode::GLM52_SHARED_EXPERT_INTERMEDIATE, + ], + )?; + let expected_tp1_shared_silu = bf16_tensor( + &tp1_fixture, + "shared_silu", + &[ROWS, crate::moe_decode::GLM52_SHARED_EXPERT_INTERMEDIATE], + )?; + let expected_tp1_shared = bf16_tensor(&tp1_fixture, "shared_hidden", &[ROWS, GLM52_HIDDEN])?; + let expected_tp1_recycle = bf16_tensor(&tp1_fixture, "recycle_hidden", &[ROWS, GLM52_HIDDEN])?; + let expected = bf16_tensor(&fixture, "raw_hidden", &[ROWS, GLM52_HIDDEN])?; + let model_path = model_path(); + let unique_id = glm52_ep_deepep_unique_id(EP_RANKS)?; + let tensors = Arc::new(LayerTensors::load(&model_path, GLM52_MTP_LAYER)?); + + let handles: Vec<_> = (1..EP_RANKS) + .map(|rank| { + let tensors = Arc::clone(&tensors); + std::thread::Builder::new() + .name(format!("mtp-ep8-gate-rank-{rank}")) + .spawn(move || -> Result<()> { + let ctx = DeviceContext::new_with_device(rank)?; + let bank = + load_rank_expert_bank(&ctx, &tensors, GLM52_MTP_LAYER, rank, EP_RANKS)?; + let mut ep8 = Glm52MoeEp8State::new(&ctx, &unique_id, EP_RANKS, rank)?; + for _ in 0..2 * ROWS { + let dispatched = + glm52_moe_ep8_routed_forward(&ctx, &mut ep8, &bank, None, EP_RANKS)?; + ensure!(!dispatched, "expert rank produced a combined output"); + } + Ok(()) + }) + .expect("spawn MTP EP8 gate rank thread") + }) + .collect(); + + let ctx = DeviceContext::new_with_device(0)?; + let manifest = Glm52WeightManifest::from_model_dir(&model_path)?; + let mtp_weights = load_mtp_head_weights(&ctx, &manifest, &model_path)?; + let positions = ctx.stream.clone_htod(&positions(&fixture)?)?; + let inputs_embeds = upload_rows::( + &ctx, + &bf16_tensor(&fixture, "inputs_embeds_raw", &[ROWS, GLM52_HIDDEN])?, + )?; + let previous_hidden = upload_rows::( + &ctx, + &bf16_tensor(&fixture, "previous_hidden_raw", &[ROWS, GLM52_HIDDEN])?, + )?; + let mut mtp_scratch = Glm52MtpScratch::new(&ctx, ROWS)?; + let mut prepared_decoder_input = Rows::zeros(&ctx, ROWS)?; + glm52_mtp_prepare_into( + &ctx, + &mtp_weights, + &positions, + &inputs_embeds, + &previous_hidden, + &mut mtp_scratch, + &mut prepared_decoder_input, + )?; + let prepared_decoder_input = ctx.stream.clone_dtoh(prepared_decoder_input.data())?; + let weights = load_decoder_layer( + &ctx, + &model_path, + GLM52_MTP_LAYER, + GateLayerMlp::MoeEp8Rank0, + )?; + let mut ep8 = Glm52MoeEp8State::new(&ctx, &unique_id, EP_RANKS, 0)?; + let actual = run_layer_prefill_ep8( + &ctx, + &weights, + &mut ep8, + &prepared_decoder_input, + ROWS, + EP_RANKS, + ); + let isolated = run_moe_ep8_rows( + &ctx, + &weights, + &mut ep8, + &expected_post_attention, + &expected_tp1_normed, + ROWS, + EP_RANKS, + ); + + drop(ep8); + for (rank, handle) in handles.into_iter().enumerate() { + handle + .join() + .expect("MTP EP8 gate rank thread panicked") + .with_context(|| format!("MTP EP8 gate rank {}", rank + 1))?; + } + let actual = actual?; + let isolated = isolated?; + let matching_topk = isolated + .topk_ids + .iter() + .zip(&expected_tp1_topk_ids) + .filter(|(actual, expected)| actual == expected) + .count(); + println!( + "mtp/layer78/topk_ids_vs_vllm_tp1: exact={matching_topk}/{} actual={:?} expected={:?}", + expected_tp1_topk_ids.len(), + isolated.topk_ids, + expected_tp1_topk_ids + ); + ensure!( + isolated.topk_ids == expected_tp1_topk_ids, + "MTP layer 78 router top-k IDs differ from official vLLM TP1" + ); + assert_close_values( + "mtp/layer78/post_attention_norm_vs_vllm_tp1", + &isolated.normed, + &expected_tp1_normed, + 1.0e-3, + 4.0e-3, + )?; + assert_close_f32_values( + "mtp/layer78/topk_weights_vs_vllm_tp1", + &isolated.topk_weights, + &expected_tp1_topk_weights, + 1.0e-6, + 1.0e-6, + )?; + assert_close_values( + "mtp/layer78/shared_gate_up_vs_vllm_tp1", + &isolated.shared_gate_up, + &expected_tp1_shared_gate_up, + 1.25e-2, + 4.0e-2, + )?; + assert_close_values( + "mtp/layer78/routed_vs_vllm_tp1", + &isolated.routed, + &expected_tp1_routed, + 7.0e-3, + 2.1e-2, + )?; + assert_close_values( + "mtp/layer78/shared_silu_vs_vllm_tp1", + &isolated.shared_silu, + &expected_tp1_shared_silu, + 6.0e-3, + 2.1e-2, + )?; + assert_close_values( + "mtp/layer78/shared_vs_vllm_tp1", + &isolated.shared, + &expected_tp1_shared, + 6.0e-3, + 1.8e-2, + )?; + assert_close_values( + "mtp/layer78/post_attention", + &actual.post_attention, + &expected_post_attention, + 1.2e-2, + 3.125e-2, + )?; + assert_close_values( + "mtp/layer78/mlp", + &actual.mlp, + &expected_tp8_mlp, + 2.5e-2, + 7.8125e-2, + )?; + assert_close_values( + "mtp/layer78/mlp_from_vllm_tp8_residual", + &isolated.mlp, + &expected_tp8_mlp, + 1.2e-2, + 3.125e-2, + )?; + assert_close_values( + "mtp/layer78/mlp_vs_vllm_tp1", + &isolated.mlp, + &expected_tp1_mlp, + 1.0e-2, + 3.125e-2, + )?; + for row in 0..ROWS { + let range = row * GLM52_HIDDEN..(row + 1) * GLM52_HIDDEN; + assert_close_values( + &format!("mtp/layer78/row{row}"), + &actual.hidden[range.clone()], + &expected[range], + 3.5e-2, + 9.375e-2, + )?; + } + let sampled_row = usize::try_from(i64_tensor(&fixture, "logits_sampled_row", &[1])?[0]) + .context("negative MTP logits sampled row")?; + ensure!( + sampled_row < ROWS, + "MTP logits sampled row {sampled_row} is outside {ROWS} rows" + ); + let lm_head = DeviceMatrix::from_safetensors( + &ctx, + &checkpoint_tensor(&manifest, &model_path, "lm_head.weight")?, + GLM52_VOCAB, + GLM52_HIDDEN, + )?; + let raw_hidden = actual + .hidden + .iter() + .copied() + .map(bf16::from_f32) + .collect::>(); + let raw_hidden = upload_rows::(&ctx, &raw_hidden)?; + let mut recycle_hidden = Rows::::zeros(&ctx, ROWS)?; + glm52_mtp_recycle_into(&ctx, &mtp_weights, &raw_hidden, &mut recycle_hidden)?; + let recycle_host = ctx + .stream + .clone_dtoh(recycle_hidden.data())? + .into_iter() + .map(bf16::to_f32) + .collect::>(); + assert_close_values( + "mtp/recycle_vllm_tp1_vs_tp8", + &expected_tp1_recycle + .iter() + .map(|value| value.to_f32()) + .collect::>(), + &bf16_tensor(&fixture, "recycle_hidden", &[ROWS, GLM52_HIDDEN])?, + 2.5e-2, + 7.8125e-2, + )?; + // The shared RMSNorm amplifies the EP8/full-attention raw-hidden delta. + // Bound that state explicitly, then let the stricter top-k checks below + // decide whether the amplified delta changes draft-token decisions. + assert_close_values( + "mtp/recycle_chained_vs_vllm_tp1", + &recycle_host, + &expected_tp1_recycle, + 1.1e-1, + 3.125e-1, + )?; + assert_close_values( + "mtp/recycle_chained_vs_vllm_tp8", + &recycle_host, + &bf16_tensor(&fixture, "recycle_hidden", &[ROWS, GLM52_HIDDEN])?, + 1.1e-1, + 3.125e-1, + )?; + let mut logits = Rows::::zeros(&ctx, ROWS)?; + glm52_lm_head_into(&ctx, &recycle_hidden, &lm_head, &mut logits)?; + let logits = ctx.stream.clone_dtoh(logits.data())?; + let logits = &logits[sampled_row * GLM52_VOCAB..(sampled_row + 1) * GLM52_VOCAB]; + assert_vllm_topk( + logits, + &i64_tensor(&fixture, "logits_topk_ids", &[1, 32])?, + &bf16_tensor(&fixture, "logits_topk_values", &[1, 32])?, + )?; + assert_close_values("mtp/layer78", &actual.hidden, &expected, 2.5e-2, 7.8125e-2) +} diff --git a/openinfer-glm52/src/oracle/mtp_production.rs b/openinfer-glm52/src/oracle/mtp_production.rs new file mode 100644 index 000000000..90fedd8ff --- /dev/null +++ b/openinfer-glm52/src/oracle/mtp_production.rs @@ -0,0 +1,181 @@ +//! Production-path regression for the target-hidden boundary consumed by MTP. + +use std::path::PathBuf; + +use anyhow::Context; +use anyhow::Result; +use openinfer_core::engine::GenerateRequest; +use openinfer_core::engine::TokenEvent; +use openinfer_core::engine::TokenSink; +use openinfer_sample::SamplingParams; + +use crate::Glm52LaunchOptions; +use crate::Glm52MoeTopo; + +const PATHOLOGICAL_PROMPT: &[u32] = &[ + 98770, 98771, 98772, 98773, 98774, 98775, 98776, 98777, 98778, 98779, 98780, 98781, 98782, + 98783, 98784, 98785, 98786, 98787, 98788, 98789, 98790, 98791, 98792, 98793, 98794, 98795, + 98796, 98797, 98798, 98799, 98800, 98801, 98802, 98803, 98804, 98805, 98806, 98807, 98808, + 98809, 98810, 98811, 98812, 98813, 98814, 98815, 98816, 98817, 98818, 98819, 98820, 98821, + 98822, 98823, 98824, 98825, 98826, 98827, 5691, 109691, 98831, 98832, 98833, 98834, 98835, + 98836, 98837, 98838, 98839, 98840, 98841, 98842, 5691, 98844, 98845, 98846, 98847, 5691, 98849, + 98850, 98851, 98852, 98853, 98854, 5691, 98856, 98857, 98858, 98859, 98860, 98861, 98862, + 98863, 98864, 98865, 98866, 98867, 98868, 98869, 98870, 98871, 98872, 98873, 98874, 98875, + 98876, 98877, 98878, 98879, 98880, 98881, 98882, 98883, 98884, 98885, 98886, 98887, 98888, + 98889, 98890, 98891, 98892, 98893, 98894, 98895, 98896, 98897, +]; + +#[test] +#[ignore = "requires 8×H200 + GLM-5.2-FP8 checkpoint + NCCL >= 2.30.4"] +fn native_mtp_uses_final_normalized_target_hidden() -> Result<()> { + let model_path = std::env::var_os("OPENINFER_TEST_MODEL_PATH") + .map(PathBuf::from) + .context("OPENINFER_TEST_MODEL_PATH must point to GLM-5.2-FP8")?; + crate::scheduler::reset_mtp_production_stats(); + let engine = crate::launch( + &model_path, + Glm52LaunchOptions { + tp_size: 1, + dp_size: 8, + drafter: crate::Glm52Drafter::NativeMtp, + max_model_len: Some(4096), + prefill_only: None, + no_prefix_cache: true, + kv_offload: None, + moe_topo: Glm52MoeTopo::Ep8, + weight_staging: true, + dump_graph_png: None, + rank_hosts: Vec::new(), + }, + )?; + let (token_tx, mut token_rx) = TokenSink::standalone(); + engine.submit(GenerateRequest { + request_id: Some(crate::scheduler::MTP_PRODUCTION_GATE_REQUEST_ID.into()), + queued_at_unix_s: None, + trace_parent: None, + data_parallel_rank: Some(0), + prompt_tokens: PATHOLOGICAL_PROMPT.to_vec(), + params: SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens: 256, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + })?; + + let mut completion = Vec::new(); + loop { + let (_, event) = token_rx + .blocking_recv() + .context("GLM5.2 engine closed the production-gate token stream")?; + match event { + TokenEvent::Token { id, .. } => completion.push(id), + TokenEvent::Finished { + completion_tokens, .. + } => { + assert_eq!(completion_tokens, 256); + break; + } + TokenEvent::Error { message, .. } | TokenEvent::Rejected { message, .. } => { + anyhow::bail!("GLM5.2 production gate failed: {message}") + } + TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. } => {} + } + } + assert_eq!(completion.len(), 256); + assert!( + completion.iter().all(|&token| token == 98824), + "selected target trajectory changed: {:?}", + &completion[..completion.len().min(16)] + ); + + // Reuse rank 0's released slot while all eight ranks process different + // prompt/output lengths. This keeps the native-MTP collectives live + // across mixed prefill, proposal, and idle rank states without paying for + // a second model load. + let prompt_lengths = [PATHOLOGICAL_PROMPT.len(), 112, 96, 80, 64, 48, 32, 16]; + let output_lengths = [256, 28, 24, 20, 16, 12, 8, 6]; + let mut receivers = Vec::with_capacity(8); + for rank in 0..8 { + let (token_tx, token_rx) = TokenSink::standalone(); + engine.submit(GenerateRequest { + request_id: Some(if rank == 0 { + crate::scheduler::MTP_SLOT_REUSE_GATE_REQUEST_ID.into() + } else { + format!("native-mtp-multirank-{rank}") + }), + queued_at_unix_s: None, + trace_parent: None, + data_parallel_rank: Some(rank), + prompt_tokens: PATHOLOGICAL_PROMPT[..prompt_lengths[rank]].to_vec(), + params: SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens: output_lengths[rank], + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + })?; + receivers.push(token_rx); + } + for (rank, mut receiver) in receivers.into_iter().enumerate() { + let mut tokens = 0; + loop { + let (_, event) = receiver + .blocking_recv() + .with_context(|| format!("rank {rank} closed its MTP gate token stream"))?; + match event { + TokenEvent::Token { .. } => tokens += 1, + TokenEvent::Finished { + completion_tokens, .. + } => { + assert_eq!(completion_tokens, output_lengths[rank]); + assert_eq!(tokens, output_lengths[rank]); + break; + } + TokenEvent::Error { message, .. } | TokenEvent::Rejected { message, .. } => { + anyhow::bail!( + "GLM5.2 multi-rank production gate failed on rank {rank}: {message}" + ) + } + TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. } => {} + } + } + } + drop(engine); + + let stats = crate::scheduler::mtp_production_stats(); + let first_proposal = stats + .first_proposal + .as_deref() + .context("native MTP did not produce a proposal")?; + assert_eq!(first_proposal.first(), Some(&98825)); + assert!(stats.rounds > 0, "native MTP did not verify any proposal"); + let mean_accepted_length = 1.0 + stats.accepted_drafts as f64 / stats.rounds as f64; + assert!( + mean_accepted_length >= 5.0, + "native MTP mean accepted length regressed to {mean_accepted_length:.3}: {stats:?}" + ); + let reuse_first_proposal = stats + .reuse_first_proposal + .as_deref() + .context("reused native MTP slot did not produce a proposal")?; + assert_eq!(reuse_first_proposal.first(), Some(&98825)); + assert!( + stats.reuse_rounds >= 32, + "reused native MTP slot produced too few rounds for a stable acceptance gate: {stats:?}" + ); + let reuse_mean_accepted_length = + 1.0 + stats.reuse_accepted_drafts as f64 / stats.reuse_rounds as f64; + assert!( + reuse_mean_accepted_length >= 5.0, + "reused native MTP slot mean accepted length regressed to \ + {reuse_mean_accepted_length:.3}: {stats:?}" + ); + Ok(()) +} diff --git a/openinfer-glm52/src/remote.rs b/openinfer-glm52/src/remote.rs index 197718293..c399e709f 100644 --- a/openinfer-glm52/src/remote.rs +++ b/openinfer-glm52/src/remote.rs @@ -57,7 +57,7 @@ use crate::weights::Glm52WeightManifest; /// Bump on ANY wire-visible change. Both ends ship in one repo at one /// commit; the handshake check turns a mixed deploy into a clean reject /// instead of a bincode decode error mid-flight. -const GLM52_WIRE_VERSION: u32 = 1; +const GLM52_WIRE_VERSION: u32 = 2; /// Frames are small (a `Step` is < 100 KiB even at max table width); anything /// bigger than this is a corrupted length prefix, not a real frame. @@ -130,7 +130,7 @@ enum WireRequest { BuildModel { max_model_len: usize, moe_topo: Glm52MoeTopo, - dspark_enabled: bool, + drafter: crate::Glm52Drafter, }, SetupComm { unique_id: Vec, @@ -476,7 +476,7 @@ impl Glm52RemoteRankWorker { &self, max_model_len: usize, moe_topo: Glm52MoeTopo, - dspark_enabled: bool, + drafter: crate::Glm52Drafter, ) -> Result>>> { let (tx, rx) = bounded(1); self.node.shared.submit( @@ -484,7 +484,7 @@ impl Glm52RemoteRankWorker { WireRequest::BuildModel { max_model_len, moe_topo, - dspark_enabled, + drafter, }, PendingResp::BuildModel(tx), )?; @@ -776,7 +776,7 @@ fn serve_connection(stream: TcpStream) -> Result<()> { fn spawn_hosted_workers(hello: &WireHello) -> Result> { let manifest = Glm52WeightManifest::from_model_dir(&hello.model_path)?; - let bundles = manifest.all_rank_load_bundles(hello.moe_topo)?; + let bundles = manifest.all_rank_load_bundles(hello.moe_topo, false)?; ensure!( hello.first_rank + hello.rank_count <= bundles.len(), "GLM5.2 rank-host asked for ranks {}..{} but {:?} has {} ranks", @@ -834,11 +834,11 @@ fn host_demux_loop( WireRequest::BuildModel { max_model_len, moe_topo, - dspark_enabled, + drafter, } => HostPending::BuildModel(worker.build_model_async( max_model_len, moe_topo, - dspark_enabled, + drafter, None, )?), WireRequest::SetupComm { @@ -938,6 +938,35 @@ mod tests { Ok(()) } + #[test] + fn build_model_frame_roundtrip_preserves_drafter_path() -> Result<()> { + let mut buf = Vec::new(); + let cmd = WireCmd { + worker: 5, + req: WireRequest::BuildModel { + max_model_len: 65_536, + moe_topo: Glm52MoeTopo::Ep8, + drafter: crate::Glm52Drafter::Dspark(PathBuf::from("/models/dspark")), + }, + }; + write_frame(&mut buf, &cmd)?; + let decoded: WireCmd = read_frame(&mut buf.as_slice())?; + ensure!(decoded.worker == 5); + match decoded.req { + WireRequest::BuildModel { + max_model_len, + moe_topo, + drafter, + } => { + ensure!(max_model_len == 65_536); + ensure!(moe_topo == Glm52MoeTopo::Ep8); + ensure!(drafter == crate::Glm52Drafter::Dspark(PathBuf::from("/models/dspark"))); + } + other => bail!("decoded wrong variant: {other:?}"), + } + Ok(()) + } + #[test] fn oversized_frame_rejected() { let mut buf = Vec::new(); diff --git a/openinfer-glm52/src/runner.rs b/openinfer-glm52/src/runner.rs index 45afdcb8c..3090f80a1 100644 --- a/openinfer-glm52/src/runner.rs +++ b/openinfer-glm52/src/runner.rs @@ -154,6 +154,60 @@ impl Glm52PrefillBatch { } } +#[derive(Clone, Copy, Debug)] +pub(crate) struct Glm52MtpAppend { + /// Row in the target step's retained final-normalized hidden buffer. + pub(crate) target_row: usize, + pub(crate) slot: usize, + /// Sequence token shifted one place to the left, matching vLLM's MTP + /// first-pass input construction. + pub(crate) input_token: u32, + pub(crate) position: usize, +} + +/// One rank's work in a fleet-wide native-MTP round. The coordinator selects +/// the same variant for every EP rank, including empty ranks, so no worker can +/// skip a collective entered by its peers. +#[derive(Debug)] +pub(crate) enum Glm52MtpRound { + Reset { + resets: Vec, + }, + Context { + source_bucket: usize, + context_bucket: usize, + resets: Vec, + appends: Vec, + }, + Propose { + source_bucket: usize, + context_bucket: usize, + draft_bucket: usize, + resets: Vec, + appends: Vec, + proposal_slots: Vec, + }, +} + +impl Glm52MtpRound { + pub(crate) fn resets(&self) -> &[usize] { + match self { + Self::Reset { resets } + | Self::Context { resets, .. } + | Self::Propose { resets, .. } => resets, + } + } + + pub(crate) fn source_bucket(&self) -> Option { + match self { + Self::Reset { .. } => None, + Self::Context { source_bucket, .. } | Self::Propose { source_bucket, .. } => { + Some(*source_bucket) + } + } + } +} + enum Glm52RankCommand { LoadWeights { model_path: PathBuf, @@ -169,7 +223,7 @@ enum Glm52RankCommand { BuildModel { max_model_len: usize, moe_topo: crate::Glm52MoeTopo, - dspark_enabled: bool, + drafter: crate::Glm52Drafter, prefill_chunk_size: Option, resp: Sender>>, }, @@ -232,6 +286,13 @@ enum Glm52RankCommand { proposals: Vec<(usize, u32, usize)>, resp: Sender>>, }, + /// Collective native-MTP round. Every EP rank receives one command, + /// including ranks with no live proposal, and uses the coordinator-agreed + /// context/draft buckets for the layer-78 MoE collectives. + MtpDraft { + round: Glm52MtpRound, + resp: Sender>>, + }, /// Rank-local, vLLM-compat P/D only: deinterleave the RoPE dims of pages /// just restored from a vLLM-written namespace (see /// glm52_vllm_rope_fixup.cu). Sent after the pegaflow H2D completed and @@ -330,7 +391,7 @@ impl Glm52RankWorker { &self, max_model_len: usize, moe_topo: crate::Glm52MoeTopo, - dspark_enabled: bool, + drafter: crate::Glm52Drafter, prefill_chunk_size: Option, ) -> Result>>> { let (resp_tx, resp_rx) = bounded(1); @@ -338,7 +399,7 @@ impl Glm52RankWorker { .send(Glm52RankCommand::BuildModel { max_model_len, moe_topo, - dspark_enabled, + drafter, prefill_chunk_size, resp: resp_tx, }) @@ -441,6 +502,20 @@ impl Glm52RankWorker { Ok(resp_rx) } + pub(crate) fn mtp_draft_async( + &self, + round: Glm52MtpRound, + ) -> Result>>> { + let (resp_tx, resp_rx) = bounded(1); + self.tx + .send(Glm52RankCommand::MtpDraft { + round, + resp: resp_tx, + }) + .map_err(|_| anyhow::anyhow!("GLM5.2 rank worker channel closed"))?; + Ok(resp_rx) + } + fn dump_decode_graph_async( &self, bucket: usize, @@ -514,18 +589,19 @@ impl Glm52Worker { &self, max_model_len: usize, moe_topo: crate::Glm52MoeTopo, - dspark_enabled: bool, + drafter: crate::Glm52Drafter, prefill_chunk_size: Option, ) -> Result>>> { match self { - Self::Local(worker) => worker.build_model_async( - max_model_len, - moe_topo, - dspark_enabled, - prefill_chunk_size, - ), + Self::Local(worker) => { + worker.build_model_async(max_model_len, moe_topo, drafter, prefill_chunk_size) + } Self::Remote(worker) => { - worker.build_model_async(max_model_len, moe_topo, dspark_enabled) + ensure!( + prefill_chunk_size.is_none(), + "GLM5.2 TP4 prefill-only execution is single-host" + ); + worker.build_model_async(max_model_len, moe_topo, drafter) } } } @@ -602,6 +678,18 @@ impl Glm52Worker { } } + pub(crate) fn mtp_draft_async( + &self, + round: Glm52MtpRound, + ) -> Result>>> { + match self { + Self::Local(worker) => worker.mtp_draft_async(round), + Self::Remote(_) => { + anyhow::bail!("GLM5.2 native MTP currently requires all EP ranks in one process") + } + } + } + pub(crate) fn dump_decode_graph_async( &self, bucket: usize, @@ -762,7 +850,7 @@ impl Glm52RankThreadState { &mut self, max_model_len: usize, moe_topo: crate::Glm52MoeTopo, - dspark_enabled: bool, + drafter: crate::Glm52Drafter, prefill_chunk_size: Option, ) -> Result> { let mut weights = self @@ -778,7 +866,7 @@ impl Glm52RankThreadState { moe_topo .uses_tensor_replicated_moe() .then_some(self.placement.rank), - dspark_enabled, + &drafter, prefill_chunk_size, )?); let arenas = model.kv_arenas(&dev_ctx.stream)?; @@ -889,6 +977,26 @@ impl Glm52RankThreadState { ) } + fn mtp_draft( + &mut self, + round: &Glm52MtpRound, + ) -> Result> { + let dev_ctx = self.ctx.device_context()?; + let runtime = self + .runtime + .as_mut() + .context("GLM5.2 native MTP draft before build_model")?; + runtime.model.mtp_propose( + &dev_ctx, + &runtime.aux_ctx, + runtime + .ep8 + .as_mut() + .context("GLM5.2 native MTP requires the EP8 collective state")?, + round, + ) + } + fn setup_comm( &mut self, unique_id: &[u8; 128], @@ -1079,14 +1187,14 @@ fn rank_worker_loop(rx: &Receiver, mut state: Glm52RankThreadS Glm52RankCommand::BuildModel { max_model_len, moe_topo, - dspark_enabled, + drafter, prefill_chunk_size, resp, } => { let _ = resp.send(state.build_model( max_model_len, moe_topo, - dspark_enabled, + drafter, prefill_chunk_size, )); } @@ -1138,6 +1246,9 @@ fn rank_worker_loop(rx: &Receiver, mut state: Glm52RankThreadS } => { let _ = resp.send(state.draft(bucket, &resets, &appends, &proposals)); } + Glm52RankCommand::MtpDraft { round, resp } => { + let _ = resp.send(state.mtp_draft(&round)); + } Glm52RankCommand::Shutdown => break, } } diff --git a/openinfer-glm52/src/scheduler/admission.rs b/openinfer-glm52/src/scheduler/admission.rs index 61f71e090..4841f3ef7 100644 --- a/openinfer-glm52/src/scheduler/admission.rs +++ b/openinfer-glm52/src/scheduler/admission.rs @@ -171,7 +171,7 @@ pub(super) fn admit_from_queue( workers: &[Glm52Worker], mirrored: bool, prefix_cache_enabled: bool, - dspark_enabled: bool, + drafter_enabled: bool, _prefill_only: bool, pending_resets: &mut [Vec], slots_changed: &mut bool, @@ -321,7 +321,7 @@ pub(super) fn admit_from_queue( req.params.ignore_eos, cached_tokens, ); - if dspark_enabled { + if drafter_enabled { pending_resets[rank].push(slot); } slots[rank][slot] = Some(ActiveRequest { req, state, kv }); diff --git a/openinfer-glm52/src/scheduler/mod.rs b/openinfer-glm52/src/scheduler/mod.rs index 1847a04a8..7e49e6a3b 100644 --- a/openinfer-glm52/src/scheduler/mod.rs +++ b/openinfer-glm52/src/scheduler/mod.rs @@ -37,6 +37,7 @@ mod admission; mod contract_tests; mod graph; mod load; +mod mtp; mod offload; mod plan; mod slot; @@ -54,6 +55,7 @@ use graph::precapture_step_graphs; use load::pending_is_empty; use load::publish_load; use load::running_counts; +use mtp::run_mtp_round; pub(crate) use offload::REMOTE_FETCH_DEADLINE; use offload::VllmPdState; use openinfer_core::engine::GenerateRequest; @@ -72,6 +74,14 @@ use plan::takes_argmax; use slot::GLM52_PADDING_STEP; use slot::Glm52SlotState; use slot::Glm52StepOutcome; +#[cfg(test)] +pub(crate) use slot::MTP_PRODUCTION_GATE_REQUEST_ID; +#[cfg(test)] +pub(crate) use slot::MTP_SLOT_REUSE_GATE_REQUEST_ID; +#[cfg(test)] +pub(crate) use slot::mtp_production_stats; +#[cfg(test)] +pub(crate) use slot::reset_mtp_production_stats; use tokio::sync::mpsc; use tokio::sync::watch; @@ -81,6 +91,7 @@ use crate::model::Glm52StepKv; use crate::model::Glm52StepShape; use crate::model::glm52_pool_blocks; use crate::model::glm52_table_width; +use crate::runner::Glm52MtpAppend; use crate::runner::Glm52PrefillBatch; use crate::runner::Glm52StepFlags; use crate::runner::Glm52Worker; @@ -134,7 +145,7 @@ pub(crate) fn run_dp8_coordinator( mut submit_rx: mpsc::UnboundedReceiver, workers: Vec, eos_token_ids: &[u32], - dspark_enabled: bool, + drafter: crate::Glm52Drafter, prefill_chunk_size: Option, max_model_len: usize, no_prefix_cache: bool, @@ -145,6 +156,8 @@ pub(crate) fn run_dp8_coordinator( graph_dump_request: Option, ) { let prefill_only = prefill_chunk_size.is_some(); + let dspark_enabled = drafter.is_dspark(); + let mtp_enabled = drafter.is_mtp(); // Tensor-replicated topology: ONE logical rank drives mirrored executors. // Every worker receives the identical step (inputs, shape, KV, seed) and // must return bit-identical outputs — the scheduler admits, plans, and @@ -215,15 +228,14 @@ pub(crate) fn run_dp8_coordinator( // Pool pages available to requests per rank (total minus the padding // page) — constant for the engine's lifetime. let usable_blocks: Vec = pools.iter().map(|pool| pool.total_blocks() - 1).collect(); - // The DSpark draft lane asserts every anchor position equals its - // committed + pending context rows — a skipped (cache-hit) prefix never - // produces the aux-hidden captures the draft consumes, so prefix - // matching is off while the drafter is on. Speculative decoding and - // prefix caching are mutually exclusive for now (the qwen3 offload path - // draws the same line). `--no-prefix-cache` is the explicit kill switch. - let prefix_cache_enabled = !dspark_enabled && !no_prefix_cache; - if dspark_enabled && !no_prefix_cache { - log::info!("GLM5.2 prefix cache disabled: the DSpark drafter is on"); + // A cache-hit prefix skips state required by either speculative lane: + // DSpark loses the aux-hidden captures it consumes, while native MTP + // loses target hidden rows and continuity in its separate KV cache. + // Prefix matching therefore stays off while any drafter is active. + // `--no-prefix-cache` remains the explicit kill switch. + let prefix_cache_enabled = !drafter.enabled() && !no_prefix_cache; + if drafter.enabled() && !no_prefix_cache { + log::info!("GLM5.2 prefix cache disabled: speculative decoding is on"); } let mut slots: Vec = (0..logical_ranks) .map(|_| std::array::from_fn(|_| None)) @@ -330,7 +342,7 @@ pub(crate) fn run_dp8_coordinator( &workers, mirrored, prefix_cache_enabled, - dspark_enabled, + drafter.enabled(), prefill_only, &mut pending_resets, &mut slots_changed, @@ -394,7 +406,7 @@ pub(crate) fn run_dp8_coordinator( leased_shapes.as_deref(), slots_changed, pending_is_empty(&pending), - dspark_enabled, + drafter.enabled(), offload.is_some(), &slots, max_model_len, @@ -404,7 +416,7 @@ pub(crate) fn run_dp8_coordinator( slots_changed = false; sample_step += 1; // One lock-step step (see [`submit_and_join_step`]). - let (outputs, span_kinds) = match submit_and_join_step( + let (outputs, span_kinds, step_inputs) = match submit_and_join_step( &workers, &pools, &mut slots, @@ -420,7 +432,7 @@ pub(crate) fn run_dp8_coordinator( } }; - let (rank_appends, mut rank_proposals) = match apply_step_outputs( + let (rank_appends, mtp_appends, mut rank_proposals) = match apply_step_outputs( &mut slots, outputs, &shapes, @@ -428,7 +440,8 @@ pub(crate) fn run_dp8_coordinator( &pools, offload.as_deref(), eos_token_ids, - dspark_enabled, + &drafter, + &step_inputs, &mut pending_resets, &mut slots_changed, ) { @@ -463,8 +476,8 @@ pub(crate) fn run_dp8_coordinator( } } - if dspark_enabled - && let Err(err) = run_draft_round( + let draft_result = if dspark_enabled { + run_draft_round( &workers, &mut slots, &shapes, @@ -473,7 +486,19 @@ pub(crate) fn run_dp8_coordinator( rank_proposals, span_drafts, ) - { + } else if mtp_enabled { + run_mtp_round( + &workers, + &mut slots, + &shapes, + &mut pending_resets, + mtp_appends, + rank_proposals, + ) + } else { + Ok(()) + }; + if let Err(err) = draft_result { fail_step(&mut slots, &err); break 'serve; } @@ -529,6 +554,7 @@ fn submit_and_join_step( ) -> anyhow::Result<( Vec<[u32; GLM52_MAX_BATCH_PER_RANK]>, Vec<[Option; GLM52_MAX_BATCH_PER_RANK]>, + Vec<[(u32, usize); GLM52_MAX_BATCH_PER_RANK]>, )> { // Logical-to-executor mapping: 1:1 under EP8, or the single logical // rank's step mirrored onto every worker under the replicated tp8 @@ -540,6 +566,7 @@ fn submit_and_join_step( .map(|_| [None; GLM52_MAX_BATCH_PER_RANK]) .collect(); let mut responses = Vec::with_capacity(workers.len()); + let mut step_inputs = Vec::with_capacity(slots.len()); let mut submit_err: Option = None; 'submit: for (rank, (rank_slots, shape)) in slots.iter_mut().zip(shapes).enumerate() { let pool = &pools[rank]; @@ -630,6 +657,7 @@ fn submit_and_join_step( pages: pages.into_boxed_slice(), slot_mapping, }; + step_inputs.push(inputs); let executors: &[Glm52Worker] = if mirrored { workers } else { @@ -685,7 +713,7 @@ fn submit_and_join_step( } outputs.truncate(1); } - Ok((outputs, span_kinds)) + Ok((outputs, span_kinds, step_inputs)) } fn submit_join_apply_prefill( @@ -865,11 +893,17 @@ fn apply_step_outputs( pools: &[BlockPool], offload: Option<&[offload::RankOffload]>, eos_token_ids: &[u32], - dspark_enabled: bool, + drafter: &crate::Glm52Drafter, + step_inputs: &[[(u32, usize); GLM52_MAX_BATCH_PER_RANK]], pending_resets: &mut [Vec], slots_changed: &mut bool, -) -> anyhow::Result<(Vec>, Vec>)> { +) -> anyhow::Result<( + Vec>, + Vec>, + Vec>, +)> { let mut rank_appends: Vec> = slots.iter().map(|_| Vec::new()).collect(); + let mut mtp_appends: Vec> = slots.iter().map(|_| Vec::new()).collect(); let mut rank_proposals: Vec> = slots.iter().map(|_| Vec::new()).collect(); for (rank, ((rank_slots, rank_outputs), shape)) in @@ -965,6 +999,10 @@ fn apply_step_outputs( } }; if freed { + #[cfg(test)] + active + .state + .record_mtp_production_gate(active.req.request_id.as_deref()); active.state.log_spec_stats(rank, slot_id); // Offload the freshly-sealed blocks BEFORE release: the // hashes and guards come off the still-assigned request @@ -982,24 +1020,42 @@ fn apply_step_outputs( (blocks return via RAII): {err:#}" ); } - if dspark_enabled { + if drafter.enabled() { pending_resets[rank].push(slot_id); } *slot = None; *slots_changed = true; - } else if dspark_enabled { - // Committed rows' captured hidden feeds the draft - // context; then re-propose from the new anchor. - rank_appends[rank].extend(span_rows.take(context_rows).map(|r| (r, slot_id))); - if active.state.wants_drafts() - && let Some((anchor, anchor_pos)) = active.state.decode_anchor() - { + } else if drafter.enabled() { + if drafter.is_dspark() { + rank_appends[rank] + .extend(span_rows.clone().take(context_rows).map(|r| (r, slot_id))); + } else { + for (offset, target_row) in span_rows.clone().take(context_rows).enumerate() { + let input_token = if offset + 1 < context_rows { + step_inputs[rank][target_row + 1].0 + } else { + active.state.next_input_at(0).token + }; + mtp_appends[rank].push(Glm52MtpAppend { + target_row, + slot: slot_id, + input_token, + position: step_inputs[rank][target_row].1, + }); + } + } + let wants_drafts = if drafter.is_mtp() { + active.state.wants_full_draft(crate::mtp::GLM52_MTP_DRAFTS) + } else { + active.state.wants_drafts() + }; + if wants_drafts && let Some((anchor, anchor_pos)) = active.state.decode_anchor() { rank_proposals[rank].push((slot_id, anchor, anchor_pos)); } } } } - Ok((rank_appends, rank_proposals)) + Ok((rank_appends, mtp_appends, rank_proposals)) } /// Draft round (rank-local, no collectives): resets, context appends from diff --git a/openinfer-glm52/src/scheduler/mtp.rs b/openinfer-glm52/src/scheduler/mtp.rs new file mode 100644 index 000000000..5bf8c4eb7 --- /dev/null +++ b/openinfer-glm52/src/scheduler/mtp.rs @@ -0,0 +1,198 @@ +//! Coordinator for the checkpoint-native MTP draft lane. + +use anyhow::Context as _; + +use super::RankSlots; +use crate::model::GLM52_DECODE_BUCKETS; +use crate::model::Glm52StepShape; +use crate::runner::Glm52MtpAppend; +use crate::runner::Glm52MtpRound; +use crate::runner::Glm52Worker; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RoundKind { + Reset, + Context, + Propose, +} + +fn select_round_kind( + rank_appends: &[Vec], + rank_proposals: &[Vec<(usize, u32, usize)>], +) -> RoundKind { + if rank_proposals.iter().any(|proposals| !proposals.is_empty()) { + RoundKind::Propose + } else if rank_appends.iter().any(|appends| !appends.is_empty()) { + RoundKind::Context + } else { + RoundKind::Reset + } +} + +/// Native MTP is an EP collective, unlike DSpark. Every worker receives every +/// round, including empty/padded ranks, and all use the same packed context +/// and proposal buckets for each of layer 78's five forwards. +pub(super) fn run_mtp_round( + workers: &[Glm52Worker], + slots: &mut [RankSlots], + shapes: &[Glm52StepShape], + pending_resets: &mut [Vec], + rank_appends: Vec>, + rank_proposals: Vec>, +) -> anyhow::Result<()> { + anyhow::ensure!( + workers.len() == slots.len() + && shapes.len() == workers.len() + && pending_resets.len() == workers.len() + && rank_appends.len() == workers.len() + && rank_proposals.len() == workers.len(), + "GLM5.2 native MTP requires one logical rank per local EP worker" + ); + let source_bucket = shapes + .first() + .context("GLM5.2 native MTP round has no target shape")? + .bucket; + anyhow::ensure!( + shapes.iter().all(|shape| shape.bucket == source_bucket), + "GLM5.2 native MTP source buckets diverged across EP ranks" + ); + + let pick_bucket = |rows: usize| { + GLM52_DECODE_BUCKETS + .into_iter() + .find(|&bucket| bucket >= rows.max(1)) + .with_context(|| format!("GLM5.2 native MTP row count {rows} exceeds bucket capacity")) + }; + let context_bucket = pick_bucket(rank_appends.iter().map(Vec::len).max().unwrap_or(0))?; + let draft_bucket = pick_bucket(rank_proposals.iter().map(Vec::len).max().unwrap_or(0))?; + let kind = select_round_kind(&rank_appends, &rank_proposals); + + let mut joins = Vec::with_capacity(workers.len()); + let mut proposal_slots = Vec::with_capacity(workers.len()); + let mut rank_errors: Vec> = (0..workers.len()).map(|_| None).collect(); + for (rank, ((worker, appends), proposals)) in workers + .iter() + .zip(rank_appends) + .zip(rank_proposals) + .enumerate() + { + let slots_for_rank = proposals + .iter() + .map(|&(slot, _, _)| slot) + .collect::>(); + let resets = std::mem::take(&mut pending_resets[rank]); + let round = match kind { + RoundKind::Reset => Glm52MtpRound::Reset { resets }, + RoundKind::Context => Glm52MtpRound::Context { + source_bucket, + context_bucket, + resets, + appends, + }, + RoundKind::Propose => Glm52MtpRound::Propose { + source_bucket, + context_bucket, + draft_bucket, + resets, + appends, + proposal_slots: slots_for_rank.clone(), + }, + }; + let response = match worker.mtp_draft_async(round) { + Ok(response) => Some(response), + Err(err) => { + let err = err.context(format!("GLM5.2 rank {rank} MTP draft submission")); + log::error!("GLM5.2 rank {rank} MTP draft submission failed: {err:#}"); + rank_errors[rank] = Some(err); + None + } + }; + joins.push(response); + proposal_slots.push(slots_for_rank); + } + + // Join every rank before returning an error. The first rank received can + // be blocked inside DeepEP and report only its device timeout; a later + // response may contain the pre-collective invariant failure that caused + // it. Preserve every error in the log and return the first in rank order. + let mut rank_spans = Vec::with_capacity(joins.len()); + for (rank, (rx, expected_slots)) in joins.iter().zip(&proposal_slots).enumerate() { + let Some(rx) = rx else { + rank_spans.push(Vec::new()); + continue; + }; + let result = rx + .recv() + .map_err(|_| anyhow::anyhow!("dropped its response")) + .and_then(|result| result) + .and_then(|spans| { + anyhow::ensure!( + spans.len() == expected_slots.len(), + "returned {} spans for {} proposals", + spans.len(), + expected_slots.len() + ); + Ok(spans) + }); + match result { + Ok(spans) => rank_spans.push(spans), + Err(err) => { + let err = err.context(format!("GLM5.2 rank {rank} MTP draft")); + log::error!("GLM5.2 rank {rank} MTP draft failed: {err:#}"); + rank_errors[rank] = Some(err); + rank_spans.push(Vec::new()); + } + } + } + if let Some(err) = rank_errors.into_iter().flatten().next() { + return Err(err); + } + + for (rank, (spans, proposal_slots)) in rank_spans.into_iter().zip(proposal_slots).enumerate() { + for (slot_id, span) in proposal_slots.into_iter().zip(spans) { + if let Some(active) = slots[rank][slot_id].as_mut() { + #[cfg(test)] + super::slot::record_mtp_proposal(active.req.request_id.as_deref(), &span); + active + .state + .set_drafts(span.to_vec(), crate::mtp::GLM52_MTP_DRAFTS); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn append() -> Glm52MtpAppend { + Glm52MtpAppend { + target_row: 0, + slot: 0, + input_token: 1, + position: 0, + } + } + + #[test] + fn any_rank_proposal_keeps_every_rank_in_the_collective_chain() { + let appends = vec![vec![append()], Vec::new()]; + let proposals = vec![vec![(0, 1, 0)], Vec::new()]; + assert_eq!(select_round_kind(&appends, &proposals), RoundKind::Propose); + } + + #[test] + fn committed_context_without_proposals_runs_only_the_first_pass() { + let appends = vec![Vec::new(), vec![append()]]; + let proposals = vec![Vec::new(), Vec::new()]; + assert_eq!(select_round_kind(&appends, &proposals), RoundKind::Context); + } + + #[test] + fn an_empty_round_only_resets_host_state() { + let appends = vec![Vec::new(), Vec::new()]; + let proposals = vec![Vec::new(), Vec::new()]; + assert_eq!(select_round_kind(&appends, &proposals), RoundKind::Reset); + } +} diff --git a/openinfer-glm52/src/scheduler/plan.rs b/openinfer-glm52/src/scheduler/plan.rs index dc73cac03..325ba0d97 100644 --- a/openinfer-glm52/src/scheduler/plan.rs +++ b/openinfer-glm52/src/scheduler/plan.rs @@ -129,14 +129,14 @@ pub(super) fn launch_ahead_flags( leased_shapes: Option<&[Glm52StepShape]>, slots_changed: bool, pending_empty: bool, - dspark_enabled: bool, + drafter_enabled: bool, offload_enabled: bool, slots: &[RankSlots], max_model_len: usize, ) -> Glm52StepFlags { let consume = !slots_changed && leased_shapes == Some(shapes); let lease = pending_empty - && !dspark_enabled + && !drafter_enabled && !offload_enabled && slots .iter() diff --git a/openinfer-glm52/src/scheduler/slot.rs b/openinfer-glm52/src/scheduler/slot.rs index 66f8e435a..8306efe04 100644 --- a/openinfer-glm52/src/scheduler/slot.rs +++ b/openinfer-glm52/src/scheduler/slot.rs @@ -8,6 +8,63 @@ use openinfer_core::engine::FinishReason; use crate::dspark::GLM52_DSPARK_DRAFTS; use crate::dspark::accept_prefix_match; +#[cfg(test)] +pub(crate) const MTP_PRODUCTION_GATE_REQUEST_ID: &str = "native-mtp-hidden-boundary-gate"; +#[cfg(test)] +pub(crate) const MTP_SLOT_REUSE_GATE_REQUEST_ID: &str = "native-mtp-slot-reuse-gate"; + +#[cfg(test)] +#[derive(Clone, Debug, Default)] +pub(crate) struct MtpProductionGateStats { + pub(crate) first_proposal: Option>, + pub(crate) rounds: u64, + pub(crate) accepted_drafts: u64, + pub(crate) reuse_first_proposal: Option>, + pub(crate) reuse_rounds: u64, + pub(crate) reuse_accepted_drafts: u64, +} + +#[cfg(test)] +static MTP_PRODUCTION_STATS: std::sync::LazyLock> = + std::sync::LazyLock::new(|| std::sync::Mutex::new(MtpProductionGateStats::default())); + +#[cfg(test)] +pub(crate) fn reset_mtp_production_stats() { + *MTP_PRODUCTION_STATS + .lock() + .expect("MTP production stats lock poisoned") = MtpProductionGateStats::default(); +} + +#[cfg(test)] +pub(super) fn record_mtp_proposal(request_id: Option<&str>, drafts: &[u32]) { + if !matches!( + request_id, + Some(MTP_PRODUCTION_GATE_REQUEST_ID | MTP_SLOT_REUSE_GATE_REQUEST_ID) + ) { + return; + } + let mut stats = MTP_PRODUCTION_STATS + .lock() + .expect("MTP production stats lock poisoned"); + match request_id { + Some(MTP_PRODUCTION_GATE_REQUEST_ID) if stats.first_proposal.is_none() => { + stats.first_proposal = Some(drafts.to_vec()); + } + Some(MTP_SLOT_REUSE_GATE_REQUEST_ID) if stats.reuse_first_proposal.is_none() => { + stats.reuse_first_proposal = Some(drafts.to_vec()); + } + _ => {} + } +} + +#[cfg(test)] +pub(crate) fn mtp_production_stats() -> MtpProductionGateStats { + MTP_PRODUCTION_STATS + .lock() + .expect("MTP production stats lock poisoned") + .clone() +} + /// What a rank forwards this step. Idle rows feed the padding input; their /// KV/index-cache writes land in the pool's reserved padding page, which no /// request is ever assigned. @@ -69,7 +126,7 @@ pub(super) struct Glm52SlotState { } /// Drafts fed per verify span under EP8: 3 drafts + anchor = a bucket-4 -/// verify step. A/B-measured on jz-38 (2026-07-04, +/// verify step. A/B-measured on 8×H200 (2026-07-04, /// docs/models/glm52/dspark-mtp.md): the bucket-4 step costs ~32 ms vs /// bucket-8's ~46, and that cheaper round beats span 8's extra accepted tail /// on EVERY tested prompt class. The drafter still proposes 7; the tail is @@ -111,6 +168,30 @@ impl Glm52SlotState { } } + #[cfg(test)] + pub(super) fn record_mtp_production_gate(&self, request_id: Option<&str>) { + if !matches!( + request_id, + Some(MTP_PRODUCTION_GATE_REQUEST_ID | MTP_SLOT_REUSE_GATE_REQUEST_ID) + ) { + return; + } + let mut stats = MTP_PRODUCTION_STATS + .lock() + .expect("MTP production stats lock poisoned"); + match request_id { + Some(MTP_PRODUCTION_GATE_REQUEST_ID) => { + stats.rounds = self.spec.rounds; + stats.accepted_drafts = self.spec.accepted_sum; + } + Some(MTP_SLOT_REUSE_GATE_REQUEST_ID) => { + stats.reuse_rounds = self.spec.rounds; + stats.reuse_accepted_drafts = self.spec.accepted_sum; + } + _ => {} + } + } + pub(super) fn completion_tokens(&self) -> usize { self.completion } @@ -206,6 +287,14 @@ impl Glm52SlotState { !self.mid_prefill() && self.completion + 1 < self.max_tokens } + /// Native MTP always executes its fixed five-token proposal chain. Do not + /// start that chain in a shorter request tail: the verifier would discard + /// most of it, and the unused speculative KV positions could cross the + /// request's launch-time model-length cap. + pub(super) fn wants_full_draft(&self, draft_tokens: usize) -> bool { + !self.mid_prefill() && self.max_tokens - self.completion >= draft_tokens + } + /// Install the draft lane's proposal for the next verify span, truncated /// to the topology's span cap ([`GLM52_DSPARK_EP8_SPAN_DRAFTS`] under EP8, /// all of [`GLM52_DSPARK_DRAFTS`] under TP8's span shape). @@ -248,7 +337,8 @@ impl Glm52SlotState { debug_assert!(drafts_fed <= self.drafts.len()); let committed = accept_prefix_match(&self.drafts[..drafts_fed], outputs); if drafts_fed > 0 { - self.spec.record(committed.len() - 1); + let accepted_drafts = committed.len() - 1; + self.spec.record(accepted_drafts); } let context_rows = committed.len(); (committed, context_rows) @@ -293,7 +383,7 @@ impl Glm52SlotState { } let mean_accepted = stats.accepted_sum as f64 / stats.rounds as f64; log::info!( - "GLM5.2 dspark: rank={rank} slot={slot} rounds={} mean_accepted_drafts={mean_accepted:.3} \ + "GLM5.2 speculative: rank={rank} slot={slot} rounds={} mean_accepted_drafts={mean_accepted:.3} \ mean_accepted_incl_bonus={:.3} hist={:?}", stats.rounds, mean_accepted + 1.0, @@ -597,6 +687,19 @@ mod tests { assert!(!state.wants_drafts(), "one-token tail is a plain row"); } + #[test] + fn fixed_mtp_chain_stops_before_a_short_tail() { + let mut state = state(vec![10], 6, false); + assert!(!state.wants_full_draft(5), "mid-prefill never drafts"); + assert_eq!(state.advance_span(&[20], EOS), commit(&[20], 1, None, 1)); + assert!(state.wants_full_draft(5)); + assert_eq!(state.advance_span(&[21], EOS), commit(&[21], 1, None, 1)); + assert!( + !state.wants_full_draft(5), + "four-token tail must not launch a five-token MTP chain" + ); + } + #[test] fn sampling_rows_are_the_committable_rows_of_the_span() { let mut state = state(vec![10, 11, 12], 8, false); diff --git a/openinfer-glm52/src/weights.rs b/openinfer-glm52/src/weights.rs index 560a15a63..71a0af2bc 100644 --- a/openinfer-glm52/src/weights.rs +++ b/openinfer-glm52/src/weights.rs @@ -19,6 +19,7 @@ use crate::config::GLM52_KV_A_OUT; use crate::config::GLM52_KV_B_OUT; use crate::config::GLM52_KV_LORA_RANK; use crate::config::GLM52_LAYERS; +use crate::config::GLM52_MTP_LAYER; use crate::config::GLM52_O_PROJ_IN; use crate::config::GLM52_Q_B_OUT; use crate::config::GLM52_Q_LORA_RANK; @@ -35,7 +36,6 @@ pub(crate) use load::Glm52RankGpuWeights; pub(crate) use load::load_rank_weights_to_gpu; const GLM52_WEIGHT_INDEX: &str = "model.safetensors.index.json"; -const GLM52_MTP_LAYER: usize = GLM52_LAYERS; /// The EP8 production partition (8 ranks × 32 experts). Serving-path code /// derives rank/expert counts from the launch topology /// (`Glm52MoeTopo::ep_local_experts`); these constants remain the manifest @@ -260,9 +260,10 @@ impl Glm52WeightManifest { pub(crate) fn all_rank_load_bundles( &self, moe_topo: crate::Glm52MoeTopo, + native_mtp: bool, ) -> Result> { (0..moe_topo.device_count()) - .map(|rank| self.rank_load_bundle(rank, moe_topo)) + .map(|rank| self.rank_load_bundle(rank, moe_topo, native_mtp)) .collect() } @@ -279,8 +280,9 @@ impl Glm52WeightManifest { &self, rank: usize, moe_topo: crate::Glm52MoeTopo, + native_mtp: bool, ) -> Result { - let names = self.rank_resident_tensor_names(rank, moe_topo)?; + let names = self.rank_resident_tensor_names(rank, moe_topo, native_mtp)?; let mut by_shard: BTreeMap> = BTreeMap::new(); for name in names { let shard = self @@ -317,9 +319,9 @@ impl Glm52WeightManifest { }) } - /// Full checkpoint coverage, including the native MTP layer. This is a - /// manifest invariant only: resident load plans below deliberately omit - /// tensors that the serving model never consumes. + /// Full checkpoint coverage, including the MTP accuracy-oracle layer. + /// This is a manifest invariant only: resident load plans below + /// deliberately omit tensors that the serving model never consumes. fn rank_tensor_names(&self, rank: usize) -> Result> { ensure!( rank < GLM52_EP_RANKS, @@ -335,14 +337,15 @@ impl Glm52WeightManifest { Ok(names) } - /// Tensors that must become GPU-resident for one serving rank. Native MTP - /// layer 78 is validation-only and never enters this list. TP8 gets all + /// Tensors that must become GPU-resident for one serving rank. MTP layer + /// 78 enters the EP plan only when native MTP is enabled. TP8 gets all /// routed + shared experts from `load_tp8_slice_layer`, so its first pass /// loads routers but not duplicate full shared-expert projections. fn rank_resident_tensor_names( &self, rank: usize, moe_topo: crate::Glm52MoeTopo, + native_mtp: bool, ) -> Result> { ensure!( rank < moe_topo.device_count(), @@ -357,6 +360,15 @@ impl Glm52WeightManifest { for layer_idx in GLM52_DENSE_LAYERS..GLM52_LAYERS { push_routed_experts(&mut names, layer_idx, expert_range.clone()); } + if native_mtp { + self.push_mtp_non_expert_names(&mut names, true); + push_routed_experts(&mut names, GLM52_MTP_LAYER, expert_range); + } + } else { + ensure!( + !native_mtp, + "GLM5.2 native MTP resident loading currently requires an EP topology" + ); } Ok(names) } @@ -382,6 +394,10 @@ impl Glm52WeightManifest { fn push_checkpoint_non_expert_names(&self, names: &mut Vec) { self.push_resident_non_expert_names(names, true); + self.push_mtp_non_expert_names(names, true); + } + + fn push_mtp_non_expert_names(&self, names: &mut Vec, include_shared_experts: bool) { names.push(format!("model.layers.{GLM52_MTP_LAYER}.enorm.weight")); names.push(format!("model.layers.{GLM52_MTP_LAYER}.hnorm.weight")); names.push(format!("model.layers.{GLM52_MTP_LAYER}.eh_proj.weight")); @@ -389,7 +405,7 @@ impl Glm52WeightManifest { "model.layers.{GLM52_MTP_LAYER}.shared_head.norm.weight" )); self.push_attention_names(names, GLM52_MTP_LAYER); - push_moe_non_expert(names, GLM52_MTP_LAYER, true); + push_moe_non_expert(names, GLM52_MTP_LAYER, include_shared_experts); } fn push_attention_names(&self, names: &mut Vec, layer_idx: usize) { @@ -740,10 +756,10 @@ mod tests { weight_map: BTreeMap::new(), }; let ep8 = manifest - .rank_resident_tensor_names(0, crate::Glm52MoeTopo::Ep8) + .rank_resident_tensor_names(0, crate::Glm52MoeTopo::Ep8, false) .unwrap(); let tp8 = manifest - .rank_resident_tensor_names(0, crate::Glm52MoeTopo::Tp8) + .rank_resident_tensor_names(0, crate::Glm52MoeTopo::Tp8, false) .unwrap(); let mtp_prefix = format!("model.layers.{GLM52_MTP_LAYER}."); @@ -761,6 +777,57 @@ mod tests { assert!(tp8.iter().all(|name| !name.contains(".mlp.experts."))); } + #[test] + fn native_mtp_ep8_plan_adds_the_local_layer_78_partition() { + let manifest = Glm52WeightManifest { + weight_map: BTreeMap::new(), + }; + let plain = manifest + .rank_resident_tensor_names(3, crate::Glm52MoeTopo::Ep8, false) + .unwrap(); + let native_mtp = manifest + .rank_resident_tensor_names(3, crate::Glm52MoeTopo::Ep8, true) + .unwrap(); + let mtp_prefix = format!("model.layers.{GLM52_MTP_LAYER}."); + + assert!(plain.iter().all(|name| !name.starts_with(&mtp_prefix))); + assert!( + native_mtp + .iter() + .any(|name| name == &format!("{mtp_prefix}eh_proj.weight")) + ); + assert!( + native_mtp + .iter() + .any(|name| name == &format!("{mtp_prefix}mlp.experts.96.gate_proj.weight")) + ); + assert!( + native_mtp + .iter() + .any(|name| name == &format!("{mtp_prefix}mlp.experts.127.down_proj.weight")) + ); + assert!(native_mtp.iter().all(|name| { + if !name.starts_with(&format!("{mtp_prefix}mlp.experts.")) { + return true; + } + name.split(".mlp.experts.") + .nth(1) + .and_then(|rest| rest.split('.').next()) + .and_then(|expert| expert.parse::().ok()) + .is_some_and(|expert| (96..128).contains(&expert)) + })); + } + + #[test] + fn native_mtp_resident_plan_rejects_non_ep_topologies() { + let manifest = Glm52WeightManifest { + weight_map: BTreeMap::new(), + }; + manifest + .rank_resident_tensor_names(0, crate::Glm52MoeTopo::Tp8, true) + .expect_err("native MTP must not construct a TP resident plan"); + } + #[test] fn ep4_resident_plan_carries_64_expert_bundles() { let manifest = Glm52WeightManifest { @@ -768,7 +835,7 @@ mod tests { }; // Rank 1 of EP4 owns whole experts 64..128 on every MoE layer. let ep4 = manifest - .rank_resident_tensor_names(1, crate::Glm52MoeTopo::Ep4) + .rank_resident_tensor_names(1, crate::Glm52MoeTopo::Ep4, false) .unwrap(); assert!( ep4.iter() @@ -789,7 +856,7 @@ mod tests { // Four EP4 ranks cover the full routed set. assert!( manifest - .rank_resident_tensor_names(4, crate::Glm52MoeTopo::Ep4) + .rank_resident_tensor_names(4, crate::Glm52MoeTopo::Ep4, false) .is_err() ); } diff --git a/openinfer-glm52/tests/fixtures/glm52-mtp-front-vllm-dcfebf93.safetensors b/openinfer-glm52/tests/fixtures/glm52-mtp-front-vllm-dcfebf93.safetensors new file mode 100644 index 000000000..d8f7fa895 Binary files /dev/null and b/openinfer-glm52/tests/fixtures/glm52-mtp-front-vllm-dcfebf93.safetensors differ diff --git a/openinfer-glm52/tests/fixtures/glm52-mtp-layer78-vllm-tp1-dcfebf93.safetensors b/openinfer-glm52/tests/fixtures/glm52-mtp-layer78-vllm-tp1-dcfebf93.safetensors new file mode 100644 index 000000000..5f056c2c6 Binary files /dev/null and b/openinfer-glm52/tests/fixtures/glm52-mtp-layer78-vllm-tp1-dcfebf93.safetensors differ diff --git a/openinfer-kernels/csrc/glm52/glm52_deepgemm_grouped.cu b/openinfer-kernels/csrc/glm52/glm52_deepgemm_grouped.cu index dab6ea5c2..0ec2b45c8 100644 --- a/openinfer-kernels/csrc/glm52/glm52_deepgemm_grouped.cu +++ b/openinfer-kernels/csrc/glm52/glm52_deepgemm_grouped.cu @@ -110,19 +110,24 @@ __global__ void deepgemm_grouped_fp8_metadata_kernel( __global__ void masked_out_to_aligned_kernel( const __nv_bfloat16* __restrict__ masked_out, const int* __restrict__ masked_m, const int64_t* __restrict__ offsets, - __nv_bfloat16* __restrict__ aligned_out, int n) { + const float* __restrict__ row_weights, + __nv_bfloat16* __restrict__ aligned_out, int aligned_rows, int n) { const int g = blockIdx.x; const int r = blockIdx.y; if (r >= masked_m[g]) { return; } - const uint2* src = reinterpret_cast( - masked_out + ((size_t)g * kMaskedCap + r) * n); - uint2* dst = reinterpret_cast( - aligned_out + ((size_t)offsets[g] + r) * n); - const int words = n / 4; // n is a multiple of 4 (6144) - for (int i = threadIdx.x; i < words; i += blockDim.x) { - dst[i] = src[i]; + const __nv_bfloat16* src = + masked_out + ((size_t)g * kMaskedCap + r) * n; + const int64_t aligned_row = offsets[g] + r; + if (aligned_row < 0 || aligned_row >= aligned_rows) { + __trap(); + } + __nv_bfloat16* dst = aligned_out + (size_t)aligned_row * n; + const float weight = __ldg(row_weights + aligned_row); + for (int i = threadIdx.x; i < n; i += blockDim.x) { + dst[i] = + __float2bfloat16_rn(__bfloat162float(src[i]) * weight); } } @@ -164,18 +169,21 @@ CUresult glm52_deepgemm_grouped_fp8_metadata_cuda( CUresult glm52_deepgemm_masked_out_to_aligned_cuda( const __nv_bfloat16* masked_out, const int* masked_m, - const int64_t* expert_offsets, __nv_bfloat16* aligned_out, int n, + const int64_t* expert_offsets, const float* row_weights, + __nv_bfloat16* aligned_out, int aligned_rows, int n, cudaStream_t stream) { if (masked_out == nullptr || masked_m == nullptr || - expert_offsets == nullptr || aligned_out == nullptr) { + expert_offsets == nullptr || row_weights == nullptr || + aligned_out == nullptr) { return CUDA_ERROR_INVALID_VALUE; } - if (n <= 0 || n % 4 != 0) { + if (aligned_rows <= 0 || n <= 0 || n % 4 != 0) { return CUDA_ERROR_INVALID_VALUE; } masked_out_to_aligned_kernel<<>>(masked_out, masked_m, - expert_offsets, aligned_out, n); + expert_offsets, row_weights, + aligned_out, aligned_rows, n); return consume_last_cuda_error(); } diff --git a/openinfer-kernels/csrc/glm52/glm52_moe_quant.cu b/openinfer-kernels/csrc/glm52/glm52_moe_quant.cu index 85c629fd1..71b6a1372 100644 --- a/openinfer-kernels/csrc/glm52/glm52_moe_quant.cu +++ b/openinfer-kernels/csrc/glm52/glm52_moe_quant.cu @@ -10,7 +10,6 @@ constexpr int kGroupSize = 128; constexpr float kFp8Min = -448.0f; constexpr float kFp8Max = 448.0f; constexpr float kPerTokenGroupQuantEps = 1.0e-10f; -constexpr float kMinSiluScale = 1.0f / (kFp8Max * 512.0f); __device__ __forceinline__ unsigned char quantize_e4m3(float value, float scale) { @@ -109,12 +108,13 @@ __global__ void fp8_per_token_group_quant_bf16_k128_kernel( } // Grid-strided over aligned receive rows. The gate|up input rows are already -// in the masked layout written by W13; route weights stay indexed by aligned -// receive row. +// in the masked layout written by W13. Router weights are deliberately not +// applied here: vLLM quantizes the unweighted SwiGLU activation and applies +// routing weights after W2. __global__ void silu_and_mul_per_token_group_quant_bf16_k128_masked_kernel( const __nv_bfloat16* __restrict__ input, - const float* __restrict__ topk_weights, unsigned char* __restrict__ output, - float* __restrict__ scales, int rows, int hidden_size, + unsigned char* __restrict__ output, float* __restrict__ scales, int rows, + int hidden_size, const long long* __restrict__ row_bound, const int* __restrict__ row_map, int masked_cap) { const int group = blockIdx.y; @@ -141,8 +141,10 @@ __global__ void silu_and_mul_per_token_group_quant_bf16_k128_masked_kernel( float gate = __bfloat162float(token_gate[tid]); float up = __bfloat162float(token_up[tid]); float sigmoid_gate = 1.0f / (1.0f + expf(-gate)); - const float route_weight = __ldg(topk_weights + row); - activated = gate * sigmoid_gate * up * route_weight; + // Match vLLM's fused Triton kernel: SiLU narrows to the input dtype + // before the BF16 multiply, then the product is quantized from F32. + __nv_bfloat16 glu = __float2bfloat16_rn(gate * sigmoid_gate); + activated = __bfloat162float(glu) * up; } shared[tid] = fabsf(activated); __syncthreads(); @@ -156,7 +158,7 @@ __global__ void silu_and_mul_per_token_group_quant_bf16_k128_masked_kernel( } if (tid == 0) { - shared[0] = fmaxf(shared[0] / kFp8Max, kMinSiluScale); + shared[0] = fmaxf(shared[0], kPerTokenGroupQuantEps) / kFp8Max; const int g = data_row / masked_cap; const int r_local = data_row % masked_cap; scales[((size_t)g * scale_cols + group) * masked_cap + r_local] = shared[0]; @@ -253,14 +255,12 @@ CUresult glm52_fp8_per_token_group_quant_bf16_masked_cuda( return consume_last_cuda_error(); } -CUresult glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_cuda( - const __nv_bfloat16* input, const float* topk_weights, - unsigned char* output, float* scales, int rows, int hidden_size, - int group_size, const long long* row_bound, const int* row_map, - int masked_cap, cudaStream_t stream) { - if (input == nullptr || topk_weights == nullptr || output == nullptr || - scales == nullptr || row_bound == nullptr || row_map == nullptr || - masked_cap <= 0) { +CUresult glm52_silu_and_mul_per_token_group_quant_bf16_masked_cuda( + const __nv_bfloat16* input, unsigned char* output, float* scales, int rows, + int hidden_size, int group_size, const long long* row_bound, + const int* row_map, int masked_cap, cudaStream_t stream) { + if (input == nullptr || output == nullptr || scales == nullptr || + row_bound == nullptr || row_map == nullptr || masked_cap <= 0) { return CUDA_ERROR_INVALID_VALUE; } if (!valid_quant_shape(rows, hidden_size, group_size)) { @@ -268,8 +268,8 @@ CUresult glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_cuda( } dim3 grid(row_grid(rows), hidden_size / kGroupSize, 1); silu_and_mul_per_token_group_quant_bf16_k128_masked_kernel - <<>>(input, topk_weights, output, scales, - rows, hidden_size, row_bound, row_map, + <<>>(input, output, scales, rows, + hidden_size, row_bound, row_map, masked_cap); return consume_last_cuda_error(); } diff --git a/openinfer-kernels/csrc/shared/elementwise.cu b/openinfer-kernels/csrc/shared/elementwise.cu index a86a8e8ba..cd7e80532 100644 --- a/openinfer-kernels/csrc/shared/elementwise.cu +++ b/openinfer-kernels/csrc/shared/elementwise.cu @@ -20,6 +20,24 @@ __global__ void add_kernel( } } +__global__ void add_scaled_bf16_kernel( + const __nv_bfloat16 *__restrict__ routed, + float scale, + const __nv_bfloat16 *__restrict__ shared, + __nv_bfloat16 *__restrict__ out, + int n) { + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; + idx < n; + idx += gridDim.x * blockDim.x) { + // Match vLLM's two BF16 operations: scale the routed output in place, + // then add the shared expert. + __nv_bfloat16 scaled = __float2bfloat16( + __bfloat162float(routed[idx]) * scale); + out[idx] = __float2bfloat16( + __bfloat162float(shared[idx]) + __bfloat162float(scaled)); + } +} + __global__ void scaled_add_rows_kernel( const __nv_bfloat16 *__restrict__ delta, float scale, @@ -83,6 +101,21 @@ __global__ void copy_hidden_rows_kernel( } } +__global__ void mask_position_zero_rows_kernel( + const __nv_bfloat16 *__restrict__ src, + const uint32_t *__restrict__ positions, + __nv_bfloat16 *__restrict__ dst, + int hidden_dim, + int rows) { + int total = hidden_dim * rows; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; + idx < total; + idx += gridDim.x * blockDim.x) { + int row = idx / hidden_dim; + dst[idx] = positions[row] == 0 ? __float2bfloat16(0.0f) : src[idx]; + } +} + __global__ void copy_hidden_token_range_kernel( const __nv_bfloat16 *__restrict__ src, __nv_bfloat16 *__restrict__ dst, @@ -355,6 +388,21 @@ CUresult add_cuda( return (CUresult)cudaGetLastError(); } +CUresult add_scaled_bf16_cuda( + const __nv_bfloat16 *routed, float scale, + const __nv_bfloat16 *shared, __nv_bfloat16 *out, + int n, cudaStream_t stream) { + if (routed == nullptr || shared == nullptr || out == nullptr || + !isfinite(scale) || n <= 0) { + return CUDA_ERROR_INVALID_VALUE; + } + int block = 256; + int grid = (n + block - 1) / block; + add_scaled_bf16_kernel<<>>( + routed, scale, shared, out, n); + return (CUresult)cudaGetLastError(); +} + CUresult scaled_add_rows_cuda( const __nv_bfloat16 *delta, float scale, @@ -422,6 +470,25 @@ CUresult copy_hidden_rows_cuda( return (CUresult)cudaGetLastError(); } +CUresult mask_position_zero_rows_cuda( + const __nv_bfloat16 *src, + const uint32_t *positions, + __nv_bfloat16 *dst, + int hidden_dim, + int rows, + cudaStream_t stream) { + if (src == nullptr || positions == nullptr || dst == nullptr || + hidden_dim <= 0 || rows <= 0) { + return CUDA_ERROR_INVALID_VALUE; + } + int total = hidden_dim * rows; + int block = 256; + int grid = (total + block - 1) / block; + mask_position_zero_rows_kernel<<>>( + src, positions, dst, hidden_dim, rows); + return (CUresult)cudaGetLastError(); +} + CUresult copy_hidden_token_range_cuda( const __nv_bfloat16 *src, __nv_bfloat16 *dst, diff --git a/openinfer-kernels/src/ffi/glm52.rs b/openinfer-kernels/src/ffi/glm52.rs index 9aae9a831..bc099d99c 100644 --- a/openinfer-kernels/src/ffi/glm52.rs +++ b/openinfer-kernels/src/ffi/glm52.rs @@ -119,7 +119,9 @@ unsafe extern "C" { masked_out: *const Half, masked_m: *const i32, expert_offsets: *const i64, + row_weights: *const f32, aligned_out: *mut Half, + aligned_rows: i32, n: i32, stream: CUstream, ) -> CUresult; @@ -282,9 +284,8 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; - pub fn glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_cuda( + pub fn glm52_silu_and_mul_per_token_group_quant_bf16_masked_cuda( input: *const Half, - topk_weights: *const f32, output: *mut u8, scales: *mut f32, rows: i32, diff --git a/openinfer-kernels/src/ffi/shared.rs b/openinfer-kernels/src/ffi/shared.rs index 7e9caafc9..449e23186 100644 --- a/openinfer-kernels/src/ffi/shared.rs +++ b/openinfer-kernels/src/ffi/shared.rs @@ -32,6 +32,15 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; + pub fn add_scaled_bf16_cuda( + routed: *const Half, + scale: f32, + shared: *const Half, + out: *mut Half, + n: i32, + stream: CUstream, + ) -> CUresult; + pub fn copy_hidden_rows_cuda( src: *const Half, dst: *mut Half, @@ -43,6 +52,15 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; + pub fn mask_position_zero_rows_cuda( + src: *const Half, + positions: *const u32, + dst: *mut Half, + hidden_dim: i32, + rows: i32, + stream: CUstream, + ) -> CUresult; + pub fn copy_hidden_token_range_cuda( src: *const Half, dst: *mut Half, diff --git a/openinfer-kernels/src/ops.rs b/openinfer-kernels/src/ops.rs index 643bec0fd..02c8e82b7 100644 --- a/openinfer-kernels/src/ops.rs +++ b/openinfer-kernels/src/ops.rs @@ -66,6 +66,7 @@ pub use elementwise::accumulate_bf16_token_scaled_to_f32_into; pub use elementwise::add_batch; pub use elementwise::add_batch_into; pub use elementwise::add_into; +pub use elementwise::add_scaled_bf16_into; pub use elementwise::bf16_bytes_to_f32_into; pub use elementwise::bf16_hidden_to_f32_into; pub use elementwise::copy_hidden_rows_into; @@ -77,6 +78,7 @@ pub use elementwise::extract_vec_ref; pub use elementwise::extract_vec_ref_into; pub use elementwise::f32_to_bf16_hidden_into; pub use elementwise::gather_hidden_tokens_into; +pub use elementwise::mask_position_zero_rows_into; pub use elementwise::repeat_f32_for_reduce_scatter_into; pub use elementwise::scale_f32_in_place; pub use elementwise::scaled_add_batch_into; diff --git a/openinfer-kernels/src/ops/elementwise.rs b/openinfer-kernels/src/ops/elementwise.rs index 2778d1202..7de853a6e 100644 --- a/openinfer-kernels/src/ops/elementwise.rs +++ b/openinfer-kernels/src/ops/elementwise.rs @@ -83,6 +83,46 @@ pub fn add_into( Ok(()) } +/// `out = bf16(shared + bf16(scale * routed))`. +/// +/// The intermediate BF16 narrowing matches vLLM's CUDA MoE path, which +/// scales the reduced routed output in place before adding the shared expert. +pub fn add_scaled_bf16_into( + ctx: &DeviceContext, + routed: &CudaSlice, + scale: f32, + shared: &CudaSlice, + n: usize, + out: &mut CudaSlice, +) -> Result<()> { + if !scale.is_finite() { + return Err(anyhow!("add_scaled_bf16_into scale must be finite")); + } + if routed.len() < n || shared.len() < n || out.len() < n { + return Err(anyhow!( + "add_scaled_bf16_into buffers too small for n={n}: routed {}, shared {}, out {}", + routed.len(), + shared.len(), + out.len() + )); + } + let (routed_ptr, _routed_guard) = routed.device_ptr(&ctx.stream); + let (shared_ptr, _shared_guard) = shared.device_ptr(&ctx.stream); + let (out_ptr, _out_guard) = out.device_ptr_mut(&ctx.stream); + let result = unsafe { + ffi::add_scaled_bf16_cuda( + routed_ptr as *const ffi::Half, + scale, + shared_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + n as i32, + crate::tensor::active_cu_stream(ctx), + ) + }; + result.result()?; + Ok(()) +} + /// In-place scaled add into a row range of `out`: out[row_offset..] += scale * delta. pub fn scaled_add_rows_into( ctx: &DeviceContext, @@ -284,6 +324,46 @@ pub fn copy_hidden_rows_raw_into( Ok(()) } +/// Copy `[rows, hidden_dim]`, replacing every row whose position is zero with +/// zeros. This is the MTP embedding-mask boundary: position zero has no +/// previous token and must not contribute its embedding. +pub fn mask_position_zero_rows_into( + ctx: &DeviceContext, + src: &CudaSlice, + positions: &CudaSlice, + hidden_dim: usize, + rows: usize, + dst: &mut CudaSlice, +) -> Result<()> { + assert!( + rows * hidden_dim <= src.len() && rows * hidden_dim <= dst.len(), + "mask_position_zero_rows_into rows {rows} x hidden_dim {hidden_dim} exceed src {} / dst {}", + src.len(), + dst.len() + ); + assert!( + rows <= positions.len(), + "mask_position_zero_rows_into rows {rows} exceed positions {}", + positions.len() + ); + + let (src_ptr, _gs) = src.device_ptr(&ctx.stream); + let (positions_ptr, _gp) = positions.device_ptr(&ctx.stream); + let (dst_ptr, _gd) = dst.device_ptr_mut(&ctx.stream); + let result = unsafe { + ffi::mask_position_zero_rows_cuda( + src_ptr as *const ffi::Half, + positions_ptr as *const u32, + dst_ptr as *mut ffi::Half, + hidden_dim as i32, + rows as i32, + crate::tensor::active_cu_stream(ctx), + ) + }; + result.result()?; + Ok(()) +} + pub fn copy_hidden_token_range_into( ctx: &DeviceContext, src: &HiddenStates, diff --git a/openinfer-kernels/src/ops/glm52/deepgemm_grouped.rs b/openinfer-kernels/src/ops/glm52/deepgemm_grouped.rs index 0b0cc9dc3..c8b644544 100644 --- a/openinfer-kernels/src/ops/glm52/deepgemm_grouped.rs +++ b/openinfer-kernels/src/ops/glm52/deepgemm_grouped.rs @@ -161,41 +161,52 @@ pub fn glm52_deepgemm_masked_grouped_fp8_launch( .map_err(|err| anyhow!("GLM5.2 DeepGEMM masked grouped FP8 {kind:?} launch failed: {err}")) } -/// Masked GEMM output `[32, 64, n]` → the aligned recv slots -/// `decode_combine` addresses (rows `offsets[g] + r` for `r < masked_m[g]`). +/// Masked W2 output `[32, 64, n]` → the aligned recv slots +/// `decode_combine` addresses, applying each aligned row's router weight +/// after the GEMM (rows `offsets[g] + r` for `r < masked_m[g]`). pub fn glm52_deepgemm_masked_out_to_aligned_launch( ctx: &DeviceContext, n: usize, masked_out: &CudaSlice, masked_m: &CudaSlice, expert_offsets: &CudaSlice, + row_weights: &CudaSlice, aligned_out: &mut CudaSlice, ) -> Result<()> { let groups = GLM52_DEEPGEMM_MASKED_GROUPS; let cap = GLM52_DEEPGEMM_MASKED_CAP; ensure!( - n > 0 && n.is_multiple_of(4), - "GLM5.2 masked-out remap needs n % 4 == 0, got {n}" + n > 0 && n.is_multiple_of(4) && aligned_out.len().is_multiple_of(n), + "GLM5.2 masked-out remap needs n % 4 == 0 and whole output rows, got n {n}, output {}", + aligned_out.len() ); + let aligned_rows = i32::try_from(aligned_out.len() / n) + .map_err(|_| anyhow!("GLM5.2 masked-out remap output row count exceeds i32"))?; ensure!( masked_out.len() >= groups * cap * n && masked_m.len() >= groups - && expert_offsets.len() > groups, - "GLM5.2 masked-out remap buffers too small: masked {}, masked_m {}, offsets {}", + && expert_offsets.len() > groups + && aligned_rows > 0 + && row_weights.len() >= aligned_rows as usize, + "GLM5.2 masked-out remap buffers too small: masked {}, masked_m {}, offsets {}, weights {}, output rows {aligned_rows}", masked_out.len(), masked_m.len(), - expert_offsets.len() + expert_offsets.len(), + row_weights.len() ); let (src_ptr, _src_guard) = masked_out.device_ptr(&ctx.stream); let (masked_ptr, _masked_guard) = masked_m.device_ptr(&ctx.stream); let (offsets_ptr, _offsets_guard) = expert_offsets.device_ptr(&ctx.stream); + let (weights_ptr, _weights_guard) = row_weights.device_ptr(&ctx.stream); let (dst_ptr, _dst_guard) = aligned_out.device_ptr_mut(&ctx.stream); let result = unsafe { ffi::glm52_deepgemm_masked_out_to_aligned_cuda( src_ptr as *const ffi::Half, masked_ptr as *const i32, offsets_ptr as *const i64, + weights_ptr as *const f32, dst_ptr as *mut ffi::Half, + aligned_rows, n as i32, ctx.stream.cu_stream(), ) diff --git a/openinfer-kernels/src/ops/glm52/moe_quant.rs b/openinfer-kernels/src/ops/glm52/moe_quant.rs index 73460f5ea..23039fb89 100644 --- a/openinfer-kernels/src/ops/glm52/moe_quant.rs +++ b/openinfer-kernels/src/ops/glm52/moe_quant.rs @@ -188,18 +188,17 @@ pub fn glm52_fp8_per_token_group_quant_bf16_masked_launch( .map_err(|err| anyhow!("GLM5.2 FP8 masked group quant launch failed: {err}")) } -/// Bounded weighted SwiGLU quant for the masked layout: the gate|up input -/// rows are already masked (the W13 masked GEMM wrote them), the route -/// weight stays indexed by the aligned recv row, output/scales land in the -/// masked layouts (see the quant twin above). +/// Bounded SwiGLU quant for the masked layout: the gate|up input rows are +/// already masked (the W13 masked GEMM wrote them), and output/scales land in +/// the masked layouts (see the quant twin above). Router weights are applied +/// after W2, matching vLLM's DeepGEMM path. #[allow(clippy::too_many_arguments)] -pub fn glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_launch( +pub fn glm52_silu_and_mul_per_token_group_quant_bf16_masked_launch( ctx: &DeviceContext, shape: Glm52MoeQuantShape, masked_groups: usize, masked_cap: usize, input: &CudaSlice, - topk_weights: &CudaSlice, output: &mut CudaSlice, scales: &mut CudaSlice, row_bound: &CudaSlice, @@ -210,27 +209,24 @@ pub fn glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_launch( let masked_rows = masked_groups * masked_cap; ensure!( input.len() >= masked_rows * shape.width * 2 - && topk_weights.len() >= shape.rows && output.len() >= masked_rows * shape.width && scales.len() >= masked_rows * shape.scale_cols()? && row_map.len() >= shape.rows, - "GLM5.2 weighted SiLU masked quant buffers too small" + "GLM5.2 SiLU masked quant buffers too small" ); ensure!( row_bound.len() > bound_index, - "GLM5.2 weighted SiLU masked quant row_bound index {bound_index} outside buffer of {}", + "GLM5.2 SiLU masked quant row_bound index {bound_index} outside buffer of {}", row_bound.len() ); let (input_ptr, _g0) = input.device_ptr(&ctx.stream); - let (weight_ptr, _g1) = topk_weights.device_ptr(&ctx.stream); - let (output_ptr, _g2) = output.device_ptr_mut(&ctx.stream); - let (scale_ptr, _g3) = scales.device_ptr_mut(&ctx.stream); - let (bound_ptr, _g4) = row_bound.device_ptr(&ctx.stream); - let (map_ptr, _g5) = row_map.device_ptr(&ctx.stream); + let (output_ptr, _g1) = output.device_ptr_mut(&ctx.stream); + let (scale_ptr, _g2) = scales.device_ptr_mut(&ctx.stream); + let (bound_ptr, _g3) = row_bound.device_ptr(&ctx.stream); + let (map_ptr, _g4) = row_map.device_ptr(&ctx.stream); let result = unsafe { - ffi::glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_cuda( + ffi::glm52_silu_and_mul_per_token_group_quant_bf16_masked_cuda( input_ptr as *const ffi::Half, - weight_ptr as *const f32, output_ptr as *mut u8, scale_ptr as *mut f32, shape.rows as i32, @@ -244,5 +240,5 @@ pub fn glm52_silu_and_mul_weighted_per_token_group_quant_bf16_masked_launch( }; result .result() - .map_err(|err| anyhow!("GLM5.2 weighted SiLU masked quant launch failed: {err}")) + .map_err(|err| anyhow!("GLM5.2 SiLU masked quant launch failed: {err}")) } diff --git a/openinfer-kernels/src/ops/glm52/router.rs b/openinfer-kernels/src/ops/glm52/router.rs index b1bd45850..2bab0b9a2 100644 --- a/openinfer-kernels/src/ops/glm52/router.rs +++ b/openinfer-kernels/src/ops/glm52/router.rs @@ -12,8 +12,9 @@ use crate::tensor::DeviceContext; const GLM52_ROUTER_HIDDEN: usize = 6144; const GLM52_ROUTER_EXPERTS: usize = 256; const GLM52_ROUTER_TOPK: usize = 8; -/// `routed_scaling_factor` from the GLM5.2 checkpoint config, folded into the -/// normalized top-k weights (the shared expert is added unscaled). +/// `routed_scaling_factor` from the GLM5.2 checkpoint config. The default +/// TP path folds it into top-k weights; the EP path uses normalized weights +/// and applies the factor after expert reduction, matching vLLM. const GLM52_ROUTED_RESIDUAL_SCALE: f32 = 2.5; #[derive(Clone, Copy, Debug, PartialEq)] @@ -34,6 +35,15 @@ impl Glm52RouterConfig { } } + pub const fn glm52_unscaled() -> Self { + Self { + hidden_dim: GLM52_ROUTER_HIDDEN, + n_experts: GLM52_ROUTER_EXPERTS, + topk: GLM52_ROUTER_TOPK, + route_scale: 1.0, + } + } + fn validate(self) -> Result<()> { ensure!( self.hidden_dim == GLM52_ROUTER_HIDDEN, diff --git a/openinfer-server/src/bin/glm52_step_bench.rs b/openinfer-server/src/bin/glm52_step_bench.rs index 8e2310628..b7865655a 100644 --- a/openinfer-server/src/bin/glm52_step_bench.rs +++ b/openinfer-server/src/bin/glm52_step_bench.rs @@ -112,7 +112,7 @@ fn main() -> Result<()> { openinfer_glm52::Glm52LaunchOptions { tp_size, dp_size, - dspark_draft_model_path: None, + drafter: openinfer_glm52::Glm52Drafter::None, max_model_len: cli.max_model_len, prefill_only: None, no_prefix_cache: false, diff --git a/openinfer-server/src/config.rs b/openinfer-server/src/config.rs index 86cb6f43f..399987e53 100644 --- a/openinfer-server/src/config.rs +++ b/openinfer-server/src/config.rs @@ -87,7 +87,7 @@ pub(crate) struct Args { /// Sealed KV blocks are saved to host pinned memory and restored into /// HBM before prefill when a prompt's prefix has fallen out of the GPU /// cache. GLM5.2 requires the prefix cache: incompatible with - /// --no-prefix-cache and the DSpark drafter. + /// --no-prefix-cache and speculative decoding. #[arg(long, default_value_t = false)] pub kv_offload: bool, @@ -179,6 +179,10 @@ pub(crate) struct Args { #[arg(long = "dflash-draft-model-path")] pub dflash_draft_model_path: Option, + /// Use GLM5.2's checkpoint-native MTP layer as the speculative drafter. + #[arg(long = "glm52-native-mtp")] + pub glm52_native_mtp: bool, + /// Cap on total prompt tokens forwarded in one scheduler step. Qwen3 and /// Qwen3.5 only (rejected for other model lines); when omitted, they use /// their own crate defaults. @@ -356,6 +360,7 @@ fn consumed_args(model_type: ModelType) -> &'static [&'static str] { "tp_size", "dp_size", "dflash_draft_model_path", + "glm52_native_mtp", "max_model_len", "glm52_prefill_only", "glm52_prefill_chunk_size", @@ -577,6 +582,9 @@ impl Args { if self.dflash_draft_model_path.is_some() { bail!("--glm52-prefill-only is incompatible with the DSpark drafter"); } + if self.glm52_native_mtp { + bail!("--glm52-prefill-only is incompatible with native MTP"); + } if self.kv_offload || self.kv_pd_vllm_seed.is_some() { bail!( "--glm52-prefill-only does not support KV offload or an external P/D peer" @@ -599,6 +607,12 @@ impl Args { self.glm52_prefill_chunk_size ); } + if self.glm52_native_mtp && self.dflash_draft_model_path.is_some() { + bail!("--glm52-native-mtp and --dflash-draft-model-path are mutually exclusive"); + } + if self.glm52_native_mtp && !matches!(moe_topo, openinfer_glm52::Glm52MoeTopo::Ep8) { + bail!("--glm52-native-mtp currently requires --moe-topo=ep8"); + } } Ok(()) } @@ -964,6 +978,41 @@ mod tests { .expect("GLM5.2 should default to DP8/EP8 when --dp-size is omitted"); } + #[cfg(feature = "glm52")] + #[test] + fn glm52_accepts_native_mtp_on_ep8() { + let (args, provided) = parse_with_provided(&["openinfer", "--glm52-native-mtp"]); + assert!(args.glm52_native_mtp); + args.validate(ModelType::Glm52, &provided) + .expect("native MTP should validate on the default EP8 topology"); + } + + #[cfg(feature = "glm52")] + #[test] + fn glm52_native_mtp_rejects_a_second_drafter() { + let (args, provided) = parse_with_provided(&[ + "openinfer", + "--glm52-native-mtp", + "--dflash-draft-model-path", + "/tmp/dspark", + ]); + let error = args + .validate(ModelType::Glm52, &provided) + .expect_err("native MTP and DSpark must be mutually exclusive"); + assert!(error.to_string().contains("mutually exclusive"), "{error}"); + } + + #[cfg(feature = "glm52")] + #[test] + fn glm52_native_mtp_rejects_non_ep8_topology() { + let (args, provided) = + parse_with_provided(&["openinfer", "--glm52-native-mtp", "--moe-topo", "ep4"]); + let error = args + .validate(ModelType::Glm52, &provided) + .expect_err("native MTP currently requires EP8"); + assert!(error.to_string().contains("--moe-topo=ep8"), "{error}"); + } + #[cfg(feature = "glm52")] #[test] fn glm52_rejects_non_dp8_for_ep8() { diff --git a/openinfer-server/src/main.rs b/openinfer-server/src/main.rs index 627c5a35c..74d98fb7c 100644 --- a/openinfer-server/src/main.rs +++ b/openinfer-server/src/main.rs @@ -174,12 +174,19 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result { let moe_topo: openinfer_glm52::Glm52MoeTopo = args.moe_topo.parse().context("--moe-topo")?; + let drafter = if args.glm52_native_mtp { + openinfer_glm52::Glm52Drafter::NativeMtp + } else if let Some(path) = &args.dflash_draft_model_path { + openinfer_glm52::Glm52Drafter::Dspark(path.clone()) + } else { + openinfer_glm52::Glm52Drafter::None + }; openinfer_glm52::launch( &args.model_path, openinfer_glm52::Glm52LaunchOptions { tp_size: args.tp_size, dp_size: args.dp_size.unwrap_or_else(|| moe_topo.default_dp_size()), - dspark_draft_model_path: args.dflash_draft_model_path.clone(), + drafter, max_model_len: args.max_model_len, prefill_only: args.glm52_prefill_only.then_some( openinfer_glm52::Glm52PrefillOnlyOptions {