diff --git a/CLAUDE.md b/CLAUDE.md index 0a872e048..5917661e2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -13,7 +13,7 @@ Every model line is behind a cargo feature; only `qwen3` is a default feature, s | Qwen3-4B / 8B | `pegainfer-qwen3` | `qwen3` (default) | Full attention, TP support | | Qwen3.5-4B / 9B / 27B · Qwen3.8-27B | `pegainfer-qwen35` | `--features qwen35` (needs build-time Python + Triton) | Hybrid Gated DeltaNet + full attention. Qwen3.8 shares the line: same `model_type`, same text geometry — see `docs/models/qwen35/support-qwen38.md` | | DeepSeek-V2-Lite | `pegainfer-deepseek-v2-lite` | `--features deepseek-v2-lite` | MoE + EP, 2-GPU | -| Gemma 4 | `pegainfer-gemma4` | `--features gemma4` | Sliding-window + global full attention, single GPU, batched decode, opt-in chunked prefill | +| Gemma 4 | `pegainfer-gemma4` | `--features gemma4` | Sliding-window + global full attention, dense sizes at TP2 (one rank per device), batched decode, opt-in chunked prefill (single rank only) | | Kimi-K2 | `pegainfer-kimi-k2` | `--features kimi-k2` | MLA + MoE + Marlin INT4, 8-GPU EP | | GLM5.2 | `pegainfer-glm52` | `--features glm52` | MLA + MoE + FP8, 8-GPU EP (bring-up) | | Kimi-K3 | `pegainfer-k3` | `--features k3` | Hybrid KDA + MLA, latent MoE + MXFP4, EP (bring-up — single-rank decode wired) | diff --git a/README.md b/README.md index 129aa4e79..0a5a21cd2 100644 --- a/README.md +++ b/README.md @@ -147,7 +147,7 @@ Only **`qwen3`** is enabled by default, including in the prebuilt binary. Build | --- | --- | --- | --- | | **Qwen3 · dense 0.6B to 32B** | Full attention, GQA | `qwen3` · default | Greedy + sampling, tensor parallel, prefix cache, KV offload; DFlash / DSpark on 4B. [Model page](https://pegainfer.org/models/qwen3-4b/) | | **Qwen3.5 · dense 0.8B to 27B** | Gated DeltaNet + full attention | `qwen35` | Text-only BF16; build-time Triton. Qwen3.8-27B rides the same feature. [Model page](https://pegainfer.org/models/qwen35/) · [Qwen3.8 support record](docs/models/qwen35/support-qwen38.md) | -| **Gemma 4 · 12B and 26B-A4B** | Sliding-window + global attention; NVFP4 routed experts on 26B | `gemma4` | Text-only, single GPU, batched decode and optional chunked prefill. [Model page](https://pegainfer.org/models/gemma4/) | +| **Gemma 4 · 12B and 26B-A4B** | Sliding-window + global attention; NVFP4 routed experts on 26B | `gemma4` | Text-only, batched decode and optional chunked prefill; the dense sizes also run at tensor parallel 2, one rank per device. [Model page](https://pegainfer.org/models/gemma4/) | | **DeepSeek-V2-Lite** | MLA + MoE | `deepseek-v2-lite` | 2-GPU EP2 correctness path. [Status and limits](docs/models/deepseek-v2-lite/status.md) | | **Kimi-K2 / K2.5** | MLA + MoE, Marlin INT4 | `kimi-k2` | 8-GPU expert parallelism. [Model roadmap](docs/models/kimi-k2/roadmap.md) | | **GLM-5.2** | Sparse MLA + MoE, FP8 | `glm52` | Blackwell; EP decode, TP4 prefill, native MTP speculative decoding, P/D disaggregation. Bring-up. [Model page](https://pegainfer.org/models/glm52/) | diff --git a/docs/benchmarks/gemma4-31b-tp2-l20.md b/docs/benchmarks/gemma4-31b-tp2-l20.md new file mode 100644 index 000000000..ecfe35a1d --- /dev/null +++ b/docs/benchmarks/gemma4-31b-tp2-l20.md @@ -0,0 +1,69 @@ +# Gemma 4 31B bf16 tensor-parallel-2 on two L20s + +TL;DR: the stock 31B bf16 checkpoint serves on two 48 GiB L20s (sm_89, PCIe) as a two-rank eager engine — 29.91 GiB of weights per rank, a `PEGAINFER_MAX_CONTEXT=8192` envelope at 8 decode slots, cold load 30 s, first token 57 ms, 49.4 ms per decode step, and 143 tok/s aggregate at 8 concurrent. The checkpoint does not fit one card (57 GiB), so there is no in-box single-rank baseline for it; numerical parity rests on the 31B-geometry synthetic gate (`models/gemma4/tp.md`). + +## Rig + +| | | +| --- | --- | +| checkpoint | `gemma-4-31b-it`, bf16, text tower only (356 vision/audio tensors skipped) | +| hardware | two NVIDIA L20 (sm_89, 44.99 GiB usable each), PCIe, one p8s node | +| build | `--features gemma4`, `PEGAINFER_CUDA_SM=89`, CUDA 12.9 | +| runtime | `--tp-size=2 --cuda-graph=false`, `NCCL_PROTO=LL128`, pod NCCL 2.18.3 | +| envelope | `PEGAINFER_MAX_CONTEXT=8192`, `PEGAINFER_DECODE_SLOTS=8` | +| checkpoint source | node-local disk (a PVC read is ~3.5 min per rank) | + +## Envelope + +| | per rank | +| --- | --- | +| weights (device) | 29.91 GiB | +| global KV pool | 2.50 GiB (1025 pages x 2.50 MiB) | +| resident after weights | 14.18 GiB free | +| load wall (cold) | 30 s (31.4 s in a second cold run) | + +The 8192 x 16-slot default point does not quite fit: 29.91 GiB of weights plus 14.38 GiB of pools (9.38 local + 5.00 global, all binary GiB) is 44.29 GiB, and the ~0.90 GiB implied by the "14.18 GiB free" row above — CUDA context, step arena, cuBLAS workspace, NCCL buffers — puts it at ≈45.2 GiB against the 44.99 GiB usable. It misses by ~0.2 GiB, not by a wide margin. 8 slots is the point that fits. `PEGAINFER_MAX_CONTEXT`/`PEGAINFER_DECODE_SLOTS` trade context against concurrency — see `models/gemma4/tp.md` for the other measured envelopes and the page arithmetic. + +## Method + +Raw `/v1/completions`, greedy, one connection per request. Time to first token is measured directly as the wall of a `max_tokens=1` request; per-step time is `(t(65) - t(1)) / 64`, which cancels the first-token cost. One warmup request is issued first to absorb the one-time buffer/pool warm; concurrency is a thread pool of `N` over `N` requests of 64 tokens each, wall-timed as a batch. + +## Results + +Prompt `The capital of France is`, 5 TTFT samples and 3 TPOT samples: + +| metric | value | +| --- | --- | +| first token (TTFT) | 57 ms (58, 57, 57, 57, 57) | +| decode step (TPOT) | 49.4 ms (49.4, 49.41, 49.40) | + +Throughput, 64 tokens per request: + +| concurrency | wall | aggregate | per request | +| ---: | ---: | ---: | ---: | +| 1 | 3.17 s | 20.2 tok/s | 3.17 s | +| 4 | 3.47 s | 73.8 tok/s | 0.87 s | +| 8 | 3.57 s | 143.5 tok/s | 0.45 s | + +A decode step costs ~49 ms for one row and ~56 ms for eight (3.57 s / 64 steps), so batch width is nearly free until the collective and the wide-batch GEMMs bite. + +## Eager vs captured decode + +The default is captured (`--cuda-graph=true`, as at TP1); both modes measured back to back on the pair: + +| mode | TTFT | step | c8 aggregate | +| --- | ---: | ---: | ---: | +| eager (`--cuda-graph=false`) | 56.5 ms | 49.4 ms | 143.7 tok/s | +| captured (default) | 57.0 ms | 51.5 ms | 146.0 tok/s | + +A wash: the collective, not launch count, dominates the step, so capture buys little and costs ~4% per step at c1. Both start, serve and shut down cleanly; the numbers in the tables above are the eager run. + +## Correctness + +The distribution gate at 31B geometry measured exact (`models/gemma4/tp.md`): the synthetic checkpoint carrying this checkpoint's shapes came out bit-identical between one rank and two (48/48 picks, worst picked-token gap `0.0000`) — a measurement, since the gate asserts against `LOGBROB_LINE` rather than against exactness. On the real checkpoint the sane answers come out through the chat template — `What is the capital of France?` → `Paris`; `Name three primary colors.` → `The three primary colors are red, yellow, and blue.` (the sampled defaults; greedy repeats, which is the model's published non-default behavior, not a TP artifact). + +## Notes + +- Raw completions without `` degenerate, as the tokenizer contract states; the chat endpoint (which carries the template's BOS) is the coherent path. +- The first GEMM of a cold two-rank process would hang on this environment's NCCL before the module-loading warm was added — see "Cold start" in `models/gemma4/tp.md`. The load figure above is with that warm in place. +- Every rank holds the whole 262144-entry vocabulary projection (~2.8 GiB at 31B), and the extra ranks compute a full batch of logits per decode step that is then discarded. The step numbers above include that waste; see "Known bounds" in `models/gemma4/tp.md`. diff --git a/docs/index.md b/docs/index.md index 395be45ef..e88255d54 100644 --- a/docs/index.md +++ b/docs/index.md @@ -46,16 +46,16 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | | `models/qwen35/roadmap.md` | Qwen3.5 dense roadmap v2 (#654): core correctness/admission/chunked-prefill/sampling/step-tail gates are landed; current 4B HTTP boundary is the retained #469 RTX 5090 sweep, which completed with zero failed requests but trails vLLM at high concurrency. Next: HTTP gap attribution, mixed-load ITL (#470), lifecycle recovery (#471), joint-state prefix reuse (#257), and design-first TP (#446). | -| `models/qwen35/load-snapshot.md` | Qwen3.5 publishes logical running, current pending, and KV load after drain/cancellation pruning and before admission through the shared single-GPU/TP scheduler backend. | +| `models/qwen35/load-snapshot.md` | Qwen3.5 publishes post-step running, pending and KV metrics through the shared driver for single-GPU and TP; running includes in-flight prefill. | | `models/qwen35/prefix-cache.md` | Qwen3.5 prefix-cache design: a hit is valid only when full-attention KV and a complete recurrent/conv snapshot exist at the same 256-token boundary; the first version uses a fixed-budget GPU snapshot pool with joint lookup, pinning, and LRU eviction. | | `models/qwen35/kv-admission.md` | Issue #254 complete: Qwen3.5 now uses full-lifetime KV admission, deferred pressure handling, impossible-request rejection, explicit error semantics, direct rejection-event coverage, RTX 5090 e2e, and real HTTP pressure/post-pressure validation. | | `models/qwen35/optimization.md` | Hybrid 24 linear + 8 full attn optimization ledger. Decode-tuning refresh fuses MLP gate/up and tunes decode cublasLt buckets, improving direct TPOT by 2-3%; vLLM still leads 1024/256 HTTP decode. | | `models/qwen35/accuracy.md` | Qwen3.5 HF bf16 logits goldens, matched to the checkpoint by `config_sha256` (0.8b/2b/4b/9b/27b + Qwen3.8-27B all committed), through `past_key_values`: short replay covers sequential graph, bucket-straddling batched graph, and slot-compaction; long replay covers 4097/8192-token prompts; full GSM8K 8-shot now matches the HF baseline within 0.15 percentage points. | | `models/qwen35/support-qwen38.md` | Qwen3.8-27B is the same architecture as Qwen3.5-27B (same `model_type`, 1,199 identically-named tensors, 0 shape deltas), so it needs no new model line: the blockers were bf16-stored GDN scalars (the f32 loader widens 1D bf16 exactly) and a fixture key that could not tell two `(5120,64)` checkpoints apart (now matched by `config_sha256`, failing closed on an unresolvable revision). Its `output_gate_type: "swish"` is inert metadata — the trained gate is sigmoid, measured by a two-hypothesis probe (NLL 1.82 vs 14.27). Serving deltas live in the chat template: `reasoning_effort` accepts only `xhigh|medium|low`, so the frontend maps the OpenAI-only extremes with the active set probed from the loaded template at startup (a template that accepts the caller's value is never rewritten); upstream `Auto` tool-parser selection resolves a `Qwen3.8-27B` path to `qwen3_xml`, so the line resolves `auto` to `qwen3_coder` itself — verified against the real checkpoint (tool calls round-trip streaming + non-streaming). Verified TP2 on 2×L20 sm_89: logits gate mean 0.0238 / p99 0.0862 (short), 0.0226 / 0.0860 (4097+8192). Traps: dump goldens through flash-linear-attention (the eager GDN fallback is all-NaN), and set `PEGAINFER_TEST_MODEL_REVISION` on hand-downloaded checkpoints. Also fixes TP executors ignoring the caller's max_batch when sizing recurrent state. | -| `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. | +| `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns the Qwen3.5/Qwen3.8 runtime and scheduler behind the `qwen35` feature; the server launches a step-contract `Engine`. Covers ownership, startup and the remaining stop-policy boundary. | | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. | | `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 eager dense TP on Qwen3's controller/worker runtime, P2a mixed-step protocol, P2b rank-local GDR sharding, P2c CUDA Graph under TP gated on the compiled decode GQA group (27B group-6 stays eager). | -| `models/qwen35/tp-implementation.md` | Qwen3.5 TP landed through P2c (2026-08, 2× RTX 4090): Phase 1/P2A lifecycle and ID contracts kept; GDR state sharded per rank (27B TP2 fits 48 GB pairs); batched eager TP decode; TP decode CUDA Graphs for 4B/9B (27B group-6 stays eager). 9B/27B TP2 HF + e2e gates pass; the throughput A/B on the merged stack is tracked in #1001. | +| `models/qwen35/tp-implementation.md` | Current TP runtime: rank-local state, start-gated mixed execution, ID-aligned results, acknowledged retirement and joint prefix reuse; batched eager decode and Graph replay for compiled GQA groups. Documents lifecycle constraints and validation boundaries. | | `models/qwen35/mixed-load-itl-470.md` | Issue #470: full cold `--max-batch 8/bg=4` matrix on RTX 4090 (24/24 valid) + starvation negative control. Qwen3.5 is not immune; chunking bounds max/per-step stall but raises p99 at low QPS (~14→~80–92ms) and pulls p99/max back from the prefill wall to the chunk wall at high load; `qps·prefill_s≳1` is a throughput wall (chunking can't fix it, and ON's +15% TTFT can trip it earlier). The old "p99 immunity" was a slot-starvation artifact. | | `models/qwen35/adaptive-scheduler-policy.md` | Issue #727 adaptive scheduler policy record: default `off`, opt-in `auto`, hard `--max-prefill-tokens` cap, TP `auto` rejection, and pre-review whole-prefill benchmark tradeoff retained as non-default evidence. | | `models/qwen35/decode-kernel-attribution.md` | nsys kernel attribution of the remaining serving gap vs vLLM 0.27 on A100: c16 kernel-busy 16.5 vs ~13.2 ms/step (GEMM family +3.5), bs1 92% busy with the deficit in FlashInfer paged decode (6x at bs1), GDN decode 2.2x per layer-step, and the once-per-step output-projection GEMM on an align-1 sm_75 kernel (1.67 ms/step, fixed by tile-aligning the selection width). | @@ -66,9 +66,10 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | | `models/gemma4/tokenizer.md` | Gemma 4 tokenizer/chat-template contracts, gated against a Hugging Face reference (token ids plus all five chat renders, content flattened to strings): BOS comes only from the standalone `chat_template.jinja` (which opens a thought channel and accepts a native system role), EOS is declared three times with three values, the published defaults are sampled rather than greedy, image/audio tokens encode straight from user text so text-only serving must reject them at admission, and one divergence stays open — the server's default content format adds a trailing space to system turns. | -| `models/gemma4/serving.md` | What the engine promises under load: iteration level scheduling with a ceiling on prompt plus output (8192 by default), prompts prefilled whole at a step boundary by default, requests beyond the configured decode slots (16 by default) queued rather than refused, and the two KV families budgeted separately (7.27 GiB sliding + 2.00 GiB global at 12B with the chunk knob off; the sliding budget shrinks to window plus segment under it). The default configuration needs a 48 GiB card at 12B (32.2 GiB resident before the first request) and one GH200 at 31B (85.1 GiB idle on the folded pool; recipe, envelope at four operating points and hardware matrix in their own section). Clients must send `` themselves or the model degenerates. A row is bit-identical when its companions' content and lengths change under a fixed batch width trajectory, and moves when arrivals or retirements change that trajectory, so greedy output is reproducible for a workload rather than across workloads. An opt-in conversation prefix cache (`PEGAINFER_PREFIX_CACHE=K`) resumes multi-turn prompts from captured prompt state at a pre-allocated page cost. An opt-in chunked walk (`PEGAINFER_MIX_CHUNK_TOKENS=N`) walks admissions through shared segment steps with round-by-round page reservation, and `PEGAINFER_MAX_CONTEXT` raises the ceiling to the checkpoint's 262144 while `PEGAINFER_DECODE_SLOTS` trades concurrency for the memory that buys. `PEGAINFER_ADMIT_COALESCE_MS=N` batches arrivals that would interrupt a live decode batch into one admission burst, defaults off, and refuses the async lane. Open: no cross-request prefix sharing, single GPU. | +| `models/gemma4/serving.md` | What the engine promises under load: iteration level scheduling with a ceiling on prompt plus output (8192 by default), prompts prefilled whole at a step boundary by default, requests beyond the configured decode slots (16 by default) queued rather than refused, and the two KV families budgeted separately (7.27 GiB sliding + 2.00 GiB global at 12B with the chunk knob off; the sliding budget shrinks to window plus segment under it). The default configuration needs a 48 GiB card at 12B (32.2 GiB resident before the first request) and one GH200 at 31B (85.1 GiB idle on the folded pool; recipe, envelope at four operating points and hardware matrix in their own section). Clients must send `` themselves or the model degenerates. A row is bit-identical when its companions' content and lengths change under a fixed batch width trajectory, and moves when arrivals or retirements change that trajectory, so greedy output is reproducible for a workload rather than across workloads. An opt-in conversation prefix cache (`PEGAINFER_PREFIX_CACHE=K`) resumes multi-turn prompts from captured prompt state at a pre-allocated page cost. An opt-in chunked walk (`PEGAINFER_MIX_CHUNK_TOKENS=N`) walks admissions through shared segment steps with round-by-round page reservation, and `PEGAINFER_MAX_CONTEXT` raises the ceiling to the checkpoint's 262144 while `PEGAINFER_DECODE_SLOTS` trades concurrency for the memory that buys. `PEGAINFER_ADMIT_COALESCE_MS=N` batches arrivals that would interrupt a live decode batch into one admission burst, defaults off, and refuses the async lane. Open: no cross-request prefix sharing; tensor parallelism serves the dense sizes one rank per device, with its own prompt cap (`tp.md`). | | `models/gemma4/hf-golden.md` | Three Hugging Face references, committed for 12B and dumpable for another size under `PEGAINFER_GEMMA4_FIXTURE_TAG` (the deep ones with the tower sharded by `--device auto`). The base fixture carries layer-boundary activations at both ends of both layer types plus top-64 logprobs, over a single-token, a nine-token and a 1024-token (exactly the sliding window) case, and pins three facts the forward path has to match — the embedding scale is bf16 62.0 rather than `sqrt(3840)`, text attention is causal, and `layer_scalar` applies to the layer output after both residual adds. The window fixture goes past the window (1023/1024/1025/4096, teacher-forced) under both sdpa and eager. The long-context fixture takes the same comparison to 16384/32768 for the raised ceiling, sdpa-only, with the window fixture's dual-backend floor on loan. Regeneration is byte-identical and checked with sha256. | | `models/gemma4/kv-cache.md` | The two-family KV contract (256 sliding / 512 full attention, two lifetimes): two coordinate systems with the never-conflate invariant frozen — slots derive in-kernel from the resident origin; storage declared along orthogonal axes (lifecycle, paged/slab live storage, reference/copy checkpoint materialization) with observable semantics frozen, not a closed pairing; staged serving — unreclaimed pool first, paged window with batching; a frozen checkpoint query — the largest position at which every required family is exactly recoverable, and a global prefix without its local window is worthless; capacity and TP accounting; four frozen behaviors with a default-off prefix-cache amendment. | +| `models/gemma4/tp.md` | Dense Gemma 4 under tensor parallelism: one rank per device, per-layer-kind row/column shards (replicated embedding, norms, head, rope), the world-size rule including the per-rank GQA integrality the three-count rule misses, one `all_reduce` after `o_proj` and one after `down` over the live extent, and a per-rank `GemmaKv` bound to its own pools — which is why one scheduler thread suffices lock-step and no host/device split was needed. Pool budget is `layers × pages × bytes-per-page`, so 31B TP2 needs 29.91 GiB of weights plus 14.38 GiB of pools at the default point — 44.29 GiB before the ~0.9 GiB of context and workspace, which is just over an L20's 44.99 GiB usable — and fits at 8192 × 8 (38.5 GiB) or 2048 × 2 (31.3 GiB). The 12B gate compares requested top-8 logprobs against a single-rank run, holds the distributions rather than the greedy tokens (a different reduction order flips near-ties), and carries a same-run-twice control. 31B bf16 then serves at 8192 × 8 on the 48 GiB-class sm_89 pair (L20/L40S/RTX 6000 Ada; measured on two L20s: 57 ms first token, 49.4 ms per step, 143 tok/s at c8) — the need is capacity, not the card; and a cold two-rank process must warm the cuBLAS kernels before the communicator exists — otherwise the first GEMM deadlocks in the driver's module loader. Decode graphs work under TP: the sweep interleaves phases across ranks and each rank releases its graphs before the communicators drop. | ## models / glm52 @@ -111,7 +112,8 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | -| `models/deepseek-v2-lite/status.md` | DeepSeek-V2-Lite EP2 status ledger: correctness, direct attribution, lifecycle reliability, and the #466 six-child host-staged/NCCL HTTP SLO report remain separate evidence buckets; no soak, parity, or production claim. | +| `models/deepseek-v2-lite/status.md` | DeepSeek-V2-Lite EP2 capability and validation ledger, with correctness, performance, and production-readiness gates tracked separately. | +| `models/deepseek-v2-lite/device-routed-moe.md` | NCCL MoE forwards with rows 1..8 keep routing and routed-expert execution on device through fixed route slots and pointer-batched N=1 GEMMs; decode performance improves while larger forwards retain the eager fallback. | | `models/deepseek-v2-lite/benchmarking.md` | Verification ladder and commands for DSV2-Lite correctness, direct decode diagnostics, retained host-staged/NCCL HTTP SLO profiles, and the separate soak/production boundary. | | `models/deepseek-v2-lite/roadmap.md` | Single-node EP2 roadmap: #466 retained SLO reporting is an evidence layer; soak, device KV/attention, long-prefill scheduling, and Stable promotion remain separate gates. | | `models/deepseek-v2-lite/benchmark-artifact-manifest.md` | Issue #467 implemented: the retained DeepSeek-V2-Lite benchmark matrix emits `artifact_manifest.json` and `regression_summary.json`, with CPU-only summarize-only tests. | @@ -173,13 +175,13 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | -| `subsystems/frontend/frontend-architecture.md` | `pegainfer-frontend` owns everything north of the model schedulers. Two contract generations coexist: the step contract (qwen3 + pegainfer-sim migrated) and the legacy `EngineHandle`/`TokenEvent` path (other five lines). Next: migrate glm52, then delete the legacy contract. | +| `subsystems/frontend/frontend-architecture.md` | `pegainfer-frontend` owns everything north of the model schedulers. Two contract generations coexist: the step contract (qwen3 + qwen35 + gemma4 + pegainfer-sim migrated) and the legacy `EngineHandle`/`TokenEvent` path for remaining model lines. Next: migrate glm52, then delete the legacy contract. | | `subsystems/frontend/simulated-inference-engine.md` | CPU-only simulated model crate on the step contract (`SimScheduler` → `LaunchedEngine::Stepped`) for vLLM/OpenAI frontend and `vllm bench serve` validation without CUDA or weights. | | `subsystems/frontend/sim-step-contract.md` | Cut `pegainfer-sim` from the legacy `EngineHandle`/`TokenEvent` path onto the step contract. | | `subsystems/frontend/sim-high-concurrency-bench.md` | Same-session A/B vs main: feat TPOT ~30–180× better, TTFT worse and linear in C; E2EL/throughput win at c=64 and c=1024. | | `subsystems/frontend/cpu-profiling-baseline.md` | Frontend CPU profiling baseline using `pegainfer-sim` with fixed TTFT=5ms/TPOT=12ms: 200 req / concurrency=16 shows ~150ms TTFT overhead (no dominant hotspot), heap allocation ~10%, stream polling ~7.5%, IPC ~1%; reproducible benchmark command and perf evidence documented. | | `subsystems/frontend/startup-time.md` | Qwen3-4B warm startup-to-ready: frontend tokenizer load runs concurrently with the engine load (HTTP still binds only after the engine registers); mmap teardown is paid at the end of load since #377; pinned-staging upload (2026-07) cuts warm ready 5.22s → 4.66s on sm_89, and the remaining floor is the engine's own post-load startup work. | -| `subsystems/frontend/prometheus-metrics.md` | `/metrics` request histograms work for every model; Qwen3, Qwen3.5, and GLM5.2 schedulers also publish running/waiting/KV engine gauges through `LoadSnapshot` watches. | +| `subsystems/frontend/prometheus-metrics.md` | `/metrics` request histograms work for every model; Qwen3, Qwen3.5 and Gemma 4 publish running/waiting/KV gauges through step-contract metrics cells; instrumented legacy schedulers use `LoadSnapshot` watches. | | `subsystems/frontend/dashboards/README.md` | Grafana 10.4-validated dashboard for PegaInfer's live `/metrics` surface: HTTP traffic, request outcomes, scheduler/KV state, token throughput, and request latency. | ## subsystems / correctness @@ -235,6 +237,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `benchmarks/qwen35-4b-serving-vllm-rtx5090.md` | Qwen3.5-4B TP1 vs vLLM 0.23.0 on RTX 5090: latest direct PegaInfer A/B improves TPOT by 2-3%; HTTP `vllm bench serve` shows prompt-len-1 decode close, but vLLM still leads 1024/256 TPOT and high-concurrency output tok/s. Includes Nsight Systems direct/HTTP gap notes. | | `benchmarks/qwen-mixed-sampling-http.md` | Issue #412 HTTP mixed-sampling evidence: Qwen3-4B `/v1/completions` completed 64/64 with 32 greedy + 32 sampled requests, failed=0/timeouts=0, TTFT/TPOT/ITL/output tok/s retained; Qwen3.5-4B passed the same workload as supplemental evidence. | | `benchmarks/gemma4-31b-gh200.md` | Gemma 4 31B bf16 on one GH200 against a pinned vLLM (FlashAttention 4): the folded-pool state leads on all twelve single-request cells (E2EL −4% at 10.6K to −15% at 163K, peak memory 19 GiB lower); 32 concurrency and QPS cells with zero failures, throughput equal at 8 and 16 in flight with TTFT halved and wide-batch TPOT 8 .. 10% behind; the memory envelope at four operating points; the 31B gate run (43 of 43, the whole-prompt 32K waypoint skipped with its arithmetic); regression thresholds. | +| `benchmarks/gemma4-31b-tp2-l20.md` | Gemma 4 31B bf16 TP2 on two L20s (sm_89, PCIe, eager, `NCCL_PROTO=LL128`): 29.91 GiB of weights per rank and an 8192-token ceiling at 8 decode slots, cold load 31.4 s, 57 ms first token, 49.4 ms per step, 20/74/143 tok/s at c1/c4/c8. No single-card baseline exists (the checkpoint is 57 GiB), so parity rests on the 31B-geometry synthetic gate; the chat template answers correctly (greedy repeats — the model's non-default mode). | | `benchmarks/bs1-4k64-vllm-pegainfer.md` | RTX 5090 single-concurrency probe: `input_len=4096`, `output_len=64`, no vLLM prefix cache. PegaInfer TTFT median `177ms` vs vLLM `198ms`; TPOT median `6.47ms` vs `6.36ms`; corrected output throughput `+6%` for PegaInfer. | | `benchmarks/mixed-load-itl.md` | Qwen3-4B + Qwen3.5 mixed-load ITL (#244, #375): chunking-off sweeps via `bench_serving mixed`. Both freeze active decode for the full prefill. Qwen3 p99 blows up with prompt/QPS; the old Qwen3.5 “p99-immune” table is a **measurement artifact** (primary: hardcoded `max_batch=4` slot starvation — see #470 / `models/qwen35/mixed-load-itl-470.md`; secondary: short `bg_output_len`). Prefix reuse defeats it on Qwen3. | | `benchmarks/accuracy-eval-results.md` | Phase 1 GSM8K: Qwen3-4B PASS (pegainfer 85.37% vs HF 85.82%, delta -0.45 pp). Qwen3.5-4B historical FAIL recovered by #250 (strict 79.38%, flexible 79.30% vs HF 79.45%). | diff --git a/docs/models/deepseek-v2-lite/device-routed-moe.md b/docs/models/deepseek-v2-lite/device-routed-moe.md new file mode 100644 index 000000000..136a24f2f --- /dev/null +++ b/docs/models/deepseek-v2-lite/device-routed-moe.md @@ -0,0 +1,86 @@ +# DeepSeek-V2-Lite Device-Routed MoE + +> **TL;DR:** NCCL MoE forwards with at most eight rows keep routing, routed-expert projection, weighted reduction, and EP combine inputs on device. Each rank executes fixed route slots through pointer-batched `N=1` GEMMs; this removes the host route plan but is not an expert-sorted grouped GEMM. + +Last touched: 2026-10 + +## Scope + +The path covers NCCL MoE forwards with one to eight rows, including decode and short prefill. DeepSeek-V2-Lite has 64 routed experts with top-6 routing: rank 0 owns experts 0..31 and the model driver, while rank 1 owns experts 32..63 and replicated MoE gates. Larger forwards and the host-router or serial-expert rollback modes continue through the eager route-plan path. + +## Design + +The eager path copies router logits to the host, performs CPU softmax/top-k, builds a `MoeRoutePlan`, and replays many small expert calls. The new path keeps routing and expert execution on GPU: + +```text +rank-local hidden + -> router logits and top-6 on GPU + -> route-slot pointer arrays on GPU + -> pointer-batched gate/up + -> SiLU x up + -> pointer-batched down + -> ordered rank-local f32 reduction + -> NCCL f32 combine + -> bf16 routed output on rank 0 +``` + +Each token reserves six slots on each rank. Owned routes point to the selected expert; non-owned routes use a zero input and are discarded during reduction. This fixed layout avoids host coordination and dynamic packing. At eight rows, each projection submits 48 independent `N=1` problems through one `cublasGemmBatchedEx` call per rank. It is pointer-batched execution, not an expert-grouped GEMM: tokens selecting the same expert are not merged into a larger matrix. + +The shared expert remains on rank 0 and is added after the routed result is combined. + +## Runtime State + +`DeviceRoutedMoeRuntime` owns per-layer expert pointer tables and reusable capacity-eight scratch for both ranks. The pointer tables reference model weights, so the runtime is destroyed before the two rank models. Mutex guards serialize scratch reuse and protect the NCCL send buffers while both rank-local reductions write them. + +`begin_forward` clears route summaries once before the MoE layers; `finish_forward` reads them once afterward. The summaries check nonfinite values, route totals, ownership counts, and matching per-layer route hashes across ranks without adding a host synchronization to every layer. For finite logits, the private router producer constructs six in-range, distinct IDs in ascending order; the parity fixtures check that invariant before the consumer relies on it. + +## Numerical Contract And Fallback + +GPU routing preserves the host selection order: stable scan over 64 logits, softmax, top-6, then ascending expert ID for ordered f32 reduction. Batched cuBLAS may produce low-bit differences from the serial GEMM loop, so acceptance is based on identical route IDs and model output accuracy rather than intermediate bitwise equality. + +Host-staged execution remains the correctness oracle. More than eight rows, host-router mode, and serial-expert mode automatically use the eager route-plan path. + +## Correctness Validation + +Validation was run on two A100 40 GB GPUs with release builds. + +| Check | Result | +| --- | --- | +| DeepSeek-V2-Lite library tests | Passed | +| Router randomized / tie / near-tie fixtures | Exact route IDs and ordering at rows 1 / 4 / 8 | +| Compute Sanitizer memcheck | Passed with `0 errors` | +| HF / host-staged / NCCL case set | Token- and text-exact | +| Mixed-request NCCL E2E | Passed | +| Host-router + serial-expert fallback | Passed | +| HTTP lifecycle scenarios | Passed with healthy follow-up requests | + +The direct A/B runs also produced matching token hashes. Batched cuBLAS projections are not bitwise-equal to the serial intermediate tensors, but the retained route checks and generated outputs passed. + +## Performance Validation + +The baseline uses GPU router logits followed by CPU softmax/top-k and a host route plan. Baseline and device-routed binaries were alternated on two A100 40 GB GPUs in the same environment. Each cell below is the median of three per-run decode-step means. + +Measurements used `dsv2_lite_ep2_decode_attribution` with the same `Hello` prompt in each batch row, 16 output tokens, attribution enabled, and CUDA Graph disabled. See [Verification And Benchmarking](benchmarking.md#direct-diagnostic-benchmark) for the run command. + +| Decode batch | Baseline | Device-routed | Improvement | +| ---: | ---: | ---: | ---: | +| 1 | 41.896 ms | 27.499 ms | 34.4% | +| 4 | 83.433 ms | 50.162 ms | 39.9% | +| 8 | 136.096 ms | 77.736 ms | 42.9% | + +This path covers the current decode scheduler limit of eight live requests and short prefill with at most eight rows. The performance numbers above measure decode only. Any forward with more than eight rows continues through the eager route-plan path. + +One same-binary HTTP A/B used 32 requests with exactly 64 input and 64 output tokens at each client concurrency. The baseline enabled the host-router rollback (`PEGAINFER_DSV2_LITE_NCCL_ROUTER=host`); the device-routed run removed only that setting. Each cell was measured once without warmup. + +| Client concurrency | Baseline TPOT p50 | Device-routed TPOT p50 | TPOT change | Output tok/s change | +| ---: | ---: | ---: | ---: | ---: | +| 1 | 57.128 ms | 49.697 ms | -13.0% | +19.7% | +| 4 | 235.396 ms | 207.178 ms | -12.0% | +18.8% | +| 8 | 448.599 ms | 362.379 ms | -19.2% | +26.8% | + +All HTTP cells completed without failures or timeouts, had complete server-trace coverage, and produced matching output hashes. + +## Follow-Up Work + +1. Split timing attribution into router, pointer construction, expert projection, reduction, and NCCL combine. +2. Treat larger decode batches and prefill separately. If they become performance targets, measure expert occupancy before choosing between wider fixed slots and a true expert-grouped GEMM. diff --git a/docs/models/deepseek-v2-lite/status.md b/docs/models/deepseek-v2-lite/status.md index bc43a486f..e6766fec2 100644 --- a/docs/models/deepseek-v2-lite/status.md +++ b/docs/models/deepseek-v2-lite/status.md @@ -1,8 +1,8 @@ # DeepSeek-V2-Lite Status And Benchmark Ledger -> **TL;DR:** DeepSeek-V2-Lite keeps correctness, direct decode diagnostics, retained HTTP SLO reports, and soak readiness as separate gates. HF/host-staged/NCCL exactness and HTTP lifecycle evidence are retained; issue #466 added fixed host-staged/NCCL SLO artifacts, and the follow-up NCCL readiness fix covers the no-selector short HTTP path without claiming production readiness. +> **TL;DR:** DeepSeek-V2-Lite keeps correctness, direct decode diagnostics, retained HTTP SLO reports, and soak readiness as separate gates. NCCL MoE forwards with rows 1..8 now use device-routed pointer-batched experts with exact retained token/text outputs; this does not change the separate soak, parity, or production-readiness boundary. -Last touched: 2026-07 +Last touched: 2026-09 ## Capability Contract @@ -18,6 +18,7 @@ Last touched: 2026-07 | Device-resident NCCL combine | Available | Issue #275 keeps NCCL combine contributions/results on reusable f32 device scratch and preserves the HF / host-staged / NCCL exact gate on 2x RTX 5090. | | Device-resident NCCL dense exchange | Available | Issue #276 reuses backend-owned bf16 dense-exchange scratch, clears rank1 zero-send every exchange, removes dense-exchange stream sync from the backend call, and preserves HF / host-staged / NCCL exactness on 2x RTX 5090. | | NCCL route-plan replay | Available | Issue #277 builds a token-major host route plan once after top-k routing, replays that plan for NCCL expert launches and device contribution accumulation, keeps route counters visible, and preserves HF / host-staged / NCCL exactness on 2x RTX 5090. This remains the eager NCCL oracle path. | +| Device-routed NCCL small-row MoE | Available for rows 1..8 | Decode and short prefill keep routing and routed-expert execution on GPU through fixed route slots and pointer-batched GEMMs; the eager route-plan path remains the rollback oracle. Correctness and performance evidence lives in `device-routed-moe.md`. | | NCCL CUDA Graph readiness | Covered-shape diagnostic | Schema-2 `cuda_graph_readiness` now includes a fail-closed `full_decode_graph_probe`. The 2026-06-20 run reports capture, instantiate, replay, and verification success with `8/8` verified replays for the retained batch-1 NCCL decode step. | | First mixed-request serving gate | Available | Issue #281 adds greedy-only request admission, FCFS deferral, explicit request-local rejection/error/finish events, and one owned `DecodeCache` per active request. The 2026-06-23 2x RTX 5090 run passed HF / host-staged / NCCL exactness and the mixed-serving E2E for host-staged and NCCL. | | Long-shape NCCL collectives | Available | Issue #280 chunks large bf16 dense-exchange and f32 combine all-reduces. The 2026-06-24 2x RTX 5090 NCCL checks preserve HF / host-staged / NCCL exactness and complete 24/64/128-word direct long-shape probes. | diff --git a/docs/models/gemma4/kv-cache.md b/docs/models/gemma4/kv-cache.md index 8eccdfadc..15c76362e 100644 --- a/docs/models/gemma4/kv-cache.md +++ b/docs/models/gemma4/kv-cache.md @@ -364,9 +364,16 @@ resident origin) from day one — swapping the storage underneath later must not ## Tensor parallelism -Still unbuilt as of 2026-09: `pegainfer-gemma4` carries no TP path — no NCCL, no world size — so -one Gemma 4 model is one GPU whatever a node holds. The constraint below is what a legal world -size would have to satisfy, not something the engine enforces today. +Built as of 2026-10: `pegainfer-gemma4` shards the dense sizes across a world size, with a +per-rank KV bound to that rank's own pools and the world-size rule enforced at launch. The +invariants and the measured 12B comparison live in [`tp.md`](tp.md); what follows is the +constraint a legal world size satisfies, plus the capacity accounting it drives. + +The pools are laid out **per rank**: each rank sizes its two families from its own head counts (the +query heads shard, and the global family's KV heads shard or replicate), so a page holds +`local_kv_heads × head_dim` columns of K and V rather than the whole model's. Page *counts* stay +identical across ranks, which is what keeps the page ids in step; the rest of the KV contract is +unchanged. The qwen3 sharding policy — refuse a world size that does not divide `num_key_value_heads`, shard by integer division — cannot be reused: the full-attention group has fewer KV heads than a diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index 0fdd33fd2..0bbdc0d7f 100644 --- a/docs/models/gemma4/serving.md +++ b/docs/models/gemma4/serving.md @@ -206,7 +206,7 @@ The weights are 57.19 GiB on the device; the global pool is slots × ceiling row **The QAT W4A16 checkpoint.** `google/gemma-4-31B-it-qat-w4a16-ct` (compressed-tensors, int4 symmetric, group 32 over every text linear) serves from the same build: its weights stay 4-bit on the device (18.66 GiB), steps of up to 16 rows run TileLang W4A16 GEMMs, and wider steps dequantize each linear to bf16 for the dense GEMM. The GEMMs are compiled for one SM count: the build takes it from the visible device or from `PEGAINFER_GEMMA4_W4A16_SMS`, and without either it links the stub tier and the checkpoint is refused at load. The same load refuses a device with a different SM count. `PEGAINFER_ASYNC_PREFILL` refuses a W4A16 checkpoint, since those GEMMs need every CTA resident on the whole device. With `PEGAINFER_GLOBAL_ATTN=tilelang`, against vLLM serving the same checkpoint through Marlin on the same GH200: 1.07x output throughput at one request (9.71 against 10.36 ms per token), 1.05x at eight and sixteen with a quarter of the first-token latency, while vLLM holds the lower per-token time there by chunking prefill into its decode steps. -**What this line does not serve at 31B, and how it says so.** The NVFP4 export (`nvidia/Gemma-4-31B-IT-NVFP4`, whose dense MLPs are NVFP4) is not served: only the bf16 and W4A16 tensor sets are in the manifest, and a checkpoint whose tensors do not match it stops at the loader's manifest check, which names every mismatching tensor (`a_disagreeing_config_names_every_faulty_tensor` in the gate suite, run against this checkpoint). Tensor parallelism is not offered for this line; the launch takes the device ordinal and the graph switch and nothing else from the shared arguments. A raised ceiling without the chunk knob, the async lane beside a raise, fp8 sliding KV beside the prefix cache, and a generated state on a build without the kernels all refuse at start-up with the message the 12B sections above describe. A prompt at or past the ceiling is refused by the frontend with HTTP 400 naming the ceiling and the prompt length (measured: 9001 and 8192 tokens at the default point both refused, the next request served); `max_tokens` past what the ceiling leaves is clamped to it, as vLLM's frontend does (a 3-token prompt asking for 8192 got 8189), and the engine's own admission check, prompt plus output within the ceiling, stands behind that. +**What this line does not serve at 31B, and how it says so.** The NVFP4 export (`nvidia/Gemma-4-31B-IT-NVFP4`, whose dense MLPs are NVFP4) is not served: only the bf16 and W4A16 tensor sets are in the manifest, and a checkpoint whose tensors do not match it stops at the loader's manifest check, which names every mismatching tensor (`a_disagreeing_config_names_every_faulty_tensor` in the gate suite, run against this checkpoint). Tensor parallelism is offered for the dense sizes: `--tp-size=P` runs one rank per device with the sharded text tower and a per-rank KV, and `P` must satisfy the world-size rule in [`tp.md`](tp.md). Under TP the async lane, the generated global-attention state, the prefix cache **and the chunked walk are all refused at start-up** — so the `PEGAINFER_MIX_CHUNK_TOKENS` knob documented above is a single-rank feature. A routed (MoE) checkpoint and a W4A16 one are refused by the sharding validator rather than by the knob checks. A raised ceiling without the chunk knob, the async lane beside a raise, fp8 sliding KV beside the prefix cache, and a generated state on a build without the kernels all refuse at start-up with the message the 12B sections above describe. A prompt at or past the ceiling is refused by the frontend with HTTP 400 naming the ceiling and the prompt length (measured: 9001 and 8192 tokens at the default point both refused, the next request served); `max_tokens` past what the ceiling leaves is clamped to it, as vLLM's frontend does (a 3-token prompt asking for 8192 got 8189), and the engine's own admission check, prompt plus output within the ceiling, stands behind that. **Hardware matrix.** @@ -221,5 +221,5 @@ The weights are 57.19 GiB on the device; the global pool is slots × ceiling row - **An admission rides the live decode batch instead of freezing it.** With streams in flight, newcomers' prompts share one eager step with the decode batch: the prompt rows sit in the row prefix as segments, every active stream advances its token in the suffix, and one sampler call covers every newcomer's first token and every active row. Measured with a streaming request underneath a flood of one-token requests: its inter-token gap stays at about 30 ms — one mixed step — whether 16, 48 or 96 requests are queued, where the frozen-prefill scheduling this replaced measured about 500 ms at the same depths. A burst of sixteen coincident ~120-token prompts folds its admission staircase — TTFT p50 and p99 both drop 20-30% — while the stream's worst gap stays bounded by the gathered step's own cost (~130 ms at the 512-row ceiling); a 16 × ~1900-token flood is untouched, since a long prompt keeps its own step. Admission work per turn stays bounded by the slot ceiling, gathered or not. - **Whole-prompt prefill by default.** A prompt runs whole in one step unless the opt-in chunked walk above bounds the step's prompt rows; with the knob set a walker reserves pages round by round instead of holding its whole prompt. - **No cross-request prefix sharing.** Two live requests with a common prefix pay for it twice; the opt-in conversation cache above serves consecutive turns of one conversation, not concurrent requests. -- **Single GPU.** No tensor parallelism for this line yet. +- **TP for the dense sizes, one rank per device.** The budget above is per rank: at 31B the default point does not quite fit a 48 GiB card even at TP2 — 29.91 GiB of weights plus 14.38 GiB of pools is 44.29 GiB, and the ~0.90 GiB of context, arena and workspace on top puts it just over the 44.99 GiB an L20 leaves usable — so TP2 serving needs a reduced envelope, and the measured point is an 8192-token ceiling at 8 decode slots (57 ms first token, 49.4 ms per step, 143 tok/s at 8 concurrent, on two L20s). See [`tp.md`](tp.md) and `docs/benchmarks/gemma4-31b-tp2-l20.md`. - **Token capacity is not reported as `EngineInfo::kv_capacity`**: the engine leaves the token-fit hint empty because the two independent page families cannot be represented by one token count. Runtime load metrics still report the sum of used and total local/global pages. diff --git a/docs/models/gemma4/tp.md b/docs/models/gemma4/tp.md new file mode 100644 index 000000000..20ad335de --- /dev/null +++ b/docs/models/gemma4/tp.md @@ -0,0 +1,157 @@ +# Gemma 4 tensor parallelism + +TL;DR: dense Gemma 4 runs as one rank per device: a per-rank weight shard, one `all_reduce` after `o_proj` and one after `down`, and a per-rank KV bound to that rank's own pools, all driven lock-step from one scheduler thread. 12B bf16 TP2 is verified against a single-rank run; 31B bf16 serves on the 48 GiB-class sm_89 pair (L20/L40S/RTX 6000 Ada) with 8 decode slots and a `PEGAINFER_MAX_CONTEXT` ceiling, measured on two L20s: 57 ms first token, 49.4 ms per step, 143 tok/s at 8 concurrent (`benchmarks/gemma4-31b-tp2-l20.md`). The need is capacity, not the card — ~59.8 GiB of weights do not fit one device under ~60 GiB usable, whatever its model. A cold two-rank process must warm the cuBLAS kernels before the communicator exists, or the first GEMM deadlocks in the driver's module loader. Decode graphs work under TP: the sweep interleaves phases across ranks, and each rank releases its graphs before the communicators drop. Gate: `engine::lane_gates_tp::the_two_rank_engine_matches_one_rank`. + +Last touched: 2026-10 + +## What is sharded, and what is not + +The text tower's linears split the qwen3 way — output-sharded rows or input-sharded columns — but the head counts differ per layer kind, so the row ranges do too. + +| tensor | shard | 12B TP2 (per rank) | 31B TP2 (per rank) | +| --- | --- | --- | --- | +| `q_proj` rows | contiguous query heads, `Q / P` | 8 × 256 | 16 × 256 sliding / 16 × 512 global | +| `k_proj` rows, sliding | `Kv / P` | 4 × 256 | 8 × 256 | +| `v_proj` rows (sliding only) | same as `k_proj` | 4 × 256 | 8 × 256 | +| `k_proj` rows, global | `G % P == 0` → shard, `P % G == 0` → replicate | 1 × 512 (replicated) | 2 × 512 (sharded) | +| `o_proj` cols | the same query-head run | 8 × 256 | 4096 sliding / 8192 global | +| `gate`, `up` rows | `intermediate / P` | 7680 | 10752 | +| `down` cols | `intermediate / P` | 7680 | 10752 | + +Published per-rank counts, from each checkpoint's own `config.json` (31B: `Q` 32, sliding `Kv` 16, `G` 4, head dims 256/512, hidden 5376, intermediate 21504, 50 sliding + 10 global layers). + +Replicated on every rank: the token embedding (and the tied head), all four layer norms, `q_norm` / `k_norm` (whole heads are held, so the vectors stay whole), `layer_scalar`, and both rope tables. The residual stream and every head width stay whole; only head *counts* and the MLP width shard. + +## The world size a shard decision has to satisfy + +`TensorParallelConfig::validate_for` refuses a launch that would silently drop or misalign heads. The three counts the kv-cache design doc states are necessary but not sufficient — they do not keep the **per-rank GQA group** integral (`Q = 8`, `Kv = 6`, `P = 2` clears them and yields group `4/3`), so that is checked as well: + +``` +P > 0 +rank < P +Q % P == 0 +Kv % P == 0 +G % P == 0 || P % G == 0 +(Q/P) % (Kv/P) == 0 # per-rank sliding group +(Q/P) % G_local == 0 # per-rank global group, sharding branch only +intermediate % P == 0 +``` + +The global check carries a branch condition because only the **sharding** branch can split a group: `G % P == 0` hands a rank `G / P` KV heads, so `(Q/P) / (G/P)` has to be integral. The replicate branch (`P % G == 0`) gives a rank one whole KV head, and `P % G == 0` already keeps its `Q / P` query heads inside that head's group — the `P / G` ranks holding that head split its `Q / G` query heads evenly, so `Q / P` is `(Q / G) / (P / G)`, not a multiple of `Q / G` — so there is nothing left to police, and the check would only compare against a group of one. + +A routed (MoE) checkpoint and a W4A16 one are refused under TP > 1: the first needs an expert-sharding design, the second projects whole matrices. + +## How a step runs + +One scheduler thread owns every rank, but a step drives them **concurrently**: rank 0's segment runs on the scheduler thread and each extra rank's on its own scoped thread, then all are joined (`EngineState::drive_ranks`). One thread per rank is what keeps a host-blocking operation inside a rank's own segment from stranding the others — see "Known bounds" for the stall it fixes. Each rank owns its `DeviceContext`, `GemmaServe` and `StepArena`; a request's KV is **one `GemmaKv` per rank**, each minted from that rank's pools (a `GemmaKv` is pool-bound by construction, so this is what `admit_tokens`' `belongs_to` check requires). `GemmaKv` derefs to rank 0's families, which is why the whole single-rank engine and serve code paths are unchanged. + +That per-rank KV is why no host/device split was needed: each rank's entry point advances its own frontier, and because the pools have identical budgets and the admission sequence is identical, the page ids stay in step across ranks. + +Per step, each rank's segment calls `activate_rank` (set the device, bind its context, make its thread-local cuBLAS handles current — creating them on first use), runs that rank's segment, and drains it: the extras on their own threads, rank 0 on the scheduler thread before the sampler, which runs on rank 0 alone. The drain is what turns a device fault on a non-primary rank into a named error instead of the primary's collective stalling forever. + +Why the extras get threads is the invariant the single-threaded driver could not hold: **a host blocking point must come after every rank's launch.** `all_reduce_rows` enqueues and returns, but a prefill's tower reaches a host-side wait of its own — rank 0's per-row readback, and a GEMM workspace allocation behind it — so with one thread driving every rank in turn, rank 0 could block on a collective whose peer call had not been issued yet and hang. Running each extra rank's segment on its own thread, and joining only after rank 0's has returned, makes that ordering impossible rather than something each call site has to arrange; it is also what lifted the old prompt ceiling (see "Known bounds"). + +The staged decode pipeline (ids written by the previous step's sampler) is disabled under TP: a non-primary rank has no sampler and no ids of its own, so every rank takes the explicit-token path. + +## The collective + +`all_reduce_in_place` on the live extent (`hidden_size × seq_len`) of the projection buffer, inserted in `attention_epilogue_into` after `o_proj` and after `down` — both are row-parallel sums, and both are reduced before the residual they feed is formed. One communicator per rank, built on that rank's compute stream so the reduction lands inside a captured graph once graphs are enabled. + +## Cold start: warm the kernels before the communicator + +Under the default `CUDA_MODULE_LOADING=LAZY` a fresh process enters the CUDA driver's module loader on the first GEMM of *each shape*. `Comm::from_devices` starts NCCL's proxy threads, which enter that same loader while the communicator comes up; a GEMM that materializes a shape for the first time after that point deadlocks — the engine thread spins at 100% CPU inside `cuLibraryGetModule` (reached through `cublasGemmEx` → `cublasLtTSTMatmul`), never issues rank 1's tower, and the request never returns. It only bites a **cold** process: any run that already did a single-rank pass in the same process (the gate, a warm box) reuses the kernels that pass loaded, which is why it can hide behind a passing gate. + +The engine therefore runs one tower pass per rank at load, **before** the communicator is built, at a spread of row counts across cublasLt's kernel-selection regions (`1, 2, 3, 4, 8, …, 1024`, capped by the serving ceiling; larger prompts and every decode bucket reuse those kernels — verified by sweeping prompts of 5…6000 tokens). With no comm alive the loads happen single-threaded, and the serving path never touches the loader again. `CUDA_MODULE_LOADING=EAGER` is the equivalent sledgehammer and also works; the warm is preferred because it is per-rank and bounded by the rows the tower actually uses. The cost is a few extra tower passes folded into load — not measured on their own, but the real 31B's cold load is 30 s with them folded in (`benchmarks/gemma4-31b-tp2-l20.md`). + +## Memory: the pools are per-layer arrays + +Each family's pool is `layers × pages × bytes-per-page`, which is the number that decides whether a 31B configuration fits a 48 GiB card. + +| | local family | global family | +| --- | --- | --- | +| bytes per token | `kv_heads × head_dim × 2 (K,V) × 2` | same | +| 31B TP2 per token | `8 × 256 × 4` = 8 KiB | `2 × 512 × 4` = 4 KiB | +| pages at ceiling 8192 / 16 slots | `15 × 17 + 128 + 1` = 384 | `16 × 128 + 1` = 2049 | +| pool bytes | 50 × 384 × 0.5 MiB = 9600 MiB = 9.38 GiB | 10 × 2049 × 0.25 MiB = 5122.5 MiB = 5.00 GiB | + +Every figure here and below is binary (1 GiB = 1024 MiB); the MiB columns are what the start-up log prints, so the two are directly comparable. + +31B TP2 therefore needs **29.91 GiB of weights plus 14.38 GiB of pools = 44.29 GiB per rank at the default point**, before the ~0.90 GiB of CUDA context, step arena, cuBLAS workspace and NCCL buffers that the same run reports as the gap between 29.91 GiB of weights and the 14.18 GiB it leaves free on a 44.99 GiB L20. That is **≈45.2 GiB against 44.99 GiB usable**, so the default point does not quite fit — by ~0.2 GiB, not by a wide margin, which is why the reduced envelopes below are the ones that were served. The working envelopes measured on L20: + +| ceiling | slots | pools | ≈ resident per rank | +| --- | ---: | ---: | ---: | +| 8192 | 8 | 8.6 GiB | 38.5 GiB | +| 4096 | 16 | 10.3 GiB | 40.2 GiB | +| 2048 | 2 | 1.4 GiB | 31.3 GiB | + +The startup error names the page counts when a budget cannot be allocated, and the global pool's size is logged once it is. + +The 8192 x 8 row is the one that has been served: 2.50 GiB global pool, 14.18 GiB free after weights, cold load 31.4 s, and 57 ms / 49.4 ms / 143 tok/s (TTFT / step / 8-way aggregate). The full run and method are in `benchmarks/gemma4-31b-tp2-l20.md`. + +## What is refused under TP today + +`PEGAINFER_ASYNC_PREFILL` (its lane stream cannot be lock-stepped across ranks), `PEGAINFER_GLOBAL_ATTN=tilelang*` (the generated kernels are compiled for the whole global family), `PEGAINFER_PREFIX_CACHE`, and `PEGAINFER_MIX_CHUNK_TOKENS` (a chunked walk's rounds gather across prompts, which the TP gates do not cover). The prompt ceiling that used to be refused here (`PEGAINFER_TP_MAX_PROMPT`) is gone with the stall it guarded — see "Known bounds". + +## Decode graphs under TP + +`--cuda-graph=true` works under TP. The sweep interleaves **phase by phase across ranks** — for each bucket, each phase, each rank — because a `Warm`/`Launch` phase executes and enqueues its rank's all-reduce, whose peer call has to be in flight for the step to finish; `Capture` only records, and a recorded collective replays when its peer replays. One scheduler thread is enough because the phases, not whole sweeps, are the unit that has to line up. + +Two environment constraints are baked into the code. The **kernels must be warm before the communicator exists** (see "Cold start"): a first-shape GEMM that materializes after `Comm::from_devices` wedges in the driver's module loader, which is what made the earlier capture attempts look like a capture defect. And **every rank must release its graphs before the communicators drop**: a captured collective bakes in NCCL kernel launches, and `ncclCommAbort` wedges while a graph that references them is still alive, so `EngineState::drop` releases them first, on each rank's own device. + +Parity is the same gate, with `PEGAINFER_TP_GRAPH=1`: two ranks captured against one rank eager, `48/48` picks and worst gap `0.0000` on the 31B-geometry checkpoint. The server starts, serves and shuts down cleanly with `--tp-size=2 --cuda-graph=true`. + +Captured against eager is a wash on two L20s at 31B (8192 x 8), back to back on a quiet box: first token 57.0 vs 56.5 ms, decode step 51.5 vs 49.4 ms (+4%), 8-way aggregate 146.0 vs 143.7 tok/s. The collective, not launch count, dominates the step here, so graphs buy little; the flag is on by default (as at TP1) and `--cuda-graph=false` is the marginally faster per-step choice at low concurrency. + +## Verification + +`engine::lane_gates_tp::the_two_rank_engine_matches_one_rank` starts a real engine twice — once with one rank, once with two — over the same 12B checkpoint, three prompts and one batch, and compares the requested top-8 logprobs. + +It holds the *distributions*, not the greedy tokens: a two-rank reduction sums the same products in a different order and NCCL writes bf16 back, so the logits differ in their last bits and a near-tie can flip the pick. The gate prints the build and the device pair it ran on (`tp2: NCCL 2.18.3, devices 1,2`), because a reading is only comparable with the environment it came from — but nothing here claims *why* two readings from different environments differ. The line comes from this gate's own readings, one procedure, one head (`45eb4a68`, four runs of the 12B on two L20s, `NCCL_PROTO=LL128`): + +| claim | result | +| --- | --- | +| one-rank run twice (control) | bit-identical, **asserted** token for token and logprob for logprob, so the comparison is not measuring harness noise | +| steps keeping the one-rank pick | 22 / 24, the same in all four runs | +| worst picked-token logprob gap | 0.5738 (top 0.2362), the same to four decimals in all four runs | +| a differing pick | always a genuine near-tie: each pick inside the other run's top-8, and both picks' own gaps folded into the bound above before the comparison stops | +| **one `all_reduce` skipped** — the fault the line is for | the gate fails: at request 0 step 0 the one-rank pick leaves the two-rank top-8 (the near-tie rule), and with that rule relaxed the top-token gap reads 2.1895 | + +The fault is the `o_proj` reduction commented out in `layer.rs`: both ranks skip it, so the run still completes, and the single-rank control is untouched (at world size 1 the call is a no-op). `worst_pick` reads `0.0000` on that run — the two top-8 lists share no token, so the picked-token bound has nothing left to measure, which is the case the near-tie rule exists to catch — so the rule and the top-token bound are what catch a fault this size, and `LOGBROB_LINE` sits between the two readings instead of under the noise. `LOGBROB_LINE` is `1.0` — it was `0.5`, which sits below this floor and so fails a correct run; 1.0 is ~1.7x over the floor and ~2.2x under the fault, and it is the value `DRIFT_LINE` below already uses. The repo's one precedent for calibrating a quantity of this kind is `serve_oracle`'s `neutral_scale` (two algorithms over one context, 0.31..5.75 observed, line 12.0), i.e. roughly 2x the worst observed drift. What the line is for is a structural error — a wrong shard, a missing reduction, a rank out of step — which moves a logprob by many nats and flips most picks, not 2 of 24; reduction noise is policed by the near-tie rule beside it. + +The **shard branch** (`G % P == 0`, which the published 12B never takes — its single global KV head is replicated) is gated with a synthetic checkpoint carrying the real 31B shapes (`Q` 32, `G` 4, head dims 256/512, hidden 5376, intermediate 21504) cut down to six layers, so the whole run is ~8 GiB and takes seconds on any pair. It **measured bit-identical**: 48/48 picks, worst picked-token gap `0.0000`. That is a measurement, not an assertion — the two-rank comparison asserts against `LOGBROB_LINE`, the same line for both branches, so a future run that drifts within the line passes while this sentence's `0.0000` no longer holds. + +The gate serves its first prompt on its own and only then the rest as one batch, so the **solo** admission path (`step` + `prefill_extra_ranks`) — where the cold-start hazard above first surfaced, and the only path a lone short request takes — is compared on every run; `PEGAINFER_TP_PROMPTS` / `PEGAINFER_TP_PROMPT_TOKENS` still widen the set. Which branch the gate covers is the checkpoint's: point `PEGAINFER_TEST_MODEL_PATH` at the 12B for the replicate branch, at a 31B-geometry checkpoint for the shard one. **A full-depth 31B cannot be gated against a *single-rank* baseline here**, not for want of a fixture but because the gate needs a single-rank control and 57 GiB of weights do not fit one card — the six-layer synthetic is exactly that shape at a depth that fits. So the numbers above are the 12B (replicate) and the synthetic (shard); the 60-layer real checkpoint is gated instead against its own Hugging Face dump (below) and served end to end on two L20s. + +The single-rank suite is unchanged by the TP path (`cargo test --release -p pegainfer-gemma4 --features gemma4 --lib`), and now also carries the prefill-fault gate named above, which runs on one GPU. + +The **real 31B checkpoint**, which no single card holds, is served end to end on two L20s; its load, envelope and serving numbers are in `benchmarks/gemma4-31b-tp2-l20.md`. + +### The real 31B against its own Hugging Face dump + +`engine::lane_gates_tp::the_two_rank_engine_matches_the_hf_reference` compares the two-rank engine directly against a Hugging Face reference for the **60-layer 31B** — the checkpoint whose single-rank baseline one card cannot hold. The fixture is that checkpoint's own dump (`tools/accuracy/dump_gemma4_hf_golden.py`, pinned by file digests), compared over the teacher-forced prompt rows' top-64 logprobs. Measured on two L20s, nine-token prompt: + +| claim | result | +| --- | --- | +| reference rows keeping the engine's pick | 7 / 8 | +| worst shared-token logprob gap | 0.6753 | +| the one differing row | a near-tie: the engine's pick sits inside the reference's top-64 | + +So the real geometry (hidden 5376, `Q` 32 / `Kv` 16 / `G` 4, head dims 256/512, the 50 sliding + 10 global split) matches the reference over all 60 layers — with the same `NCCL_PROTO=LL128` requirement. The tolerance here is **not** the synthetic gate's: that one is bit-identical (`0.0000` against `LOGBROB_LINE`), this one is `0.6753` against a `1.0` line. `DRIFT_LINE` is a first cut — calibrating it from this comparison on the 12B at one rank is a follow-up. + +## Known bounds + +- **A rank-0 step failure stops the engine — but only with extra ranks.** Rank 0's tower runs on the scheduler thread while each extra rank's runs on its own (see "How a step runs"): on failure rank 0's graphs are released and its communicator detached *before* the join, so a peer that is already waiting on a call which will never come is unblocked instead of hanging, and every later step refuses — so the engine cannot serve a comm-less reduction that would return partial sums. A failure on an extra rank is fatal too. At **world size 1** there is no peer to unblock and no frontier to keep in step, so the same failure stays what it has always been: that one request fails and the driver keeps serving (`drive_ranks` marks the engine broken only when `more` is non-empty). `engine::lane_gates_logprobs::a_failed_prefill_costs_that_request_not_the_engine` pins the single-rank half on both prefill paths, using a token id outside the embedding — `prepare_single`'s `validate_tokens` refuses it before the tower allocates anything, so the fault needs no injection hook. **The multi-rank half is still untested**: no gate injects a rank-0 or extra-rank failure at TP2, so the abort itself — destroying every rank's captured graph execs and dropping its communicator, with no stream drained first — is argued from the teardown order rather than exercised. +- **The long-prompt prefill stall is fixed by the concurrent driver.** A prefill's tower reaches a host-side device sync on rank 0 (a GEMM workspace allocation, behind its per-row readback) at longer prompts; a single driving thread would then wait on rank 0's reduction before any peer had launched its collectives, and hang. The ranks are now driven **concurrently** (`EngineState::drive_ranks`: rank 0 on the scheduler thread, one scoped thread per extra rank, then join), so every peer's collectives are in flight before rank 0 can block. Verified on the 60-layer 31B at TP2: a scored prefill of 128, 256 and 1024 tokens all pass where 128 and up used to stall; the three TP gates (one-rank comparison `48/48` at gap `0.0000`, the scored probe, the HF reference) all pass on the 2×L20 pair. The cost is one thread spawn per rank per step. The prompt ceiling of the single-threaded driver (`PEGAINFER_TP_MAX_PROMPT`) and its refusal are gone with the stall, rather than left as dead scaffolding. +- **Page-id agreement is inspected, not enforced.** Every rank's pools must stay identical in free-page order — that is the whole basis for "identical page ids". A debug build compares free-page counts after every admission, and every prefill compares the ranks' frontiers; a release build only does the frontiers, and nothing repairs a divergence. Three of the four admission sites self-heal, because the request's `GemmaKv` is dropped and both ranks' pages come back with it; the fourth (`ready_rows_pinned`, which keeps the row in the batch) cannot, so an extra rank that refuses there stops the engine instead of carrying the divergence into a step. +- **The chunked walk's TP code has no runtime.** `PEGAINFER_MIX_CHUNK_TOKENS` is refused under TP, so the walker's extra-rank admission and its `tp_broken` check are unreachable with a non-empty `more` — kept because the concurrent-driver follow-up re-enables the knob, untested until then. +- **The vocabulary projection is replicated.** `embed_tokens` and the tied head are `Whole` on each rank (31B: 262144 × 5376 bf16, ~2.8 GiB per rank), and every extra rank computes a whole batch of `vocab × rows` logits per decode step that is then discarded. Dropping it needs a head-free tower for the extra ranks — in graph mode that means a second per-bucket capture, because the head sits inside `decode_gpu_body`. Not implemented, not measured. +- **The 1024-token reference comparison is not gated yet — its drift is the known chaotic band.** With the concurrent driver a 1024-token (window-width) prompt no longer stalls, but compared against the HF dump its worst shared-token logprob gap is **8.37** and four rows have a pick outside the other's top-64. That magnitude is the shape-dependence the repo already measures: `serve_oracle`'s `neutral_scale` scores the same context two ways and records **0.31..5.75 raw logits**, with its line at **12.0** because the quantity is chaotic and a tighter one only flakes. A prefix-consistency check confirms it is that and not a page/position error: the same 512-token prefix scored in a 1024-token run differs by 4.05 on the 31B and 1.69 on the four-layer synthetic, from row 0 on. So the nine-token case is gated (7/8 rows, worst gap 0.6753) and the 1024-token case needs the tolerance-plus-top-1-share discipline of `serve_oracle` — its stored per-case tolerance and `top1_share` floor — rather than a strict top-k containment, which cannot hold at this depth; that calibration is the follow-up. +- **`NCCL_PROTO=LL128` is a requirement the code reports rather than enforces.** Startup logs the NCCL version and the effective protocol and warns when it is not LL128; it does not refuse. On this container's stock NCCL the default protocol corrupts the all-reduce buffer (see the operational notes). + +## Operational notes + +- **A pair, not a card.** Two ranks are two devices; the gate runner claims both (`require_devices 2`, `PEGAINFER_GATE_GPU=a,b`) and exports them as `CUDA_VISIBLE_DEVICES=a,b`, so the ranks inside the process are devices 0 and 1. +- **Set `NCCL_PROTO=LL128` in this container.** With the pod's stock NCCL 2.18.3 (2023, CUDA 12.2 vintage) on CUDA 12.9 and sm_89, the default `LL` protocol — and `SIMPLE` — make the all-reduce kernel write far outside its buffer (`compute-sanitizer` names `ncclKernel_AllReduce_RING_LL_Sum`, ~95 GB past the allocation). A non-primary rank faults asynchronously, so the visible symptom is the primary's collective stalling forever rather than an error. `LL128` is sound, and the repo's own L20 precedent ran NCCL 2.32.3. Prefer a newer NCCL over the variable where the environment allows it. +- **The checkpoint is read once per rank.** 62 GB does not stay in a 32 GB page cache, so a 31B load is ~3.5 min per rank from the PVC and ~15 s from node-local disk — copy it local for iteration. +- **In a container, disable the fabrics NCCL cannot reach** (`NCCL_IB_DISABLE=1`, and `NCCL_P2P_DISABLE=1` if P2P is unavailable). The log otherwise shows `Unable to open device mlx5_*` warnings. +- **Check the pair is healthy.** On a shared box, an uncorrectable DRAM fault shows up as a pair-dependent illegal memory access, which is not a TP defect. diff --git a/docs/models/qwen35/batched-step-tail.md b/docs/models/qwen35/batched-step-tail.md index 344350f36..e3a338044 100644 --- a/docs/models/qwen35/batched-step-tail.md +++ b/docs/models/qwen35/batched-step-tail.md @@ -1,8 +1,8 @@ # Qwen3.5 Batched Step Tail -> **TL;DR:** Qwen3.5 issue #353 is implemented in the local branch: prefill gathers per-request last hidden rows, runs batched offset RMSNorm + one lm_head GEMM, and scheduler/executor decode sample from batched logits. Full-vocab host copies now happen only for `logprobs > 0`. HF logits + scheduler e2e pass; benchmark evidence supports a first-token/short-output TTFT claim only. +> **TL;DR:** Implementation record for Qwen3.5 issue #353: prefill gathers per-request last hidden rows, runs batched offset RMSNorm + one lm_head GEMM, and scheduler/executor decode sample from batched logits. Full-vocab host copies happen only for `logprobs > 0`. HF logits + scheduler e2e passed; benchmark evidence supports a first-token/short-output TTFT claim only. > -> **Last touched:** 2026-06 +> **Last touched:** 2026-10 ## Scope @@ -72,4 +72,4 @@ Long-output diagnostic: - Read-only DeepSeek diff review found no blocker. - Follow-up review cleaned duplicate logprobs helpers into `src/logprobs.rs` and kept docs public-facing with placeholders rather than machine paths. - Fail-fast logprobs snapshot errors are intentional; silent logprobs loss is worse than an explicit request error in this correctness-sensitive path. -- Final review flagged a `concurrency=64` shutdown SIGSEGV in `cublas_destroy`; qwen35 now returns an `EngineHandle` that joins the scheduler thread before process teardown. +- Final review flagged a `concurrency=64` shutdown SIGSEGV in `cublas_destroy`. The server must join the scheduler thread before process teardown; the current `Engine` carries that join handle in `LiveScheduler`. diff --git a/docs/models/qwen35/load-snapshot.md b/docs/models/qwen35/load-snapshot.md index dcbad0ab7..034d7478a 100644 --- a/docs/models/qwen35/load-snapshot.md +++ b/docs/models/qwen35/load-snapshot.md @@ -1,165 +1,65 @@ -# Qwen3.5 Scheduler LoadSnapshot +# Qwen3.5 Scheduler Metrics -> **TL;DR:** Qwen3.5 publishes one logical post-drain/post-prune `LoadSnapshot` stream from its shared single-GPU/TP scheduler: running counts active and prefilling requests, waiting counts all current pending work, and KV usage is request-page capacity minus available pages. +> **TL;DR:** Single-GPU and TP Qwen3.5 expose one logical scheduler's post-step `SchedulerMetrics` through the shared driver; in-flight prefill counts as running work. > -> **Last touched:** 2026-08 +> **Last touched:** 2026-10 -## Preparation - -- **Read**: - - `docs/index.md` — located the Qwen3.5 model and frontend observability documentation. - - `docs/models/qwen35/model-crate.md` — confirmed that the model crate owns the scheduler and exposes it through the generic `EngineHandle`. - - `docs/models/qwen35/roadmap.md` — confirmed the serving and lifecycle observability context. - - `docs/subsystems/frontend/prometheus-metrics.md` — confirmed the existing `LoadSnapshot` bridge contract. - - `pegainfer-qwen35/src/scheduler.rs` before and after the metrics change — compared the wiring with the shared single-GPU/TP scheduler flow. -- **Relevant history**: - - Qwen3 established the `LoadSnapshot` watch path consumed by the frontend bridge. - - Qwen3.5 shares one scheduler loop between single-GPU and TP backends, so KV accounting must come from `SchedulerBackend`, not directly from `Qwen35Model`. -- **Plan**: - 1. Publish backend-neutral snapshots from existing Qwen3.5 scheduler boundaries. - 2. Attach one load watch to the single-GPU and TP engine handles. - 3. Validate field accounting and idle reset with the existing scheduler E2E, HTTP benchmark, and raw `/metrics` sampling. - -## Design - -The data path reuses the existing frontend contract: +## Publication boundary ```text -Qwen3.5 SchedulerBackend - -> LoadSnapshot watch - -> EngineHandle - -> LocalEngineBridge +Qwen35Scheduler::metrics + -> shared driver metrics cell + -> SchedulerHandle::metrics + -> SteppedEngineBridge -> SchedulerStats -> /metrics ``` -Both Qwen3.5 execution modes own one logical request stream, so single-GPU and TP each attach one `EngineHandle::with_load_watch` receiver. The frontend bridge, metric names, labels, and scheduler-stat conversion remain unchanged. +The driver drains submissions before calling `Scheduler::step`. The scheduler +settles ready asynchronous prefill and prunes cancelled pending, active and +prefilling requests before admission and planning. A ready prefill completion +ends the step after cancellation cleanup, so its updates are committed before +the next model execution. Admission and planning resume on the next step. +The driver publishes a post-step snapshot before committing that step's +request updates. The asynchronous +bridge reads the latest metrics cell when forwarding a batch; the driver may +have advanced again by then, so metrics are not an immutable snapshot tied +to that batch. -Each scheduler tick first merges deferred work with every submission currently available, then prunes closed pending, active, and prefilling requests before publishing. The fixed boundary is `drain -> prune -> publish load -> admission -> plan`. If the idle scheduler wakes through `blocking_recv()`, it drains, prunes, and publishes again before admission so work closed before admission never consumes a slot or appears in the snapshot. +The metric names and labels are shared with the other model lines; see +[Prometheus metrics](../../subsystems/frontend/prometheus-metrics.md) for the +wire mapping. -Snapshot accounting is: +## Accounting -| Metric field | Existing Qwen3.5 state | +| Metric field | Live scheduler state | | --- | --- | -| `num_running_reqs` | `active.len() + prefilling.len()` | -| `num_waiting_reqs` | the merged pending queue: prior deferred work plus newly drained submissions | -| `kv_used_blocks` | request KV capacity minus currently available request pages | -| `kv_total_blocks` | backend request KV capacity, excluding the CUDA Graph padding page | - -Publication reads the scheduler's queues and KV allocator after closed resident state has gone through its normal retirement path. The snapshot therefore describes the state used by the following admission decision: cancelled residents no longer count as running or hold capacity, while live pending requests count as waiting even if they were submitted during the current tick. - -The live gate uses `scripts/bench_http_serving.py` to create overlapping HTTP traffic and a 100 ms `curl /metrics` sampler to retain the three labeled gauges. - -## Execution Log - -- Added load watches to `start_with_capacity` and `start_tp_with_capacity` and attached each receiver to its engine handle. -- Added direct, backend-neutral `LoadSnapshot` publication in the shared scheduler loop, following Qwen3's instrumentation shape. -- Derived KV capacity and availability through `SchedulerBackend`, so the same publication logic serves single-GPU and TP. -- Validated the single-GPU path with the existing scheduler E2E and live HTTP pressure: running and KV usage rose during generation, waiting reached three at `--max-batch 1`, and every gauge returned to zero after drain and recovery. -- Updated the shared Prometheus documentation for Qwen3.5's one-logical-engine contract. -- P2A step 2 moved publication after drain and cancellation pruning. Focused CPU tests cover closed pending and resident work, and a real TP1 `max_batch=1` gate proves a cancelled resident disappears from the post-prune load, frees capacity for same-tick admission, and leaves running, waiting, and KV usage at zero after recovery. - -## Validation Boundary - -The single-GPU NVIDIA run was captured at `a033258c1de1944469d6c6335d4a36d4a80192cf` on an RTX 5090 with driver `580.105.08`, CUDA toolkit `12.8.93`, Rust nightly `1.99.0`, Triton `3.6.0`, and model revision `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a`. Commit `6c9b7d2b8464a846414605fdbde9020887f18ee7` was not rerun; its scheduler diff only inlines the snapshot construction and preserves its fields and publication point. - -Build and existing scheduler E2E: - -```bash -export PEGAINFER_CUDA_SM=120 -export PEGAINFER_TRITON_PYTHON="$PWD/.venv/bin/python" -export PEGAINFER_TEST_MODEL_PATH="$PWD/models/Qwen3.5-4B" - -cargo build --release -p pegainfer-server --features qwen35 -cargo test --release -p pegainfer-qwen35 --features qwen35 \ - --test e2e_scheduler test_e2e_qwen35_scheduler -- --exact --nocapture -``` - -The release build passed and the existing E2E reported `1 passed; 0 failed`. - -Server and 100 ms metric sampler: - -```bash -RUST_LOG=info target/release/pegainfer \ - --model-path models/Qwen3.5-4B \ - --served-model-name qwen35-metrics \ - --port 18080 --device-ordinal 0 --tp-size 1 \ - --cuda-graph=true --max-batch 1 --max-prefill-tokens 1024 - -while :; do - date -Ins - curl -fsS http://127.0.0.1:18080/metrics \ - | grep -E '^vllm:(num_requests_running|num_requests_waiting|kv_cache_usage_perc)\{' \ - | grep -F 'engine="0"' \ - | grep -F 'model_name="qwen35-metrics"' - sleep 0.1 -done > metrics-pressure.log -``` - -Real batch-slot pressure used the repository's existing benchmark. The provenance arguments below are preserved exactly as recorded in the raw artifact: - -```bash -python3 scripts/bench_http_serving.py \ - --base-url http://127.0.0.1:18080 \ - --model qwen35-metrics \ - --num-requests 4 --concurrency 4 --warmup 0 \ - --prompt-words 32 --max-tokens 512 \ - --temperature 0 --top-k 0 --top-p 1 --ignore-eos \ - --timeout 300 \ - --model-path models/Qwen3.5-4B \ - --commit a033258c1de1944469d6c6335d4a36d4a80192cf \ - --source-revision a033258c1de1944469d6c6335d4a36d4a80192cf \ - --model-revision 851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a \ - --server-binary target/release/pegainfer \ - --claim-boundary "Qwen3.5 LoadSnapshot live metrics pressure and recovery only" \ - --out pressure.json -``` - -Traffic completed without failures or timeouts: - -```text -completed=4 -failed=0 -timeouts=0 -wall_s=12.9674 -``` - -One pressure sample and the observed peaks were: - -```text -vllm:num_requests_running{model_name="qwen35-metrics",engine="0"} 1 -vllm:num_requests_waiting{model_name="qwen35-metrics",engine="0"} 3 -vllm:kv_cache_usage_perc{model_name="qwen35-metrics",engine="0"} 0.00008846686915750052 - -max_running=1 -max_waiting=3 -max_kv_cache_usage_perc=0.0010026245171183392 -``` - -After the workload drained, all three gauges returned to zero. A follow-up request then validated recovery: - -```bash -curl -fsS http://127.0.0.1:18080/v1/completions \ - -H 'Content-Type: application/json' \ - --data-binary '{"model":"qwen35-metrics","prompt":"Hello","max_tokens":8,"temperature":0,"top_k":0,"top_p":1,"ignore_eos":true,"stream":false}' -``` - -```text -usage={"completion_tokens":8,"prompt_tokens":1,"total_tokens":9} - -vllm:num_requests_running{model_name="qwen35-metrics",engine="0"} 0 -vllm:num_requests_waiting{model_name="qwen35-metrics",engine="0"} 0 -vllm:kv_cache_usage_perc{model_name="qwen35-metrics",engine="0"} 0.0 -``` - -The server exited cleanly and the metric sampler reported no errors. - -## Debrief - -- **Outcome**: Qwen3.5 feeds one logical `LoadSnapshot` stream to the frontend for both single-GPU and TP. Publication now occurs after current submissions are drained and closed work is pruned, so the snapshot is the admission boundary rather than a view of only the previous tick. -- **Pitfalls encountered**: - - The TP scheduler rebase required KV accounting through `SchedulerBackend`; retaining model-specific `model.kv_pool()` access would not compile against the shared loop. -- **Lessons learned**: - - A shared scheduler loop should expose observability through `SchedulerBackend` so one implementation covers both execution topologies. - - Post-drain/post-prune publication captures both newly waiting work and capacity returned by cancellation before the next admission decision. - - The existing HTTP benchmark plus raw metric sampling covers the live gauge contract. +| `num_running_reqs` | Active decode requests, requests still prefilling and requests owned by the in-flight prefill | +| `num_waiting_reqs` | Pending requests, including deferred work and newly drained submissions | +| `kv_used_blocks` | Backend request-page capacity minus available request pages | +| `kv_total_blocks` | Backend request-page capacity, excluding the CUDA Graph padding page | + +TP counts logical requests and pages once, rather than summing replicas. +Including in-flight prefill also prevents the driver's shutdown check from +mistaking an unfinished asynchronous chunk for an idle scheduler. + +A change in metrics can commit an empty step. The bridge forwards its stats, +so cancelling the last request can reset HTTP gauges even though cancellation +retires the ledger account without a terminal update. Unchanged idle +iterations send no step. + +## Waiting and failure + +The shared driver polls when idle and can occupy a CPU thread even without +requests. With overlap enabled, Qwen3.5 polls the prefill event while decode +remains active. If no decoder remains, it waits for that event inside the +step and then finishes the chunk. + +On an engine-fatal error, the scheduler drains asynchronous work before +clearing its request owners. The shared driver fails open ledger accounts +and exits. The bridge observes the exit and reports the dead engine to the +frontend, which shuts down the HTTP service. The final metrics retain total +capacity and display zero running, waiting and used KV for the dead engine. +That zero is not evidence that TP pages were returned: a poisoned executor +is not retried through healthy `DropRequest` cleanup, and its remaining +resources are released during backend teardown. diff --git a/docs/models/qwen35/model-crate.md b/docs/models/qwen35/model-crate.md index 6d8d4d6a7..28385709d 100644 --- a/docs/models/qwen35/model-crate.md +++ b/docs/models/qwen35/model-crate.md @@ -1,114 +1,66 @@ -# Qwen3.5-4B Model Crate +# Qwen3.5 Model Crate -**Created**: 2026-05-05 -**TL;DR**: `pegainfer-qwen35` now owns Qwen3.5 config, weights, prefill/decode/unified forward, recurrent state, scheduler, recurrent op wrappers, scheduler integration tests, and Qwen3.5 op benches. The whole crate is behind the `qwen35` feature (`--features qwen35` on `pegainfer-server`) because its GDR prefill kernels are Triton AOT-generated — this keeps the default Qwen3 build Python-free. Root `pegainfer` loads Qwen3.5 through `pegainfer_qwen35::start_engine(...)` / generic `EngineHandle`; root no longer exposes `pegainfer::model::Qwen35Model` or `pegainfer::scheduler_qwen35`. The original exact-text e2e/regen tests described in this migration record were later retired by the HF logits gate in `docs/models/qwen35/accuracy.md`. -**Last touched**: 2026-07 +> **TL;DR:** `pegainfer-qwen35` owns the hybrid model runtime and scheduler; the server launches it through the shared step-contract `Engine`. Qwen3.8 uses the same model line. +> +> **Last touched:** 2026-10 -## Feature gate (2026-06) +## Ownership and entry points -Qwen3.5 is the only model line whose kernels need Python at build time (Triton -AOT for the GDR chunkwise prefill). To make the stock build pure Rust + CUDA, -the crate is feature-gated end to end: +The crate owns config and weight loading, recurrent/conv state, paged KV, +prefill/decode/unified execution, scheduler policy, TP workers and model-local +tests. Shared CUDA wrappers and kernel builds remain in `pegainfer-core` and +`pegainfer-kernels`; HTTP and the engine contract belong to +`pegainfer-frontend`. -- `pegainfer-kernels/qwen35` gates the Triton AOT build step and the GDR - chunk FFI declarations — without it, `build.rs` never probes for Python. -- `pegainfer-qwen35` compiles to an empty crate without its `qwen35` - feature (crate-root `#![cfg]`), so `cargo test --workspace --lib` stays - Python-free; its tests/benches carry `required-features` and fail with an - actionable message instead of a link error. -- `pegainfer-server` defaults to `qwen3` only; serve Qwen3.5 with - `cargo run --release --features qwen35 -- --model-path models/Qwen3.5-4B`. +`start_engine*` and `launch_with_options_policy_and_overlap` return +`Result`. `Qwen35Line::launch` returns `LaunchedEngine::Stepped`, +exposing one logical scheduler for either a single GPU or a TP replica. +`EngineInfo` carries the servable context limit and logical KV capacity. -The unused Triton HD256 prefill kernel (replaced by the native paged -`batch_prefill_paged_cuda_hd256`) was deleted in the same change. +The low-level `runtime` module exposes executors and model state for +model-local tests, debugging and benchmarks. `runtime_ops` exposes the +operator wrappers needed by the crate's benches. The server uses the model +line's launch boundary rather than these execution internals. -## Preparation +Qwen3.8 shares the config identity and text geometry used by this line. Its +weight-loading, checkpoint-fixture and serving-template differences are +recorded in [Qwen3.8 support](support-qwen38.md). -- **Read**: - - `docs/index.md` - identified the existing core split, Qwen3 model crate split, and Qwen3.5 accuracy/optimization docs. - - `docs/models/qwen3/model-crate.md` - Qwen3 already owns its scheduler, executor/runtime API, tests, benches, and root-facing `EngineHandle` entry. - - `docs/models/qwen35/accuracy.md` - at the time of this migration, Qwen3.5 e2e tests were regression guards against `test_data/Qwen3.5-4B.json`; current accuracy coverage is the HF logits gate recorded there. - - `docs/models/qwen35/optimization.md` - Qwen3.5 should keep its hybrid linear/full-attention scheduler/state architecture. - - GitHub issue #79 - acceptance criteria require `pegainfer-qwen35`, removal of root `pegainfer::model::Qwen35Model` and `pegainfer::scheduler_qwen35`, generic root `bench_serving`, and CUDA validation. - - `Cargo.toml`, `src/lib.rs`, `src/main.rs`, `src/ops.rs`, `src/scheduler.rs`, `src/model/qwen35.rs`, and `pegainfer-qwen3/src/lib.rs` - mapped the current root Qwen3.5 surface and the Qwen3 crate interface to copy. -- **Relevant history**: - - `docs/models/qwen3/model-crate.md` - root should load model crates through `EngineHandle`; model-owned execution details should move behind crate-local modules. -- **Plan**: - 1. Add `pegainfer-qwen35` to the workspace with dependencies mirroring the Qwen3 crate plus the root dependencies Qwen3.5 currently uses. - 2. Move `src/model/qwen35.rs`, `src/model/qwen35/*`, `src/scheduler_qwen35.rs`, and Qwen3.5 recurrent op wrappers into the new crate, keeping CUDA/Triton kernel sources and FFI in `pegainfer-kernels`. - 3. Rewrite imports so the new crate depends on `pegainfer-core` and `pegainfer-kernels`, not on root `pegainfer`. - 4. Expose `start_engine` and a deliberate `runtime` module from `pegainfer-qwen35`. - 5. Update root `main.rs` and `src/bin/bench_serving.rs` to call `pegainfer_qwen35::start_engine`. - 6. Move Qwen3.5 e2e tests and regen test into the model crate; adjust model/test-data paths after the move. - 7. Remove root Qwen3.5 modules and compatibility exports, then audit root with `rg`. - 8. Verify with `cargo fmt --all --check`, `cargo metadata --no-deps --format-version 1`, and the CUDA-capable build/test commands available on this machine. -- **Risks / open questions**: - - Some root operator tests cover Qwen3.5 recurrent wrappers; they may need to move with the wrappers or be split so root no longer imports model-specific scratch types. - - Accuracy docs reference historical `qwen35_dump_*` and `tools/accuracy/*` files that are not present in the current tree; this migration can document the current test locations but cannot move absent tools. +## Scheduler boundary -## Execution Log +The shared driver drains submissions, calls `Qwen35Scheduler::step`, publishes +metrics and commits ledger updates. `RequestLedger` owns request identity, +admission/terminal updates and emitted-token counts. The Qwen3.5 scheduler +retains admission, planning and backend-state ownership; its existing +backends retain chunked prefill, prefix caching, TP, CUDA Graph and single-GPU +overlap execution. -### Step 1: Add model crate and move Qwen3.5 runtime -- Added `pegainfer-qwen35` to the workspace and root dependencies. -- Moved Qwen3.5-owned runtime files out of root: - - `src/model/qwen35.rs` - - `src/model/qwen35/*` - - `src/scheduler_qwen35.rs` - - `src/ops/recurrent.rs` -- The new crate exposes: - - `start_engine(model_path, EngineLoadOptions, max_batch, max_prefill_tokens) -> Result` - - `runtime::{Qwen35Model, MAX_BATCH}` for model-local tests/debugging - - `runtime_ops` for Qwen3.5-local operator benches. +Single-GPU launch waits for CUDA context binding and thread-local cuBLAS +initialization on the scheduler thread before returning an engine. TP workers +initialize their own CUDA/cuBLAS/NCCL resources before the controller starts +its scheduler, so that scheduler needs no second startup handshake. See +[TP implementation](tp-implementation.md) for the worker protocol and +[scheduler metrics](load-snapshot.md) for publication, cancellation and idle +polling. -### Step 2: Move tests and benches -- Moved root Qwen3.5 tests to the model crate at the time: - - `pegainfer-qwen35/tests/e2e.rs` - - `pegainfer-qwen35/tests/e2e_scheduler.rs` - - `pegainfer-qwen35/tests/regen_test_data.rs` -- The exact-text `e2e.rs` and `regen_test_data.rs` were later removed by the Qwen3.5 HF logits gate work; `e2e_scheduler.rs` remains as request-flow coverage. -- Moved Qwen3.5-specific op benches to `pegainfer-qwen35/benches/qwen35_ops.rs`. -- Moved the `conv1d_prefill_handoff_matches_single_prefill` operator test into `pegainfer-qwen35/src/recurrent.rs`, next to the wrapper it validates. -- Removed Qwen3.5-specific GEMV shapes from the root generic `ops_bench`; the model-specific benches now live with Qwen3.5. +Qwen3.5 suppresses model EOS unless `ignore_eos` is set, and a stop finishes +without a typed cause. The bridge supplies its legacy sentinel so HTTP usage +still counts suppressed EOS; ledger counts cover only emitted tokens. +Request-scoped stop IDs and retained trigger/logprob metadata remain +follow-up work. Qwen3.5 rejects non-null `prompt_logprobs` before backend +allocation and does not provide prompt echo. -### Step 3: Remove root Qwen3.5 compatibility surface -- Removed root exports/modules: - - `pub mod model` - - `pub mod scheduler_qwen35` - - `src/model.rs` - - `src/ffi.rs` - - `src/kv_pool.rs` -- Root `main.rs` now calls `pegainfer_qwen35::start_engine(...)` for Qwen3.5. -- Root `bench_serving` now calls `pegainfer_qwen35::start_engine(...)` and still benchmarks via generic `EngineHandle`. -- The Qwen3.5 engine entry honors a single `EngineLoadOptions.device_ordinals` value and rejects multi-device input, matching the current single-GPU implementation instead of silently ignoring the option. -- `rg` confirms there are no root references to `pegainfer::model::Qwen35Model`, `pegainfer::scheduler_qwen35`, or `src/model/qwen35`. +## Build and validation -### Step 4: Validation -- Passed: - - `cargo metadata --no-deps --format-version 1` - - `cargo fmt --all --check` - - `PEGAINFER_CUDA_SM=120 cargo check --release --workspace --all-targets` - - `PEGAINFER_CUDA_SM=120 cargo clippy --release --workspace --all-targets -- -D warnings` - - `PEGAINFER_CUDA_SM=120 cargo build --release` - - `PEGAINFER_CUDA_SM=120 cargo test --release -p pegainfer-qwen35 recurrent::tests::conv1d_prefill_handoff_matches_single_prefill -- --nocapture` - - `PEGAINFER_CUDA_SM=120 cargo run --release --bin bench_serving -- --model-path $LOCAL_PEGAINFER_DIR/models/Qwen3.5-4B request --prompt-len 1 --output-len 1 --warmup 0 --iters 1` -- Initial Qwen3.5 e2e failure: - - `PEGAINFER_CUDA_SM=120 PEGAINFER_TEST_MODEL_PATH=$LOCAL_PEGAINFER_DIR/models/Qwen3.5-4B cargo test --release -p pegainfer-qwen35 --test e2e -- --nocapture` - - `PEGAINFER_CUDA_SM=120 PEGAINFER_TEST_MODEL_PATH=$LOCAL_PEGAINFER_DIR/models/Qwen3.5-4B cargo test --release -p pegainfer-qwen35 --test e2e_scheduler -- --nocapture` - - Both initially produced all-case gibberish-output mismatches. -- Control run: - - A temporary old-HEAD worktree at `$RESULT_ROOT/pegainfer-head` ran `PEGAINFER_CUDA_SM=120 PEGAINFER_TRITON_PYTHON=$LOCAL_PEGAINFER_DIR/.venv/bin/python PEGAINFER_TEST_MODEL_PATH=$LOCAL_PEGAINFER_DIR/models/Qwen3.5-4B CARGO_TARGET_DIR=$RESULT_ROOT/pegainfer-head-target cargo test --release --test e2e_qwen35 -- --nocapture`. - - Old HEAD failed the same way on all 10 Qwen3.5 cases, so the e2e mismatch predated this crate split. -- Follow-up fix: - - `docs/lessons/exact-match-gate-thread-cublas.md` identified the first gibberish commit as `6a5b826`, fixed Qwen3.5 scheduler thread CUDA/cuBLAS binding, kept greedy sampling on FlashInfer top1, and refreshed the exact Qwen3.5 golden for the default engine shape. - - After that fix, both Qwen3.5 e2e commands above pass. +The `qwen35` feature enables the crate and its Triton AOT GDR prefill kernels. +`PEGAINFER_TRITON_PYTHON` selects the build-time Python/Triton interpreter. +The default Qwen3 build does not enable this kernel path; the Qwen3.5 tests +and benches declare their required feature explicitly. -## Debrief - -- **Outcome**: Qwen3.5 is now an independent model crate with the same root-facing engine style as Qwen3-4B. Root retains model detection/frontend/bench orchestration, but not Qwen3.5 model internals. The follow-up e2e corruption fix restored the then-current exact-text e2e and scheduler e2e; the exact-text gate was later retired in favor of the HF logits gate. -- **Pitfalls encountered**: - - The first e2e run used a relative `PEGAINFER_TEST_MODEL_PATH`; package tests execute with a crate-oriented working directory, so absolute model paths are safer for crate-local tests. - - Qwen3.5 e2e initially looked like a crate-split regression, but git history showed the corruption started earlier when cuBLAS handles became thread-local without equivalent Qwen3.5 scheduler thread binding. - - Moving recurrent wrappers out of root exposed stale root compatibility re-exports (`src/ffi.rs`, `src/kv_pool.rs`, and root Qwen3.5 ops bench shapes), which were removed. -- **Lessons learned**: - - Model-local benches need a deliberate public surface. `runtime_ops` is intentionally narrow and only exposes the Qwen3.5 operator wrappers needed by Qwen3.5 benches. - - Qwen3.5 test docs should use absolute `PEGAINFER_TEST_MODEL_PATH` examples when run from the workspace, because package test working directories can make relative paths misleading. +Use the model crate's scheduler, chunked-prefill, sampling, prefix-cache and +TP2 serving integration tests for request-flow changes. The HF logits gate +in [accuracy](accuracy.md) is the numerical oracle; an E2E test completing or +producing plausible text is not an accuracy measurement. The shared model +fixture uses `PEGAINFER_TEST_MODEL_PATH`; use an absolute path to avoid +package working-directory ambiguity. TP2 gates require two devices and must +run serially. diff --git a/docs/models/qwen35/prefix-cache.md b/docs/models/qwen35/prefix-cache.md index 928ec969b..bab5aa672 100644 --- a/docs/models/qwen35/prefix-cache.md +++ b/docs/models/qwen35/prefix-cache.md @@ -2,7 +2,7 @@ > **TL;DR:** Qwen3.5-4B uses a GPU-only, content-hashed joint prefix cache: full-attention KV is reusable only with a matching complete recurrent/conv snapshot at the same 256-token boundary, otherwise the request is cold. KV stays in kvbm's reclaimable inactive pool, while `SnapshotCache` owns snapshot lookup, pinning, LRU, and publication. TP1/TP2 correctness and serving tests pass, and a 4,160-token shared prefix cuts warm TTFT by 94.3% (TP1) / 95.1% (TP2). > -> **Last touched:** 2026-09 +> **Last touched:** 2026-10 ## Preparation @@ -70,7 +70,7 @@ Joint lookup adds these exact-boundary operations to the shared cache: 5. Advance the new request's KV position as part of the same attachment. 6. Keep at least one prompt token uncached so prefill can produce the first generated token. -`TokenEvent::Scheduled.cached_tokens` reports the selected joint boundary. A KV-only tail is never reported as a hit. +The scheduler records the selected joint boundary with `RequestLedger::set_cached_tokens`; the committed `RequestUpdate.cached_tokens` reports it to the frontend. A KV-only tail is never reported as a hit. ## Valid cache hit diff --git a/docs/models/qwen35/support-qwen38.md b/docs/models/qwen35/support-qwen38.md index cab433811..02256dc2b 100644 --- a/docs/models/qwen35/support-qwen38.md +++ b/docs/models/qwen35/support-qwen38.md @@ -52,9 +52,11 @@ like a numerics break. It is not, on three independent grounds: release matching the checkpoint's declared `5.8.0.dev0`) and 5.17.0. The only consumer in the whole wheel is `qwen4_exp`, a different architecture. 2. **Reading it would change nothing.** swish ≡ silu, and the gated RMSNorm - silu is exactly what PegaInfer already applies - (`pegainfer-kernels/csrc/shared/norm.cu:58`), matching HF's - `Qwen3_5RMSNormGated`. + SiLU branch of `rms_norm_gated_kernel` + (`pegainfer-kernels/csrc/shared/norm.cu`) is exactly what PegaInfer already + applies, matching HF's `Qwen3_5RMSNormGated`. Scoped to `qwen3_5`: that + kernel now also carries a sigmoid entry point for the `qwen4_exp` + architecture of item 1, and Qwen3.8-27B never reaches it. 3. **The weights say so.** Swap `torch.sigmoid` for `x·sigmoid(x)` at the module's single call site and teacher-force both hypotheses: Qwen3.8-27B NLL **1.82** (sigmoid, coherent prose) vs 14.27 (silu, token soup); a diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index bb0dc66de..f05957060 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -1,6 +1,6 @@ # Qwen3.5 Tensor Parallelism Design -> **TL;DR:** Qwen3.5 tensor parallelism reuses Qwen3's controller/worker TP runtime and stays degree-parametric. Phases 1, 2a, 2b, and 2c are implemented (see `tp-implementation.md` for the landing record, including the rebase onto #870): eager dense TP, TP mixed-step unified execution, sharded linear-attention/GDR state, and decode CUDA Graphs under TP gated on the compiled decode GQA group (4B/9B TP2 capture; 27B group-6 stays eager). Remaining design work: group-6 batch-decode kernels so 27B TP2 can capture, and TP perf gates. +> **TL;DR:** Qwen3.5 tensor parallelism reuses Qwen3's controller/worker TP runtime and stays degree-parametric. Phases 1, 2a, 2b, and 2c are implemented (see `tp-implementation.md` for the current runtime): eager dense TP, TP mixed-step unified execution, sharded linear-attention/GDR state, and decode CUDA Graphs under TP gated on the compiled decode GQA group (4B/9B TP2 capture; 27B group-6 stays eager). Remaining design work: group-6 batch-decode kernels so 27B TP2 can capture, and TP perf gates. > > **Last touched:** 2026-09 @@ -164,7 +164,7 @@ Execution: State ownership: -- scheduler owns request admission, request identity, logical page allocation, streaming handles, sampling params, generation counters, and finish bookkeeping +- scheduler owns request admission, backend request identity, logical page allocation, and sampling params; `RequestLedger` owns output, generation counts, and terminal bookkeeping - rank workers own rank-local model shards, rank-local physical KV buffers, rank-local decode buffers, and rank-local recurrent/conv state - rank 0 is not special for state mutation; it follows the same worker command protocol as other ranks - non-primary workers may return acknowledgement or step failure only, while the primary worker returns artifacts for scheduler-side result resolution @@ -204,24 +204,24 @@ State and artifacts are keyed by `RequestId`: - finish, explicit drop, cancellation, and client disconnect broadcast the same `DropRequest(RequestId)` lifecycle command to every rank; - each worker returns `DropAck { existed: bool }`, where `existed` reports whether that rank actually found and removed the request state; - every controller-side drop carries a `DropExpectation::{MustBeAbsent, MustExist}` derived from scheduler-owned lifecycle state, not from a separate `worker_state_materialized` flag; -- cancellation before the first successful prefill dispatch requires `MustBeAbsent` and an exact-rank all-false `DropAck` set; partially prefetched, active, disconnected, and completion-candidate requests require `MustExist` and an exact-rank all-true set; +- pending requests have no worker state and retire without a drop; admitted prefill requests, including those at cursor zero, and active or completion-candidate requests require `MustExist` and an exact-rank all-true `DropAck` set; - uniformity alone is insufficient. All-false for `MustExist`, all-true for `MustBeAbsent`, or mixed true/false values prove lifecycle divergence and poison the whole TP executor even if the drop leaves every rank absent. Dispatch failure, a missing acknowledgement, or a rank-local drop failure is likewise fatal; the scheduler must not warn and continue promotion or serving. -The scheduler allocates a TP `RequestId` before the first prefill command materializes worker state, so cancellation in that interval legitimately requires `MustBeAbsent`; every successfully prefetched or active request requires `MustExist`. If execution mutates state and then fails, the executor is already poisoned and no healthy-path drop is attempted. Worker replies prove cross-rank existence, while scheduler lifecycle determines which value is valid. +Admission calls `begin_request`, which materializes rank-local state even on a cold prefix miss. Cancellation before the first prefill dispatch therefore still requires `MustExist`. The lower-level `MustBeAbsent` expectation applies only when worker state has not been created. If execution mutates state and then fails, the executor is already poisoned and no healthy-path drop is attempted. Worker replies prove cross-rank existence, while scheduler lifecycle determines which value is valid. -Successful request completion has a fail-closed commit boundary. After computing a terminal artifact, the scheduler keeps the request unresolved in a local completion candidate and withholds its user-visible terminal events. An EOS candidate buffers only `Finished`; a length-limited candidate buffers the final `Token` followed by `Finished`; completion on the first prefill token, including `max_tokens <= 1`, follows the same rule. A completion candidate is materialized, so the scheduler first requires a valid `MustExist` all-rank all-true `DropAck` contract, then removes the logical request from scheduler state, and only then publishes the buffered events in order. A client that observes `Finished` can therefore rely on consistent rank-local cleanup acknowledgement for that request. +Successful request completion has a fail-closed commit boundary. After computing a terminal artifact, the scheduler requires a valid `MustExist` all-rank all-true `DropAck` before recording completion in `RequestLedger`. EOS remains suppressed; a length-limited completion records its final token before `finish`. First-prefill completion follows the same rule. The shared driver commits the update at the end of the step, so a client observing `Terminal::Finished` can rely on consistent rank-local cleanup acknowledgement. -If completion drop fails, the scheduler publishes neither the buffered final `Token` nor `Finished`, poisons the replica, and keeps the candidate unresolved for the complete terminal error fan-out. The client may already have received earlier streamed tokens, so this boundary makes only successful termination atomic; it does not make the complete streamed response transactional. The completion candidate is a scheduler-local prepare/commit abstraction, not a distributed `ValidateUnified -> ExecuteUnified` protocol. A process-level fail-stop may still prevent both the acknowledgement and the terminal `Error` from being delivered. +If completion drop fails, the scheduler records neither the final token nor successful completion, poisons the replica, and leaves the ledger account unresolved for the driver's fatal error fan-out. The client may already have received earlier streamed tokens, so this boundary makes only successful termination atomic; it does not make the complete streamed response transactional or introduce a distributed `ValidateUnified -> ExecuteUnified` protocol. A process-level fail-stop may still prevent both the acknowledgement and `Terminal::Failed` from being delivered. Controller response collection validates an exact rank set, not only a message count. Every response rank must be in `0..world_size` and appear exactly once. Ping requires `Ack` from every rank; drop requires a `DropAck` from every rank whose existence values match the controller's `DropExpectation`; prefill, decode, and unified execution require exactly one matching typed result from rank 0 and `Ack` from every non-primary rank. Missing, duplicate, out-of-range, wrong-variant, mixed drop-existence, or uniformly unexpected drop-existence responses are protocol failures. If workers may already have mutated state, any such response-set failure poisons the complete TP replica. -Cancellation cleanup is a scheduler-tick boundary, not a late planning cleanup. At the start of each tick the scheduler first merges deferred work with newly received submissions, then prunes every already-closed sink from active, prefilling, and pending work. Healthy active/prefilling TP removals complete their all-rank `DropRequest` before the scheduler publishes load or computes admission capacity; pending requests have no worker state and are discarded directly. The required order is `drain -> prune -> publish load -> admission -> plan`. This keeps cancelled requests out of running/waiting metrics, decode slots, future-KV reservation, and the current tick's prefill budget. A cancellation racing after the prune check is still retired by the existing token-send failure path. +Cancellation cleanup is a scheduler-tick boundary, not a late planning cleanup. The shared driver drains submissions, then the scheduler prunes aborted active, prefilling, and pending requests. Healthy active/prefilling TP removals complete their all-rank `DropRequest` before ledger retirement and capacity admission; pending requests have no worker state and retire directly. The order is `drain -> prune -> admission -> plan/execute -> publish metrics/commit`. This keeps pruned requests out of decode slots, future-KV reservation, and the current step's prefill budget. Cancellation racing after pruning is checked when processing results or in the next step. Artifact alignment follows the ordered plan rather than returned row position. Prefill expects artifacts only for `finish_prefill == true` IDs and preserves explicit absence for non-final rows; decode expects one artifact for every row. Unknown, duplicate, or missing IDs are fatal after execution. This TP adapter contract does not require the TP1/single-GPU logits-and-sampling path to adopt sparse artifacts. P2a separates recoverable plan rejection from fatal replica failure. Controller-provable structural errors reject before dispatch without poisoning. Worker-local existence, phase, materialization, or capacity mismatches are not protected by an all-rank validation barrier and are therefore replica-fatal, as is any CUDA/NCCL, response-set, artifact, or lifecycle failure after execution is released. -When fatal failure returns control, the scheduler closes and drains submissions, emits exactly one terminal `Error` for every unresolved stable or tick-local request, publishes an idle load snapshot, exits, and begins whole-executor teardown. Exactly-once fan-out comes from consuming mutually exclusive request owners; pre-admission requests do not yet have a TP `RequestId`, so that ID is not a global deduplication key. The scheduler does not retry per-request drops or claim rank-local cleanup after poison. Collective or teardown timeout remains process-level fail-stop, where client fan-out cannot be guaranteed; rollback, communicator recovery, and in-place restart remain out of scope. +When fatal failure returns control, the scheduler drains asynchronous work and clears its owned requests. The shared driver publishes zero running, waiting and used KV metrics, fails every unresolved ledger account, and exits. Submissions not yet registered in the ledger receive `Terminal::Failed` from `RequestEnvelope` when the submit channel is dropped. Frontend identity covers pre-admission and tick-local work independently of the TP `RequestId`. The scheduler does not retry per-request drops or claim rank-local cleanup after poison. Collective or teardown timeout remains process-level fail-stop, where client fan-out cannot be guaranteed; rollback, communicator recovery, and in-place restart remain out of scope. The worker's internal request table may remove and reinsert entries as an implementation detail, but that is not CUDA Graph slot semantics and must not require state copying between request identities. diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index 48e410b23..d83106757 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -1,689 +1,166 @@ -# Qwen3.5 TP Implementation Record +# Qwen3.5 TP Implementation -> **TL;DR:** Qwen3.5 TP is complete through Phase 2b plus batched eager TP decode and P2c CUDA Graph under TP: TP2 supports start-gated eager unified prefill+decode with strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load CUDA ordinal validation (#870); linear-attention/GDR state is sharded per rank (27B TP2 fits on 2×48 GB); TP decode rows run as one batched forward per step; and 4B/9B TP2 decode replays pre-captured CUDA Graphs (27B group-6 stays eager by gate). Remaining TP work: group-6 batch-decode kernels for 27B graphs, perf gates. +> **TL;DR:** Qwen3.5 TP uses rank-local model state behind one step-contract scheduler, with start-gated mixed execution, ID-aligned results and acknowledged retirement. Decode is batched eager or pre-captured CUDA Graph replay when the decode GQA group is compiled. Remaining TP work: group-6 batch-decode kernels for 27B graphs, perf gates. > -> **Last touched:** 2026-09 - -## Scope - -This is the implementation record for Qwen3.5 tensor parallelism. The stable architecture contract lives in `docs/models/qwen35/tp-design.md`; this file records what actually landed, what was verified, and what should carry into later phases. - -Keep Phase 1 and Phase 2 in this file until the Phase 2 implementation becomes large enough to split. The same state ownership risks continue across phases, so keeping the history together is useful: Phase 1 proves dense TP and worker-owned request state, while Phase 2 builds on that boundary for mixed prefill/decode and sharded recurrent state. - -Out of scope for this file: - -- local machine paths, NCCL symlink details, and temporary environment setup -- raw command transcripts unless they are part of retained evidence -- benchmark/performance claims -- prompt echo support; Qwen3.5 accepts completion-only requests - -## Phase 1 Outcome - -Phase 1 is complete as a correctness/runtime milestone. - -Implemented: - -- TP config validation for rank/world size, dense divisibility, and `TP > 1 && CUDA Graph` fail-closed startup. -- Dense TP weight loading for full-attention projections, full-attention KV heads, and MLP projections. -- Rank-local worker executor with worker-owned model shards, KV state, recurrent/conv state, CUDA context, cuBLAS, and NCCL comms. -- Eager TP prefill, chunked prefill, and eager decode. -- Scheduler TP backend that routes chunked prefill and eager decode through TP workers while keeping logical request/page accounting in the scheduler. -- Public multi-device Qwen3.5 engine path and server launch path for `tp_size > 1` with CUDA Graph disabled. -- Real HTTP TP2 serving smoke through the vLLM/OpenAI-compatible frontend. - -Not implemented in Phase 1: - -- TP CUDA Graph capture/replay. -- TP `RunUnifiedStep` mixed prefill+decode execution. -- Sharded linear-attention/GDR weights, kernels, conv state, or recurrent state. -- Vocab-parallel embedding or `lm_head`. -- Prefix-cache or recurrent-state snapshot support. -- Performance claims. - -## Post-Phase 1 Follow-up: RequestKv and Joint Prefix Cache - -This follow-up unifies the TP and single-GPU request lifecycle around `RequestKv` and adds joint full-attention KV plus recurrent/conv prefix reuse. - -1. **Unified KV lifecycle** - - The controller uses one `KvCacheManager` for prefill/decode scheduling and commit. - - Immutable `KvView`s carry the logical page ids to every worker; each rank writes its local KV shard into its own `KvBuffer`. -2. **TP capacity and layout** - - Startup validates identical KV geometry and snapshot-slot counts across ranks. - - The logical pool is capped by the smallest rank-local physical capacity. -3. **Joint recurrent snapshots** - - Key, pin, and LRU metadata is centralized; recurrent/conv tensors remain rank-local. - - Publication reserves one common slot, saves it on every rank, verifies the committed boundary, and only then publishes the key. - - Restore follows the same all-rank rule and reports a hit only after every rank confirms the selected boundary. -4. **Opt-in budget** - - `--qwen35-prefix-cache-mib` reserves the snapshot budget independently on each rank. - - `0` keeps cold serving. - -## Important Fixes - -### Gated q projection layout - -The major numeric blocker was the full-attention gated `q_proj` TP shard layout. - -The wrong assumption was that `q_proj.weight` rows were physically arranged as: - -```text -[all q rows][all gate rows] -``` - -The actual Qwen3.5 kernel contract is per-head interleaved: - -```text -[head0 q][head0 gate][head1 q][head1 gate]... -``` - -For TP2, the fixed loader preserves contiguous head-interleaved ranges: - -- rank 0 loads rows `0..4096` -- rank 1 loads rows `4096..8192` - -The old loader gathered local q rows and local gate rows separately, then rebuilt a `[q][gate]` fused matrix. That corrupted the first full-attention contribution and failed the TP2 HF gate from prefill position `0`. - -### Per-device Triton AOT handles - -Real TP2 prefill exposed that Qwen3.5 GDR Triton AOT C stubs could not cache `CUmodule` / `CUfunction` in process-global state. With two CUDA devices, the rank that loaded a GDR kernel first could leave the other rank with an invalid function handle. - -The generated stubs now cache module/function handles per CUDA device ordinal. This is an implementation constraint worth remembering for future multi-GPU users of generated Triton C stubs. - -Follow-up review tightened this path: the generated stubs now fail closed before indexing the fixed per-device handle tables if `cuCtxGetDevice` returns an ordinal outside the table size. This preserves the Phase 1 static-table implementation while avoiding out-of-bounds writes on high CUDA ordinals. - -### Worker-local NCCL setup - -NCCL comms are initialized inside rank worker threads after each worker binds its CUDA context and initializes thread-local cuBLAS. Creating comms on the controller thread and moving them into workers led to invalid-handle symptoms and hangs. - -This matches the design contract: TP workers own rank-local CUDA/NCCL execution resources. - -### Current-main API compatibility - -Rebasing Phase 1 onto current `main` required preserving the TP execution boundary while adopting newer shared contracts: - -- Hybrid batch decode now builds `Vec<&mut RecurrentState>` from graph-owned slots before entering the common linear-attention helper. This keeps request state in place while satisfying the helper's mutable-reference slice contract. -- `pegainfer_sample::select_batch` now requires request-local sampling steps. Phase 1 TP still samples one row at a time and has no request-local sampling counter, so it passes step `0` and retains its existing per-row `sample_seed` offset. Do not substitute batch row indices for request-local steps: that would make seeded output depend on batch composition. -- Qwen3.5 launch and tests use the current `EngineLoadOptions` surface; the removed `enable_prefill_profile` field is no longer supplied. -- TP scheduler tests explicitly set the newer `GenerateRequest::data_parallel_rank` field to `None` because Phase 1 is TP-only, not DP. -- Synthetic TP config/loader fixtures include `tie_word_embeddings`, matching the current `Config35` contract without changing production config loading. -- TP2 short/long HF gates use `Golden::load_for(model_path, long)` and pass the complete `Golden` to metadata validation, matching the model-selected fixture flow used by TP1. - -These are compatibility changes, not extensions of Phase 1 scope. In particular, full seeded-sampling replay under TP should add a request-local completion counter rather than overloading batch position. - -## Validation Evidence - -Phase 1 acceptance coverage: - -- TP2 short HF logits gate passes: - - sequential eager: `108` positions, mean `0.0258`, p99 `0.0801`, max `0.1298` - - batched eager: `72` positions, mean `0.0257`, p99 `0.0809`, max `0.1298` -- TP2 long HF logits gate passes: - - prompts `4097` and `8192`, sequential eager: `18` positions, mean `0.0232`, p99 `0.0792`, max `0.1035` -- TP2 scheduler e2e passes and covers: - - context-window rejection - - greedy/logprobs paths - - sequential requests - - repeated request reuse - - concurrent mixed greedy/sampling requests - - consumer drop - - post-drop scheduler health -- TP2 HTTP serving smoke passes through `pegainfer_vllm_frontend::serve`: - - `/v1/models` - - non-streaming `/v1/completions` - - streaming `/v1/completions` - - concurrent completions - - finite logprobs - - chunked prefill forced with `max_prefill_tokens=1` - - `TP2 + CUDA Graph` fail-closed startup -- TP1 regression gates pass after the TP2 additions: - - TP1 short/long HF logits gates - - TP1 scheduler e2e -- Current-main rebase verification passes: - - formatting check - - Qwen3.5 release compilation for all test targets - - `pegainfer-server` release compilation with only the `qwen35` model feature - -Known validation constraints: - -- TP2 tests remain ignored by default because they require two CUDA devices, NCCL, and real Qwen3.5 weights. -- Long TP2 HF replay is GPU-memory-sensitive; choose a sufficiently free device pair. -- Qwen3.5 HF golden integration tests should run serially on memory-constrained hosts to avoid unrelated KV-capacity failures from concurrent model loads. - -Stable test knobs: - -- `PEGAINFER_TEST_MODEL_PATH`: real Qwen3.5 weights path for HF, scheduler, and serving tests. -- `PEGAINFER_TEST_TP_DEVICES`: comma-separated TP2 CUDA ordinals. Defaults to `0,1`; examples: `1,2`, `2,3`. TP2 tests require exactly two distinct ordinals. -- `PEGAINFER_TEST_FRONTEND_MODEL_PATH`: optional tokenizer/config metadata path for HTTP serving tests. Defaults to `PEGAINFER_TEST_MODEL_PATH` when unset. - -## Phase 2 Progress - -Phase 2 is locked in `docs/models/qwen35/tp-design.md` as two separate implementation series: P2a is eager mixed unified execution on the replicated Phase 1 GDR path; P2b shards the head-indexed linear-attention/GDR weight and state surface. P2a protocol/lifecycle gates are complete, and P2b's core sharding has landed on top of them without weakening the P2A lifecycle and ID contracts (see below). Batched eager TP decode (#1004) and P2c TP decode CUDA Graphs (#1005, compiled decode GQA groups only — 4B/9B TP2 capture, 27B group-6 stays eager) have landed on top of P2b (see below); the remaining Phase 2 work is group-6 batch-decode kernels so 27B TP2 can capture, plus TP perf gates. - -### P2a: TP mixed-step unified execution - -P2A implements eager `RunUnifiedStep` under TP while retaining Phase 1's replicated linear-attention/GDR weights, kernels, conv state, recurrent state, and scratch shapes. - -#### Step 1: lifecycle cleanup gates - -Completed as test and fault-observation infrastructure before changing production lifecycle semantics. - -Why it exists: - -- A successful controller-side `DropRequest` return did not directly prove that every rank removed the same `RequestId` or released its KV/recurrent/conv ownership. -- Later cancellation, partial-dispatch, drop-acknowledgement, and unified-step work needs direct evidence of rank-local state before and after each transition. -- Cleanup must prove both capacity recovery and fresh numeric state; scheduler bookkeeping alone cannot detect a stale rank-local recurrent or KV allocation. - -Implemented in `tp_executor.rs` under `#[cfg(test)]` only: - -- `WorkerStateSnapshot { rank, request_count, requests }`, where each request entry carries its `RequestId` and `Prefilling`/`Decoding` phase. -- A healthy snapshot API that requires an unpoisoned executor and an unchecked test-only API that bypasses only the controller poison guard. The latter is reserved for later synthetic failure tests while every worker channel remains connected. -- Exact-rank collection that rejects duplicate or out-of-range ranks, payload/response rank disagreement, missing responses, wrong reply variants, and inconsistent request counts. Valid snapshots are returned in rank order. -- An ignored TP2 capacity gate that fills configured `max_batch=2`, observes every ID in `Decoding` on both ranks, drops every ID, requires zero requests on every rank, refills the complete capacity, executes decode, and proves the second cleanup is also empty. -- An ignored TP2 numeric gate that records a prompt's deterministic first token and five requested logprobs, drops the request, verifies every rank is empty, re-admits the same prompt under a new `RequestId`, and requires an exact artifact match. - -Verification: - -- Qwen3.5 release all-target check and clippy with `-D warnings` pass. -- The regular library suite passes: `71 passed`, `0 failed`, `8 ignored`. -- Both new TP2 GPU gates pass independently. -- Formatting and `git diff --check` pass. - -This step deliberately adds no production command, scheduler state transition, drop semantics, or serving claim. It supplies the observability needed to prove the later P2A changes rather than trusting controller-visible acknowledgements alone. - -#### Step 2: cancellation pruning before admission - -Completed as the first production scheduler lifecycle change in P2A. - -The shared TP1/TP scheduler tick now follows the fixed order `drain -> prune -> publish load -> admission -> plan`. It merges deferred and newly submitted work, removes requests whose `TokenSink` is already closed, publishes the resulting state, and only then computes slot/KV budgets and admits work. The idle wakeup path repeats drain/prune/publication after `blocking_recv()` because the waking request can close before admission and other submissions can race with the receive. - -Change to the existing scheduler skeleton: - -- Previously, the loop published the state left by the prior tick before draining `submit_rx`; `num_waiting_reqs` therefore represented only `deferred.len()`. It then drained submissions and entered admission. -- The loop now first takes `deferred`, drains every currently available submission into the same `pending` vector, and calls one backend-neutral `prune_closed_requests` helper across pending, active, and prefilling ownership. -- `publish_load` remains the same internal watch publication surface and `LoadSnapshot` retains the same public fields and types. A small `logical_load_counts` helper makes the post-prune running/waiting calculation directly testable; waiting is now `pending.len()` at the admission boundary. -- The empty-loop blocking skeleton is retained, but wakeup now has a second drain/prune/publish boundary before admission. If every received request is already closed, the loop returns to blocking without entering admission or planning. -- The admission and plan implementations are otherwise unchanged. This commit changes which settled state they consume, not their policies or the public `EngineHandle::load_watch()` API. - -Effect on the pre-existing TP1 (`SchedulerBackend::Single`) path: - -- A closed pending request is removed before `alloc_prefill_state`, so TP1 no longer allocates its KV/recurrent state and then discovers the disconnected consumer during token delivery. -- A closed prefilling request is removed before plan construction. Consuming its existing `PrefillBackendState::Single` drops the owned `KvState` and `RecurrentState` through their existing RAII path; no TP1 release routine or allocator rule was added. -- A closed active request is retired before the next model step through the existing `compact_single_slot` path. The same `swap_remove` and graph-slot copy used by EOS, length, and token-send failure are reused; only the retirement timing moves earlier, so the cancelled row no longer executes and samples one unnecessary decode step. -- TP1 page availability, slot budgets, and plan construction now observe that cleanup in the same tick. Consequently, the executed batch width and load samples can differ under cancellation, but KV accounting formulas, slot-compaction mechanics, admission policy, plan builders, kernels, and sampling implementation are unchanged. - -Cleanup deliberately reuses existing backend ownership paths: - -- pending cancellation uses stable `retain`, preserving FIFO order among live requests; -- active cancellation uses normal request retirement, including TP drop and single-GPU slot compaction; -- prefilling cancellation removes the entry and drops its backend state through the existing adapter. - -This ordering makes cancellation visible to admission in the same tick. A closed resident is absent from the post-prune running/KV state, a live replacement is present in waiting, and the capacity released by the resident can admit that replacement immediately. - -Verification: - -- Three focused CPU tests pass for pending FIFO pruning, post-prune logical load, and same-tick capacity reuse. -- The ignored real TP1 `max_batch=1` cancellation/replacement gate passes on an RTX 3090. It observes the post-prune `running=0, waiting=1` boundary, completes the replacement, and returns running, waiting, and KV usage to zero. -- The existing real TP1 scheduler integration E2E passes (`1 passed`, 29.06 seconds). -- Release all-target clippy with `-D warnings` passes. -- The regular library suite passes: `74 passed`, `0 failed`, `9 ignored`. - -This step does not strengthen the TP drop reply protocol or add scheduler-wide fail-closed propagation. It uses the Phase 1 cleanup calls as they exist; exact-rank `DropExpectation` acknowledgement belongs to step 3, and mandatory replica-fatal propagation belongs to step 4. - -#### Step 3: TP worker protocol hardening - -Completed as the distributed command/reply foundation for later scheduler fail-closed handling and unified execution. - -State-mutating prefill, decode, drop, and future unified envelopes now share one `Pending -> Execute | Cancel` start gate. The controller enqueues an envelope for every rank before resolving `Execute`; if any enqueue fails, it resolves `Cancel` for the delivered prefix and poisons the executor. Workers wait on the gate before request-state mutation, kernel launch, or NCCL entry. Ping and test-only snapshots remain ungated because they do not mutate request state or enter collectives. - -Response handling now separates timed transport collection from pure response-set validation. Ping, prefill, decode, and drop all require exactly one in-range response from every rank: - -- ping requires `Ack` from every rank; -- prefill requires one `Prefill` result from rank 0 and `Ack` from every non-primary rank; -- decode requires one `Decode` result from rank 0 and `Ack` from every non-primary rank; -- drop requires `DropAck { existed }` from every rank and validates the complete existence vector against the scheduler's `DropExpectation`. - -Duplicate, missing, out-of-range, wrong-variant, non-primary typed, or missing-primary responses poison the executor after dispatch. Drop accepts only exact-rank all-false for `MustBeAbsent` and exact-rank all-true for `MustExist`; mixed values and uniformly unexpected values are lifecycle divergence even though every rank is absent after the command. - -The low-level `Qwen35TpExecutor::drop_request` API now requires `DropExpectation` and returns `Result<()>`; `DropExpectation` is re-exported through the model-local `runtime` module for tests and debugging. The server-facing `EngineHandle` API is unchanged. Scheduler callers derive expectations from owned lifecycle state: - -- active and successfully dispatched/final-prefill state use `MustExist`; -- prefilling cancellation at `cursor == 0` uses `MustBeAbsent`; -- prefilling cancellation after progress uses `MustExist`; -- execution failure performs no healthy-path drop because the executor is already poisoned. - -This commit intentionally retains the temporary scheduler adapter that logs a returned healthy-drop failure. Executor poison prevents subsequent normal commands, but step 4 still owns mandatory propagation into one scheduler-wide terminal path, fail-closed completion publication, and exactly-once unresolved-request error fan-out. - -Verification: - -- Release all-target check and clippy with `-D warnings` pass. -- Pure protocol/gate tests cover exact-rank reply matrices, lifecycle-expectation mismatch, controller structural rejection with zero dispatch, all-enqueued `Execute`, prefix-only `Cancel`, and poison preservation. -- The complete regular library suite passes on SM86: `82 passed`, `0 failed`, `12 ignored`. -- Three new real TP2 gates pass: prefix-only dispatch failure leaves every rank empty (214.38 seconds under unrelated four-GPU training contention), lifecycle expectations plus mixed-rank divergence pass (8.05 seconds), and a real receiver disconnect poisons without claiming an unavailable all-rank snapshot (6.96 seconds). -- Existing healthy TP2 prefill/decode/drop and scheduler chunked-prefill/decode smoke tests pass (6.45 and 6.54 seconds). - -Implementation pitfall: the first release-only gate test hung because `start.execute()` was placed inside `debug_assert!`; release builds remove the complete assertion expression, including side effects. The state transition now executes unconditionally and only its boolean result is debug-asserted. Protocol state transitions must never live inside debug-only assertions. - -#### Step 4: fail-closed TP scheduler lifecycle - -Completed as the scheduler recovery and user-visible completion boundary required before TP unified execution. - -The temporary step-3 logging adapter is gone. TP active retirement and prefill-state cleanup now return `Result`, and cancellation pruning, standalone prefill/decode execution, token dispatch, and final-prefill promotion propagate any returned TP lifecycle failure to the scheduler loop. A returned TP prefill/decode or post-execution artifact-alignment error follows the same fatal path. The scheduler does not retry per-request cleanup after poison. - -Successful TP completion is now prepared and committed explicitly: - -- decode EOS buffers only `Finished(Stop)`; -- decode length buffers the final `Token` followed by `Finished(Length)`; -- immediate final-prefill EOS/length uses the same buffering rule, so its first token is not exposed before cleanup; -- the candidate keeps logical request ownership while `DropRequest(MustExist)` runs, and publishes the buffered events only after exact-rank all-true `DropAck` succeeds; -- drop failure discards the buffered success events and transfers the unresolved candidate to terminal error fan-out. - -Fatal errors carry any tick-local request ownership back to the scheduler loop instead of consuming it in a step helper. The single terminal helper closes `submit_rx`, drains every request whose send completed before close, consumes transient candidates/scheduled work plus active, prefilling, current pending, and deferred owners exactly once, attempts one `TokenEvent::Error` per request, publishes a zero-running/zero-waiting/zero-KV load snapshot, and exits. It does not deduplicate by TP or external request ID; exclusivity comes from moving each request out of its prior owner. - -The existing single-GPU boundary remains explicit. TP completion requires cleanup acknowledgement before publication, while the non-TP decode path retains its original `Token`/`Finished` publication before slot retirement. Single-GPU prefill state still drops through its existing RAII ownership, and `EngineHandle` plus `TokenEvent` public types are unchanged. - -Verification: - -- Focused scheduler lifecycle coverage passes: `16 passed`, `0 failed`, `2 ignored`. It covers TP EOS/length and immediate-prefill event buffering, active/prefill drop failure with no leaked success event, remaining scheduled ownership preservation, prune failure propagation, non-TP publication order, and barrier-controlled submit close/drain fan-out across every owner class. -- Release all-target check and clippy with `-D warnings` pass; formatting and `git diff --check` pass. -- The complete regular library suite passes on SM86 outside the GPU-isolated sandbox: `90 passed`, `0 failed`, `12 ignored`. -- The real-weight TP2 scheduler chunked-prefill/decode smoke and the complete TP2 scheduler E2E pass on physical GPUs 1/2. The complete E2E covers context rejection, greedy/logprobs, sequential/repeated/concurrent requests, consumer drop, and post-drop health. -- A same-session TP1 scheduler E2E attempt could not pass model startup because an unrelated training process left insufficient memory for the loader's default 64-slot graph allocation. No training process was stopped. The non-TP ordering CPU gate and complete regular suite are green; the previously established real TP1 E2E evidence remains unchanged. - -Step 4 does not add unified plans, change worker command payloads, alter TP1 sampling, or modify GDR weights/state/kernel shapes. Step 5 owns `RunUnifiedStep`, strict ID-aligned unified artifacts, ordinal validation, and the final TP1/TP2 regression ladder. - -#### Step 5: TP unified execution - -Completed as the final P2A execution-protocol milestone. When active decode and scheduled prefill coexist, TP now uses the same `plan::build_next_plan` decision as the single-GPU scheduler and emits one start-gated `RunUnifiedStep` to every rank. The canonical plan carries ordered prefill/decode items plus separate seeds. The scheduler selects the decode seed first and the prefill seed second to preserve its prior RNG order; workers execute prefill first and decode second on every rank. Rank 0 returns `TpUnifiedResult`, while every non-primary rank must return `Ack`. - -Internal protocol/API changes: - -- Added crate-private `TpUnifiedPlan`, `TpUnifiedResult`, the complete `RunUnifiedStep` payload, and `TpWorkerReply::Unified`. These are not re-exported through `runtime`; the public server `EngineHandle`, `TokenEvent`, and request contract are unchanged. -- Refactored worker prefill/decode into typed inner row operations so standalone and unified commands share the same eager kernels and sampling behavior. Existing public low-level `execute_prefill`/`execute_decode` results retain their shapes. -- Added controller structural validation before enqueue: both halves must be non-empty, row counts must fit scheduler capacity, prefill/decode IDs must be internally unique and mutually disjoint, and prefill chunks must be non-empty. Worker-local request existence, phase, and actual capacity are revalidated after gate release and remain replica-fatal rather than adding a two-phase validation protocol. -- Extended exact-rank reply validation so unified accepts exactly one rank-0 `Unified` result plus non-primary acknowledgements. Returned worker or reply-set failure continues through the step-4 poison and terminal scheduler path. - -Change to the existing scheduler skeleton: - -- Removed the TP-only `build_eager_only_plan` branch. Both backends now use the normal planner, so active decode plus scheduled prefill becomes `ExecutionPlan::Unified` instead of two serialized scheduler ticks. -- Added a TP-only execution adapter that builds the canonical plan, aligns returned artifacts by `RequestId`, processes decode results first, and only then promotes or requeues prefill. Decode completion can therefore release capacity through `DropRequest(MustExist)` before final-prefill promotion. -- Replaced TP prefill token-`0` placeholders and positional decode matching with explicit artifacts. Prefill alignment is a chunk-length `Vec>`: outer `None` means non-final, while `Some { logprob: None, .. }` is a valid final artifact. Decode requires one artifact for every active ID. Shuffled valid results are accepted; unknown, duplicate, non-final, or missing IDs poison the replica after execution. -- The shared `promote_or_requeue` skeleton now accepts a backend-specific artifact wrapper. The single-GPU branch still consumes its original dense sampled tokens/logprobs, uses the same logits path and RNG calls, and retains its step-4 publish/retire semantics. - -Startup now validates TP CUDA ordinals against the generated Triton AOT handle-table contract before model/filesystem/CUDA access. Ordinals must be distinct and below the named table length `16`; the generated wrapper's runtime bounds check remains defense in depth. - -Verification: - -- Release all-target check and clippy with `-D warnings`, formatting, and `git diff --check` pass. -- Complete non-ignored library suite passes on SM86: `96 passed`, `0 failed`, `14 ignored`. -- Pure tests cover ordinal bounds/duplicates, unified structural rejection plus subsequent healthy dispatch, worker-local missing/phase/capacity failures, exact-rank unified replies, shuffled artifact alignment, unknown/duplicate/missing IDs, and final-without-logprobs versus non-final absence. -- The real-weight TP2 executor mixed-step gate passes (`1 passed`, 7.27 seconds), confirms both ranks retain two decoding requests after the combined operation, and returns both ranks to empty after lifecycle cleanup. -- The deterministic real-weight TP2 scheduler case with `max_batch=2` and `max_prefill_tokens=1` passes (`1 passed`, 7.15 seconds). -- All `13` TP2/lifecycle ignored library gates pass together. The fourteenth ignored test is TP1-only and could not load under the unrelated training process: the loader reported `8101 MB` free versus `3248 MB` prefill scratch plus `6288 MB` recurrent state and minimal KV needs. No training process was stopped. -- Complete TP2 scheduler E2E passes (`1 passed`, 35.89 seconds). -- TP2 short/long HF golden gates pass (`2 passed`): short sequential mean `0.0260`, p99 `0.1081`; short batched mean `0.0267`, p99 `0.1167`; long sequential mean `0.0228`, p99 `0.0689`. -- TP2 OpenAI-compatible HTTP serving smoke passes (`1 passed`, 10.41 seconds), including streaming, non-streaming, concurrent completions, logprobs, and TP+CUDA Graph rejection. - -P2A does not shard GDR weights/state/scratch, add a post-GDR all-reduce, enable TP CUDA Graph, or claim a performance improvement. TP1's earlier real-weight evidence remains applicable, but a same-session full TP1 model E2E rerun is still resource-blocked by external GPU occupancy. - -#### Merge-gate test surface - -The pre-merge review reduced P2A coverage to distributed behavior that cannot be -established by restating local helpers. Removed tests enumerated start-gate -booleans, planner choices, CUDA-ordinal bounds, unified-plan shapes, reply -variants, artifact-alignment variants, snapshot-collector variants, and basic -TP2 startup/prefill/chunk/decode smoke already exercised by higher-level gates. - -The retained P2A library surface is seven real TP2 gates: - -- partial dispatch cancels the delivered prefix before rank-local mutation; -- rank-local lifecycle divergence is detected and poisons the replica; -- a disconnected worker receiver fails closed; -- executor and scheduler mixed prefill/decode both complete; -- drop-all restores complete request capacity; -- re-admission under a fresh `RequestId` reproduces the clean first-token artifact. - -Completion cleanup and terminal error fan-out remain covered by the focused -scheduler lifecycle tests. The production-facing regression ladder remains the -TP2 scheduler E2E, short/long HF logits gates, and OpenAI-compatible HTTP -serving lifecycle test. - -On the local 2x RTX 3090 SM86 fixture, the complete explicit ignored run passed -`11/11`: seven library gates, one scheduler E2E, two HF gates, and one HTTP gate. -The regular release library suite passed `94` tests with those seven real TP2 -tests ignored by default. The HF results remained within the established -tolerances: short sequential/batched mean deltas `0.0260`/`0.0267`, and long -sequential mean delta `0.0228`. - -#### Qwen3.5 feature CI gate - -The CUDA workflow now has separate matrix results for Qwen3.5 compile and -Clippy. Both install the repository's CUDA 13.0.2 toolchain, Python 3.10, and -pinned Triton 3.7.1, then build `pegainfer-qwen35` with the actual `qwen35` -feature enabled for all targets on `sm_80`: - -```text -cargo check --release --locked -p pegainfer-qwen35 --features qwen35 --all-targets -cargo clippy --release --locked -p pegainfer-qwen35 --features qwen35 --all-targets -- -D warnings -``` - -This is a compile/lint merge gate, not a substitute for the ignored real-GPU -TP2 tests above. The same target surface passed locally on SM86 with nightly -2026-07-10 and Triton 3.4.0; both commands compiled the Qwen3.5 Triton AOT path. -The GitHub jobs provide the independent Ubuntu, SM80, and pinned Triton 3.7.1 -verification. - -#### Real TP2 multi-turn serving gate - -The Phase 2A production-path gate starts the real OpenAI-compatible Qwen3.5 -server on two RTX 3090 GPUs with TP=2, eager execution, `max_batch=2`, and -`max_prefill_tokens=64`. A pinned upstream Rust `vllm-bench` client runs -dependent `openai-chat` conversations at concurrency 4, so the workload exceeds -resident capacity and forces cleanup followed by fresh admission. - -The primary workload completed 12/12 conversations and 44/44 measured turns -with zero failures. Conversation lengths varied from 2-5 turns, later prompts -carried accumulated history, and every turn returned its configured 24 tokens. -Without restarting the server, a second 4-conversation, 8-turn probe also -completed with zero failures. Graceful shutdown released both TP devices. - -The exact client/server revisions, build and run commands, request counts, -concurrency, prompt/output lengths, per-turn history evidence, pass criteria, -limitations, and raw JSON results are published in -[`docs/benchmarks/qwen35-tp2-phase2a-multiturn.md`](../../benchmarks/qwen35-tp2-phase2a-multiturn.md). - -#### Serving scope: prompt echo is unsupported - -Qwen3.5 serving does not support `echo=true`. Both the legacy and stepped vLLM -bridges submit `echo: false`, so prompt echo has never been part of the current -HTTP serving contract. Direct engine callers can still construct the shared -request type with `echo=true`; the Qwen3.5 scheduler now rejects those requests -immediately after queue drain and cancellation pruning, before load accounting, -capacity admission, backend-state allocation, or TP command dispatch. - -Qwen3.5 no longer emits `TokenEvent::PromptTokens` after prefill. The shared -event variant remains in `pegainfer-frontend` because other model lines still -use it, but it is sealed off from the Qwen3.5 scheduler path. A focused release -test requires a `Rejected` event with the original prompt length and proves the -request is absent from the vector eligible for backend admission. - -Verification on the local SM86 environment passed: - -- focused unsupported-echo release test: `1 passed`; -- complete Qwen3.5 release library suite: `95 passed`, `0 failed`, `7 ignored`; -- all retained real TP2 executor/scheduler gates on two RTX 3090 GPUs: - `7 passed`, `0 failed`; -- release Clippy with `-D warnings` and formatting checks passed. - -#### Model-backed test fixtures are explicit - -Qwen3.5 model-backed tests no longer embed a developer-local absolute weights -path or fall back to a repository-relative model directory. The shared test -fixture resolver requires `PEGAINFER_TEST_MODEL_PATH`, reads and parses its -`config.json`, and confirms the Qwen3.5 model identity before any model or CUDA -initialization. A missing, non-UTF-8, empty, unreadable, malformed, or non-Qwen3.5 -fixture prints one test-specific `SKIP` diagnostic and returns without GPU work. -The optional `PEGAINFER_TEST_FRONTEND_MODEL_PATH` uses the same validation when -set and otherwise reuses the validated engine fixture. - -The resolver source lives under `tests/common` and is also included only for -crate unit-test builds, so the six retained TP executor gates, the mixed-step -scheduler gate, and the integration targets share the same contract without -adding a production API. Static review confirms that no Qwen3.5 test retains a -private model path or implicit `DEFAULT_MODEL_PATH` fallback. - -Verification on the explicit public-compatible SM86 fixture passed: - -- missing fixture, missing `config.json`, and malformed JSON probes emitted the - expected diagnostics without GPU initialization; -- release all-target compile and Clippy with `-D warnings` passed; -- regular release library suite: `95 passed`, `0 failed`, `7 ignored`; -- retained real TP2 library gates: `7 passed`, `0 failed`; -- TP2 scheduler E2E, short/long HF gates, and HTTP serving gate passed; HF - mean/p99 deltas remained `0.0260/0.1081` sequential, `0.0267/0.1167` - batched, and `0.0228/0.0689` long; -- the pinned multi-turn gate rerun completed `12/12` conversations and `44/44` - turns, followed without restart by `4/4` conversations and `8/8` turns, all - with zero failures; graceful shutdown returned both TP GPUs to idle. - -Delivered constraints: - -- Support mixed prefill+decode scheduler steps under TP. -- Preserve deterministic collective ordering across ranks. -- Return mixed prefill/decode artifacts from the primary rank. -- Use the `RequestId`-keyed `UnifiedPlan` and artifact contract in `tp-design.md`; P2a does not introduce CUDA Graph padded-slot or compaction semantics. -- Validate finish/drop/client-disconnect cleanup under mixed-step execution. -- Keep TP CUDA Graph disabled unless a separate graph design is completed. -- Keep the fixed 16-device Triton AOT handle table and add startup-time validation for unsupported logical CUDA ordinals. - -Why this should be separated from GDR sharding: - -- Mixed-step scheduling is an execution-protocol problem. -- Sharded linear-attention/GDR is a model-state-shape problem. -- Combining them would make failures hard to attribute. - -### P2b: sharded linear-attention/GDR state - -Landed as #946 split 1/4 (#1003), after P2a had established the mixed-step and state-lifecycle contract. Each TP rank now owns a rank-local slice of the linear-attention/GDR surface instead of replicating it: - -- `LocalGeometry` computes rank-local linear dims (`local_linear_num_key_heads`, `local_linear_num_value_heads`, `local_linear_v_dim`, `local_linear_qkv_dim`, `local_linear_z_dim`) and fails closed with `ConfigError::TpIndivisible` when `linear_num_key_heads` does not divide by `world_size`; there is no silent replication fallback. Value-head divisibility needs no second guard: `Config35` already validates `linear_num_value_heads % linear_num_key_heads == 0`. -- Weight loading shards the head-indexed tensors: the fused QKV projection and the depthwise conv1d are stitched head-locally per segment (`load_linear_in_proj_qkv_shard` / `load_linear_conv1d_shard`; Q/K segments follow key-head ranges, V/conv follow value-head ranges), z/beta/alpha are row shards, `dt_bias`/`A_log` are 1-D shards (A_log stays f32), and linear `out_proj` is column sharded as a row-parallel `[hidden, local_z]` matrix. -- `RecurrentState` (`[local_value_heads, K, V]` f32), conv state (`[local_qkv x (kernel_dim - 1)]` bf16), and all prefill/decode scratch (`GdrChunkwiseScratch35`, prefill/decode buffers) size themselves from the local geometry; worker capacity math uses the same locals. -- The hidden-residual all-reduce happens once after the local linear-attention `out_proj` (`all_reduce_hidden`), on prefill and decode alike; the column-sharded `out_proj` is what makes that reduction point sufficient. -- Full-attention decode-group supportability uses the config-level GQA group (`Config35::decode_group_is_compiled`): head sharding leaves the q-per-kv group size unchanged, so the predicate is identical on every rank and the reroute adds no collectives. The 27B case leaves `q/kv = 6`, which has no compiled FlashInfer batch-decode kernel, so those layers reroute decode through the eager/paged fallback. -- TP1 contract is unchanged: at `world_size == 1` every local dim equals the global dim, so kernels, buffers, and fixture behavior are byte-identical to pre-P2b. - -Non-negotiable invariant (still held): - -- Never all-reduce GDR recurrent state or conv state. These states are owned by rank-local request state. - -Acceptance on the #1003 branch before its final rebase (27B TP2 on 2x RTX 4090 48GB, sm_89; fixture-pinned 27B revision `fc05daec`): - -- TP2 short HF logits gate passes: - - sequential eager: `108` positions, mean `0.0210`, p99 `0.0749`, max `0.1240` - - batched eager: `72` positions, mean `0.0201`, p99 `0.0749`, max `0.0803`; the batched leg includes drop -> re-prefill slot cycles -- TP2 long HF logits gate passes with prompts `4097` and `8192`: sequential eager, `18` positions, mean `0.0177`, p99 `0.0660` -- TP2 scheduler E2E and TP2 HTTP serving gate pass. -- Peak per-rank HBM is `35,988` / `36,822` MiB of `49,140` MiB: 27B TP2 now fits the 2x48GB pair that Phase-1 replicated state OOMed, and memory fully releases between test processes. - -Not in this step: batching the TP decode loop across rows, TP CUDA Graph capture, and the matched Phase-1-vs-P2b HBM/latency/throughput A/B promised in #1001; no performance claim is made until that rerun lands on the merged stack. - -### Batched eager TP decode (#1004) - -Landed as #946 split 2/4 (#1004). #870's `execute_decode_rows` looped per -request with bs=1 forwards and capacity-1 per-request pointer tables. -`run_decode_batch` now runs one `batch_decode_eager_logits` over every decode -row in command order on each rank, then one batched rank-0 `select_batch` with -per-row fan-out; #870's validation and response contracts are unchanged and -`TpRequestState.linear_pointer_tables` is gone. Seeded rows keep per-row -semantics (`select_batch` isolates each seeded row into its own philox call -keyed on request seed and step), so seeded output is independent of batch -composition; unseeded rows decorrelate through the per-step command seed, as on -the single-GPU batched path. - -GDR pointer tables honour `LinearStatePointerTables`' build-once contract: each -worker allocates one table set at `max_batch` capacity (`with_capacity`) and -refills rows `0..bs` from the live decode rows every step -(`refill_from_recurrent_refs`: H2D copies into the existing allocations). -Refilling every step means swap_remove retirement between steps can never leave -a stale row addressed; allocating once removes the `2 x linear_layers` device -allocations per token per rank that a per-step `from_recurrent_refs` cost -(review follow-up on #1004). The 27B TP2 group-6 full-attention layers still -route decode through the one-token paged plan. - -Acceptance (27B TP2 on 2x RTX 4090 48GB, sm_89; revision `fc05daec`): - -- TP2 short HF logits gate passes: sequential eager `108` positions, mean `0.0210`, p99 `0.0749`, max `0.1240` (identical to the P2b acceptance run above); batched eager `72` positions, mean `0.0199`, p99 `0.0692`, max `0.0887`, the batched leg including drop -> re-prefill slot cycles. -- TP2 long HF logits gate passes with prompts `4097` and `8192`: sequential eager `18` positions, mean `0.0177`, p99 `0.0660`, max `0.0660`. -- Both run with the persistent, per-step refilled pointer tables (`2 passed`, 65 s). - -Not in this step: TP CUDA Graph capture (#1005) and the same-context -per-request-vs-batched throughput A/B tracked in #1001; no performance claim is -made before that rerun lands. - -## Rebase onto #870 (2026-08-20) - -#870 landed its own Phase 1/2a upstream while the parallel line on -`feat/qwen35-tp2-batched-decode` had implemented its own Phase 1/2a/2b -plus a batched-decode fix on the old main. The rebase onto #870 -ported only the two deltas #870 lacks, with #870's code as the base: - -1. **Phase 2b — sharded linear-attention/GDR state** (commit - `feat(qwen35): shard linear-attention/GDR state per TP rank`). #870's - `recurrent_state.rs` had no rank sharding and its linear-attention - weights loaded replicated; the port adds rank-local slices end to end - (`weight_loader` stitch/shard loaders, `config/tp.rs` `local_linear_*` - accessors, `weights/layers.rs` per-rank stitched qkv/conv1d + row/col - shards carried on `WeightSource`, - `recurrent_state`/`decode_buffers`/`prefill_buffers` at local sizes, - `batch_decode`/`prefill` local head counts + all-reduce after linear - `out_proj`, TP-local `batch_decode_full_attention_via_prefill` so 27B - TP2 group-6 eager decode routes through prefill). `tp_executor.rs` - only took the capacity-math and `RecurrentState::new` signature - changes; #870's worker protocol untouched. -2. **Step 3 — batched eager TP decode** (commit - `perf(qwen35): batch eager decode rows under TP`). #870's - `execute_decode_rows` looped per request with bs=1 forwards and - capacity-1 per-request pointer tables. The port adds `run_decode_batch` - (one `batch_decode_eager_logits` over all decode rows, one persistent - `LinearStatePointerTables` per worker refilled each step — see the - batched eager TP decode (#1004) section above — one batched rank-0 - `select_batch`, per-row fan-out in command order) inside - #870's `execute_decode_rows`, keeping its validation and response - contracts; `TpRequestState.linear_pointer_tables` removed. - -What #870 already covered (not ported): Phase 1 dense TP, the Phase 2a -unified command/scheduler surface (`TpUnifiedPlan`, command start gates, -dispatch/response validators, drop-expectation lifecycle proofs), and the -scheduler planner-gate/test updates — ours' scheduler, -`scheduler/tests.rs`, and `e2e_scheduler.rs` deltas were subsumed -upstream, so those files resolved to #870's versions except the -`alloc_recurrent` signature change. #967/#968 have since split the -scheduler into `scheduler/{mod,backend,plan,tp}.rs`, so this branch's -slot-tracking deltas land there: the `ActiveBackendState::Tp` slot field -and dispatch arms in `mod.rs`, `TpSchedulerBackend` slot compaction in -`backend.rs`, and the decode-item/alignment helpers in `tp.rs`. - -Validation on 2× RTX 4090 (venv NCCL on `LD_LIBRARY_PATH`): - -- `cargo check --release -p pegainfer-qwen35 --features qwen35` clean; - `cargo fmt --check -p pegainfer-qwen35` clean. -- Lib unit suite 101/101 (includes the four sharding layout tests: - segment tables, conv kernel-dim scaling, TP1 identity, synthetic - safetensors stitch contract). -- 9B TP2 HF short+long gates PASS (24.7 s); 9B TP2 scheduler e2e - (`test_e2e_qwen35_scheduler_tp2`) PASS (27.1 s). -- 27B TP2 HF short+long gates PASS (64.7 s) — 27B TP2 fits on 2×48 GB - only because of the Phase-2b sharding (the acceptance criterion for the - port). 27B TP2 scheduler e2e PASS (68.8 s). - -## P2c — CUDA Graph under TP (2026-08-20) - -Implemented the locked P2c design from `tp-design.md`: decode CUDA Graphs -under TP, gated on the TP-local decode GQA group. - -**Gate.** Graph mode is active iff `enable_cuda_graph && -config.decode_group_is_compiled()`. The group ratio is TP-invariant, -so 27B TP2 (group 6, not in `SUPPORTED_GQA_GROUP_SIZES`) keeps the batched -eager path byte-for-byte while 4B/9B TP2 capture. The old fail-closed -rejections (`config.rs` `validate_for`, `tp_executor.rs` startup ensure, -`lib.rs` TP branch) were replaced by the gate; startup logs once when graph -was requested but the group gate keeps decode eager. - -**State model.** Scheduler owns slots, workers execute: - -- `ActiveBackendState::Tp` gains `slot_idx` (dense `active` position, assigned - at promote via `slot_for_new_request`, updated on compaction); - `tp_decode_items` emits rows with explicit `slot_idx` and workers assert - `slot_idx == row`. -- Graph workers hold a `BatchDecodeGraphState` at - `bucket_for(effective_max_batch)` fixed-address slots plus a - `slot_map: Vec>`. On a request's first decode row the - worker D2D-copies its prefill recurrent state into the slot - (`copy_state_to_slot`) and drops the per-request allocation. -- `DropRequest` carries `compaction: Option`; the worker - validates occupancy against `slot_map` (`slot_compact`), applies the D2D - move (`BatchDecodeGraphState::move_slot_within`), and poisons on mismatch. - Requests retired between promotion and their first decode row legitimately - have no materialized slot; `slot_compact` tolerates exactly that case and - skips the GPU move. -- Convenience `execute_prefill`/`execute_decode`/`drop_request` (model-local - tests) keep a Mutex-guarded slot tracker mirroring the single-GPU - `Qwen35Executor`; scheduler flows pass explicit slots via - `execute_decode_items`/`drop_request_with_compaction` and never touch the - tracker. - -**Capture/replay.** Startup pre-capture sweep ported from qwen3: -`TpWorkerCommand::Precapture { phase }` over Warmup (new -`Qwen35Model::warmup_tp_collective` — one all-reduce per bucket message size; -without it the lazy NCCL connect wedges inside capture), Capture + Launch per -bucket `[1,2,4,8,16,32,64]` up to `bucket_for(max_batch)` with synthetic rows, -Finalize asserting all buckets captured. Dedicated 600 s abort watchdog (the -60 s NCCL startup timeout is too small for the sweep). -`batch_decode_graph` gained `DecodeGraphUse` (Serve lazy / CaptureOnly / -Replay); TP serve time is Replay-only. Workers tune decode GEMM algos on the -worker thread before capture (cuBLASLt plans are thread-local). Graph state -is declared before `model` in `TpWorkerState` so graphs drop before the NCCL -comm. Sampling/logprobs stay rank-0 host-side outside the graph. Eager -workers ignore `slot_idx`/`compaction`, keeping 27B TP2 byte-identical. - -**Gotcha fixed during validation:** `track_retired_slot` used -`bool::then_some`, which evaluates eagerly and indexed past the tail when the -retired request was the last slot — use `then` for the lazy closure. - -**Validation (2× RTX 4090):** - -- `cargo check --release -p pegainfer-qwen35 --features qwen35` clean; lib - suite 105/105 (new: slot map admit/compact/mismatch CPU tests); - `cargo fmt --check` clean. -- 9B TP2 HF gates (`--test-threads=1`): eager sequential+batched PASS; - graph sequential replay (identical fingerprints across reruns), - bucket-straddling batched replay (5→bucket 8, 3→bucket 4), and - post-compaction replay after a mid-batch drop all within the existing TP2 - tolerances (worst graph arm: mean 0.0228, p99 0.1002 against MEAN_TOL 0.06 - / P99_TOL 0.20; eager arm mean 0.0227/0.0230). -- 9B TP2 scheduler e2e eager + graph variants PASS; 9B TP2 serving smoke now - launches with `--cuda-graph true` (graph acceptance replaced the old - rejection assertion). -- 27B TP2 HF short+long and scheduler e2e PASS unchanged — the gate log - confirms group 6 keeps decode eager (the graph test variant skips itself - via `graph_enabled()`). -- **Historical, not re-measured on this stack:** 9B TP2 serving benchmark - from the pre-rebase #946 branch (2026-08-20; `pegainfer-server --tp-size 2 --port 18093`, - vllm-bench `openai` backend, random 128-in/256-out, 64 prompts at - concurrency 16, greedy, ignore_eos, seed 42): - - | arm | steady output tok/s | mean TPOT (ms) | total tok/s | - |-----|--------------------:|---------------:|------------:| - | CUDA Graph on | 767.15 | 20.04 | 1146.65 | - | CUDA Graph off | 705.86 | 21.99 | 1054.56 | - - On that branch graph decode measured +8.7% steady output tok/s (-8.8% TPOT) - at 16 concurrent. No performance claim is made for this stack until the - same-context A/B is rerun on it (tracked in #1001). - The design's "record in `bench_snapshots/`" step was skipped: the - in-process snapshot gate is retired (`docs/conventions/bench-regression.md`), - so the HTTP bench numbers live here instead. - -**Test-isolation note:** TP2 GPU tests must run with `--test-threads=1`. Two -TP executors sharing the GPUs perturb cuBLASLt algorithm selection (workspace -pressure), which once flipped a sequential-replay fingerprint comparison in -the *eager* test while the graph test ran concurrently. - -## Follow-Ups - -- P2c CUDA Graph under TP landed for compiled decode GQA groups (4B/9B); 27B - TP2 graphs stay gated off until group-6 batch-decode kernels are compiled - (`SUPPORTED_GQA_GROUP_SIZES`). The eager path is the 27B fallback and must - not regress. -- 27B TP2 knowledge-benchmark parity (2026-08-20, validated pre-rebase on - the f4c66780 line; `docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md`): - C-Eval 88.11 vs official 90.5 (full 1346, thinking-cap truncation - rerun-merged) is inside the cross-harness band and is the current - evidence; the MMLU-Redux 94.09 figure was scored with the pre-fix - first-capital extractor — rerun-pending, do not cite. MMLU-Pro / - SuperGPQA sampled runs remain outstanding; rerun on this rebased branch - before citing parity. -- P2B sharded linear-attention/GDR state landed (see "Rebase onto #870"); keep the completed P2A lifecycle and ID contracts unweakened. -- Promote any stable contract changes discovered here back into `tp-design.md` through the design-doc branch. -- Decide whether Qwen3.5 server CLI should accept arbitrary TP device ordinals instead of only `0..tp_size`. -- Consider lifting the per-device Triton AOT handle lesson into a kernels or runtime subsystem doc if another model hits the same issue. +> **Last touched:** 2026-10 + +## Runtime ownership + +The architecture contract is in [TP design](tp-design.md). The implementation +is split between `scheduler/{mod,backend,tp}.rs` and `tp_executor.rs`: + +- `Qwen35Scheduler` owns logical admission, prefill/decode planning and dense + decode-slot order. It writes request updates through `RequestLedger`; the + shared driver publishes metrics and commits those updates. +- `TpSchedulerBackend` maps frontend request ownership to model-local + `RequestId`s and builds commands for `Qwen35TpExecutor`. The two ID types + serve different boundaries; client labels are not lifecycle keys. +- Rank workers own model shards, CUDA contexts, thread-local cuBLAS, NCCL + comms, KV buffers and recurrent/conv tensors. They bind their contexts and + initialize CUDA/NCCL resources on their own threads. The controller starts + its scheduler after worker startup completes, without a second handshake. +- The controller's `KvCacheManager` owns logical request pages. Immutable + `KvView`s describe the same page IDs on every rank; each worker writes its + own KV shard. Capacity is capped by the smallest rank-local pool after + validating consistent geometry and snapshot slots. + +TP ordinal validation precedes model loading. Ordinals must be distinct and +fit the generated Triton AOT handle table. The generated stubs cache module +and function handles per device and check the device index before accessing +the table; process-global CUDA function handles are invalid across devices. + +## Shards and execution + +Full-attention and MLP projections use dense TP sharding. Gated `q_proj` +weights are interleaved by head (`[head0 q][head0 gate][head1 q][head1 gate]`), +so loaders preserve head ranges rather than rebuilding separate q/gate +halves. Embeddings and `lm_head` remain replicated. + +Linear-attention/GDR weights, recurrent/conv state and scratch shapes are +rank-local. The local linear output is followed by an all-reduce. Recurrent +state and KV memory sizing use the configured `max_batch`, rather than +reserving an unrelated default number of slots. + +Decode rows run as one batched forward on each rank, with sampling and +logprobs returned by rank 0. Eager GDR pointer tables are allocated once at +scheduler capacity and refilled from live rows each step: slot compaction +must not leave stale row addresses. Request-local sampling steps must remain +independent of a request's position in the batch. + +When prefill and decode coexist, the common planner produces one +`RunUnifiedStep`. Every rank executes prefill before decode in the same +collective order. The scheduler consumes decode results before promoting +prefill results, allowing acknowledged retirement to free slots first. +TP does not support the single-GPU stream-overlap or adaptive `auto` policies. + +## Dispatch and result contracts + +State-mutating commands share a `Pending -> Execute | Cancel` start gate. +The controller enqueues to every rank before choosing `Execute`; a partial +enqueue failure cancels the delivered prefix and poisons the executor. +Workers cannot mutate request state or enter kernels/collectives before the +gate opens. Gate transitions execute unconditionally, never only inside a +`debug_assert!`. + +Structural plan validation runs before dispatch. Unified plans require both +halves, fit the configured capacity and contain unique, mutually disjoint +prefill/decode IDs and non-empty chunks. Workers revalidate local existence, +phase and capacity after the gate opens. A post-dispatch inconsistency is +replica-fatal; it is not retried as an individual request failure. + +Replies must cover every rank exactly once with the expected variant. Rank +0 supplies sampled prefill/decode/unified artifacts; other ranks acknowledge. +Drop uses `DropAck { existed }` from every rank. Duplicated or missing ranks, +out-of-range ranks, unexpected variants and inconsistent existence results +poison the executor. + +The scheduler aligns returned artifacts by model-local `RequestId`, not row +position. Decode needs one result per active ID. Prefill needs a result only +for final chunks: a non-final chunk has no artifact, while a final artifact +may legitimately have no requested logprobs. Unknown, duplicated, non-final +or missing result IDs after execution poison the replica. + +## Cancellation, completion and failure + +Cancellation is the frontend abort flag observed through the ledger. Pending +requests retire before allocation; resident requests use normal backend +cleanup before planning, so released capacity is available to admission in +the same step. + +Admission creates rank-local recurrent state even on a cold prefix miss. +Consequently, cancelling an admitted prefill at cursor zero still requires +`DropExpectation::MustExist`, as do active retirement and final-prefill +cleanup. +The exact-rank drop contract accepts all-true for `MustExist` and all-false +for `MustBeAbsent`; mixed or unexpected answers are lifecycle divergence. + +Successful completion is recorded in the ledger only after backend +retirement succeeds. In TP this requires all-rank drop acknowledgement, so a +cleanup failure cannot leak a successful terminal or its final token. +Remaining open accounts are failed by the shared driver when the scheduler +returns an engine-fatal error. Poisoned workers are not retried through +healthy cleanup; backend teardown owns their remaining resources. + +## Joint prefix cache + +Prefix reuse requires full-attention KV and a complete recurrent/conv +snapshot at the same boundary. Key, pin and LRU metadata is centralized; +snapshot tensors remain rank-local. Publication reserves a common slot, +saves it on every rank and verifies the boundary before publishing the key. +Restore reports a hit only after every rank confirms that boundary. + +The snapshot budget is reserved independently on each rank through +`qwen35-prefix-cache-mib`; zero keeps cold serving. See +[prefix-cache design](prefix-cache.md) for boundary and eviction rules. + +## P2c — CUDA Graph under TP + +Graph mode requires both the launch option and a compiled decode GQA group. +4B/9B TP2 can capture; group-6 27B execution stays eager until that kernel is +available. A graph request falling back to eager is logged at startup. + +The scheduler assigns dense `slot_idx` values and supplies compaction moves; +workers do not infer slot ownership. Each graph worker owns fixed-address +state for the rounded batch capacity. A request's first decode row copies +its prefill recurrent state into its assigned slot. Retirement validates +occupancy and applies the explicit D2D compaction; a request retired before +its first decode legitimately has no materialized slot. + +Startup warms NCCL collectives, captures and launches each required bucket, +then verifies all buckets are captured. Lazy NCCL connection inside capture +can hang, so warmup and the capture watchdog are required. Serving replays +existing graphs; it never captures during a live request. Sampling and +logprobs stay outside capture. Graph state is destroyed before the model's +NCCL resources during worker teardown. + +Low-level convenience executor calls maintain their own guarded slot +tracker. The scheduler uses explicit-slot commands and must not also update +that tracker. + +## Validation surface + +The tests exercise different ownership boundaries: + +- TP executor gates observe rank-local state to prove partial-dispatch + cancellation, lifecycle divergence, disconnected-worker failure, mixed + execution, full capacity recovery and clean readmission under a fresh ID. +- Scheduler lifecycle tests cover acknowledged completion, cancellation, + request rejection and cleanup failures. +- `e2e_scheduler` covers eager and Graph request flow, mixed sampling and + post-cancellation health. `prefix_cache` covers joint restore and eviction. +- `hf_golden_gate` checks short/long numeric replay, batched buckets and + post-compaction replay. Fixtures match `config_sha256` and the checkpoint + revision, including Qwen3.8; see [accuracy](accuracy.md). +- `serving_tp2` covers real HTTP model discovery, JSON/SSE completions, + concurrency and logprobs. Historical dependent multi-turn evidence is + retained in the [TP2 serving benchmark](../../benchmarks/qwen35-tp2-phase2a-multiturn.md). + +The shared model-fixture resolver uses `PEGAINFER_TEST_MODEL_PATH`; an absent or invalid +fixture emits `SKIP`, which is not a GPU pass. TP2 tests are ignored by +default and use `PEGAINFER_TEST_TP_DEVICES` (two distinct ordinals, default +`0,1`). `PEGAINFER_TEST_FRONTEND_MODEL_PATH` optionally selects frontend +metadata; otherwise it uses the engine fixture. Hand-downloaded HF fixtures +need a resolvable revision or `PEGAINFER_TEST_MODEL_REVISION`. + +Run GPU gates serially to avoid capacity pressure and cuBLASLt tuning +interference between independent executors. diff --git a/docs/subsystems/frontend/frontend-architecture.md b/docs/subsystems/frontend/frontend-architecture.md index 045c9626d..17411c3d8 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -1,12 +1,12 @@ # Frontend architecture: pegainfer-frontend and the engine boundary -**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Gemma 4 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite still launch through it). The vLLM Rust dependency revision is recorded in `Cargo.toml` and `Cargo.lock`. **Next step: migrate glm52, then delete the legacy contract.** +**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Qwen3.5, Gemma 4 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/kimi-k2/deepseek-v2-lite still launch through it). The vLLM Rust dependency revision is recorded in `Cargo.toml` and `Cargo.lock`. **Next step: migrate glm52, then delete the legacy contract.** Last touched: 2026-10 ## The boundary, in one sentence -An engine is a set of schedulers, each a `Scheduler` implementation driven by the contract's polling loop: *the frontend submits `Request`s into a scheduler and receives one `StepOutputs` message per scheduler step, in which every touched request has exactly one flat `RequestUpdate`; every request ends in exactly one terminal.* Tokenizer, chat templates, HTTP, metrics, LoRA routing live north of the scheduler; KV, batching, CUDA live south. The contract contains no CUDA types, and the frontend holds no model-layer structs — everything it touches is a contract type (channels, `JoinHandle<()>`, POD info). How many schedulers an engine exposes and what each one means (DP rank, P/D role) is the model line's decision; the contract attaches no rank semantics, and TP/EP lockstep is encapsulated below the `Scheduler` impl. Placement across schedulers is frontend policy over the load feed. +An engine is a set of schedulers, each a `Scheduler` implementation driven by the contract's polling loop: *the frontend submits `Request`s into a scheduler and receives a `StepOutputs` message for each step with request updates or metric changes, with each request's updates merged into one flat `RequestUpdate`; requests finish, fail or are rejected with one terminal, while cancellation retires them silently.* Tokenizer, chat templates, HTTP, metrics, LoRA routing live north of the scheduler; KV, batching, CUDA live south. The contract contains no CUDA types, and the frontend holds no model-layer structs — everything it touches is a contract type (channels, `JoinHandle<()>`, POD info). How many schedulers an engine exposes and what each one means (DP rank, P/D role) is the model line's decision; the contract attaches no rank semantics, and TP/EP lockstep is encapsulated below the `Scheduler` impl. Placement across schedulers is frontend policy over the load feed. ## The step contract @@ -20,7 +20,7 @@ pegainfer-frontend/src/engine/ │ # DeferredFinish remains available for P/D handoff ├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, │ # prompt/completion tallies, one merged update per touched id -├── wiring.rs # scheduler_pair, SchedulerHandle (submit/take_steps/load), +├── wiring.rs # scheduler_pair, SchedulerHandle (submit/take_steps/metrics), │ # Engine { schedulers, info, lora }, LiveScheduler, │ # EngineInfo, LaunchedEngine { Handle | Stepped } ├── control.rs # LoRA capability — outside the contract: LoraControl vocabulary, @@ -31,14 +31,14 @@ pegainfer-frontend/src/engine/ Design decisions worth knowing before touching it: -- **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. +- **Step-batched wire.** One message per scheduler step with request updates or metric changes, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. - **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step. -- **Typed stop contract.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in the emitted output and carries `Terminal::Finished.stop_cause`: The vLLM primary EOS maps to `Eos(id)` with no wire `stop_reason`; secondary EOS IDs join the explicit stop set and map to `Token(id)`, reporting the actual matching ID. `ignore_eos` disables only model EOS. Legacy producers may leave the cause empty while they are migrated individually. +- **Typed stop contract.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in the emitted output and carries `Terminal::Finished.stop_cause`: The vLLM primary EOS maps to `Eos(id)` with no wire `stop_reason`; secondary EOS IDs join the explicit stop set and map to `Token(id)`, reporting the actual matching ID. `ignore_eos` disables only model EOS. Qwen3.5 still suppresses EOS and leaves the cause empty for the bridge's legacy sentinel; its stop-policy migration remains follow-up work. - **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver. - **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration. -- **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. +- **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. Qwen3.5 likewise polls its in-flight prefill while decode is active, but waits on the prefill event inside the step when no decoder remains. - **Abort is a flag, not channel teardown.** `SchedulerHandle::submit` returns a `RequestControl`; the frontend flips its boolean abort flag and the scheduler retires the request silently on its next touch (no terminal — the frontend already dropped its state for that id). -- **Channels:** the submit channel is crossbeam (sync consumer on the scheduler thread), steps are tokio mpsc (async consumer in the bridge); load is a shared cell read via `SchedulerHandle::load()` — pull-only by design, "notify me on load change" is deliberately unrepresentable (the driver busy-polls, so a subscription edge would fire per spin). All channels unbounded on purpose — admission control is the scheduler's job, expressed as `Rejected`, never as backpressure on submit. +- **Channels:** the submit channel is crossbeam (sync consumer on the scheduler thread), steps are tokio mpsc (async consumer in the bridge); load is a shared cell read via `SchedulerHandle::metrics()`. Metric changes without request updates send an empty step, so cancellation can reset HTTP gauges; unchanged polling iterations send nothing. All channels unbounded on purpose — admission control is the scheduler's job, expressed as `Rejected`, never as backpressure on submit. - **Control plane lives outside the contract.** `Scheduler` has no control method and the contract carries no control channel. A capability like LoRA is a private channel the model crate mints *before* `spawn_scheduler` — the scheduler closes over the receiver, the `LoraClient` sender surfaces as `Engine.lora: Option`, and the `Option` *is* the capability (no `bool` flag, no registry until a second capability exists). The vocabulary (`LoraControl`, `LoraClient`) is still defined in the frontend crate because the frontend must speak it without holding model structs; only the wiring is the model's business. - **A scheduler that exits stops the server.** Once the driver returns or its thread unwinds, `SchedulerHandle::exited` resolves and the stepped bridge fails after forwarding what is already on the step stream. A failing stepped bridge gives its output sender up to a second to flush, ending with the `ENGINE_CORE_DEAD` sentinel, which makes the engine core client fail every request it still tracks, including ones no bridge has read. The server then shuts down so it can be restarted, and every other bridge aborts its requests and flushes its output before the scheduler threads are joined. @@ -49,14 +49,14 @@ Design decisions worth knowing before touching it: - `step(ledger)` — one scheduling step: admit/reject, execute, push tokens, finish/fail/retire. Return `Err` only for engine-fatal states. - `metrics()` — KV occupancy + running/waiting counts for routers and shutdown draining. - Extra capabilities (LoRA etc.) are not trait methods: mint the private channel before spawning, close the scheduler over the receiver, drain it inside `step`. -2. `spawn_scheduler(name, scheduler)` per scheduler; return `Engine { schedulers, info: EngineInfo { kv_capacity, servable_len }, lora }` — the required-metadata fields are the checklist, and `lora: Some(client)` only when the line actually serves adapter control. +2. `spawn_scheduler(name, scheduler)` per scheduler, or `scheduler_pair` plus `drive` when launch must await thread-local initialization, as single-GPU Qwen3.5 does for CUDA/cuBLAS; return `Engine { schedulers, info: EngineInfo { kv_capacity, servable_len }, lora }` — the required-metadata fields are the checklist, and `lora: Some(client)` only when the line actually serves adapter control. 3. `ModelLine::launch` returns `LaunchedEngine::Stepped(engine)`. `pegainfer-sim` is the CPU-only reference (`SimScheduler` in `pegainfer-sim/src/lib.rs`); it has no `ModelLine` and hands the `Engine` to `vllm::serve` directly. The contract's own invariants are tested in `ledger.rs`/`driver.rs` tests; the qwen3 adapter's contract tests (`frontend_adapter/tests.rs`) are the CPU reference for testing a model's protocol behaviour end to end with a fake executor. Gemma 4's ignored GPU lifecycle gates drive the same `spawn_scheduler`/`StepOutputs` surface. ## The legacy handle contract (migration pending) -`request.rs`/`event.rs`/`sink.rs`/`kv.rs`/`handle.rs` still carry the previous generation: `launch -> EngineHandle`, per-request `TokenSink` events (`Scheduled … Token* … terminal` by convention), send-failure-as-cancellation. glm52, qwen35, kimi-k2 and deepseek-v2-lite launch through it (`LaunchedEngine::Handle`), and the vllm stack keeps both bridge paths (`bridge.rs` for handles, `bridge/stepped.rs` for step engines). KV-prefix resolution (`KvPrefix`, `submit_resolved`) currently exists only on the legacy path; step schedulers can refuse only the features represented by `Request`, including non-null `kv_transfer_params`. +`request.rs`/`event.rs`/`sink.rs`/`kv.rs`/`handle.rs` still carry the previous generation: `launch -> EngineHandle`, per-request `TokenSink` events (`Scheduled … Token* … terminal` by convention), send-failure-as-cancellation. glm52, kimi-k2 and deepseek-v2-lite launch through it (`LaunchedEngine::Handle`), and the vllm stack keeps both bridge paths (`bridge.rs` for handles, `bridge/stepped.rs` for step engines). KV-prefix resolution (`KvPrefix`, `submit_resolved`) currently exists only on the legacy path; step schedulers can refuse only the features represented by `Request`, including non-null `kv_transfer_params`. ## Crate layout @@ -93,7 +93,7 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the ## Protocol stacks -**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a `Finished{Stop}` appends the stop sentinel token, which is how usage keeps counting the suppressed EOS). +**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a `Finished{Stop}` without a typed cause appends the legacy stop sentinel token, which is how usage keeps counting Qwen3.5's suppressed EOS). **`dynamo` (planned second stack).** dynamo's `lib/llm` in-process path removes the wire protocol entirely (`EngineConfig::InProcessTokens` + `run_input`). The step contract was shaped so this stack can consume `StepOutputs` directly without impersonation overhead. Decision gate: prototype, A/B against the vllm stack, let TTFT/step-overhead numbers pick the default. @@ -111,7 +111,7 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the ## Next step -Migrate glm52 onto the step contract (the multi-scheduler pilot; brings P/D and EP requirements), then qwen35/kimi-k2/deepseek-v2-lite, then delete the legacy contract modules and `LaunchedEngine::Handle`. +Migrate glm52 onto the step contract (the multi-scheduler pilot; brings P/D and EP requirements), then kimi-k2/deepseek-v2-lite, then delete the legacy contract modules and `LaunchedEngine::Handle`. ## Upstream vLLM dependency pin diff --git a/docs/subsystems/frontend/prometheus-metrics.md b/docs/subsystems/frontend/prometheus-metrics.md index e17df01c8..f8a33a71e 100644 --- a/docs/subsystems/frontend/prometheus-metrics.md +++ b/docs/subsystems/frontend/prometheus-metrics.md @@ -1,19 +1,19 @@ # Prometheus /metrics via the vLLM frontend -**TL;DR:** `/metrics` exposes request histograms for every model and engine gauges for schedulers that publish `LoadSnapshot`: Qwen3 and Qwen3.5 use one logical engine, while GLM5.2 EP8/DP8 uses eight rank-local engines and GLM5.2 TP8 uses one logical engine. The bridge forwards each partition's stats under the same identity the vLLM frontend uses for least-load routing. +**TL;DR:** `/metrics` exposes request histograms for every model and engine gauges for schedulers that publish load: Qwen3, Qwen3.5 and Gemma 4 use one logical engine, while GLM5.2 EP8/DP8 uses eight rank-local engines and GLM5.2 TP8 uses one logical engine. The bridge forwards each partition's stats under the same identity the vLLM frontend uses for least-load routing. -Last touched: 2026-07 +Last touched: 2026-10 ## How the numbers flow Two independent paths feed the upstream Prometheus registry (`vllm-metrics`, served by `vllm-server` at `/metrics` with its HTTP middleware counters): 1. **Per-request path (works for every model crate).** The bridge stamps each request's first output with `Queued`/`Scheduled` timestamps and `PrefillStats` (prompt/computed/cached token split). The upstream `RequestMetricsTracker` turns those into `time_to_first_token_seconds`, `inter_token_latency_seconds`, `request_queue_time_seconds`, `prompt_tokens_total`, `generation_tokens_total`, `request_success_total`, `prompt_tokens_by_source_total`, … unconditionally — `disable_log_stats` only gates the periodic *text* logger, not Prometheus. -2. **Engine-gauge path (needs one `LoadSnapshot` watch per scheduler partition).** The scheduler publishes `LoadSnapshot { kv_used_blocks, kv_total_blocks, num_running_reqs, num_waiting_reqs }` at scheduler boundaries; one bridge identity per partition forwards its snapshot as a stats-only `RequestBatchOutputs`. The enclosing `engine_index` is both the routing identity and the Prometheus `engine` label. Watches coalesce to ≤1 message per scheduler step, and the scheduler's final idle publish settles the gauges back to 0. +2. **Engine-gauge path.** Step-contract schedulers publish `SchedulerMetrics { kv_used_blocks, kv_total_blocks, num_running_reqs, num_waiting_reqs }` through the shared driver's post-step metrics cell. The bridge reads the latest cell at startup and when forwarding a step; it is not a snapshot bound to that step's tokens. Metric-only changes send an empty step, so cancellation of the last request can reset the gauges. Unchanged idle iterations send nothing. Legacy schedulers use `LoadSnapshot` watches. Both bridges forward stats in `RequestBatchOutputs`, whose `engine_index` is both the routing identity and the Prometheus `engine` label. -For a single-partition model, `EngineHandle::with_load_watch` keeps the original one-engine contract. Qwen3.5 uses that contract for both its single-GPU backend and its TP backend because both execute one logical request stream through one scheduler. A partitioned scheduler uses `with_load_watches`, and the frontend launch declares the same engine count; a mismatch fails startup. GLM5.2 EP8 therefore registers engines 0–7, each bound to its own pending queue and KV pool. TP8 registers only engine 0 because its eight workers mirror one logical request stream. +Qwen3.5 exposes one step-contract scheduler for both single-GPU and TP. For legacy engines, `EngineHandle::with_load_watch` keeps the original one-engine contract. A partitioned scheduler uses `with_load_watches`, and the frontend launch declares the same engine count; a mismatch fails startup. GLM5.2 EP8 therefore registers engines 0–7, each bound to its own pending queue and KV pool. TP8 registers only engine 0 because its eight workers mirror one logical request stream. -Measured cost is noise in both covered configurations: +Historical watch-path measurements found noise-level cost in both covered configurations: - Qwen3 TPOT: 10.6387 ms (main) vs 10.6395 ms (metrics branch) over 828 tokens. - GLM5.2 EP8, three-run median at concurrency 64: 1268.58 vs 1264.82 output tok/s (-0.30%); TPOT p50 41.76 vs 41.35 ms. @@ -21,11 +21,9 @@ Measured cost is noise in both covered configurations: ## What deliberately reads zero (state at capture time) - `prefix_cache_queries/hits` and the by-reason waiting split (`reason="deferred"` is driven by a skipped-request counter we don't report; all waiting shows as `reason="capacity"`). -- Spec-decode counters, per-GPU FLOPs/bytes estimates, KV-block residency histograms, cudagraph stats — the bridge sends `SchedulerStats::default()` for these fields. -- Every model crate whose scheduler doesn't publish a `LoadSnapshot` watch (currently deepseek and kimi) gets path 1 only; its engine gauges are absent, not lying-zero — the bridge skips the stats task for that partition when no watch exists. +- Per-GPU FLOPs/bytes estimates, KV-block residency histograms, cudagraph stats — the bridge sends `SchedulerStats::default()` for these fields. +- Every legacy model crate whose scheduler doesn't publish a `LoadSnapshot` watch (currently deepseek and kimi) gets path 1 only; its engine gauges are absent, not lying-zero — the bridge skips the stats task for that partition when no watch exists. -## Validated coverage and next step - -Qwen3.5 single-GPU live RTX 5090 validation confirmed that running and KV gauges rise during generation, waiting rises under batch-slot pressure, and all three return to zero after drain and recovery. The commands and metric samples are recorded in [Qwen3.5 Scheduler LoadSnapshot](../../models/qwen35/load-snapshot.md#validation-boundary). TP uses the same scheduler publication path but was not part of that live run. +## Next step Next, wire the DeepSeek-V2-Lite and Kimi-K2 schedulers using the same recipe, and report real prefix-cache query/hit counters instead of zeros. A future partitioned model must expose its logical scheduler partitions instead of averaging them behind engine 0. diff --git a/pegainfer-core/src/ops.rs b/pegainfer-core/src/ops.rs index a0a683663..daf3c1963 100644 --- a/pegainfer-core/src/ops.rs +++ b/pegainfer-core/src/ops.rs @@ -8,16 +8,14 @@ mod paged_plan; #[cfg(feature = "kernel-call-trace")] mod traced; -pub use attention::paged_attention_batch_decode_hd256_into; pub use attention::paged_attention_batch_decode_into; pub use attention::paged_attention_batch_decode_split_kv_into; -pub use attention::paged_attention_batch_decode_via_prefill_hd256_into; pub use attention::prefill_attention_paged_into; pub use paged_plan::PrefillPagedPlan; pub use paged_plan::SplitKvCsr; pub use paged_plan::build_split_kv_csr; pub use pegainfer_kernels::ops::GEMM_LT_MAX_N; -pub use pegainfer_kernels::ops::Hd512DecodeMetadata; +pub use pegainfer_kernels::ops::GatedNormActivation; pub use pegainfer_kernels::ops::LoraDecodeGroupedProjection; pub use pegainfer_kernels::ops::SUPPORTED_GQA_GROUP_SIZES; pub use pegainfer_kernels::ops::SuppressIds; @@ -27,8 +25,6 @@ pub use pegainfer_kernels::ops::add_batch_into; pub use pegainfer_kernels::ops::advance_decode_metadata; pub use pegainfer_kernels::ops::argmax; pub use pegainfer_kernels::ops::argmax_batch_bf16_into; -pub use pegainfer_kernels::ops::batch_prefill_paged_hd512_into; -pub use pegainfer_kernels::ops::batch_prefill_paged_window_hd256_into; pub use pegainfer_kernels::ops::bf16_hidden_to_f32_into; pub use pegainfer_kernels::ops::copy_hidden_rows_into; pub use pegainfer_kernels::ops::copy_hidden_token_range_into; @@ -69,16 +65,8 @@ pub use pegainfer_kernels::ops::linear; pub use pegainfer_kernels::ops::lora_decode_fused_delta_group3_into; pub use pegainfer_kernels::ops::lora_decode_fused_delta_into; pub use pegainfer_kernels::ops::pack_lora_b_rows_into; -pub use pegainfer_kernels::ops::paged_attention_batch_decode_split_kv_hd512_into; -pub use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd256_into; -pub use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd512_into; -pub use pegainfer_kernels::ops::qk_norm_partial_rope_paged_decode_hd512_into; -pub use pegainfer_kernels::ops::qk_norm_partial_rope_paged_prefill_hd512_into; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::qk_norm_rope_batch_decode_into; -pub use pegainfer_kernels::ops::qk_norm_rope_prefill_hd256_plain_into; -pub use pegainfer_kernels::ops::qkv_norm_rope_paged_decode_hd256_plain_into; -pub use pegainfer_kernels::ops::qkv_norm_rope_paged_prefill_hd256_plain_into; pub use pegainfer_kernels::ops::rms_norm; pub use pegainfer_kernels::ops::rms_norm_add_rms_norm_round_batch_into; pub use pegainfer_kernels::ops::rms_norm_add_scale_batch_into; @@ -100,8 +88,6 @@ pub use pegainfer_kernels::ops::silu_mul_batch_into; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::silu_mul_fused_batch_into; pub use pegainfer_kernels::ops::single_decode_nhd_into; -pub use pegainfer_kernels::ops::single_prefill_hd256_into; -pub use pegainfer_kernels::ops::single_prefill_hd512_into; pub use pegainfer_kernels::ops::single_prefill_nhd_causal_into; pub use pegainfer_kernels::ops::single_prefill_nhd_noncausal_into; pub use pegainfer_kernels::ops::softcap_bf16_in_place; diff --git a/pegainfer-core/src/ops/attention.rs b/pegainfer-core/src/ops/attention.rs index 1fecc3215..74bbdf451 100644 --- a/pegainfer-core/src/ops/attention.rs +++ b/pegainfer-core/src/ops/attention.rs @@ -144,75 +144,3 @@ pub fn paged_attention_batch_decode_split_kv_into( batch_size, ) } - -#[allow(clippy::too_many_arguments)] -pub fn paged_attention_batch_decode_hd256_into( - ctx: &DeviceContext, - q: &HiddenStates, - k: &HiddenStates, - v: &HiddenStates, - kv_buffer: &CudaSlice, - layout: &KvLayout, - layer: usize, - page_indices_d: &CudaSlice, - page_indptr_d: &CudaSlice, - last_page_len_d: &CudaSlice, - positions_d: &CudaSlice, - request_indices_d: &CudaSlice, - kv_tile_indices_d: &CudaSlice, - kv_chunk_size_d: &CudaSlice, - output: &mut HiddenStates, - num_qo_heads: usize, - batch_size: usize, -) -> Result<()> { - pegainfer_kernels::ops::paged_attention_batch_decode_hd256_into( - ctx, - q, - k, - v, - kv_buffer, - &layout.kernel_layout(), - layer, - page_indices_d, - page_indptr_d, - last_page_len_d, - positions_d, - request_indices_d, - kv_tile_indices_d, - kv_chunk_size_d, - output, - num_qo_heads, - batch_size, - ) -} - -#[allow(clippy::too_many_arguments)] -pub fn paged_attention_batch_decode_via_prefill_hd256_into( - ctx: &DeviceContext, - q: &HiddenStates, - k: &HiddenStates, - v: &HiddenStates, - kv_buffer: &CudaSlice, - layout: &KvLayout, - layer: usize, - plan: &PrefillPagedPlan, - positions_d: &CudaSlice, - output: &mut HiddenStates, - num_qo_heads: usize, - batch_size: usize, -) -> Result<()> { - pegainfer_kernels::ops::paged_attention_batch_decode_via_prefill_hd256_into( - ctx, - q, - k, - v, - kv_buffer, - &layout.kernel_layout(), - layer, - plan, - positions_d, - output, - num_qo_heads, - batch_size, - ) -} diff --git a/pegainfer-deepseek-v2-lite/src/model.rs b/pegainfer-deepseek-v2-lite/src/model.rs index 1ae10517d..272db86ff 100644 --- a/pegainfer-deepseek-v2-lite/src/model.rs +++ b/pegainfer-deepseek-v2-lite/src/model.rs @@ -377,6 +377,19 @@ fn routed_expert_from_slice<'a>( Ok(expert) } +pub(crate) fn routed_expert_projections<'a>( + layout: &ExpertParallelLayout, + mut expert: impl FnMut(usize) -> Result<&'a ExpertMlp>, +) -> Result> { + layout + .owned_experts() + .map(|global_expert| { + let expert = expert(global_expert)?; + Ok((&expert.dense.gate_up_proj, &expert.dense.down_proj)) + }) + .collect() +} + fn load_dense_mlp( ctx: &DeviceContext, shards: &[safetensors::SafeTensors<'_>], diff --git a/pegainfer-deepseek-v2-lite/src/nccl_backend.rs b/pegainfer-deepseek-v2-lite/src/nccl_backend.rs index 9290e4584..f896d5ffd 100644 --- a/pegainfer-deepseek-v2-lite/src/nccl_backend.rs +++ b/pegainfer-deepseek-v2-lite/src/nccl_backend.rs @@ -224,6 +224,16 @@ impl DenseExchangeOutput<'_> { } } +pub(crate) struct DeviceCombineOutput<'a> { + scratch: MutexGuard<'a, DeviceCombineScratch>, +} + +impl DeviceCombineOutput<'_> { + pub(crate) fn rank_send_mut(&mut self, rank: usize) -> Result<&mut CudaSlice> { + self.scratch.send_mut(rank) + } +} + #[derive(Default)] struct DeviceCombineScratch { hidden_dim: usize, @@ -1096,6 +1106,19 @@ impl NaiveNcclEp2Backend { op_result?; end_result } + + pub(crate) fn prepare_device_combine_output( + &self, + rank0: &DeviceContext, + rank1: &DeviceContext, + hidden_dim: usize, + seq_len: usize, + ) -> Result> { + let mut scratch = self.combine_scratch()?; + scratch.ensure(rank0, rank1, hidden_dim, seq_len)?; + scratch.ensure_shape(hidden_dim, seq_len)?; + Ok(DeviceCombineOutput { scratch }) + } } fn combine_elems(hidden_dim: usize, seq_len: usize) -> Result { diff --git a/pegainfer-deepseek-v2-lite/src/runtime.rs b/pegainfer-deepseek-v2-lite/src/runtime.rs index a76cde1c0..4e7cf5f6d 100644 --- a/pegainfer-deepseek-v2-lite/src/runtime.rs +++ b/pegainfer-deepseek-v2-lite/src/runtime.rs @@ -10,6 +10,8 @@ mod routing; mod tests; mod types; +use std::sync::Mutex; + use backend::EpBackendRuntime; pub use types::BatchedGenerationResult; pub use types::DecodeGraphReadinessReport; @@ -23,6 +25,7 @@ use crate::model::ExpertRankModel; pub struct DeepSeekV2LiteEp2Generator { device_ordinals: Vec, config: Config, + device_routed_moe: Option>, rank0: DriverRankModel, rank1: ExpertRankModel, backend: EpBackendRuntime, diff --git a/pegainfer-deepseek-v2-lite/src/runtime/generation.rs b/pegainfer-deepseek-v2-lite/src/runtime/generation.rs index 4248d6720..f4638d390 100644 --- a/pegainfer-deepseek-v2-lite/src/runtime/generation.rs +++ b/pegainfer-deepseek-v2-lite/src/runtime/generation.rs @@ -16,6 +16,7 @@ use super::helpers::append_generated_token; use super::helpers::duration_micros; use super::helpers::ensure_same_prompt_batch_rows_match; use super::helpers::token_sha256; +use super::moe::DeviceRoutedMoeRuntime; use super::types::BatchedGenerationResult; use super::types::GenerationResult; use super::types::GenerationStats; @@ -136,13 +137,21 @@ impl DeepSeekV2LiteEp2Generator { started.elapsed().as_secs_f64() ); - Ok(Self { + let mut generator = Self { device_ordinals: options.device_ordinals, config, rank0, rank1, backend, - }) + device_routed_moe: None, + }; + if generator.backend.kind() == super::backend::EpBackendKind::Nccl { + generator.device_routed_moe = Some(std::sync::Mutex::new( + DeviceRoutedMoeRuntime::new(&generator) + .context("initialize DeepSeek-V2-Lite device-routed MoE")?, + )); + } + Ok(generator) } pub fn generate_greedy( diff --git a/pegainfer-deepseek-v2-lite/src/runtime/layers.rs b/pegainfer-deepseek-v2-lite/src/runtime/layers.rs index a344c5d7c..a90114b8b 100644 --- a/pegainfer-deepseek-v2-lite/src/runtime/layers.rs +++ b/pegainfer-deepseek-v2-lite/src/runtime/layers.rs @@ -57,6 +57,7 @@ impl DeepSeekV2LiteEp2Generator { cache.layers.len() == self.rank0.layers.len(), "decode cache layer count mismatch" ); + let deferred_device_routes = self.begin_device_routed_forward(hidden.seq_len)?; for layer_idx in 0..self.rank0.layers.len() { hidden = self .forward_layer( @@ -71,6 +72,10 @@ impl DeepSeekV2LiteEp2Generator { ) .with_context(|| format!("DeepSeek-V2-Lite layer {layer_idx}"))?; } + if deferred_device_routes { + let (local_routes, remote_routes) = self.finish_device_routed_forward()?; + stats.record_routes(EpBackendKind::Nccl, local_routes, remote_routes); + } Ok(hidden) } @@ -95,6 +100,7 @@ impl DeepSeekV2LiteEp2Generator { .all(|cache| cache.layers.len() == self.rank0.layers.len()), "batched decode cache layer count mismatch" ); + let deferred_device_routes = self.begin_device_routed_forward(hidden.seq_len)?; for layer_idx in 0..self.rank0.layers.len() { let mut layer_caches: Vec<_> = caches .iter_mut() @@ -112,6 +118,10 @@ impl DeepSeekV2LiteEp2Generator { ) .with_context(|| format!("DeepSeek-V2-Lite batched layer {layer_idx}"))?; } + if deferred_device_routes { + let (local_routes, remote_routes) = self.finish_device_routed_forward()?; + stats.record_routes(EpBackendKind::Nccl, local_routes, remote_routes); + } Ok(hidden) } diff --git a/pegainfer-deepseek-v2-lite/src/runtime/moe.rs b/pegainfer-deepseek-v2-lite/src/runtime/moe.rs index 199041639..8bb642966 100644 --- a/pegainfer-deepseek-v2-lite/src/runtime/moe.rs +++ b/pegainfer-deepseek-v2-lite/src/runtime/moe.rs @@ -11,12 +11,16 @@ use half::bf16; use pegainfer_core::ops; use pegainfer_core::tensor::HiddenStates; use pegainfer_core::tensor::HiddenStatesRef; +use pegainfer_kernels::ops::DSV2_ROUTED_MOE_MAX_ROWS; +use pegainfer_kernels::ops::Dsv2ExpertPointerTable; use pegainfer_kernels::ops::Dsv2LiteRouterOutput; +use pegainfer_kernels::ops::Dsv2RoutedMoeScratch; use pegainfer_kernels::ops::dsv2_lite_router_logits_into; use pegainfer_kernels::ops::dsv2_lite_router_softmax_topk_into; use pegainfer_kernels::ops::dsv2_lite_router_softmax_topk_ref_into; use super::DeepSeekV2LiteEp2Generator; +use super::backend::EpBackendKind; use super::backend::EpBackendRuntime; use super::routing::MoeRouteEntry; use super::routing::MoeRoutePlan; @@ -35,6 +39,7 @@ use crate::model::dense_mlp_forward; use crate::model::dense_mlp_forward_per_token; use crate::model::dense_mlp_forward_preallocated_into; use crate::model::dense_mlp_forward_preallocated_ref_into; +use crate::model::routed_expert_projections; use crate::nccl_backend::NaiveNcclEp2Backend; fn parse_rollback_value( @@ -84,7 +89,6 @@ static NCCL_SERIAL: LazyLock> = LazyLock::new( }); static NCCL_HOST_ROUTER: LazyLock> = LazyLock::new(|| load_rollback_value("PEGAINFER_DSV2_LITE_NCCL_ROUTER", "device", "host")); - fn group_route_indices( keys: impl IntoIterator, ) -> BTreeMap<(usize, usize), Vec> { @@ -100,6 +104,105 @@ struct NcclRouteReplayBuffers { _outputs: Vec, } +pub(super) struct DeviceRoutedMoeRuntime { + rank0_tables: Vec>, + rank1_tables: Vec>, + rank0_scratch: Dsv2RoutedMoeScratch, + rank1_scratch: Dsv2RoutedMoeScratch, +} + +impl DeviceRoutedMoeRuntime { + pub(super) fn new(generator: &DeepSeekV2LiteEp2Generator) -> Result { + let mut rank0_tables = Vec::with_capacity(generator.config.num_hidden_layers); + let mut rank1_tables = Vec::with_capacity(generator.config.num_hidden_layers); + for layer_idx in 0..generator.config.num_hidden_layers { + if generator.config.is_moe_layer(layer_idx) { + activate(&generator.rank0.ctx)?; + rank0_tables.push(Some(Dsv2ExpertPointerTable::new( + &generator.rank0.ctx, + routed_expert_projections(&generator.rank0.layout, |expert| { + generator.rank0.routed_expert(layer_idx, expert) + })?, + generator.rank0.layout.owned_experts().start, + )?)); + activate(&generator.rank1.ctx)?; + rank1_tables.push(Some(Dsv2ExpertPointerTable::new( + &generator.rank1.ctx, + routed_expert_projections(&generator.rank1.layout, |expert| { + generator.rank1.routed_expert(layer_idx, expert) + })?, + generator.rank1.layout.owned_experts().start, + )?)); + } else { + rank0_tables.push(None); + rank1_tables.push(None); + } + } + activate(&generator.rank0.ctx)?; + let rank0_scratch = Dsv2RoutedMoeScratch::new( + &generator.rank0.ctx, + generator.config.hidden_size, + generator.config.moe_intermediate_size, + DSV2_ROUTED_MOE_MAX_ROWS, + generator.config.num_hidden_layers, + )?; + activate(&generator.rank1.ctx)?; + let rank1_scratch = Dsv2RoutedMoeScratch::new( + &generator.rank1.ctx, + generator.config.hidden_size, + generator.config.moe_intermediate_size, + DSV2_ROUTED_MOE_MAX_ROWS, + generator.config.num_hidden_layers, + )?; + Ok(Self { + rank0_tables, + rank1_tables, + rank0_scratch, + rank1_scratch, + }) + } + + fn begin_forward(&mut self, generator: &DeepSeekV2LiteEp2Generator) -> Result<()> { + activate(&generator.rank0.ctx)?; + self.rank0_scratch.begin_forward(&generator.rank0.ctx)?; + activate(&generator.rank1.ctx)?; + self.rank1_scratch.begin_forward(&generator.rank1.ctx) + } + + fn finish_forward(&self, generator: &DeepSeekV2LiteEp2Generator) -> Result<(usize, usize)> { + activate(&generator.rank0.ctx)?; + let rank0 = self.rank0_scratch.finish_forward(&generator.rank0.ctx)?; + activate(&generator.rank1.ctx)?; + let rank1 = self.rank1_scratch.finish_forward(&generator.rank1.ctx)?; + ensure!( + rank0.errors == 0 && rank1.errors == 0, + "device-routed MoE reported nonfinite router input: rank0={}, rank1={}", + rank0.errors, + rank1.errors + ); + ensure!( + rank0.total_routes == rank1.total_routes, + "device-routed MoE route totals differ across ranks: rank0={}, rank1={}", + rank0.total_routes, + rank1.total_routes + ); + ensure!( + rank0.layer_hashes == rank1.layer_hashes, + "device-routed MoE route IDs differ across ranks by layer: rank0={:?}, rank1={:?}", + rank0.layer_hashes, + rank1.layer_hashes + ); + ensure!( + rank0.local_routes + rank1.local_routes == rank0.total_routes, + "device-routed MoE ownership accounting drift: rank0_local={}, rank1_local={}, total={}", + rank0.local_routes, + rank1.local_routes, + rank0.total_routes + ); + Ok((rank0.local_routes, rank1.local_routes)) + } +} + fn accumulate_host_staged_route_output( dst: &mut [f32], route: &MoeRouteEntry, @@ -424,30 +527,35 @@ impl DeepSeekV2LiteEp2Generator { ) -> Result<(HiddenStates, usize, usize)> { activate(&self.rank0.ctx)?; let host_router = rollback_enabled(&NCCL_HOST_ROUTER)?; - let route_section = if host_router { - "ep_route_host" + let device_routed = self.device_routed_moe_enabled(input.seq_len)?; + let route_plan = if device_routed { + None } else { - "ep_route_device" + let route_section = if host_router { + "ep_route_host" + } else { + "ep_route_device" + }; + Some(attribution.record_result( + phase, + route_section, + || format!("layer.{layer_idx}.nccl.route"), + Some(layer_idx), + token_index, + || { + if host_router { + let input_host = hidden_to_bf16(&self.rank0.ctx, input)?; + let route_logits_host = + gate_logits_host(&self.config, &input_host, &moe.gate_host); + let routes = + topk_softmax_routes(&self.config, &route_logits_host, input.seq_len); + MoeRoutePlan::from_topk_routes(&routes, &self.rank0.layout) + } else { + self.build_nccl_route_plan_device(input, &moe.gate_device) + } + }, + )?) }; - let route_plan = attribution.record_result( - phase, - route_section, - || format!("layer.{layer_idx}.nccl.route"), - Some(layer_idx), - token_index, - || { - if host_router { - let input_host = hidden_to_bf16(&self.rank0.ctx, input)?; - let route_logits_host = - gate_logits_host(&self.config, &input_host, &moe.gate_host); - let routes = - topk_softmax_routes(&self.config, &route_logits_host, input.seq_len); - MoeRoutePlan::from_topk_routes(&routes, &self.rank0.layout) - } else { - self.build_nccl_route_plan_device(input, &moe.gate_device) - } - }, - )?; let shared = attribution.record_gpu_result( &self.rank0.ctx, @@ -475,52 +583,126 @@ impl DeepSeekV2LiteEp2Generator { || nccl.dense_all_reduce_rank0_hidden_to_rank1(&self.rank0.ctx, &self.rank1.ctx, input), )?; let rank1_hidden = rank1_input.rank1_hidden()?; - attribution.record_gpu_pair_result( - &self.rank0.ctx, - &self.rank1.ctx, - phase, - "nccl_combine_clear", - || format!("layer.{layer_idx}.nccl.combine_clear"), - Some(layer_idx), - token_index, - || { - nccl.clear_device_combine( - &self.rank0.ctx, - &self.rank1.ctx, - input.hidden_dim, - input.seq_len, - ) - }, - )?; - let live_expert_outputs = self.replay_nccl_route_plan( - nccl, - layer_idx, - input, - rank1_hidden, - &route_plan, - attribution, - phase, - token_index, - )?; - - let routed = attribution.record_gpu_pair_result( - &self.rank0.ctx, - &self.rank1.ctx, - phase, - "nccl_combine", - || format!("layer.{layer_idx}.nccl.combine"), - Some(layer_idx), - token_index, - || { - nccl.combine_device_contributions_to_rank0( - &self.rank0.ctx, - &self.rank1.ctx, - input.hidden_dim, - input.seq_len, - ) - }, - )?; - drop(live_expert_outputs); + let (routed, local_routes, remote_routes) = if device_routed { + attribution.record_gpu_pair_result( + &self.rank0.ctx, + &self.rank1.ctx, + phase, + "nccl_device_routed_experts", + || format!("layer.{layer_idx}.nccl.device_routed_experts"), + Some(layer_idx), + token_index, + || { + let runtime = self + .device_routed_moe + .as_ref() + .context("NCCL device-routed MoE runtime is not initialized")?; + let mut runtime = runtime + .lock() + .map_err(|_| anyhow::anyhow!("device-routed MoE mutex poisoned"))?; + let DeviceRoutedMoeRuntime { + rank0_tables, + rank1_tables, + rank0_scratch, + rank1_scratch, + } = &mut *runtime; + let rank0_table = rank0_tables + .get(layer_idx) + .and_then(Option::as_ref) + .context("rank0 device-routed MoE pointer table missing")?; + let rank1_table = rank1_tables + .get(layer_idx) + .and_then(Option::as_ref) + .context("rank1 device-routed MoE pointer table missing")?; + let mut contributions = nccl.prepare_device_combine_output( + &self.rank0.ctx, + &self.rank1.ctx, + input.hidden_dim, + input.seq_len, + )?; + activate(&self.rank0.ctx)?; + rank0_scratch.enqueue_into( + &self.rank0.ctx, + input.as_ref(), + &moe.gate_device, + rank0_table, + contributions.rank_send_mut(0)?, + layer_idx, + )?; + activate(&self.rank1.ctx)?; + rank1_scratch.enqueue_into( + &self.rank1.ctx, + rank1_hidden, + self.rank1.gate_device(layer_idx)?, + rank1_table, + contributions.rank_send_mut(1)?, + layer_idx, + )?; + drop(contributions); + let routed = nccl.combine_device_contributions_to_rank0( + &self.rank0.ctx, + &self.rank1.ctx, + input.hidden_dim, + input.seq_len, + )?; + Ok((routed, 0, 0)) + }, + )? + } else { + let route_plan = route_plan + .as_ref() + .context("NCCL fallback route plan is missing")?; + attribution.record_gpu_pair_result( + &self.rank0.ctx, + &self.rank1.ctx, + phase, + "nccl_combine_clear", + || format!("layer.{layer_idx}.nccl.combine_clear"), + Some(layer_idx), + token_index, + || { + nccl.clear_device_combine( + &self.rank0.ctx, + &self.rank1.ctx, + input.hidden_dim, + input.seq_len, + ) + }, + )?; + let live_expert_outputs = self.replay_nccl_route_plan( + nccl, + layer_idx, + input, + rank1_hidden, + route_plan, + attribution, + phase, + token_index, + )?; + let routed = attribution.record_gpu_pair_result( + &self.rank0.ctx, + &self.rank1.ctx, + phase, + "nccl_combine", + || format!("layer.{layer_idx}.nccl.combine"), + Some(layer_idx), + token_index, + || { + nccl.combine_device_contributions_to_rank0( + &self.rank0.ctx, + &self.rank1.ctx, + input.hidden_dim, + input.seq_len, + ) + }, + )?; + drop(live_expert_outputs); + ( + routed, + route_plan.local_routes(), + route_plan.remote_routes(), + ) + }; activate(&self.rank0.ctx)?; let hidden = attribution.record_gpu_result( &self.rank0.ctx, @@ -531,11 +713,7 @@ impl DeepSeekV2LiteEp2Generator { token_index, || ops::add_batch(&self.rank0.ctx, &routed, &shared), )?; - Ok(( - hidden, - route_plan.local_routes(), - route_plan.remote_routes(), - )) + Ok((hidden, local_routes, remote_routes)) } fn build_nccl_route_plan_device( @@ -549,7 +727,12 @@ impl DeepSeekV2LiteEp2Generator { .checked_mul(self.config.n_routed_experts) .context("NCCL device router logits element count overflow")?; let mut route_logits = self.rank0.ctx.stream.alloc_zeros::(logits_elems)?; - dsv2_lite_router_logits_into(&self.rank0.ctx, input, gate_device, &mut route_logits)?; + dsv2_lite_router_logits_into( + &self.rank0.ctx, + input.as_ref(), + gate_device, + &mut route_logits, + )?; let route_logits = self.rank0.ctx.stream.clone_dtoh(&route_logits)?; self.rank0.ctx.sync()?; let routes = topk_softmax_routes(&self.config, &route_logits, input.seq_len); @@ -966,6 +1149,39 @@ impl DeepSeekV2LiteEp2Generator { ), } } + + fn device_routed_moe_enabled(&self, rows: usize) -> Result { + Ok(self.backend.kind() == EpBackendKind::Nccl + && rows <= DSV2_ROUTED_MOE_MAX_ROWS + && !rollback_enabled(&NCCL_HOST_ROUTER)? + && !rollback_enabled(&NCCL_SERIAL)?) + } + + pub(super) fn begin_device_routed_forward(&self, rows: usize) -> Result { + if !self.device_routed_moe_enabled(rows)? { + return Ok(false); + } + let runtime = self + .device_routed_moe + .as_ref() + .context("NCCL device-routed MoE runtime is not initialized")?; + runtime + .lock() + .map_err(|_| anyhow::anyhow!("device-routed MoE mutex poisoned"))? + .begin_forward(self)?; + Ok(true) + } + + pub(super) fn finish_device_routed_forward(&self) -> Result<(usize, usize)> { + let runtime = self + .device_routed_moe + .as_ref() + .context("NCCL device-routed MoE runtime is not initialized")?; + runtime + .lock() + .map_err(|_| anyhow::anyhow!("device-routed MoE mutex poisoned"))? + .finish_forward(self) + } } fn expert_forward_device( diff --git a/pegainfer-deepseek-v2-lite/src/runtime/moe/tests.rs b/pegainfer-deepseek-v2-lite/src/runtime/moe/tests.rs index 0c257315f..71ade4192 100644 --- a/pegainfer-deepseek-v2-lite/src/runtime/moe/tests.rs +++ b/pegainfer-deepseek-v2-lite/src/runtime/moe/tests.rs @@ -120,7 +120,7 @@ fn device_router_logits_match_host_accumulation_bitwise() { .alloc_zeros::(hidden.seq_len * config.n_routed_experts) .expect("router logits"); - dsv2_lite_router_logits_into(&ctx, &hidden, &gate, &mut logits).expect("device logits"); + dsv2_lite_router_logits_into(&ctx, hidden.as_ref(), &gate, &mut logits).expect("device logits"); let got = ctx.stream.clone_dtoh(&logits).expect("logits D2H"); ctx.sync().expect("sync logits"); @@ -174,7 +174,7 @@ fn router_logits_rejects_inflated_hidden_metadata() { .alloc_zeros::(hidden.seq_len * gate.rows) .expect("logits"); - let err = dsv2_lite_router_logits_into(&ctx, &hidden, &gate, &mut logits) + let err = dsv2_lite_router_logits_into(&ctx, hidden.as_ref(), &gate, &mut logits) .expect_err("inflated hidden metadata must fail before CUDA launch"); assert_error_contains(&err, "router hidden backing buffer too small"); } @@ -203,7 +203,7 @@ fn router_logits_rejects_inflated_gate_metadata() { .alloc_zeros::(hidden.seq_len * gate.rows) .expect("logits"); - let err = dsv2_lite_router_logits_into(&ctx, &hidden, &gate, &mut logits) + let err = dsv2_lite_router_logits_into(&ctx, hidden.as_ref(), &gate, &mut logits) .expect_err("inflated gate metadata must fail before CUDA launch"); assert_error_contains(&err, "router gate backing buffer too small"); } @@ -670,3 +670,70 @@ fn assert_error_contains(err: &anyhow::Error, needle: &str) { "expected error containing {needle:?}, got {message:?}" ); } + +#[test] +fn device_routed_router_matches_host_at_supported_boundaries() { + let ctx = DeviceContext::new().expect("create CUDA context"); + let mut config = test_lite_config(); + config.n_routed_experts = 64; + config.num_experts_per_token = 6; + for batch in [1, 4, 8] { + for fixture in 0..3 { + let logits: Vec<_> = (0..batch) + .flat_map(|token| { + (0..config.n_routed_experts).map(move |expert| match fixture { + 0 => ((expert * 37 + token * 13) % 101) as f32 / 17.0 - 3.0, + 1 => 0.0, + _ if expert < 5 => 10.0 - expert as f32, + _ if expert == 5 => 1.0, + _ if expert == 6 => 1.0 + 1.0e-5, + _ => -(expert as f32), + }) + }) + .collect(); + let expected = topk_softmax_routes(&config, &logits, batch); + let logits = ctx.stream.clone_htod(&logits).expect("logits H2D"); + let mut ids = ctx + .stream + .alloc_zeros::(batch * config.num_experts_per_token) + .unwrap(); + let mut weights = ctx + .stream + .alloc_zeros::(batch * config.num_experts_per_token) + .unwrap(); + let mut errors = ctx.stream.alloc_zeros::(batch).unwrap(); + let mut summary = ctx.stream.alloc_zeros::(5).unwrap(); + pegainfer_kernels::ops::dsv2_lite_route_logits_into( + &ctx, + &logits, + batch, + 1, + &mut Dsv2LiteRouterOutput { + topk_idx: &mut ids, + topk_weight: &mut weights, + }, + &mut errors, + &mut summary, + ) + .expect("device routed router"); + let ids = ctx.stream.clone_dtoh(&ids).unwrap(); + let weights = ctx.stream.clone_dtoh(&weights).unwrap(); + let errors = ctx.stream.clone_dtoh(&errors).unwrap(); + let summary = ctx.stream.clone_dtoh(&summary).unwrap(); + ctx.sync().unwrap(); + for (token, routes) in expected.iter().enumerate() { + for (route, &(expert, weight)) in routes.iter().enumerate() { + let slot = token * config.num_experts_per_token + route; + assert_eq!( + ids[slot], expert as i32, + "batch={batch}, fixture={fixture}, slot={slot}" + ); + assert_close(weights[slot], weight, 1.0e-6); + } + } + assert!(errors.iter().all(|&error| error == 0)); + assert_eq!(summary[3], 0); + assert_ne!(summary[4], 0); + } + } +} diff --git a/pegainfer-frontend/src/engine/driver.rs b/pegainfer-frontend/src/engine/driver.rs index d42313bc0..3d4ec9ab4 100644 --- a/pegainfer-frontend/src/engine/driver.rs +++ b/pegainfer-frontend/src/engine/driver.rs @@ -80,15 +80,15 @@ pub fn drive(mut scheduler: S, backend: SchedulerBackend) { } let step = scheduler.step(&mut ledger); let snapshot = scheduler.metrics(); - metrics.publish(&snapshot); - ledger.commit_step(); + let metrics_changed = metrics.publish(&snapshot); + ledger.commit_step(metrics_changed); if let Err(error) = step { // The ledger holds an account for every unanswered request, so // the write-off reaches them all — including any the scheduler // lost track of — with the real error attached. log::error!("scheduler fatal, engine winding down: {error:#}"); ledger.fail_all(&format!("engine fatal: {error:#}")); - ledger.commit_step(); + ledger.commit_step(false); return; } if snapshot.num_running_reqs == 0 && snapshot.num_waiting_reqs == 0 { @@ -212,6 +212,96 @@ mod tests { partition.join.join().expect("driver thread exits cleanly"); } + #[test] + fn cancellation_notifies_metrics_once_without_request_output() { + struct PausedScheduler { + queued: Vec, + running: Vec, + entered: std::sync::mpsc::Sender<()>, + resume: std::sync::mpsc::Receiver<()>, + } + + impl Scheduler for PausedScheduler { + fn submit(&mut self, request: QueuedRequest) { + self.queued.push(request.id); + } + + fn step(&mut self, ledger: &mut RequestLedger) -> anyhow::Result<()> { + let _ = self.entered.send(()); + let _ = self.resume.recv(); + for id in self.queued.drain(..) { + ledger.admit(id); + self.running.push(id); + } + self.running.retain(|&id| { + if ledger.is_aborted(id) { + ledger.retire(id); + false + } else { + true + } + }); + Ok(()) + } + + fn metrics(&self) -> SchedulerMetrics { + SchedulerMetrics { + num_running_reqs: self.running.len() as u64, + ..SchedulerMetrics::default() + } + } + } + + let (mut handle, backend) = scheduler_pair(); + let mut steps = handle.take_steps().unwrap(); + let control = handle.submit(request(8)); + let (entered_tx, entered_rx) = std::sync::mpsc::channel(); + let (resume_tx, resume_rx) = std::sync::mpsc::channel(); + let join = std::thread::spawn(move || { + drive( + PausedScheduler { + queued: Vec::new(), + running: Vec::new(), + entered: entered_tx, + resume: resume_rx, + }, + backend, + ); + }); + let next_step = || { + entered_rx + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("driver advanced to its next step"); + }; + + next_step(); + resume_tx.send(()).unwrap(); + next_step(); + steps.try_recv().expect("admission step"); + + control.abort(); + resume_tx.send(()).unwrap(); + next_step(); + assert!( + steps + .try_recv() + .expect("cancellation load change") + .updates + .is_empty() + ); + assert_eq!(handle.metrics().num_running_reqs, 0); + + resume_tx.send(()).unwrap(); + next_step(); + assert!(matches!( + steps.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + drop(handle); + drop(resume_tx); + join.join().expect("driver drains after cancellation"); + } + #[test] fn fatal_step_fails_in_flight_requests_with_the_error() { let partition = spawn_scheduler( diff --git a/pegainfer-frontend/src/engine/ledger.rs b/pegainfer-frontend/src/engine/ledger.rs index d3bed79e3..c0ab40f66 100644 --- a/pegainfer-frontend/src/engine/ledger.rs +++ b/pegainfer-frontend/src/engine/ledger.rs @@ -339,12 +339,10 @@ impl RequestLedger { } } - /// Ship the step's statement as one message; a step that touched nothing - /// ships nothing. Called once per driver iteration — model code never - /// calls this (the driver owns the cadence). - pub(crate) fn commit_step(&mut self) { + /// Commit the buffered updates as described by [`StepOutputs`]. + pub(crate) fn commit_step(&mut self, metrics_changed: bool) { let updates = self.statement.take_updates(); - if updates.is_empty() { + if updates.is_empty() && !metrics_changed { return; } // A closed receiver means the frontend is gone; the driver notices @@ -430,7 +428,7 @@ mod tests { backend .ledger .finish_with_cause(id, FinishReason::Stop, Some(StopCause::Token(11))); - backend.ledger.commit_step(); + backend.ledger.commit_step(false); let mut steps = handle_steps(handle); let step = steps.try_recv().expect("one step message"); @@ -468,7 +466,7 @@ mod tests { limit: 4, }, ); - backend.ledger.commit_step(); + backend.ledger.commit_step(false); let step = handle_steps(handle).try_recv().expect("step"); let update = &step.updates[0]; @@ -492,7 +490,7 @@ mod tests { backend.ledger.admit(id); backend.ledger.push_tokens(id, &[9], &[]); backend.ledger.retire(id); - backend.ledger.commit_step(); + backend.ledger.commit_step(false); // Scheduled was buffered before the retire extracted the entry, so // nothing observable remains this step. @@ -510,7 +508,7 @@ mod tests { let deferred = backend .ledger .defer_finish_with_cause(id, FinishReason::Length, None); - backend.ledger.commit_step(); + backend.ledger.commit_step(false); let mut steps = handle_steps(handle); // The request's whole record rode into the deferred finish; the step diff --git a/pegainfer-frontend/src/engine/step.rs b/pegainfer-frontend/src/engine/step.rs index 07670d367..e0855bd0d 100644 --- a/pegainfer-frontend/src/engine/step.rs +++ b/pegainfer-frontend/src/engine/step.rs @@ -78,9 +78,9 @@ pub struct QueuedRequest { pub request: Request, } -/// Everything one scheduler step produced, in one message. The scheduler-side -/// ledger sends exactly one per step that touched any request; an idle step -/// sends nothing. +/// Everything one scheduler step produced, in one message. Empty `updates` +/// notify the frontend of changed metrics, for example after cancellation. +/// A step with no request updates or metric changes sends nothing. #[derive(Debug, Default)] pub struct StepOutputs { pub updates: Vec, diff --git a/pegainfer-frontend/src/engine/wiring.rs b/pegainfer-frontend/src/engine/wiring.rs index 5067a1c6a..8cc7431b5 100644 --- a/pegainfer-frontend/src/engine/wiring.rs +++ b/pegainfer-frontend/src/engine/wiring.rs @@ -8,8 +8,8 @@ //! crossbeam (sync consumer on the scheduler thread; senders never block on //! unbounded channels), the step stream is tokio (async consumer in the //! protocol stack; the sync producer's send never blocks either), metrics are -//! a shared cell (read-only pull, deliberately unsubscribable — see -//! [`MetricsPublisher`]). +//! a shared cell, with changes signaled on the step stream (see +//! [`super::StepOutputs`]). //! //! How many schedulers an engine runs and what each one means (DP replicas, //! anything else) is the model line's decision; the contract carries the @@ -51,18 +51,16 @@ pub struct SchedulerBackend { /// Sole writer of a scheduler's metrics cell; the driver publishes once per /// iteration from [`super::Scheduler::metrics`]. /// -/// Deliberately a plain cell and not a `watch` channel: the driver busy-polls, -/// so a subscription edge (`changed()`) would fire per spin and turn any -/// subscriber into a message flood at idle. With only -/// [`SchedulerHandle::metrics`] to read it, "notify me on metrics change" is -/// unrepresentable — consumers pull the snapshot at the moment they need one. /// A `Mutex` (not per-field atomics) so a reader never sees fields torn /// across two steps; both sides touch it uncontended for nanoseconds. pub struct MetricsPublisher(Arc>); impl MetricsPublisher { - pub(crate) fn publish(&self, snapshot: &SchedulerMetrics) { - *self.0.lock().expect("metrics cell poisoned") = *snapshot; + pub(crate) fn publish(&self, snapshot: &SchedulerMetrics) -> bool { + let mut current = self.0.lock().expect("metrics cell poisoned"); + let changed = *current != *snapshot; + *current = *snapshot; + changed } } @@ -102,9 +100,8 @@ impl SchedulerHandle { self.steps.take() } - /// The scheduler's most recent metrics snapshot. Pull-only by design (see - /// [`MetricsPublisher`]): read it at the moment you need one — routing a - /// request, stamping stats onto an outgoing batch, serving a scrape. + /// The scheduler's most recent metrics snapshot. Read it when routing a + /// request, receiving a step, or serving a scrape. pub fn metrics(&self) -> SchedulerMetrics { *self.metrics.lock().expect("metrics cell poisoned") } diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 9500e2690..70a42ef31 100644 --- a/pegainfer-frontend/src/vllm/bridge/stepped.rs +++ b/pegainfer-frontend/src/vllm/bridge/stepped.rs @@ -87,9 +87,6 @@ impl SteppedEngineBridge { .take_steps() .context("partition step stream already taken")?; let mut spec = SpecDecodeTracker::default(); - // Stats are pull-at-send: no push task, the load cell is read when a - // batch goes out (and once here, so the frontend's gauges initialize - // before any traffic). An idle engine publishes nothing. let BridgeLink { mut input, output_tx, @@ -254,28 +251,6 @@ impl SteppedEngineBridge { } } - if outputs.is_empty() { - // A drafted step with no batch to ride would strand its increment - // until the next batch, which may never come. - let stats = self.stats(spec); - if stats.spec_decoding_stats.is_some() { - send_outputs( - output_tx, - RequestBatchOutputs { - engine_index: self.engine_index, - scheduler_stats: Some(Box::new(stats)), - timestamp: now_secs_f64(), - ..Default::default() - } - .into(), - )?; - } - return Ok(()); - } - // The cell already holds this step's snapshot (the driver publishes - // load before committing the step), so the batch carries stats that - // match its own tokens — a finishing batch reports the drained state - // and the gauges settle instead of freezing at the last busy value. send_outputs( output_tx, RequestBatchOutputs { @@ -467,7 +442,7 @@ struct SteppedStream { request_id: String, control: RequestControl, /// Queued/Scheduled wire events, held until the request's first shipped - /// output (a scheduled-only update ships nothing on its own). + /// output (a scheduled-only update produces no request output). first_token_events: Option>, /// Set by `Scheduled`; a request refused or failed while still queued /// never prefilled and reports no prefill stats. @@ -687,6 +662,7 @@ mod tests { use super::*; use crate::engine::PromptEcho; use crate::engine::RejectReason; + use crate::engine::ScheduledInfo; use crate::engine::StopPolicy; use crate::engine::TokenLogprob; use crate::engine::scheduler_pair; @@ -747,6 +723,49 @@ mod tests { } } + #[test] + fn scheduled_and_metrics_only_steps_publish_load_without_request_outputs() { + let (handle, backend) = scheduler_pair(); + let control = handle.submit(request()); + let mut scheduled = RequestUpdate::empty(control.id()); + let now = Instant::now(); + scheduled.scheduled = Some(ScheduledInfo { + queued_at: now, + scheduled_at: now, + prompt_tokens: 2, + }); + let mut streams = HashMap::from([( + control.id(), + SteppedStream::new("scheduled".into(), control, Span::noop(), None), + )]); + let bridge = bridge(handle); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let mut spec = SpecDecodeTracker::default(); + for (running, updates) in [(1, vec![scheduled]), (0, Vec::new())] { + backend.metrics.publish(&crate::engine::SchedulerMetrics { + num_running_reqs: running, + ..Default::default() + }); + bridge + .dispatch_step( + StepOutputs { updates }, + &UnixAnchor::now(), + &mut streams, + &mut HashMap::new(), + &mut spec, + &tx, + ) + .unwrap(); + let EngineCoreOutputs::RequestBatch(batch) = rx.try_recv().expect("load batch") else { + panic!("expected scheduler stats"); + }; + assert!(batch.outputs.is_empty()); + let stats = batch.scheduler_stats.expect("scheduler stats"); + assert_eq!(stats.num_running_reqs, running); + } + assert!(rx.try_recv().is_err()); + } + fn wire_request(completion: Option, prompt: Option) -> EngineCoreRequest { let mut params = EngineCoreSamplingParams::for_test(); params.logprobs = completion; diff --git a/pegainfer-gemma4/src/config.rs b/pegainfer-gemma4/src/config.rs index bbcba61f7..068457a41 100644 --- a/pegainfer-gemma4/src/config.rs +++ b/pegainfer-gemma4/src/config.rs @@ -1,9 +1,11 @@ //! Typed text-tower config, read only from a file the probe has accepted. -#[cfg(feature = "gemma4")] +#[cfg(any(feature = "gemma4", test))] use anyhow::Result; -#[cfg(feature = "gemma4")] +#[cfg(any(feature = "gemma4", test))] use anyhow::bail; +#[cfg(any(feature = "gemma4", test))] +use anyhow::ensure; // Only the loader reads a config off disk. #[cfg(feature = "gemma4")] @@ -85,6 +87,175 @@ impl MoeConfig { } } +/// Which rank of a tensor-parallel group this process serves, and how wide the +/// group is. The default is the whole model on one rank. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct TensorParallelConfig { + pub(crate) rank: usize, + pub(crate) world_size: usize, +} + +impl Default for TensorParallelConfig { + fn default() -> Self { + Self::new(0, 1) + } +} + +/// The head width a layer kind's queries and keys use; the two families differ +/// only here and in their KV head counts. +pub(crate) fn head_dim_of(config: &Gemma4Config, kind: LayerKind) -> usize { + match kind { + LayerKind::Sliding => config.head_dim, + LayerKind::Global => config.global_head_dim, + } +} + +impl TensorParallelConfig { + /// The single-rank configuration the tests build geometries at; serving + /// takes it from [`Default`]. + #[cfg(test)] + pub(crate) const SINGLE: Self = Self { + rank: 0, + world_size: 1, + }; + + pub(crate) fn new(rank: usize, world_size: usize) -> Self { + Self { rank, world_size } + } + + pub(crate) fn is_single(&self) -> bool { + self.world_size == 1 + } + + pub(crate) fn local_q_heads(&self, config: &Gemma4Config) -> Result { + let q = config.num_attention_heads; + ensure!(q > 0, "query head count must be positive"); + ensure!( + q.is_multiple_of(self.world_size), + "{q} query heads do not divide over world size {}", + self.world_size + ); + Ok(q / self.world_size) + } + + pub(crate) fn local_sliding_kv_heads(&self, config: &Gemma4Config) -> Result { + let kv = config.num_key_value_heads; + ensure!(kv > 0, "sliding KV head count must be positive"); + ensure!( + kv.is_multiple_of(self.world_size), + "{kv} sliding KV heads do not divide over world size {}", + self.world_size + ); + Ok(kv / self.world_size) + } + + /// The global family's KV head range for this rank as `(head_offset, + /// head_count)`. That family has few enough KV heads that a world size can + /// exceed them: `G % P == 0` shards contiguous runs of `G / P` heads, and + /// `P % G == 0` replicates each head onto `P / G` contiguous ranks, one + /// head apiece. Replicating KV is exact under a sum reduction only because + /// query rows and `o_proj` columns still split `1 / P` — this branch never + /// touches them. + pub(crate) fn global_kv_head_range(&self, config: &Gemma4Config) -> Result<(usize, usize)> { + let g = config.num_global_key_value_heads; + let p = self.world_size; + ensure!( + g > 0 && p > 0, + "global KV heads ({g}) and world size ({p}) must be positive" + ); + if g.is_multiple_of(p) { + Ok((self.rank * (g / p), g / p)) + } else if p.is_multiple_of(g) { + Ok((self.rank / (p / g), 1)) + } else { + bail!( + "world size {p} is legal neither for sharding nor for replicating {g} global KV \ + heads" + ) + } + } + + pub(crate) fn local_q_dim(&self, config: &Gemma4Config, kind: LayerKind) -> Result { + Ok(self.local_q_heads(config)? * head_dim_of(config, kind)) + } + + /// The fused K (and, on sliding layers, V) width this rank holds. Global + /// layers carry no `v_proj`, so the sliding value covers both their K and + /// V rows. + pub(crate) fn local_kv_dim(&self, config: &Gemma4Config, kind: LayerKind) -> Result { + Ok(match kind { + LayerKind::Sliding => self.local_sliding_kv_heads(config)? * config.head_dim, + LayerKind::Global => self.global_kv_head_range(config)?.1 * config.global_head_dim, + }) + } + + pub(crate) fn local_intermediate(&self, config: &Gemma4Config) -> Result { + let i = config.intermediate_size; + ensure!(i > 0, "intermediate size must be positive"); + ensure!( + i.is_multiple_of(self.world_size), + "intermediate size {i} does not divide over world size {}", + self.world_size + ); + Ok(i / self.world_size) + } + + /// Everything a tensor-parallel launch must satisfy before a multi-GiB + /// load. The three-count rule the design doc states (`Q % P`, `Kv % P`, + /// `G % P || P % G`) is necessary but not sufficient: it does not keep the + /// per-rank GQA group integral (`Q = 8, Kv = 6, P = 2` clears it and yields + /// group `4/3`), so that is checked here too. + pub(crate) fn validate_for(&self, config: &Gemma4Config) -> Result<()> { + // A single rank is the incumbent path: it has no shard to police, and + // the MoE/W4A16 refusals below are about sharding those families, not + // about those checkpoints. + if self.is_single() { + return Ok(()); + } + let p = self.world_size; + ensure!(p > 0, "tensor-parallel world size must be positive"); + ensure!( + self.rank < p, + "tensor-parallel rank {} is outside world size {p}", + self.rank + ); + ensure!( + config.moe.is_none(), + "gemma4 tensor parallelism does not shard the routed experts; this checkpoint routes" + ); + ensure!( + !config.w4a16, + "gemma4 tensor parallelism does not shard the W4A16 GEMMs, which project whole matrices" + ); + let local_q = self.local_q_heads(config)?; + let local_sliding_kv = self.local_sliding_kv_heads(config)?; + ensure!( + local_q.is_multiple_of(local_sliding_kv), + "per-rank sliding GQA group {local_q}/{local_sliding_kv} is not integral at world \ + size {p}" + ); + // Resolving the range also refuses a world size that can neither shard + // nor replicate the global heads. + let (_, local_global_kv) = self.global_kv_head_range(config)?; + // Only the sharding branch can split a GQA group: it hands a rank + // `g / p` KV heads, so the per-rank group `(q / p) / (g / p)` has to be + // integral. The replicate branch gives a rank one whole KV head, shared + // with the `p / g` ranks holding that same head, and those ranks split + // its `q / g` query heads evenly (`q / p` is `(q / g) / (p / g)`), so + // there is nothing left to police there — `local_global_kv` is 1 on that + // branch, and the check below would divide by one. + if config.num_global_key_value_heads.is_multiple_of(p) { + ensure!( + local_q.is_multiple_of(local_global_kv), + "per-rank global GQA group {local_q}/{local_global_kv} is not integral at world \ + size {p}" + ); + } + self.local_intermediate(config)?; + Ok(()) + } +} + #[cfg(feature = "gemma4")] impl Gemma4Config { pub(crate) fn from_file(model_path: &str) -> Result { @@ -322,3 +493,183 @@ fn rope_type_field(rope_group: &serde_json::Value, ctx: &str, implemented: &str) ); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + /// A text geometry with the four counts a shard decision reads varied; + /// the rest is fixed filler. Built directly, so it is not probe-backed — + /// only the shard arithmetic is under test here. + fn config(q: usize, kv: usize, g: usize, intermediate: usize) -> Gemma4Config { + Gemma4Config { + hidden_size: 2560, + intermediate_size: intermediate, + vocab_size: 262_144, + num_attention_heads: q, + num_key_value_heads: kv, + num_global_key_value_heads: g, + head_dim: 256, + global_head_dim: 512, + layer_types: vec![LayerKind::Sliding; 12], + tie_word_embeddings: true, + moe: None, + w4a16: false, + rms_norm_eps: 1e-6, + sliding_rope_theta: 10_000.0, + sliding_window: 1024, + global_rope_theta: 1_000_000.0, + global_rotary_dim: 256, + final_logit_softcapping: 30.0, + max_position_embeddings: 262_144, + } + } + + /// 12B-like: 16 query heads, 8 sliding KV heads, the single global KV head + /// the design doc's table names. + fn config_12b() -> Gemma4Config { + config(16, 8, 1, 15_360) + } + + /// 31B-like: 32 query heads, 16 sliding KV heads, 4 global KV heads. + fn config_31b() -> Gemma4Config { + config(32, 16, 4, 16_384) + } + + #[test] + fn single_is_the_identity() { + let tp = TensorParallelConfig::SINGLE; + let cfg = config_31b(); + assert!(tp.is_single()); + assert_eq!(tp.local_q_heads(&cfg).unwrap(), 32); + assert_eq!(tp.local_sliding_kv_heads(&cfg).unwrap(), 16); + assert_eq!(tp.global_kv_head_range(&cfg).unwrap(), (0, 4)); + assert_eq!(tp.local_q_dim(&cfg, LayerKind::Sliding).unwrap(), 32 * 256); + assert_eq!(tp.local_q_dim(&cfg, LayerKind::Global).unwrap(), 32 * 512); + assert_eq!(tp.local_kv_dim(&cfg, LayerKind::Global).unwrap(), 4 * 512); + assert_eq!(tp.local_intermediate(&cfg).unwrap(), 16_384); + tp.validate_for(&cfg).unwrap(); + } + + #[test] + fn twelve_b_replicates_its_single_global_kv_head() { + let cfg = config_12b(); + for rank in 0..2 { + let tp = TensorParallelConfig::new(rank, 2); + assert_eq!(tp.local_q_heads(&cfg).unwrap(), 8); + assert_eq!(tp.local_sliding_kv_heads(&cfg).unwrap(), 4); + // G = 1 < P = 2: the one head lives on both ranks, at offset 0. + assert_eq!(tp.global_kv_head_range(&cfg).unwrap(), (0, 1)); + tp.validate_for(&cfg).unwrap(); + } + } + + #[test] + fn thirty_one_b_shards_two_global_kv_heads_per_rank() { + let cfg = config_31b(); + assert_eq!( + TensorParallelConfig::new(0, 2) + .global_kv_head_range(&cfg) + .unwrap(), + (0, 2) + ); + assert_eq!( + TensorParallelConfig::new(1, 2) + .global_kv_head_range(&cfg) + .unwrap(), + (2, 2) + ); + } + + #[test] + fn local_dims_shrink_by_the_world_size() { + let cfg = config_31b(); + let tp = TensorParallelConfig::new(1, 2); + assert_eq!(tp.local_q_dim(&cfg, LayerKind::Sliding).unwrap(), 16 * 256); + assert_eq!(tp.local_q_dim(&cfg, LayerKind::Global).unwrap(), 16 * 512); + assert_eq!(tp.local_kv_dim(&cfg, LayerKind::Sliding).unwrap(), 8 * 256); + assert_eq!(tp.local_kv_dim(&cfg, LayerKind::Global).unwrap(), 2 * 512); + assert_eq!(tp.local_intermediate(&cfg).unwrap(), 8192); + } + + #[test] + fn world_size_zero_is_rejected() { + let err = TensorParallelConfig::new(0, 0) + .validate_for(&config_31b()) + .unwrap_err() + .to_string(); + assert!(err.contains("positive"), "{err}"); + } + + #[test] + fn rank_outside_the_world_is_rejected() { + let err = TensorParallelConfig::new(2, 2) + .validate_for(&config_31b()) + .unwrap_err() + .to_string(); + assert!(err.contains("rank"), "{err}"); + } + + #[test] + fn query_heads_that_do_not_divide_are_rejected() { + // 12B has 16 query heads, which do not divide over world size 3. + let err = TensorParallelConfig::new(0, 3) + .validate_for(&config_12b()) + .unwrap_err() + .to_string(); + assert!(err.contains("query heads"), "{err}"); + } + + #[test] + fn per_rank_gqa_group_must_be_integral() { + // 8 query heads over 6 sliding KV heads at world size 2 clears + // `Q % P` and `Kv % P` yet leaves the per-rank group 4/3. + let err = TensorParallelConfig::new(0, 2) + .validate_for(&config(8, 6, 2, 15_360)) + .unwrap_err() + .to_string(); + assert!(err.contains("GQA group"), "{err}"); + } + + #[test] + fn global_heads_that_neither_shard_nor_replicate_are_rejected() { + // 3 global KV heads over world size 2: neither 3 % 2 nor 2 % 3 is 0. + let err = TensorParallelConfig::new(0, 2) + .validate_for(&config(8, 4, 3, 15_360)) + .unwrap_err() + .to_string(); + assert!(err.contains("global KV heads"), "{err}"); + } + + #[test] + fn intermediate_size_must_divide() { + let err = TensorParallelConfig::new(0, 2) + .validate_for(&config(16, 8, 1, 15_361)) + .unwrap_err() + .to_string(); + assert!(err.contains("intermediate"), "{err}"); + } + + #[test] + fn routed_and_w4a16_checkpoints_are_refused() { + let mut cfg = config_31b(); + cfg.moe = Some(MoeConfig { + num_experts: 128, + top_k: 8, + intermediate_size: 704, + }); + let err = TensorParallelConfig::new(0, 2) + .validate_for(&cfg) + .unwrap_err() + .to_string(); + assert!(err.contains("routed experts"), "{err}"); + + let mut cfg = config_31b(); + cfg.w4a16 = true; + let err = TensorParallelConfig::new(0, 2) + .validate_for(&cfg) + .unwrap_err() + .to_string(); + assert!(err.contains("W4A16"), "{err}"); + } +} diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 134343c69..9b60430bf 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -6,6 +6,7 @@ //! advances one token per batched decode step, sharing each layer's weight //! pass. +use std::cell::Cell; use std::collections::VecDeque; use std::path::Path; @@ -33,14 +34,18 @@ use pegainfer_frontend::engine::spawn_scheduler; use pegainfer_sample::LogprobRequest; use pegainfer_sample::SampleScratch; +use crate::config::TensorParallelConfig; use crate::forward::MULTIMODAL_PLACEHOLDER_IDS; use crate::kv::GLOBAL_PAGE_SIZE; use crate::kv::GemmaKv; use crate::kv::LOCAL_PAGE_SIZE; +use crate::kv::RankKv; use crate::kv::admit_tokens; +use crate::layer::LayerGeometry; use crate::prefix_cache::PrefixCache; use crate::serve::GemmaServe; use crate::serve::GlobalAttn; +use crate::serve::PrecapturePhase; use crate::serve::StepArena; use crate::weights::Gemma4Weights; @@ -125,6 +130,17 @@ fn parse_decode_slots(raw: &str) -> Result { } } +/// The NCCL build this process linked, as `major.minor.patch`. Start-up logs it +/// and the tensor-parallel gates print it, because NCCL picks the reduction +/// order and so a measured drift has to be attributed to the build it ran on +/// before it can be compared against another run's. +pub(crate) fn nccl_version() -> String { + cudarc::nccl::result::get_nccl_version().map_or_else( + |_| "unknown".to_string(), + |v| format!("{}.{}.{}", v / 10000, (v / 100) % 100, v % 100), + ) +} + /// GEMM and attention tiles consume whole 128-row blocks, so a width that is /// not a multiple of 128 pays for a tile it does not fill. fn parse_mix_chunk_tokens(raw: &str, max_context: usize) -> Result> { @@ -335,28 +351,88 @@ fn start_with_knobs( .to_str() .context("model path is not valid UTF-8")? .to_string(); + let ordinals = options.device_ordinals.clone(); anyhow::ensure!( - options.device_ordinals.len() == 1, - "gemma4 is single-device; got device_ordinals {:?}", - options.device_ordinals + !ordinals.is_empty(), + "gemma4 needs at least one device ordinal" ); - anyhow::ensure!( - options.parallel_config.is_none(), - "gemma4 has no parallel topology support yet" - ); - let device = options.device_ordinals[0]; + let world = ordinals.len(); + if let Some(parallel) = options.parallel_config.as_ref() { + anyhow::ensure!( + parallel.tp_world() == world + && parallel.dp_world() == 1 + && parallel.ep_world() == world, + "gemma4 tensor parallelism takes one rank per device (tp_size = ep_size = {world}, \ + dp_size = 1); got tp {} dp {} ep {}", + parallel.tp_world(), + parallel.dp_world(), + parallel.ep_world() + ); + } let base_seed = options.seed; let graph_enabled = options.enable_cuda_graph; let config = crate::config::Gemma4Config::from_file(&dir)?; let knobs = ServingKnobs::resolve(lookup, &config)?; let policy = generation_policy(&dir)?; + if world > 1 { + anyhow::ensure!( + knobs.lane_mode.is_none(), + "{ASYNC_PREFILL_ENV} is unsupported under tensor parallelism: its lane stream cannot \ + be lock-stepped across ranks" + ); + anyhow::ensure!( + !knobs.global_attn.tilelang(), + "{GLOBAL_ATTN_ENV} is unsupported under tensor parallelism: the generated kernels are \ + compiled for the whole global family" + ); + anyhow::ensure!( + knobs.prefix_cache.is_none(), + "{PREFIX_CACHE_ENV} is unsupported under tensor parallelism" + ); + anyhow::ensure!( + knobs.mix_chunk.is_none(), + "{MIX_CHUNK_TOKENS_ENV} is unsupported under tensor parallelism: a chunked walk's \ + rounds gather across prompts and are not covered by the TP gates" + ); + // The ranks take ordinals `0..world` of what this process can see, so a + // world size past the visible devices can only fail later, less clearly. + // `cuDeviceGetCount` needs the driver up and on a fresh process this is + // the first CUDA call; `init` is idempotent and process-wide. + cudarc::driver::result::init().map_err(|e| { + anyhow::anyhow!("CUDA driver init for a {world}-rank launch failed: {e:?}") + })?; + let visible = cudarc::driver::result::device::get_count().map_err(|e| { + anyhow::anyhow!("a {world}-rank launch needs to count CUDA devices: {e:?}") + })?; + anyhow::ensure!( + visible as usize >= world, + "a {world}-rank launch needs {world} CUDA devices, this process sees {visible}" + ); + // `docs/models/gemma4/tp.md` records that this container's stock NCCL + // corrupts memory under the default `LL`/`SIMPLE` protocols: the + // all-reduce kernel writes far past its buffer, and the symptom is the + // primary's collective stalling forever rather than an error. Nothing in + // the stack can detect that, so the configuration in force is said out + // loud rather than assumed. + let proto = std::env::var("NCCL_PROTO").unwrap_or_else(|_| "unset".to_string()); + let version = nccl_version(); + if proto.to_ascii_uppercase().contains("LL128") { + log::info!("tensor parallel: {world} ranks, NCCL {version}, NCCL_PROTO={proto}"); + } else { + log::warn!( + "tensor parallel: {world} ranks, NCCL {version}, NCCL_PROTO={proto} — with the \ + LL and SIMPLE protocols the all-reduce kernel writes past its buffer on the \ + NCCL 2.18.3 this container ships; use LL128 (see docs/models/gemma4/tp.md)" + ); + } + } let state = EngineState::load( &dir, - config, + &config, knobs, - device, + &ordinals, policy, base_seed, graph_enabled, @@ -1079,6 +1155,97 @@ fn send_scheduled(request: &QueuedRequest, cached_tokens: usize, ledger: &mut Re } } +/// One tensor-parallel rank's device state. Rank 0's lives in `EngineState`'s +/// own fields; ranks 1.. world size ride in `EngineState::more`. +struct RankState { + ctx: DeviceContext, + serve: GemmaServe, + arena: StepArena, +} + +/// The graph-before-comm release [`Drop for EngineState`] gives the success +/// path, for a `load` that fails after the capture sweep and so never builds an +/// `EngineState`: this rank's captured graphs must go before its own +/// communicator (a field of `serve`) drops. +impl Drop for RankState { + fn drop(&mut self) { + let _ = select_device(&self.ctx); + self.arena.release_graphs(); + } +} + +/// Make `ctx`'s device current on this thread. The thread-local cuBLAS handles +/// are keyed by device, so switching the current device is enough once every +/// rank has been bound at the start of a step. +fn activate_rank(ctx: &DeviceContext) -> Result<()> { + select_device(ctx)?; + // `cublas_init` is the only call that can create a device's handle pair on + // this thread, and it also selects them, so it both covers the first use + // of a device and the switch back to one already used. + unsafe { pegainfer_core::ffi::cublas_init() }; + let err = unsafe { pegainfer_core::ffi::cublas_activate_device_handles() }; + anyhow::ensure!( + err == 0, + "cuBLAS handle activation on device {} failed: cudaError={err}", + ctx.device_ordinal + ); + Ok(()) +} + +/// Set `ctx`'s device current and bind its primary context to this thread. +fn select_device(ctx: &DeviceContext) -> Result<()> { + let err = unsafe { pegainfer_core::ffi::cuda_set_device(ctx.device_ordinal as i32) }; + anyhow::ensure!( + err == 0, + "cudaSetDevice({}) on the scheduler thread failed: cudaError={err}", + ctx.device_ordinal + ); + ctx.ctx.bind_to_thread().map_err(|e| { + anyhow::anyhow!( + "bind device {} to the scheduler thread: {e}", + ctx.device_ordinal + ) + })?; + Ok(()) +} + +/// One `Vec` of rank-`r` families per rank for a batch: `payload[r][i]` is request +/// `i`'s rank-`r` families. Transposed once with a disjoint borrow per request, +/// which is what lets each rank's thread take its own slice. +fn rank_payloads(active: &mut [Active]) -> Vec> { + let world = active.first().map_or(1, |entry| entry.kv.world()); + let mut per_rank: Vec> = (0..world) + .map(|_| Vec::with_capacity(active.len())) + .collect(); + for entry in active.iter_mut() { + let (core, twins) = entry.kv.split(); + per_rank[0].push(core); + for (index, twin) in twins.iter_mut().enumerate() { + per_rank[index + 1].push(twin); + } + } + per_rank +} + +/// The prefill side of [`rank_payloads`]: the same transposition over the mixed +/// step's `(kv, prompt)` pairs. +fn prefill_payloads<'p, 't>( + prefills: &'p mut [(&mut GemmaKv, &'t [u32])], +) -> Vec> { + let world = prefills.first().map_or(1, |(kv, _)| kv.world()); + let mut per_rank: Vec> = (0..world) + .map(|_| Vec::with_capacity(prefills.len())) + .collect(); + for (kv, tokens) in prefills.iter_mut() { + let (core, twins) = kv.split(); + per_rank[0].push((core, *tokens)); + for (index, twin) in twins.iter_mut().enumerate() { + per_rank[index + 1].push((twin, *tokens)); + } + } + per_rank +} + /// Everything the contract-owned scheduler thread owns for the life of the /// engine. Loading completes before the driver thread is spawned, so launch /// failures return synchronously to the caller. @@ -1086,6 +1253,9 @@ struct EngineState { ctx: DeviceContext, serve: GemmaServe, arena: StepArena, + /// Ranks 1.. world size, empty at world size 1. Same shape as rank 0's + /// fields above. + more: Vec, scratch: SampleScratch, /// Conversation-tail prefix cache; `None` unless /// `PEGAINFER_PREFIX_CACHE=K` opted in at startup. @@ -1117,11 +1287,231 @@ struct EngineState { /// The decode-slot count the pools are budgeted for; requests past it /// queue. slots: usize, + /// Set once the tensor-parallel comms are aborted; every later step then + /// refuses rather than run a comm-less reduction that would return partial + /// sums. `Cell` so the abort path can set it while holding `&self.ctx`. + tp_broken: Cell, /// The admission coalesce window; `None` unless /// `PEGAINFER_ADMIT_COALESCE_MS` opted in at startup. admit_coalesce: Option, } +/// Captured collective graphs bake in NCCL kernel launches, and NCCL's +/// communicator abort wedges while a graph that references them is still +/// alive — so every rank releases its graphs here, on its own device, before +/// the fields (and with them the communicators) drop. +impl Drop for EngineState { + fn drop(&mut self) { + let _ = select_device(&self.ctx); + self.arena.release_graphs(); + for state in &mut self.more { + let _ = select_device(&state.ctx); + state.arena.release_graphs(); + } + } +} + +impl EngineState { + /// A fresh per-request KV across every rank: rank 0 from the primary + /// serve, each other rank from its own serve. Each entry is bound to its + /// rank's pools, which is what `admit_tokens`' `belongs_to` check needs. + fn alloc_kv(&self) -> GemmaKv { + GemmaKv::multi( + self.serve.alloc_rank_kv(), + self.more + .iter() + .map(|rank| rank.serve.alloc_rank_kv()) + .collect(), + ) + } + + /// Every rank must stand on rank 0's frontier. The page-id contract — + /// identical budgets and an identical admission sequence keep the ranks' + /// page ids in step — is checked nowhere else, so a divergence caught here + /// is the last chance before it becomes a silent read of another rank's + /// page. The free-page counts carry the same signal but take the pool locks, + /// so they are compared in a debug build only. + fn assert_ranks_in_step(&self, kv: &GemmaKv) -> Result<()> { + let local = kv.local.seq_len(); + let global = kv.global.seq_len(); + for rank in 1..=self.more.len() { + let theirs = kv.core(rank); + anyhow::ensure!( + theirs.local.seq_len() == local && theirs.global.seq_len() == global, + "rank {rank} stands on frontier {}/{} where rank 0 stands on {local}/{global}", + theirs.local.seq_len(), + theirs.global.seq_len() + ); + } + #[cfg(debug_assertions)] + for (index, state) in self.more.iter().enumerate() { + let rank = index + 1; + anyhow::ensure!( + state.serve.local_pool.available_pages() == self.serve.local_pool.available_pages(), + "rank {rank}'s local pool holds different free pages than rank 0's" + ); + anyhow::ensure!( + state.serve.global_pool.available_pages() + == self.serve.global_pool.available_pages(), + "rank {rank}'s global pool holds different free pages than rank 0's" + ); + } + Ok(()) + } + + /// Drive one step's segment on every rank at once — rank 0 on this thread, + /// each extra rank on its own — then join, and require that they all end on + /// rank 0's frontier. `payload` carries one item per rank, rank 0's first. + /// + /// One thread per extra rank is what lets rank 0's segment contain a + /// **host-blocking** operation (a device sync, an allocation in a GEMM's + /// workspace): by the time rank 0 reaches it, every peer's collectives are + /// already in flight, so the drain rank 0 waits on can pair — where a single + /// thread that ran rank 0's whole segment first would wait on a call it had + /// not yet issued. A rank-0 failure aborts rank 0's communicator so the + /// peers' pending collectives error out instead of hanging, and the extras' + /// are aborted after the join. An extra rank's failure is fatal: the ranks' + /// frontiers must not drift apart. + fn drive_ranks(&mut self, payload: Vec

, rank0: R, extra: E) -> Result + where + P: Send, + R: FnOnce(&DeviceContext, &GemmaServe, &mut StepArena, P) -> Result, + E: Fn(&mut RankState, P) -> Result<()> + Sync, + { + debug_assert_eq!( + payload.len(), + self.more.len() + 1, + "one payload item per rank" + ); + let EngineState { + ctx, + serve, + arena, + more, + tp_broken, + .. + } = self; + let mut payload = payload.into_iter(); + let core = payload.next().expect("a rank-0 payload item"); + let twins: Vec

= payload.collect(); + let extra_ranks = more.len(); + let outcome = std::thread::scope(|scope| -> Result { + let extra = &extra; + let mut handles = Vec::with_capacity(more.len()); + for (index, (state, item)) in more.iter_mut().zip(twins).enumerate() { + let rank = index + 1; + handles.push(scope.spawn(move || -> Result<()> { + activate_rank(&state.ctx)?; + extra(state, item) + .with_context(|| format!("tensor-parallel rank {rank} segment"))?; + state + .ctx + .sync() + .with_context(|| format!("drain rank {rank}")) + })); + } + activate_rank(ctx)?; + let ranked = rank0(ctx, serve, arena, core); + if ranked.is_err() && extra_ranks > 0 { + // Rank 0's sequence may have stopped short, so a peer's matching + // call will never come. Rank 0's own fields are not borrowed by + // the threads, so its communicator and graphs can be torn down + // here to unblock them; the extras' follow the join. + // + // Only with extra ranks: with none there is no collective to + // abort and no peer frontier to keep in step, so this stays the + // ordinary single-rank failure the caller answers by failing that + // one request. Marking the engine broken there would turn a + // request-local error — a per-request scratch allocation — into + // the end of the engine. + tp_broken.set(true); + if select_device(ctx).is_ok() { + arena.release_graphs(); + } + serve.detach_tp_comm(); + } + let mut extras: Result<()> = Ok(()); + for handle in handles { + let joined = handle + .join() + .map_err(|_| anyhow::anyhow!("a tensor-parallel rank thread panicked"))?; + if extras.is_ok() { + extras = joined; + } + } + // A rank-0 failure is the one to report; otherwise an extra rank's + // failure is fatal. + match ranked { + Err(err) => Err(err), + Ok(value) => { + extras?; + Ok(value) + } + } + }); + if outcome.is_err() { + self.break_extras(); + } + // The caller's next device work is rank 0's. + activate_rank(&self.ctx)?; + outcome + } + + /// Both halves of one request's KV, rank 0 first, as `drive_ranks`'s payload. + fn rank_payload(kv: &mut GemmaKv) -> Vec<&mut RankKv> { + let (core, twins) = kv.split(); + let mut payload = Vec::with_capacity(twins.len() + 1); + payload.push(core); + payload.extend(twins.iter_mut()); + payload + } + + /// Abort every extra rank's communicator and release its graphs, after a + /// failed step has already stopped the engine. + fn break_extras(&mut self) { + for state in &mut self.more { + if select_device(&state.ctx).is_ok() { + state.arena.release_graphs(); + } + state.serve.detach_tp_comm(); + } + } + + /// Reserve on every extra rank's pools, so their fronts stay in step with + /// rank 0's. Purely host-side, so no activation is needed. + fn admit_extra_ranks(&self, kv: &mut GemmaKv, tokens: usize) -> Result<()> { + for (rank, state) in self.more.iter().enumerate() { + admit_tokens( + &state.serve.local_pool, + &state.serve.global_pool, + kv.core_mut(rank + 1), + tokens, + )?; + } + // A rank that reserved different pages here has already lost the page-id + // race, and nothing after this can repair it. + #[cfg(debug_assertions)] + { + let free = ( + self.serve.local_pool.available_pages(), + self.serve.global_pool.available_pages(), + ); + for (rank, state) in self.more.iter().enumerate() { + debug_assert_eq!( + free, + ( + state.serve.local_pool.available_pages(), + state.serve.global_pool.available_pages() + ), + "rank {}'s pools reserved different pages than rank 0's", + rank + 1 + ); + } + } + Ok(()) + } +} + /// The scheduler thread is not the thread that loaded the engine: the /// primary context must be made current there and the thread-local cuBLAS /// handles created, or the first eager GEMM fails with an invalid handle. @@ -1211,7 +1601,10 @@ impl EngineState { loop { let refusal = match need { AdmissionNeed::Tokens(tokens) => { + // Every rank's pools reserve the same pages, so the fronts + // stay in step; only then does the admission succeed. admit_tokens(&self.serve.local_pool, &self.serve.global_pool, kv, tokens) + .and_then(|()| self.admit_extra_ranks(kv, tokens)) .err() .map(|err| format!("admission refused: {err:#}")) } @@ -1257,10 +1650,10 @@ impl EngineState { Ok(kv) => (kv, Some(entry.id)), Err(err) => { log::warn!("prefix-cache restore failed (falling back): {err:#}"); - (self.serve.alloc_kv(), None) + (self.alloc_kv(), None) } }, - None => (self.serve.alloc_kv(), None), + None => (self.alloc_kv(), None), } } @@ -1318,16 +1711,25 @@ impl EngineState { fn load( dir: &str, - config: crate::config::Gemma4Config, + config: &crate::config::Gemma4Config, knobs: ServingKnobs, - device: usize, + ordinals: &[usize], policy: GenerationPolicy, base_seed: u64, graph_enabled: bool, ) -> Result { + anyhow::ensure!( + !ordinals.is_empty(), + "a gemma4 engine needs at least one device ordinal" + ); + let world = ordinals.len(); + let device = ordinals[0]; + let tp = TensorParallelConfig::new(0, world); // Refuse an unservable global GQA shape or device before the // multi-GiB load. - let global_split = crate::serve::global_split_factor(&config)?; + tp.validate_for(config)?; + let global_split = + crate::serve::global_split_factor(&LayerGeometry::global_of(config, tp)?)?; let ServingKnobs { max_context, lane_mode, @@ -1343,7 +1745,7 @@ impl EngineState { if global_attn.tilelang() { ensure_tilelang_device(device)?; } - let weights = Gemma4Weights::from_safetensors(dir, device, config)?; + let weights = Gemma4Weights::from_safetensors(dir, device, config.clone(), tp)?; let ctx = DeviceContext::new_with_device(device)?; let vocab = weights.embed_tokens.rows; policy.check_against_vocab(vocab)?; @@ -1412,7 +1814,7 @@ impl EngineState { derives page or row counts past the i32 metadata domain (the global family's pseudo \ tables carry {global_split} copies of every page)" ); - let serve = GemmaServe::new( + let mut serve = GemmaServe::new( &ctx, weights, max_context, @@ -1446,7 +1848,11 @@ impl EngineState { let prefix_cache = cache_cap.map(|k| PrefixCache::new(k, sliding_window)); let mut scratch = SampleScratch::new(&ctx, vocab, arena_rows)?; let mut arena = serve.alloc_step_arena(&ctx, arena_rows, graph_enabled)?; - serve.precapture_decode_graphs(&ctx, &mut arena)?; + // One rank sweeps here; more than one sweeps in the interleaved driver + // below, which owns the buckets for every rank. + if world == 1 { + serve.precapture_decode_graphs(&ctx, &mut arena)?; + } let suppress_ids = ops::SuppressIds::upload(&ctx, &policy.suppress, vocab)?; let mut sampler_graphs = Vec::new(); if graph_enabled { @@ -1473,10 +1879,141 @@ impl EngineState { let lane = lane_mode .map(|mode| AsyncPrefillLane::new(&ctx, mode)) .transpose()?; + // The other ranks: their own sharded weights, context, pools and + // arena. The pool budget is rank-independent (identical pools with + // identical counts), so it is computed once above and handed to each. + let mut more = Vec::with_capacity(world - 1); + for (rank, ordinal) in ordinals.iter().copied().enumerate().skip(1) { + let rank_tp = TensorParallelConfig::new(rank, world); + let weights = Gemma4Weights::from_safetensors(dir, ordinal, config.clone(), rank_tp)?; + let rank_ctx = DeviceContext::new_with_device(ordinal)?; + let rank_serve = GemmaServe::new( + &rank_ctx, + weights, + max_context, + local_kv_storage, + local_pages, + global_pages, + global_attn, + ) + .map_err(|err| { + err.context(format!( + "rank {rank} (device {ordinal}) sized its pools to {local_pages} local / \ + {global_pages} global pages" + )) + })?; + let rank_arena = rank_serve.alloc_step_arena(&rank_ctx, arena_rows, graph_enabled)?; + more.push(RankState { + ctx: rank_ctx, + serve: rank_serve, + arena: rank_arena, + }); + } + // Materialize every rank's lazily-loaded cuBLAS/cublasLt kernels before + // the communicator exists. Under the default CUDA_MODULE_LOADING=LAZY + // a fresh process enters the driver's module loader on the first GEMM + // of each shape; NCCL's proxy threads enter that same loader while the + // communicator comes up, and the two deadlock — the eager prefill then + // never returns while the engine spin-waits. One tower pass per rank + // here, with no comm live, pays those loads single-threaded. The row + // counts span cublasLt's kernel-selection regions: a single row takes + // a GEMV-like kernel, the rest tile GEMMs whose choice shifts at small + // row counts. It also warms the decode capture sweep that used to run + // below under graphs. + if world > 1 { + let warm_rows = [1usize, 2, 3, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + .into_iter() + .filter(|rows| *rows <= max_context); + let warm = |serve: &GemmaServe, ctx: &DeviceContext, rows: usize| -> Result<()> { + let mut kv = serve.alloc_kv(); + admit_tokens(&serve.local_pool, &serve.global_pool, &mut kv, rows)?; + activate_rank(ctx)?; + serve.step(ctx, &mut kv, &vec![0u32; rows])?; + Ok(()) + }; + for rows in warm_rows.clone() { + warm(&serve, &ctx, rows)?; + } + for state in &more { + for rows in warm_rows.clone() { + warm(&state.serve, &state.ctx, rows)?; + } + } + activate_rank(&ctx)?; + } + // One communicator per rank, built on the stream the decode graph would + // run on, so an all-reduce lands inside that graph. + if world > 1 { + let mut streams = Vec::with_capacity(world); + streams.push(ctx.stream.clone()); + for state in &more { + streams.push(state.ctx.stream.clone()); + } + let mut comms = cudarc::nccl::safe::Comm::from_devices(streams) + .map_err(|e| anyhow::anyhow!("failed to initialize NCCL comms: {e:?}"))? + .into_iter(); + serve.attach_tp_comm(comms.next().expect("one comm per rank")); + for state in &mut more { + state + .serve + .attach_tp_comm(comms.next().expect("one comm per rank")); + } + } + // With graphs on, every rank captures its decode graphs, phase by phase: + // `Warm` and `Launch` execute and enqueue the all-reduce, whose peer + // call must be in flight, so the phases interleave across ranks instead + // of one rank finishing its whole sweep. `Capture` only records, and a + // recorded collective replays when its peer replays. Each rank sweeps + // its own arena and its own single-rank dummy. + if graph_enabled && world > 1 { + let mut dummies: Vec = Vec::with_capacity(world); + let mut primary = serve.alloc_kv(); + admit_tokens(&serve.local_pool, &serve.global_pool, &mut primary, 1)?; + dummies.push(primary); + for state in &more { + let mut dummy = state.serve.alloc_kv(); + admit_tokens( + &state.serve.local_pool, + &state.serve.global_pool, + &mut dummy, + 1, + )?; + dummies.push(dummy); + } + let mut bucket = 1usize; + while bucket <= arena.bucket_ceiling() { + for phase in PrecapturePhase::ALL { + activate_rank(&ctx)?; + serve.precapture_bucket(&ctx, &mut arena, &mut dummies[0], bucket, phase)?; + for (rank, state) in more.iter_mut().enumerate() { + activate_rank(&state.ctx)?; + state.serve.precapture_bucket( + &state.ctx, + &mut state.arena, + &mut dummies[rank + 1], + bucket, + phase, + )?; + } + } + bucket *= 2; + } + // Every rank's floor goes back to 1: the sweep left each arena's + // `min_bucket` at the ceiling, and a rank that kept it would pad a + // decode step to a different bucket than rank 0 — different graphs + // on different ranks for the same step. + arena.reset_min_bucket(); + for state in &mut more { + state.arena.reset_min_bucket(); + } + activate_rank(&ctx)?; + ctx.sync()?; + } Ok(Self { ctx, serve, arena, + more, scratch, prefix_cache, policy, @@ -1494,6 +2031,7 @@ impl EngineState { max_context, score_ceiling, slots, + tp_broken: Cell::new(false), admit_coalesce, }) } @@ -1668,6 +2206,20 @@ impl EngineState { let stepped = if let Some(top_k) = request.request.prompt_logprobs { let prompt = &request.request.prompt_tokens; let mut scores: Vec> = vec![None]; + // Rank 0's tower and the peers' run at once, so the per-row readback + // below — which blocks on rank 0's own collective — starts only once + // every rank's tower is in flight. + let mut tower = self.drive_ranks( + Self::rank_payload(&mut kv), + |ctx, serve, _arena, core| serve.launch_prompt_tower(ctx, core, prompt), + |state, rank_kv| { + state + .serve + .step_scoring(&state.ctx, rank_kv, prompt, None) + .map(|_| ()) + }, + )?; + self.assert_ranks_in_step(&kv)?; let (ctx, suppress_ids) = (&self.ctx, &self.suppress_ids); let mut score = |logits: &mut HiddenStates, start: usize| -> Result<()> { // The same logits a sampled token is scored on: softcapped, @@ -1686,28 +2238,40 @@ impl EngineState { scores.extend(scored.into_iter().map(Some)); Ok(()) }; - let stepped = self - .serve - .step_scoring(ctx, &mut kv, prompt, Some(&mut score)); + self.serve + .score_prompt_tower(ctx, &mut tower, &mut score) + .context("prompt logprobs")?; echo = Some(PromptEcho { ids: prompt.clone(), logprobs: scores, }); - stepped + Ok(tower.into_logits()) } else if let Some(chunk) = self.mix_chunk { // Under the chunk knob a solo prompt walks its own segments too: // residency stays window plus segment whatever the prompt length. self.walk_plain_prompt(&mut kv, &request.request.prompt_tokens, chunk) } else { let resume = kv.local.seq_len(); - self.serve - .step(&self.ctx, &mut kv, &request.request.prompt_tokens[resume..]) + let tokens = &request.request.prompt_tokens[resume..]; + let logits = self.drive_ranks( + Self::rank_payload(&mut kv), + |ctx, serve, _arena, core| serve.step(ctx, core, tokens), + |state, rank_kv| state.serve.step(&state.ctx, rank_kv, tokens).map(|_| ()), + )?; + self.assert_ranks_in_step(&kv)?; + Ok(logits) }; let mut logits = match stepped { Ok(logits) => logits, Err(err) => { - // This prompt's prefill failed; its pages return with `kv` - // and the engine keeps serving. + // A broken tensor-parallel engine cannot serve: the comms are + // gone, so every later step would skip its reduction and return + // partial sums. Stop rather than fail this request and continue. + if self.tp_broken.get() { + return Err(err.context("tensor-parallel ranks diverged; engine stopped")); + } + // This prompt's prefill failed on its own; its pages return with + // `kv` and the engine keeps serving. log::error!("solo prefill failed: {err:#}"); ledger.fail(request.id, format!("prefill failed: {err:#}")); return Ok(Admitted::Done); @@ -1877,7 +2441,7 @@ impl EngineState { /// A prompt's unseen suffix on the plain path, one `chunk` segment at a /// time. fn walk_plain_prompt( - &self, + &mut self, kv: &mut GemmaKv, prompt: &[u32], chunk: usize, @@ -1890,14 +2454,21 @@ impl EngineState { self.step_plain_segment(kv, &prompt[offset..]) } - fn step_plain_segment(&self, kv: &mut GemmaKv, tokens: &[u32]) -> Result { + fn step_plain_segment(&mut self, kv: &mut GemmaKv, tokens: &[u32]) -> Result { admit_tokens( &self.serve.local_pool, &self.serve.global_pool, kv, tokens.len(), )?; - self.serve.step(&self.ctx, kv, tokens) + self.admit_extra_ranks(kv, tokens.len())?; + let logits = self.drive_ranks( + Self::rank_payload(kv), + |ctx, serve, _arena, core| serve.step(ctx, core, tokens), + |state, rank_kv| state.serve.step(&state.ctx, rank_kv, tokens).map(|_| ()), + )?; + self.assert_ranks_in_step(kv)?; + Ok(logits) } fn finish_plain_walker( @@ -1914,6 +2485,12 @@ impl EngineState { ) { Ok(logits) => logits, Err(err) => { + // A walk tail after a broken pair would score its first token + // from partial sums; stop the engine instead (see + // `drive_ranks`' rank-0 teardown). + if self.tp_broken.get() { + return Err(err.context("tensor-parallel ranks diverged; engine stopped")); + } ledger.fail(walker.request.id, format!("walk tail failed: {err:#}")); walker.failed = true; return Ok(()); @@ -2009,7 +2586,9 @@ impl EngineState { &self.serve.global_pool, &mut walker.kv, take, - ) { + ) + .and_then(|()| self.admit_extra_ranks(&mut walker.kv, take)) + { ledger.fail( walker.request.id, format!("walk segment admission failed: {err:#}"), @@ -2128,7 +2707,9 @@ impl EngineState { &self.serve.global_pool, &mut entry.kv, 1, - ) { + ) + .and_then(|()| self.admit_extra_ranks(&mut entry.kv, 1)) + { ledger.fail( entry.request.id, format!("decode KV admission failed: {err:#}"), @@ -2160,17 +2741,33 @@ impl EngineState { } /// Reserve the next token without retiring or reordering a row. - fn ready_rows_pinned(&self, active: &mut [Active], ledger: &RequestLedger) -> bool { - active.iter_mut().all(|entry| { - !ledger.is_aborted(entry.request.id) - && admit_tokens( - &self.serve.local_pool, - &self.serve.global_pool, - &mut entry.kv, - 1, - ) - .is_ok() - }) + /// + /// `Ok(false)` is rank 0's own shortfall, which the caller answers by + /// falling back to [`Self::ready_decode_rows`] — that path fails the row and + /// drops it, so its pages come back. `Err` is only reachable with extra + /// ranks: a rank that cannot reserve what rank 0 just reserved has lost the + /// page-id race, and nothing here drops the row, so the divergence would + /// ride into the step and read one rank's KV at another's frontier. By + /// `admit_extra_ranks`' own contract nothing after that can repair it, so + /// the engine stops instead. + fn ready_rows_pinned(&self, active: &mut [Active], ledger: &RequestLedger) -> Result { + for entry in active.iter_mut() { + if ledger.is_aborted(entry.request.id) { + return Ok(false); + } + if admit_tokens( + &self.serve.local_pool, + &self.serve.global_pool, + &mut entry.kv, + 1, + ) + .is_err() + { + return Ok(false); + } + self.admit_extra_ranks(&mut entry.kv, 1)?; + } + Ok(true) } fn fence(&self) -> Result<()> { @@ -2191,15 +2788,41 @@ impl EngineState { slot: usize, ) -> Result { let rows = active.len(); + // The staged pipeline hands rank 1 no ids of its own (it has no + // sampler), so under tensor parallelism every rank takes the + // explicit-token path. + let resident = resident && self.more.is_empty(); let tokens = (!resident).then(|| active.iter().map(|entry| entry.next).collect::>()); { - let mut kvs: Vec<&mut GemmaKv> = active.iter_mut().map(|entry| &mut entry.kv).collect(); + let payload = rank_payloads(active); if let Some(tokens) = tokens.as_deref() { - self.serve - .decode_batch_step(&self.ctx, &mut self.arena, &mut kvs, tokens)?; + // Every rank carries the same tokens and page ids; only rank 0 + // samples, so the extras' logits are discarded. + self.drive_ranks( + payload, + |ctx, serve, arena, mut kvs| { + serve + .decode_batch_step(ctx, arena, &mut kvs, tokens) + .map(|_| ()) + }, + |state, mut kvs| { + state + .serve + .decode_batch_step(&state.ctx, &mut state.arena, &mut kvs, tokens) + .map(|_| ()) + }, + )?; } else { - self.serve - .decode_batch_step_resident(&self.ctx, &mut self.arena, &mut kvs)?; + // World size 1: the resident path, with no extra rank to drive. + self.drive_ranks( + payload, + |ctx, serve, arena, mut kvs| { + serve + .decode_batch_step_resident(ctx, arena, &mut kvs) + .map(|_| ()) + }, + |_state, _kvs| Ok(()), + )?; } } let graph_slot = crate::serve::decode_bucket_slot(rows); @@ -2281,23 +2904,39 @@ impl EngineState { ledger: &mut RequestLedger, ) -> Result { let decode_tokens: Vec = active.iter().map(|entry| entry.next).collect(); - let stepped = { - let mut kvs: Vec<&mut GemmaKv> = active.iter_mut().map(|entry| &mut entry.kv).collect(); - self.serve.mixed_prefill_decode_step( - &self.ctx, - &mut self.arena, - prefills, - &mut kvs, - &decode_tokens, - ) - }; - let logits = match stepped { - Ok(logits) => logits, - Err(err) => { - fail_active_batch(active, "mixed step", &err, ledger); - return Err(err.context("gemma4 mixed step")); - } - }; + // One payload per rank: that rank's slice of the decode batch and of the + // gathered prompts. `mixed_prefill_decode_step` runs them as one step, so + // every rank is in flight at once (see `drive_ranks`). + let payload: Vec<_> = rank_payloads(active) + .into_iter() + .zip(prefill_payloads(prefills)) + .collect(); + let stepped = self.drive_ranks( + payload, + |ctx, serve, arena, (mut kvs, mut prompts)| { + serve + .mixed_prefill_decode_step(ctx, arena, &mut prompts, &mut kvs, &decode_tokens) + .map(|_| ()) + }, + |state, (mut kvs, mut prompts)| { + state + .serve + .mixed_prefill_decode_step( + &state.ctx, + &mut state.arena, + &mut prompts, + &mut kvs, + &decode_tokens, + ) + .map(|_| ()) + }, + ); + if let Err(err) = stepped { + fail_active_batch(active, "mixed step", &err, ledger); + return Err(err.context("gemma4 mixed step")); + } + // The mixed step's logits live in rank 0's arena for this step. + let (logits, _) = self.arena.logits_and_ids(); mixed_head_flow( &self.ctx, &self.suppress_ids, @@ -2375,20 +3014,27 @@ impl EngineState { ledger: &mut RequestLedger, ) -> Result<()> { let tokens: Vec = active.iter().map(|entry| entry.next).collect(); - let logits = { - let mut kvs: Vec<&mut GemmaKv> = active.iter_mut().map(|entry| &mut entry.kv).collect(); - match self - .serve - .decode_batch_step(&self.ctx, &mut self.arena, &mut kvs, &tokens) - { - Ok(logits) => logits, - Err(err) => { - self.fence()?; - fail_active_batch(active, "batched decode", &err, ledger); - return Err(err.context("batched decode")); - } - } - }; + // Every rank carries the same tokens and page ids; only rank 0 samples, so + // the extras' logits are discarded. A failure on any rank is fatal. + if let Err(err) = self.drive_ranks( + rank_payloads(active), + |ctx, serve, arena, mut kvs| { + serve + .decode_batch_step(ctx, arena, &mut kvs, &tokens) + .map(|_| ()) + }, + |state, mut kvs| { + state + .serve + .decode_batch_step(&state.ctx, &mut state.arena, &mut kvs, &tokens) + .map(|_| ()) + }, + ) { + let _ = self.fence(); + fail_active_batch(active, "batched decode", &err, ledger); + return Err(err.context("batched decode")); + } + let (logits, _) = self.arena.logits_and_ids(); let sampled = { let rows: Vec> = active .iter() @@ -2419,7 +3065,7 @@ impl EngineState { /// successor in flight and drain before any row-order change. fn decode_round(&mut self, active: &mut Vec, ledger: &mut RequestLedger) -> Result<()> { if let Some(pending) = self.pipeline.take() { - if self.pipeline_eligible(active, ledger) && self.ready_rows_pinned(active, ledger) { + if self.pipeline_eligible(active, ledger) && self.ready_rows_pinned(active, ledger)? { let next_slot = (pending.slot + 1) % DECODE_PIPELINE_DEPTH; match self.launch_staged(active, true, next_slot) { Ok(rows) => { @@ -2491,8 +3137,18 @@ impl Scheduler for Gemma4Scheduler { } fn step(&mut self, ledger: &mut RequestLedger) -> Result<()> { + if self.state.tp_broken.get() { + anyhow::bail!( + "tensor-parallel ranks diverged; the communicators are aborted and the engine \ + cannot serve" + ); + } if self.cublas.is_none() { + // Rank 0's handles, and the guard that tears every device's handles + // down on this thread. Each other rank's handles are created the + // first time `activate_rank` reaches it. self.cublas = Some(bind_engine_thread(&self.state.ctx)?); + activate_rank(&self.state.ctx)?; } if self.walk.is_some() { return self.advance_walk(ledger); @@ -2931,3 +3587,7 @@ mod lane_gates_walk; #[cfg(test)] #[path = "engine/lane_gates_logprobs.rs"] mod lane_gates_logprobs; + +#[cfg(test)] +#[path = "engine/lane_gates_tp.rs"] +mod lane_gates_tp; diff --git a/pegainfer-gemma4/src/engine/lane_gates_logprobs.rs b/pegainfer-gemma4/src/engine/lane_gates_logprobs.rs index 560fd3e8c..fb1e6403b 100644 --- a/pegainfer-gemma4/src/engine/lane_gates_logprobs.rs +++ b/pegainfer-gemma4/src/engine/lane_gates_logprobs.rs @@ -118,10 +118,11 @@ fn a_low_slot_chunked_pool_scores_up_to_what_it_holds() { harness.shutdown(&[]); } -#[test] -fn only_a_scored_prompt_is_bound_to_the_whole_prompt_ceiling() { - let request = |prompt_logprobs| Request { - prompt_tokens: ids(super::MAX_CONTEXT + 1, 9), +/// A probe whose prompt is exactly `prompt_len` tokens, so a ceiling can be +/// tested against a length known without reading the prompt back. +fn ceiling_probe(prompt_len: usize, prompt_logprobs: Option) -> Request { + Request { + prompt_tokens: ids(prompt_len, 9), params: pegainfer_frontend::sampler::SamplingParams::default(), stop_policy: pegainfer_frontend::engine::StopPolicy::default(), max_tokens: 4, @@ -131,11 +132,62 @@ fn only_a_scored_prompt_is_bound_to_the_whole_prompt_ceiling() { prompt_logprobs, trace_parent: None, client_label: None, - }; + } +} + +#[test] +fn only_a_scored_prompt_is_bound_to_the_whole_prompt_ceiling() { let raised = 4 * super::MAX_CONTEXT; - assert!(super::validate_request(&request(None), raised, super::MAX_CONTEXT).is_ok()); + let prompt_len = super::MAX_CONTEXT + 1; + assert!( + super::validate_request(&ceiling_probe(prompt_len, None), raised, super::MAX_CONTEXT) + .is_ok() + ); assert!(matches!( - super::validate_request(&request(Some(0)), raised, super::MAX_CONTEXT), + super::validate_request( + &ceiling_probe(prompt_len, Some(0)), + raised, + super::MAX_CONTEXT + ), Err(RejectReason::EchoPrefillTokens { .. }) )); } + +/// A token id outside the embedding fails in `prepare_single`'s +/// `validate_tokens`, before the tower allocates anything or — under tensor +/// parallelism — before rank 0 issues a single collective. So a prompt carrying +/// one is a deterministic prefill fault that needs no injection hook, and the +/// property under test is its **scope**: the driver contract makes `Err` from +/// `Scheduler::step` mean "the engine is beyond use", so a request-local fault +/// that escaped as `Err` would close the step stream and write off every open +/// account. Both prefill paths have to answer it by failing that one request. +#[test] +#[ignore = "requires a Gemma 4 checkpoint, a GPU, and --test-threads=1"] +fn a_failed_prefill_costs_that_request_not_the_engine() { + let mut harness = launch(&[]); + for (name, salt, prompt_logprobs) in [("scored", 1u32, Some(8)), ("plain", 2, None)] { + // `u32::MAX` is outside any vocabulary this line ships. + let bad = harness.submit_scored(vec![9, u32::MAX, 11 + salt], 4, None, prompt_logprobs); + match harness.steps.terminal(bad.id()) { + Terminal::Failed { message, .. } => assert!( + message.contains("prefill failed"), + "{name}: the request should fail as a prefill failure, got: {message}" + ), + other => panic!("{name}: an unusable prompt must fail that request, not {other:?}"), + } + // The engine is still serving, and the failed request's pages came back + // with its KV: a good request is admitted and completes afterwards. + let good = harness.submit_scored(ids(12, 7 + salt), 4, None, prompt_logprobs); + let drained = harness.steps.drain(good.id(), name); + assert_eq!( + drained.tokens, 4, + "{name}: the request after a failed prefill still decodes to its budget" + ); + assert_eq!( + drained.finish, + FinishReason::Length, + "{name}: and finishes on its budget" + ); + } + harness.shutdown(&[]); +} diff --git a/pegainfer-gemma4/src/engine/lane_gates_tp.rs b/pegainfer-gemma4/src/engine/lane_gates_tp.rs new file mode 100644 index 000000000..b690d32ac --- /dev/null +++ b/pegainfer-gemma4/src/engine/lane_gates_tp.rs @@ -0,0 +1,681 @@ +//! Tensor-parallel gates: a two-rank engine must agree with one rank on the +//! same checkpoint, prompts and batch. This is the check that the sharded +//! loader, the two reductions and the per-rank KV stay in step; it drives the +//! real launch path, not a hand-built stack. +//! +//! It compares *distributions*, not greedy tokens. A two-rank reduction sums +//! the same products in a different order than one rank does, and NCCL writes +//! bf16 back, so the logits differ in their last bits and a near-tie can flip +//! the greedy pick. The engine's own docs say as much — greedy is reproducible +//! for a fixed workload, not across two arithmetic orders — so the gate holds +//! the top-k logprobs to a line and requires a differing pick to be a genuine +//! near-tie (each pick inside the other run's top-k). +//! +//! Two GPUs are mandatory, so these run one at a time on a box that has them. + +use std::collections::HashMap; + +use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::TokenLogprob; +use pegainfer_frontend::parallel::ParallelConfig; + +use super::lane_tests::Drained; +use super::lane_tests::Harness; +use super::lane_tests::ids; +use super::lane_tests::launch_with; +use crate::testkit::f32_tensor; +use crate::testkit::golden_bytes; +use crate::testkit::i32_tensor; +use crate::testkit::u32_tensor; + +/// How many top logprobs each run keeps, and the largest absolute logprob gap +/// two runs may show on a token they both kept. The gap is the bf16 +/// reduction-order drift accumulated over the tower, not a shape difference. +/// +/// The line is set from this gate's own readings — one procedure, one head: +/// four runs of the 12B on two L20s (`NCCL_PROTO=LL128`) each held +/// `worst_pick` to 0.5738, the same to four decimals every run, and one run +/// with the `o_proj` reduction skipped — the structural error this line is for +/// — drove the one-rank pick out of the two-rank top-8 (the near-tie rule fires +/// before the line is reached) and read `worst_top` 2.1895 with that rule +/// relaxed; `worst_pick` read 0 there, the two lists sharing no token, which is +/// the case the near-tie rule exists for. 1.0 is ~1.7x over the floor and ~2.2x +/// under the fault, and is the value `DRIFT_LINE` below already uses. The shape +/// follows the repo's one precedent for calibrating a quantity of this kind, +/// `serve_oracle`'s `neutral_scale`: two algorithms over one context, 0.31..5.75 +/// observed and a line at 12.0. +/// +/// What the line is for: a structural tensor-parallel error — a wrong shard, a +/// missing reduction, a rank out of step — moves a logprob by many nats and +/// flips most picks, not 2 of 24. Reduction noise is what the near-tie rule +/// beside it polices. +const TOP_K: usize = 8; +const LOGBROB_LINE: f32 = 1.0; + +/// The two devices a TP2 gate runs across, defaulting to 0 and 1. +fn devices() -> (usize, usize) { + let read = |name: &str, fallback: usize| { + std::env::var(name) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(fallback) + }; + (read("PEGAINFER_TP_DEVICE", 0), read("PEGAINFER_TP_PEER", 1)) +} + +/// One rank, one device, eager. +fn single_options(device: usize) -> EngineLoadOptions { + EngineLoadOptions { + enable_cuda_graph: false, + device_ordinals: vec![device], + ..EngineLoadOptions::default() + } +} + +/// Two ranks, two devices. Graphs are on when `PEGAINFER_TP_GRAPH` is set, which +/// turns the gate into a parity check: one rank eager against two ranks +/// captured. +fn tp2_options(a: usize, b: usize) -> EngineLoadOptions { + EngineLoadOptions { + enable_cuda_graph: std::env::var("PEGAINFER_TP_GRAPH").is_ok(), + device_ordinals: vec![a, b], + parallel_config: Some(ParallelConfig::new(2, 1)), + ..EngineLoadOptions::default() + } +} + +/// Serve the first prompt on its own, then the rest as one concurrent batch, +/// and shut the engine down once. +/// +/// A single prompt in flight is the **solo** admission path (`step` plus +/// `prefill_extra_ranks`), which a batch never takes; draining it before the +/// others are submitted is what makes it a lone arrival. Both runs drive this +/// same sequence, so the per-prompt comparison stays index-aligned. +fn serve_batch( + harness: &mut Harness, + prompts: &[Vec], + max_tokens: usize, + logprobs: usize, +) -> Vec { + let mut controls = Vec::with_capacity(prompts.len()); + let mut drained = Vec::with_capacity(prompts.len()); + if let Some(first) = prompts.first() { + let control = harness.submit_scored(first.clone(), max_tokens, Some(logprobs), None); + drained.push(harness.steps.drain(control.id(), "greedy solo")); + controls.push(control); + } + for prompt in &prompts[1..] { + controls.push(harness.submit_scored(prompt.clone(), max_tokens, Some(logprobs), None)); + } + for control in &controls[1..] { + drained.push(harness.steps.drain(control.id(), "greedy")); + } + let controls: Vec<&_> = controls.iter().collect(); + harness.shutdown(&controls); + drained +} + +fn prompts() -> Vec> { + // `PEGAINFER_TP_PROMPT_TOKENS` widens or narrows the prompts; the default + // is the three-prompt set the line compares at. `PEGAINFER_TP_PROMPTS` + // narrows the count to exercise the solo admission path. + let base: usize = std::env::var("PEGAINFER_TP_PROMPT_TOKENS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(20); + let count: usize = std::env::var("PEGAINFER_TP_PROMPTS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(3); + (0..count) + .map(|i| ids(base + i * 7, 11 + i as u32)) + .collect() +} + +fn top_of(lp: &TokenLogprob) -> HashMap { + lp.top_logprobs + .iter() + .map(|(token, value)| (*token, *value)) + .collect() +} + +/// Every value a gap is taken from has to be finite. One NaN logit folds the +/// row's logsumexp — and with it every logprob on that row — into NaN while the +/// top-k *ids* stay right, and both `f32::max` and `gap > worst` step over a NaN +/// without moving the running bound, so a gate would pass on a row it exists to +/// fail. `serve_oracle::compare_row` asserts the same over both arms. +fn assert_finite(lp: &TokenLogprob, at: &str) { + assert!( + lp.logprob.is_finite() && lp.top_logprobs.iter().all(|(_, value)| value.is_finite()), + "{at} scored a non-finite logprob: {} with top {:?}", + lp.logprob, + lp.top_logprobs + ); +} + +/// Two runs of the same one-rank workload must come out the same token for token +/// and bit for bit. This is the control's bar: a gap metric that only bounds a +/// tolerance can pass while two tokens trade places, so the control compares the +/// runs themselves and the two-rank comparison below is left measuring tensor +/// parallelism rather than harness noise. +fn assert_identical(one: &[Drained], two: &[Drained], what: &str) { + assert_eq!( + one.len(), + two.len(), + "{what}: {} vs {} requests", + one.len(), + two.len() + ); + for (index, (a, b)) in one.iter().zip(two).enumerate() { + assert_eq!( + a.ids, b.ids, + "{what}: request {index} decoded a different token" + ); + assert_eq!( + a.logprobs, b.logprobs, + "{what}: request {index} scored different logprobs" + ); + } +} + +/// The worst shared top-k logprob gap over the run, and how many steps kept +/// the same pick. `Detail` names the step that produced the gap so a failure +/// can be read without a re-run. +struct Detail { + request: usize, + step: usize, + token: u32, + one: f32, + two: f32, + one_top: Vec<(u32, f32)>, + two_top: Vec<(u32, f32)>, +} + +struct Gaps { + /// Largest gap over tokens both runs kept in their top-k. Inflated by a + /// flattened distribution, where a shared token can sit at very different + /// ranks in the two lists — read it beside `worst_pick`. + worst_shared: f32, + /// Largest gap on a token *both* runs picked, and on either token when the + /// picks differ: the like-for-like measure of how far the two distributions + /// moved. A flip is where a large move hides — the two tops can still read + /// equal — so the flipped picks are measured too, not skipped. + worst_pick: f32, + /// Largest gap between the two runs' *top* logprobs, whichever token holds + /// the top. A flip keeps this at zero, so it is a bound beside `worst_pick`, + /// not what catches one. + worst_top: f32, + same_pick: usize, + compared: usize, + detail: Detail, +} + +fn distribution_gap(one: &[Drained], two: &[Drained], what: &str) -> Gaps { + let mut worst_shared = 0.0f32; + let mut worst_pick = 0.0f32; + let mut worst_top = 0.0f32; + let mut same_pick = 0usize; + let mut compared = 0usize; + let mut detail = Detail { + request: 0, + step: 0, + token: 0, + one: 0.0, + two: 0.0, + one_top: Vec::new(), + two_top: Vec::new(), + }; + for (index, (a, b)) in one.iter().zip(two).enumerate() { + assert_eq!( + a.logprobs.len(), + b.logprobs.len(), + "{what}: request {index} produced {} vs {} scored steps", + a.logprobs.len(), + b.logprobs.len() + ); + for (step, (a, b)) in a.logprobs.iter().zip(&b.logprobs).enumerate() { + let (Some(a), Some(b)) = (a, b) else { + continue; + }; + compared += 1; + assert_finite( + a, + &format!("{what}: request {index} step {step} (one rank)"), + ); + assert_finite( + b, + &format!("{what}: request {index} step {step} (two rank)"), + ); + let (left, right) = (top_of(a), top_of(b)); + worst_top = worst_top.max((a.top_logprobs[0].1 - b.top_logprobs[0].1).abs()); + for (token, value) in &left { + if let Some(other) = right.get(token) { + let gap = (value - other).abs(); + if gap > worst_shared { + worst_shared = gap; + detail = Detail { + request: index, + step, + token: *token, + one: *value, + two: *other, + one_top: a.top_logprobs.clone(), + two_top: b.top_logprobs.clone(), + }; + } + } + } + let (pa, pb) = (a.top_logprobs[0].0, b.top_logprobs[0].0); + if pa == pb { + same_pick += 1; + // The same token in both runs: the clean like-for-like measure + // of how far the two distributions moved. + if let (Some(va), Some(vb)) = (left.get(&pa), right.get(&pa)) { + worst_pick = worst_pick.max((va - vb).abs()); + } + } else { + // A different pick has to be a near-tie: each pick inside the + // other run's top-k. Otherwise the two runs disagree about the + // distribution, not about its last bits. + assert!( + right.contains_key(&pa), + "{what}: request {index} step {step}: one-rank pick is outside the two-rank top-{TOP_K}" + ); + assert!( + left.contains_key(&pb), + "{what}: request {index} step {step}: two-rank pick is outside the one-rank top-{TOP_K}" + ); + // Both picks now sit in both lists, so each one's own gap is + // measurable, and a flip is exactly where a large move hides: + // the two tops can read equal while the same token moved far. + // Fold both in before the break, or the flip is the one step + // that leaves the run with no bound at all. + for token in [pa, pb] { + if let (Some(va), Some(vb)) = (left.get(&token), right.get(&token)) { + worst_pick = worst_pick.max((va - vb).abs()); + } + } + // The two runs now decode from different prefixes, so every later + // step of this request compares different text; stop here. + break; + } + } + } + Gaps { + worst_shared, + worst_pick, + worst_top, + same_pick, + compared, + detail, + } +} + +/// Two ranks must agree with one rank on the same checkpoint. Which +/// tensor-parallel branch it exercises is the checkpoint's business: the shipped +/// 12B (a single global KV head) takes the **replicate** branch, a +/// 31B-geometry checkpoint (`G % P == 0`) takes the **shard** branch that +/// production serves. Point `PEGAINFER_TEST_MODEL_PATH` at whichever is under +/// test, and note that a full-depth 31B cannot be gated here at all: the +/// single-rank control has to fit one card, which 57 GiB of weights do not. +#[test] +#[ignore = "needs two GPUs and --test-threads=1; checkpoint from PEGAINFER_TEST_MODEL_PATH"] +fn the_two_rank_engine_matches_one_rank() { + let (device, peer) = devices(); + assert_ne!(device, peer, "TP2 needs two distinct device ordinals"); + let prompts = prompts(); + // Serving knobs may be forced from the environment so the same gate runs + // at a reduced envelope as well as at the default point. + let overrides = envelope_overrides(); + let overrides = as_refs(&overrides); + eprintln!("tp2: knob overrides {overrides:?}"); + + // One rank first, and torn down before the two-rank engine takes the same + // devices. + let mut one = launch_with(&single_options(device), &overrides); + let single = serve_batch(&mut one, &prompts, 16, TOP_K); + drop(one); + // The control: the same one-rank run again, and it must come out + // **bit-identical**. A metric that only holds a tolerance can pass while a + // token trades places with another, so compare the runs themselves. + let mut one_again = launch_with(&single_options(device), &overrides); + let repeat = serve_batch(&mut one_again, &prompts, 16, TOP_K); + drop(one_again); + assert_identical(&single, &repeat, "tp1-repeat"); + + let mut two = launch_with(&tp2_options(device, peer), &overrides); + let tp2 = serve_batch(&mut two, &prompts, 16, TOP_K); + drop(two); + + // `single == repeat` bit for bit, so comparing tp2 against the repeat run + // would repeat this comparison; one of them is enough. + let gaps = distribution_gap(&single, &tp2, "tp2"); + // A reading is only comparable with the build and pair it came from, so + // name both. + eprintln!( + "tp2: NCCL {}, devices {device},{peer}", + super::nccl_version() + ); + eprintln!( + "tp2: {}/{} steps keep the one-rank pick; worst picked-token logprob gap {:.4} (top {:.4})", + gaps.same_pick, gaps.compared, gaps.worst_pick, gaps.worst_top + ); + eprintln!( + "tp2: worst shared-token {:.4} at request {} step {} token {} ({:.4} vs {:.4})", + gaps.worst_shared, + gaps.detail.request, + gaps.detail.step, + gaps.detail.token, + gaps.detail.one, + gaps.detail.two + ); + eprintln!("tp2: one-rank top-{TOP_K} {:?}", gaps.detail.one_top); + eprintln!("tp2: two-rank top-{TOP_K} {:?}", gaps.detail.two_top); + assert!( + gaps.worst_pick < LOGBROB_LINE && gaps.worst_top < LOGBROB_LINE, + "two-rank and one-rank logprobs differ by {} on a picked token / {} at the top \ + (line {LOGBROB_LINE})", + gaps.worst_pick, + gaps.worst_top + ); +} + +/// The serving knobs a gate may force from the environment, so the same gate +/// runs at a reduced envelope as well as at the default point — a 31B needs a +/// smaller one to fit a 48 GiB card at all. +fn envelope_overrides() -> Vec<(&'static str, String)> { + let mut overrides = Vec::new(); + for (var, knob) in [ + ("PEGAINFER_TP_SLOTS", super::DECODE_SLOTS_ENV), + ("PEGAINFER_TP_CTX", super::MAX_CONTEXT_ENV), + ] { + if let Ok(value) = std::env::var(var) { + overrides.push((knob, value)); + } + } + overrides +} + +fn as_refs<'a>(overrides: &'a [(&'static str, String)]) -> Vec<(&'static str, &'a str)> { + overrides + .iter() + .map(|(knob, value)| (*knob, value.as_str())) + .collect() +} + +/// A request that asks for its prompt's logprobs is scored through a path the +/// sampled gate never touches — rank 0 reads the head's rows back to the host +/// inside its own step — so it gets its own two-rank gate. +#[test] +#[ignore = "needs two GPUs and --test-threads=1"] +fn the_two_rank_engine_scores_prompt_logprobs() { + const TOP_K: usize = 8; + // A mid-length prompt, so the scored path runs a real prefill per rank rather + // than a single-token one. `PEGAINFER_TP_PROMPT_TOKENS` probes another point. + let len: usize = std::env::var("PEGAINFER_TP_PROMPT_TOKENS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(64); + let (device, peer) = devices(); + assert_ne!(device, peer, "TP2 needs two distinct device ordinals"); + let mut harness = launch_with(&tp2_options(device, peer), &as_refs(&envelope_overrides())); + // Both submitted before either is drained, so the second scored admission + // arrives while the first still holds KV — the path a single submit misses. + let prompts: Vec> = [11u32, 26].into_iter().map(|seed| ids(len, seed)).collect(); + let controls: Vec<_> = prompts + .iter() + .enumerate() + .map(|(index, prompt)| { + // The first request asks for more than one token, so it is still + // decoding — a live batch — when the second is admitted. + let max_tokens = if index == 0 { 4 } else { 1 }; + harness.submit_scored(prompt.clone(), max_tokens, Some(TOP_K), Some(TOP_K)) + }) + .collect(); + for (control, prompt) in controls.iter().zip(&prompts) { + let drained = harness.steps.drain(control.id(), "scored prompt"); + let echo = drained.prompt_echo.expect("the engine echoes the prompt"); + assert_eq!( + echo.logprobs.len(), + prompt.len(), + "one scored row per prompt token" + ); + } + let refs: Vec<&_> = controls.iter().collect(); + harness.shutdown(&refs); +} + +/// A causal model's row `r` depends only on tokens `..= r`, so the same prompt's +/// early rows must score the same whether the prompt is run whole or cut short. +/// They will not be bit-identical — the prefill GEMM's tiling, and so its +/// accumulation order, depends on the row count — so the gate bounds the shared +/// rows by the same chaotic-shape line `serve_oracle` calibrates (`PREFIX_LINE`), +/// not by the one-rank comparison's tighter one. A **large** gap here is the +/// engine's own long-context bug (a page, position or rope error), with no +/// reference implementation involved. +/// +/// `PEGAINFER_TP_PROMPT_TOKENS` sets the long length (default 1024, the window +/// width); the short run is its first half. +#[test] +#[ignore = "needs two GPUs and --test-threads=1; checkpoint from PEGAINFER_TEST_MODEL_PATH"] +fn the_two_rank_engine_is_prefix_consistent() { + const SALT: u32 = 11; + let long: usize = std::env::var("PEGAINFER_TP_PROMPT_TOKENS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(1024); + let short = long / 2; + let (device, peer) = devices(); + assert_ne!(device, peer, "TP2 needs two distinct device ordinals"); + let mut harness = launch_with(&tp2_options(device, peer), &as_refs(&envelope_overrides())); + let mut controls = Vec::new(); + let mut echoes = Vec::new(); + // `ids` is position-indexed, so the short prompt is a prefix of the long one. + for len in [short, long] { + let prompt = ids(len, SALT); + let control = harness.submit_scored(prompt.clone(), 1, Some(TOP_K), Some(TOP_K)); + let drained = harness.steps.drain(control.id(), "prefix"); + echoes.push(drained.prompt_echo.expect("the engine echoes the prompt")); + controls.push(control); + } + let (short_echo, long_echo) = (&echoes[0], &echoes[1]); + let mut compared = 0usize; + let mut worst = 0.0f32; + let mut worst_at = String::new(); + let mut flipped = 0usize; + let mut first_flip = String::new(); + for row in 0..short - 1 { + let (Some(a), Some(b)) = ( + short_echo + .logprobs + .get(row + 1) + .and_then(|entry| entry.as_ref()), + long_echo + .logprobs + .get(row + 1) + .and_then(|entry| entry.as_ref()), + ) else { + continue; + }; + compared += 1; + let (left, right) = (top_of(a), top_of(b)); + for (token, value) in &left { + if let Some(other) = right.get(token) { + let gap = (value - other).abs(); + if gap > worst { + worst = gap; + worst_at = format!("row {row} token {token}"); + } + } + } + let (pa, pb) = (a.top_logprobs[0].0, b.top_logprobs[0].0); + if pa != pb { + flipped += 1; + if first_flip.is_empty() { + first_flip = + format!("row {row}: {short}-token picks {pa}, {long}-token picks {pb}"); + } + } + } + eprintln!( + "prefix: {compared} shared rows; worst shared-token gap {worst:.4} at {worst_at}; \ + {flipped} argmax flips (first {first_flip})" + ); + let refs: Vec<&_> = controls.iter().collect(); + harness.shutdown(&refs); + assert!(compared > 0, "no shared rows were scored"); + // The line is `serve_oracle`'s calibration, not a tight one: scoring the same + // context two ways (there a greedy walk against a single prefill, here a + // 512-row prefill against a 1024-row one) moves the logits by a chaotically + // amplifying amount — measured in-repo at 0.31..5.75 raw logits — because the + // prefill GEMM's tiling depends on the row count. Measured here: 1.69 on the + // four-layer synthetic, 4.05 on the 60-layer 31B. A grosser gap (a page, + // position or rope error) still fails. + assert!( + worst < PREFIX_LINE, + "the same prefix scored in a longer run differs by {worst} at {worst_at} (line \ + {PREFIX_LINE}), so the engine is not prefix-consistent" + ); +} + +/// The prefix gate's line, on the `serve_oracle` scale (`DRIFT_LINE = 12.0` +/// there) rather than the one-rank comparison's `LOGBROB_LINE`: the drift it +/// looks for is the same chaotic shape-dependence, bounded the same way. +const PREFIX_LINE: f32 = 12.0; + +/// Two ranks against the Hugging Face reference, for the size the gate above +/// cannot afford a single-rank control on: a 31B's whole tower is 57 GiB, so its +/// baseline is impossible on one card and the *reference* is the baseline here. +/// The fixture is that checkpoint's own HF dump +/// (`tools/accuracy/dump_gemma4_hf_golden.py`, selected with +/// `PEGAINFER_GEMMA4_GOLDEN`), compared over the teacher-forced prompt rows' +/// top-64 logprobs. +/// +/// What it does not cover: the fixture's layer-boundary probes (the serving path +/// exposes logits, not activations) and its `single` case — row `r` predicts +/// token `r + 1`, so a one-token prompt has no row to score. +#[test] +#[ignore = "needs two GPUs, a 31B fixture in PEGAINFER_GEMMA4_GOLDEN, and --test-threads=1"] +fn the_two_rank_engine_matches_the_hf_reference() { + /// The fixture's rows are the reference's `log_softmax` over softcapped + /// logits. A shared token may sit this far apart: bf16 reduction order and a + /// different attention backend, nothing structural. + const DRIFT_LINE: f32 = 1.0; + // Only the nine-token case is gated. The 1024-token "edge" case runs (the + // concurrent driver carries it), but at window width its drift is the + // shape-dependent chaos `serve_oracle`'s `neutral_scale` already measures at + // 0.31..5.75 raw logits — worst shared-token logprob gap 8.37 here — so a + // strict top-k containment cannot hold and the case needs that gate's + // tolerance-plus-top-1-share discipline. See docs/models/gemma4/tp.md "Known + // bounds"; the calibration is the follow-up. + const CASES: [&str; 1] = ["short"]; + + let (device, peer) = devices(); + assert_ne!(device, peer, "TP2 needs two distinct device ordinals"); + let dir = crate::testkit::model_path(); + // `golden_bytes` holds the fixture against the checkpoint's own file digests + // and fails loudly on a mismatch, so the wrong pair cannot be compared. + let (bytes, manifest) = golden_bytes(&dir); + let fixture = safetensors::SafeTensors::deserialize(&bytes).expect("fixture"); + eprintln!( + "hf: {} at revision {}, {}", + manifest["model_class"], manifest["revision"], manifest["dtypes"] + ); + + let mut harness = launch_with(&tp2_options(device, peer), &as_refs(&envelope_overrides())); + let mut controls = Vec::new(); + let mut compared = 0usize; + let mut same_pick = 0usize; + let mut outside = 0usize; + let mut offenders: Vec = Vec::new(); + let mut worst = 0.0f32; + let mut worst_at = String::new(); + for case in CASES { + let (_, tokens) = u32_tensor(&fixture, &format!("{case}_tokens")); + let (shape, ids) = i32_tensor(&fixture, &format!("{case}_topk_ids")); + let (_, lps) = f32_tensor(&fixture, &format!("{case}_topk_logprobs")); + let top_k = shape[1]; + assert_eq!( + shape[0], + tokens.len(), + "{case}: one reference row per token" + ); + let control = harness.submit_scored(tokens.clone(), 1, Some(top_k), Some(top_k)); + let drained = harness.steps.drain(control.id(), case); + let echo = drained.prompt_echo.expect("the engine echoes the prompt"); + for row in 0..tokens.len() - 1 { + // Row 0 of the echo is the head's own placeholder, so the row that + // predicts token `r + 1` sits at `r + 1`. + let ours = echo + .logprobs + .get(row + 1) + .and_then(|entry| entry.as_ref()) + .unwrap_or_else(|| panic!("{case}: row {row} was not scored")); + let theirs: Vec<(u32, f32)> = (0..top_k) + .map(|k| (ids[row * top_k + k] as u32, lps[row * top_k + k])) + .collect(); + assert!( + theirs.iter().all(|(_, value)| value.is_finite()), + "{case} row {row}: the reference fixture carries a non-finite logprob" + ); + assert_finite(ours, &format!("{case} row {row} (engine)")); + let ours_top = top_of(ours); + let (pa, pb) = (ours.top_logprobs[0].0, theirs[0].0); + compared += 1; + if pa == pb { + same_pick += 1; + } else { + // A differing pick has to be a near-tie: each pick inside the + // other run's top-k. Every offender is collected rather than + // panicking on the first, so one run names them all. + if !theirs.iter().any(|(token, _)| *token == pa) { + outside += 1; + if offenders.len() < 10 { + offenders.push(format!( + "{case} row {row}: engine pick {pa} at {:.4}, reference top-1 {pb} at \ + {:.4}, reference top-{top_k} floor {:.4}", + ours.top_logprobs[0].1, + theirs[0].1, + theirs.last().expect("non-empty top-k").1 + )); + } + } + if !ours_top.contains_key(&pb) { + outside += 1; + if offenders.len() < 10 { + offenders.push(format!( + "{case} row {row}: reference pick {pb} outside the engine's top-{top_k}" + )); + } + } + } + for (token, value) in &theirs { + if let Some(other) = ours_top.get(token) { + let gap = (value - other).abs(); + if gap > worst { + worst = gap; + worst_at = format!("{case} row {row} token {token}"); + } + } + } + } + controls.push(control); + } + eprintln!( + "hf: {same_pick}/{compared} rows keep the reference's pick; {outside} out-of-top-k; worst \ + shared-token logprob gap {worst:.4} at {worst_at}" + ); + for offender in &offenders { + eprintln!("hf: outside: {offender}"); + } + let controls: Vec<&_> = controls.iter().collect(); + harness.shutdown(&controls); + assert!( + outside == 0, + "{outside} rows have a pick outside the other's top-k; first up to 10: {offenders:#?}" + ); + assert!( + worst < DRIFT_LINE, + "the two-rank engine and the reference differ on {worst_at} by {worst} (line {DRIFT_LINE})" + ); +} diff --git a/pegainfer-gemma4/src/engine/lane_tests.rs b/pegainfer-gemma4/src/engine/lane_tests.rs index d06f238af..e4768cb4d 100644 --- a/pegainfer-gemma4/src/engine/lane_tests.rs +++ b/pegainfer-gemma4/src/engine/lane_tests.rs @@ -130,13 +130,15 @@ pub(super) fn knob_table<'a>( } pub(super) fn launch(overrides: &[(&str, &str)]) -> Harness { + launch_with(&EngineLoadOptions::default(), overrides) +} + +/// Start a real engine from explicit load options — the seam a tensor-parallel +/// gate uses to ask for two device ordinals and a topology. +pub(super) fn launch_with(options: &EngineLoadOptions, overrides: &[(&str, &str)]) -> Harness { let dir = crate::testkit::model_path(); - let engine = super::start_with_knobs( - Path::new(&dir), - &EngineLoadOptions::default(), - &knob_table(overrides), - ) - .expect("engine start"); + let engine = super::start_with_knobs(Path::new(&dir), options, &knob_table(overrides)) + .expect("engine start"); Harness::from_engine(engine) } @@ -145,7 +147,7 @@ pub(super) fn load_state(overrides: &[(&str, &str)]) -> anyhow::Result Vec { diff --git a/pegainfer-gemma4/src/kv.rs b/pegainfer-gemma4/src/kv.rs index ec11153a9..32611c528 100644 --- a/pegainfer-gemma4/src/kv.rs +++ b/pegainfer-gemma4/src/kv.rs @@ -210,20 +210,27 @@ impl SlidingLocalKv { } } -pub(crate) struct GemmaKv { +/// One rank's KV families. The per-rank entry points take this, so a step +/// drives exactly the families of the rank it was activated on. +pub(crate) struct RankKv { pub(crate) local: SlidingLocalKv, pub(crate) global: KvState, - /// Distinct for every state built in this process. + /// The request's id, shared by every rank of one multi-rank state. id: u64, } -impl GemmaKv { +/// The id a [`RankKv`] carries until [`GemmaKv::multi`] hands every rank of one +/// request the same one. It is distinct from every id `multi` issues, so a state +/// built outside `multi` can neither collide with — nor fool the steady-decode +/// fingerprint of — a real request. +const NO_REQUEST_ID: u64 = u64::MAX; + +impl RankKv { pub(crate) fn new(local: SlidingLocalKv, global: KvState) -> Self { - static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); Self { local, global, - id: NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed), + id: NO_REQUEST_ID, } } @@ -232,6 +239,90 @@ impl GemmaKv { } } +/// A request's KV across the tensor-parallel group: rank 0's families in +/// `core`, the other ranks' in `twins`. [`Deref`](std::ops::Deref) exposes +/// `core`, so single-rank code and rank-0 steps read `kv.local` / `kv.global` +/// unchanged; the engine drives a rank-`r` step through +/// [`GemmaKv::core_mut`] with that rank's own serve. +pub(crate) struct GemmaKv { + core: RankKv, + twins: Vec, +} + +impl std::ops::Deref for GemmaKv { + type Target = RankKv; + + fn deref(&self) -> &RankKv { + &self.core + } +} + +impl std::ops::DerefMut for GemmaKv { + fn deref_mut(&mut self) -> &mut RankKv { + &mut self.core + } +} + +impl GemmaKv { + pub(crate) fn new(local: SlidingLocalKv, global: KvState) -> Self { + Self::multi(RankKv::new(local, global), Vec::new()) + } + + /// One `RankKv` per rank: `core` is rank 0's, `twins` the rest. Every rank + /// carries the same request id. The caller hands in exactly one twin per + /// extra rank, in rank order — the same world the geometry was built for. + pub(crate) fn multi(mut core: RankKv, mut twins: Vec) -> Self { + static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + let id = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + core.id = id; + for twin in &mut twins { + twin.id = id; + } + Self { core, twins } + } + + /// The family pair `rank` owns, for reading (0 is `core`). A rank past the + /// twins is a caller bug rather than a runtime condition, so say which rank + /// was asked for instead of an index panic. + pub(crate) fn core(&self, rank: usize) -> &RankKv { + if rank == 0 { + &self.core + } else { + match self.twins.get(rank - 1) { + Some(twin) => twin, + None => panic!( + "rank {rank} is outside the {}-rank KV", + self.twins.len() + 1 + ), + } + } + } + + /// The family pair `rank` owns (0 is `core`). + pub(crate) fn core_mut(&mut self, rank: usize) -> &mut RankKv { + if rank == 0 { + &mut self.core + } else { + let twins = self.twins.len() + 1; + match self.twins.get_mut(rank - 1) { + Some(twin) => twin, + None => panic!("rank {rank} is outside the {twins}-rank KV"), + } + } + } + + /// Both halves at once: rank 0's `core` and one `&mut` per twin, so a driver + /// that hands each rank to its own thread gets disjoint borrows. + pub(crate) fn split(&mut self) -> (&mut RankKv, &mut [RankKv]) { + (&mut self.core, &mut self.twins) + } + + /// How many ranks this request's KV spans. + pub(crate) fn world(&self) -> usize { + self.twins.len() + 1 + } +} + /// Pages a family still has to reserve to cover `kv_len` tokens, given what /// its account already holds — the exact frontier account, a ceiling over /// the post-step kv_len. `None` means the account is already past the @@ -245,7 +336,7 @@ fn pages_to_reserve(kv_len: usize, accounted: usize, page_size: usize) -> Option pub(crate) fn admit_tokens( local_pool: &KvPool, global_pool: &KvPool, - kv: &mut GemmaKv, + kv: &mut RankKv, new_tokens: usize, ) -> Result<()> { anyhow::ensure!( diff --git a/pegainfer-gemma4/src/layer.rs b/pegainfer-gemma4/src/layer.rs index 3594d26ba..4350d6897 100644 --- a/pegainfer-gemma4/src/layer.rs +++ b/pegainfer-gemma4/src/layer.rs @@ -22,11 +22,16 @@ use pegainfer_core::tensor::HiddenStates; use crate::config::Gemma4Config; use crate::config::MoeConfig; +use crate::config::TensorParallelConfig; use crate::moe::MoeScratch; use crate::weights::Gemma4Layer; use crate::weights::Linear; use crate::weights::LinearScratch; +/// The NCCL communicator of this rank's tensor-parallel group, or `None` at +/// world size 1, where every collective is a no-op. +pub(crate) type TpComm = cudarc::nccl::safe::Comm; + /// The geometry a layer runs at, read off the validated config — the local /// and global kinds differ only in head width and KV head count. pub(crate) struct LayerGeometry { @@ -40,28 +45,34 @@ pub(crate) struct LayerGeometry { } impl LayerGeometry { - pub(crate) fn local_of(config: &Gemma4Config) -> Self { - Self { + /// This rank's geometry for the sliding family. The residual stream and + /// the head width stay whole — every rank holds whole heads — while the + /// head counts and the MLP width shard. + pub(crate) fn local_of(config: &Gemma4Config, tp: TensorParallelConfig) -> Result { + Ok(Self { hidden_size: config.hidden_size, - intermediate_size: config.intermediate_size, - num_q_heads: config.num_attention_heads, - num_kv_heads: config.num_key_value_heads, + intermediate_size: tp.local_intermediate(config)?, + num_q_heads: tp.local_q_heads(config)?, + num_kv_heads: tp.local_sliding_kv_heads(config)?, head_dim: config.head_dim, rms_norm_eps: config.rms_norm_eps, moe: config.moe, - } + }) } - pub(crate) fn global_of(config: &Gemma4Config) -> Self { - Self { + /// This rank's geometry for the global family. Its KV heads take the + /// two-branch range — shard when the world size divides them, replicate + /// when it does not — which the config resolves. + pub(crate) fn global_of(config: &Gemma4Config, tp: TensorParallelConfig) -> Result { + Ok(Self { hidden_size: config.hidden_size, - intermediate_size: config.intermediate_size, - num_q_heads: config.num_attention_heads, - num_kv_heads: config.num_global_key_value_heads, + intermediate_size: tp.local_intermediate(config)?, + num_q_heads: tp.local_q_heads(config)?, + num_kv_heads: tp.global_kv_head_range(config)?.1, head_dim: config.global_head_dim, rms_norm_eps: config.rms_norm_eps, moe: config.moe, - } + }) } } @@ -251,14 +262,51 @@ impl EpilogueScratch { } } +/// Sum this rank's partial projection across the tensor-parallel group, in +/// place. Only the live rows are handed to NCCL: a decode arena is padded to +/// its power-of-two bucket and the rows past `seq_len` are zero padding no +/// rank reads back, so reducing them would move bytes for nothing. The extent +/// is `hidden_size * seq_len`, the same on every rank of the step; inside a +/// captured graph `seq_len` is the bucket, so the recorded shape stays +/// constant across replays. At world size 1 `comm` is `None` and this is free. +fn all_reduce_rows( + comm: Option<&TpComm>, + geom: &LayerGeometry, + buf: &mut HiddenStates, + seq_len: usize, +) -> Result<()> { + let Some(comm) = comm else { + return Ok(()); + }; + let elems = geom + .hidden_size + .checked_mul(seq_len) + .context("tensor-parallel reduction extent overflows")?; + anyhow::ensure!( + buf.data.len() >= elems, + "a tensor-parallel reduction needs {elems} elements, the buffer holds {}", + buf.data.len() + ); + comm.all_reduce_in_place( + &mut buf.data.slice_mut(..elems), + &cudarc::nccl::safe::ReduceOp::Sum, + ) + .map_err(|e| anyhow::anyhow!("gemma4 tensor-parallel all-reduce failed: {e:?}"))?; + Ok(()) +} + /// Everything downstream of attention — o_proj through the `layer_scalar` /// multiply (applied after both residual adds, not either branch) — is /// identical for both layer kinds; one implementation keeps the two /// forwards' numerics from drifting apart. +/// +/// Under tensor parallelism both projections here are row-parallel sums, so +/// each is reduced across the group before the residual it feeds is formed. pub(crate) fn attention_epilogue_into( ctx: &DeviceContext, layer: &Gemma4Layer, geom: &LayerGeometry, + comm: Option<&TpComm>, x: &HiddenStates, attn: &HiddenStates, scratch: &mut EpilogueScratch, @@ -285,6 +333,10 @@ pub(crate) fn attention_epilogue_into( .attention .o_proj .project_into(ctx, attn, &scratch.linear, &mut scratch.attn_proj)?; + // o_proj is row-parallel: its output is this rank's partial sum over + // the whole hidden width, so it is summed across the group before the + // residual add reads it. + all_reduce_rows(comm, geom, &mut scratch.attn_proj, seq_len)?; // The first normalized value and its residual sum still round to bf16 // before the second reduction reads them. ops::rms_norm_add_rms_norm_round_batch_into( @@ -315,7 +367,7 @@ pub(crate) fn attention_epilogue_into( .mlp .down .project_into(ctx, &scratch.act, &scratch.linear, &mut scratch.down)?; - let feed_forward = match (&layer.moe, &mut scratch.moe) { + let feed_forward: &mut HiddenStates = match (&layer.moe, &mut scratch.moe) { (Some(moe), Some(moe_scratch)) => { crate::moe::moe_into( ctx, @@ -328,11 +380,14 @@ pub(crate) fn attention_epilogue_into( // and has the same shape, so the routed block's result reuses it. &mut scratch.attn_proj, )?; - &scratch.attn_proj + &mut scratch.attn_proj } - (None, _) => &scratch.down, + (None, _) => &mut scratch.down, (Some(_), None) => anyhow::bail!("Gemma 4: a routed layer met a dense epilogue scratch"), }; + // `down` is row-parallel as well, and a routed block's output is summed + // the same way. + all_reduce_rows(comm, geom, feed_forward, seq_len)?; ops::rms_norm_add_scale_batch_into( ctx, feed_forward, diff --git a/pegainfer-gemma4/src/model_line.rs b/pegainfer-gemma4/src/model_line.rs index 1b179485c..721a223a6 100644 --- a/pegainfer-gemma4/src/model_line.rs +++ b/pegainfer-gemma4/src/model_line.rs @@ -4,6 +4,7 @@ use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::LaunchedEngine; use pegainfer_frontend::model_line::LaunchContext; use pegainfer_frontend::model_line::ModelLine; +use pegainfer_frontend::parallel::ParallelConfig; pub static MODEL_LINE: Gemma4Line = Gemma4Line; @@ -19,15 +20,25 @@ impl ModelLine for Gemma4Line { } fn consumed_shared_args(&self) -> &'static [&'static str] { - &["device_ordinal", "cuda_graph"] + &["device_ordinal", "cuda_graph", "tp_size"] } fn launch(&self, ctx: &LaunchContext<'_>) -> anyhow::Result { + // One tensor-parallel rank per device: `--tp-size=2` on two L20s is + // the whole device list. + let tp = ctx.shared.tp_size; + let device_ordinals: Vec = if tp <= 1 { + vec![ctx.shared.device_ordinal] + } else { + (0..tp).collect() + }; + let parallel_config = (tp > 1).then(|| ParallelConfig::new(tp, 1)); crate::start_engine( ctx.model_path, &EngineLoadOptions { enable_cuda_graph: ctx.shared.cuda_graph, - device_ordinals: vec![ctx.shared.device_ordinal], + device_ordinals, + parallel_config, ..EngineLoadOptions::default() }, ) diff --git a/pegainfer-gemma4/src/moe/tests.rs b/pegainfer-gemma4/src/moe/tests.rs index adc68dcfa..8b2a471ca 100644 --- a/pegainfer-gemma4/src/moe/tests.rs +++ b/pegainfer-gemma4/src/moe/tests.rs @@ -468,7 +468,8 @@ fn the_routed_block_matches_the_reference_formulas() { let config = crate::config::Gemma4Config::from_file(&model).expect("config"); let eps = config.rms_norm_eps; let manifest = Manifest::from_config(&config).expect("manifest"); - let geom = LayerGeometry::local_of(&config); + let geom = LayerGeometry::local_of(&config, crate::config::TensorParallelConfig::SINGLE) + .expect("geometry"); let routed = geom.moe.expect("the checkpoint routes"); let hidden = geom.hidden_size; let width = routed.intermediate_size; @@ -476,8 +477,13 @@ fn the_routed_block_matches_the_reference_formulas() { assert_eq!(super::marlin_block(NARROW_EDGE_ROWS * top_k), 16); assert_eq!(super::marlin_block(COARSE_ROWS * top_k), 64); - let weights = - crate::weights::Gemma4Weights::from_safetensors(&model, 0, config).expect("weights"); + let weights = crate::weights::Gemma4Weights::from_safetensors( + &model, + 0, + config, + crate::config::TensorParallelConfig::SINGLE, + ) + .expect("weights"); let ctx = DeviceContext::new_with_device(0).expect("device"); let layer = &weights.layers[0]; let moe = layer.moe.as_ref().expect("layer 0 routes"); diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index dd7ab9042..63e311752 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -30,7 +30,6 @@ use pegainfer_core::tensor::DeviceContext; use pegainfer_core::tensor::DeviceVec; use pegainfer_core::tensor::HiddenStates; -use crate::config::Gemma4Config; use crate::config::LayerKind; use crate::forward::embed_scale_bf16; use crate::forward::logits_tail; @@ -39,10 +38,12 @@ use crate::forward::validate_tokens; use crate::kv::GLOBAL_PAGE_SIZE; use crate::kv::GemmaKv; use crate::kv::LOCAL_PAGE_SIZE; +use crate::kv::RankKv; use crate::kv::SlidingLocalKv; use crate::kv::admit_tokens; use crate::layer::EpilogueScratch; use crate::layer::LayerGeometry; +use crate::layer::TpComm; use crate::layer::attention_epilogue_into; use crate::layer::build_proportional_rope_tables; use crate::weights::Gemma4Layer; @@ -695,7 +696,7 @@ struct SplitKvSpec { } struct SplitKvLaunch<'a> { - metadata: ops::Hd512DecodeMetadata<'a>, + metadata: pegainfer_kernels::ops::Hd512DecodeMetadata<'a>, o_indptr_d: &'a CudaSlice, valid_mask_d: &'a CudaSlice, tmp_v: &'a mut CudaSlice, @@ -754,7 +755,7 @@ impl SplitKvState { last_page_len: &'a CudaSlice, ) -> SplitKvLaunch<'a> { SplitKvLaunch { - metadata: ops::Hd512DecodeMetadata::new( + metadata: pegainfer_kernels::ops::Hd512DecodeMetadata::new( page_indices, page_indptr, last_page_len, @@ -843,6 +844,25 @@ impl StepArena { pub(crate) fn logits_and_ids(&mut self) -> (&mut HiddenStates, &mut CudaSlice) { (&mut self.logits, &mut self.ids) } + + /// The largest power-of-two decode bucket the capture sweep reaches. + pub(crate) fn bucket_ceiling(&self) -> usize { + self.max_rows + } + + /// Restore the serving floor after a capture sweep, which walks every + /// bucket and would otherwise leave `min_bucket` at the ceiling. + pub(crate) fn reset_min_bucket(&mut self) { + self.min_bucket = 1; + } + + /// Release the captured graphs. Called on teardown, before the + /// communicators go away: a captured collective bakes in NCCL kernel + /// launches, and NCCL's communicator abort wedges while a graph that + /// references them is still alive. + pub(crate) fn release_graphs(&mut self) { + self.graphs.clear(); + } } /// The allocators wrap the driver error in a message rather than keeping the @@ -905,7 +925,7 @@ impl GemmaServe { if self.tilelang_global_attn { pegainfer_kernels::ops::gemma4_hd512_decode_split_kv_into } else { - ops::paged_attention_batch_decode_split_kv_hd512_into + pegainfer_kernels::ops::paged_attention_batch_decode_split_kv_hd512_into } } } @@ -944,12 +964,17 @@ impl GlobalAttn { /// as. FlashInfer's decode dispatcher compiles GQA groups {1,2,3,4,8}: a /// dispatchable group passes through whole, and a non-dispatchable group /// over one KV head halves into pseudo-requests — an exact memory identity -/// only because MQA gives every query head the same KV head (the 12B -/// global family's 16 over 1). Anything else fails loud. -pub(crate) fn global_split_factor(config: &Gemma4Config) -> Result { +/// only because MQA gives every query head the same KV head (12B's global +/// family is 16 over 1 at TP1). Anything else fails loud. +/// +/// This reads the *resolved per-rank* geometry, not the config: query heads +/// shard with the world size while KV heads follow their own range, so the +/// same group is 16 at TP1, 8 at TP2, 4 at TP4 — and lands back inside the +/// compiled decode set above TP1, where the split factor returns 1. +pub(crate) fn global_split_factor(geom: &LayerGeometry) -> Result { const DISPATCHABLE: [usize; 5] = [1, 2, 3, 4, 8]; - let q = config.num_attention_heads; - let kv = config.num_global_key_value_heads; + let q = geom.num_q_heads; + let kv = geom.num_kv_heads; anyhow::ensure!( kv > 0 && q.is_multiple_of(kv), "global family of {q} query heads over {kv} KV heads is not a whole GQA group" @@ -968,6 +993,21 @@ pub(crate) fn global_split_factor(config: &Gemma4Config) -> Result { ) } +/// One phase of the decode-graph pre-capture sweep. `Warm` and `Launch` +/// execute, and so enqueue collectives that need every rank's matching call; +/// `Capture` only records. +#[derive(Clone, Copy)] +pub(crate) enum PrecapturePhase { + Warm, + Capture, + Launch, +} + +impl PrecapturePhase { + /// The sweep's phases in order. + pub(crate) const ALL: [Self; 3] = [Self::Warm, Self::Capture, Self::Launch]; +} + /// Everything a serving step needs that outlives requests. pub(crate) struct GemmaServe { /// The weights these pools, rope tables and layer numbering were built @@ -994,12 +1034,25 @@ pub(crate) struct GemmaServe { /// Model layer index -> index within its family's pool layer axis. family_index: Vec, linear: LinearScratch, + /// This rank's tensor-parallel communicator, `None` at world size 1 where + /// every reduction is a no-op. Attached after construction, once every + /// rank's compute stream exists. + tp_comm: Option, } +// SAFETY: A serve is pinned to one CUDA device and driven by one thread at a +// time — the thread that loads it (the load-time warm and the decode-graph +// sweep walk the ranks one after another, never in parallel) and afterwards the +// engine's scheduler thread, which owns the whole `EngineState` and touches a +// rank's serve only while that rank's device is current. Nothing shares a serve +// between threads, so the raw NCCL communicator it may hold never crosses a +// thread while in use. (Same assertions `Qwen3Model` makes.) +unsafe impl Send for GemmaServe {} + impl GemmaServe { fn decode_row( &self, - kv: &GemmaKv, + kv: &RankKv, local_pages: &mut Vec, global_pages: &mut Vec, ) -> Result { @@ -1030,6 +1083,9 @@ impl GemmaServe { ) -> Result { // One source of truth for geometry, rope tables and layer numbering. let config = &weights.config; + // The weights carry the shard they were cut as, so the geometry and the + // pools sized below cannot belong to a different rank than they hold. + let tp = weights.tp; let device_ordinal = ctx.device_ordinal; // One tensor answers for the set: the loader materializes every // weight through a single context. @@ -1039,7 +1095,9 @@ impl GemmaServe { "weights live on device {weights_ordinal} but this context \ allocates on device {device_ordinal}" ); - let global_split_factor = global_split_factor(config)?; + let local_geom = LayerGeometry::local_of(config, tp)?; + let global_geom = LayerGeometry::global_of(config, tp)?; + let global_split_factor = global_split_factor(&global_geom)?; let (mut locals, mut globals) = (0usize, 0usize); let family_index = config .layer_types @@ -1055,11 +1113,14 @@ impl GemmaServe { } }) .collect(); + // Both pools are laid out at the *per-rank* head counts: the geometry + // above is what this rank's kernels read, so a pool sized to the full + // config would disagree with them. let local_pool = KvPool::with_storage( ctx, locals, - config.num_key_value_heads, - config.head_dim, + local_geom.num_kv_heads, + local_geom.head_dim, LOCAL_PAGE_SIZE, local_pages, local_kv_storage, @@ -1067,15 +1128,13 @@ impl GemmaServe { let global_pool = KvPool::with_storage_and_format( ctx, globals, - config.num_global_key_value_heads, - config.global_head_dim, + global_geom.num_kv_heads, + global_geom.head_dim, GLOBAL_PAGE_SIZE, global_pages, KvStorage::Bf16, global_attn.format(config.global_rotary_dim), )?; - let local_geom = LayerGeometry::local_of(config); - let global_geom = LayerGeometry::global_of(config); let (sliding_cos, sliding_sin) = pegainfer_core::rope::precompute_rope( ctx, &RopeTableSpec { @@ -1112,9 +1171,25 @@ impl GemmaServe { family_index, tilelang_global_attn: global_attn.tilelang(), linear, + tp_comm: None, }) } + /// Attach this rank's tensor-parallel communicator. Every decode graph must + /// be captured after this, so the reduction lands inside the graph. + pub(crate) fn attach_tp_comm(&mut self, comm: TpComm) { + self.tp_comm = Some(comm); + } + + /// Take the communicator out, dropping — and so aborting — it. A cudarc + /// NCCL comm aborts on drop, which unblocks a peer already waiting on a + /// call this rank will never issue. Call it once the pair's collective + /// sequence is broken beyond repair; a comm-less step would otherwise + /// reduce to a no-op, so the engine must stop afterwards. + pub(crate) fn detach_tp_comm(&mut self) -> Option { + self.tp_comm.take() + } + /// The global family's prefill, through whichever kernel this engine was /// started with. Both are the same fn type, so a drift between them stops /// compiling rather than computing something else. @@ -1130,7 +1205,7 @@ impl GemmaServe { let attend = if self.tilelang_global_attn { pegainfer_kernels::ops::gemma4_hd512_prefill_varlen_into } else { - ops::batch_prefill_paged_hd512_into + pegainfer_kernels::ops::batch_prefill_paged_hd512_into }; attend( ctx, @@ -1316,6 +1391,14 @@ impl GemmaServe { ) } + /// This rank's families alone, for the rank-`r` step of a multi-rank KV. + pub(crate) fn alloc_rank_kv(&self) -> RankKv { + RankKv::new( + SlidingLocalKv::new(self.local_pool.clone()), + self.global_pool.alloc(), + ) + } + /// Copy a request's post-prefill KV into cache-owned pages — the /// capture half of the conversation-tail prefix cache. Returns `None` /// when either pool cannot spare the pages: capture is strictly @@ -1323,10 +1406,14 @@ impl GemmaServe { /// the copies safe — the prompt's KV writes were enqueued on /// `ctx.stream` before this call, and the copies enqueue after them on /// the same stream. + /// + /// Rank 0's families only: `PEGAINFER_PREFIX_CACHE` is refused under tensor + /// parallelism, so this never runs beside `twins`, and a multi-rank capture + /// would have to take every rank's shard to be resumable. pub(crate) fn capture_checkpoint( &self, ctx: &DeviceContext, - kv: &GemmaKv, + kv: &RankKv, token_ids: &[u32], ) -> Option { // Slack guard: cache pages must never push either pool to the edge, @@ -1471,11 +1558,11 @@ impl GemmaServe { )) } - fn advance_local(&self, kv: &mut GemmaKv, tokens: usize) -> Result<()> { + fn advance_local(&self, kv: &mut RankKv, tokens: usize) -> Result<()> { kv.local.advance_and_release(tokens, self.sliding_window) } - fn check_step_bounds(&self, kv: &GemmaKv, kv_len: usize) -> Result<()> { + fn check_step_bounds(&self, kv: &RankKv, kv_len: usize) -> Result<()> { anyhow::ensure!( kv.local.seq_len() == kv.global.seq_len(), "the two families' frontiers diverged: local {} global {}", @@ -1498,7 +1585,7 @@ impl GemmaServe { fn plan_step( &self, ctx: &DeviceContext, - kv: &GemmaKv, + kv: &RankKv, start_pos: usize, seq_len: usize, ) -> Result { @@ -1587,7 +1674,7 @@ impl GemmaServe { local_page_origin, .. } => { - ops::qkv_norm_rope_paged_prefill_hd256_plain_into( + pegainfer_kernels::ops::qkv_norm_rope_paged_prefill_hd256_plain_into( ctx, projected.q, projected.k, @@ -1616,7 +1703,7 @@ impl GemmaServe { local_origins, global_tables, } => { - ops::qkv_norm_rope_paged_decode_hd256_plain_into( + pegainfer_kernels::ops::qkv_norm_rope_paged_decode_hd256_plain_into( ctx, projected.q, projected.k, @@ -1681,7 +1768,7 @@ impl GemmaServe { } => { // One per-token launch covers every row — each row writes // its position's slot in its own one-page window. - ops::qkv_norm_rope_paged_decode_hd256_plain_into( + pegainfer_kernels::ops::qkv_norm_rope_paged_decode_hd256_plain_into( ctx, projected.q, projected.k, @@ -1727,7 +1814,7 @@ impl GemmaServe { window_left, )?; } else { - ops::batch_prefill_paged_window_hd256_into( + pegainfer_kernels::ops::batch_prefill_paged_window_hd256_into( ctx, &scratch.q_prep, self.local_pool.buffer(), @@ -1741,7 +1828,16 @@ impl GemmaServe { )?; } } - attention_epilogue_into(ctx, layer, geom, x, &scratch.attn, epilogue, out) + attention_epilogue_into( + ctx, + layer, + geom, + self.tp_comm.as_ref(), + x, + &scratch.attn, + epilogue, + out, + ) } #[allow(clippy::too_many_arguments)] @@ -1795,7 +1891,7 @@ impl GemmaServe { global_plan, .. } => { - ops::qk_norm_partial_rope_paged_prefill_hd512_into( + pegainfer_kernels::ops::qk_norm_partial_rope_paged_prefill_hd512_into( ctx, projected.q, projected.k, @@ -1827,7 +1923,7 @@ impl GemmaServe { } PrepRef::Batched { global_tables, .. } => { let split = split.context("batched global decode needs the split-KV state")?; - ops::qk_norm_partial_rope_paged_decode_hd512_into( + pegainfer_kernels::ops::qk_norm_partial_rope_paged_decode_hd512_into( ctx, projected.q, projected.k, @@ -1897,7 +1993,7 @@ impl GemmaServe { let split = split.context("mixed global decode needs the split-KV state")?; // One per-token launch covers every row — each row writes // its position's slot in its own one-page window. - ops::qk_norm_partial_rope_paged_decode_hd512_into( + pegainfer_kernels::ops::qk_norm_partial_rope_paged_decode_hd512_into( ctx, projected.q, projected.k, @@ -1968,7 +2064,16 @@ impl GemmaServe { scratch.attn.seq_len = seq_len; } } - attention_epilogue_into(ctx, layer, geom, x, &scratch.attn, epilogue, out) + attention_epilogue_into( + ctx, + layer, + geom, + self.tp_comm.as_ref(), + x, + &scratch.attn, + epilogue, + out, + ) } /// Expand each global row into `global_split_factor` pseudo requests, build @@ -2014,7 +2119,7 @@ impl GemmaServe { global_split.upload_csr(ctx, &csr) } - fn decode_fingerprint(kvs: &[&mut GemmaKv], padded: usize) -> Option { + fn decode_fingerprint(kvs: &[&mut RankKv], padded: usize) -> Option { let rows = kvs .iter() .map(|kv| { @@ -2039,7 +2144,7 @@ impl GemmaServe { &self, ctx: &DeviceContext, arena: &mut StepArena, - kvs: &[&mut GemmaKv], + kvs: &[&mut RankKv], padded: usize, ) -> Result<()> { let batch = kvs.len(); @@ -2232,7 +2337,7 @@ impl GemmaServe { &self, ctx: &DeviceContext, arena: &'a mut StepArena, - kvs: &mut [&mut GemmaKv], + kvs: &mut [&mut RankKv], tokens: &[u32], ) -> Result<&'a mut HiddenStates> { self.decode_batch_step_inner(ctx, arena, kvs, Some(tokens)) @@ -2244,7 +2349,7 @@ impl GemmaServe { &self, ctx: &DeviceContext, arena: &'a mut StepArena, - kvs: &mut [&mut GemmaKv], + kvs: &mut [&mut RankKv], ) -> Result<&'a mut HiddenStates> { self.decode_batch_step_inner(ctx, arena, kvs, None) } @@ -2253,7 +2358,7 @@ impl GemmaServe { &self, ctx: &DeviceContext, arena: &'a mut StepArena, - kvs: &mut [&mut GemmaKv], + kvs: &mut [&mut RankKv], tokens: Option<&[u32]>, ) -> Result<&'a mut HiddenStates> { let batch = kvs.len(); @@ -2371,8 +2476,8 @@ impl GemmaServe { &self, ctx: &DeviceContext, arena: &'a mut StepArena, - prefills: &mut [(&mut GemmaKv, &[u32])], - decode_kvs: &mut [&mut GemmaKv], + prefills: &mut [(&mut RankKv, &[u32])], + decode_kvs: &mut [&mut RankKv], decode_tokens: &[u32], ) -> Result<&'a mut HiddenStates> { let prompts = prefills.len(); @@ -2601,7 +2706,7 @@ impl GemmaServe { &self, ctx: &DeviceContext, arena: &mut StepArena, - kvs: &[&mut GemmaKv], + kvs: &[&mut RankKv], tokens: Option<&[u32]>, ) -> Result { let batch = kvs.len(); @@ -2656,11 +2761,33 @@ impl GemmaServe { self.step(ctx, &mut kv, &[0])?; let mut bucket = 1usize; while bucket <= arena.max_rows { - arena.min_bucket = bucket; - arena.steady = None; - admit_tokens(&self.local_pool, &self.global_pool, &mut kv, 1)?; - { - let mut kvs: [&mut GemmaKv; 1] = [&mut kv]; + for phase in PrecapturePhase::ALL { + self.precapture_bucket(ctx, arena, &mut kv, bucket, phase)?; + } + bucket *= 2; + } + arena.min_bucket = 1; + ctx.sync() + } + + /// One bucket of the sweep, one phase at a time. `Capture` only records, but + /// `Warm` and `Launch` execute, and the collectives they enqueue need every + /// rank's matching call — so a driver with more than one rank interleaves + /// the phases across ranks rather than running one rank's whole sweep. + pub(crate) fn precapture_bucket( + &self, + ctx: &DeviceContext, + arena: &mut StepArena, + kv: &mut RankKv, + bucket: usize, + phase: PrecapturePhase, + ) -> Result<()> { + match phase { + PrecapturePhase::Warm => { + arena.min_bucket = bucket; + arena.steady = None; + admit_tokens(&self.local_pool, &self.global_pool, kv, 1)?; + let kvs: [&mut RankKv; 1] = [&mut *kv]; let padded = self.prepare_decode_step(ctx, arena, &kvs, Some(&[0]))?; let StepArena { tower, @@ -2672,7 +2799,6 @@ impl GemmaServe { ids, head_normed, logits, - graphs, .. } = arena; self.decode_gpu_body( @@ -2687,13 +2813,28 @@ impl GemmaServe { local_origins, head_normed, logits, - )?; - graphs[decode_bucket_slot(padded)].capture_only(ctx, || { + ) + } + PrecapturePhase::Capture => { + let StepArena { + tower, + local_plan, + global_tables, + global_split, + local_split, + local_origins, + ids, + head_normed, + logits, + graphs, + .. + } = arena; + graphs[decode_bucket_slot(bucket)].capture_only(ctx, || { self.decode_gpu_body( ctx, tower, ids, - padded, + bucket, local_plan, global_tables, global_split, @@ -2702,13 +2843,14 @@ impl GemmaServe { head_normed, logits, ) - })?; + }) + } + PrecapturePhase::Launch => { + let mut kvs: [&mut RankKv; 1] = [&mut *kv]; self.decode_batch_step(ctx, arena, &mut kvs, &[0])?; + Ok(()) } - bucket *= 2; } - arena.min_bucket = 1; - ctx.sync() } /// The host-side prologue every single-request prompt pass shares. All of it @@ -2717,7 +2859,7 @@ impl GemmaServe { fn prepare_single( &self, ctx: &DeviceContext, - kv: &GemmaKv, + kv: &RankKv, tokens: &[u32], what: &str, ) -> Result { @@ -2786,7 +2928,7 @@ impl GemmaServe { pub(crate) fn step( &self, ctx: &DeviceContext, - kv: &mut GemmaKv, + kv: &mut RankKv, tokens: &[u32], ) -> Result { self.step_scoring(ctx, kv, tokens, None) @@ -2794,14 +2936,32 @@ impl GemmaServe { /// [`GemmaServe::step`] that also hands `score` the head's logits for /// every row but the last, `PROMPT_SCORE_ROWS` rows at a time with the - /// first row's index: row `r` predicts token `r + 1`. + /// first row's index: row `r` predicts token `r + 1`. The readback and the + /// tower are split so a driver sharing one thread across ranks can launch + /// every rank's tower before rank 0's readback waits on its collective. pub(crate) fn step_scoring( &self, ctx: &DeviceContext, - kv: &mut GemmaKv, + kv: &mut RankKv, tokens: &[u32], score: Option>, ) -> Result { + let mut tower = self.launch_prompt_tower(ctx, kv, tokens)?; + if let Some(score) = score { + self.score_prompt_tower(ctx, &mut tower, score)?; + } + Ok(tower.into_logits()) + } + + /// Run the prompt's tower and head, and hold the rows a later + /// [`GemmaServe::score_prompt_tower`] reads. Nothing is read back to the + /// host here, so the launch never waits on the rank's collective. + pub(crate) fn launch_prompt_tower( + &self, + ctx: &DeviceContext, + kv: &mut RankKv, + tokens: &[u32], + ) -> Result { let seq_len = tokens.len(); let mut pass = self.prepare_single(ctx, kv, tokens, "step")?; log::debug!( @@ -2811,36 +2971,6 @@ impl GemmaServe { kv.global.held_pages() ); let src = self.run_single_tower(ctx, &mut pass, seq_len)?; - if let Some(score) = score - && seq_len > 1 - { - let rows = PROMPT_SCORE_ROWS.min(seq_len - 1); - let mut hidden = HiddenStates::zeros(ctx, self.local_geom.hidden_size, rows)?; - let mut normed = HiddenStates::zeros(ctx, self.local_geom.hidden_size, rows)?; - let mut logits = HiddenStates::zeros(ctx, self.weights.embed_tokens.rows, rows)?; - for start in (0..seq_len - 1).step_by(rows) { - let n = rows.min(seq_len - 1 - start); - hidden.seq_len = n; - ops::copy_hidden_token_range_into( - ctx, - &pass.tower.hidden[src], - start, - &mut hidden, - 0, - n, - )?; - logits_tail_into( - ctx, - &self.weights, - &hidden, - self.local_geom.rms_norm_eps, - self.final_logit_softcapping, - &mut normed, - &mut logits, - )?; - score(&mut logits, start)?; - } - } let logits = logits_tail( ctx, &self.weights, @@ -2855,7 +2985,53 @@ impl GemmaServe { // leave one family ahead of the other. self.advance_local(kv, seq_len)?; kv.global.advance(seq_len); - Ok(logits) + Ok(PromptTower { + pass, + src, + seq_len, + logits, + }) + } + + /// The per-row readback a scored prompt owes rank 0. It blocks on this + /// rank's collectives, so call it only once every rank's tower is launched. + pub(crate) fn score_prompt_tower( + &self, + ctx: &DeviceContext, + tower: &mut PromptTower, + score: PromptScorer<'_>, + ) -> Result<()> { + let seq_len = tower.seq_len; + if seq_len <= 1 { + return Ok(()); + } + let rows = PROMPT_SCORE_ROWS.min(seq_len - 1); + let mut hidden = HiddenStates::zeros(ctx, self.local_geom.hidden_size, rows)?; + let mut normed = HiddenStates::zeros(ctx, self.local_geom.hidden_size, rows)?; + let mut logits = HiddenStates::zeros(ctx, self.weights.embed_tokens.rows, rows)?; + for start in (0..seq_len - 1).step_by(rows) { + let n = rows.min(seq_len - 1 - start); + hidden.seq_len = n; + ops::copy_hidden_token_range_into( + ctx, + &tower.pass.tower.hidden[tower.src], + start, + &mut hidden, + 0, + n, + )?; + logits_tail_into( + ctx, + &self.weights, + &hidden, + self.local_geom.rms_norm_eps, + self.final_logit_softcapping, + &mut normed, + &mut logits, + )?; + score(&mut logits, start)?; + } + Ok(()) } /// Whole-prompt prefill left in flight on the active stream — the @@ -2873,7 +3049,7 @@ impl GemmaServe { pub(crate) fn prefill_into_logits( &self, ctx: &DeviceContext, - kv: &mut GemmaKv, + kv: &mut RankKv, tokens: &[u32], ) -> Result { let seq_len = tokens.len(); @@ -2903,11 +3079,27 @@ impl GemmaServe { /// The deferred append-then-attend release for an overlapped prefill: /// call once its completion event has fired, before the request joins /// the decode batch. - pub(crate) fn release_prefill_window(&self, kv: &mut GemmaKv) -> Result<()> { + pub(crate) fn release_prefill_window(&self, kv: &mut RankKv) -> Result<()> { self.advance_local(kv, 0) } } +/// A launched prompt tower and the head rows a later scored pass reads. Kept +/// alive across the peer ranks' tower launches so rank 0's own readback — the +/// only part that waits on its collective — runs after every rank is in flight. +pub(crate) struct PromptTower { + pass: SinglePass, + src: usize, + seq_len: usize, + logits: HiddenStates, +} + +impl PromptTower { + pub(crate) fn into_logits(self) -> HiddenStates { + self.logits + } +} + /// What one single-request prompt pass allocates before any layer runs and /// keeps alive until its kernels retire. struct SinglePass { diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index 1c6f3e600..44b4f7acb 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -31,9 +31,14 @@ fn stack_with_storage( storage: KvStorage, ) -> (DeviceContext, GemmaServe, String) { let dir = model_path(); - let config = Gemma4Config::from_file(&dir).expect("config"); - let weights = - Gemma4Weights::from_safetensors(&dir, 0, config).expect("load checkpoint weights"); + let config = crate::config::Gemma4Config::from_file(&dir).expect("config"); + let weights = Gemma4Weights::from_safetensors( + &dir, + 0, + config, + crate::config::TensorParallelConfig::SINGLE, + ) + .expect("load checkpoint weights"); let ctx = DeviceContext::new_with_device(0).expect("device context"); // The oracle measures the incumbent kernel; the opt-in one has its own // gate. @@ -1208,7 +1213,7 @@ fn mixed_gate_decode_rounds( } let tokens = gate_step_tokens(serve, lanes); let (vocab, host) = { - let mut kvs: Vec<&mut GemmaKv> = lanes.iter_mut().map(|(_, kv, _)| kv).collect(); + let mut kvs: Vec<&mut RankKv> = lanes.iter_mut().map(|(_, kv, _)| &mut **kv).collect(); let logits = serve .decode_batch_step(ctx, arena, &mut kvs, &tokens) .expect("batched decode"); @@ -1272,10 +1277,10 @@ fn assert_mixed_admissions_match_serial(ctx: &DeviceContext, serve: &GemmaServe) let mut kv_c = gate_admit_kv(serve, &prompts[2], "prompt c"); let tokens = gate_step_tokens(serve, &mut lanes); let (vocab, host) = { - let mut kvs: Vec<&mut GemmaKv> = lanes.iter_mut().map(|(_, kv, _)| kv).collect(); + let mut kvs: Vec<&mut RankKv> = lanes.iter_mut().map(|(_, kv, _)| &mut **kv).collect(); let mut prefills = [ - (&mut kv_b, prompts[1].as_slice()), - (&mut kv_c, prompts[2].as_slice()), + (&mut *kv_b, prompts[1].as_slice()), + (&mut *kv_c, prompts[2].as_slice()), ]; let logits = serve .mixed_prefill_decode_step(ctx, &mut arena, &mut prefills, &mut kvs, &tokens) @@ -1363,8 +1368,9 @@ fn assert_mixed_window_crossing_matches_serial(ctx: &DeviceContext, serve: &Gemm // window inside the mixed step. let tokens = gate_step_tokens(serve, &mut lanes); let (vocab, host) = { - let mut kvs: Vec<&mut GemmaKv> = lanes.iter_mut().map(|(_, kv, _)| kv).collect(); - let mut prefills = [(&mut kv, long_prompt.as_slice())]; + let mut kvs: Vec<&mut RankKv> = + lanes.iter_mut().map(|(_, kv, _)| &mut **kv).collect(); + let mut prefills = [(&mut *kv, long_prompt.as_slice())]; let logits = serve .mixed_prefill_decode_step(ctx, &mut arena, &mut prefills, &mut kvs, &tokens) .expect("mixed step"); @@ -1487,10 +1493,10 @@ fn fp8_plain_mixed_walk(ctx: &DeviceContext, serve: &GemmaServe) { let mut kv_c = gate_admit_kv(serve, &prompts[2], "prompt c"); let tokens = gate_step_tokens(serve, &mut lanes); let (vocab, host) = { - let mut kvs: Vec<&mut GemmaKv> = lanes.iter_mut().map(|(_, kv, _)| kv).collect(); + let mut kvs: Vec<&mut RankKv> = lanes.iter_mut().map(|(_, kv, _)| &mut **kv).collect(); let mut prefills = [ - (&mut kv_b, prompts[1].as_slice()), - (&mut kv_c, prompts[2].as_slice()), + (&mut *kv_b, prompts[1].as_slice()), + (&mut *kv_c, prompts[2].as_slice()), ]; let logits = serve .mixed_prefill_decode_step(ctx, &mut arena, &mut prefills, &mut kvs, &tokens) @@ -1556,8 +1562,8 @@ fn fp8_window_mixed_walk(ctx: &DeviceContext, serve: &GemmaServe) { let mut kv_long = gate_admit_kv(serve, &long_prompt, "long prompt"); let tokens = gate_step_tokens(serve, &mut lanes); let (vocab, host) = { - let mut kvs: Vec<&mut GemmaKv> = lanes.iter_mut().map(|(_, kv, _)| kv).collect(); - let mut prefills = [(&mut kv_long, long_prompt.as_slice())]; + let mut kvs: Vec<&mut RankKv> = lanes.iter_mut().map(|(_, kv, _)| &mut **kv).collect(); + let mut prefills = [(&mut *kv_long, long_prompt.as_slice())]; let logits = serve .mixed_prefill_decode_step(ctx, &mut arena, &mut prefills, &mut kvs, &tokens) .expect("window-crossing mixed step"); @@ -1855,7 +1861,7 @@ fn decode_serving( token: u32, ) -> Result> { admit_tokens(&serve.local_pool, &serve.global_pool, kv, 1)?; - let mut borrowed: [&mut GemmaKv; 1] = [kv]; + let mut borrowed: [&mut RankKv; 1] = [&mut **kv]; let logits = serve.decode_batch_step(ctx, arena, &mut borrowed, &[token])?; logits.to_host(ctx) } @@ -1965,7 +1971,8 @@ fn a_ragged_batch_does_not_depend_on_row_order() { for kv in &mut kvs { admit_tokens(&serve.local_pool, &serve.global_pool, kv, 1).expect("admit token"); } - let mut slots: Vec> = kvs.iter_mut().map(Some).collect(); + let mut slots: Vec> = + kvs.iter_mut().map(|kv| Some(&mut **kv)).collect(); let mut borrowed = Vec::with_capacity(order.len()); let mut tokens = Vec::with_capacity(order.len()); for &request in order { diff --git a/pegainfer-gemma4/src/weights.rs b/pegainfer-gemma4/src/weights.rs index c0ee19792..aed0b0796 100644 --- a/pegainfer-gemma4/src/weights.rs +++ b/pegainfer-gemma4/src/weights.rs @@ -15,11 +15,16 @@ use pegainfer_kernels::ops::W4a16Matrix; use pegainfer_kernels::ops::W4a16Scratch; use crate::config::Gemma4Config; +use crate::config::TensorParallelConfig; mod load; pub(crate) struct Gemma4Weights { pub(crate) config: Gemma4Config, + /// The shard these weights are, and the only source of the tensor-parallel + /// geometry: a serve built around them reads its rank and world size here + /// instead of taking them as a second argument, so the two cannot disagree. + pub(crate) tp: TensorParallelConfig, pub(crate) embed_tokens: DeviceMatrix, pub(crate) norm: DeviceVec, pub(crate) layers: Vec, diff --git a/pegainfer-gemma4/src/weights/load.rs b/pegainfer-gemma4/src/weights/load.rs index f222f779b..55fceae49 100644 --- a/pegainfer-gemma4/src/weights/load.rs +++ b/pegainfer-gemma4/src/weights/load.rs @@ -26,6 +26,9 @@ use super::Gemma4Weights; use super::Linear; use super::StackedProjection; use crate::config::Gemma4Config; +use crate::config::LayerKind; +use crate::config::TensorParallelConfig; +use crate::config::head_dim_of; use crate::manifest::schema::ExpertTensors; use crate::manifest::schema::Manifest; use crate::manifest::schema::Matrix2d; @@ -110,27 +113,59 @@ struct RecordedPlan { layers: Vec, } -fn record_matrix(loader: &mut StagedWeightLoader, tensor: &Matrix2d) -> Result { - loader.matrix(&tensor.name, tensor.rows, tensor.cols) +/// How much of a checkpoint matrix this rank stages. `Whole` is the +/// single-rank path and must stay reachable so TP1 numerics cannot move. +#[derive(Clone, Copy)] +enum Shard { + Whole, + /// A column range: the input-sharded linears (o_proj, down). + Cols { + offset: usize, + cols: usize, + }, } -/// The rows of `parts` stacked in order into one device matrix. -fn record_stacked<'m>( +/// `Whole` at world size 1 so the single-rank load path is untouched. +fn col_shard(tp: TensorParallelConfig, offset: usize, cols: usize) -> Shard { + if tp.is_single() { + Shard::Whole + } else { + Shard::Cols { offset, cols } + } +} + +fn record_matrix( loader: &mut StagedWeightLoader, - parts: impl IntoIterator, + tensor: &Matrix2d, + shard: Shard, +) -> Result { + match shard { + Shard::Whole => loader.matrix(&tensor.name, tensor.rows, tensor.cols), + Shard::Cols { offset, cols } => { + loader.col_shard(&tensor.name, tensor.rows, tensor.cols, offset, cols) + } + } +} + +/// The rows of `parts` stacked in order into one device matrix. Each part +/// carries its own `(row_offset, rows)` so a rank stages its own head run from +/// every input; at world size 1 those ranges are the whole matrix, which is the +/// path this loader has always taken. +fn record_stacked( + loader: &mut StagedWeightLoader, + parts: &[(&Matrix2d, usize, usize)], ) -> Result { - let parts: Vec<&Matrix2d> = parts.into_iter().collect(); let cols = parts .first() - .map(|part| part.cols) + .map(|(part, _, _)| part.cols) .ok_or_else(|| anyhow::anyhow!("Gemma 4: a stacked load needs at least one part"))?; let fused: Vec = parts .iter() - .map(|part| FusedPart { + .map(|(part, offset, rows)| FusedPart { name: part.name.as_str(), src_rows: part.rows, - row_offset: 0, - rows: part.rows, + row_offset: *offset, + rows: *rows, }) .collect(); loader.fused_rows(cols, &fused) @@ -202,6 +237,9 @@ fn tensor_bytes<'a>(shards: &'a [SafeTensors<'a>], name: &str) -> Result<&'a [u8 /// Upload every layer's W4A16 linears in the GEMMs' layout: one entry per /// layer, empty on a bf16 checkpoint. +/// +/// Reached only at world size 1: `validate_for` refuses a W4A16 checkpoint under +/// tensor parallelism precisely because these linears are staged whole. fn upload_w4a16( ctx: &DeviceContext, shards: &[SafeTensors], @@ -317,6 +355,9 @@ struct StackedExperts { /// Upload every routed layer's experts. Returns one entry per layer, empty on /// the sizes that do not route. +/// +/// Reached only at world size 1: a routed checkpoint is refused under tensor +/// parallelism, which shards neither the experts nor these whole-tensor tiles. fn upload_experts( ctx: &DeviceContext, shards: &[SafeTensors], @@ -478,12 +519,33 @@ fn record_plan( loader: &mut StagedWeightLoader, shards: &[SafeTensors], manifest: &Manifest, + config: &Gemma4Config, + tp: TensorParallelConfig, ) -> Result { - let embed_tokens = record_matrix(loader, &manifest.embed_tokens)?; + let embed_tokens = record_matrix(loader, &manifest.embed_tokens, Shard::Whole)?; let norm = record_vector(loader, &manifest.norm)?; + let local_intermediate = tp.local_intermediate(config)?; + // The global family's KV heads take the two-branch range the config + // resolves: shard when the world size divides them, else replicate. + let (global_kv_offset, _) = tp.global_kv_head_range(config)?; let mut layers = Vec::with_capacity(manifest.layers.len()); - for layer in &manifest.layers { + for (index, layer) in manifest.layers.iter().enumerate() { let attention = &layer.attention; + let kind = config.layer_types[index]; + // Query rows are a contiguous run of this rank's query heads; the same + // run shards o_proj's columns, which is what makes the sum-reduce + // exact. + let local_q_dim = tp.local_q_dim(config, kind)?; + let q_offset = tp.rank * local_q_dim; + // K rows follow their family's own range (V, on sliding layers, matches + // K): a shard of the sliding family's KV heads, or the global family's + // two-branch range. + let local_kv_dim = tp.local_kv_dim(config, kind)?; + let kv_offset = match kind { + LayerKind::Sliding => tp.rank * local_kv_dim, + LayerKind::Global => global_kv_offset * head_dim_of(config, kind), + }; + let inter_offset = tp.rank * local_intermediate; layers.push(LayerSlots { input_layernorm: record_vector(loader, &layer.input_layernorm)?, post_attention_layernorm: record_vector(loader, &layer.post_attention_layernorm)?, @@ -493,16 +555,32 @@ fn record_plan( linears: if attention.q_proj.w4a16.is_some() { None } else { + let mut qkv: Vec<(&Matrix2d, usize, usize)> = vec![ + (&attention.q_proj, q_offset, local_q_dim), + (&attention.k_proj, kv_offset, local_kv_dim), + ]; + if let Some(v_proj) = attention.v_proj.as_ref() { + qkv.push((v_proj, kv_offset, local_kv_dim)); + } Some(LinearSlots { - qkv: record_stacked( + qkv: record_stacked(loader, &qkv)?, + o_proj: record_matrix( + loader, + &attention.o_proj, + col_shard(tp, q_offset, local_q_dim), + )?, + gate_up: record_stacked( + loader, + &[ + (&layer.mlp.gate, inter_offset, local_intermediate), + (&layer.mlp.up, inter_offset, local_intermediate), + ], + )?, + down: record_matrix( loader, - [&attention.q_proj, &attention.k_proj] - .into_iter() - .chain(attention.v_proj.as_ref()), + &layer.mlp.down, + col_shard(tp, inter_offset, local_intermediate), )?, - o_proj: record_matrix(loader, &attention.o_proj)?, - gate_up: record_stacked(loader, [&layer.mlp.gate, &layer.mlp.up])?, - down: record_matrix(loader, &layer.mlp.down)?, }) }, q_norm: record_vector(loader, &attention.q_norm)?, @@ -524,7 +602,7 @@ fn record_plan( loader, &moe.post_feedforward_layernorm_2, )?, - router_proj: record_matrix(loader, &moe.router.proj)?, + router_proj: record_matrix(loader, &moe.router.proj, Shard::Whole)?, router_scale: record_vector(loader, &moe.router.scale)?, router_per_expert_scale: record_vector( loader, @@ -547,6 +625,7 @@ fn materialize( loader: &mut StagedWeightLoader, plan: RecordedPlan, config: Gemma4Config, + tp: TensorParallelConfig, experts: Vec>, w4a16: Vec>, ) -> Result { @@ -626,6 +705,7 @@ fn materialize( }) .collect::>>()?, config, + tp, }) } @@ -654,6 +734,7 @@ impl Gemma4Weights { model_path: &str, device_ordinal: usize, config: Gemma4Config, + tp: TensorParallelConfig, ) -> Result { let started = Instant::now(); let manifest = Manifest::from_config(&config)?; @@ -671,7 +752,7 @@ impl Gemma4Weights { let mut loader = StagedWeightLoader::new(&ctx, &shards, &weight_map)?; let recording = Instant::now(); - let plan = record_plan(&mut loader, &shards, &manifest)?; + let plan = record_plan(&mut loader, &shards, &manifest, &config, tp)?; let record_api_wall_ms = elapsed_ms(recording); let uploading = Instant::now(); @@ -680,7 +761,7 @@ impl Gemma4Weights { let experts = upload_experts(&ctx, &shards, &manifest)?; let w4a16 = upload_w4a16(&ctx, &shards, &manifest)?; - let weights = materialize(&mut loader, plan, config, experts, w4a16)?; + let weights = materialize(&mut loader, plan, config, tp, experts, w4a16)?; drop(loader); drop(prefetch); let device_free_bytes = free_device_bytes()?; @@ -770,10 +851,15 @@ mod tests { let staged_path = staged.path().to_str().context("temp path is not UTF-8")?; let parsed = Gemma4Config::from_file(staged_path).expect("config"); - let err = Gemma4Weights::from_safetensors(staged_path, UNOPENABLE_DEVICE, parsed) - .err() - .context("a config that disagrees with the checkpoint was accepted")? - .to_string(); + let err = Gemma4Weights::from_safetensors( + staged_path, + UNOPENABLE_DEVICE, + parsed, + TensorParallelConfig::SINGLE, + ) + .err() + .context("a config that disagrees with the checkpoint was accepted")? + .to_string(); let layers = config["text_config"]["layer_types"] .as_array() diff --git a/pegainfer-k3/src/executor/forward/gemm.rs b/pegainfer-k3/src/executor/forward/gemm.rs index 255702785..8dfc8719f 100644 --- a/pegainfer-k3/src/executor/forward/gemm.rs +++ b/pegainfer-k3/src/executor/forward/gemm.rs @@ -29,6 +29,7 @@ use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; use half::bf16; use pegainfer_kernels::ffi; +use pegainfer_kernels::ops::ffi_status_detail; use pegainfer_kernels::tensor::DeviceContext; use pegainfer_kernels::tensor::DeviceMatrix; use pegainfer_kernels::tensor::active_cu_stream; @@ -118,11 +119,7 @@ pub(crate) fn k3_gemm_partial( ensure!( status == 0, "K3 dense GEMM ({out_features}x{k}, rows={rows}) failed: {}", - if status >= 100_000 { - format!("cublasStatus={}", status - 100_000) - } else { - format!("cudaError={status}") - } + ffi_status_detail(status) ); Ok(()) } diff --git a/pegainfer-kernels/KERNELS.md b/pegainfer-kernels/KERNELS.md index 57b207b2d..61fdb4350 100644 --- a/pegainfer-kernels/KERNELS.md +++ b/pegainfer-kernels/KERNELS.md @@ -227,13 +227,13 @@ reference implementation. ## Non-Qwen3 Compatibility -The crate still builds CUDA/Triton symbols needed by the current root binary: +The `qwen35` feature builds the Qwen3.5 CUDA/Triton symbols: -- Qwen3.5 HD256 full-attention kernels: `csrc/qwen35/prefill_attention_hd256.cu`, `csrc/shared/paged_attention.cu`. +- Qwen3.5 HD256 full-attention kernels: `csrc/qwen35/prefill_attention_hd256.cu`, `csrc/qwen35/paged_attention_hd256.cu`. - Qwen3.5 linear-attention decode kernels: `csrc/qwen35/conv1d.cu`, `csrc/qwen35/gated_delta_rule.cu`. - Qwen3.5 chunk-wise GDR prefill Triton AOT kernels: `tools/triton/gated_delta_rule_chunkwise_kernels.py`. -These are preserved for build compatibility. They are not part of the Qwen3-4B Phase 1 API surface. +They are not part of the Qwen3-4B Phase 1 API surface. ## Editing Rule diff --git a/pegainfer-kernels/build.rs b/pegainfer-kernels/build.rs index 24cc523fe..eee6a9e7e 100644 --- a/pegainfer-kernels/build.rs +++ b/pegainfer-kernels/build.rs @@ -304,7 +304,12 @@ fn nvcc_job_count() -> usize { fn nvcc_task_priority(cu_file: &Path) -> usize { match cu_file.file_stem().and_then(|stem| stem.to_str()) { Some("paged_attention") => 0, - Some("paged_attention_hd512" | "flashinfer_sampling") => 1, + Some( + "paged_attention_hd512" + | "paged_attention_hd256" + | "paged_attention_window_hd256" + | "flashinfer_sampling", + ) => 1, Some("flashinfer_top1") => 2, Some("flashinfer_norm") => 3, _ => 10, @@ -752,6 +757,17 @@ fn is_gemma4_source(csrc_dir: &Path, path: &Path) -> bool { } } +/// Qwen3.5 model-local CUDA: the linear-attention recurrence and the +/// head_dim 256 full-attention path. +fn is_qwen35_source(csrc_dir: &Path, path: &Path) -> bool { + match path.strip_prefix(csrc_dir) { + Ok(relative) => relative + .components() + .any(|part| part.as_os_str() == "qwen35"), + Err(_) => false, + } +} + /// GLM5.2 model-local CUDA shims over DeepGEMM/FlashMLA contracts. fn is_glm52_source(csrc_dir: &Path, path: &Path) -> bool { match path.strip_prefix(csrc_dir) { @@ -2262,6 +2278,9 @@ fn main() { if !kimi_k2_enabled && is_kimi_k2_source(&csrc_dir, path) { return None; } + if !qwen35_enabled && is_qwen35_source(&csrc_dir, path) { + return None; + } // --- k3 --- if !k3_enabled && is_k3_source(&csrc_dir, path) { return None; @@ -2459,6 +2478,8 @@ fn main() { // Files that include FlashInfer headers (C++17, header-only) if stem == "paged_attention" || stem == "paged_attention_hd512" + || stem == "paged_attention_hd256" + || stem == "paged_attention_window_hd256" || stem == "flashinfer_norm" || stem == "flashinfer_sampling" || stem == "flashinfer_top1" diff --git a/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_attention.cu b/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_attention.cu index 3618316a0..5f044a913 100644 --- a/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_attention.cu +++ b/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_attention.cu @@ -1,4 +1,5 @@ #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -293,44 +294,6 @@ __global__ void decode_attention_kernel( (void)max_seq_len; } -CUresult map_cuda_error(cudaError_t err) { - switch (err) { - case cudaSuccess: - return CUDA_SUCCESS; - case cudaErrorInvalidValue: - case cudaErrorInvalidDevicePointer: - return CUDA_ERROR_INVALID_VALUE; - case cudaErrorInvalidDevice: - return CUDA_ERROR_INVALID_DEVICE; - case cudaErrorInvalidResourceHandle: - return CUDA_ERROR_INVALID_HANDLE; - case cudaErrorMemoryAllocation: - return CUDA_ERROR_OUT_OF_MEMORY; - case cudaErrorNotSupported: - return CUDA_ERROR_NOT_SUPPORTED; - case cudaErrorIllegalAddress: - return CUDA_ERROR_ILLEGAL_ADDRESS; - case cudaErrorLaunchOutOfResources: - return CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES; - case cudaErrorLaunchTimeout: - return CUDA_ERROR_LAUNCH_TIMEOUT; - case cudaErrorLaunchFailure: - return CUDA_ERROR_LAUNCH_FAILED; - case cudaErrorAssert: - return CUDA_ERROR_ASSERT; - case cudaErrorIllegalInstruction: - return CUDA_ERROR_ILLEGAL_INSTRUCTION; - case cudaErrorMisalignedAddress: - return CUDA_ERROR_MISALIGNED_ADDRESS; - case cudaErrorInvalidAddressSpace: - return CUDA_ERROR_INVALID_ADDRESS_SPACE; - case cudaErrorInvalidPc: - return CUDA_ERROR_INVALID_PC; - default: - return CUDA_ERROR_UNKNOWN; - } -} - CUresult consume_last_cuda_error() { cudaError_t err = cudaGetLastError(); return map_cuda_error(err); diff --git a/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_moe.cu b/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_moe.cu index cb0d0be73..0f4f66b12 100644 --- a/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_moe.cu +++ b/pegainfer-kernels/csrc/deepseek_v2_lite/dsv2_lite_moe.cu @@ -1,15 +1,142 @@ #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include +#include #include +extern thread_local cublasHandle_t g_cublas_handle; + namespace { constexpr int kMaxExperts = 64; constexpr int kMaxTopk = 8; +constexpr int kRoutesPerToken = 6; +constexpr int kLocalExperts = kMaxExperts / 2; +constexpr int kMaxRows = 8; constexpr int kRouterThreads = 128; constexpr int kAccumThreads = 256; +__device__ __forceinline__ bool better_prob_choice(float value, int expert, float best_value, + int best_expert) { + return value > best_value || (value == best_value && expert < best_expert); +} + +__device__ __forceinline__ void select_and_sort_routes( + const float *probs, int n_experts, int topk, int *ids, float *weights) { + bool selected[kMaxExperts] = {}; + for (int route = 0; route < topk; ++route) { + int best_expert = n_experts; + float best_value = -CUDART_INF_F; + for (int expert = 0; expert < n_experts; ++expert) { + if (!selected[expert] && + better_prob_choice(probs[expert], expert, best_value, best_expert)) { + best_value = probs[expert]; + best_expert = expert; + } + } + selected[best_expert] = true; + ids[route] = best_expert; + weights[route] = best_expert < n_experts ? probs[best_expert] : 0.0f; + } + for (int i = 1; i < topk; ++i) { + int id = ids[i]; + float weight = weights[i]; + int j = i - 1; + while (j >= 0 && ids[j] > id) { + ids[j + 1] = ids[j]; + weights[j + 1] = weights[j]; + --j; + } + ids[j + 1] = id; + weights[j + 1] = weight; + } +} + +// One thread per row deliberately retains the host softmax summation order. +__global__ void route_from_logits_kernel(const float *logits, float *weights, + int *ids, int *errors, + unsigned long long *summary, + int layer_idx) { + const int token = blockIdx.x; + float probs[kMaxExperts]; + float maximum = -CUDART_INF_F; + bool valid = true; + for (int e = 0; e < kMaxExperts; ++e) { + float x = logits[token * kMaxExperts + e]; + valid = valid && isfinite(x); + maximum = fmaxf(maximum, x); + } + errors[token] = valid ? 0 : 1; + if (!valid) { + atomicAdd(summary + 2, 1ULL); + for (int r = 0; r < kRoutesPerToken; ++r) { + ids[token * kRoutesPerToken + r] = -1; + weights[token * kRoutesPerToken + r] = 0.0f; + } + return; + } + float sum = 0.0f; + for (int e = 0; e < kMaxExperts; ++e) { + probs[e] = expf(logits[token * kMaxExperts + e] - maximum); + sum = __fadd_rn(sum, probs[e]); + } + for (int e = 0; e < kMaxExperts; ++e) probs[e] = __fdiv_rn(probs[e], sum); + int chosen[kMaxTopk]; + float chosen_weights[kMaxTopk]; + select_and_sort_routes(probs, kMaxExperts, kRoutesPerToken, chosen, chosen_weights); + for (int r = 0; r < kRoutesPerToken; ++r) { + ids[token * kRoutesPerToken + r] = chosen[r]; + weights[token * kRoutesPerToken + r] = chosen_weights[r]; + } + unsigned long long route_hash = 1469598103934665603ULL; + for (int r = 0; r < kRoutesPerToken; ++r) { + route_hash ^= static_cast(chosen[r] + 1 + 67 * r + 4099 * token); + route_hash *= 1099511628211ULL; + } + atomicAdd(summary + 3 + layer_idx, route_hash); +} + +__global__ void route_pointers_kernel( + const int *ids, const __nv_bfloat16 *hidden, const __nv_bfloat16 *zero, + const __nv_bfloat16 *const *w13, const __nv_bfloat16 *const *w2, + __nv_bfloat16 *gate, __nv_bfloat16 *act, __nv_bfloat16 *rows, + const __nv_bfloat16 **a13, const __nv_bfloat16 **x13, __nv_bfloat16 **y13, + const __nv_bfloat16 **a2, const __nv_bfloat16 **x2, __nv_bfloat16 **y2, + unsigned long long *summary, int routes, int first_expert, int hidden_dim, + int intermediate) { + int r = threadIdx.x; + if (r >= routes) return; + int global_expert = ids[r]; + bool local = global_expert >= first_expert && global_expert < first_expert + kLocalExperts; + int expert = local ? global_expert - first_expert : 0; + atomicAdd(summary + 1, 1ULL); + if (local) atomicAdd(summary, 1ULL); + // Dummy reads use a valid weight and zero input; every output has its own slot. + a13[r] = w13[local ? expert : 0]; + x13[r] = local ? hidden + (r / kRoutesPerToken) * hidden_dim : zero; + y13[r] = gate + r * 2 * intermediate; + a2[r] = w2[local ? expert : 0]; + x2[r] = act + r * intermediate; + y2[r] = rows + r * hidden_dim; +} + +__global__ void route_reduce_kernel(const __nv_bfloat16 *rows, const int *ids, + const float *weights, const int *errors, + float *out, int first_expert, int hidden_dim) { + int token = blockIdx.x; + for (int d = threadIdx.x; d < hidden_dim; d += blockDim.x) { + float acc = 0.0f; + for (int r = 0; r < kRoutesPerToken; ++r) { + int slot = token * kRoutesPerToken + r; + int e = ids[slot] - first_expert; + if (e >= 0 && e < kLocalExperts) + acc += __bfloat162float(rows[slot * hidden_dim + d]) * weights[slot]; + } + out[token * hidden_dim + d] = errors[token] ? CUDART_NAN_F : acc; + } +} + __global__ void router_logits_kernel( const __nv_bfloat16 *__restrict__ hidden, const __nv_bfloat16 *__restrict__ gate_weight, @@ -33,11 +160,6 @@ __global__ void router_logits_kernel( logits[token * n_experts + expert] = acc; } -__device__ __forceinline__ bool better_prob_choice(float value, int expert, float best_value, - int best_expert) { - return value > best_value || (value == best_value && expert < best_expert); -} - __global__ void router_softmax_topk_kernel( const __nv_bfloat16 *__restrict__ hidden, const __nv_bfloat16 *__restrict__ gate_weight, @@ -89,36 +211,7 @@ __global__ void router_softmax_topk_kernel( probs[expert] *= inv_denom; } - bool selected[kMaxExperts]; - for (int expert = 0; expert < kMaxExperts; ++expert) selected[expert] = false; - for (int route = 0; route < topk; ++route) { - int best_expert = n_experts; - float best_value = -CUDART_INF_F; - for (int expert = 0; expert < n_experts; ++expert) { - if (selected[expert]) continue; - float value = probs[expert]; - if (better_prob_choice(value, expert, best_value, best_expert)) { - best_value = value; - best_expert = expert; - } - } - selected[best_expert] = true; - selected_idx[route] = best_expert; - selected_weight[route] = best_expert < n_experts ? probs[best_expert] : 0.0f; - } - - for (int i = 1; i < topk; ++i) { - int id = selected_idx[i]; - float weight = selected_weight[i]; - int j = i - 1; - while (j >= 0 && selected_idx[j] > id) { - selected_idx[j + 1] = selected_idx[j]; - selected_weight[j + 1] = selected_weight[j]; - --j; - } - selected_idx[j + 1] = id; - selected_weight[j + 1] = weight; - } + select_and_sort_routes(probs, n_experts, topk, selected_idx, selected_weight); int out_base = token * topk; for (int route = 0; route < topk; ++route) { @@ -154,44 +247,6 @@ __global__ void accumulate_fixed_expert_kernel( } } -CUresult map_cuda_error(cudaError_t err) { - switch (err) { - case cudaSuccess: - return CUDA_SUCCESS; - case cudaErrorInvalidValue: - case cudaErrorInvalidDevicePointer: - return CUDA_ERROR_INVALID_VALUE; - case cudaErrorInvalidDevice: - return CUDA_ERROR_INVALID_DEVICE; - case cudaErrorInvalidResourceHandle: - return CUDA_ERROR_INVALID_HANDLE; - case cudaErrorMemoryAllocation: - return CUDA_ERROR_OUT_OF_MEMORY; - case cudaErrorNotSupported: - return CUDA_ERROR_NOT_SUPPORTED; - case cudaErrorIllegalAddress: - return CUDA_ERROR_ILLEGAL_ADDRESS; - case cudaErrorLaunchOutOfResources: - return CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES; - case cudaErrorLaunchTimeout: - return CUDA_ERROR_LAUNCH_TIMEOUT; - case cudaErrorLaunchFailure: - return CUDA_ERROR_LAUNCH_FAILED; - case cudaErrorAssert: - return CUDA_ERROR_ASSERT; - case cudaErrorIllegalInstruction: - return CUDA_ERROR_ILLEGAL_INSTRUCTION; - case cudaErrorMisalignedAddress: - return CUDA_ERROR_MISALIGNED_ADDRESS; - case cudaErrorInvalidAddressSpace: - return CUDA_ERROR_INVALID_ADDRESS_SPACE; - case cudaErrorInvalidPc: - return CUDA_ERROR_INVALID_PC; - default: - return CUDA_ERROR_UNKNOWN; - } -} - CUresult consume_last_cuda_error() { cudaError_t err = cudaGetLastError(); return map_cuda_error(err); @@ -201,6 +256,71 @@ CUresult consume_last_cuda_error() { extern "C" { +// DSV2-Lite route consumer. Pointer arrays reside on device; each entry is an +// independent N=1 problem, not a multi-row expert GEMM. +int dsv2_lite_pointer_gemm_cuda(const void *const *weights, + const void *const *inputs, void *const *outputs, + int m, int k, int routes, cudaStream_t stream) { + if (!g_cublas_handle) return static_cast(cudaErrorInvalidResourceHandle); + if (!weights || !inputs || !outputs || m <= 0 || k <= 0 || routes <= 0 || routes > kMaxRows * kRoutesPerToken) + return static_cast(cudaErrorInvalidValue); + cublasStatus_t status = cublasSetStream(g_cublas_handle, stream); + if (status != CUBLAS_STATUS_SUCCESS) { + return PEGAINFER_CUBLAS_STATUS_BASE + static_cast(status); + } + const float alpha = 1.0f, beta = 0.0f; + status = cublasGemmBatchedEx( + g_cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, m, 1, k, + &alpha, weights, CUDA_R_16BF, k, inputs, CUDA_R_16BF, k, + &beta, outputs, CUDA_R_16BF, m, routes, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP); + return status == CUBLAS_STATUS_SUCCESS ? static_cast(cudaPeekAtLastError()) + : PEGAINFER_CUBLAS_STATUS_BASE + static_cast(status); +} + +CUresult dsv2_lite_route_logits_cuda(const float *logits, float *weights, + int *ids, int *errors, + unsigned long long *summary, int batch, + int layer_idx, cudaStream_t stream) { + if (!logits || !weights || !ids || !errors || !summary || batch < 1 || batch > kMaxRows || + layer_idx < 0) + return CUDA_ERROR_INVALID_VALUE; + route_from_logits_kernel<<>>( + logits, weights, ids, errors, summary, layer_idx); + return map_cuda_error(cudaPeekAtLastError()); +} + +CUresult dsv2_lite_route_pointers_cuda( + const int *ids, const __nv_bfloat16 *hidden, const __nv_bfloat16 *zero, + const __nv_bfloat16 *const *w13, const __nv_bfloat16 *const *w2, + __nv_bfloat16 *gate, __nv_bfloat16 *act, __nv_bfloat16 *rows, + const __nv_bfloat16 **a13, const __nv_bfloat16 **x13, __nv_bfloat16 **y13, + const __nv_bfloat16 **a2, const __nv_bfloat16 **x2, __nv_bfloat16 **y2, + unsigned long long *summary, int batch, int first_expert, int hidden_dim, + int intermediate, cudaStream_t stream) { + if (!ids || !hidden || !zero || !w13 || !w2 || !gate || !act || !rows || + !a13 || !x13 || !y13 || !a2 || !x2 || !y2 || !summary || + batch < 1 || batch > kMaxRows || + (first_expert != 0 && first_expert != kLocalExperts) || hidden_dim <= 0 || intermediate <= 0) + return CUDA_ERROR_INVALID_VALUE; + route_pointers_kernel<<<1, kMaxRows * kRoutesPerToken, 0, stream>>>( + ids, hidden, zero, w13, w2, gate, act, rows, a13, x13, y13, a2, x2, y2, + summary, batch * kRoutesPerToken, first_expert, hidden_dim, intermediate); + return map_cuda_error(cudaPeekAtLastError()); +} + +CUresult dsv2_lite_route_reduce_cuda(const __nv_bfloat16 *rows, const int *ids, + const float *weights, const int *errors, + float *out, int batch, int first_expert, + int hidden_dim, cudaStream_t stream) { + if (!rows || !ids || !weights || !errors || !out || batch < 1 || batch > kMaxRows || + (first_expert != 0 && first_expert != kLocalExperts) || hidden_dim <= 0) + return CUDA_ERROR_INVALID_VALUE; + route_reduce_kernel<<>>( + rows, ids, weights, errors, out, first_expert, hidden_dim); + return map_cuda_error(cudaPeekAtLastError()); +} + CUresult dsv2_lite_router_logits_cuda( const __nv_bfloat16 *hidden, const __nv_bfloat16 *gate_weight, diff --git a/pegainfer-kernels/csrc/gemma4/marlin_nvfp4_prepare.cu b/pegainfer-kernels/csrc/gemma4/marlin_nvfp4_prepare.cu index 9a44a11d7..0acc04ed1 100644 --- a/pegainfer-kernels/csrc/gemma4/marlin_nvfp4_prepare.cu +++ b/pegainfer-kernels/csrc/gemma4/marlin_nvfp4_prepare.cu @@ -15,6 +15,8 @@ #include #include +#include "../shared/ffi_guard.cuh" + namespace { constexpr int kGroup = 16; @@ -95,7 +97,7 @@ CUresult gemma4_marlin_nvfp4_prepare_scales_cuda(const unsigned char *checkpoint } nvfp4_prepare_scales_kernel<<<(int)grid, block, 0, stream>>>( checkpoint, prepared, out_dim, scale_k, rescale, total); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } } // extern "C" diff --git a/pegainfer-kernels/csrc/gemma4/moe.cu b/pegainfer-kernels/csrc/gemma4/moe.cu index a74299750..703bdc6df 100644 --- a/pegainfer-kernels/csrc/gemma4/moe.cu +++ b/pegainfer-kernels/csrc/gemma4/moe.cu @@ -6,6 +6,8 @@ #include #include +#include "../shared/ffi_guard.cuh" + namespace { // The strided lane scan and shuffle tree preserve lower-expert tie breaking. @@ -145,7 +147,7 @@ CUresult gemma4_moe_router_topk_cuda(const __nv_bfloat16 *logits, } gemma4_moe_router_topk_kernel<<>>( logits, per_expert_scale, experts, top_k, index_out, weight_out); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } // The routed GEMM leaves one row per (token, pick); this folds the picks back @@ -165,7 +167,7 @@ CUresult gemma4_moe_sum_topk_cuda(const __nv_bfloat16 *routed, int rows, } gemma4_moe_sum_topk_kernel<<<(int)grid, block, 0, stream>>>( routed, top_k, hidden, total, out); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } } // extern "C" diff --git a/pegainfer-kernels/csrc/shared/paged_attention_hd512.cu b/pegainfer-kernels/csrc/gemma4/paged_attention_hd512.cu similarity index 55% rename from pegainfer-kernels/csrc/shared/paged_attention_hd512.cu rename to pegainfer-kernels/csrc/gemma4/paged_attention_hd512.cu index 77ae98c92..43079f50d 100644 --- a/pegainfer-kernels/csrc/shared/paged_attention_hd512.cu +++ b/pegainfer-kernels/csrc/gemma4/paged_attention_hd512.cu @@ -2,26 +2,32 @@ // attention layers). Separate TU so the long FlashInfer template builds // compile in parallel with paged_attention.cu. -#include "paged_launch.cuh" +#include "../shared/paged_launch.cuh" extern "C" { -int single_prefill_cuda_hd512( - void* q, - void* output, - void* k_cache, - void* v_cache, - int32_t num_qo_heads, - int32_t num_kv_heads, - int32_t seq_len, // number of Q tokens (qo_len) - int32_t kv_len, // total KV length (start_pos + seq_len) - int32_t max_seq_len, // allocated cache rows (for HND stride) - float sm_scale, - void* stream) +// Split-KV decode — the global family's decode read. The non-partitioned +// grid is (pseudo-requests, kv heads) CTAs and starves the device; chunking +// the KV brings the grid to occupancy. +int paged_attention_decode_split_kv_cuda_hd512( + void* q, void* output, void* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, + int32_t* request_indices, int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, + int32_t* o_indptr, uint8_t* block_valid_mask, + void* tmp_v, float* tmp_s, + int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, + int32_t page_size, int32_t batch_size, int32_t padded_batch_size, + int64_t stride_page, float sm_scale, void* stream) { - return single_prefill_launch( - q, output, k_cache, v_cache, num_qo_heads, num_kv_heads, - /*head_dim=*/512, seq_len, kv_len, max_seq_len, sm_scale, stream); + return decode_split_kv_launch( + q, output, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indptr, last_page_len_d, + request_indices, kv_tile_indices, kv_chunk_size_ptr, + o_indptr, block_valid_mask, tmp_v, tmp_s, + num_qo_heads, num_kv_heads, head_dim, page_size, + batch_size, padded_batch_size, stride_page, sm_scale, + /*window_left=*/-1, stream); } // Takes no cta_tile_q override parameter because hd512 only ever selects 16|32 diff --git a/pegainfer-kernels/csrc/gemma4/paged_attention_window_hd256.cu b/pegainfer-kernels/csrc/gemma4/paged_attention_window_hd256.cu new file mode 100644 index 000000000..3f9805686 --- /dev/null +++ b/pegainfer-kernels/csrc/gemma4/paged_attention_window_hd256.cu @@ -0,0 +1,29 @@ +// FlashInfer paged prefill at HEAD_DIM=256 with the sliding-window mask +// (Gemma 4 local attention layers), over a bf16 pool. The e4m3 pool has its +// own entry in local_kv_fp8.cu. + +#include "../shared/paged_launch.cuh" + +extern "C" { + +int batch_prefill_paged_window_cuda_hd256( + void* q, void* output, void* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, + int32_t* q_indptr, int32_t* request_indices, int32_t* qo_tile_indices, + int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, uint32_t* total_num_rows, + int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, + int32_t page_size, int32_t seq_len, int32_t batch_size, + int32_t padded_batch_size, int64_t stride_page, float sm_scale, + int32_t window_left, void* stream) +{ + return prefill_paged_launch( + q, output, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indptr, last_page_len_d, q_indptr, + request_indices, qo_tile_indices, kv_tile_indices, kv_chunk_size_ptr, + total_num_rows, num_qo_heads, num_kv_heads, head_dim, page_size, + seq_len, batch_size, padded_batch_size, stride_page, sm_scale, + /*cta_tile_q_override=*/0, window_left, stream); +} + +} // extern "C" diff --git a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu b/pegainfer-kernels/csrc/gemma4/prefill_attention_hd256_plain.cu similarity index 78% rename from pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu rename to pegainfer-kernels/csrc/gemma4/prefill_attention_hd256_plain.cu index ce9c516e9..515f2df2a 100644 --- a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu +++ b/pegainfer-kernels/csrc/gemma4/prefill_attention_hd256_plain.cu @@ -3,9 +3,6 @@ // computes the 1+w offset norm, assumes a gated Q layout twice as wide, and // ships no v_norm. // -// Two entry points share the math: the contiguous oracle form and the -// paged serving form. -// // rotary_dim is a runtime argument, checked at the launcher for positive, // even and <= HD256. Gemma 4 local layers rotate the full head (256); at // that width the pass-through tail is empty. Evenness is load-bearing: with @@ -15,11 +12,11 @@ // Positions derive from host-known start_pos, so the launcher rejects any // out-of-range window before launch; the device trap before the cos read is // a second layer, not the contract. Page ids are device data, so for them -// the paged form's trap is the only check. +// the kernel's trap is the only check. #include "common.cuh" -#include "ffi_guard.cuh" -#include "qk_prep.cuh" +#include "../shared/ffi_guard.cuh" +#include "../shared/qk_prep.cuh" #include #define HD256_PLAIN 256 @@ -38,100 +35,7 @@ __device__ __forceinline__ __nv_bfloat16 rms_norm_elem_hd256_plain( return __float2bfloat16(__bfloat162float(x) * rms_inv * w); } -__global__ void qk_norm_rope_prefill_hd256_plain_kernel( - const __nv_bfloat16* __restrict__ q_batch, // [q_dim, seq_len] - const __nv_bfloat16* __restrict__ k_batch, // [kv_dim, seq_len] - const __nv_bfloat16* __restrict__ q_norm_weight, // [HD256_PLAIN] - const __nv_bfloat16* __restrict__ k_norm_weight, // [HD256_PLAIN] - const __nv_bfloat16* __restrict__ cos_cache, // [max_seq * rotary_dim] - const __nv_bfloat16* __restrict__ sin_cache, - __nv_bfloat16* __restrict__ q_batch_out, // [q_dim, seq_len] - __nv_bfloat16* __restrict__ k_batch_out, // [kv_dim, seq_len] - int num_q_heads, - int num_kv_heads, - int start_pos, // host base position - int cos_max_pos, // rows in cos/sin tables - int rotary_dim, - float rms_eps -) { - // seq_len is mapped onto grid.x (limit ~2^31) and the head index onto - // grid.y so prompts longer than the 65535 grid.y limit still launch. - int token = blockIdx.x; - int head_global = blockIdx.y; - int d = threadIdx.x; - - bool is_q = head_global < num_q_heads; - int head_local = is_q ? head_global : (head_global - num_q_heads); - int q_dim = num_q_heads * HD256_PLAIN; - int kv_dim = num_kv_heads * HD256_PLAIN; - - int src_offset = is_q - ? token * q_dim + head_local * HD256_PLAIN + d - : token * kv_dim + head_local * HD256_PLAIN + d; - __nv_bfloat16 x = is_q ? q_batch[src_offset] : k_batch[src_offset]; - const __nv_bfloat16* norm_w = is_q ? q_norm_weight : k_norm_weight; - - float sq = __bfloat162float(x); - sq *= sq; - float sq_sum = warp_reduce_sum(sq); - - int warp_id = d / WARP_SIZE; - int lane_id = d % WARP_SIZE; - __shared__ float warp_sums[NUM_WARPS_HD256_PLAIN]; - __shared__ float inv_rms; - __shared__ __nv_bfloat16 smem[HD256_PLAIN]; - - if (lane_id == 0) warp_sums[warp_id] = sq_sum; - __syncthreads(); - - if (d == 0) { - float total = 0.0f; - for (int i = 0; i < NUM_WARPS_HD256_PLAIN; i++) total += warp_sums[i]; - inv_rms = 1.0f / sqrtf(total / HD256_PLAIN + rms_eps); - } - __syncthreads(); - - smem[d] = rms_norm_elem_hd256_plain(x, inv_rms, norm_w[d]); - __syncthreads(); - - int pos = start_pos + token; - // Reject before reading the cos/sin tables. - if (pos < 0 || pos >= cos_max_pos) __trap(); - int half_rotary = rotary_dim / 2; - - if (d < half_rotary) { - __nv_bfloat16 lo = smem[d]; - __nv_bfloat16 hi = smem[d + half_rotary]; - apply_rope_pair( - lo, - hi, - cos_cache[pos * rotary_dim + d], - sin_cache[pos * rotary_dim + d] - ); - - if (is_q) { - int dst = token * q_dim + head_local * HD256_PLAIN; - q_batch_out[dst + d] = lo; - q_batch_out[dst + d + half_rotary] = hi; - } else { - int dst = token * kv_dim + head_local * HD256_PLAIN; - k_batch_out[dst + d] = lo; - k_batch_out[dst + d + half_rotary] = hi; - } - } - - if (d >= rotary_dim) { - if (is_q) { - int dst = token * q_dim + head_local * HD256_PLAIN; - q_batch_out[dst + d] = smem[d]; - } else { - int dst = token * kv_dim + head_local * HD256_PLAIN; - k_batch_out[dst + d] = smem[d]; - } - } -} - -// Paged serving prep. grid.y carries three bands: [0, num_q_heads) Q, +// Paged prep. grid.y carries three bands: [0, num_q_heads) Q, // then num_kv_heads K, then num_kv_heads V. Q and K are plain-w normed // and rotated; V is weightless-normed over its own head vector (v_proj // output — a separate reduction, unlike the hd512 K=V fork) and never @@ -320,80 +224,6 @@ __global__ void qkv_norm_rope_paged_prefill_hd256_plain_kernel( } } -extern "C" { - -int qk_norm_rope_prefill_hd256_plain_cuda( - const __nv_bfloat16* q_batch, - const __nv_bfloat16* k_batch, - const __nv_bfloat16* q_norm_weight, - const __nv_bfloat16* k_norm_weight, - const __nv_bfloat16* cos_cache, - const __nv_bfloat16* sin_cache, - __nv_bfloat16* q_batch_out, - __nv_bfloat16* k_batch_out, - int num_q_heads, - int num_kv_heads, - int seq_len, - int start_pos, - int cos_max_pos, - int rotary_dim, - float rms_eps, - cudaStream_t stream -) { - PEGAINFER_FFI_GUARD_BEGIN - if (rotary_dim <= 0 || (rotary_dim & 1) != 0 || rotary_dim > HD256_PLAIN) { - pegainfer_ffi_set_last_error( - "qk_norm_rope_prefill_hd256_plain_cuda: rotary_dim must be " - "positive, even and <= 256"); - return -1; - } - if (q_batch == nullptr || k_batch == nullptr || q_norm_weight == nullptr || - k_norm_weight == nullptr || cos_cache == nullptr || sin_cache == nullptr || - q_batch_out == nullptr || k_batch_out == nullptr) { - pegainfer_ffi_set_last_error( - "qk_norm_rope_prefill_hd256_plain_cuda: null pointer argument"); - return -1; - } - if (num_q_heads <= 0 || num_kv_heads <= 0 || seq_len <= 0) { - pegainfer_ffi_set_last_error( - "qk_norm_rope_prefill_hd256_plain_cuda: num_q_heads, num_kv_heads " - "and seq_len must be positive"); - return -1; - } - if (start_pos < 0 || start_pos + seq_len > cos_max_pos) { - pegainfer_ffi_set_last_error( - "qk_norm_rope_prefill_hd256_plain_cuda: start_pos + seq_len must " - "be <= cos_max_pos"); - return -1; - } - dim3 prep_grid(seq_len, num_q_heads + num_kv_heads); - qk_norm_rope_prefill_hd256_plain_kernel<<>>( - q_batch, - k_batch, - q_norm_weight, - k_norm_weight, - cos_cache, - sin_cache, - q_batch_out, - k_batch_out, - num_q_heads, - num_kv_heads, - start_pos, - cos_max_pos, - rotary_dim, - rms_eps - ); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - pegainfer_ffi_set_last_error(cudaGetErrorString(err)); - return -1; - } - return 0; - PEGAINFER_FFI_GUARD_END(-1) -} - -} // extern "C" - template static int qkv_prep_paged_prefill_launch( const __nv_bfloat16* q_batch, diff --git a/pegainfer-kernels/csrc/shared/prefill_attention_hd512.cu b/pegainfer-kernels/csrc/gemma4/prefill_attention_hd512.cu similarity index 76% rename from pegainfer-kernels/csrc/shared/prefill_attention_hd512.cu rename to pegainfer-kernels/csrc/gemma4/prefill_attention_hd512.cu index 647e71956..02f2e452c 100644 --- a/pegainfer-kernels/csrc/shared/prefill_attention_hd512.cu +++ b/pegainfer-kernels/csrc/gemma4/prefill_attention_hd512.cu @@ -20,8 +20,8 @@ // host would require a D2H synchronization. #include "common.cuh" -#include "ffi_guard.cuh" -#include "qk_prep.cuh" +#include "../shared/ffi_guard.cuh" +#include "../shared/qk_prep.cuh" #define HD512 512 #define HALF_HD512 (HD512 / 2) @@ -225,100 +225,6 @@ __global__ void qk_norm_partial_rope_paged_prefill_hd512_kernel( } } -// Batched decode prep: Q is written to contiguous q_batch_out; K is updated -// in place. Paged scatter is caller-owned. -__global__ void qk_norm_partial_rope_batched_decode_hd512_kernel( - const __nv_bfloat16* __restrict__ q_batch, // [q_dim, batch] - __nv_bfloat16* __restrict__ k_batch, // [kv_dim, batch] in-place - const __nv_bfloat16* __restrict__ q_norm_weight, // [HD512] - const __nv_bfloat16* __restrict__ k_norm_weight, // [HD512] - const __nv_bfloat16* __restrict__ cos_cache, // [max_seq * rotary_dim] - const __nv_bfloat16* __restrict__ sin_cache, - const int* __restrict__ positions, // [batch] - int cos_max_pos, // rows in cos/sin tables - __nv_bfloat16* __restrict__ q_batch_out, // [q_dim, batch] - int num_q_heads, - int num_kv_heads, - int batch_size, - int rotary_dim, - float rms_eps -) { - int head_global = blockIdx.x; - int token = blockIdx.y; - int d = threadIdx.x; - - bool is_q = head_global < num_q_heads; - int head_local = is_q ? head_global : (head_global - num_q_heads); - int q_dim = num_q_heads * HD512; - int kv_dim = num_kv_heads * HD512; - - int src_offset = is_q - ? token * q_dim + head_local * HD512 + d - : token * kv_dim + head_local * HD512 + d; - - __nv_bfloat16 x = is_q ? q_batch[src_offset] : k_batch[src_offset]; - const __nv_bfloat16* norm_w = is_q ? q_norm_weight : k_norm_weight; - - float sq = __bfloat162float(x); - sq *= sq; - float sq_sum = warp_reduce_sum(sq); - - int warp_id = d / WARP_SIZE; - int lane_id = d % WARP_SIZE; - __shared__ float warp_sums[NUM_WARPS_HD512]; - __shared__ float inv_rms; - __shared__ __nv_bfloat16 smem[HD512]; - - if (lane_id == 0) warp_sums[warp_id] = sq_sum; - __syncthreads(); - - if (d == 0) { - float total = 0.0f; - for (int i = 0; i < NUM_WARPS_HD512; i++) total += warp_sums[i]; - inv_rms = 1.0f / sqrtf(total / HD512 + rms_eps); - } - __syncthreads(); - - smem[d] = rms_norm_elem_hd512(x, inv_rms, norm_w[d]); - __syncthreads(); - - int pos = positions[token]; - // Reject before reading the cos/sin tables. - if (pos < 0 || pos >= cos_max_pos) __trap(); - int half_rotary = rotary_dim / 2; - - if (d < half_rotary) { - __nv_bfloat16 lo = smem[d]; - __nv_bfloat16 hi = smem[d + half_rotary]; - apply_rope_pair( - lo, - hi, - cos_cache[pos * rotary_dim + d], - sin_cache[pos * rotary_dim + d] - ); - - if (is_q) { - int dst = token * q_dim + head_local * HD512; - q_batch_out[dst + d] = lo; - q_batch_out[dst + d + half_rotary] = hi; - } else { - int dst = token * kv_dim + head_local * HD512; - k_batch[dst + d] = lo; - k_batch[dst + d + half_rotary] = hi; - } - } - - if (d >= rotary_dim) { - if (is_q) { - int dst = token * q_dim + head_local * HD512; - q_batch_out[dst + d] = smem[d]; - } else { - int dst = token * kv_dim + head_local * HD512; - k_batch[dst + d] = smem[d]; - } - } -} - extern "C" { int qk_norm_partial_rope_paged_prefill_hd512_cuda( @@ -432,74 +338,6 @@ int qk_norm_partial_rope_paged_prefill_hd512_cuda( PEGAINFER_FFI_GUARD_END(-1) } -int qk_norm_partial_rope_batched_decode_hd512_cuda( - const __nv_bfloat16* q_batch, - __nv_bfloat16* k_batch, - const __nv_bfloat16* q_norm_weight, - const __nv_bfloat16* k_norm_weight, - const __nv_bfloat16* cos_cache, - const __nv_bfloat16* sin_cache, - const int* positions, - int cos_max_pos, - __nv_bfloat16* q_batch_out, - int num_q_heads, - int num_kv_heads, - int batch_size, - int rotary_dim, - float rms_eps, - cudaStream_t stream -) { - PEGAINFER_FFI_GUARD_BEGIN - if (rotary_dim <= 0 || (rotary_dim & 1) != 0 || rotary_dim > HD512) { - pegainfer_ffi_set_last_error( - "qk_norm_partial_rope_batched_decode_hd512_cuda: rotary_dim must be " - "positive, even and <= 512"); - return -1; - } - if (q_batch == nullptr || k_batch == nullptr || q_norm_weight == nullptr || - k_norm_weight == nullptr || cos_cache == nullptr || sin_cache == nullptr || - positions == nullptr || q_batch_out == nullptr) { - pegainfer_ffi_set_last_error( - "qk_norm_partial_rope_batched_decode_hd512_cuda: null pointer argument"); - return -1; - } - if (num_q_heads <= 0 || num_kv_heads <= 0 || batch_size <= 0) { - pegainfer_ffi_set_last_error( - "qk_norm_partial_rope_batched_decode_hd512_cuda: num_q_heads, " - "num_kv_heads and batch_size must be positive"); - return -1; - } - if (cos_max_pos <= 0) { - pegainfer_ffi_set_last_error( - "qk_norm_partial_rope_batched_decode_hd512_cuda: cos_max_pos must be positive"); - return -1; - } - dim3 grid(num_q_heads + num_kv_heads, batch_size); - qk_norm_partial_rope_batched_decode_hd512_kernel<<>>( - q_batch, - k_batch, - q_norm_weight, - k_norm_weight, - cos_cache, - sin_cache, - positions, - cos_max_pos, - q_batch_out, - num_q_heads, - num_kv_heads, - batch_size, - rotary_dim, - rms_eps - ); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - pegainfer_ffi_set_last_error(cudaGetErrorString(err)); - return -1; - } - return 0; - PEGAINFER_FFI_GUARD_END(-1) -} - int qk_norm_partial_rope_paged_decode_hd512_cuda( const __nv_bfloat16* q_batch, const __nv_bfloat16* k_batch, diff --git a/pegainfer-kernels/csrc/gemma4/w4a16_layout.cu b/pegainfer-kernels/csrc/gemma4/w4a16_layout.cu index c1d2e21a3..8fe1c39a8 100644 --- a/pegainfer-kernels/csrc/gemma4/w4a16_layout.cu +++ b/pegainfer-kernels/csrc/gemma4/w4a16_layout.cu @@ -20,6 +20,8 @@ #include #include +#include "../shared/ffi_guard.cuh" + namespace pegainfer_gemma4_w4a16 { constexpr int kGroup = 32; @@ -120,7 +122,7 @@ CUresult gemma4_w4a16_pack_cuda(const uint32_t* packed, const uint16_t* scales, return CUDA_ERROR_INVALID_VALUE; } pack_kernel<<<1024, 256, 0, stream>>>(packed, scales, wq, sq, n, k, split); - return cudaGetLastError() == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(cudaGetLastError()); } CUresult gemma4_w4a16_dequant_cuda(const uint32_t* wq, const uint32_t* sq, __nv_bfloat16* out, int n, int k, @@ -129,7 +131,7 @@ CUresult gemma4_w4a16_dequant_cuda(const uint32_t* wq, const uint32_t* sq, __nv_ return CUDA_ERROR_INVALID_VALUE; } dequant_kernel<<<1024, 256, 0, stream>>>(wq, sq, out, n, k, split); - return cudaGetLastError() == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(cudaGetLastError()); } } // extern "C" diff --git a/pegainfer-kernels/csrc/glm52/glm52_deepgemm_grouped_sm100.cu b/pegainfer-kernels/csrc/glm52/glm52_deepgemm_grouped_sm100.cu index ee70d6697..cba675b65 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_deepgemm_grouped_sm100.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_deepgemm_grouped_sm100.cu @@ -35,6 +35,7 @@ // baked into the kernel template, so each width gets its own instantiation. #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -128,16 +129,6 @@ __global__ void deepgemm_sm100_masked_out_to_aligned_kernel( } } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } } // namespace diff --git a/pegainfer-kernels/csrc/glm52/glm52_deepgemm_mqa.cu b/pegainfer-kernels/csrc/glm52/glm52_deepgemm_mqa.cu index a5a023e4c..7eeadbeba 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_deepgemm_mqa.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_deepgemm_mqa.cu @@ -16,6 +16,7 @@ // defined via build.rs. #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -100,7 +101,7 @@ CUresult launch_aot(const void* func, dim3 grid_dim, dim3 block_dim, int smem_si if (attr_err != cudaSuccess) { fprintf(stderr, "glm52_deepgemm_mqa: cudaFuncSetAttribute failed: %s\n", cudaGetErrorString(attr_err)); - return CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(attr_err); } } @@ -119,7 +120,7 @@ CUresult launch_aot(const void* func, dim3 grid_dim, dim3 block_dim, int smem_si const cudaError_t err = cudaLaunchKernelExC(&config, func, args); if (err != cudaSuccess) { fprintf(stderr, "glm52_deepgemm_mqa: launch failed: %s\n", cudaGetErrorString(err)); - return CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } return CUDA_SUCCESS; } @@ -136,7 +137,7 @@ CUresult launch_aot_plain(const void* func, dim3 grid_dim, dim3 block_dim, int s if (attr_err != cudaSuccess) { fprintf(stderr, "glm52_deepgemm_mqa: cudaFuncSetAttribute failed: %s\n", cudaGetErrorString(attr_err)); - return CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(attr_err); } } cudaLaunchConfig_t config = {}; @@ -147,7 +148,7 @@ CUresult launch_aot_plain(const void* func, dim3 grid_dim, dim3 block_dim, int s const cudaError_t err = cudaLaunchKernelExC(&config, func, args); if (err != cudaSuccess) { fprintf(stderr, "glm52_deepgemm_mqa: launch failed: %s\n", cudaGetErrorString(err)); - return CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } return CUDA_SUCCESS; } diff --git a/pegainfer-kernels/csrc/glm52/glm52_flashinfer_sparse.cu b/pegainfer-kernels/csrc/glm52/glm52_flashinfer_sparse.cu index cc6344f0d..fa81e056c 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_flashinfer_sparse.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_flashinfer_sparse.cu @@ -23,13 +23,7 @@ constexpr int kHeadDimV = 512; constexpr int kPageSize = 64; constexpr size_t kCounterBytes = 8192ull * 256 * sizeof(uint32_t); -CUresult consume_last_cuda_error() { - cudaError_t err = cudaGetLastError(); - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue) return CUDA_ERROR_INVALID_VALUE; - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - return CUDA_ERROR_LAUNCH_FAILED; -} +CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } } // namespace diff --git a/pegainfer-kernels/csrc/glm52/glm52_flashmla_sparse.cu b/pegainfer-kernels/csrc/glm52/glm52_flashmla_sparse.cu index 620f342aa..94467cda3 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_flashmla_sparse.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_flashmla_sparse.cu @@ -6,6 +6,7 @@ #include #include +#include "../shared/ffi_guard.cuh" #include "params.h" #include "sm100/decode/head64/kernel.cuh" #include "sm100/prefill/sparse/fwd/head64/phase1.cuh" @@ -29,16 +30,6 @@ constexpr int kBlockSizeTopk = 64; constexpr int kFixedOverheadBlocks = 5; constexpr float kLog2E = 1.4426950408889634f; -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } template diff --git a/pegainfer-kernels/csrc/glm52/glm52_fp8_gemm.cu b/pegainfer-kernels/csrc/glm52/glm52_fp8_gemm.cu index f0d8cb1ba..2e7b39703 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_fp8_gemm.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_fp8_gemm.cu @@ -10,16 +10,6 @@ namespace { -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - // Grouped-GEMM workspace split. FlashInfer's grouped entry takes two buffers: // an int_buffer it carves with AlignedAllocator into eleven per-group argument // arrays (problem_sizes 12 B, five pointer arrays 8 B, three cute packed diff --git a/pegainfer-kernels/csrc/glm52/glm52_indexer.cu b/pegainfer-kernels/csrc/glm52/glm52_indexer.cu index 008648fcc..2518df1db 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_indexer.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_indexer.cu @@ -9,6 +9,7 @@ // int32 index-remap: page = block_table[t, off//bs]; slot = page*bs + off%bs. #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include "glm52_min_gemv.cuh" #include @@ -161,16 +162,6 @@ __global__ void indexer_weights_fold_kernel( } } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } bool valid_cache_layout(int head_dim, int quant_block_size, diff --git a/pegainfer-kernels/csrc/glm52/glm52_indexer_k_gather.cu b/pegainfer-kernels/csrc/glm52/glm52_indexer_k_gather.cu index 6f2801dbf..c7d747c51 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_indexer_k_gather.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_indexer_k_gather.cu @@ -16,6 +16,7 @@ // table, mirroring the -1-padding and topk_lens semantics of the paged twin. #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include @@ -133,16 +134,6 @@ __global__ void topk_to_slots_lut_kernel( } } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } } // namespace diff --git a/pegainfer-kernels/csrc/glm52/glm52_indexer_rope.cu b/pegainfer-kernels/csrc/glm52/glm52_indexer_rope.cu index 441eeb5db..ded84e19f 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_indexer_rope.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_indexer_rope.cu @@ -18,6 +18,7 @@ // left unchanged. #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -102,11 +103,7 @@ CUresult glm52_indexer_rope_cuda(__nv_bfloat16* q, // [tokens, n_heads, hea } const dim3 grid(n_heads, tokens, 1); glm52_indexer_rope_kernel<<>>(q, k, n_heads, cos, sin); - cudaError_t err = cudaGetLastError(); - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue) return CUDA_ERROR_INVALID_VALUE; - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - return CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(cudaGetLastError()); } } // extern "C" diff --git a/pegainfer-kernels/csrc/glm52/glm52_mla_assembly.cu b/pegainfer-kernels/csrc/glm52/glm52_mla_assembly.cu index bd3550d1b..a6264644b 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_mla_assembly.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_mla_assembly.cu @@ -17,6 +17,7 @@ // device function; validated bit-for-bit against the HF oracle (q_rot/k_rot). #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -256,13 +257,7 @@ __global__ void glm52_mla_ckv_split_kernel( } } -CUresult consume_last_cuda_error() { - cudaError_t err = cudaGetLastError(); - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue) return CUDA_ERROR_INVALID_VALUE; - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - return CUDA_ERROR_LAUNCH_FAILED; -} +CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } } // namespace diff --git a/pegainfer-kernels/csrc/glm52/glm52_moe_gemv.cu b/pegainfer-kernels/csrc/glm52/glm52_moe_gemv.cu index d6ecf5e5e..80a604af8 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_moe_gemv.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_moe_gemv.cu @@ -16,6 +16,7 @@ // fills the single grid column with 1 row/warp. #include "../common.cuh" // warp_reduce_sum +#include "../shared/ffi_guard.cuh" #include #include @@ -40,14 +41,6 @@ constexpr int kRowsPlainShortKHopper = 4; // k <= 2048, measured on SM90 constexpr int kRowsPlainBlackwell = 1; // all K, measured on SM103 constexpr int kRowsPlainLongK = 1; // k > 2048 on every architecture -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) - return CUDA_ERROR_INVALID_VALUE; - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } // The activation is read with 16-byte float4/uint4 loads straight from global (no shared diff --git a/pegainfer-kernels/csrc/glm52/glm52_moe_quant.cu b/pegainfer-kernels/csrc/glm52/glm52_moe_quant.cu index 1b8091dc1..db243a9bd 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_moe_quant.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_moe_quant.cu @@ -1,4 +1,5 @@ #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -180,16 +181,6 @@ __global__ void silu_and_mul_per_token_group_quant_bf16_k128_masked_kernel( } } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - // f32 mn-major scales [groups, scale_cols, cap] → packed UE8M0 i32 // [groups, scale_cols/4, cap], LSB-first exponent bytes // (u32 = b0 | b1<<8 | b2<<16 | b3<<24 — smxx_layout.cuh's diff --git a/pegainfer-kernels/csrc/glm52/glm52_prefill_moe_route.cu b/pegainfer-kernels/csrc/glm52/glm52_prefill_moe_route.cu index e115300d8..f9a71c39d 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_prefill_moe_route.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_prefill_moe_route.cu @@ -34,16 +34,6 @@ constexpr int kScanThreads = 256; constexpr int kMaxBlocks = 65535; constexpr int kVecWidth = 8; // bf16 per float4 -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } int grid_blocks(long long work) { diff --git a/pegainfer-kernels/csrc/glm52/glm52_prefill_unpack.cu b/pegainfer-kernels/csrc/glm52/glm52_prefill_unpack.cu index ee3f4d4f8..8876b7c22 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_prefill_unpack.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_prefill_unpack.cu @@ -57,6 +57,6 @@ extern "C" CUresult glm52_prefill_unpack_pages_cuda( reinterpret_cast(stream)>>>( packed, block_ids, blocks, packed_bytes, packed_block_stride, max_slots, unpacked); - return static_cast(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) } diff --git a/pegainfer-kernels/csrc/glm52/glm52_router.cu b/pegainfer-kernels/csrc/glm52/glm52_router.cu index 02a1cd204..f62b30827 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_router.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_router.cu @@ -1,4 +1,5 @@ #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include "glm52_min_gemv.cuh" #include @@ -78,16 +79,6 @@ __global__ void router_scores_topk_normalize_kernel( } } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { cudaError_t err = cudaGetLastError(); return map_cuda_error(err); diff --git a/pegainfer-kernels/csrc/glm52/glm52_topk.cu b/pegainfer-kernels/csrc/glm52/glm52_topk.cu index c573022d8..4e1f46264 100644 --- a/pegainfer-kernels/csrc/glm52/glm52_topk.cu +++ b/pegainfer-kernels/csrc/glm52/glm52_topk.cu @@ -46,7 +46,7 @@ extern "C" int glm52_flashinfer_topk_2048_cuda( if (err != cudaSuccess) { fprintf(stderr, "glm52_flashinfer_topk_2048_cuda: FilteredTopK failed: %s\n", cudaGetErrorString(err)); - return static_cast(CUDA_ERROR_LAUNCH_FAILED); + return static_cast(map_cuda_error(err)); } err = flashinfer::sampling::LaunchFinalizeTopKIndices< /*SORT_LOCAL_INDICES=*/true, flashinfer::sampling::FilteredTopKMode::Plain, float, int>( @@ -57,7 +57,7 @@ extern "C" int glm52_flashinfer_topk_2048_cuda( fprintf(stderr, "glm52_flashinfer_topk_2048_cuda: LaunchFinalizeTopKIndices failed: %s\n", cudaGetErrorString(err)); - return static_cast(CUDA_ERROR_LAUNCH_FAILED); + return static_cast(map_cuda_error(err)); } return static_cast(CUDA_SUCCESS); PEGAINFER_FFI_GUARD_END(-1) diff --git a/pegainfer-kernels/csrc/k3/k3_deepgemm_fp8_fp4_grouped_sm100.cu b/pegainfer-kernels/csrc/k3/k3_deepgemm_fp8_fp4_grouped_sm100.cu index 2d6098737..a2a6c2914 100644 --- a/pegainfer-kernels/csrc/k3/k3_deepgemm_fp8_fp4_grouped_sm100.cu +++ b/pegainfer-kernels/csrc/k3/k3_deepgemm_fp8_fp4_grouped_sm100.cu @@ -101,16 +101,6 @@ constexpr int kFp4SfGroupK = 32; constexpr int kFp4SfPerWord = 4; constexpr int kFp4SfWordK = kFp4SfGroupK * kFp4SfPerWord; // 128 -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } // Checkpoint MXFP4 weight scales -> the runtime SFB tensor. diff --git a/pegainfer-kernels/csrc/k3/k3_flash_kda.cu b/pegainfer-kernels/csrc/k3/k3_flash_kda.cu index c2ef2156e..d613e8f32 100644 --- a/pegainfer-kernels/csrc/k3/k3_flash_kda.cu +++ b/pegainfer-kernels/csrc/k3/k3_flash_kda.cu @@ -28,6 +28,8 @@ #include #include +#include "../shared/ffi_guard.cuh" + #ifdef K3_FLASH_KDA_SM90A #include "fwd_launch.cu" #include "k3_flash_kda_md.cuh" @@ -35,16 +37,6 @@ // Not in an anonymous namespace: nvcc's device-stub generation trips over an // anonymous-namespace __global__ when an included header (cute) opens its own. -static CUresult k3_flash_kda_map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - static __global__ void k3_flash_kda_beta_transpose_kernel( const __nv_bfloat16* __restrict__ beta_th, __nv_bfloat16* __restrict__ beta_ht, @@ -90,7 +82,7 @@ CUresult k3_flash_kda_beta_transpose( t_total, h ); - return k3_flash_kda_map_cuda_error(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); } // One sequence of `t_total` tokens through the chunkwise KDA forward, @@ -145,7 +137,7 @@ CUresult k3_flash_kda_fwd( gate_scale, stream ); - return k3_flash_kda_map_cuda_error(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); #endif } @@ -198,7 +190,7 @@ CUresult k3_flash_kda_fwd_md( gate_scale, stream ); - return k3_flash_kda_map_cuda_error(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); #endif } diff --git a/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu b/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu index a03a17170..d23a28644 100644 --- a/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu +++ b/pegainfer-kernels/csrc/k3/k3_flash_mla_prefill.cu @@ -44,6 +44,8 @@ #include +#include "../shared/ffi_guard.cuh" + #ifdef K3_FLASH_MLA_SM100F #include "collective/fmha_fusion.hpp" #include "collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp" @@ -59,18 +61,8 @@ // Not in an anonymous namespace: nvcc's device-stub generation trips over an // anonymous-namespace __global__ when an included header (cute) opens its own. -static CUresult k3_flash_mla_map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - static CUresult k3_flash_mla_consume_last_cuda_error() { - return k3_flash_mla_map_cuda_error(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); } // The paged latent row: post-norm kv latent | shared per-token rope half. @@ -383,9 +375,9 @@ static CUresult k3_flash_mla_prefill_run( using Operation = typename Cfg::Operation; int device = 0; cudaError_t err = cudaGetDevice(&device); - if (err != cudaSuccess) return k3_flash_mla_map_cuda_error(err); + if (err != cudaSuccess) return map_cuda_error(err); err = k3_flash_mla_prefill_ensure_smem_attr(device); - if (err != cudaSuccess) return k3_flash_mla_map_cuda_error(err); + if (err != cudaSuccess) return map_cuda_error(err); cutlass::KernelHardwareInfo hw_info; hw_info.device_id = device; hw_info.sm_count = diff --git a/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100.cu b/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100.cu index c50403758..cb78545c2 100644 --- a/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100.cu +++ b/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100.cu @@ -113,18 +113,8 @@ constexpr int kTransformSfWordK = kTransformSfGroupK * kTransformSfPerWord; // // Weight-row interleave granularity for the fused gate|up projection. constexpr int kInterleaveGran = 8; -CUresult transform_map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { - return transform_map_cuda_error(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); } // --------------------------------------------------------------------------- @@ -417,26 +407,26 @@ CUresult k3_mega_open_peer_access(int self_ordinal, int peer_ordinal) { if (self_ordinal == peer_ordinal) return CUDA_SUCCESS; int reachable = 0; const cudaError_t query = cudaDeviceCanAccessPeer(&reachable, self_ordinal, peer_ordinal); - if (query != cudaSuccess) return transform_map_cuda_error(query); + if (query != cudaSuccess) return map_cuda_error(query); if (reachable == 0) return CUDA_ERROR_PEER_ACCESS_UNSUPPORTED; const cudaError_t bind = cudaSetDevice(self_ordinal); - if (bind != cudaSuccess) return transform_map_cuda_error(bind); + if (bind != cudaSuccess) return map_cuda_error(bind); const cudaError_t enable = cudaDeviceEnablePeerAccess(peer_ordinal, 0); if (enable == cudaErrorPeerAccessAlreadyEnabled) { (void)cudaGetLastError(); } else if (enable != cudaSuccess) { - return transform_map_cuda_error(enable); + return map_cuda_error(enable); } cudaMemPool_t pool = nullptr; const cudaError_t pool_err = cudaDeviceGetDefaultMemPool(&pool, self_ordinal); - if (pool_err != cudaSuccess) return transform_map_cuda_error(pool_err); + if (pool_err != cudaSuccess) return map_cuda_error(pool_err); cudaMemAccessDesc desc{}; desc.location.type = cudaMemLocationTypeDevice; desc.location.id = peer_ordinal; desc.flags = cudaMemAccessFlagsProtReadWrite; - return transform_map_cuda_error(cudaMemPoolSetAccess(pool, &desc, 1)); + return map_cuda_error(cudaMemPoolSetAccess(pool, &desc, 1)); PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) } diff --git a/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100_common.cuh b/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100_common.cuh index 431167e37..4103740d3 100644 --- a/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100_common.cuh +++ b/pegainfer-kernels/csrc/k3/k3_mega_moe_sm100_common.cuh @@ -17,6 +17,8 @@ #include #include +#include "../shared/ffi_guard.cuh" + // `layout::Data`'s constructor uses DG_UNIFIED_ASSERT, which expands to a raw // `asm("trap;")`. That is device-only asm, but `MegaMoEBuffer` is also built on // the host here (and inside `CUTLASS_HOST_DEVICE` constexpr code), so give the @@ -90,16 +92,6 @@ constexpr int kNumNonEpilogueThreads = 128; constexpr int cdiv(int a, int b) { return (a + b - 1) / b; } constexpr int alignup(int a, int b) { return cdiv(a, b) * b; } -inline CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - // `get_symm_buffer_size_for_mega_moe`'s ring-capacity loop. constexpr int mega_ring_tokens(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_topk, int hidden, int intermediate, int num_sms) { diff --git a/pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu b/pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu index dba4c45ef..d1d9c41bf 100644 --- a/pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu +++ b/pegainfer-kernels/csrc/k3/k3_mla_paged_attn.cu @@ -252,16 +252,6 @@ __global__ void mla_paged_absorbed_attn_kernel( o[((size_t)bb * heads + bh) * (size_t)kVd + dv] = __float2bfloat16_rn(acc); } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - } // namespace extern "C" { diff --git a/pegainfer-kernels/csrc/k3/k3_moe_chain.cu b/pegainfer-kernels/csrc/k3/k3_moe_chain.cu index ae35cc917..623f55943 100644 --- a/pegainfer-kernels/csrc/k3/k3_moe_chain.cu +++ b/pegainfer-kernels/csrc/k3/k3_moe_chain.cu @@ -341,16 +341,6 @@ __global__ void moe_weighted_combine_kernel( out[(size_t)token * hidden + col] = __float2bfloat16_rn(acc); } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { return map_cuda_error(cudaGetLastError()); } int entry_grid(int entries) { diff --git a/pegainfer-kernels/csrc/k3/k3_router_topk.cu b/pegainfer-kernels/csrc/k3/k3_router_topk.cu index 1bddfbe72..f1bf07c5d 100644 --- a/pegainfer-kernels/csrc/k3/k3_router_topk.cu +++ b/pegainfer-kernels/csrc/k3/k3_router_topk.cu @@ -113,16 +113,6 @@ __global__ void router_topk_kernel(const float* __restrict__ s, } } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - } // namespace extern "C" { diff --git a/pegainfer-kernels/csrc/k3/k3_whale_doorbell.cu b/pegainfer-kernels/csrc/k3/k3_whale_doorbell.cu index b5af2cde6..931d374d5 100644 --- a/pegainfer-kernels/csrc/k3/k3_whale_doorbell.cu +++ b/pegainfer-kernels/csrc/k3/k3_whale_doorbell.cu @@ -55,9 +55,7 @@ CUresult k3_whale_doorbell_ring(const unsigned long long* flag_addrs, int flag_c args.value = value; args.flag_count = flag_count; whale_doorbell_kernel<<<1, 1, 0, stream>>>(args); - cudaError_t err = cudaGetLastError(); - if (err == cudaSuccess) return CUDA_SUCCESS; - return err == cudaErrorInvalidValue ? CUDA_ERROR_INVALID_VALUE : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(cudaGetLastError()); PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) } diff --git a/pegainfer-kernels/csrc/kimi_k2/kimi_experts.cu b/pegainfer-kernels/csrc/kimi_k2/kimi_experts.cu index 3cdc6c8cd..7f6d01983 100644 --- a/pegainfer-kernels/csrc/kimi_k2/kimi_experts.cu +++ b/pegainfer-kernels/csrc/kimi_k2/kimi_experts.cu @@ -1,5 +1,6 @@ #include "../common.cuh" #include "../marlin/ffi.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -75,7 +76,7 @@ CUresult kimi_add_f32_bf16_to_bf16_cuda( int blocks = (n + threads - 1) / threads; kimi_add_f32_bf16_to_bf16_kernel<<>>(a, b, out, n); cudaError_t err = cudaGetLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } CUresult kimi_residual_add_scaled_f32_cuda( @@ -96,7 +97,7 @@ CUresult kimi_residual_add_scaled_f32_cuda( kimi_residual_add_scaled_f32_kernel<<>>( hidden, projected, routed_f32, scale, out, n); cudaError_t err = cudaGetLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } __global__ void kimi_residual_add_scaled_bf16_kernel( @@ -136,7 +137,7 @@ CUresult kimi_residual_add_scaled_bf16_cuda( kimi_residual_add_scaled_bf16_kernel<<>>( hidden, projected, routed, scale, out, n); cudaError_t err = cudaGetLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } // Marlin routing metadata from a DeepEP post-epilogue expert prefix sum. @@ -249,7 +250,7 @@ CUresult kimi_deepep_build_marlin_routing_on_stream( max_padded_tokens, max_m_blocks); cudaError_t err = cudaGetLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } } // extern "C" diff --git a/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_int4.cu b/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_int4.cu index 8a01d3eff..4ed393669 100644 --- a/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_int4.cu +++ b/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_int4.cu @@ -1,4 +1,5 @@ #include "../marlin/ffi.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -139,7 +140,7 @@ CUresult kimi_marlin_int4_reorder_scale_cuda( kimi_marlin_reorder_scale_kernel<<>>( weight_scale_checkpoint, weight_scale_marlin, out_dim, scale_k, total_elements); cudaError_t err = cudaPeekAtLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_INVALID_VALUE; + return map_cuda_error(err); } CUresult kimi_marlin_int4_fuse_w13_cuda( @@ -173,7 +174,7 @@ CUresult kimi_marlin_int4_fuse_w13_cuda( reinterpret_cast(up_weight_packed_marlin), reinterpret_cast(w13_weight_packed_marlin), in_dim, intermediate_dim); cudaError_t err = cudaPeekAtLastError(); - if (err != cudaSuccess) return CUDA_ERROR_INVALID_VALUE; + if (err != cudaSuccess) return map_cuda_error(err); size_t scale_elements = static_cast(local_experts) * (in_dim / group_size) * intermediate_dim; @@ -181,7 +182,7 @@ CUresult kimi_marlin_int4_fuse_w13_cuda( kimi_marlin_fuse_w13_scale_kernel<<>>( gate_scale_marlin, up_scale_marlin, w13_scale_marlin, in_dim, intermediate_dim); err = cudaPeekAtLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_INVALID_VALUE; + return map_cuda_error(err); } } // extern "C" diff --git a/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_wna16.cu b/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_wna16.cu index b52ad9f4b..e6ba6b257 100644 --- a/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_wna16.cu +++ b/pegainfer-kernels/csrc/kimi_k2/kimi_marlin_wna16.cu @@ -212,7 +212,7 @@ CUresult kimi_marlin_w13_swiglu_cuda( int blocks = (total + threads - 1) / threads; pegainfer_kimi_marlin_moe_wna16::swiglu_w13_kernel<<>>( w13, out, rows, intermediate_dim); - return pegainfer_kimi_marlin_moe_wna16::last_error_to_cu(cudaPeekAtLastError()); + return map_cuda_error(cudaPeekAtLastError()); } CUresult kimi_marlin_w13_swiglu_expanded_cuda( @@ -233,9 +233,10 @@ CUresult kimi_marlin_w13_swiglu_expanded_cuda( if (sm_count <= 0) { int dev = 0; cudaError_t err = cudaGetDevice(&dev); - if (err != cudaSuccess) return CUDA_ERROR_INVALID_VALUE; + if (err != cudaSuccess) return map_cuda_error(err); err = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); - if (err != cudaSuccess || sm_count <= 0) return CUDA_ERROR_INVALID_VALUE; + if (err != cudaSuccess) return map_cuda_error(err); + if (sm_count <= 0) return CUDA_ERROR_INVALID_VALUE; } constexpr int threads = 256; // Fixed occupancy-sized grid; the kernel grid-strides up to the actual @@ -245,7 +246,7 @@ CUresult kimi_marlin_w13_swiglu_expanded_cuda( if (blocks > max_blocks) blocks = max_blocks; pegainfer_kimi_marlin_moe_wna16::swiglu_w13_expanded_kernel<<>>( w13, out, num_tokens_post_padded, max_rows, intermediate_dim); - return pegainfer_kimi_marlin_moe_wna16::last_error_to_cu(cudaPeekAtLastError()); + return map_cuda_error(cudaPeekAtLastError()); } CUresult kimi_marlin_sum_topk_rows_f32_cuda( @@ -265,7 +266,7 @@ CUresult kimi_marlin_sum_topk_rows_f32_cuda( int blocks = (total + threads - 1) / threads; pegainfer_kimi_marlin_moe_wna16::sum_topk_rows_kernel<<>>( route_output, out, active_tokens, topk, hidden_dim); - return pegainfer_kimi_marlin_moe_wna16::last_error_to_cu(cudaPeekAtLastError()); + return map_cuda_error(cudaPeekAtLastError()); } } // extern "C" diff --git a/pegainfer-kernels/csrc/kimi_k2/kimi_mla_cublaslt.cu b/pegainfer-kernels/csrc/kimi_k2/kimi_mla_cublaslt.cu index 8a8ee9b99..802d1d908 100644 --- a/pegainfer-kernels/csrc/kimi_k2/kimi_mla_cublaslt.cu +++ b/pegainfer-kernels/csrc/kimi_k2/kimi_mla_cublaslt.cu @@ -5,6 +5,8 @@ #include #include +#include "../shared/ffi_guard.cuh" + namespace { using DType = __nv_bfloat16; @@ -13,7 +15,6 @@ constexpr int kKvLoraRank = 512; constexpr int kNopeDim = 128; constexpr int kVHeadDim = 128; constexpr int kKvBHeadDim = kNopeDim + kVHeadDim; -constexpr int kCublasStatusErrorOffset = 100000; constexpr int kMlaLtLocalHeads = 64; constexpr int kMlaLtMaxBatch = 8; @@ -37,7 +38,7 @@ int cublas_status_to_error(cublasStatus_t status) { if (status == CUBLAS_STATUS_SUCCESS) { return static_cast(cudaSuccess); } - return kCublasStatusErrorOffset + static_cast(status); + return PEGAINFER_CUBLAS_STATUS_BASE + static_cast(status); } void destroy_mla_lt_plan(MlaLtPlan& plan) { diff --git a/pegainfer-kernels/csrc/kimi_k2/kimi_o_proj.cu b/pegainfer-kernels/csrc/kimi_k2/kimi_o_proj.cu index 951364c17..ef0bff2de 100644 --- a/pegainfer-kernels/csrc/kimi_k2/kimi_o_proj.cu +++ b/pegainfer-kernels/csrc/kimi_k2/kimi_o_proj.cu @@ -5,7 +5,8 @@ #include #include -static constexpr int CUBLAS_STATUS_ERROR_OFFSET = 100000; +#include "../shared/ffi_guard.cuh" + static constexpr int KIMI_O_PROJ_M = 7168; static constexpr int KIMI_O_PROJ_K = 8192; static constexpr int KIMI_O_PROJ_MAX_BATCH = 64; @@ -14,7 +15,7 @@ static int cublas_status_to_error(cublasStatus_t status) { if (status == CUBLAS_STATUS_SUCCESS) { return static_cast(cudaSuccess); } - return CUBLAS_STATUS_ERROR_OFFSET + static_cast(status); + return PEGAINFER_CUBLAS_STATUS_BASE + static_cast(status); } struct KimiOProjLtPlan { diff --git a/pegainfer-kernels/csrc/kimi_k2/kimi_router.cu b/pegainfer-kernels/csrc/kimi_k2/kimi_router.cu index ba10cff67..272505818 100644 --- a/pegainfer-kernels/csrc/kimi_k2/kimi_router.cu +++ b/pegainfer-kernels/csrc/kimi_k2/kimi_router.cu @@ -1,4 +1,5 @@ #include "../common.cuh" +#include "../shared/ffi_guard.cuh" #include #include @@ -90,16 +91,6 @@ __global__ void router_scores_topk_normalize_kernel( } } -CUresult map_cuda_error(cudaError_t err) { - if (err == cudaSuccess) return CUDA_SUCCESS; - if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) { - return CUDA_ERROR_INVALID_VALUE; - } - if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY; - if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED; - return CUDA_ERROR_LAUNCH_FAILED; -} - CUresult consume_last_cuda_error() { cudaError_t err = cudaGetLastError(); return map_cuda_error(err); diff --git a/pegainfer-kernels/csrc/kimi_k2/kimi_shared_gate_up.cu b/pegainfer-kernels/csrc/kimi_k2/kimi_shared_gate_up.cu index 8208b2af9..05a5ce7f8 100644 --- a/pegainfer-kernels/csrc/kimi_k2/kimi_shared_gate_up.cu +++ b/pegainfer-kernels/csrc/kimi_k2/kimi_shared_gate_up.cu @@ -5,7 +5,8 @@ #include #include -static constexpr int CUBLAS_STATUS_ERROR_OFFSET = 100000; +#include "../shared/ffi_guard.cuh" + static constexpr int KIMI_SHARED_GATE_UP_M = 4096; static constexpr int KIMI_SHARED_GATE_UP_K = 7168; static constexpr int KIMI_SHARED_GATE_UP_MAX_BATCH = 64; @@ -14,7 +15,7 @@ static int cublas_status_to_error(cublasStatus_t status) { if (status == CUBLAS_STATUS_SUCCESS) { return static_cast(cudaSuccess); } - return CUBLAS_STATUS_ERROR_OFFSET + static_cast(status); + return PEGAINFER_CUBLAS_STATUS_BASE + static_cast(status); } struct KimiSharedGateUpLtPlan { diff --git a/pegainfer-kernels/csrc/lora_fused.cu b/pegainfer-kernels/csrc/lora_fused.cu index 233b2ef6f..86ba89168 100644 --- a/pegainfer-kernels/csrc/lora_fused.cu +++ b/pegainfer-kernels/csrc/lora_fused.cu @@ -1,4 +1,5 @@ #include "common.cuh" +#include "shared/ffi_guard.cuh" #include #include @@ -465,7 +466,7 @@ static CUresult launch_lora_decode_fused_delta_rank( lora_decode_fused_delta_rank_kernel<<>>( a_packed, b_packed, scales, token_slots, input, out, batch, max_loras, max_rank, in_dim, out_dim, out_hidden_dim, row_offset); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } static CUresult launch_lora_decode_fused_delta( @@ -528,7 +529,7 @@ static CUresult launch_lora_decode_fused_delta( lora_decode_fused_delta_kernel<<>>( a_packed, b_packed, scales, token_slots, input, out, batch, max_loras, max_rank, rank, in_dim, out_dim, out_hidden_dim, row_offset); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } } @@ -571,7 +572,7 @@ static CUresult launch_lora_decode_fused_delta_group3_rank( scales1, out1, rank1, out_dim1, out_hidden_dim1, a2, b2, scales2, out2, rank2, out_dim2, out_hidden_dim2, token_slots, input, batch, max_loras, max_rank, in_dim); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } static CUresult launch_lora_decode_fused_delta_group3( @@ -668,7 +669,7 @@ static CUresult launch_lora_decode_fused_delta_group3( scales1, out1, rank1, out_dim1, out_hidden_dim1, a2, b2, scales2, out2, rank2, out_dim2, out_hidden_dim2, token_slots, input, batch, max_loras, max_rank, in_dim); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } } @@ -691,7 +692,7 @@ CUresult lora_pack_b_rows_cuda( dim3 grid((total + block.x - 1) / block.x); lora_pack_b_rows_kernel<<>>( src, dst, rank, max_rank, out_dim); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult lora_decode_fused_delta_cuda( diff --git a/pegainfer-kernels/csrc/marlin/moe_align.cu b/pegainfer-kernels/csrc/marlin/moe_align.cu index c865f64d6..a9caa429c 100644 --- a/pegainfer-kernels/csrc/marlin/moe_align.cu +++ b/pegainfer-kernels/csrc/marlin/moe_align.cu @@ -9,6 +9,7 @@ // The contract is vLLM's `moe_align_block_size`. #include "ffi.cuh" +#include "../shared/ffi_guard.cuh" namespace { @@ -205,7 +206,7 @@ CUresult marlin_moe_align_block_size_cuda( topk_idx, sorted_token_ids, expert_ids, num_tokens_post_padded, expert_offsets, route_elems, global_start, local_experts, block_size); cudaError_t err = cudaGetLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } int clear_elems = max_padded_tokens; @@ -216,24 +217,24 @@ CUresult marlin_moe_align_block_size_cuda( sorted_token_ids, expert_ids, num_tokens_post_padded, expert_offsets, route_elems, local_experts, max_padded_tokens, max_m_blocks); cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) return CUDA_ERROR_LAUNCH_FAILED; + if (err != cudaSuccess) return map_cuda_error(err); int route_blocks = (route_elems + threads - 1) / threads; marlin_moe_align_count_kernel<<>>( topk_idx, expert_offsets, route_elems, global_start, local_experts); err = cudaGetLastError(); - if (err != cudaSuccess) return CUDA_ERROR_LAUNCH_FAILED; + if (err != cudaSuccess) return map_cuda_error(err); marlin_moe_align_prefix_kernel<<<1, 1, 0, stream>>>( expert_ids, num_tokens_post_padded, expert_offsets, local_experts, block_size); err = cudaGetLastError(); - if (err != cudaSuccess) return CUDA_ERROR_LAUNCH_FAILED; + if (err != cudaSuccess) return map_cuda_error(err); int expert_blocks = (local_experts + threads - 1) / threads; marlin_moe_align_stable_fill_kernel<<>>( topk_idx, sorted_token_ids, expert_offsets, route_elems, global_start, local_experts); err = cudaGetLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; + return map_cuda_error(err); } } // extern "C" diff --git a/pegainfer-kernels/csrc/marlin/moe_wna16_launch.cuh b/pegainfer-kernels/csrc/marlin/moe_wna16_launch.cuh index c25f04ace..e3b9713ef 100644 --- a/pegainfer-kernels/csrc/marlin/moe_wna16_launch.cuh +++ b/pegainfer-kernels/csrc/marlin/moe_wna16_launch.cuh @@ -12,6 +12,8 @@ #pragma once +#include "../shared/ffi_guard.cuh" + namespace MARLIN_NAMESPACE_NAME { __global__ void MarlinDefault(MARLIN_KERNEL_PARAMS) {} @@ -53,10 +55,6 @@ struct ThreadConfigs { int large_count; }; -inline CUresult last_error_to_cu(cudaError_t err) { - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_LAUNCH_FAILED; -} - inline int get_scales_cache_size( ThreadConfig const& th_config, int prob_k, @@ -267,14 +265,16 @@ inline CUresult launch_marlin_moe_gemm( } int dev = 0; cudaError_t err = cudaGetDevice(&dev); - if (err != cudaSuccess) return CUDA_ERROR_INVALID_VALUE; + if (err != cudaSuccess) return map_cuda_error(err); if (sm_count <= 0) { err = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); - if (err != cudaSuccess || sm_count <= 0) return CUDA_ERROR_INVALID_VALUE; + if (err != cudaSuccess) return map_cuda_error(err); + if (sm_count <= 0) return CUDA_ERROR_INVALID_VALUE; } int max_shared_mem = 0; err = cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - if (err != cudaSuccess || max_shared_mem <= 0) return CUDA_ERROR_INVALID_VALUE; + if (err != cudaSuccess) return map_cuda_error(err); + if (max_shared_mem <= 0) return CUDA_ERROR_INVALID_VALUE; int thread_m_blocks = div_ceil(moe_block_size, 16); bool m_block_size_8 = moe_block_size == 8; @@ -306,7 +306,7 @@ inline CUresult launch_marlin_moe_gemm( err = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); - if (err != cudaSuccess) return CUDA_ERROR_INVALID_VALUE; + if (err != cudaSuccess) return map_cuda_error(err); const int4* A_ptr = reinterpret_cast(input); const int4* B_ptr = reinterpret_cast(b_qweight); @@ -320,7 +320,7 @@ inline CUresult launch_marlin_moe_gemm( nullptr, nullptr, sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, top_k, mul_topk_weights, size_k / group_size, size_m, size_n, size_k, locks, false, use_atomic_add, use_fp32_reduce); - return last_error_to_cu(cudaPeekAtLastError()); + return map_cuda_error(cudaPeekAtLastError()); } } // namespace MARLIN_NAMESPACE_NAME diff --git a/pegainfer-kernels/csrc/marlin/repack_4bit.cu b/pegainfer-kernels/csrc/marlin/repack_4bit.cu index b3573df93..9ce1dbc0e 100644 --- a/pegainfer-kernels/csrc/marlin/repack_4bit.cu +++ b/pegainfer-kernels/csrc/marlin/repack_4bit.cu @@ -6,6 +6,7 @@ // plane per expert, which is how both checkpoints store it. #include "ffi.cuh" +#include "../shared/ffi_guard.cuh" namespace { @@ -114,7 +115,7 @@ CUresult marlin_repack_4bit_cuda( reinterpret_cast(src), reinterpret_cast(dst), in_dim, out_dim); cudaError_t err = cudaPeekAtLastError(); - return err == cudaSuccess ? CUDA_SUCCESS : CUDA_ERROR_INVALID_VALUE; + return map_cuda_error(err); } } // extern "C" diff --git a/pegainfer-kernels/csrc/qwen35/paged_attention_hd256.cu b/pegainfer-kernels/csrc/qwen35/paged_attention_hd256.cu new file mode 100644 index 000000000..819403a3b --- /dev/null +++ b/pegainfer-kernels/csrc/qwen35/paged_attention_hd256.cu @@ -0,0 +1,45 @@ +// FlashInfer paged attention at HEAD_DIM=256 (Qwen3.5 full-attention +// layers): batch decode and batch prefill, causal, no window. + +#include "../shared/paged_launch.cuh" + +extern "C" { + +int paged_attention_decode_cuda_hd256( + void* q, void* output, void* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, + int32_t* request_indices, int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, + int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, + int32_t page_size, int32_t batch_size, int64_t stride_page, + float sm_scale, void* stream) +{ + return decode_launch( + q, output, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indptr, last_page_len_d, + request_indices, kv_tile_indices, kv_chunk_size_ptr, + num_qo_heads, num_kv_heads, head_dim, page_size, batch_size, + stride_page, sm_scale, /*window_left=*/-1, stream); +} + +int batch_prefill_paged_cuda_hd256( + void* q, void* output, void* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, + int32_t* q_indptr, int32_t* request_indices, int32_t* qo_tile_indices, + int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, uint32_t* total_num_rows, + int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, + int32_t page_size, int32_t seq_len, int32_t batch_size, + int32_t padded_batch_size, int64_t stride_page, float sm_scale, + void* stream) +{ + return prefill_paged_launch( + q, output, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indptr, last_page_len_d, q_indptr, + request_indices, qo_tile_indices, kv_tile_indices, kv_chunk_size_ptr, + total_num_rows, num_qo_heads, num_kv_heads, head_dim, page_size, + seq_len, batch_size, padded_batch_size, stride_page, sm_scale, + /*cta_tile_q_override=*/0, /*window_left=*/-1, stream); +} + +} // extern "C" diff --git a/pegainfer-kernels/csrc/shared/elementwise.cu b/pegainfer-kernels/csrc/shared/elementwise.cu index f0506e564..18cc174fc 100644 --- a/pegainfer-kernels/csrc/shared/elementwise.cu +++ b/pegainfer-kernels/csrc/shared/elementwise.cu @@ -1,4 +1,5 @@ #include "common.cuh" +#include "ffi_guard.cuh" #include #include #include @@ -555,7 +556,7 @@ CUresult add_cuda( int block = 256; int grid = (n + block - 1) / block; add_kernel<<>>(a, b, out, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult add_scaled_bf16_cuda( @@ -570,7 +571,7 @@ CUresult add_scaled_bf16_cuda( int grid = (n + block - 1) / block; add_scaled_bf16_kernel<<>>( routed, scale, shared, out, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult scaled_add_rows_cuda( @@ -595,7 +596,7 @@ CUresult scaled_add_rows_cuda( dim3 grid(grid_x, grid_y); scaled_add_rows_kernel<<>>( delta, scale, out, out_hidden_dim, row_offset, rows, seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult gather_hidden_tokens_cuda( @@ -615,7 +616,7 @@ CUresult gather_hidden_tokens_cuda( int grid = (total + block - 1) / block; gather_hidden_tokens_kernel<<>>( input, token_indices, out, hidden_dim, token_count, input_seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult extract_hidden_rows_cuda( @@ -637,7 +638,7 @@ CUresult extract_hidden_rows_cuda( int grid = (total + block - 1) / block; extract_hidden_rows_kernel<<>>( src, dst, src_hidden_dim, dst_hidden_dim, col_offset, rows, seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult copy_hidden_rows_cuda( @@ -659,7 +660,7 @@ CUresult copy_hidden_rows_cuda( int grid = (total + block - 1) / block; copy_hidden_rows_kernel<<>>( src, dst, src_hidden_dim, dst_hidden_dim, row_offset, rows, seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult mask_position_zero_rows_cuda( @@ -678,7 +679,7 @@ CUresult mask_position_zero_rows_cuda( int grid = (total + block - 1) / block; mask_position_zero_rows_kernel<<>>( src, positions, dst, hidden_dim, rows); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult copy_hidden_token_range_cuda( @@ -703,7 +704,7 @@ CUresult copy_hidden_token_range_cuda( int grid = (total + block - 1) / block; copy_hidden_token_range_kernel<<>>( src, dst, hidden_dim, src_token_offset, dst_token_offset, token_count); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult scaled_add_rows_indexed_cuda( @@ -732,7 +733,7 @@ CUresult scaled_add_rows_indexed_cuda( scaled_add_rows_indexed_kernel<<>>( delta, scale, token_indices, out, out_hidden_dim, row_offset, rows, token_count, out_seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult store_rows_indexed_cuda( @@ -760,7 +761,7 @@ CUresult store_rows_indexed_cuda( store_rows_indexed_kernel<<>>( src, token_indices, out, out_hidden_dim, row_offset, rows, token_count, out_seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult bf16_to_f32_cuda( @@ -768,7 +769,7 @@ CUresult bf16_to_f32_cuda( int block = 256; int grid = (n + block - 1) / block; bf16_to_f32_kernel<<>>(input, output, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult f32_to_bf16_cuda( @@ -776,14 +777,14 @@ CUresult f32_to_bf16_cuda( int block = 256; int grid = (n + block - 1) / block; f32_to_bf16_kernel<<>>(input, output, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult scale_f32_cuda(float *values, float scale, int n, cudaStream_t stream) { int block = 256; int grid = (n + block - 1) / block; scale_f32_kernel<<>>(values, scale, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult accumulate_bf16_token_scaled_to_f32_cuda( @@ -802,7 +803,7 @@ CUresult accumulate_bf16_token_scaled_to_f32_cuda( int grid = (hidden_dim + block - 1) / block; accumulate_bf16_token_scaled_to_f32_kernel<<>>( token, scale, out, hidden_dim, token_idx, seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult repeat_f32_for_reduce_scatter_cuda( @@ -813,7 +814,7 @@ CUresult repeat_f32_for_reduce_scatter_cuda( int grid = (total + block - 1) / block; repeat_f32_rows_for_reduce_scatter_kernel<<>>( local, repeated, local_elems, world_size); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult silu_mul_triton_aot_cuda( @@ -822,7 +823,7 @@ CUresult silu_mul_triton_aot_cuda( int block = 256; int grid = (n + block - 1) / block; silu_mul_kernel<<>>(gate, up, out, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } // Eight elements a thread through 16-byte loads and stores, the same @@ -883,7 +884,7 @@ CUresult gelu_tanh_mul_cuda( gelu_tanh_mul_kernel<<>>( gate, up, out, cols, gate_stride, up_stride, n); } - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult scale_bf16_in_place_cuda( @@ -891,7 +892,7 @@ CUresult scale_bf16_in_place_cuda( int block = 256; int grid = n / block + (n % block != 0); scale_bf16_kernel<<>>(buf, scale, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult softcap_bf16_in_place_cuda( @@ -899,7 +900,7 @@ CUresult softcap_bf16_in_place_cuda( int block = 256; int grid = n / block + (n % block != 0); softcap_bf16_kernel<<>>(buf, cap, n); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult suppress_logits_bf16_in_place_cuda( @@ -914,7 +915,7 @@ CUresult suppress_logits_bf16_in_place_cuda( int grid = total / block + (total % block != 0); suppress_logits_bf16_kernel<<>>( logits, ids, vocab, id_count, total); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult embedding_decode_cuda( @@ -923,7 +924,7 @@ CUresult embedding_decode_cuda( int block = 256; int grid = (hidden_size + block - 1) / block; embedding_decode_kernel<<>>(embed, token_id, out, hidden_size); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } // Whether the vectorized row-copy embedding kernels can serve this call: @@ -952,13 +953,13 @@ CUresult embedding_batched_cuda( embedding_batched_vec4_kernel<<>>( reinterpret_cast(embed), token_ids, reinterpret_cast(out), hidden_vec, seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } int block = 256; int total = hidden_size * seq_len; int grid = (total + block - 1) / block; embedding_batched_kernel<<>>(embed, token_ids, out, hidden_size, seq_len); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult embedding_batched_vocab_shard_cuda( @@ -975,14 +976,14 @@ CUresult embedding_batched_vocab_shard_cuda( reinterpret_cast(embed), token_ids, reinterpret_cast(out), hidden_vec, seq_len, vocab_start, part_vocab_size); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } int block = 256; int total = hidden_size * seq_len; int grid = (total + block - 1) / block; embedding_batched_vocab_shard_kernel<<>>( embed, token_ids, out, hidden_size, seq_len, vocab_start, part_vocab_size); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } CUresult advance_decode_metadata_cuda( @@ -996,7 +997,7 @@ CUresult advance_decode_metadata_cuda( int grid = 1 + (rows - 1) / block; advance_decode_metadata_kernel<<>>( positions, local_last, pseudo_last, kv_chunk, rows, factor); - return (CUresult)cudaGetLastError(); + return map_cuda_error(cudaGetLastError()); } } // extern "C" diff --git a/pegainfer-kernels/csrc/shared/ffi_guard.cuh b/pegainfer-kernels/csrc/shared/ffi_guard.cuh index 805121171..944cd523d 100644 --- a/pegainfer-kernels/csrc/shared/ffi_guard.cuh +++ b/pegainfer-kernels/csrc/shared/ffi_guard.cuh @@ -3,10 +3,133 @@ // process with no message. #pragma once +#include +#include + #include void pegainfer_ffi_set_last_error(const char* what); +// Entry points that call cuBLAS return this plus the cublasStatus_t on a +// cuBLAS failure, and the cudaError_t otherwise; Rust's +// `ops::ffi_status_detail` decodes the pair. +constexpr int PEGAINFER_CUBLAS_STATUS_BASE = 100000; + +// The driver result an entry point that returns CUresult gives for a runtime +// error; every such conversion goes through here, so the same fault reads the +// same in Rust whichever kernel hit it. Each runtime error with a driver +// counterpart maps to it (the pairs share a value in the 12.8 and 13.1 +// headers). The runtime-only errors have no driver code and report +// CUDA_ERROR_UNKNOWN, apart from the deprecated invalid device pointer. The +// guarded cases are newer than CUDA 12.3, the oldest supported toolkit. +inline CUresult map_cuda_error(cudaError_t err) { + switch (err) { + case cudaSuccess: return CUDA_SUCCESS; + case cudaErrorInvalidValue: return CUDA_ERROR_INVALID_VALUE; + case cudaErrorInvalidDevicePointer: return CUDA_ERROR_INVALID_VALUE; + case cudaErrorMemoryAllocation: return CUDA_ERROR_OUT_OF_MEMORY; + case cudaErrorInitializationError: return CUDA_ERROR_NOT_INITIALIZED; + case cudaErrorCudartUnloading: return CUDA_ERROR_DEINITIALIZED; + case cudaErrorProfilerDisabled: return CUDA_ERROR_PROFILER_DISABLED; + case cudaErrorProfilerNotInitialized: return CUDA_ERROR_PROFILER_NOT_INITIALIZED; + case cudaErrorProfilerAlreadyStarted: return CUDA_ERROR_PROFILER_ALREADY_STARTED; + case cudaErrorProfilerAlreadyStopped: return CUDA_ERROR_PROFILER_ALREADY_STOPPED; + case cudaErrorStubLibrary: return CUDA_ERROR_STUB_LIBRARY; + case cudaErrorDevicesUnavailable: return CUDA_ERROR_DEVICE_UNAVAILABLE; + case cudaErrorNoDevice: return CUDA_ERROR_NO_DEVICE; + case cudaErrorInvalidDevice: return CUDA_ERROR_INVALID_DEVICE; + case cudaErrorDeviceNotLicensed: return CUDA_ERROR_DEVICE_NOT_LICENSED; + case cudaErrorInvalidKernelImage: return CUDA_ERROR_INVALID_IMAGE; + case cudaErrorDeviceUninitialized: return CUDA_ERROR_INVALID_CONTEXT; + case cudaErrorMapBufferObjectFailed: return CUDA_ERROR_MAP_FAILED; + case cudaErrorUnmapBufferObjectFailed: return CUDA_ERROR_UNMAP_FAILED; + case cudaErrorArrayIsMapped: return CUDA_ERROR_ARRAY_IS_MAPPED; + case cudaErrorAlreadyMapped: return CUDA_ERROR_ALREADY_MAPPED; + case cudaErrorNoKernelImageForDevice: return CUDA_ERROR_NO_BINARY_FOR_GPU; + case cudaErrorAlreadyAcquired: return CUDA_ERROR_ALREADY_ACQUIRED; + case cudaErrorNotMapped: return CUDA_ERROR_NOT_MAPPED; + case cudaErrorNotMappedAsArray: return CUDA_ERROR_NOT_MAPPED_AS_ARRAY; + case cudaErrorNotMappedAsPointer: return CUDA_ERROR_NOT_MAPPED_AS_POINTER; + case cudaErrorECCUncorrectable: return CUDA_ERROR_ECC_UNCORRECTABLE; + case cudaErrorUnsupportedLimit: return CUDA_ERROR_UNSUPPORTED_LIMIT; + case cudaErrorDeviceAlreadyInUse: return CUDA_ERROR_CONTEXT_ALREADY_IN_USE; + case cudaErrorPeerAccessUnsupported: return CUDA_ERROR_PEER_ACCESS_UNSUPPORTED; + case cudaErrorInvalidPtx: return CUDA_ERROR_INVALID_PTX; + case cudaErrorInvalidGraphicsContext: return CUDA_ERROR_INVALID_GRAPHICS_CONTEXT; + case cudaErrorNvlinkUncorrectable: return CUDA_ERROR_NVLINK_UNCORRECTABLE; + case cudaErrorJitCompilerNotFound: return CUDA_ERROR_JIT_COMPILER_NOT_FOUND; + case cudaErrorUnsupportedPtxVersion: return CUDA_ERROR_UNSUPPORTED_PTX_VERSION; + case cudaErrorJitCompilationDisabled: return CUDA_ERROR_JIT_COMPILATION_DISABLED; + case cudaErrorUnsupportedExecAffinity: return CUDA_ERROR_UNSUPPORTED_EXEC_AFFINITY; + case cudaErrorUnsupportedDevSideSync: return CUDA_ERROR_UNSUPPORTED_DEVSIDE_SYNC; + case cudaErrorInvalidSource: return CUDA_ERROR_INVALID_SOURCE; + case cudaErrorFileNotFound: return CUDA_ERROR_FILE_NOT_FOUND; + case cudaErrorSharedObjectSymbolNotFound: return CUDA_ERROR_SHARED_OBJECT_SYMBOL_NOT_FOUND; + case cudaErrorSharedObjectInitFailed: return CUDA_ERROR_SHARED_OBJECT_INIT_FAILED; + case cudaErrorOperatingSystem: return CUDA_ERROR_OPERATING_SYSTEM; + case cudaErrorInvalidResourceHandle: return CUDA_ERROR_INVALID_HANDLE; + case cudaErrorIllegalState: return CUDA_ERROR_ILLEGAL_STATE; + case cudaErrorLossyQuery: return CUDA_ERROR_LOSSY_QUERY; + case cudaErrorSymbolNotFound: return CUDA_ERROR_NOT_FOUND; + case cudaErrorNotReady: return CUDA_ERROR_NOT_READY; + case cudaErrorIllegalAddress: return CUDA_ERROR_ILLEGAL_ADDRESS; + case cudaErrorLaunchOutOfResources: return CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES; + case cudaErrorLaunchTimeout: return CUDA_ERROR_LAUNCH_TIMEOUT; + case cudaErrorLaunchIncompatibleTexturing: return CUDA_ERROR_LAUNCH_INCOMPATIBLE_TEXTURING; + case cudaErrorPeerAccessAlreadyEnabled: return CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED; + case cudaErrorPeerAccessNotEnabled: return CUDA_ERROR_PEER_ACCESS_NOT_ENABLED; + case cudaErrorSetOnActiveProcess: return CUDA_ERROR_PRIMARY_CONTEXT_ACTIVE; + case cudaErrorContextIsDestroyed: return CUDA_ERROR_CONTEXT_IS_DESTROYED; + case cudaErrorAssert: return CUDA_ERROR_ASSERT; + case cudaErrorTooManyPeers: return CUDA_ERROR_TOO_MANY_PEERS; + case cudaErrorHostMemoryAlreadyRegistered: return CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED; + case cudaErrorHostMemoryNotRegistered: return CUDA_ERROR_HOST_MEMORY_NOT_REGISTERED; + case cudaErrorHardwareStackError: return CUDA_ERROR_HARDWARE_STACK_ERROR; + case cudaErrorIllegalInstruction: return CUDA_ERROR_ILLEGAL_INSTRUCTION; + case cudaErrorMisalignedAddress: return CUDA_ERROR_MISALIGNED_ADDRESS; + case cudaErrorInvalidAddressSpace: return CUDA_ERROR_INVALID_ADDRESS_SPACE; + case cudaErrorInvalidPc: return CUDA_ERROR_INVALID_PC; + case cudaErrorLaunchFailure: return CUDA_ERROR_LAUNCH_FAILED; + case cudaErrorCooperativeLaunchTooLarge: return CUDA_ERROR_COOPERATIVE_LAUNCH_TOO_LARGE; + case cudaErrorNotPermitted: return CUDA_ERROR_NOT_PERMITTED; + case cudaErrorNotSupported: return CUDA_ERROR_NOT_SUPPORTED; + case cudaErrorSystemNotReady: return CUDA_ERROR_SYSTEM_NOT_READY; + case cudaErrorSystemDriverMismatch: return CUDA_ERROR_SYSTEM_DRIVER_MISMATCH; + case cudaErrorCompatNotSupportedOnDevice: return CUDA_ERROR_COMPAT_NOT_SUPPORTED_ON_DEVICE; + case cudaErrorMpsConnectionFailed: return CUDA_ERROR_MPS_CONNECTION_FAILED; + case cudaErrorMpsRpcFailure: return CUDA_ERROR_MPS_RPC_FAILURE; + case cudaErrorMpsServerNotReady: return CUDA_ERROR_MPS_SERVER_NOT_READY; + case cudaErrorMpsMaxClientsReached: return CUDA_ERROR_MPS_MAX_CLIENTS_REACHED; + case cudaErrorMpsMaxConnectionsReached: return CUDA_ERROR_MPS_MAX_CONNECTIONS_REACHED; + case cudaErrorMpsClientTerminated: return CUDA_ERROR_MPS_CLIENT_TERMINATED; + case cudaErrorCdpNotSupported: return CUDA_ERROR_CDP_NOT_SUPPORTED; + case cudaErrorCdpVersionMismatch: return CUDA_ERROR_CDP_VERSION_MISMATCH; + case cudaErrorStreamCaptureUnsupported: return CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED; + case cudaErrorStreamCaptureInvalidated: return CUDA_ERROR_STREAM_CAPTURE_INVALIDATED; + case cudaErrorStreamCaptureMerge: return CUDA_ERROR_STREAM_CAPTURE_MERGE; + case cudaErrorStreamCaptureUnmatched: return CUDA_ERROR_STREAM_CAPTURE_UNMATCHED; + case cudaErrorStreamCaptureUnjoined: return CUDA_ERROR_STREAM_CAPTURE_UNJOINED; + case cudaErrorStreamCaptureIsolation: return CUDA_ERROR_STREAM_CAPTURE_ISOLATION; + case cudaErrorStreamCaptureImplicit: return CUDA_ERROR_STREAM_CAPTURE_IMPLICIT; + case cudaErrorCapturedEvent: return CUDA_ERROR_CAPTURED_EVENT; + case cudaErrorStreamCaptureWrongThread: return CUDA_ERROR_STREAM_CAPTURE_WRONG_THREAD; + case cudaErrorTimeout: return CUDA_ERROR_TIMEOUT; + case cudaErrorGraphExecUpdateFailure: return CUDA_ERROR_GRAPH_EXEC_UPDATE_FAILURE; + case cudaErrorExternalDevice: return CUDA_ERROR_EXTERNAL_DEVICE; + case cudaErrorInvalidClusterSize: return CUDA_ERROR_INVALID_CLUSTER_SIZE; +#if CUDART_VERSION >= 12060 + case cudaErrorFunctionNotLoaded: return CUDA_ERROR_FUNCTION_NOT_LOADED; + case cudaErrorInvalidResourceType: return CUDA_ERROR_INVALID_RESOURCE_TYPE; + case cudaErrorInvalidResourceConfiguration: return CUDA_ERROR_INVALID_RESOURCE_CONFIGURATION; +#endif +#if CUDART_VERSION >= 12080 + case cudaErrorContained: return CUDA_ERROR_CONTAINED; + case cudaErrorTensorMemoryLeak: return CUDA_ERROR_TENSOR_MEMORY_LEAK; +#endif + default: return CUDA_ERROR_UNKNOWN; + } +} + // Entering a guard clears the previous message so a -1 seen by Rust never // reads a stale what() from an earlier call on the same thread. #define PEGAINFER_FFI_GUARD_BEGIN \ diff --git a/pegainfer-kernels/csrc/shared/flashinfer_norm.cu b/pegainfer-kernels/csrc/shared/flashinfer_norm.cu index e6c6511b7..de977d996 100644 --- a/pegainfer-kernels/csrc/shared/flashinfer_norm.cu +++ b/pegainfer-kernels/csrc/shared/flashinfer_norm.cu @@ -23,6 +23,8 @@ #include +#include "ffi_guard.cuh" + using DType = __nv_bfloat16; namespace pegainfer { @@ -580,7 +582,7 @@ CUresult rms_norm_batched_dual_cuda(const DType *x, const DType *weight_a, const pegainfer::norm::DualRMSNormKernel<<>>( x, weight_a, weight_b, out_a, out_b, d, eps, scale_a); }); - return static_cast(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); } CUresult dual_rms_norm_add_batched_cuda(const DType *a, const DType *weight_a, const DType *b, @@ -597,7 +599,7 @@ CUresult dual_rms_norm_add_batched_cuda(const DType *a, const DType *weight_a, c pegainfer::norm::DualRMSNormAddKernel<<>>( a, weight_a, b, weight_b, out, d, eps); }); - return static_cast(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); } CUresult rms_norm_add_rms_norm_round_batched_cuda(const DType *x, const DType *weight_post, @@ -625,7 +627,7 @@ CUresult rms_norm_add_rms_norm_round_batched_cuda(const DType *x, const DType *w err = cudaGetLastError(); } }); - return static_cast(err); + return map_cuda_error(err); } CUresult rms_norm_add_scale_batched_cuda(const DType *x, const DType *weight, @@ -643,7 +645,7 @@ CUresult rms_norm_add_scale_batched_cuda(const DType *x, const DType *weight, pegainfer::norm::RMSNormAddScaleKernel<<>>( x, weight, residual, out, d, eps, scale); }); - return static_cast(cudaGetLastError()); + return map_cuda_error(cudaGetLastError()); } // ============================================================================ @@ -698,7 +700,7 @@ CUresult fused_add_rms_norm_round_batched_cuda(DType *hidden, const DType *resid /*stride_residual=*/static_cast(hidden_dim), /*stride_out=*/static_cast(hidden_dim), eps, stream); - return static_cast(err); + return map_cuda_error(err); } // ============================================================================ @@ -822,7 +824,7 @@ CUresult layer_norm_cuda(const DType *x, const float *gamma, const float *beta, size_t shmem_size = std::min(block_size / 32, (n + 31) / 32) * sizeof(float); layer_norm_kernel<<>>(x, gamma, beta, out, n, eps); cudaError_t err = cudaGetLastError(); - return static_cast(err); + return map_cuda_error(err); } // ============================================================================ diff --git a/pegainfer-kernels/csrc/shared/linear.cu b/pegainfer-kernels/csrc/shared/linear.cu index 29f2e3d1e..9690d4058 100644 --- a/pegainfer-kernels/csrc/shared/linear.cu +++ b/pegainfer-kernels/csrc/shared/linear.cu @@ -2,6 +2,8 @@ #include #include +#include "ffi_guard.cuh" + #include #include @@ -16,13 +18,11 @@ #include #include -static constexpr int CUBLAS_STATUS_ERROR_OFFSET = 100000; - static int cublas_status_to_error(cublasStatus_t status) { if (status == CUBLAS_STATUS_SUCCESS) { return static_cast(cudaSuccess); } - return CUBLAS_STATUS_ERROR_OFFSET + static_cast(status); + return PEGAINFER_CUBLAS_STATUS_BASE + static_cast(status); } // cuBLAS handle management. diff --git a/pegainfer-kernels/csrc/shared/norm.cu b/pegainfer-kernels/csrc/shared/norm.cu index ebfed8f60..0fc86c80e 100644 --- a/pegainfer-kernels/csrc/shared/norm.cu +++ b/pegainfer-kernels/csrc/shared/norm.cu @@ -3,17 +3,23 @@ // The standard RMSNorm / FusedAddRMSNorm / Gemma variants have been migrated // to flashinfer_norm.cu which delegates to FlashInfer's header-only templates. // -// This file retains only: -// rms_norm_gated_cuda — per-head RMSNorm with SiLU gate (Qwen3.5 linear attention output). +// This file retains only the per-head gated RMSNorm used on a linear-attention +// output, in two activations: +// rms_norm_gated_cuda — SiLU gate (Qwen3.5 Gated DeltaNet) +// rms_norm_gated_sigmoid_cuda — sigmoid gate (Qwen3.8-Flash-Next, whose +// `output_gate_type` selects this) #include "common.cuh" // ============================================================================ // Gated RMSNorm for linear attention output: -// out = rms_norm(x, f32_weight) * silu(gate) +// out = rms_norm(x, f32_weight) * activation(gate) // Per-head normalization: x is [num_heads * head_dim], weight is [head_dim] (broadcast). // Grid: num_heads blocks, head_dim threads. // ============================================================================ +enum class GatedNormActivation { kSilu, kSigmoid }; + +template __global__ void rms_norm_gated_kernel( const __nv_bfloat16 *__restrict__ x, const float *__restrict__ weight, @@ -55,9 +61,15 @@ __global__ void rms_norm_gated_kernel( float normed = x_val * s_inv_rms * weight[tid]; float g = __bfloat162float(gate[offset]); - float silu_g = g / (1.0f + expf(-g)); + // SiLU keeps its original `g / (1 + exp(-g))` form rather than the equal + // `g * sigmoid(g)`: the two are not bit-identical in fp32, and Qwen3.5's gated + // path must stay bit-for-bit unchanged. Its golden gates would not catch the + // swap — they bound logprob drift, not this expression. + float activated = (kActivation == GatedNormActivation::kSilu) + ? g / (1.0f + expf(-g)) + : 1.0f / (1.0f + expf(-g)); - out[offset] = __float2bfloat16(normed * silu_g); + out[offset] = __float2bfloat16(normed * activated); } extern "C" { @@ -65,7 +77,13 @@ extern "C" { void rms_norm_gated_cuda(const __nv_bfloat16 *x, const float *weight, const __nv_bfloat16 *gate, __nv_bfloat16 *out, int num_heads, int head_dim, float eps, cudaStream_t stream) { - rms_norm_gated_kernel<<>>(x, weight, gate, out, head_dim, eps); + rms_norm_gated_kernel<<>>(x, weight, gate, out, head_dim, eps); +} + +void rms_norm_gated_sigmoid_cuda(const __nv_bfloat16 *x, const float *weight, + const __nv_bfloat16 *gate, __nv_bfloat16 *out, + int num_heads, int head_dim, float eps, cudaStream_t stream) { + rms_norm_gated_kernel<<>>(x, weight, gate, out, head_dim, eps); } } // extern "C" diff --git a/pegainfer-kernels/csrc/shared/paged_attention.cu b/pegainfer-kernels/csrc/shared/paged_attention.cu index 4f100a372..a88404046 100644 --- a/pegainfer-kernels/csrc/shared/paged_attention.cu +++ b/pegainfer-kernels/csrc/shared/paged_attention.cu @@ -2,11 +2,10 @@ // // We include FlashInfer headers (header-only C++) and instantiate only the // template variants needed: bf16 Q/KV/O, NHD layout, no RoPE, at HEAD_DIM 128 -// and 256 — the latter both with and without a sliding-window mask — plus the -// hd512 split-KV decode entry the global family reads through. +// (plus the head_dim 64 non-causal prefill below). // // FlashInfer's dispatchers internally instantiate multiple GQA group sizes -// (1,2,3,4,8) — this covers both Qwen3-4B (GQA=4) and Qwen3.5-4B (GQA=8). +// (1,2,3,4,8). #include "paged_launch.cuh" @@ -561,137 +560,4 @@ int single_decode_nhd_cuda( PEGAINFER_FFI_GUARD_END(-1) } -// --------------------------------------------------------------------------- -// Single-request prefill for HEAD_DIM=256 — wraps FlashInfer SinglePrefillWithKVCache. -// -// Identical to single_prefill_cuda but instantiated with HEAD_DIM_QK/VO=256. -// Reads Q from col-major [q_dim, seq_len] (HiddenStates layout). -// Reads K/V from contiguous HND cache: k[head, pos, dim]. -// No RoPE inside (caller does QK norm + partial RoPE beforehand). -// Causal mask, no split-KV. -// -// Used by Qwen3.5-4B multi-token prefill. Single-token decode still routes to -// the Triton AOT path (CUDA-Graph safe) until Phase 2d introduces paged decode. -// --------------------------------------------------------------------------- -int single_prefill_cuda_hd256( - void* q, - void* output, - void* k_cache, - void* v_cache, - int32_t num_qo_heads, - int32_t num_kv_heads, - int32_t seq_len, // number of Q tokens (qo_len) - int32_t kv_len, // total KV length (start_pos + seq_len) - int32_t max_seq_len, // allocated cache rows (for HND stride) - float sm_scale, - void* stream) -{ - return single_prefill_launch( - q, output, k_cache, v_cache, num_qo_heads, num_kv_heads, - /*head_dim=*/256, seq_len, kv_len, max_seq_len, - sm_scale, stream); -} - -// --------------------------------------------------------------------------- -// HEAD_DIM=256 paged entry points. Qwen3.5-4B uses the full-attention pair; -// the windowed pair applies the sliding-window mask for Gemma 4's local layers. -// --------------------------------------------------------------------------- -int paged_attention_decode_cuda_hd256( - void* q, void* output, void* kv_data, - int64_t k_offset_elems, int64_t v_offset_elems, - int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, - int32_t* request_indices, int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, - int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, - int32_t page_size, int32_t batch_size, int64_t stride_page, - float sm_scale, void* stream) -{ - return decode_launch( - q, output, kv_data, k_offset_elems, v_offset_elems, - page_indices, page_indptr, last_page_len_d, - request_indices, kv_tile_indices, kv_chunk_size_ptr, - num_qo_heads, num_kv_heads, head_dim, page_size, batch_size, - stride_page, sm_scale, /*window_left=*/-1, stream); -} - -int paged_attention_decode_window_cuda_hd256( - void* q, void* output, void* kv_data, - int64_t k_offset_elems, int64_t v_offset_elems, - int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, - int32_t* request_indices, int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, - int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, - int32_t page_size, int32_t batch_size, int64_t stride_page, - float sm_scale, int32_t window_left, void* stream) -{ - return decode_launch( - q, output, kv_data, k_offset_elems, v_offset_elems, - page_indices, page_indptr, last_page_len_d, - request_indices, kv_tile_indices, kv_chunk_size_ptr, - num_qo_heads, num_kv_heads, head_dim, page_size, batch_size, - stride_page, sm_scale, window_left, stream); -} - -// Split-KV decode at head_dim 512 — the Gemma global family's decode -// read. The non-partitioned grid is (pseudo-requests, kv heads) CTAs and -// starves the device; chunking the KV brings the grid to occupancy. -int paged_attention_decode_split_kv_cuda_hd512( - void* q, void* output, void* kv_data, - int64_t k_offset_elems, int64_t v_offset_elems, - int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, - int32_t* request_indices, int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, - int32_t* o_indptr, uint8_t* block_valid_mask, - void* tmp_v, float* tmp_s, - int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, - int32_t page_size, int32_t batch_size, int32_t padded_batch_size, - int64_t stride_page, float sm_scale, void* stream) -{ - return decode_split_kv_launch( - q, output, kv_data, k_offset_elems, v_offset_elems, - page_indices, page_indptr, last_page_len_d, - request_indices, kv_tile_indices, kv_chunk_size_ptr, - o_indptr, block_valid_mask, tmp_v, tmp_s, - num_qo_heads, num_kv_heads, head_dim, page_size, - batch_size, padded_batch_size, stride_page, sm_scale, - /*window_left=*/-1, stream); -} - -int batch_prefill_paged_cuda_hd256( - void* q, void* output, void* kv_data, - int64_t k_offset_elems, int64_t v_offset_elems, - int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, - int32_t* q_indptr, int32_t* request_indices, int32_t* qo_tile_indices, - int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, uint32_t* total_num_rows, - int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, - int32_t page_size, int32_t seq_len, int32_t batch_size, - int32_t padded_batch_size, int64_t stride_page, float sm_scale, - void* stream) -{ - return prefill_paged_launch( - q, output, kv_data, k_offset_elems, v_offset_elems, - page_indices, page_indptr, last_page_len_d, q_indptr, - request_indices, qo_tile_indices, kv_tile_indices, kv_chunk_size_ptr, - total_num_rows, num_qo_heads, num_kv_heads, head_dim, page_size, - seq_len, batch_size, padded_batch_size, stride_page, sm_scale, - /*cta_tile_q_override=*/0, /*window_left=*/-1, stream); -} - -int batch_prefill_paged_window_cuda_hd256( - void* q, void* output, void* kv_data, - int64_t k_offset_elems, int64_t v_offset_elems, - int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, - int32_t* q_indptr, int32_t* request_indices, int32_t* qo_tile_indices, - int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, uint32_t* total_num_rows, - int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, - int32_t page_size, int32_t seq_len, int32_t batch_size, - int32_t padded_batch_size, int64_t stride_page, float sm_scale, - int32_t window_left, void* stream) -{ - return prefill_paged_launch( - q, output, kv_data, k_offset_elems, v_offset_elems, - page_indices, page_indptr, last_page_len_d, q_indptr, - request_indices, qo_tile_indices, kv_tile_indices, kv_chunk_size_ptr, - total_num_rows, num_qo_heads, num_kv_heads, head_dim, page_size, - seq_len, batch_size, padded_batch_size, stride_page, sm_scale, - /*cta_tile_q_override=*/0, window_left, stream); -} - } // extern "C" diff --git a/pegainfer-kernels/src/ffi.rs b/pegainfer-kernels/src/ffi.rs index 7f981d8f8..c8d911488 100644 --- a/pegainfer-kernels/src/ffi.rs +++ b/pegainfer-kernels/src/ffi.rs @@ -21,6 +21,7 @@ mod k3_tilelang; #[cfg(feature = "kimi-k2")] mod kimi; mod lora; +#[cfg(feature = "qwen35")] mod qwen35; mod shared; #[cfg(feature = "moe")] @@ -40,5 +41,6 @@ pub use k3_tilelang::*; #[cfg(feature = "kimi-k2")] pub use kimi::*; pub use lora::*; +#[cfg(feature = "qwen35")] pub use qwen35::*; pub use shared::*; diff --git a/pegainfer-kernels/src/ffi/deepseek_v2_lite.rs b/pegainfer-kernels/src/ffi/deepseek_v2_lite.rs index 1fcc1fbd9..527c0d11a 100644 --- a/pegainfer-kernels/src/ffi/deepseek_v2_lite.rs +++ b/pegainfer-kernels/src/ffi/deepseek_v2_lite.rs @@ -6,6 +6,59 @@ use super::Half; // DeepSeek-V2-Lite private kernels (feature `deepseek-v2-lite`). // Sources: csrc/deepseek_v2_lite/*.cu. unsafe extern "C" { + pub fn dsv2_lite_route_logits_cuda( + logits: *const f32, + weights: *mut f32, + ids: *mut i32, + errors: *mut i32, + summary: *mut u64, + batch: i32, + layer_idx: i32, + stream: CUstream, + ) -> CUresult; + pub fn dsv2_lite_route_pointers_cuda( + ids: *const i32, + hidden: *const Half, + zero: *const Half, + w13: *const u64, + w2: *const u64, + gate: *mut Half, + act: *mut Half, + rows: *mut Half, + a13: *mut u64, + x13: *mut u64, + y13: *mut u64, + a2: *mut u64, + x2: *mut u64, + y2: *mut u64, + summary: *mut u64, + batch: i32, + first_expert: i32, + hidden_dim: i32, + intermediate: i32, + stream: CUstream, + ) -> CUresult; + pub fn dsv2_lite_route_reduce_cuda( + rows: *const Half, + ids: *const i32, + weights: *const f32, + errors: *const i32, + out: *mut f32, + batch: i32, + first_expert: i32, + hidden_dim: i32, + stream: CUstream, + ) -> CUresult; + pub fn dsv2_lite_pointer_gemm_cuda( + weights: *const u64, + inputs: *const u64, + outputs: *const u64, + m: i32, + k: i32, + routes: i32, + stream: CUstream, + ) -> i32; + pub fn dsv2_lite_router_logits_cuda( hidden: *const Half, gate_weight: *const Half, diff --git a/pegainfer-kernels/src/ffi/gemma4.rs b/pegainfer-kernels/src/ffi/gemma4.rs index 703283eee..c59d37b9a 100644 --- a/pegainfer-kernels/src/ffi/gemma4.rs +++ b/pegainfer-kernels/src/ffi/gemma4.rs @@ -98,7 +98,7 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; - /// fp8-KV twin of `batch_prefill_paged_window_hd256_cuda`: same launch + /// fp8-KV twin of `batch_prefill_paged_window_cuda_hd256`: same launch /// body over an e4m3 pool (scale 1.0), Q and output bf16. pub fn gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda( q: *const Half, @@ -156,3 +156,306 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; } + +// hd512 attention (global layers): csrc/gemma4/paged_attention_hd512.cu +unsafe extern "C" { + pub fn paged_attention_decode_split_kv_cuda_hd512( + q: *const Half, + output: *mut Half, + kv_data: *const Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indptr: *const i32, + last_page_len_d: *const i32, + request_indices: *const i32, + kv_tile_indices: *const i32, + kv_chunk_size_ptr: *const i32, + o_indptr: *const i32, + block_valid_mask: *const u8, + tmp_v: *mut Half, + tmp_s: *mut f32, + num_qo_heads: i32, + num_kv_heads: i32, + head_dim: i32, + page_size: i32, + batch_size: i32, + padded_batch_size: i32, + stride_page: i64, + sm_scale: f32, + stream: CUstream, + ) -> i32; + + pub fn batch_prefill_paged_cuda_hd512( + q: *const Half, + output: *mut Half, + kv_data: *const Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indptr: *const i32, + last_page_len_d: *const i32, + q_indptr: *const i32, + request_indices: *const i32, + qo_tile_indices: *const i32, + kv_tile_indices: *const i32, + kv_chunk_size_ptr: *const i32, + total_num_rows: *const u32, + num_qo_heads: i32, + num_kv_heads: i32, + head_dim: i32, + page_size: i32, + seq_len: i32, + batch_size: i32, + padded_batch_size: i32, + stride_page: i64, + sm_scale: f32, + stream: CUstream, + ) -> i32; +} + +// Windowed hd256 prefill (local layers): csrc/gemma4/paged_attention_window_hd256.cu. +// `window_left` is an inclusive distance: an N-token window passes N - 1, and +// -1 degrades to full attention. +unsafe extern "C" { + pub fn batch_prefill_paged_window_cuda_hd256( + q: *const Half, + output: *mut Half, + kv_data: *const Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indptr: *const i32, + last_page_len_d: *const i32, + q_indptr: *const i32, + request_indices: *const i32, + qo_tile_indices: *const i32, + kv_tile_indices: *const i32, + kv_chunk_size_ptr: *const i32, + total_num_rows: *const u32, + num_qo_heads: i32, + num_kv_heads: i32, + head_dim: i32, + page_size: i32, + seq_len: i32, + batch_size: i32, + padded_batch_size: i32, + stride_page: i64, + sm_scale: f32, + window_left: i32, + stream: CUstream, + ) -> i32; +} + +// hd256 plain-w QK-norm + RoPE prep (local layers): +// csrc/gemma4/prefill_attention_hd256_plain.cu. Contract and validation live +// on the Rust wrappers in ops::gemma4_attention; each entry returns 0 on +// success, -1 with a diagnostic on failure. +unsafe extern "C" { + // Q → contiguous q_batch_out; K (normed + rotated) and V (weightless-normed, + // never rotated) → straight into the paged KV pool at k_offset_elems / + // v_offset_elems. + pub fn qkv_norm_rope_paged_prefill_hd256_plain_cuda( + q_batch: *const Half, + k_batch: *const Half, + v_batch: *const Half, + q_stride: i32, + k_stride: i32, + v_stride: i32, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + page_origin: i32, + num_q_heads: i32, + num_kv_heads: i32, + seq_len: i32, + start_pos: i32, + cos_max_pos: i32, + rotary_dim: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; + + pub fn qkv_norm_rope_paged_decode_hd256_plain_cuda( + q_batch: *const Half, + k_batch: *const Half, + v_batch: *const Half, + q_stride: i32, + k_stride: i32, + v_stride: i32, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + page_indptr: *const i32, + page_origins: *const i32, + positions: *const i32, + num_q_heads: i32, + num_kv_heads: i32, + batch: i32, + cos_max_pos: i32, + rotary_dim: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; + + /// E4m3 KV twin. + pub fn qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda( + q_batch: *const Half, + k_batch: *const Half, + v_batch: *const Half, + q_stride: i32, + k_stride: i32, + v_stride: i32, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + page_origin: i32, + num_q_heads: i32, + num_kv_heads: i32, + seq_len: i32, + start_pos: i32, + cos_max_pos: i32, + rotary_dim: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; + + /// E4m3 KV twin. + pub fn qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda( + q_batch: *const Half, + k_batch: *const Half, + v_batch: *const Half, + q_stride: i32, + k_stride: i32, + v_stride: i32, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + page_indptr: *const i32, + page_origins: *const i32, + positions: *const i32, + num_q_heads: i32, + num_kv_heads: i32, + batch: i32, + cos_max_pos: i32, + rotary_dim: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; +} + +// hd512 QK-norm + partial RoPE prep (global layers): +// csrc/gemma4/prefill_attention_hd512.cu. Contract and validation live on +// the Rust wrappers in ops::gemma4_attention; both entries return 0 on +// success, -1 with a diagnostic on failure. +unsafe extern "C" { + // Prefill: Q → contiguous q_batch_out; K → straight into the paged KV + // pool at k_offset_elems (feeds batch_prefill_paged). + // V is the K=V fork — the weightless norm of the same raw K, sharing + // its denominator — written to v_offset_elems in the same pass. The + // pool row is described as bands: `row_width` columns per (token, kv + // head) and `fold_rotary`, zero for the split K|V format and the + // rotated columns K keeps in the folded one, where both offsets name + // the layer's single block. + pub fn qk_norm_partial_rope_paged_prefill_hd512_cuda( + q_batch: *const Half, + k_batch: *const Half, + q_stride: i32, + k_stride: i32, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + num_q_heads: i32, + num_kv_heads: i32, + seq_len: i32, + start_pos: i32, + cos_max_pos: i32, + row_width: i32, + fold_rotary: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; + + // Batched decode straight into the pool: per-token position and page + // window; V is the K=V fork written alongside K. + pub fn qk_norm_partial_rope_paged_decode_hd512_cuda( + q_batch: *const Half, + k_batch: *const Half, + q_stride: i32, + k_stride: i32, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + page_indptr: *const i32, + page_origins: *const i32, + positions: *const i32, + num_q_heads: i32, + num_kv_heads: i32, + batch: i32, + cos_max_pos: i32, + row_width: i32, + fold_rotary: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; +} diff --git a/pegainfer-kernels/src/ffi/qwen35.rs b/pegainfer-kernels/src/ffi/qwen35.rs index c949dc186..e03062325 100644 --- a/pegainfer-kernels/src/ffi/qwen35.rs +++ b/pegainfer-kernels/src/ffi/qwen35.rs @@ -1,13 +1,60 @@ -#[cfg(feature = "qwen35")] use cudarc::driver::sys::CUresult; use cudarc::driver::sys::CUstream; use super::Half; -// Qwen3.5-4B private kernels (hybrid linear + HD256 full attention). -// Sources: csrc/qwen35/*.cu. The paged HD256 attention entry points are shared -// with Gemma 4 and are declared in `shared.rs`. +// Qwen3.5 private kernels (hybrid linear + HD256 full attention). +// Sources: csrc/qwen35/*.cu. unsafe extern "C" { + pub fn paged_attention_decode_cuda_hd256( + q: *const Half, + output: *mut Half, + kv_data: *const Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indptr: *const i32, + last_page_len_d: *const i32, + request_indices: *const i32, + kv_tile_indices: *const i32, + kv_chunk_size_ptr: *const i32, + num_qo_heads: i32, + num_kv_heads: i32, + head_dim: i32, + page_size: i32, + batch_size: i32, + stride_page: i64, + sm_scale: f32, + stream: CUstream, + ) -> i32; + + pub fn batch_prefill_paged_cuda_hd256( + q: *const Half, + output: *mut Half, + kv_data: *const Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indptr: *const i32, + last_page_len_d: *const i32, + q_indptr: *const i32, + request_indices: *const i32, + qo_tile_indices: *const i32, + kv_tile_indices: *const i32, + kv_chunk_size_ptr: *const i32, + total_num_rows: *const u32, + num_qo_heads: i32, + num_kv_heads: i32, + head_dim: i32, + page_size: i32, + seq_len: i32, + batch_size: i32, + padded_batch_size: i32, + stride_page: i64, + sm_scale: f32, + stream: CUstream, + ) -> i32; + // Qwen3.5 full-attention prefill prep that writes K/V directly into paged KV. pub fn prefill_attention_hd256_prep_paged_cuda( q_full_batch: *const Half, @@ -115,10 +162,7 @@ unsafe extern "C" { ); } -// Chunk-wise GDR prefill kernels, Triton AOT-generated at build time. The -// `qwen35` feature is what pulls Python + Triton into the build; without it -// these symbols don't exist. -#[cfg(feature = "qwen35")] +// Chunk-wise GDR prefill kernels, Triton AOT-generated at build time. unsafe extern "C" { pub fn gated_delta_rule_prefill_chunk_prepare_cuda( qkv: *const Half, diff --git a/pegainfer-kernels/src/ffi/shared.rs b/pegainfer-kernels/src/ffi/shared.rs index 54b19908c..4347b4443 100644 --- a/pegainfer-kernels/src/ffi/shared.rs +++ b/pegainfer-kernels/src/ffi/shared.rs @@ -490,6 +490,19 @@ unsafe extern "C" { stream: CUstream, ); + // The same kernel with a sigmoid gate, for the lines whose + // `output_gate_type` selects it. + pub fn rms_norm_gated_sigmoid_cuda( + x: *const Half, + weight: *const f32, + gate: *const Half, + out: *mut Half, + num_heads: i32, + head_dim: i32, + eps: f32, + stream: CUstream, + ); + // ======================================================================== // Paged attention (FlashInfer) // ======================================================================== @@ -788,139 +801,6 @@ unsafe extern "C" { ) -> i32; } -// HEAD_DIM=256 paged attention. Qwen3.5-4B calls the full-attention pair; the -// windowed pair carries the sliding-window mask for Gemma 4's local layers. -// `window_left` is an inclusive distance: an N-token window passes N - 1, and -// -1 degrades to full attention. -unsafe extern "C" { - pub fn paged_attention_decode_cuda_hd256( - q: *const Half, - output: *mut Half, - kv_data: *const Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indptr: *const i32, - last_page_len_d: *const i32, - request_indices: *const i32, - kv_tile_indices: *const i32, - kv_chunk_size_ptr: *const i32, - num_qo_heads: i32, - num_kv_heads: i32, - head_dim: i32, - page_size: i32, - batch_size: i32, - stride_page: i64, - sm_scale: f32, - stream: CUstream, - ) -> i32; - - pub fn paged_attention_decode_window_cuda_hd256( - q: *const Half, - output: *mut Half, - kv_data: *const Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indptr: *const i32, - last_page_len_d: *const i32, - request_indices: *const i32, - kv_tile_indices: *const i32, - kv_chunk_size_ptr: *const i32, - num_qo_heads: i32, - num_kv_heads: i32, - head_dim: i32, - page_size: i32, - batch_size: i32, - stride_page: i64, - sm_scale: f32, - window_left: i32, - stream: CUstream, - ) -> i32; - - pub fn paged_attention_decode_split_kv_cuda_hd512( - q: *const Half, - output: *mut Half, - kv_data: *const Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indptr: *const i32, - last_page_len_d: *const i32, - request_indices: *const i32, - kv_tile_indices: *const i32, - kv_chunk_size_ptr: *const i32, - o_indptr: *const i32, - block_valid_mask: *const u8, - tmp_v: *mut Half, - tmp_s: *mut f32, - num_qo_heads: i32, - num_kv_heads: i32, - head_dim: i32, - page_size: i32, - batch_size: i32, - padded_batch_size: i32, - stride_page: i64, - sm_scale: f32, - stream: CUstream, - ) -> i32; - - pub fn batch_prefill_paged_cuda_hd256( - q: *const Half, - output: *mut Half, - kv_data: *const Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indptr: *const i32, - last_page_len_d: *const i32, - q_indptr: *const i32, - request_indices: *const i32, - qo_tile_indices: *const i32, - kv_tile_indices: *const i32, - kv_chunk_size_ptr: *const i32, - total_num_rows: *const u32, - num_qo_heads: i32, - num_kv_heads: i32, - head_dim: i32, - page_size: i32, - seq_len: i32, - batch_size: i32, - padded_batch_size: i32, - stride_page: i64, - sm_scale: f32, - stream: CUstream, - ) -> i32; - - pub fn batch_prefill_paged_window_cuda_hd256( - q: *const Half, - output: *mut Half, - kv_data: *const Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indptr: *const i32, - last_page_len_d: *const i32, - q_indptr: *const i32, - request_indices: *const i32, - qo_tile_indices: *const i32, - kv_tile_indices: *const i32, - kv_chunk_size_ptr: *const i32, - total_num_rows: *const u32, - num_qo_heads: i32, - num_kv_heads: i32, - head_dim: i32, - page_size: i32, - seq_len: i32, - batch_size: i32, - padded_batch_size: i32, - stride_page: i64, - sm_scale: f32, - window_left: i32, - stream: CUstream, - ) -> i32; -} - // Added during rebase onto main: generic dtype/scale helpers, batched argmax/top1, // rms-norm-round variant, gemm-per-token. unsafe extern "C" { @@ -1121,323 +1001,6 @@ unsafe extern "C" { } -// hd256 single prefill (qwen35 full-attention layers, Gemma 4 local layers): -// csrc/shared/paged_attention.cu. K/V are a contiguous HND cache — -// k[head, pos, dim] with max_seq_len rows per head — not token-major NHD. -unsafe extern "C" { - pub fn single_prefill_cuda_hd256( - q: *const Half, - output: *mut Half, - k_cache: *const Half, - v_cache: *const Half, - num_qo_heads: i32, - num_kv_heads: i32, - seq_len: i32, - kv_len: i32, - max_seq_len: i32, - sm_scale: f32, - stream: CUstream, - ) -> i32; -} - -// hd512 (Gemma 4 global layers): csrc/shared/paged_attention_hd512.cu -unsafe extern "C" { - pub fn single_prefill_cuda_hd512( - q: *const Half, - output: *mut Half, - k_cache: *const Half, - v_cache: *const Half, - num_qo_heads: i32, - num_kv_heads: i32, - seq_len: i32, - kv_len: i32, - max_seq_len: i32, - sm_scale: f32, - stream: CUstream, - ) -> i32; - - pub fn batch_prefill_paged_cuda_hd512( - q: *const Half, - output: *mut Half, - kv_data: *const Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indptr: *const i32, - last_page_len_d: *const i32, - q_indptr: *const i32, - request_indices: *const i32, - qo_tile_indices: *const i32, - kv_tile_indices: *const i32, - kv_chunk_size_ptr: *const i32, - total_num_rows: *const u32, - num_qo_heads: i32, - num_kv_heads: i32, - head_dim: i32, - page_size: i32, - seq_len: i32, - batch_size: i32, - padded_batch_size: i32, - stride_page: i64, - sm_scale: f32, - stream: CUstream, - ) -> i32; -} - -// hd256 plain-w QK-norm + RoPE prep (Gemma 4 local layers): -// csrc/shared/prefill_attention_hd256_plain.cu. Contract and validation live -// on the Rust wrapper in ops::attention; the entry returns 0 on success, -1 -// with a diagnostic on failure. -unsafe extern "C" { - // Oracle form: Q and K both land in contiguous buffers shaped like - // their inputs; the paged serving form is the qkv_ entry below. - pub fn qk_norm_rope_prefill_hd256_plain_cuda( - q_batch: *const Half, - k_batch: *const Half, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - q_batch_out: *mut Half, - k_batch_out: *mut Half, - num_q_heads: i32, - num_kv_heads: i32, - seq_len: i32, - start_pos: i32, - cos_max_pos: i32, - rotary_dim: i32, - rms_eps: f32, - stream: CUstream, - ) -> i32; - - // Paged serving form: Q → contiguous q_batch_out; K (normed + rotated) - // and V (weightless-normed, never rotated) → straight into the paged KV - // pool at k_offset_elems / v_offset_elems. - pub fn qkv_norm_rope_paged_prefill_hd256_plain_cuda( - q_batch: *const Half, - k_batch: *const Half, - v_batch: *const Half, - q_stride: i32, - k_stride: i32, - v_stride: i32, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - q_batch_out: *mut Half, - kv_data: *mut Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indices_len: i32, - page_origin: i32, - num_q_heads: i32, - num_kv_heads: i32, - seq_len: i32, - start_pos: i32, - cos_max_pos: i32, - rotary_dim: i32, - rms_eps: f32, - page_size: i32, - num_pages: i32, - stride_page: i64, - stream: CUstream, - ) -> i32; - - pub fn qkv_norm_rope_paged_decode_hd256_plain_cuda( - q_batch: *const Half, - k_batch: *const Half, - v_batch: *const Half, - q_stride: i32, - k_stride: i32, - v_stride: i32, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - q_batch_out: *mut Half, - kv_data: *mut Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indices_len: i32, - page_indptr: *const i32, - page_origins: *const i32, - positions: *const i32, - num_q_heads: i32, - num_kv_heads: i32, - batch: i32, - cos_max_pos: i32, - rotary_dim: i32, - rms_eps: f32, - page_size: i32, - num_pages: i32, - stride_page: i64, - stream: CUstream, - ) -> i32; - - /// E4m3 KV twin. - pub fn qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda( - q_batch: *const Half, - k_batch: *const Half, - v_batch: *const Half, - q_stride: i32, - k_stride: i32, - v_stride: i32, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - q_batch_out: *mut Half, - kv_data: *mut Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indices_len: i32, - page_origin: i32, - num_q_heads: i32, - num_kv_heads: i32, - seq_len: i32, - start_pos: i32, - cos_max_pos: i32, - rotary_dim: i32, - rms_eps: f32, - page_size: i32, - num_pages: i32, - stride_page: i64, - stream: CUstream, - ) -> i32; - - /// E4m3 KV twin. - pub fn qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda( - q_batch: *const Half, - k_batch: *const Half, - v_batch: *const Half, - q_stride: i32, - k_stride: i32, - v_stride: i32, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - q_batch_out: *mut Half, - kv_data: *mut Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indices_len: i32, - page_indptr: *const i32, - page_origins: *const i32, - positions: *const i32, - num_q_heads: i32, - num_kv_heads: i32, - batch: i32, - cos_max_pos: i32, - rotary_dim: i32, - rms_eps: f32, - page_size: i32, - num_pages: i32, - stride_page: i64, - stream: CUstream, - ) -> i32; -} - -// hd512 QK-norm + partial RoPE prep (Gemma 4 global layers): -// csrc/shared/prefill_attention_hd512.cu. Contract and validation live on -// the Rust wrappers in ops::attention; both entries return 0 on success, -// -1 with a diagnostic on failure. -unsafe extern "C" { - // Prefill: Q → contiguous q_batch_out; K → straight into the paged KV - // pool at k_offset_elems (feeds batch_prefill_paged, not single_prefill). - // V is the K=V fork — the weightless norm of the same raw K, sharing - // its denominator — written to v_offset_elems in the same pass. The - // pool row is described as bands: `row_width` columns per (token, kv - // head) and `fold_rotary`, zero for the split K|V format and the - // rotated columns K keeps in the folded one, where both offsets name - // the layer's single block. - pub fn qk_norm_partial_rope_paged_prefill_hd512_cuda( - q_batch: *const Half, - k_batch: *const Half, - q_stride: i32, - k_stride: i32, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - q_batch_out: *mut Half, - kv_data: *mut Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indices_len: i32, - num_q_heads: i32, - num_kv_heads: i32, - seq_len: i32, - start_pos: i32, - cos_max_pos: i32, - row_width: i32, - fold_rotary: i32, - rms_eps: f32, - page_size: i32, - num_pages: i32, - stride_page: i64, - stream: CUstream, - ) -> i32; - - // Batched decode straight into the pool: per-token position and page - // window; V is the K=V fork written alongside K. - pub fn qk_norm_partial_rope_paged_decode_hd512_cuda( - q_batch: *const Half, - k_batch: *const Half, - q_stride: i32, - k_stride: i32, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - q_batch_out: *mut Half, - kv_data: *mut Half, - k_offset_elems: i64, - v_offset_elems: i64, - page_indices: *const i32, - page_indices_len: i32, - page_indptr: *const i32, - page_origins: *const i32, - positions: *const i32, - num_q_heads: i32, - num_kv_heads: i32, - batch: i32, - cos_max_pos: i32, - row_width: i32, - fold_rotary: i32, - rms_eps: f32, - page_size: i32, - num_pages: i32, - stride_page: i64, - stream: CUstream, - ) -> i32; - - // Batched decode: Q → contiguous q_batch_out; K normalised + partially - // rotated in place (caller scatters it into the paged pool afterwards). - pub fn qk_norm_partial_rope_batched_decode_hd512_cuda( - q_batch: *const Half, - k_batch: *mut Half, - q_norm_weight: *const Half, - k_norm_weight: *const Half, - cos_cache: *const Half, - sin_cache: *const Half, - positions: *const i32, - cos_max_pos: i32, - q_batch_out: *mut Half, - num_q_heads: i32, - num_kv_heads: i32, - batch_size: i32, - rotary_dim: i32, - rms_eps: f32, - stream: CUstream, - ) -> i32; -} - unsafe extern "C" { pub fn pegainfer_kernels_last_error() -> *const std::os::raw::c_char; } diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index 49ef93117..104ee2821 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/pegainfer-kernels/src/ops.rs @@ -5,11 +5,25 @@ mod attention; mod deepep; #[cfg(feature = "deepseek-v2-lite")] mod deepseek_v2_lite; +#[cfg(feature = "deepseek-v2-lite")] +mod dsv2_routed_moe; +#[cfg(feature = "deepseek-v2-lite")] +pub use dsv2_routed_moe::DSV2_ROUTED_MOE_MAX_ROWS; +#[cfg(feature = "deepseek-v2-lite")] +pub use dsv2_routed_moe::Dsv2ExpertPointerTable; +#[cfg(feature = "deepseek-v2-lite")] +pub use dsv2_routed_moe::Dsv2RouteSummary; +#[cfg(feature = "deepseek-v2-lite")] +pub use dsv2_routed_moe::Dsv2RoutedMoeScratch; +#[cfg(feature = "deepseek-v2-lite")] +pub use dsv2_routed_moe::dsv2_lite_route_logits_into; mod elementwise; mod embedding; #[cfg(feature = "gemma4")] mod gemma4; #[cfg(feature = "gemma4")] +mod gemma4_attention; +#[cfg(feature = "gemma4")] mod gemma4_tilelang; #[cfg(feature = "gemma4")] mod gemma4_w4a16; @@ -26,33 +40,19 @@ mod lora; #[cfg(any(feature = "gemma4", feature = "kimi-k2"))] mod marlin_face; mod norm; +#[cfg(feature = "qwen35")] +mod qwen35; mod sampling; -pub use attention::Hd512DecodeMetadata; pub use attention::PrefillPagedPlan; pub use attention::SUPPORTED_GQA_GROUP_SIZES; -pub use attention::batch_prefill_paged_hd512_into; -pub use attention::batch_prefill_paged_window_hd256_into; pub use attention::dflash_qk_norm_rope_into; pub use attention::eagle3_rope_into; -pub use attention::paged_attention_batch_decode_hd256_into; pub use attention::paged_attention_batch_decode_into; -pub use attention::paged_attention_batch_decode_split_kv_hd512_into; pub use attention::paged_attention_batch_decode_split_kv_into; -pub use attention::paged_attention_batch_decode_via_prefill_hd256_into; -pub use attention::paged_attention_batch_decode_via_prefill_hd512_into; pub use attention::prefill_attention_paged_into; -pub use attention::qk_norm_partial_rope_batched_decode_hd256_into; -pub use attention::qk_norm_partial_rope_batched_decode_hd512_into; -pub use attention::qk_norm_partial_rope_paged_decode_hd512_into; -pub use attention::qk_norm_partial_rope_paged_prefill_hd512_into; pub use attention::qk_norm_rope_batch_decode_into; -pub use attention::qk_norm_rope_prefill_hd256_plain_into; -pub use attention::qkv_norm_rope_paged_decode_hd256_plain_into; -pub use attention::qkv_norm_rope_paged_prefill_hd256_plain_into; pub use attention::single_decode_nhd_into; -pub use attention::single_prefill_hd256_into; -pub use attention::single_prefill_hd512_into; pub use attention::single_prefill_nhd_causal_into; pub use attention::single_prefill_nhd_noncausal_into; #[cfg(feature = "moe")] @@ -129,6 +129,22 @@ pub use embedding::embedding_rows_into; #[cfg(feature = "gemma4")] pub use gemma4::*; #[cfg(feature = "gemma4")] +pub use gemma4_attention::Hd512DecodeMetadata; +#[cfg(feature = "gemma4")] +pub use gemma4_attention::batch_prefill_paged_hd512_into; +#[cfg(feature = "gemma4")] +pub use gemma4_attention::batch_prefill_paged_window_hd256_into; +#[cfg(feature = "gemma4")] +pub use gemma4_attention::paged_attention_batch_decode_split_kv_hd512_into; +#[cfg(feature = "gemma4")] +pub use gemma4_attention::qk_norm_partial_rope_paged_decode_hd512_into; +#[cfg(feature = "gemma4")] +pub use gemma4_attention::qk_norm_partial_rope_paged_prefill_hd512_into; +#[cfg(feature = "gemma4")] +pub use gemma4_attention::qkv_norm_rope_paged_decode_hd256_plain_into; +#[cfg(feature = "gemma4")] +pub use gemma4_attention::qkv_norm_rope_paged_prefill_hd256_plain_into; +#[cfg(feature = "gemma4")] pub use gemma4_tilelang::GlobalDecodeAttend; #[cfg(feature = "gemma4")] pub use gemma4_tilelang::gemma4_hd256_decode_window_into; @@ -195,6 +211,7 @@ pub use lora::LoraDecodeGroupedProjection; pub use lora::lora_decode_fused_delta_group3_into; pub use lora::lora_decode_fused_delta_into; pub use lora::pack_lora_b_rows_into; +pub use norm::GatedNormActivation; pub use norm::dual_rms_norm_add_batch_into; pub use norm::fused_add_rms_norm_batch_into; pub use norm::fused_add_rms_norm_into; @@ -211,6 +228,12 @@ pub use norm::rms_norm_gated_batch_into; pub use norm::rms_norm_into; pub use norm::rms_norm_offset_into; pub use norm::rms_norm_rows_into; +#[cfg(feature = "qwen35")] +pub use qwen35::paged_attention_batch_decode_hd256_into; +#[cfg(feature = "qwen35")] +pub use qwen35::paged_attention_batch_decode_via_prefill_hd256_into; +#[cfg(feature = "qwen35")] +pub use qwen35::qk_norm_partial_rope_batched_decode_hd256_into; pub use sampling::BatchSamplingRow; pub use sampling::BatchSamplingScratch; pub use sampling::argmax; @@ -254,3 +277,17 @@ pub fn ffi_exception_message(result: i32) -> String { format!(": {text}") } } + +/// `PEGAINFER_CUBLAS_STATUS_BASE` in `csrc/shared/ffi_guard.cuh`: an entry +/// point that calls cuBLAS returns it plus the `cublasStatus_t` on a cuBLAS +/// failure, and the `cudaError_t` otherwise. +pub const CUBLAS_STATUS_BASE: i32 = 100_000; + +/// Name a nonzero status from such an entry point, for an error message. +pub fn ffi_status_detail(status: i32) -> String { + if status >= CUBLAS_STATUS_BASE { + format!("cublas_status={}", status - CUBLAS_STATUS_BASE) + } else { + format!("cuda_status={status}") + } +} diff --git a/pegainfer-kernels/src/ops/attention.rs b/pegainfer-kernels/src/ops/attention.rs index 62cad6da0..1fa52e04f 100644 --- a/pegainfer-kernels/src/ops/attention.rs +++ b/pegainfer-kernels/src/ops/attention.rs @@ -9,7 +9,6 @@ use crate::paged_kv::KvFormat; use crate::paged_kv::KvStorage; use crate::paged_kv::PagedKvLayout; use crate::paged_kv::derive_strides; -use crate::tensor::Columns; use crate::tensor::DeviceContext; use crate::tensor::DeviceVec; use crate::tensor::HiddenStates; @@ -27,30 +26,30 @@ pub const SUPPORTED_GQA_GROUP_SIZES: &[usize] = &[1, 2, 3, 4, 8]; /// Built once per prefill call, shared across all layers. /// Supports both single-request (`new`) and multi-request (`new_batch`) prefill. pub struct PrefillPagedPlan { - page_indices_d: CudaSlice, - page_indptr_d: CudaSlice, - last_page_len_d: CudaSlice, - batch_indices_d: CudaSlice, + pub(super) page_indices_d: CudaSlice, + pub(super) page_indptr_d: CudaSlice, + pub(super) last_page_len_d: CudaSlice, + pub(super) batch_indices_d: CudaSlice, positions_d: CudaSlice, - q_indptr_d: CudaSlice, + pub(super) q_indptr_d: CudaSlice, /// The same boundaries the device copy holds, kept because a launcher /// that sizes its own grid needs them on the host and because the batch /// size is this vector's length, not a second thing to keep in step. /// Never empty: a plan with no requests still carries the leading zero. q_indptr_host: Vec, - request_indices_d: CudaSlice, - qo_tile_indices_d: CudaSlice, - kv_tile_indices_d: CudaSlice, - kv_chunk_size_d: CudaSlice, - total_num_rows_d: CudaSlice, - num_tiles: i32, - total_tokens: usize, + pub(super) request_indices_d: CudaSlice, + pub(super) qo_tile_indices_d: CudaSlice, + pub(super) kv_tile_indices_d: CudaSlice, + pub(super) kv_chunk_size_d: CudaSlice, + pub(super) total_num_rows_d: CudaSlice, + pub(super) num_tiles: i32, + pub(super) total_tokens: usize, cta_tile_q: i32, - head_dim: usize, - num_qo_heads: usize, - num_kv_heads: usize, - max_page_index: i32, - max_last_page_len: i32, + pub(super) head_dim: usize, + pub(super) num_qo_heads: usize, + pub(super) num_kv_heads: usize, + pub(super) max_page_index: i32, + pub(super) max_last_page_len: i32, } /// GQA group size is an integer division by `num_kv_heads`, guarded by neither @@ -140,7 +139,7 @@ impl PrefillPagedPlan { pub fn num_tiles(&self) -> i32 { self.num_tiles } - fn cta_tile_q(&self) -> i32 { + pub(super) fn cta_tile_q(&self) -> i32 { self.cta_tile_q } @@ -907,39 +906,6 @@ pub(crate) fn checked_row_offset( Ok(bytes as u64) } -/// Byte offset of row `offset`'s first column of the band `c`, after checking -/// the row window and that the band lies within the row. -pub(crate) fn checked_band_offset( - c: &Columns<'_>, - offset: usize, - span: usize, - name: &str, -) -> Result { - anyhow::ensure!( - c.col - .checked_add(c.width) - .is_some_and(|end| end <= c.states.hidden_dim), - "{name} columns [{}..+{}) exceed the row's {}", - c.col, - c.width, - c.states.hidden_dim - ); - let rows = checked_row_offset(c.states, offset, span, name)?; - Ok(rows + (c.col * std::mem::size_of::()) as u64) -} - -/// `DeviceVec` fields are public too: reject a logical `len` past the backing -/// allocation before a kernel indexes through it. -fn ensure_vec_backed(v: &DeviceVec, name: &str) -> Result<()> { - anyhow::ensure!( - v.data.len() >= v.len, - "{name} backing len {} < len {}", - v.data.len(), - v.len - ); - Ok(()) -} - /// Plain RoPE (no QK-norm) for one EAGLE-3 draft step, because EAGLE-3 has no per-head q/k norm /// we implement a new kernel for EAGLE-3 #[allow(clippy::too_many_arguments)] @@ -1213,59 +1179,6 @@ pub fn single_prefill_nhd_causal_into( Ok(()) } -/// Batched QK RMSNorm + partial RoPE for Qwen3.5 HD256 decode. -/// -/// Reads Q from interleaved `q_full` ([q, gate] per head), writes prepared Q into `q`, -/// and normalizes/applies partial RoPE to `k` in-place using per-request positions. -#[allow(clippy::too_many_arguments)] -pub fn qk_norm_partial_rope_batched_decode_hd256_into( - ctx: &DeviceContext, - q_full: &HiddenStates, - q: &mut HiddenStates, - k: &mut HiddenStates, - q_norm_weight: &DeviceVec, - k_norm_weight: &DeviceVec, - cos_cache: &DeviceVec, - sin_cache: &DeviceVec, - positions_d: &CudaSlice, - num_q_heads: usize, - num_kv_heads: usize, - rotary_dim: usize, - rms_eps: f32, -) { - let batch_size = q.seq_len; - debug_assert_eq!(q_full.seq_len, batch_size); - debug_assert_eq!(k.seq_len, batch_size); - - let (qf_ptr, _gqf) = q_full.data.device_ptr(&ctx.stream); - let (q_ptr, _gq) = q.data.device_ptr_mut(&ctx.stream); - let (k_ptr, _gk) = k.data.device_ptr_mut(&ctx.stream); - let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); - let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); - let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); - let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); - let (pos_ptr, _gp) = positions_d.device_ptr(&ctx.stream); - - unsafe { - ffi::qk_norm_partial_rope_batched_decode_hd256_cuda( - qf_ptr as *const ffi::Half, - k_ptr as *mut ffi::Half, - qn_ptr as *const ffi::Half, - kn_ptr as *const ffi::Half, - cos_ptr as *const ffi::Half, - sin_ptr as *const ffi::Half, - pos_ptr as *const i32, - q_ptr as *mut ffi::Half, - num_q_heads as i32, - num_kv_heads as i32, - batch_size as i32, - rotary_dim as i32, - rms_eps, - crate::tensor::active_cu_stream(ctx), - ); - } -} - /// Batched paged attention decode: append K/V + FlashInfer BatchDecode for batch_size >= 1. /// /// Q: HiddenStates [q_dim, batch_size], output: HiddenStates [q_dim, batch_size]. @@ -1564,2730 +1477,140 @@ pub fn paged_attention_batch_decode_split_kv_into( Ok(()) } -#[allow(clippy::too_many_arguments)] -fn scatter_decode_kv_into_paged( - ctx: &DeviceContext, - k: &HiddenStates, - v: &HiddenStates, - kv_buffer: &CudaSlice, - layout: &PagedKvLayout, - layer: usize, - page_indices_d: &CudaSlice, - page_indptr_d: &CudaSlice, - last_page_len_d: &CudaSlice, - positions_d: &CudaSlice, - request_indices_d: &CudaSlice, - batch_size: usize, - op_name: &str, -) -> Result<()> { - let num_kv_heads = layout.num_kv_heads; - let head_dim = layout.head_dim; - let page_size = layout.page_size; - - let PagedGeometry { - k_offset_elems: k_offset, - v_offset_elems: v_offset, - stride_page, - .. - } = checked_paged_geometry( - op_name, - layout, - kv_buffer.len(), - layer, - head_dim, - num_kv_heads, - false, - )?; - - let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); - let (k_ptr, _gk) = k.data.device_ptr(&ctx.stream); - let (v_ptr, _gv) = v.data.device_ptr(&ctx.stream); - let (pi_ptr, _gpi) = page_indices_d.device_ptr(&ctx.stream); - let (pip_ptr, _gpip) = page_indptr_d.device_ptr(&ctx.stream); - let (lpl_ptr, _glpl) = last_page_len_d.device_ptr(&ctx.stream); - let (pos_ptr, _gpos) = positions_d.device_ptr(&ctx.stream); - let (ri_ptr, _gri) = request_indices_d.device_ptr(&ctx.stream); - - let src_stride_n = (num_kv_heads * head_dim) as i64; - let src_stride_h = head_dim as i64; - let result = unsafe { - ffi::paged_kv_scatter_cuda( - buf_ptr as *const ffi::Half, - k_offset, - v_offset, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - k_ptr as *const ffi::Half, - v_ptr as *const ffi::Half, - ri_ptr as *const i32, - pos_ptr as *const i32, - batch_size as i32, - num_kv_heads as i32, - head_dim as i32, - page_size as i32, - stride_page, - src_stride_n, - src_stride_h, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "paged_kv_scatter_cuda ({op_name}) failed for layer {layer}, bs={batch_size}, \ - kv_heads={num_kv_heads}, head_dim={head_dim}, page_size={page_size}: {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -pub fn paged_attention_batch_decode_hd256_into( - ctx: &DeviceContext, - q: &HiddenStates, - k: &HiddenStates, - v: &HiddenStates, - kv_buffer: &CudaSlice, - layout: &PagedKvLayout, - layer: usize, - page_indices_d: &CudaSlice, - page_indptr_d: &CudaSlice, - last_page_len_d: &CudaSlice, - positions_d: &CudaSlice, - request_indices_d: &CudaSlice, - kv_tile_indices_d: &CudaSlice, - kv_chunk_size_d: &CudaSlice, - output: &mut HiddenStates, - num_qo_heads: usize, - batch_size: usize, -) -> Result<()> { - let num_kv_heads = layout.num_kv_heads; - let head_dim = layout.head_dim; - debug_assert_eq!(head_dim, 256); - let page_size = layout.page_size; - - let PagedGeometry { - k_offset_elems: k_offset, - v_offset_elems: v_offset, - stride_page, - .. - } = checked_paged_geometry( - "batch hd256 decode", - layout, - kv_buffer.len(), - layer, - head_dim, - num_kv_heads, - false, - )?; - - let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let (pi_ptr, _gpi) = page_indices_d.device_ptr(&ctx.stream); - let (pip_ptr, _gpip) = page_indptr_d.device_ptr(&ctx.stream); - let (lpl_ptr, _glpl) = last_page_len_d.device_ptr(&ctx.stream); - let (ri_ptr, _gri) = request_indices_d.device_ptr(&ctx.stream); - let (kti_ptr, _gkti) = kv_tile_indices_d.device_ptr(&ctx.stream); - let (kcs_ptr, _gkcs) = kv_chunk_size_d.device_ptr(&ctx.stream); - - let stream = crate::tensor::active_cu_stream(ctx); - - scatter_decode_kv_into_paged( - ctx, - k, - v, - kv_buffer, - layout, - layer, - page_indices_d, - page_indptr_d, - last_page_len_d, - positions_d, - request_indices_d, - batch_size, - "batch hd256 decode", - )?; - - let sm_scale = 1.0f32 / (head_dim as f32).sqrt(); - let result = unsafe { - ffi::paged_attention_decode_cuda_hd256( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const ffi::Half, - k_offset, - v_offset, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - ri_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - num_qo_heads as i32, - num_kv_heads as i32, - head_dim as i32, - page_size as i32, - batch_size as i32, - stride_page, - sm_scale, - stream, - ) - }; - if result != 0 { - anyhow::bail!( - "paged_attention_decode_cuda_hd256 (batch) failed with error {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - - Ok(()) +/// Everything a paged-pool launch (prep write or attention read) derives +/// from a caller-supplied `PagedKvLayout`, re-derived with overflow checks +/// and narrowed through +/// checked conversions: the layout's fields are public, so a contradictory +/// or wrapping layout must fail here, before anything reaches an unsafe +/// launch. +pub(super) struct PagedGeometry { + pub(super) k_offset_elems: i64, + pub(super) v_offset_elems: i64, + #[cfg_attr(not(feature = "gemma4"), allow(dead_code))] + pub(super) page_size: i32, + #[cfg_attr(not(feature = "gemma4"), allow(dead_code))] + pub(super) num_pages: i32, + pub(super) stride_page: i64, } -#[allow(clippy::too_many_arguments)] -pub fn paged_attention_batch_decode_via_prefill_hd256_into( - ctx: &DeviceContext, - q: &HiddenStates, - k: &HiddenStates, - v: &HiddenStates, - kv_buffer: &CudaSlice, +pub(super) fn checked_paged_geometry( + what: &str, layout: &PagedKvLayout, + pool_len: usize, layer: usize, - plan: &PrefillPagedPlan, - positions_d: &CudaSlice, - output: &mut HiddenStates, - num_qo_heads: usize, - batch_size: usize, -) -> Result<()> { - let num_kv_heads = layout.num_kv_heads; - let head_dim = layout.head_dim; - debug_assert_eq!(head_dim, 256); + head_dim: usize, + num_kv_heads: usize, + fp8_capable: bool, +) -> Result { + // This is the split geometry: a K block and a V block per layer, which + // is what FlashInfer's paged view and the prep scatter address through + // `k_offset`/`v_offset`. Neither can express another format, so one is + // refused here rather than read at the wrong rows. anyhow::ensure!( - batch_size == plan.total_tokens && batch_size == plan.batch_size() as usize, - "decode-via-prefill plan shape mismatch: bs={batch_size}, total_tokens={}, plan_batch={}", - plan.total_tokens, - plan.batch_size() + layout.format == KvFormat::Split, + "{what} reads the split K|V format; the layout is {:?}", + layout.format ); - - scatter_decode_kv_into_paged( - ctx, - k, - v, - kv_buffer, - layout, - layer, - &plan.page_indices_d, - &plan.page_indptr_d, - &plan.last_page_len_d, - positions_d, - &plan.batch_indices_d, - batch_size, - "batch hd256 decode via prefill", - )?; - - let PagedGeometry { - k_offset_elems: k_offset, - v_offset_elems: v_offset, - stride_page, - .. - } = checked_paged_geometry( - "batch hd256 decode via prefill", - layout, - kv_buffer.len(), - layer, - head_dim, - num_kv_heads, - false, - )?; - let sm_scale = 1.0f32 / (head_dim as f32).sqrt(); - - let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let (pi_ptr, _gpi) = plan.page_indices_d.device_ptr(&ctx.stream); - let (pip_ptr, _gpip) = plan.page_indptr_d.device_ptr(&ctx.stream); - let (lpl_ptr, _glpl) = plan.last_page_len_d.device_ptr(&ctx.stream); - let (qi_ptr, _gqi) = plan.q_indptr_d.device_ptr(&ctx.stream); - let (ri_ptr, _gri) = plan.request_indices_d.device_ptr(&ctx.stream); - let (qti_ptr, _gqti) = plan.qo_tile_indices_d.device_ptr(&ctx.stream); - let (kti_ptr, _gkti) = plan.kv_tile_indices_d.device_ptr(&ctx.stream); - let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); - let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); - - let result = unsafe { - ffi::batch_prefill_paged_cuda_hd256( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const ffi::Half, - k_offset, - v_offset, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - qi_ptr as *const i32, - ri_ptr as *const i32, - qti_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - tnr_ptr as *const u32, - num_qo_heads as i32, - num_kv_heads as i32, - head_dim as i32, - layout.page_size as i32, - batch_size as i32, - plan.batch_size(), - plan.num_tiles, - stride_page, - sm_scale, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "batch_prefill_paged_cuda_hd256 (decode via prefill) failed for layer {layer}, \ - bs={batch_size}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}: {result}{}", - plan.num_tiles, - crate::ops::ffi_exception_message(result) - ); - } - - Ok(()) -} - -// ============================================================================ -// hd512 (Gemma 4 global layers) -// ============================================================================ - -/// Decode-metadata aggregate for the hd512 split-KV decode path, validated -/// at use: the wrapper calls [`Self::validate`] with the batch size it -/// derives. -/// -/// `page_indptr` needs at least `batch_size + 1` elements and `last_page_len` -/// at least `batch_size`; `request_indices` and `kv_tile_indices` are -/// per-split-slot arrays sized by the caller's padded slot count, and -/// `kv_chunk_size` holds the one chunk-size entry the kernel reads. -/// `page_indices` must be non-empty (its per-request reach is device-side -/// data and cannot be verified here). -/// -/// The pages referenced by `page_indices` must lie within `kv_buffer`; the -/// wrapper cannot verify device-side contents — caller contract. -pub struct Hd512DecodeMetadata<'a> { - page_indices: &'a CudaSlice, - page_indptr: &'a CudaSlice, - last_page_len: &'a CudaSlice, - request_indices: &'a CudaSlice, - kv_tile_indices: &'a CudaSlice, - kv_chunk_size: &'a CudaSlice, - /// The chunk the tile indices count in, on the host as well: the device - /// slot holds the same number for the incumbent's kernel, and a launcher - /// that sizes its own grid from the tiles needs it where it can read it. - chunk_tokens: usize, -} - -impl<'a> Hd512DecodeMetadata<'a> { - pub fn new( - page_indices: &'a CudaSlice, - page_indptr: &'a CudaSlice, - last_page_len: &'a CudaSlice, - request_indices: &'a CudaSlice, - kv_tile_indices: &'a CudaSlice, - kv_chunk_size: &'a CudaSlice, - chunk_tokens: usize, - ) -> Self { - Self { - page_indices, - page_indptr, - last_page_len, - request_indices, - kv_tile_indices, - kv_chunk_size, - chunk_tokens, - } - } - - pub fn page_indices(&self) -> &'a CudaSlice { - self.page_indices - } - - pub fn page_indptr(&self) -> &'a CudaSlice { - self.page_indptr - } - - pub fn last_page_len(&self) -> &'a CudaSlice { - self.last_page_len - } - - pub fn request_indices(&self) -> &'a CudaSlice { - self.request_indices - } - - pub fn kv_tile_indices(&self) -> &'a CudaSlice { - self.kv_tile_indices - } - - pub fn chunk_tokens(&self) -> usize { - self.chunk_tokens - } - - pub fn validate(&self, batch_size: usize) -> anyhow::Result<()> { - anyhow::ensure!( - !self.page_indices.is_empty(), - "Hd512DecodeMetadata: page_indices must be non-empty" - ); - anyhow::ensure!( - self.page_indptr.len() > batch_size, - "Hd512DecodeMetadata: page_indptr length {} must be >= batch_size + 1 = {}", - self.page_indptr.len(), - batch_size + 1 - ); - for (name, len) in [ - ("last_page_len", self.last_page_len.len()), - ("request_indices", self.request_indices.len()), - ("kv_tile_indices", self.kv_tile_indices.len()), - ] { - anyhow::ensure!( - len >= batch_size, - "Hd512DecodeMetadata: {name} length {len} must be >= batch_size {batch_size}", - ); - } - // The kernel reads one chunk size for the whole launch. - anyhow::ensure!( - !self.kv_chunk_size.is_empty(), - "Hd512DecodeMetadata: kv_chunk_size must hold its one entry" - ); - anyhow::ensure!( - self.chunk_tokens > 0, - "Hd512DecodeMetadata: chunk_tokens must be positive" - ); - Ok(()) - } -} - -#[allow(clippy::too_many_arguments)] -pub fn paged_attention_batch_decode_split_kv_hd512_into( - ctx: &DeviceContext, - q: &HiddenStates, - row_offset: usize, - kv_buffer: &CudaSlice, - layout: &PagedKvLayout, - layer: usize, - meta: &Hd512DecodeMetadata, - split_o_indptr_d: &CudaSlice, - split_valid_mask_d: &CudaSlice, - split_tmp_v: &mut CudaSlice, - split_tmp_s: &mut CudaSlice, - split_padded_slots: usize, - output: &mut HiddenStates, - num_qo_heads: usize, - sm_scale: f32, -) -> Result<()> { - anyhow::ensure!( - sm_scale.is_finite(), - "paged_attention_batch_decode_hd512 sm_scale {sm_scale} must be finite" - ); - let num_kv_heads = layout.num_kv_heads; - let geometry = checked_paged_geometry( - "hd512 split decode", - layout, - kv_buffer.len(), - layer, - 512, - num_kv_heads, - false, - )?; - anyhow::ensure!( - row_offset < q.seq_len, - "hd512 split decode row_offset {row_offset} leaves no rows of {}", - q.seq_len - ); - let batch_size = q.seq_len - row_offset; - meta.validate(batch_size)?; - anyhow::ensure!( - output.seq_len == q.seq_len, - "hd512 decode output.seq_len {} != q.seq_len {}", - output.seq_len, - q.seq_len - ); - let qo_dim = num_qo_heads.checked_mul(512).ok_or_else(|| { - anyhow::anyhow!("hd512 split decode num_qo_heads {num_qo_heads} * 512 overflows") - })?; - anyhow::ensure!( - q.hidden_dim == qo_dim, - "hd512 decode q.hidden_dim {} != num_qo_heads {num_qo_heads} * 512", - q.hidden_dim - ); - anyhow::ensure!( - output.hidden_dim == qo_dim, - "hd512 decode output.hidden_dim {} != num_qo_heads {num_qo_heads} * 512", - output.hidden_dim - ); - anyhow::ensure!( - split_padded_slots >= batch_size, - "hd512 split decode padded_slots {split_padded_slots} < batch {batch_size}" - ); - anyhow::ensure!( - meta.request_indices.len() >= split_padded_slots - && meta.kv_tile_indices.len() >= split_padded_slots - && split_valid_mask_d.len() >= split_padded_slots, - "hd512 split decode plan arrays shorter than padded_slots {split_padded_slots}" - ); - anyhow::ensure!( - split_o_indptr_d.len() > batch_size, - "hd512 split decode o_indptr len {} < batch {batch_size} + 1", - split_o_indptr_d.len() - ); - let tmp_v_need = split_padded_slots.checked_mul(qo_dim).ok_or_else(|| { - anyhow::anyhow!("hd512 split decode padded_slots {split_padded_slots} * {qo_dim} overflows") - })?; - let tmp_s_need = split_padded_slots - .checked_mul(num_qo_heads) - .ok_or_else(|| { - anyhow::anyhow!( - "hd512 split decode padded_slots {split_padded_slots} * {num_qo_heads} overflows" - ) - })?; + // `pool_len` counts the pool's bf16 backing slots; an e4m3 pool packs two + // elements per slot. Only wrappers with an fp8 kernel twin may see one. anyhow::ensure!( - split_tmp_v.len() >= tmp_v_need && split_tmp_s.len() >= tmp_s_need, - "hd512 split decode workspace shorter than padded_slots {split_padded_slots}" + fp8_capable || layout.storage == KvStorage::Bf16, + "{what} has no fp8 KV path; the layout carries {}-byte elements", + layout.storage.elem_bytes() ); - let q_elems = q.checked_extent("batch_decode_hd512 q")?; - let out_elems = output.checked_extent("batch_decode_hd512 output")?; - crate::ops::checked_i32(q_elems, "hd512 split decode q extent")?; - crate::ops::checked_i32(out_elems, "hd512 split decode output extent")?; - let batch_i32 = crate::ops::checked_i32(batch_size, "hd512 split decode batch")?; - let padded_i32 = - crate::ops::checked_i32(split_padded_slots, "hd512 split decode padded slots")?; - let qo_heads_i32 = crate::ops::checked_i32(num_qo_heads, "hd512 split decode qo heads")?; - let kv_heads_i32 = crate::ops::checked_i32(num_kv_heads, "hd512 split decode kv heads")?; - - let q_row_bytes = checked_row_offset(q, row_offset, batch_size, "hd512 split decode q")?; - let out_row_bytes = - checked_row_offset(output, row_offset, batch_size, "hd512 split decode output")?; - let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let q_ptr = q_ptr + q_row_bytes; - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let out_ptr = out_ptr + out_row_bytes; - let (pi_ptr, _gpi) = meta.page_indices.device_ptr(&ctx.stream); - let (pip_ptr, _gpip) = meta.page_indptr.device_ptr(&ctx.stream); - let (lpl_ptr, _glpl) = meta.last_page_len.device_ptr(&ctx.stream); - let (ri_ptr, _gri) = meta.request_indices.device_ptr(&ctx.stream); - let (kti_ptr, _gkti) = meta.kv_tile_indices.device_ptr(&ctx.stream); - let (kcs_ptr, _gkcs) = meta.kv_chunk_size.device_ptr(&ctx.stream); - let (soi_ptr, _gsoi) = split_o_indptr_d.device_ptr(&ctx.stream); - let (svm_ptr, _gsvm) = split_valid_mask_d.device_ptr(&ctx.stream); - let (stv_ptr, _gstv) = split_tmp_v.device_ptr_mut(&ctx.stream); - let (sts_ptr, _gsts) = split_tmp_s.device_ptr_mut(&ctx.stream); - - let stream = crate::tensor::active_cu_stream(ctx); - - // No KV scatter here: the serving prep kernel has already written this - // step's normed+roped K and its V fork; this only attends over the - // resident pages. - let result = unsafe { - ffi::paged_attention_decode_split_kv_cuda_hd512( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const ffi::Half, - geometry.k_offset_elems, - geometry.v_offset_elems, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - ri_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - soi_ptr as *const i32, - svm_ptr as *const u8, - stv_ptr as *mut ffi::Half, - sts_ptr as *mut f32, - qo_heads_i32, - kv_heads_i32, - 512, - geometry.page_size, - batch_i32, - padded_i32, - geometry.stride_page, - sm_scale, - stream, - ) + let pool_len = match layout.storage { + KvStorage::E4m3 => pool_len + .checked_mul(2) + .ok_or_else(|| anyhow::anyhow!("{what} fp8 pool element count overflows"))?, + KvStorage::Bf16 => pool_len, }; - if result != 0 { - anyhow::bail!( - "paged_attention_decode_split_kv_cuda_hd512 (batch) failed with error {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -pub fn paged_attention_batch_decode_via_prefill_hd512_into( - ctx: &DeviceContext, - q: &HiddenStates, - k: &HiddenStates, - v: &HiddenStates, - kv_buffer: &CudaSlice, - layout: &PagedKvLayout, - layer: usize, - plan: &PrefillPagedPlan, - positions_d: &CudaSlice, - output: &mut HiddenStates, - num_qo_heads: usize, - sm_scale: f32, -) -> Result<()> { anyhow::ensure!( - sm_scale.is_finite(), - "paged_attention_batch_decode_via_prefill_hd512 sm_scale {sm_scale} must be finite" + layout.head_dim == head_dim, + "{what} layout.head_dim {} != {head_dim}", + layout.head_dim ); - let num_kv_heads = layout.num_kv_heads; - let head_dim = layout.head_dim; anyhow::ensure!( - head_dim == 512, - "hd512 decode expects head_dim 512, got {}", - head_dim + layout.num_kv_heads == num_kv_heads, + "{what} layout.num_kv_heads {} != num_kv_heads {num_kv_heads}", + layout.num_kv_heads ); + // A zero head_dim makes every stride zero, and `pool_len / page_stride` + // below would panic on a layout that is otherwise self-consistent. anyhow::ensure!( - q.hidden_dim == num_qo_heads * 512, - "hd512 decode-via-prefill q.hidden_dim {} != num_qo_heads {} * 512", - q.hidden_dim, - num_qo_heads + layout.page_size > 0 + && layout.num_kv_heads > 0 + && layout.num_layers > 0 + && layout.head_dim > 0, + "{what} layout dims must be positive: page_size {} num_kv_heads {} num_layers {} \ + head_dim {}", + layout.page_size, + layout.num_kv_heads, + layout.num_layers, + layout.head_dim ); + // The layout's fields are public, so its strides are re-derived from its + // primitives and held to what it carries -- through the one derivation + // the layout itself uses, not a second statement of the arithmetic. + let (kv_block_len, layer_stride, page_stride) = derive_strides( + layout.page_size, + layout.num_kv_heads, + layout.head_dim, + layout.num_layers, + layout.format, + ) + .ok_or_else(|| { + anyhow::anyhow!( + "{what} strides overflow: page_size {} num_kv_heads {} head_dim {} num_layers {}", + layout.page_size, + layout.num_kv_heads, + layout.head_dim, + layout.num_layers + ) + })?; anyhow::ensure!( - output.hidden_dim == num_qo_heads * 512, - "hd512 decode-via-prefill output.hidden_dim {} != num_qo_heads {} * 512", - output.hidden_dim, - num_qo_heads + layout.kv_block_len == kv_block_len, + "{what} layout.kv_block_len {} != page_size * num_kv_heads * head_dim {kv_block_len}", + layout.kv_block_len ); anyhow::ensure!( - q.seq_len >= plan.total_tokens, - "hd512 decode-via-prefill q.seq_len {} < plan.total_tokens {}", - q.seq_len, - plan.total_tokens + layout.layer_stride == layer_stride, + "{what} layout.layer_stride {} != the format's {layer_stride}", + layout.layer_stride ); anyhow::ensure!( - output.seq_len >= plan.total_tokens, - "hd512 decode-via-prefill output.seq_len {} < plan.total_tokens {}", - output.seq_len, - plan.total_tokens + layout.page_stride == page_stride, + "{what} layout.page_stride {} != num_layers {} * layer_stride {layer_stride}", + layout.page_stride, + layout.num_layers ); anyhow::ensure!( layer < layout.num_layers, - "hd512 decode-via-prefill layer {layer} >= layout.num_layers {}", + "{what} layer {layer} >= layout.num_layers {}", layout.num_layers ); - q.checked_extent("decode_via_prefill_hd512 q")?; - output.checked_extent("decode_via_prefill_hd512 output")?; anyhow::ensure!( - positions_d.len() >= plan.batch_size() as usize, - "hd512 decode-via-prefill positions_d length {} must be >= plan.batch_size {}", - positions_d.len(), - plan.batch_size(), - ); - let batch_size = plan.batch_size() as usize; - anyhow::ensure!( - k.hidden_dim == num_kv_heads * 512, - "hd512 decode-via-prefill k.hidden_dim {} != num_kv_heads {} * 512", - k.hidden_dim, - num_kv_heads - ); - anyhow::ensure!( - v.hidden_dim == num_kv_heads * 512, - "hd512 decode-via-prefill v.hidden_dim {} != num_kv_heads {} * 512", - v.hidden_dim, - num_kv_heads - ); - anyhow::ensure!( - k.seq_len >= batch_size, - "hd512 decode-via-prefill k.seq_len {} < batch_size {}", - k.seq_len, - batch_size - ); - anyhow::ensure!( - v.seq_len >= batch_size, - "hd512 decode-via-prefill v.seq_len {} < batch_size {}", - v.seq_len, - batch_size - ); - k.checked_extent("decode_via_prefill_hd512 k")?; - v.checked_extent("decode_via_prefill_hd512 v")?; - anyhow::ensure!( - plan.head_dim == 512, - "plan built for head_dim {}, expected 512", - plan.head_dim + pool_len.is_multiple_of(page_stride), + "{what} kv_pool.len {pool_len} is not a multiple of layout.page_stride {page_stride}" ); + let num_pages = pool_len / page_stride; anyhow::ensure!( - plan.num_kv_heads == num_kv_heads, - "plan built for num_kv_heads {}, expected {}", - plan.num_kv_heads, - num_kv_heads + num_pages >= 1, + "{what} kv_pool.len {pool_len} holds no whole page (page_stride {page_stride})" ); - anyhow::ensure!( - plan.num_qo_heads == num_qo_heads, - "plan built for num_qo_heads {}, expected {}", - plan.num_qo_heads, - num_qo_heads - ); - anyhow::ensure!( - plan.total_tokens == plan.batch_size() as usize, - "decode-via-prefill plan shape mismatch: total_tokens={}, plan_batch={}", - plan.total_tokens, - plan.batch_size() - ); - - // Plan tiles laid out for a cta_tile_q override would be misread by the - // kernel's automatic derivation. - let kernel_cta_tile_q = unsafe { - ffi::batch_prefill_cta_tile_q_with_override( - plan.total_tokens as i32, - plan.num_qo_heads as i32, - plan.num_kv_heads as i32, - plan.head_dim as i32, - 0, // auto-detect, matching the hd512 kernel - ) - }; - anyhow::ensure!( - kernel_cta_tile_q == plan.cta_tile_q(), - "hd512 decode-via-prefill cta_tile_q: plan recorded {}, kernel derives {}; \ - overridden plans are not supported at hd512", - plan.cta_tile_q(), - kernel_cta_tile_q - ); - - scatter_decode_kv_into_paged( - ctx, - k, - v, - kv_buffer, - layout, - layer, - &plan.page_indices_d, - &plan.page_indptr_d, - &plan.last_page_len_d, - positions_d, - &plan.batch_indices_d, - plan.batch_size() as usize, - "batch hd512 decode via prefill", - )?; - - let PagedGeometry { - k_offset_elems: k_offset, - v_offset_elems: v_offset, - stride_page, - .. - } = checked_paged_geometry( - "batch hd512 decode via prefill", - layout, - kv_buffer.len(), - layer, - head_dim, - num_kv_heads, - false, - )?; - - let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let (pi_ptr, _gpi) = plan.page_indices_d.device_ptr(&ctx.stream); - let (pip_ptr, _gpip) = plan.page_indptr_d.device_ptr(&ctx.stream); - let (lpl_ptr, _glpl) = plan.last_page_len_d.device_ptr(&ctx.stream); - let (qi_ptr, _gqi) = plan.q_indptr_d.device_ptr(&ctx.stream); - let (ri_ptr, _gri) = plan.request_indices_d.device_ptr(&ctx.stream); - let (qti_ptr, _gqti) = plan.qo_tile_indices_d.device_ptr(&ctx.stream); - let (kti_ptr, _gkti) = plan.kv_tile_indices_d.device_ptr(&ctx.stream); - let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); - let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); - - let result = unsafe { - ffi::batch_prefill_paged_cuda_hd512( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const ffi::Half, - k_offset, - v_offset, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - qi_ptr as *const i32, - ri_ptr as *const i32, - qti_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - tnr_ptr as *const u32, - num_qo_heads as i32, - num_kv_heads as i32, - head_dim as i32, - layout.page_size as i32, - plan.total_tokens as i32, - plan.batch_size(), - plan.num_tiles, - stride_page, - sm_scale, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "batch_prefill_paged_cuda_hd512 (decode via prefill) failed for layer {layer}, \ - bs={}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}: {result}{}", - plan.batch_size(), - plan.num_tiles, - crate::ops::ffi_exception_message(result) - ); - } - - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -/// The pages referenced by `plan.page_indices_d()` must lie within -/// `kv_buffer`; the wrapper cannot verify device-side contents — caller -/// contract. -pub fn batch_prefill_paged_hd512_into( - ctx: &DeviceContext, - q: &HiddenStates, - kv_buffer: &CudaSlice, - layout: &PagedKvLayout, - layer: usize, - plan: &PrefillPagedPlan, - output: &mut HiddenStates, - num_qo_heads: usize, - sm_scale: f32, -) -> Result<()> { - anyhow::ensure!( - sm_scale.is_finite(), - "batch_prefill_paged_hd512 sm_scale {sm_scale} must be finite" - ); - let num_kv_heads = layout.num_kv_heads; - let head_dim = layout.head_dim; - anyhow::ensure!( - head_dim == 512, - "hd512 prefill expects head_dim 512, got {}", - head_dim - ); - anyhow::ensure!( - q.hidden_dim == num_qo_heads * 512, - "hd512 prefill q.hidden_dim {} != num_qo_heads {} * 512", - q.hidden_dim, - num_qo_heads - ); - anyhow::ensure!( - output.hidden_dim == num_qo_heads * 512, - "hd512 prefill output.hidden_dim {} != num_qo_heads {} * 512", - output.hidden_dim, - num_qo_heads - ); - anyhow::ensure!( - q.seq_len >= plan.total_tokens, - "hd512 prefill q.seq_len {} < plan.total_tokens {}", - q.seq_len, - plan.total_tokens - ); - anyhow::ensure!( - output.seq_len >= plan.total_tokens, - "hd512 prefill output.seq_len {} < plan.total_tokens {}", - output.seq_len, - plan.total_tokens - ); - anyhow::ensure!( - layer < layout.num_layers, - "hd512 prefill layer {layer} >= layout.num_layers {}", - layout.num_layers - ); - q.checked_extent("batch_prefill_paged_hd512 q")?; - output.checked_extent("batch_prefill_paged_hd512 output")?; - anyhow::ensure!( - plan.head_dim == 512, - "plan built for head_dim {}, expected 512", - plan.head_dim - ); - anyhow::ensure!( - plan.num_kv_heads == num_kv_heads, - "plan built for num_kv_heads {}, expected {}", - plan.num_kv_heads, - num_kv_heads - ); - anyhow::ensure!( - plan.num_qo_heads == num_qo_heads, - "plan built for num_qo_heads {}, expected {}", - plan.num_qo_heads, - num_qo_heads - ); - - // Plan tiles laid out for a cta_tile_q override would be misread by the - // kernel's automatic derivation. - let kernel_cta_tile_q = unsafe { - ffi::batch_prefill_cta_tile_q_with_override( - plan.total_tokens as i32, - plan.num_qo_heads as i32, - plan.num_kv_heads as i32, - plan.head_dim as i32, - 0, // auto-detect, matching the hd512 kernel - ) - }; - anyhow::ensure!( - kernel_cta_tile_q == plan.cta_tile_q(), - "hd512 prefill cta_tile_q: plan recorded {}, kernel derives {}; \ - overridden plans are not supported at hd512", - plan.cta_tile_q(), - kernel_cta_tile_q - ); - - let PagedGeometry { - k_offset_elems: k_offset, - v_offset_elems: v_offset, - stride_page, - .. - } = checked_paged_geometry( - "hd512 prefill", - layout, - kv_buffer.len(), - layer, - head_dim, - num_kv_heads, - false, - )?; - - let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let (pi_ptr, _gpi) = plan.page_indices_d.device_ptr(&ctx.stream); - let (pip_ptr, _gpip) = plan.page_indptr_d.device_ptr(&ctx.stream); - let (lpl_ptr, _glpl) = plan.last_page_len_d.device_ptr(&ctx.stream); - let (qi_ptr, _gqi) = plan.q_indptr_d.device_ptr(&ctx.stream); - let (ri_ptr, _gri) = plan.request_indices_d.device_ptr(&ctx.stream); - let (qti_ptr, _gqti) = plan.qo_tile_indices_d.device_ptr(&ctx.stream); - let (kti_ptr, _gkti) = plan.kv_tile_indices_d.device_ptr(&ctx.stream); - let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); - let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); - - let result = unsafe { - ffi::batch_prefill_paged_cuda_hd512( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const ffi::Half, - k_offset, - v_offset, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - qi_ptr as *const i32, - ri_ptr as *const i32, - qti_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - tnr_ptr as *const u32, - num_qo_heads as i32, - num_kv_heads as i32, - head_dim as i32, - layout.page_size as i32, - plan.total_tokens as i32, - plan.batch_size(), - plan.num_tiles, - stride_page, - sm_scale, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "batch_prefill_paged_cuda_hd512 (prefill) failed for layer {layer}, \ - bs={}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}: {result}{}", - plan.batch_size(), - plan.num_tiles, - crate::ops::ffi_exception_message(result) - ); - } - - Ok(()) -} - -/// Windowed batch prefill over paged KV at head_dim 256 (Gemma 4 local -/// layers): attention read only — the prep kernel has already written K/V -/// into the pool. `window_left` is an inclusive distance: an N-token window -/// passes N - 1, and -1 degrades to full attention. `sm_scale` is the -/// caller's; Gemma 4 runs unscaled attention (1.0). -#[allow(clippy::too_many_arguments)] -pub fn batch_prefill_paged_window_hd256_into( - ctx: &DeviceContext, - q: &HiddenStates, - kv_buffer: &CudaSlice, - layout: &PagedKvLayout, - layer: usize, - plan: &PrefillPagedPlan, - output: &mut HiddenStates, - num_qo_heads: usize, - sm_scale: f32, - window_left: i32, -) -> Result<()> { - anyhow::ensure!( - sm_scale.is_finite(), - "batch_prefill_paged_window_hd256 sm_scale {sm_scale} must be finite" - ); - anyhow::ensure!( - window_left >= -1, - "batch_prefill_paged_window_hd256 window_left {window_left} must be >= -1" - ); - let num_kv_heads = layout.num_kv_heads; - let geometry = checked_paged_geometry( - "hd256 window prefill", - layout, - kv_buffer.len(), - layer, - 256, - num_kv_heads, - true, - )?; - anyhow::ensure!( - plan.max_page_index < geometry.num_pages, - "hd256 window prefill plan references page {} but the pool holds {} pages; the attention \ - kernel computes addresses from page ids with no device-side bounds check", - plan.max_page_index, - geometry.num_pages - ); - anyhow::ensure!( - plan.max_last_page_len <= geometry.page_size, - "hd256 window prefill plan last-page length {} exceeds page_size {}; the kernel would \ - read a page-table entry that does not exist", - plan.max_last_page_len, - geometry.page_size - ); - let qo_dim = num_qo_heads.checked_mul(256).ok_or_else(|| { - anyhow::anyhow!("hd256 window prefill num_qo_heads {num_qo_heads} * 256 overflows") - })?; - anyhow::ensure!( - q.hidden_dim == qo_dim, - "hd256 window prefill q.hidden_dim {} != num_qo_heads {num_qo_heads} * 256", - q.hidden_dim - ); - anyhow::ensure!( - output.hidden_dim == qo_dim, - "hd256 window prefill output.hidden_dim {} != num_qo_heads {num_qo_heads} * 256", - output.hidden_dim - ); - anyhow::ensure!( - q.seq_len >= plan.total_tokens, - "hd256 window prefill q.seq_len {} < plan.total_tokens {}", - q.seq_len, - plan.total_tokens - ); - anyhow::ensure!( - output.seq_len >= plan.total_tokens, - "hd256 window prefill output.seq_len {} < plan.total_tokens {}", - output.seq_len, - plan.total_tokens - ); - let q_elems = q.checked_extent("batch_prefill_paged_window_hd256 q")?; - let out_elems = output.checked_extent("batch_prefill_paged_window_hd256 output")?; - crate::ops::checked_i32(q_elems, "hd256 window prefill q extent")?; - crate::ops::checked_i32(out_elems, "hd256 window prefill output extent")?; - anyhow::ensure!( - plan.head_dim == 256, - "plan built for head_dim {}, expected 256", - plan.head_dim - ); - anyhow::ensure!( - plan.num_kv_heads == num_kv_heads, - "plan built for num_kv_heads {}, expected {}", - plan.num_kv_heads, - num_kv_heads - ); - anyhow::ensure!( - plan.num_qo_heads == num_qo_heads, - "plan built for num_qo_heads {}, expected {}", - plan.num_qo_heads, - num_qo_heads - ); - - // Plan tiles laid out for a cta_tile_q override would be misread by the - // kernel's automatic derivation. - let num_qo_heads_i32 = - crate::ops::checked_i32(num_qo_heads, "hd256 window prefill num_qo_heads")?; - let num_kv_heads_i32 = - crate::ops::checked_i32(num_kv_heads, "hd256 window prefill num_kv_heads")?; - let total_tokens_i32 = - crate::ops::checked_i32(plan.total_tokens, "hd256 window prefill total_tokens")?; - let kernel_cta_tile_q = unsafe { - ffi::batch_prefill_cta_tile_q_with_override( - total_tokens_i32, - num_qo_heads_i32, - num_kv_heads_i32, - 256, - 0, // auto-detect, matching the hd256 kernel - ) - }; - anyhow::ensure!( - kernel_cta_tile_q == plan.cta_tile_q(), - "hd256 window prefill cta_tile_q: plan recorded {}, kernel derives {}; \ - overridden plans are not supported here", - plan.cta_tile_q(), - kernel_cta_tile_q - ); - - let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let (pi_ptr, _gpi) = plan.page_indices_d.device_ptr(&ctx.stream); - let (pip_ptr, _gpip) = plan.page_indptr_d.device_ptr(&ctx.stream); - let (lpl_ptr, _glpl) = plan.last_page_len_d.device_ptr(&ctx.stream); - let (qi_ptr, _gqi) = plan.q_indptr_d.device_ptr(&ctx.stream); - let (ri_ptr, _gri) = plan.request_indices_d.device_ptr(&ctx.stream); - let (qti_ptr, _gqti) = plan.qo_tile_indices_d.device_ptr(&ctx.stream); - let (kti_ptr, _gkti) = plan.kv_tile_indices_d.device_ptr(&ctx.stream); - let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); - let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); - - let (result, entry_point) = match layout.storage { - KvStorage::E4m3 => { - #[cfg(feature = "gemma4")] - let result = unsafe { - ffi::gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const core::ffi::c_void, - geometry.k_offset_elems, - geometry.v_offset_elems, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - qi_ptr as *const i32, - ri_ptr as *const i32, - qti_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - tnr_ptr as *const u32, - num_qo_heads_i32, - num_kv_heads_i32, - 256, - geometry.page_size, - total_tokens_i32, - plan.batch_size(), - plan.num_tiles, - geometry.stride_page, - sm_scale, - 0, - window_left, - crate::tensor::active_cu_stream(ctx), - ) - }; - #[cfg(not(feature = "gemma4"))] - anyhow::bail!("hd256 windowed batch prefill: fp8 KV needs the gemma4 feature"); - #[cfg(feature = "gemma4")] - (result, "gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda") - } - KvStorage::Bf16 => ( - unsafe { - ffi::batch_prefill_paged_window_cuda_hd256( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const ffi::Half, - geometry.k_offset_elems, - geometry.v_offset_elems, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - qi_ptr as *const i32, - ri_ptr as *const i32, - qti_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - tnr_ptr as *const u32, - num_qo_heads_i32, - num_kv_heads_i32, - 256, - geometry.page_size, - total_tokens_i32, - plan.batch_size(), - plan.num_tiles, - geometry.stride_page, - sm_scale, - window_left, - crate::tensor::active_cu_stream(ctx), - ) - }, - "batch_prefill_paged_window_cuda_hd256", - ), - }; - if result != 0 { - anyhow::bail!( - "{entry_point} failed for layer {layer}, \ - bs={}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}, \ - window_left={window_left}: {result}{}", - plan.batch_size(), - plan.num_tiles, - crate::ops::ffi_exception_message(result) - ); - } - - Ok(()) -} - -/// `q` rows are NHD `[seq, heads, 512]`; `k_cache`/`v_cache` are HND -/// `[heads, max_seq_len, 512]` — the head axis is strided by `max_seq_len`. -pub fn single_prefill_hd512_into( - ctx: &DeviceContext, - q: &HiddenStates, - row_offset: usize, - q_seq_len: usize, - k_cache: &HiddenStates, - v_cache: &HiddenStates, - output: &mut HiddenStates, - num_q_heads: usize, - num_kv_heads: usize, - kv_len: usize, - sm_scale: f32, -) -> Result<()> { - anyhow::ensure!( - sm_scale.is_finite(), - "single_prefill_hd512 sm_scale {sm_scale} must be finite" - ); - assert_eq!(q.hidden_dim, num_q_heads * 512); - assert_eq!(output.hidden_dim, q.hidden_dim); - assert_eq!(output.seq_len, q.seq_len); - assert_eq!(k_cache.hidden_dim, num_kv_heads * 512); - assert_eq!(v_cache.hidden_dim, k_cache.hidden_dim); - assert_eq!(v_cache.seq_len, k_cache.seq_len); - assert!(kv_len <= k_cache.seq_len); - assert!( - q_seq_len <= kv_len, - "causal prefill q_seq_len {q_seq_len} exceeds kv_len {kv_len}" - ); - let byte_offset = checked_row_offset(q, row_offset, q_seq_len, "single_prefill_hd512 q")?; - output.checked_extent("single_prefill_hd512 output")?; - k_cache.checked_extent("single_prefill_hd512 k_cache")?; - v_cache.checked_extent("single_prefill_hd512 v_cache")?; - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let q_ptr = q_ptr + byte_offset; - let (k_ptr, _gk) = k_cache.data.device_ptr(&ctx.stream); - let (v_ptr, _gv) = v_cache.data.device_ptr(&ctx.stream); - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let out_ptr = out_ptr + byte_offset; - let result = unsafe { - ffi::single_prefill_cuda_hd512( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - k_ptr as *const ffi::Half, - v_ptr as *const ffi::Half, - num_q_heads as i32, - num_kv_heads as i32, - q_seq_len as i32, - kv_len as i32, - k_cache.seq_len as i32, - sm_scale, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "single_prefill_cuda_hd512 failed with error {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -/// Plain-w QK RMSNorm + RoPE prep at head_dim 256 (Gemma 4 local layers), -/// oracle form: Q and K each land in a contiguous token-major output shaped -/// like its input; the serving form is -/// `qkv_norm_rope_paged_prefill_hd256_plain_into`. No gate, no V, and the -/// plain weight multiply — not the (1+w) offset of the qwen35 hd256 -/// sibling. `single_prefill` consumes a contiguous **HND** cache, so these -/// outputs are reassembled per head before attention, not fed directly. -/// `rotary_dim` must be positive, even and <= 256 (launcher-enforced, -/// surfaced here as `Err`); Gemma 4 local layers pass 256, and a smaller -/// even width keeps the hd512 sibling's partial-RoPE tail semantics. -#[allow(clippy::too_many_arguments)] -pub fn qk_norm_rope_prefill_hd256_plain_into( - ctx: &DeviceContext, - q: &HiddenStates, - k: &HiddenStates, - q_out: &mut HiddenStates, - k_out: &mut HiddenStates, - q_norm_weight: &DeviceVec, - k_norm_weight: &DeviceVec, - cos_cache: &DeviceVec, - sin_cache: &DeviceVec, - start_pos: usize, - cos_max_pos: usize, - num_q_heads: usize, - num_kv_heads: usize, - rotary_dim: usize, - rms_eps: f32, -) -> Result<()> { - let seq_len = q.seq_len; - let q_dim = num_q_heads - .checked_mul(256) - .ok_or_else(|| anyhow::anyhow!("hd256 plain prep num_q_heads {num_q_heads} overflows"))?; - let kv_dim = num_kv_heads - .checked_mul(256) - .ok_or_else(|| anyhow::anyhow!("hd256 plain prep num_kv_heads {num_kv_heads} overflows"))?; - anyhow::ensure!( - q.hidden_dim == q_dim, - "hd256 plain prep q.hidden_dim {} != num_q_heads {} * 256", - q.hidden_dim, - num_q_heads - ); - anyhow::ensure!( - k.hidden_dim == kv_dim, - "hd256 plain prep k.hidden_dim {} != num_kv_heads {} * 256", - k.hidden_dim, - num_kv_heads - ); - anyhow::ensure!( - q_out.hidden_dim == q.hidden_dim, - "hd256 plain prep q_out.hidden_dim {} != q.hidden_dim {}", - q_out.hidden_dim, - q.hidden_dim - ); - anyhow::ensure!( - k_out.hidden_dim == k.hidden_dim, - "hd256 plain prep k_out.hidden_dim {} != k.hidden_dim {}", - k_out.hidden_dim, - k.hidden_dim - ); - anyhow::ensure!( - k.seq_len == seq_len, - "hd256 plain prep k.seq_len {} != q.seq_len {seq_len}", - k.seq_len - ); - anyhow::ensure!( - q_out.seq_len == seq_len, - "hd256 plain prep q_out.seq_len {} != q.seq_len {seq_len}", - q_out.seq_len - ); - anyhow::ensure!( - k_out.seq_len == seq_len, - "hd256 plain prep k_out.seq_len {} != q.seq_len {seq_len}", - k_out.seq_len - ); - anyhow::ensure!( - q_norm_weight.len == 256, - "hd256 plain prep q_norm_weight len {} != 256", - q_norm_weight.len - ); - anyhow::ensure!( - k_norm_weight.len == 256, - "hd256 plain prep k_norm_weight len {} != 256", - k_norm_weight.len - ); - let end_pos = start_pos.checked_add(seq_len).ok_or_else(|| { - anyhow::anyhow!("hd256 plain prep start_pos {start_pos} + seq_len {seq_len} overflows") - })?; - anyhow::ensure!( - end_pos <= cos_max_pos, - "hd256 plain prep start_pos {start_pos} + seq_len {seq_len} > \ - cos_max_pos {cos_max_pos}" - ); - let table_len = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { - anyhow::anyhow!( - "hd256 plain prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" - ) - })?; - anyhow::ensure!( - cos_cache.len >= table_len, - "hd256 plain prep cos_cache len {} < cos_max_pos {cos_max_pos} * \ - rotary_dim {rotary_dim}", - cos_cache.len - ); - anyhow::ensure!( - sin_cache.len >= table_len, - "hd256 plain prep sin_cache len {} < cos_max_pos {cos_max_pos} * \ - rotary_dim {rotary_dim}", - sin_cache.len - ); - ensure_vec_backed(q_norm_weight, "hd256 plain prep q_norm_weight")?; - ensure_vec_backed(k_norm_weight, "hd256 plain prep k_norm_weight")?; - ensure_vec_backed(cos_cache, "hd256 plain prep cos_cache")?; - ensure_vec_backed(sin_cache, "hd256 plain prep sin_cache")?; - // The kernel indexes token-major offsets and `pos * rotary_dim + d` in - // i32, so the extents and the rope table extent must fit. - let q_extent = q.checked_extent("hd256 plain prep q")?; - let k_extent = k.checked_extent("hd256 plain prep k")?; - q_out.checked_extent("hd256 plain prep q_out")?; - k_out.checked_extent("hd256 plain prep k_out")?; - super::checked_i32(q_extent, "hd256 plain prep q extent")?; - super::checked_i32(k_extent, "hd256 plain prep k extent")?; - super::checked_i32(table_len, "hd256 plain prep rope table extent")?; - let num_q_heads = super::checked_i32(num_q_heads, "hd256 plain prep num_q_heads")?; - let num_kv_heads = super::checked_i32(num_kv_heads, "hd256 plain prep num_kv_heads")?; - let seq_len = super::checked_i32(seq_len, "hd256 plain prep seq_len")?; - let start_pos = super::checked_i32(start_pos, "hd256 plain prep start_pos")?; - let cos_max_pos = super::checked_i32(cos_max_pos, "hd256 plain prep cos_max_pos")?; - let rotary_dim = super::checked_i32(rotary_dim, "hd256 plain prep rotary_dim")?; - - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let (k_ptr, _gk) = k.data.device_ptr(&ctx.stream); - let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); - let (ko_ptr, _gko) = k_out.data.device_ptr_mut(&ctx.stream); - let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); - let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); - let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); - let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); - - let result = unsafe { - ffi::qk_norm_rope_prefill_hd256_plain_cuda( - q_ptr as *const ffi::Half, - k_ptr as *const ffi::Half, - qn_ptr as *const ffi::Half, - kn_ptr as *const ffi::Half, - cos_ptr as *const ffi::Half, - sin_ptr as *const ffi::Half, - qo_ptr as *mut ffi::Half, - ko_ptr as *mut ffi::Half, - num_q_heads, - num_kv_heads, - seq_len, - start_pos, - cos_max_pos, - rotary_dim, - rms_eps, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "qk_norm_rope_prefill_hd256_plain_cuda failed with error \ - {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -/// Everything a paged-pool launch (prep write or attention read) derives -/// from a caller-supplied `PagedKvLayout`, re-derived with overflow checks -/// and narrowed through -/// checked conversions: the layout's fields are public, so a contradictory -/// or wrapping layout must fail here, before anything reaches an unsafe -/// launch. -struct PagedGeometry { - k_offset_elems: i64, - v_offset_elems: i64, - page_size: i32, - num_pages: i32, - stride_page: i64, -} - -fn checked_paged_geometry( - what: &str, - layout: &PagedKvLayout, - pool_len: usize, - layer: usize, - head_dim: usize, - num_kv_heads: usize, - fp8_capable: bool, -) -> Result { - // This is the split geometry: a K block and a V block per layer, which - // is what FlashInfer's paged view and the prep scatter address through - // `k_offset`/`v_offset`. Neither can express another format, so one is - // refused here rather than read at the wrong rows. - anyhow::ensure!( - layout.format == KvFormat::Split, - "{what} reads the split K|V format; the layout is {:?}", - layout.format - ); - // `pool_len` counts the pool's bf16 backing slots; an e4m3 pool packs two - // elements per slot. Only wrappers with an fp8 kernel twin may see one. - anyhow::ensure!( - fp8_capable || layout.storage == KvStorage::Bf16, - "{what} has no fp8 KV path; the layout carries {}-byte elements", - layout.storage.elem_bytes() - ); - let pool_len = match layout.storage { - KvStorage::E4m3 => pool_len - .checked_mul(2) - .ok_or_else(|| anyhow::anyhow!("{what} fp8 pool element count overflows"))?, - KvStorage::Bf16 => pool_len, - }; - anyhow::ensure!( - layout.head_dim == head_dim, - "{what} layout.head_dim {} != {head_dim}", - layout.head_dim - ); - anyhow::ensure!( - layout.num_kv_heads == num_kv_heads, - "{what} layout.num_kv_heads {} != num_kv_heads {num_kv_heads}", - layout.num_kv_heads - ); - // A zero head_dim makes every stride zero, and `pool_len / page_stride` - // below would panic on a layout that is otherwise self-consistent. - anyhow::ensure!( - layout.page_size > 0 - && layout.num_kv_heads > 0 - && layout.num_layers > 0 - && layout.head_dim > 0, - "{what} layout dims must be positive: page_size {} num_kv_heads {} num_layers {} \ - head_dim {}", - layout.page_size, - layout.num_kv_heads, - layout.num_layers, - layout.head_dim - ); - // The layout's fields are public, so its strides are re-derived from its - // primitives and held to what it carries -- through the one derivation - // the layout itself uses, not a second statement of the arithmetic. - let (kv_block_len, layer_stride, page_stride) = derive_strides( - layout.page_size, - layout.num_kv_heads, - layout.head_dim, - layout.num_layers, - layout.format, - ) - .ok_or_else(|| { - anyhow::anyhow!( - "{what} strides overflow: page_size {} num_kv_heads {} head_dim {} num_layers {}", - layout.page_size, - layout.num_kv_heads, - layout.head_dim, - layout.num_layers - ) - })?; - anyhow::ensure!( - layout.kv_block_len == kv_block_len, - "{what} layout.kv_block_len {} != page_size * num_kv_heads * head_dim {kv_block_len}", - layout.kv_block_len - ); - anyhow::ensure!( - layout.layer_stride == layer_stride, - "{what} layout.layer_stride {} != the format's {layer_stride}", - layout.layer_stride - ); - anyhow::ensure!( - layout.page_stride == page_stride, - "{what} layout.page_stride {} != num_layers {} * layer_stride {layer_stride}", - layout.page_stride, - layout.num_layers - ); - anyhow::ensure!( - layer < layout.num_layers, - "{what} layer {layer} >= layout.num_layers {}", - layout.num_layers - ); - anyhow::ensure!( - pool_len.is_multiple_of(page_stride), - "{what} kv_pool.len {pool_len} is not a multiple of layout.page_stride {page_stride}" - ); - let num_pages = pool_len / page_stride; - anyhow::ensure!( - num_pages >= 1, - "{what} kv_pool.len {pool_len} holds no whole page (page_stride {page_stride})" - ); - let k_offset = layer.checked_mul(layer_stride).ok_or_else(|| { - anyhow::anyhow!("{what} layer {layer} * layer_stride {layer_stride} overflows") - })?; - let v_offset = k_offset.checked_add(kv_block_len).ok_or_else(|| { - anyhow::anyhow!("{what} K offset {k_offset} + kv_block_len {kv_block_len} overflows") - })?; - Ok(PagedGeometry { - k_offset_elems: i64::try_from(k_offset) - .map_err(|_| anyhow::anyhow!("{what} K offset {k_offset} does not fit i64"))?, - v_offset_elems: i64::try_from(v_offset) - .map_err(|_| anyhow::anyhow!("{what} V offset {v_offset} does not fit i64"))?, - page_size: crate::ops::checked_i32(layout.page_size, "paged prep page_size")?, - num_pages: crate::ops::checked_i32(num_pages, "paged prep num_pages")?, - stride_page: i64::try_from(page_stride) - .map_err(|_| anyhow::anyhow!("{what} page_stride {page_stride} does not fit i64"))?, - }) -} - -/// What the hd512 prep writes into, from the layout's format: the layer's -/// K and V blocks and the row's bands. For the split format the blocks are -/// the layer's two and the row is the head; for the folded one both offsets -/// name the layer's single block, the row is `head_dim + rotary` wide and -/// `fold_rotary` is the rotated columns the kernel keeps of K. -struct Hd512PrepTarget { - k_offset_elems: i64, - v_offset_elems: i64, - row_width: i32, - fold_rotary: i32, - page_size: i32, - num_pages: i32, - stride_page: i64, -} - -fn hd512_prep_target( - what: &str, - layout: &PagedKvLayout, - pool_len: usize, - layer: usize, - num_kv_heads: usize, -) -> Result { - anyhow::ensure!( - layout.storage == KvStorage::Bf16, - "{what} has no fp8 KV path; the layout carries {}-byte elements", - layout.storage.elem_bytes() - ); - anyhow::ensure!( - layout.head_dim == 512, - "{what} layout.head_dim {} != 512", - layout.head_dim - ); - anyhow::ensure!( - layout.num_kv_heads == num_kv_heads, - "{what} layout.num_kv_heads {} != num_kv_heads {num_kv_heads}", - layout.num_kv_heads - ); - let block = layout.block_geometry(pool_len, layer, what)?; - let v_offset = match layout.format { - KvFormat::Split => block - .block_offset_elems - .checked_add(layout.kv_block_len) - .ok_or_else(|| { - anyhow::anyhow!( - "{what} K offset {} + kv_block_len {} overflows", - block.block_offset_elems, - layout.kv_block_len - ) - })?, - KvFormat::Folded { .. } => block.block_offset_elems, - }; - let to_i64 = |value: usize, name: &str| { - i64::try_from(value).map_err(|_| anyhow::anyhow!("{what} {name} {value} does not fit i64")) - }; - Ok(Hd512PrepTarget { - k_offset_elems: to_i64(block.block_offset_elems, "K offset")?, - v_offset_elems: to_i64(v_offset, "V offset")?, - row_width: crate::ops::checked_i32( - layout.format.row_width(layout.head_dim), - &format!("{what} row_width"), - )?, - fold_rotary: crate::ops::checked_i32( - layout.format.fold_rotary(), - &format!("{what} fold_rotary"), - )?, - page_size: crate::ops::checked_i32(block.page_size, &format!("{what} page_size"))?, - num_pages: crate::ops::checked_i32(block.num_pages, &format!("{what} num_pages"))?, - stride_page: to_i64(block.stride_page, "page_stride")?, - }) -} - -/// Plain-w QKV prep at head_dim 256 (Gemma 4 local layers), serving form. -/// Q is normalised + rotated into a contiguous `q_out`; K is normalised + -/// rotated straight into the paged KV pool at layer `layer`'s K block, and -/// V — a separate v_proj head vector, unlike the hd512 K=V fork — is -/// weightless-normalised (never rotated) into the layer's V block, all in -/// one kernel with no intermediate scatter. `q`, `k` and `v` are column -/// bands read at their own row strides, so they may share one fused Q|K|V -/// projection row. -/// -/// The kernel `__trap()`s on any out-of-range pos or page id as the second -/// layer of the host validation's defence. `page_indices` is the resident -/// page row and `page_origin` the absolute page its first entry covers, so -/// a caller that released the front passes the released count; the origin is -/// page-aligned, which keeps in-page offsets position-invariant and shifts -/// only the row index. RoPE still runs on absolute positions. -#[allow(clippy::too_many_arguments)] -pub fn qkv_norm_rope_paged_prefill_hd256_plain_into<'a>( - ctx: &DeviceContext, - q: impl Into>, - k: impl Into>, - v: impl Into>, - q_out: &mut HiddenStates, - row_offset: usize, - kv_pool: &CudaSlice, - layout: &PagedKvLayout, - q_norm_weight: &DeviceVec, - k_norm_weight: &DeviceVec, - cos_cache: &DeviceVec, - sin_cache: &DeviceVec, - layer: usize, - page_indices: &CudaSlice, - pages_offset: usize, - page_origin: usize, - start_pos: usize, - cos_max_pos: usize, - num_q_heads: usize, - num_kv_heads: usize, - rotary_dim: usize, - rms_eps: f32, -) -> Result<()> { - let (q, k, v) = (q.into(), k.into(), v.into()); - // The prompt segment is the row suffix `[row_offset..seq_len)` with its - // page table `pages_offset` elements into the (possibly concatenated) - // table — a multi-prompt mixed step parks earlier prompts and their - // tables in the prefixes. - anyhow::ensure!( - row_offset < q.states.seq_len, - "hd256 paged prep row_offset {row_offset} leaves no rows of {}", - q.states.seq_len - ); - let seq_len = q.states.seq_len - row_offset; - anyhow::ensure!( - pages_offset < page_indices.len(), - "hd256 paged prep pages_offset {pages_offset} exceeds table len {}", - page_indices.len() - ); - let q_dim = num_q_heads.checked_mul(256).ok_or_else(|| { - anyhow::anyhow!("hd256 paged prep num_q_heads {num_q_heads} * 256 overflows") - })?; - let kv_dim = num_kv_heads.checked_mul(256).ok_or_else(|| { - anyhow::anyhow!("hd256 paged prep num_kv_heads {num_kv_heads} * 256 overflows") - })?; - anyhow::ensure!( - q.width == q_dim, - "hd256 paged prep q width {} != num_q_heads {num_q_heads} * 256", - q.width - ); - anyhow::ensure!( - q_out.hidden_dim == q.width, - "hd256 paged prep q_out.hidden_dim {} != q width {}", - q_out.hidden_dim, - q.width - ); - anyhow::ensure!( - q_out.seq_len == q.states.seq_len, - "hd256 paged prep q_out.seq_len {} != q rows {}", - q_out.seq_len, - q.states.seq_len - ); - anyhow::ensure!( - k.width == kv_dim, - "hd256 paged prep k width {} != num_kv_heads {num_kv_heads} * 256", - k.width - ); - anyhow::ensure!( - v.width == k.width, - "hd256 paged prep v width {} != k width {}", - v.width, - k.width - ); - anyhow::ensure!( - k.states.seq_len == q.states.seq_len, - "hd256 paged prep k rows {} != q rows {}", - k.states.seq_len, - q.states.seq_len - ); - anyhow::ensure!( - v.states.seq_len == q.states.seq_len, - "hd256 paged prep v rows {} != q rows {}", - v.states.seq_len, - q.states.seq_len - ); - let geometry = checked_paged_geometry( - "hd256 paged prep", - layout, - kv_pool.len(), - layer, - 256, - num_kv_heads, - true, - )?; - ensure_vec_backed(q_norm_weight, "hd256 paged prep q_norm_weight")?; - ensure_vec_backed(k_norm_weight, "hd256 paged prep k_norm_weight")?; - ensure_vec_backed(cos_cache, "hd256 paged prep cos_cache")?; - ensure_vec_backed(sin_cache, "hd256 paged prep sin_cache")?; - anyhow::ensure!( - q_norm_weight.len == 256, - "hd256 paged prep q_norm_weight len {} != 256", - q_norm_weight.len - ); - anyhow::ensure!( - k_norm_weight.len == 256, - "hd256 paged prep k_norm_weight len {} != 256", - k_norm_weight.len - ); - let end_pos = start_pos.checked_add(seq_len).ok_or_else(|| { - anyhow::anyhow!("hd256 paged prep start_pos {start_pos} + seq_len {seq_len} overflows") - })?; - let covered = (page_indices.len() - pages_offset) - .checked_mul(layout.page_size) - .ok_or_else(|| { - anyhow::anyhow!( - "hd256 paged prep page window {} * page_size {} overflows", - page_indices.len() - pages_offset, - layout.page_size - ) - })?; - let row_start = page_origin.checked_mul(layout.page_size).ok_or_else(|| { - anyhow::anyhow!( - "hd256 paged prep page_origin {page_origin} * page_size {} overflows", - layout.page_size - ) - })?; - let row_end = row_start.checked_add(covered).ok_or_else(|| { - anyhow::anyhow!("hd256 paged prep row start {row_start} + span {covered} overflows") - })?; - anyhow::ensure!( - row_start <= start_pos, - "hd256 paged prep page_origin {page_origin} starts after start_pos {start_pos}" - ); - anyhow::ensure!( - row_end >= end_pos, - "hd256 paged prep row covers tokens {row_start}..{row_end}, need start_pos \ - {start_pos} + seq_len {seq_len}" - ); - anyhow::ensure!( - end_pos <= cos_max_pos, - "hd256 paged prep start_pos {start_pos} + seq_len {seq_len} > cos_max_pos {cos_max_pos}" - ); - let table_len = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { - anyhow::anyhow!( - "hd256 paged prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" - ) - })?; - anyhow::ensure!( - cos_cache.len >= table_len, - "hd256 paged prep cos_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", - cos_cache.len - ); - anyhow::ensure!( - sin_cache.len >= table_len, - "hd256 paged prep sin_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", - sin_cache.len - ); - let q_elems = q.states.checked_extent("hd256 paged prep q")?; - q_out.checked_extent("hd256 paged prep q_out")?; - let k_elems = k.states.checked_extent("hd256 paged prep k")?; - v.states.checked_extent("hd256 paged prep v")?; - crate::ops::checked_i32(q_elems, "hd256 paged prep q extent")?; - crate::ops::checked_i32(k_elems, "hd256 paged prep k extent")?; - crate::ops::checked_i32(table_len, "hd256 paged prep rope table extent")?; - - let page_origin_i32 = crate::ops::checked_i32(page_origin, "hd256 paged prep page_origin")?; - let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd256 paged prep num_q_heads")?; - let num_kv_heads_i32 = crate::ops::checked_i32(num_kv_heads, "hd256 paged prep num_kv_heads")?; - let seq_len_i32 = crate::ops::checked_i32(seq_len, "hd256 paged prep seq_len")?; - let start_pos_i32 = crate::ops::checked_i32(start_pos, "hd256 paged prep start_pos")?; - let cos_max_pos_i32 = crate::ops::checked_i32(cos_max_pos, "hd256 paged prep cos_max_pos")?; - let rotary_dim_i32 = crate::ops::checked_i32(rotary_dim, "hd256 paged prep rotary_dim")?; - - let page_indices_len = crate::ops::checked_i32( - page_indices.len() - pages_offset, - "hd256 paged prep page window len", - )?; - let q_row_bytes = checked_band_offset(&q, row_offset, seq_len, "hd256 paged prep q")?; - let k_row_bytes = checked_band_offset(&k, row_offset, seq_len, "hd256 paged prep k")?; - let v_row_bytes = checked_band_offset(&v, row_offset, seq_len, "hd256 paged prep v")?; - let qo_row_bytes = checked_row_offset(q_out, row_offset, seq_len, "hd256 paged prep q_out")?; - let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); - let q_ptr = q_ptr + q_row_bytes; - let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); - let k_ptr = k_ptr + k_row_bytes; - let (v_ptr, _gv) = v.states.data.device_ptr(&ctx.stream); - let v_ptr = v_ptr + v_row_bytes; - let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); - let qo_ptr = qo_ptr + qo_row_bytes; - // The KV state owns mutation; this call borrows its shared pool handle. - let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); - let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); - let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); - let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); - let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); - let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); - let pi_ptr = pi_ptr + (pages_offset * std::mem::size_of::()) as u64; - - // Two `if`s, not a tuple-building match: a fn item only coerces to a - // fn pointer across plain `if` arms. - let fp8 = layout.storage == KvStorage::E4m3; - let launch = if fp8 { - ffi::qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda - } else { - ffi::qkv_norm_rope_paged_prefill_hd256_plain_cuda - }; - let entry_point = if fp8 { - "qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda" - } else { - "qkv_norm_rope_paged_prefill_hd256_plain_cuda" - }; - let q_stride = crate::ops::checked_i32(q.states.hidden_dim, "hd256 paged prep q stride")?; - let k_stride = crate::ops::checked_i32(k.states.hidden_dim, "hd256 paged prep k stride")?; - let v_stride = crate::ops::checked_i32(v.states.hidden_dim, "hd256 paged prep v stride")?; - let result = unsafe { - launch( - q_ptr as *const ffi::Half, - k_ptr as *const ffi::Half, - v_ptr as *const ffi::Half, - q_stride, - k_stride, - v_stride, - qn_ptr as *const ffi::Half, - kn_ptr as *const ffi::Half, - cos_ptr as *const ffi::Half, - sin_ptr as *const ffi::Half, - qo_ptr as *mut ffi::Half, - pool_ptr as *mut ffi::Half, - geometry.k_offset_elems, - geometry.v_offset_elems, - pi_ptr as *const i32, - page_indices_len, - page_origin_i32, - num_q_heads_i32, - num_kv_heads_i32, - seq_len_i32, - start_pos_i32, - cos_max_pos_i32, - rotary_dim_i32, - rms_eps, - geometry.page_size, - geometry.num_pages, - geometry.stride_page, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "{entry_point} failed with error {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -/// Per-request hd256 prep; `page_origins` is the released local-window front. -/// The decode batch is the row suffix `[row_offset..seq_len)` — 0 covers the -/// whole tensor, a mixed step parks its prefill segment in the prefix. -/// Invalid device metadata traps before the first paged-pool access. -#[allow(clippy::too_many_arguments)] -pub fn qkv_norm_rope_paged_decode_hd256_plain_into<'a>( - ctx: &DeviceContext, - q: impl Into>, - k: impl Into>, - v: impl Into>, - q_out: &mut HiddenStates, - row_offset: usize, - kv_pool: &CudaSlice, - layout: &PagedKvLayout, - q_norm_weight: &DeviceVec, - k_norm_weight: &DeviceVec, - cos_cache: &DeviceVec, - sin_cache: &DeviceVec, - layer: usize, - page_indices: &CudaSlice, - page_indptr: &CudaSlice, - page_origins: &CudaSlice, - positions: &CudaSlice, - cos_max_pos: usize, - num_q_heads: usize, - num_kv_heads: usize, - rotary_dim: usize, - rms_eps: f32, -) -> Result<()> { - let (q, k, v) = (q.into(), k.into(), v.into()); - anyhow::ensure!( - row_offset < q.states.seq_len, - "hd256 paged decode prep row_offset {row_offset} leaves no rows of {}", - q.states.seq_len - ); - let batch = q.states.seq_len - row_offset; - anyhow::ensure!( - Some(q.width) == num_q_heads.checked_mul(256), - "hd256 paged decode prep q width {} != num_q_heads {} * 256", - q.width, - num_q_heads - ); - anyhow::ensure!( - q_out.hidden_dim == q.width && q_out.seq_len == q.states.seq_len, - "hd256 paged decode prep q_out [{} x {}] != q [{} x {}]", - q_out.hidden_dim, - q_out.seq_len, - q.width, - q.states.seq_len - ); - anyhow::ensure!( - Some(k.width) == num_kv_heads.checked_mul(256) && k.states.seq_len == q.states.seq_len, - "hd256 paged decode prep k [{} x {}] != [num_kv_heads {} * 256 x {}]", - k.width, - k.states.seq_len, - num_kv_heads, - q.states.seq_len - ); - anyhow::ensure!( - v.width == k.width && v.states.seq_len == q.states.seq_len, - "hd256 paged decode prep v [{} x {}] != k [{} x {}]", - v.width, - v.states.seq_len, - k.width, - q.states.seq_len - ); - let geometry = checked_paged_geometry( - "hd256 paged decode prep", - layout, - kv_pool.len(), - layer, - 256, - num_kv_heads, - true, - )?; - ensure_vec_backed(q_norm_weight, "hd256 paged decode prep q_norm_weight")?; - ensure_vec_backed(k_norm_weight, "hd256 paged decode prep k_norm_weight")?; - ensure_vec_backed(cos_cache, "hd256 paged decode prep cos_cache")?; - ensure_vec_backed(sin_cache, "hd256 paged decode prep sin_cache")?; - anyhow::ensure!( - q_norm_weight.len == 256 && k_norm_weight.len == 256, - "hd256 paged decode prep norm weight lens {} / {} != 256", - q_norm_weight.len, - k_norm_weight.len - ); - let table_rows = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { - anyhow::anyhow!( - "hd256 paged decode prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" - ) - })?; - crate::ops::checked_i32(table_rows, "hd256 paged decode prep rope table extent")?; - anyhow::ensure!( - cos_cache.len >= table_rows && sin_cache.len >= table_rows, - "hd256 paged decode prep cos/sin lens {} / {} < cos_max_pos {cos_max_pos} * \ - rotary_dim {rotary_dim}", - cos_cache.len, - sin_cache.len - ); - anyhow::ensure!( - positions.len() >= batch && page_origins.len() >= batch && page_indptr.len() > batch, - "hd256 paged decode prep metadata lens (positions {}, origins {}, indptr {}) \ - do not cover batch {batch}", - positions.len(), - page_origins.len(), - page_indptr.len() - ); - let page_indices_len = crate::ops::checked_i32( - page_indices.len(), - "hd256 paged decode prep page_indices len", - )?; - let q_elems = q.states.checked_extent("hd256 paged decode prep q")?; - q_out.checked_extent("hd256 paged decode prep q_out")?; - let k_elems = k.states.checked_extent("hd256 paged decode prep k")?; - v.states.checked_extent("hd256 paged decode prep v")?; - crate::ops::checked_i32(q_elems, "hd256 paged decode prep q extent")?; - crate::ops::checked_i32(k_elems, "hd256 paged decode prep k extent")?; - let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd256 paged decode prep q heads")?; - let num_kv_heads_i32 = - crate::ops::checked_i32(num_kv_heads, "hd256 paged decode prep kv heads")?; - let batch_i32 = crate::ops::checked_i32(batch, "hd256 paged decode prep batch")?; - let cos_max_pos_i32 = - crate::ops::checked_i32(cos_max_pos, "hd256 paged decode prep cos_max_pos")?; - let rotary_dim_i32 = crate::ops::checked_i32(rotary_dim, "hd256 paged decode prep rotary_dim")?; - - let q_row_bytes = checked_band_offset(&q, row_offset, batch, "hd256 paged decode prep q")?; - let k_row_bytes = checked_band_offset(&k, row_offset, batch, "hd256 paged decode prep k")?; - let v_row_bytes = checked_band_offset(&v, row_offset, batch, "hd256 paged decode prep v")?; - let qo_row_bytes = - checked_row_offset(q_out, row_offset, batch, "hd256 paged decode prep q_out")?; - let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); - let q_ptr = q_ptr + q_row_bytes; - let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); - let k_ptr = k_ptr + k_row_bytes; - let (v_ptr, _gv) = v.states.data.device_ptr(&ctx.stream); - let v_ptr = v_ptr + v_row_bytes; - let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); - let qo_ptr = qo_ptr + qo_row_bytes; - // The KV states own mutation; this call borrows the shared pool handle. - let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); - let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); - let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); - let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); - let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); - let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); - let (ip_ptr, _gip) = page_indptr.device_ptr(&ctx.stream); - let (og_ptr, _gog) = page_origins.device_ptr(&ctx.stream); - let (ps_ptr, _gps) = positions.device_ptr(&ctx.stream); - - // Two `if`s, not a tuple-building match: a fn item only coerces to a - // fn pointer across plain `if` arms. - let fp8 = layout.storage == KvStorage::E4m3; - let launch = if fp8 { - ffi::qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda - } else { - ffi::qkv_norm_rope_paged_decode_hd256_plain_cuda - }; - let entry_point = if fp8 { - "qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda" - } else { - "qkv_norm_rope_paged_decode_hd256_plain_cuda" - }; - let q_stride = - crate::ops::checked_i32(q.states.hidden_dim, "hd256 paged decode prep q stride")?; - let k_stride = - crate::ops::checked_i32(k.states.hidden_dim, "hd256 paged decode prep k stride")?; - let v_stride = - crate::ops::checked_i32(v.states.hidden_dim, "hd256 paged decode prep v stride")?; - let result = unsafe { - launch( - q_ptr as *const ffi::Half, - k_ptr as *const ffi::Half, - v_ptr as *const ffi::Half, - q_stride, - k_stride, - v_stride, - qn_ptr as *const ffi::Half, - kn_ptr as *const ffi::Half, - cos_ptr as *const ffi::Half, - sin_ptr as *const ffi::Half, - qo_ptr as *mut ffi::Half, - pool_ptr as *mut ffi::Half, - geometry.k_offset_elems, - geometry.v_offset_elems, - pi_ptr as *const i32, - page_indices_len, - ip_ptr as *const i32, - og_ptr as *const i32, - ps_ptr as *const i32, - num_q_heads_i32, - num_kv_heads_i32, - batch_i32, - cos_max_pos_i32, - rotary_dim_i32, - rms_eps, - geometry.page_size, - geometry.num_pages, - geometry.stride_page, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "{entry_point} failed with error {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -/// Causal single-sequence prefill at head_dim 256 over a contiguous **HND** -/// K/V cache — `k[head, pos, dim]` with `k_cache.seq_len` allocated rows per -/// head — matching `single_prefill_cuda_hd256`'s stride contract. Q and the -/// output stay token-major (`[num_q_heads * 256, seq_len]`). -/// -/// `sm_scale` is explicit because it is model policy, not geometry: qwen35 -/// wants the conventional `rsqrt(head_dim)`, but Gemma 4 runs unscaled -/// attention (`scaling = 1.0` in the HF reference), and a hardcoded rsqrt -/// would silently mis-scale its logits by 16x. -#[allow(clippy::too_many_arguments)] -pub fn single_prefill_hd256_into( - ctx: &DeviceContext, - q: &HiddenStates, - row_offset: usize, - q_seq_len: usize, - k_cache: &HiddenStates, - v_cache: &HiddenStates, - output: &mut HiddenStates, - num_q_heads: usize, - num_kv_heads: usize, - kv_len: usize, - sm_scale: f32, -) -> Result<()> { - anyhow::ensure!( - sm_scale.is_finite(), - "single_prefill_hd256 sm_scale {sm_scale} must be finite" - ); - anyhow::ensure!( - num_q_heads > 0 && num_kv_heads > 0 && q_seq_len > 0, - "single_prefill_hd256 num_q_heads {num_q_heads}, num_kv_heads \ - {num_kv_heads} and q_seq_len {q_seq_len} must be positive" - ); - anyhow::ensure!( - num_q_heads.is_multiple_of(num_kv_heads) - && SUPPORTED_GQA_GROUP_SIZES.contains(&(num_q_heads / num_kv_heads)), - "single_prefill_hd256 GQA group {num_q_heads} q / {num_kv_heads} kv \ - heads not in the instantiated sizes {SUPPORTED_GQA_GROUP_SIZES:?}" - ); - let q_dim = num_q_heads.checked_mul(256).ok_or_else(|| { - anyhow::anyhow!("single_prefill_hd256 num_q_heads {num_q_heads} overflows") - })?; - let kv_dim = num_kv_heads.checked_mul(256).ok_or_else(|| { - anyhow::anyhow!("single_prefill_hd256 num_kv_heads {num_kv_heads} overflows") - })?; - anyhow::ensure!( - q.hidden_dim == q_dim, - "single_prefill_hd256 q.hidden_dim {} != num_q_heads {} * 256", - q.hidden_dim, - num_q_heads - ); - anyhow::ensure!( - output.hidden_dim == q.hidden_dim && output.seq_len == q.seq_len, - "single_prefill_hd256 output {}x{} != q {}x{}", - output.hidden_dim, - output.seq_len, - q.hidden_dim, - q.seq_len - ); - anyhow::ensure!( - k_cache.hidden_dim == kv_dim, - "single_prefill_hd256 k_cache.hidden_dim {} != num_kv_heads {} * 256", - k_cache.hidden_dim, - num_kv_heads - ); - anyhow::ensure!( - v_cache.hidden_dim == k_cache.hidden_dim && v_cache.seq_len == k_cache.seq_len, - "single_prefill_hd256 v_cache {}x{} != k_cache {}x{}", - v_cache.hidden_dim, - v_cache.seq_len, - k_cache.hidden_dim, - k_cache.seq_len - ); - anyhow::ensure!( - kv_len <= k_cache.seq_len, - "single_prefill_hd256 kv_len {kv_len} exceeds cache rows {}", - k_cache.seq_len - ); - anyhow::ensure!( - q_seq_len <= kv_len, - "causal prefill q_seq_len {q_seq_len} exceeds kv_len {kv_len}" - ); - let byte_offset = checked_row_offset(q, row_offset, q_seq_len, "single_prefill_hd256 q")?; - // The C entry builds its 32-bit strides q_stride_n = num_q_heads * 256 - // and kv_stride_h = cache rows * 256; with positive head counts and - // rows, each stride is bounded by its buffer's extent, so extents - // fitting i32 cover the stride products too. - let q_extent = q.checked_extent("single_prefill_hd256 q")?; - let out_extent = output.checked_extent("single_prefill_hd256 output")?; - let k_extent = k_cache.checked_extent("single_prefill_hd256 k_cache")?; - let v_extent = v_cache.checked_extent("single_prefill_hd256 v_cache")?; - super::checked_i32(q_extent, "single_prefill_hd256 q extent")?; - super::checked_i32(out_extent, "single_prefill_hd256 output extent")?; - super::checked_i32(k_extent, "single_prefill_hd256 k_cache extent")?; - super::checked_i32(v_extent, "single_prefill_hd256 v_cache extent")?; - let num_q_heads = super::checked_i32(num_q_heads, "single_prefill_hd256 num_q_heads")?; - let num_kv_heads = super::checked_i32(num_kv_heads, "single_prefill_hd256 num_kv_heads")?; - let q_seq_len = super::checked_i32(q_seq_len, "single_prefill_hd256 q_seq_len")?; - let kv_len = super::checked_i32(kv_len, "single_prefill_hd256 kv_len")?; - let kv_stride = super::checked_i32(k_cache.seq_len, "single_prefill_hd256 cache rows")?; - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let q_ptr = q_ptr + byte_offset; - let (k_ptr, _gk) = k_cache.data.device_ptr(&ctx.stream); - let (v_ptr, _gv) = v_cache.data.device_ptr(&ctx.stream); - let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); - let out_ptr = out_ptr + byte_offset; - let result = unsafe { - ffi::single_prefill_cuda_hd256( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - k_ptr as *const ffi::Half, - v_ptr as *const ffi::Half, - num_q_heads, - num_kv_heads, - q_seq_len, - kv_len, - kv_stride, - sm_scale, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "single_prefill_cuda_hd256 failed with error {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -/// QK RMSNorm + partial RoPE for the hd512 paged-prefill prep (Gemma 4 -/// global layers). Q is normalised + partially rotated into a contiguous -/// `q_out`; K is normalised + partially rotated straight into the paged KV -/// pool at layer `layer` (feeds `batch_prefill_paged`, not -/// `single_prefill`); V — the K=V fork — is the weightless RMS norm of the -/// same raw K, written in the same pass. No gate; plain-w RMSNorm. -/// -/// The rotation is the engine's proportional one over `cos_max_pos` rows of -/// 512: pairs `(d, d + 256)` with the live angles first and the identity -/// past them. How the row is laid out — K and V as two blocks, or the -/// folded row of `KvFormat::Folded` — is the layout's format, and the -/// kernel is told it as bands rather than asked to know the format. `q` and -/// `k` are column bands read at their own row strides, so they may share one -/// fused Q|K projection row. -/// -/// The kernel `__trap()`s on any out-of-range pos or page id as the -/// second layer of the host validation's defence. -#[allow(clippy::too_many_arguments)] -pub fn qk_norm_partial_rope_paged_prefill_hd512_into<'a>( - ctx: &DeviceContext, - q: impl Into>, - k: impl Into>, - q_out: &mut HiddenStates, - row_offset: usize, - kv_pool: &CudaSlice, - layout: &PagedKvLayout, - q_norm_weight: &DeviceVec, - k_norm_weight: &DeviceVec, - cos_cache: &DeviceVec, - sin_cache: &DeviceVec, - layer: usize, - page_indices: &CudaSlice, - pages_offset: usize, - start_pos: usize, - cos_max_pos: usize, - num_q_heads: usize, - num_kv_heads: usize, - rms_eps: f32, -) -> Result<()> { - let (q, k) = (q.into(), k.into()); - // Same suffix-window contract as the hd256 prefill prep: the segment - // is the rows `[row_offset..seq_len)` with its table at `pages_offset`. - anyhow::ensure!( - row_offset < q.states.seq_len, - "hd512 prefill prep row_offset {row_offset} leaves no rows of {}", - q.states.seq_len - ); - let seq_len = q.states.seq_len - row_offset; - anyhow::ensure!( - pages_offset < page_indices.len(), - "hd512 prefill prep pages_offset {pages_offset} exceeds table len {}", - page_indices.len() - ); - let q_dim = num_q_heads.checked_mul(512).ok_or_else(|| { - anyhow::anyhow!("hd512 prefill prep num_q_heads {num_q_heads} * 512 overflows") - })?; - let kv_dim = num_kv_heads.checked_mul(512).ok_or_else(|| { - anyhow::anyhow!("hd512 prefill prep num_kv_heads {num_kv_heads} * 512 overflows") - })?; - anyhow::ensure!( - q.width == q_dim, - "hd512 prefill prep q width {} != num_q_heads {num_q_heads} * 512", - q.width - ); - anyhow::ensure!( - q_out.hidden_dim == q.width, - "hd512 prefill prep q_out.hidden_dim {} != q width {}", - q_out.hidden_dim, - q.width - ); - anyhow::ensure!( - q_out.seq_len == q.states.seq_len, - "hd512 prefill prep q_out.seq_len {} != q rows {}", - q_out.seq_len, - q.states.seq_len - ); - anyhow::ensure!( - k.width == kv_dim, - "hd512 prefill prep k width {} != num_kv_heads {num_kv_heads} * 512", - k.width - ); - anyhow::ensure!( - k.states.seq_len == q.states.seq_len, - "hd512 prefill prep k rows {} != q rows {}", - k.states.seq_len, - q.states.seq_len - ); - let target = hd512_prep_target( - "hd512 prefill prep", - layout, - kv_pool.len(), - layer, - num_kv_heads, - )?; - anyhow::ensure!( - q_norm_weight.len == 512, - "hd512 prefill prep q_norm_weight len {} != 512", - q_norm_weight.len - ); - anyhow::ensure!( - k_norm_weight.len == 512, - "hd512 prefill prep k_norm_weight len {} != 512", - k_norm_weight.len - ); - let end_pos = start_pos.checked_add(seq_len).ok_or_else(|| { - anyhow::anyhow!("hd512 prefill prep start_pos {start_pos} + seq_len {seq_len} overflows") - })?; - let covered = (page_indices.len() - pages_offset) - .checked_mul(layout.page_size) - .ok_or_else(|| { - anyhow::anyhow!( - "hd512 prefill prep page window {} * page_size {} overflows", - page_indices.len() - pages_offset, - layout.page_size - ) - })?; - anyhow::ensure!( - covered >= end_pos, - "hd512 prefill prep pages cover {covered} tokens, need start_pos {start_pos} + seq_len {seq_len}" - ); - anyhow::ensure!( - end_pos <= cos_max_pos, - "hd512 prefill prep start_pos {start_pos} + seq_len {seq_len} > cos_max_pos {cos_max_pos}" - ); - let table_len = cos_max_pos.checked_mul(512).ok_or_else(|| { - anyhow::anyhow!("hd512 prefill prep cos_max_pos {cos_max_pos} * 512 overflows") + let k_offset = layer.checked_mul(layer_stride).ok_or_else(|| { + anyhow::anyhow!("{what} layer {layer} * layer_stride {layer_stride} overflows") })?; - anyhow::ensure!( - cos_cache.len >= table_len, - "hd512 prefill prep cos_cache len {} < cos_max_pos {cos_max_pos} * 512", - cos_cache.len - ); - anyhow::ensure!( - sin_cache.len >= table_len, - "hd512 prefill prep sin_cache len {} < cos_max_pos {cos_max_pos} * 512", - sin_cache.len - ); - let q_elems = q.states.checked_extent("hd512 prefill prep q")?; - q_out.checked_extent("hd512 prefill prep q_out")?; - let k_elems = k.states.checked_extent("hd512 prefill prep k")?; - crate::ops::checked_i32(q_elems, "hd512 prefill prep q extent")?; - crate::ops::checked_i32(k_elems, "hd512 prefill prep k extent")?; - crate::ops::checked_i32(table_len, "hd512 prefill prep rope table extent")?; - ensure_vec_backed(q_norm_weight, "hd512 prefill prep q_norm_weight")?; - ensure_vec_backed(k_norm_weight, "hd512 prefill prep k_norm_weight")?; - ensure_vec_backed(cos_cache, "hd512 prefill prep cos_cache")?; - ensure_vec_backed(sin_cache, "hd512 prefill prep sin_cache")?; - - let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd512 prefill prep num_q_heads")?; - let num_kv_heads_i32 = - crate::ops::checked_i32(num_kv_heads, "hd512 prefill prep num_kv_heads")?; - let seq_len_i32 = crate::ops::checked_i32(seq_len, "hd512 prefill prep seq_len")?; - let start_pos_i32 = crate::ops::checked_i32(start_pos, "hd512 prefill prep start_pos")?; - let cos_max_pos_i32 = crate::ops::checked_i32(cos_max_pos, "hd512 prefill prep cos_max_pos")?; - - let page_indices_len = crate::ops::checked_i32( - page_indices.len() - pages_offset, - "hd512 paged prep page window len", - )?; - let q_row_bytes = checked_band_offset(&q, row_offset, seq_len, "hd512 prefill prep q")?; - let k_row_bytes = checked_band_offset(&k, row_offset, seq_len, "hd512 prefill prep k")?; - let qo_row_bytes = checked_row_offset(q_out, row_offset, seq_len, "hd512 prefill prep q_out")?; - let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); - let q_ptr = q_ptr + q_row_bytes; - let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); - let k_ptr = k_ptr + k_row_bytes; - let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); - let qo_ptr = qo_ptr + qo_row_bytes; - // The KV state owns mutation; this call borrows its shared pool handle. - let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); - let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); - let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); - let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); - let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); - let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); - let pi_ptr = pi_ptr + (pages_offset * std::mem::size_of::()) as u64; - - let q_stride = crate::ops::checked_i32(q.states.hidden_dim, "hd512 prefill prep q stride")?; - let k_stride = crate::ops::checked_i32(k.states.hidden_dim, "hd512 prefill prep k stride")?; - let result = unsafe { - ffi::qk_norm_partial_rope_paged_prefill_hd512_cuda( - q_ptr as *const ffi::Half, - k_ptr as *const ffi::Half, - q_stride, - k_stride, - qn_ptr as *const ffi::Half, - kn_ptr as *const ffi::Half, - cos_ptr as *const ffi::Half, - sin_ptr as *const ffi::Half, - qo_ptr as *mut ffi::Half, - pool_ptr as *mut ffi::Half, - target.k_offset_elems, - target.v_offset_elems, - pi_ptr as *const i32, - page_indices_len, - num_q_heads_i32, - num_kv_heads_i32, - seq_len_i32, - start_pos_i32, - cos_max_pos_i32, - target.row_width, - target.fold_rotary, - rms_eps, - target.page_size, - target.num_pages, - target.stride_page, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "qk_norm_partial_rope_paged_prefill_hd512_cuda failed with error \ - {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -/// Per-request hd512 prep for the non-evicting global cache; V is the K fork. -/// Invalid device metadata traps before the first paged-pool access. -#[allow(clippy::too_many_arguments)] -pub fn qk_norm_partial_rope_paged_decode_hd512_into<'a>( - ctx: &DeviceContext, - q: impl Into>, - k: impl Into>, - q_out: &mut HiddenStates, - row_offset: usize, - kv_pool: &CudaSlice, - layout: &PagedKvLayout, - q_norm_weight: &DeviceVec, - k_norm_weight: &DeviceVec, - cos_cache: &DeviceVec, - sin_cache: &DeviceVec, - layer: usize, - page_indices: &CudaSlice, - page_indptr: &CudaSlice, - page_origins: &CudaSlice, - positions: &CudaSlice, - cos_max_pos: usize, - num_q_heads: usize, - num_kv_heads: usize, - rms_eps: f32, -) -> Result<()> { - let (q, k) = (q.into(), k.into()); - anyhow::ensure!( - row_offset < q.states.seq_len, - "hd512 paged decode prep row_offset {row_offset} leaves no rows of {}", - q.states.seq_len - ); - let batch = q.states.seq_len - row_offset; - anyhow::ensure!( - Some(q.width) == num_q_heads.checked_mul(512), - "hd512 paged decode prep q width {} != num_q_heads {} * 512", - q.width, - num_q_heads - ); - anyhow::ensure!( - q_out.hidden_dim == q.width && q_out.seq_len == q.states.seq_len, - "hd512 paged decode prep q_out [{} x {}] != q [{} x {}]", - q_out.hidden_dim, - q_out.seq_len, - q.width, - q.states.seq_len - ); - anyhow::ensure!( - Some(k.width) == num_kv_heads.checked_mul(512) && k.states.seq_len == q.states.seq_len, - "hd512 paged decode prep k [{} x {}] != [num_kv_heads {} * 512 x {}]", - k.width, - k.states.seq_len, - num_kv_heads, - q.states.seq_len - ); - let target = hd512_prep_target( - "hd512 paged decode prep", - layout, - kv_pool.len(), - layer, - num_kv_heads, - )?; - ensure_vec_backed(q_norm_weight, "hd512 paged decode prep q_norm_weight")?; - ensure_vec_backed(k_norm_weight, "hd512 paged decode prep k_norm_weight")?; - ensure_vec_backed(cos_cache, "hd512 paged decode prep cos_cache")?; - ensure_vec_backed(sin_cache, "hd512 paged decode prep sin_cache")?; - anyhow::ensure!( - q_norm_weight.len == 512 && k_norm_weight.len == 512, - "hd512 paged decode prep norm weight lens {} / {} != 512", - q_norm_weight.len, - k_norm_weight.len - ); - let table_rows = cos_max_pos.checked_mul(512).ok_or_else(|| { - anyhow::anyhow!("hd512 paged decode prep cos_max_pos {cos_max_pos} * 512 overflows") + let v_offset = k_offset.checked_add(kv_block_len).ok_or_else(|| { + anyhow::anyhow!("{what} K offset {k_offset} + kv_block_len {kv_block_len} overflows") })?; - crate::ops::checked_i32(table_rows, "hd512 paged decode prep rope table extent")?; - anyhow::ensure!( - cos_cache.len >= table_rows && sin_cache.len >= table_rows, - "hd512 paged decode prep cos/sin lens {} / {} < cos_max_pos {cos_max_pos} * 512", - cos_cache.len, - sin_cache.len - ); - anyhow::ensure!( - positions.len() >= batch && page_indptr.len() > batch && page_origins.len() >= batch, - "hd512 paged decode prep metadata lens (positions {}, indptr {}, origins {}) do \ - not cover batch {batch}", - positions.len(), - page_indptr.len(), - page_origins.len() - ); - let page_indices_len = crate::ops::checked_i32( - page_indices.len(), - "hd512 paged decode prep page_indices len", - )?; - let q_elems = q.states.checked_extent("hd512 paged decode prep q")?; - q_out.checked_extent("hd512 paged decode prep q_out")?; - let k_elems = k.states.checked_extent("hd512 paged decode prep k")?; - crate::ops::checked_i32(q_elems, "hd512 paged decode prep q extent")?; - crate::ops::checked_i32(k_elems, "hd512 paged decode prep k extent")?; - let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd512 paged decode prep q heads")?; - let num_kv_heads_i32 = - crate::ops::checked_i32(num_kv_heads, "hd512 paged decode prep kv heads")?; - let batch_i32 = crate::ops::checked_i32(batch, "hd512 paged decode prep batch")?; - let cos_max_pos_i32 = - crate::ops::checked_i32(cos_max_pos, "hd512 paged decode prep cos_max_pos")?; - let q_row_bytes = checked_band_offset(&q, row_offset, batch, "hd512 paged decode prep q")?; - let k_row_bytes = checked_band_offset(&k, row_offset, batch, "hd512 paged decode prep k")?; - let qo_row_bytes = - checked_row_offset(q_out, row_offset, batch, "hd512 paged decode prep q_out")?; - let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); - let q_ptr = q_ptr + q_row_bytes; - let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); - let k_ptr = k_ptr + k_row_bytes; - let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); - let qo_ptr = qo_ptr + qo_row_bytes; - // The KV states own mutation; this call borrows the shared pool handle. - let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); - let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); - let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); - let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); - let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); - let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); - let (ip_ptr, _gip) = page_indptr.device_ptr(&ctx.stream); - let (po_ptr, _gpo) = page_origins.device_ptr(&ctx.stream); - let (ps_ptr, _gps) = positions.device_ptr(&ctx.stream); - - let q_stride = - crate::ops::checked_i32(q.states.hidden_dim, "hd512 paged decode prep q stride")?; - let k_stride = - crate::ops::checked_i32(k.states.hidden_dim, "hd512 paged decode prep k stride")?; - let result = unsafe { - ffi::qk_norm_partial_rope_paged_decode_hd512_cuda( - q_ptr as *const ffi::Half, - k_ptr as *const ffi::Half, - q_stride, - k_stride, - qn_ptr as *const ffi::Half, - kn_ptr as *const ffi::Half, - cos_ptr as *const ffi::Half, - sin_ptr as *const ffi::Half, - qo_ptr as *mut ffi::Half, - pool_ptr as *mut ffi::Half, - target.k_offset_elems, - target.v_offset_elems, - pi_ptr as *const i32, - page_indices_len, - ip_ptr as *const i32, - po_ptr as *const i32, - ps_ptr as *const i32, - num_q_heads_i32, - num_kv_heads_i32, - batch_i32, - cos_max_pos_i32, - target.row_width, - target.fold_rotary, - rms_eps, - target.page_size, - target.num_pages, - target.stride_page, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "qk_norm_partial_rope_paged_decode_hd512_cuda failed with error \ - {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) -} - -/// Batched QK RMSNorm + partial RoPE for the hd512 decode prep (Gemma 4 -/// global layers). Q is normalised + partially rotated into a contiguous -/// `q_out`; K is normalised + partially rotated in place (the caller -/// scatters it into the paged pool afterwards). No gate, no V; plain-w -/// RMSNorm. -/// -/// Reading `positions_d` on the host would require a D2H synchronization, -/// so the kernel `__trap()`s on positions outside [0, cos_max_pos). -/// Validated here: table lengths, `positions_d.len() == batch_size`, and -/// 512-element norm weights. `rotary_dim` must be positive, even and <= 512 -/// — enforced by the launcher; propagated here as `Err`. -#[allow(clippy::too_many_arguments)] -pub fn qk_norm_partial_rope_batched_decode_hd512_into( - ctx: &DeviceContext, - q: &HiddenStates, - q_out: &mut HiddenStates, - k: &mut HiddenStates, - q_norm_weight: &DeviceVec, - k_norm_weight: &DeviceVec, - cos_cache: &DeviceVec, - sin_cache: &DeviceVec, - positions_d: &CudaSlice, - cos_max_pos: usize, - num_q_heads: usize, - num_kv_heads: usize, - rotary_dim: usize, - rms_eps: f32, -) -> Result<()> { - let batch_size = q.seq_len; - anyhow::ensure!( - q.hidden_dim == num_q_heads * 512, - "hd512 decode prep q.hidden_dim {} != num_q_heads {} * 512", - q.hidden_dim, - num_q_heads - ); - anyhow::ensure!( - q_out.hidden_dim == q.hidden_dim, - "hd512 decode prep q_out.hidden_dim {} != q.hidden_dim {}", - q_out.hidden_dim, - q.hidden_dim - ); - anyhow::ensure!( - q_out.seq_len == batch_size, - "hd512 decode prep q_out.seq_len {} != q.seq_len {batch_size}", - q_out.seq_len - ); - anyhow::ensure!( - k.hidden_dim == num_kv_heads * 512, - "hd512 decode prep k.hidden_dim {} != num_kv_heads {} * 512", - k.hidden_dim, - num_kv_heads - ); - anyhow::ensure!( - k.seq_len == batch_size, - "hd512 decode prep k.seq_len {} != q.seq_len {batch_size}", - k.seq_len - ); - anyhow::ensure!( - q_norm_weight.len == 512, - "hd512 decode prep q_norm_weight len {} != 512", - q_norm_weight.len - ); - anyhow::ensure!( - k_norm_weight.len == 512, - "hd512 decode prep k_norm_weight len {} != 512", - k_norm_weight.len - ); - anyhow::ensure!( - positions_d.len() == batch_size, - "hd512 decode prep positions len {} != batch_size {batch_size}", - positions_d.len() - ); - anyhow::ensure!( - cos_cache.len >= cos_max_pos * rotary_dim, - "hd512 decode prep cos_cache len {} < cos_max_pos {cos_max_pos} * \ - rotary_dim {rotary_dim}", - cos_cache.len - ); - anyhow::ensure!( - sin_cache.len >= cos_max_pos * rotary_dim, - "hd512 decode prep sin_cache len {} < cos_max_pos {cos_max_pos} * \ - rotary_dim {rotary_dim}", - sin_cache.len - ); - q.checked_extent("hd512 decode prep q")?; - q_out.checked_extent("hd512 decode prep q_out")?; - k.checked_extent("hd512 decode prep k")?; - ensure_vec_backed(q_norm_weight, "hd512 decode prep q_norm_weight")?; - ensure_vec_backed(k_norm_weight, "hd512 decode prep k_norm_weight")?; - ensure_vec_backed(cos_cache, "hd512 decode prep cos_cache")?; - ensure_vec_backed(sin_cache, "hd512 decode prep sin_cache")?; - - let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); - let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); - let (k_ptr, _gk) = k.data.device_ptr_mut(&ctx.stream); - let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); - let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); - let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); - let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); - let (pos_ptr, _gp) = positions_d.device_ptr(&ctx.stream); - - let result = unsafe { - ffi::qk_norm_partial_rope_batched_decode_hd512_cuda( - q_ptr as *const ffi::Half, - k_ptr as *mut ffi::Half, - qn_ptr as *const ffi::Half, - kn_ptr as *const ffi::Half, - cos_ptr as *const ffi::Half, - sin_ptr as *const ffi::Half, - pos_ptr as *const i32, - cos_max_pos as i32, - qo_ptr as *mut ffi::Half, - num_q_heads as i32, - num_kv_heads as i32, - batch_size as i32, - rotary_dim as i32, - rms_eps, - crate::tensor::active_cu_stream(ctx), - ) - }; - if result != 0 { - anyhow::bail!( - "qk_norm_partial_rope_batched_decode_hd512_cuda failed with error \ - {result}{}", - crate::ops::ffi_exception_message(result) - ); - } - Ok(()) + Ok(PagedGeometry { + k_offset_elems: i64::try_from(k_offset) + .map_err(|_| anyhow::anyhow!("{what} K offset {k_offset} does not fit i64"))?, + v_offset_elems: i64::try_from(v_offset) + .map_err(|_| anyhow::anyhow!("{what} V offset {v_offset} does not fit i64"))?, + page_size: crate::ops::checked_i32(layout.page_size, "paged prep page_size")?, + num_pages: crate::ops::checked_i32(num_pages, "paged prep num_pages")?, + stride_page: i64::try_from(page_stride) + .map_err(|_| anyhow::anyhow!("{what} page_stride {page_stride} does not fit i64"))?, + }) } diff --git a/pegainfer-kernels/src/ops/deepseek_v2_lite.rs b/pegainfer-kernels/src/ops/deepseek_v2_lite.rs index 2e6d0763a..6cfd97e39 100644 --- a/pegainfer-kernels/src/ops/deepseek_v2_lite.rs +++ b/pegainfer-kernels/src/ops/deepseek_v2_lite.rs @@ -66,7 +66,7 @@ pub struct Dsv2LiteRouterOutput<'a> { pub fn dsv2_lite_router_logits_into( ctx: &DeviceContext, - hidden: &HiddenStates, + hidden: HiddenStatesRef<'_>, gate_weight: &DeviceMatrix, logits: &mut CudaSlice, ) -> Result<()> { diff --git a/pegainfer-kernels/src/ops/dsv2_routed_moe.rs b/pegainfer-kernels/src/ops/dsv2_routed_moe.rs new file mode 100644 index 000000000..6e6e9c518 --- /dev/null +++ b/pegainfer-kernels/src/ops/dsv2_routed_moe.rs @@ -0,0 +1,352 @@ +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use half::bf16; + +use crate::ffi; +use crate::ops::ffi_status_detail; +use crate::tensor::DeviceContext; +use crate::tensor::DeviceMatrix; +use crate::tensor::HiddenStates; +use crate::tensor::HiddenStatesRef; + +const ROUTES_PER_TOKEN: usize = 6; +const ROUTED_EXPERTS: usize = 64; +const LOCAL_EXPERTS: usize = ROUTED_EXPERTS / 2; +pub const DSV2_ROUTED_MOE_MAX_ROWS: usize = 8; + +pub struct Dsv2ExpertPointerTable { + first_expert: i32, + hidden_dim: usize, + intermediate: usize, + gate_up: CudaSlice, + down: CudaSlice, +} + +impl Dsv2ExpertPointerTable { + pub fn new( + ctx: &DeviceContext, + experts: Vec<(&DeviceMatrix, &DeviceMatrix)>, + first_expert: usize, + ) -> Result { + ensure!( + first_expert == 0 || first_expert == LOCAL_EXPERTS, + "invalid DeepSeek-V2-Lite EP2 expert base {first_expert}" + ); + ensure!( + experts.len() == LOCAL_EXPERTS, + "DeepSeek-V2-Lite routed MoE requires {LOCAL_EXPERTS} local experts, got {}", + experts.len() + ); + let hidden_dim = experts[0].0.cols; + let intermediate = experts[0].1.cols; + ensure!(hidden_dim > 0 && intermediate > 0, "empty expert matrix"); + + let mut gate_up_ptrs = Vec::with_capacity(LOCAL_EXPERTS); + let mut down_ptrs = Vec::with_capacity(LOCAL_EXPERTS); + for (gate_up, down) in experts { + ensure!( + gate_up.rows == 2 * intermediate + && gate_up.cols == hidden_dim + && down.rows == hidden_dim + && down.cols == intermediate, + "inconsistent DeepSeek-V2-Lite expert projection shape" + ); + gate_up_ptrs.push(gate_up.data.device_ptr(&ctx.stream).0); + down_ptrs.push(down.data.device_ptr(&ctx.stream).0); + } + + Ok(Self { + first_expert: first_expert as i32, + hidden_dim, + intermediate, + gate_up: ctx.stream.clone_htod(&gate_up_ptrs)?, + down: ctx.stream.clone_htod(&down_ptrs)?, + }) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct Dsv2RouteSummary { + pub local_routes: usize, + pub total_routes: usize, + pub errors: usize, + pub layer_hashes: Vec, +} + +pub struct Dsv2RoutedMoeScratch { + capacity: usize, + hidden_dim: usize, + intermediate: usize, + logits: CudaSlice, + ids: CudaSlice, + weights: CudaSlice, + errors: CudaSlice, + pointers: [CudaSlice; 6], + zero: CudaSlice, + gate_up: HiddenStates, + activation: HiddenStates, + rows: HiddenStates, + summary: CudaSlice, +} + +impl Dsv2RoutedMoeScratch { + pub fn new( + ctx: &DeviceContext, + hidden_dim: usize, + intermediate: usize, + capacity: usize, + num_layers: usize, + ) -> Result { + ensure!( + (1..=DSV2_ROUTED_MOE_MAX_ROWS).contains(&capacity), + "DeepSeek-V2-Lite routed MoE capacity must be in 1..={DSV2_ROUTED_MOE_MAX_ROWS}" + ); + ensure!(hidden_dim > 0 && intermediate > 0, "empty routed MoE shape"); + ensure!(num_layers > 0, "routed MoE requires at least one layer"); + for dim in [hidden_dim, intermediate] { + ensure!( + dim <= i32::MAX as usize / (2 * ROUTES_PER_TOKEN * capacity), + "routed MoE dimension exceeds CUDA indexing range" + ); + } + let routes = capacity * ROUTES_PER_TOKEN; + Ok(Self { + capacity, + hidden_dim, + intermediate, + logits: ctx.stream.alloc_zeros(capacity * ROUTED_EXPERTS)?, + ids: ctx.stream.alloc_zeros(routes)?, + weights: ctx.stream.alloc_zeros(routes)?, + errors: ctx.stream.alloc_zeros(capacity)?, + pointers: [ + ctx.stream.alloc_zeros(routes)?, + ctx.stream.alloc_zeros(routes)?, + ctx.stream.alloc_zeros(routes)?, + ctx.stream.alloc_zeros(routes)?, + ctx.stream.alloc_zeros(routes)?, + ctx.stream.alloc_zeros(routes)?, + ], + zero: ctx.stream.alloc_zeros(hidden_dim)?, + gate_up: HiddenStates::zeros(ctx, 2 * intermediate, routes)?, + activation: HiddenStates::zeros(ctx, intermediate, routes)?, + rows: HiddenStates::zeros(ctx, hidden_dim, routes)?, + summary: ctx.stream.alloc_zeros(3 + num_layers)?, + }) + } + + pub fn begin_forward(&mut self, ctx: &DeviceContext) -> Result<()> { + ensure!( + Arc::ptr_eq(self.summary.stream(), &ctx.stream), + "routed MoE summary stream mismatch" + ); + ctx.stream.memset_zeros(&mut self.summary)?; + Ok(()) + } + + pub fn finish_forward(&self, ctx: &DeviceContext) -> Result { + let summary = ctx.stream.clone_dtoh(&self.summary)?; + ctx.sync()?; + Ok(Dsv2RouteSummary { + local_routes: usize::try_from(summary[0]).context("local route count overflow")?, + total_routes: usize::try_from(summary[1]).context("total route count overflow")?, + errors: usize::try_from(summary[2]).context("route error count overflow")?, + layer_hashes: summary[3..].to_vec(), + }) + } + + pub fn enqueue_into( + &mut self, + ctx: &DeviceContext, + hidden: HiddenStatesRef<'_>, + gate_weight: &DeviceMatrix, + experts: &Dsv2ExpertPointerTable, + output: &mut CudaSlice, + layer_idx: usize, + ) -> Result<()> { + let batch = hidden.seq_len; + ensure!( + (1..=self.capacity).contains(&batch), + "routed MoE rows {batch} exceed capacity {}", + self.capacity + ); + ensure!( + hidden.hidden_dim == self.hidden_dim + && experts.hidden_dim == self.hidden_dim + && experts.intermediate == self.intermediate, + "routed MoE hidden or expert shape mismatch" + ); + ensure!( + gate_weight.rows == ROUTED_EXPERTS && gate_weight.cols == self.hidden_dim, + "routed MoE gate shape mismatch" + ); + ensure!( + hidden.data.len() >= batch * self.hidden_dim && output.len() >= batch * self.hidden_dim, + "routed MoE input or output backing buffer too small" + ); + ensure!( + Arc::ptr_eq(hidden.data.stream(), &ctx.stream) + && Arc::ptr_eq(gate_weight.data.stream(), &ctx.stream) + && Arc::ptr_eq(output.stream(), &ctx.stream) + && Arc::ptr_eq(experts.gate_up.stream(), &ctx.stream), + "routed MoE stream mismatch" + ); + + super::dsv2_lite_router_logits_into(ctx, hidden, gate_weight, &mut self.logits)?; + dsv2_lite_route_logits_into( + ctx, + &self.logits, + batch, + layer_idx, + &mut super::Dsv2LiteRouterOutput { + topk_weight: &mut self.weights, + topk_idx: &mut self.ids, + }, + &mut self.errors, + &mut self.summary, + )?; + let stream = ctx.stream.cu_stream(); + let (ids, _gi) = self.ids.device_ptr(&ctx.stream); + let (weights, _gw) = self.weights.device_ptr(&ctx.stream); + let (errors, _ge) = self.errors.device_ptr(&ctx.stream); + let (summary, _gs) = self.summary.device_ptr_mut(&ctx.stream); + + let (input, _gx) = hidden.data.device_ptr(&ctx.stream); + let (zero, _gz) = self.zero.device_ptr(&ctx.stream); + let (gate_up_weights, _g13) = experts.gate_up.device_ptr(&ctx.stream); + let (down_weights, _g2) = experts.down.device_ptr(&ctx.stream); + let (gate_up, gate_up_guard) = self.gate_up.data.device_ptr_mut(&ctx.stream); + let (activation, activation_guard) = self.activation.data.device_ptr_mut(&ctx.stream); + let (rows, rows_guard) = self.rows.data.device_ptr_mut(&ctx.stream); + let [a13, x13, y13, a2, x2, y2] = &mut self.pointers; + let (a13, _pa13) = a13.device_ptr_mut(&ctx.stream); + let (x13, _px13) = x13.device_ptr_mut(&ctx.stream); + let (y13, _py13) = y13.device_ptr_mut(&ctx.stream); + let (a2, _pa2) = a2.device_ptr_mut(&ctx.stream); + let (x2, _px2) = x2.device_ptr_mut(&ctx.stream); + let (y2, _py2) = y2.device_ptr_mut(&ctx.stream); + let routes = batch * ROUTES_PER_TOKEN; + unsafe { + ffi::dsv2_lite_route_pointers_cuda( + ids as _, + input as _, + zero as _, + gate_up_weights as _, + down_weights as _, + gate_up as _, + activation as _, + rows as _, + a13 as _, + x13 as _, + y13 as _, + a2 as _, + x2 as _, + y2 as _, + summary as _, + batch as i32, + experts.first_expert, + self.hidden_dim as i32, + self.intermediate as i32, + stream, + ) + .result()?; + let code = ffi::dsv2_lite_pointer_gemm_cuda( + a13 as _, + x13 as _, + y13 as _, + (2 * self.intermediate) as i32, + self.hidden_dim as i32, + routes as i32, + stream, + ); + ensure!( + code == 0, + "routed MoE gate/up pointer GEMM failed: {}", + ffi_status_detail(code) + ); + } + drop((gate_up_guard, activation_guard, rows_guard)); + super::silu_mul_fused_batch_into(ctx, &self.gate_up, &mut self.activation)?; + + let (output, _go) = output.device_ptr_mut(&ctx.stream); + unsafe { + let code = ffi::dsv2_lite_pointer_gemm_cuda( + a2 as _, + x2 as _, + y2 as _, + self.hidden_dim as i32, + self.intermediate as i32, + routes as i32, + stream, + ); + ensure!( + code == 0, + "routed MoE down pointer GEMM failed: {}", + ffi_status_detail(code) + ); + ffi::dsv2_lite_route_reduce_cuda( + rows as _, + ids as _, + weights as _, + errors as _, + output as _, + batch as i32, + experts.first_expert, + self.hidden_dim as i32, + stream, + ) + .result()?; + } + Ok(()) + } +} + +pub fn dsv2_lite_route_logits_into( + ctx: &DeviceContext, + logits: &CudaSlice, + batch: usize, + layer_idx: usize, + output: &mut super::Dsv2LiteRouterOutput<'_>, + errors: &mut CudaSlice, + summary: &mut CudaSlice, +) -> Result<()> { + ensure!( + (1..=DSV2_ROUTED_MOE_MAX_ROWS).contains(&batch), + "invalid routed MoE batch" + ); + ensure!( + summary.len() > 3 && layer_idx < summary.len() - 3 && layer_idx <= i32::MAX as usize, + "routed MoE layer exceeds summary capacity" + ); + ensure!( + logits.len() >= batch * ROUTED_EXPERTS + && output.topk_idx.len() >= batch * ROUTES_PER_TOKEN + && output.topk_weight.len() >= batch * ROUTES_PER_TOKEN + && errors.len() >= batch, + "routed MoE router buffer too small" + ); + let (logits, _gl) = logits.device_ptr(&ctx.stream); + let (ids, _gi) = output.topk_idx.device_ptr_mut(&ctx.stream); + let (weights, _gw) = output.topk_weight.device_ptr_mut(&ctx.stream); + let (errors, _ge) = errors.device_ptr_mut(&ctx.stream); + let (summary, _gs) = summary.device_ptr_mut(&ctx.stream); + unsafe { + ffi::dsv2_lite_route_logits_cuda( + logits as _, + weights as _, + ids as _, + errors as _, + summary as _, + batch as i32, + layer_idx as i32, + ctx.stream.cu_stream(), + ) + .result()?; + } + Ok(()) +} diff --git a/pegainfer-kernels/src/ops/gemma4_attention.rs b/pegainfer-kernels/src/ops/gemma4_attention.rs new file mode 100644 index 000000000..a6eb8fc2d --- /dev/null +++ b/pegainfer-kernels/src/ops/gemma4_attention.rs @@ -0,0 +1,1635 @@ +//! Gemma 4 attention wrappers: the global family at head dim 512 (split-KV +//! decode, paged prefill, QK prep) and the sliding family at head dim 256 +//! (windowed paged prefill, plain-w QKV prep). + +use anyhow::Result; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use half::bf16; + +use super::PrefillPagedPlan; +use super::attention::PagedGeometry; +use super::attention::checked_paged_geometry; +use super::attention::checked_row_offset; +use crate::ffi; +use crate::paged_kv::KvFormat; +use crate::paged_kv::KvStorage; +use crate::paged_kv::PagedKvLayout; +use crate::tensor::Columns; +use crate::tensor::DeviceContext; +use crate::tensor::DeviceVec; +use crate::tensor::HiddenStates; + +/// Byte offset of row `offset`'s first column of the band `c`, after checking +/// the row window and that the band lies within the row. +fn checked_band_offset(c: &Columns<'_>, offset: usize, span: usize, name: &str) -> Result { + anyhow::ensure!( + c.col + .checked_add(c.width) + .is_some_and(|end| end <= c.states.hidden_dim), + "{name} columns [{}..+{}) exceed the row's {}", + c.col, + c.width, + c.states.hidden_dim + ); + let rows = checked_row_offset(c.states, offset, span, name)?; + Ok(rows + (c.col * std::mem::size_of::()) as u64) +} + +/// `DeviceVec` fields are public too: reject a logical `len` past the backing +/// allocation before a kernel indexes through it. +fn ensure_vec_backed(v: &DeviceVec, name: &str) -> Result<()> { + anyhow::ensure!( + v.data.len() >= v.len, + "{name} backing len {} < len {}", + v.data.len(), + v.len + ); + Ok(()) +} + +/// Decode-metadata aggregate for the hd512 split-KV decode path, validated +/// at use: the wrapper calls [`Self::validate`] with the batch size it +/// derives. +/// +/// `page_indptr` needs at least `batch_size + 1` elements and `last_page_len` +/// at least `batch_size`; `request_indices` and `kv_tile_indices` are +/// per-split-slot arrays sized by the caller's padded slot count, and +/// `kv_chunk_size` holds the one chunk-size entry the kernel reads. +/// `page_indices` must be non-empty (its per-request reach is device-side +/// data and cannot be verified here). +/// +/// The pages referenced by `page_indices` must lie within `kv_buffer`; the +/// wrapper cannot verify device-side contents — caller contract. +pub struct Hd512DecodeMetadata<'a> { + page_indices: &'a CudaSlice, + page_indptr: &'a CudaSlice, + last_page_len: &'a CudaSlice, + request_indices: &'a CudaSlice, + kv_tile_indices: &'a CudaSlice, + kv_chunk_size: &'a CudaSlice, + /// The chunk the tile indices count in, on the host as well: the device + /// slot holds the same number for the incumbent's kernel, and a launcher + /// that sizes its own grid from the tiles needs it where it can read it. + chunk_tokens: usize, +} + +impl<'a> Hd512DecodeMetadata<'a> { + pub fn new( + page_indices: &'a CudaSlice, + page_indptr: &'a CudaSlice, + last_page_len: &'a CudaSlice, + request_indices: &'a CudaSlice, + kv_tile_indices: &'a CudaSlice, + kv_chunk_size: &'a CudaSlice, + chunk_tokens: usize, + ) -> Self { + Self { + page_indices, + page_indptr, + last_page_len, + request_indices, + kv_tile_indices, + kv_chunk_size, + chunk_tokens, + } + } + + pub fn page_indices(&self) -> &'a CudaSlice { + self.page_indices + } + + pub fn page_indptr(&self) -> &'a CudaSlice { + self.page_indptr + } + + pub fn last_page_len(&self) -> &'a CudaSlice { + self.last_page_len + } + + pub fn request_indices(&self) -> &'a CudaSlice { + self.request_indices + } + + pub fn kv_tile_indices(&self) -> &'a CudaSlice { + self.kv_tile_indices + } + + pub fn chunk_tokens(&self) -> usize { + self.chunk_tokens + } + + pub fn validate(&self, batch_size: usize) -> anyhow::Result<()> { + anyhow::ensure!( + !self.page_indices.is_empty(), + "Hd512DecodeMetadata: page_indices must be non-empty" + ); + anyhow::ensure!( + self.page_indptr.len() > batch_size, + "Hd512DecodeMetadata: page_indptr length {} must be >= batch_size + 1 = {}", + self.page_indptr.len(), + batch_size + 1 + ); + for (name, len) in [ + ("last_page_len", self.last_page_len.len()), + ("request_indices", self.request_indices.len()), + ("kv_tile_indices", self.kv_tile_indices.len()), + ] { + anyhow::ensure!( + len >= batch_size, + "Hd512DecodeMetadata: {name} length {len} must be >= batch_size {batch_size}", + ); + } + // The kernel reads one chunk size for the whole launch. + anyhow::ensure!( + !self.kv_chunk_size.is_empty(), + "Hd512DecodeMetadata: kv_chunk_size must hold its one entry" + ); + anyhow::ensure!( + self.chunk_tokens > 0, + "Hd512DecodeMetadata: chunk_tokens must be positive" + ); + Ok(()) + } +} + +#[allow(clippy::too_many_arguments)] +pub fn paged_attention_batch_decode_split_kv_hd512_into( + ctx: &DeviceContext, + q: &HiddenStates, + row_offset: usize, + kv_buffer: &CudaSlice, + layout: &PagedKvLayout, + layer: usize, + meta: &Hd512DecodeMetadata, + split_o_indptr_d: &CudaSlice, + split_valid_mask_d: &CudaSlice, + split_tmp_v: &mut CudaSlice, + split_tmp_s: &mut CudaSlice, + split_padded_slots: usize, + output: &mut HiddenStates, + num_qo_heads: usize, + sm_scale: f32, +) -> Result<()> { + anyhow::ensure!( + sm_scale.is_finite(), + "paged_attention_batch_decode_hd512 sm_scale {sm_scale} must be finite" + ); + let num_kv_heads = layout.num_kv_heads; + let geometry = checked_paged_geometry( + "hd512 split decode", + layout, + kv_buffer.len(), + layer, + 512, + num_kv_heads, + false, + )?; + anyhow::ensure!( + row_offset < q.seq_len, + "hd512 split decode row_offset {row_offset} leaves no rows of {}", + q.seq_len + ); + let batch_size = q.seq_len - row_offset; + meta.validate(batch_size)?; + anyhow::ensure!( + output.seq_len == q.seq_len, + "hd512 decode output.seq_len {} != q.seq_len {}", + output.seq_len, + q.seq_len + ); + let qo_dim = num_qo_heads.checked_mul(512).ok_or_else(|| { + anyhow::anyhow!("hd512 split decode num_qo_heads {num_qo_heads} * 512 overflows") + })?; + anyhow::ensure!( + q.hidden_dim == qo_dim, + "hd512 decode q.hidden_dim {} != num_qo_heads {num_qo_heads} * 512", + q.hidden_dim + ); + anyhow::ensure!( + output.hidden_dim == qo_dim, + "hd512 decode output.hidden_dim {} != num_qo_heads {num_qo_heads} * 512", + output.hidden_dim + ); + anyhow::ensure!( + split_padded_slots >= batch_size, + "hd512 split decode padded_slots {split_padded_slots} < batch {batch_size}" + ); + anyhow::ensure!( + meta.request_indices.len() >= split_padded_slots + && meta.kv_tile_indices.len() >= split_padded_slots + && split_valid_mask_d.len() >= split_padded_slots, + "hd512 split decode plan arrays shorter than padded_slots {split_padded_slots}" + ); + anyhow::ensure!( + split_o_indptr_d.len() > batch_size, + "hd512 split decode o_indptr len {} < batch {batch_size} + 1", + split_o_indptr_d.len() + ); + let tmp_v_need = split_padded_slots.checked_mul(qo_dim).ok_or_else(|| { + anyhow::anyhow!("hd512 split decode padded_slots {split_padded_slots} * {qo_dim} overflows") + })?; + let tmp_s_need = split_padded_slots + .checked_mul(num_qo_heads) + .ok_or_else(|| { + anyhow::anyhow!( + "hd512 split decode padded_slots {split_padded_slots} * {num_qo_heads} overflows" + ) + })?; + anyhow::ensure!( + split_tmp_v.len() >= tmp_v_need && split_tmp_s.len() >= tmp_s_need, + "hd512 split decode workspace shorter than padded_slots {split_padded_slots}" + ); + let q_elems = q.checked_extent("batch_decode_hd512 q")?; + let out_elems = output.checked_extent("batch_decode_hd512 output")?; + crate::ops::checked_i32(q_elems, "hd512 split decode q extent")?; + crate::ops::checked_i32(out_elems, "hd512 split decode output extent")?; + let batch_i32 = crate::ops::checked_i32(batch_size, "hd512 split decode batch")?; + let padded_i32 = + crate::ops::checked_i32(split_padded_slots, "hd512 split decode padded slots")?; + let qo_heads_i32 = crate::ops::checked_i32(num_qo_heads, "hd512 split decode qo heads")?; + let kv_heads_i32 = crate::ops::checked_i32(num_kv_heads, "hd512 split decode kv heads")?; + + let q_row_bytes = checked_row_offset(q, row_offset, batch_size, "hd512 split decode q")?; + let out_row_bytes = + checked_row_offset(output, row_offset, batch_size, "hd512 split decode output")?; + let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); + let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); + let q_ptr = q_ptr + q_row_bytes; + let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); + let out_ptr = out_ptr + out_row_bytes; + let (pi_ptr, _gpi) = meta.page_indices.device_ptr(&ctx.stream); + let (pip_ptr, _gpip) = meta.page_indptr.device_ptr(&ctx.stream); + let (lpl_ptr, _glpl) = meta.last_page_len.device_ptr(&ctx.stream); + let (ri_ptr, _gri) = meta.request_indices.device_ptr(&ctx.stream); + let (kti_ptr, _gkti) = meta.kv_tile_indices.device_ptr(&ctx.stream); + let (kcs_ptr, _gkcs) = meta.kv_chunk_size.device_ptr(&ctx.stream); + let (soi_ptr, _gsoi) = split_o_indptr_d.device_ptr(&ctx.stream); + let (svm_ptr, _gsvm) = split_valid_mask_d.device_ptr(&ctx.stream); + let (stv_ptr, _gstv) = split_tmp_v.device_ptr_mut(&ctx.stream); + let (sts_ptr, _gsts) = split_tmp_s.device_ptr_mut(&ctx.stream); + + let stream = crate::tensor::active_cu_stream(ctx); + + // No KV scatter here: the serving prep kernel has already written this + // step's normed+roped K and its V fork; this only attends over the + // resident pages. + let result = unsafe { + ffi::paged_attention_decode_split_kv_cuda_hd512( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const ffi::Half, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + ri_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + soi_ptr as *const i32, + svm_ptr as *const u8, + stv_ptr as *mut ffi::Half, + sts_ptr as *mut f32, + qo_heads_i32, + kv_heads_i32, + 512, + geometry.page_size, + batch_i32, + padded_i32, + geometry.stride_page, + sm_scale, + stream, + ) + }; + if result != 0 { + anyhow::bail!( + "paged_attention_decode_split_kv_cuda_hd512 (batch) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +/// The pages referenced by `plan.page_indices_d()` must lie within +/// `kv_buffer`; the wrapper cannot verify device-side contents — caller +/// contract. +pub fn batch_prefill_paged_hd512_into( + ctx: &DeviceContext, + q: &HiddenStates, + kv_buffer: &CudaSlice, + layout: &PagedKvLayout, + layer: usize, + plan: &PrefillPagedPlan, + output: &mut HiddenStates, + num_qo_heads: usize, + sm_scale: f32, +) -> Result<()> { + anyhow::ensure!( + sm_scale.is_finite(), + "batch_prefill_paged_hd512 sm_scale {sm_scale} must be finite" + ); + let num_kv_heads = layout.num_kv_heads; + let head_dim = layout.head_dim; + anyhow::ensure!( + head_dim == 512, + "hd512 prefill expects head_dim 512, got {}", + head_dim + ); + anyhow::ensure!( + q.hidden_dim == num_qo_heads * 512, + "hd512 prefill q.hidden_dim {} != num_qo_heads {} * 512", + q.hidden_dim, + num_qo_heads + ); + anyhow::ensure!( + output.hidden_dim == num_qo_heads * 512, + "hd512 prefill output.hidden_dim {} != num_qo_heads {} * 512", + output.hidden_dim, + num_qo_heads + ); + anyhow::ensure!( + q.seq_len >= plan.total_tokens, + "hd512 prefill q.seq_len {} < plan.total_tokens {}", + q.seq_len, + plan.total_tokens + ); + anyhow::ensure!( + output.seq_len >= plan.total_tokens, + "hd512 prefill output.seq_len {} < plan.total_tokens {}", + output.seq_len, + plan.total_tokens + ); + anyhow::ensure!( + layer < layout.num_layers, + "hd512 prefill layer {layer} >= layout.num_layers {}", + layout.num_layers + ); + q.checked_extent("batch_prefill_paged_hd512 q")?; + output.checked_extent("batch_prefill_paged_hd512 output")?; + anyhow::ensure!( + plan.head_dim == 512, + "plan built for head_dim {}, expected 512", + plan.head_dim + ); + anyhow::ensure!( + plan.num_kv_heads == num_kv_heads, + "plan built for num_kv_heads {}, expected {}", + plan.num_kv_heads, + num_kv_heads + ); + anyhow::ensure!( + plan.num_qo_heads == num_qo_heads, + "plan built for num_qo_heads {}, expected {}", + plan.num_qo_heads, + num_qo_heads + ); + + // Plan tiles laid out for a cta_tile_q override would be misread by the + // kernel's automatic derivation. + let kernel_cta_tile_q = unsafe { + ffi::batch_prefill_cta_tile_q_with_override( + plan.total_tokens as i32, + plan.num_qo_heads as i32, + plan.num_kv_heads as i32, + plan.head_dim as i32, + 0, // auto-detect, matching the hd512 kernel + ) + }; + anyhow::ensure!( + kernel_cta_tile_q == plan.cta_tile_q(), + "hd512 prefill cta_tile_q: plan recorded {}, kernel derives {}; \ + overridden plans are not supported at hd512", + plan.cta_tile_q(), + kernel_cta_tile_q + ); + + let PagedGeometry { + k_offset_elems: k_offset, + v_offset_elems: v_offset, + stride_page, + .. + } = checked_paged_geometry( + "hd512 prefill", + layout, + kv_buffer.len(), + layer, + head_dim, + num_kv_heads, + false, + )?; + + let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); + let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); + let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); + let (pi_ptr, _gpi) = plan.page_indices_d.device_ptr(&ctx.stream); + let (pip_ptr, _gpip) = plan.page_indptr_d.device_ptr(&ctx.stream); + let (lpl_ptr, _glpl) = plan.last_page_len_d.device_ptr(&ctx.stream); + let (qi_ptr, _gqi) = plan.q_indptr_d.device_ptr(&ctx.stream); + let (ri_ptr, _gri) = plan.request_indices_d.device_ptr(&ctx.stream); + let (qti_ptr, _gqti) = plan.qo_tile_indices_d.device_ptr(&ctx.stream); + let (kti_ptr, _gkti) = plan.kv_tile_indices_d.device_ptr(&ctx.stream); + let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); + let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); + + let result = unsafe { + ffi::batch_prefill_paged_cuda_hd512( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const ffi::Half, + k_offset, + v_offset, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + qi_ptr as *const i32, + ri_ptr as *const i32, + qti_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + tnr_ptr as *const u32, + num_qo_heads as i32, + num_kv_heads as i32, + head_dim as i32, + layout.page_size as i32, + plan.total_tokens as i32, + plan.batch_size(), + plan.num_tiles, + stride_page, + sm_scale, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "batch_prefill_paged_cuda_hd512 (prefill) failed for layer {layer}, \ + bs={}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}: {result}{}", + plan.batch_size(), + plan.num_tiles, + crate::ops::ffi_exception_message(result) + ); + } + + Ok(()) +} + +/// Windowed batch prefill over paged KV at head_dim 256 (Gemma 4 local +/// layers): attention read only — the prep kernel has already written K/V +/// into the pool. `window_left` is an inclusive distance: an N-token window +/// passes N - 1, and -1 degrades to full attention. `sm_scale` is the +/// caller's; Gemma 4 runs unscaled attention (1.0). +#[allow(clippy::too_many_arguments)] +pub fn batch_prefill_paged_window_hd256_into( + ctx: &DeviceContext, + q: &HiddenStates, + kv_buffer: &CudaSlice, + layout: &PagedKvLayout, + layer: usize, + plan: &PrefillPagedPlan, + output: &mut HiddenStates, + num_qo_heads: usize, + sm_scale: f32, + window_left: i32, +) -> Result<()> { + anyhow::ensure!( + sm_scale.is_finite(), + "batch_prefill_paged_window_hd256 sm_scale {sm_scale} must be finite" + ); + anyhow::ensure!( + window_left >= -1, + "batch_prefill_paged_window_hd256 window_left {window_left} must be >= -1" + ); + let num_kv_heads = layout.num_kv_heads; + let geometry = checked_paged_geometry( + "hd256 window prefill", + layout, + kv_buffer.len(), + layer, + 256, + num_kv_heads, + true, + )?; + anyhow::ensure!( + plan.max_page_index < geometry.num_pages, + "hd256 window prefill plan references page {} but the pool holds {} pages; the attention \ + kernel computes addresses from page ids with no device-side bounds check", + plan.max_page_index, + geometry.num_pages + ); + anyhow::ensure!( + plan.max_last_page_len <= geometry.page_size, + "hd256 window prefill plan last-page length {} exceeds page_size {}; the kernel would \ + read a page-table entry that does not exist", + plan.max_last_page_len, + geometry.page_size + ); + let qo_dim = num_qo_heads.checked_mul(256).ok_or_else(|| { + anyhow::anyhow!("hd256 window prefill num_qo_heads {num_qo_heads} * 256 overflows") + })?; + anyhow::ensure!( + q.hidden_dim == qo_dim, + "hd256 window prefill q.hidden_dim {} != num_qo_heads {num_qo_heads} * 256", + q.hidden_dim + ); + anyhow::ensure!( + output.hidden_dim == qo_dim, + "hd256 window prefill output.hidden_dim {} != num_qo_heads {num_qo_heads} * 256", + output.hidden_dim + ); + anyhow::ensure!( + q.seq_len >= plan.total_tokens, + "hd256 window prefill q.seq_len {} < plan.total_tokens {}", + q.seq_len, + plan.total_tokens + ); + anyhow::ensure!( + output.seq_len >= plan.total_tokens, + "hd256 window prefill output.seq_len {} < plan.total_tokens {}", + output.seq_len, + plan.total_tokens + ); + let q_elems = q.checked_extent("batch_prefill_paged_window_hd256 q")?; + let out_elems = output.checked_extent("batch_prefill_paged_window_hd256 output")?; + crate::ops::checked_i32(q_elems, "hd256 window prefill q extent")?; + crate::ops::checked_i32(out_elems, "hd256 window prefill output extent")?; + anyhow::ensure!( + plan.head_dim == 256, + "plan built for head_dim {}, expected 256", + plan.head_dim + ); + anyhow::ensure!( + plan.num_kv_heads == num_kv_heads, + "plan built for num_kv_heads {}, expected {}", + plan.num_kv_heads, + num_kv_heads + ); + anyhow::ensure!( + plan.num_qo_heads == num_qo_heads, + "plan built for num_qo_heads {}, expected {}", + plan.num_qo_heads, + num_qo_heads + ); + + // Plan tiles laid out for a cta_tile_q override would be misread by the + // kernel's automatic derivation. + let num_qo_heads_i32 = + crate::ops::checked_i32(num_qo_heads, "hd256 window prefill num_qo_heads")?; + let num_kv_heads_i32 = + crate::ops::checked_i32(num_kv_heads, "hd256 window prefill num_kv_heads")?; + let total_tokens_i32 = + crate::ops::checked_i32(plan.total_tokens, "hd256 window prefill total_tokens")?; + let kernel_cta_tile_q = unsafe { + ffi::batch_prefill_cta_tile_q_with_override( + total_tokens_i32, + num_qo_heads_i32, + num_kv_heads_i32, + 256, + 0, // auto-detect, matching the hd256 kernel + ) + }; + anyhow::ensure!( + kernel_cta_tile_q == plan.cta_tile_q(), + "hd256 window prefill cta_tile_q: plan recorded {}, kernel derives {}; \ + overridden plans are not supported here", + plan.cta_tile_q(), + kernel_cta_tile_q + ); + + let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); + let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); + let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); + let (pi_ptr, _gpi) = plan.page_indices_d.device_ptr(&ctx.stream); + let (pip_ptr, _gpip) = plan.page_indptr_d.device_ptr(&ctx.stream); + let (lpl_ptr, _glpl) = plan.last_page_len_d.device_ptr(&ctx.stream); + let (qi_ptr, _gqi) = plan.q_indptr_d.device_ptr(&ctx.stream); + let (ri_ptr, _gri) = plan.request_indices_d.device_ptr(&ctx.stream); + let (qti_ptr, _gqti) = plan.qo_tile_indices_d.device_ptr(&ctx.stream); + let (kti_ptr, _gkti) = plan.kv_tile_indices_d.device_ptr(&ctx.stream); + let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); + let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); + + let (result, entry_point) = match layout.storage { + KvStorage::E4m3 => { + let result = unsafe { + ffi::gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const core::ffi::c_void, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + qi_ptr as *const i32, + ri_ptr as *const i32, + qti_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + tnr_ptr as *const u32, + num_qo_heads_i32, + num_kv_heads_i32, + 256, + geometry.page_size, + total_tokens_i32, + plan.batch_size(), + plan.num_tiles, + geometry.stride_page, + sm_scale, + 0, + window_left, + crate::tensor::active_cu_stream(ctx), + ) + }; + (result, "gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda") + } + KvStorage::Bf16 => ( + unsafe { + ffi::batch_prefill_paged_window_cuda_hd256( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const ffi::Half, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + qi_ptr as *const i32, + ri_ptr as *const i32, + qti_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + tnr_ptr as *const u32, + num_qo_heads_i32, + num_kv_heads_i32, + 256, + geometry.page_size, + total_tokens_i32, + plan.batch_size(), + plan.num_tiles, + geometry.stride_page, + sm_scale, + window_left, + crate::tensor::active_cu_stream(ctx), + ) + }, + "batch_prefill_paged_window_cuda_hd256", + ), + }; + if result != 0 { + anyhow::bail!( + "{entry_point} failed for layer {layer}, \ + bs={}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}, \ + window_left={window_left}: {result}{}", + plan.batch_size(), + plan.num_tiles, + crate::ops::ffi_exception_message(result) + ); + } + + Ok(()) +} + +/// What the hd512 prep writes into, from the layout's format: the layer's +/// K and V blocks and the row's bands. For the split format the blocks are +/// the layer's two and the row is the head; for the folded one both offsets +/// name the layer's single block, the row is `head_dim + rotary` wide and +/// `fold_rotary` is the rotated columns the kernel keeps of K. +struct Hd512PrepTarget { + k_offset_elems: i64, + v_offset_elems: i64, + row_width: i32, + fold_rotary: i32, + page_size: i32, + num_pages: i32, + stride_page: i64, +} + +fn hd512_prep_target( + what: &str, + layout: &PagedKvLayout, + pool_len: usize, + layer: usize, + num_kv_heads: usize, +) -> Result { + anyhow::ensure!( + layout.storage == KvStorage::Bf16, + "{what} has no fp8 KV path; the layout carries {}-byte elements", + layout.storage.elem_bytes() + ); + anyhow::ensure!( + layout.head_dim == 512, + "{what} layout.head_dim {} != 512", + layout.head_dim + ); + anyhow::ensure!( + layout.num_kv_heads == num_kv_heads, + "{what} layout.num_kv_heads {} != num_kv_heads {num_kv_heads}", + layout.num_kv_heads + ); + let block = layout.block_geometry(pool_len, layer, what)?; + let v_offset = match layout.format { + KvFormat::Split => block + .block_offset_elems + .checked_add(layout.kv_block_len) + .ok_or_else(|| { + anyhow::anyhow!( + "{what} K offset {} + kv_block_len {} overflows", + block.block_offset_elems, + layout.kv_block_len + ) + })?, + KvFormat::Folded { .. } => block.block_offset_elems, + }; + let to_i64 = |value: usize, name: &str| { + i64::try_from(value).map_err(|_| anyhow::anyhow!("{what} {name} {value} does not fit i64")) + }; + Ok(Hd512PrepTarget { + k_offset_elems: to_i64(block.block_offset_elems, "K offset")?, + v_offset_elems: to_i64(v_offset, "V offset")?, + row_width: crate::ops::checked_i32( + layout.format.row_width(layout.head_dim), + &format!("{what} row_width"), + )?, + fold_rotary: crate::ops::checked_i32( + layout.format.fold_rotary(), + &format!("{what} fold_rotary"), + )?, + page_size: crate::ops::checked_i32(block.page_size, &format!("{what} page_size"))?, + num_pages: crate::ops::checked_i32(block.num_pages, &format!("{what} num_pages"))?, + stride_page: to_i64(block.stride_page, "page_stride")?, + }) +} + +/// Plain-w QKV prep at head_dim 256 (Gemma 4 local layers). +/// Q is normalised + rotated into a contiguous `q_out`; K is normalised + +/// rotated straight into the paged KV pool at layer `layer`'s K block, and +/// V — a separate v_proj head vector, unlike the hd512 K=V fork — is +/// weightless-normalised (never rotated) into the layer's V block, all in +/// one kernel with no intermediate scatter. `q`, `k` and `v` are column +/// bands read at their own row strides, so they may share one fused Q|K|V +/// projection row. +/// +/// The kernel `__trap()`s on any out-of-range pos or page id as the second +/// layer of the host validation's defence. `page_indices` is the resident +/// page row and `page_origin` the absolute page its first entry covers, so +/// a caller that released the front passes the released count; the origin is +/// page-aligned, which keeps in-page offsets position-invariant and shifts +/// only the row index. RoPE still runs on absolute positions. +#[allow(clippy::too_many_arguments)] +pub fn qkv_norm_rope_paged_prefill_hd256_plain_into<'a>( + ctx: &DeviceContext, + q: impl Into>, + k: impl Into>, + v: impl Into>, + q_out: &mut HiddenStates, + row_offset: usize, + kv_pool: &CudaSlice, + layout: &PagedKvLayout, + q_norm_weight: &DeviceVec, + k_norm_weight: &DeviceVec, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + layer: usize, + page_indices: &CudaSlice, + pages_offset: usize, + page_origin: usize, + start_pos: usize, + cos_max_pos: usize, + num_q_heads: usize, + num_kv_heads: usize, + rotary_dim: usize, + rms_eps: f32, +) -> Result<()> { + let (q, k, v) = (q.into(), k.into(), v.into()); + // The prompt segment is the row suffix `[row_offset..seq_len)` with its + // page table `pages_offset` elements into the (possibly concatenated) + // table — a multi-prompt mixed step parks earlier prompts and their + // tables in the prefixes. + anyhow::ensure!( + row_offset < q.states.seq_len, + "hd256 paged prep row_offset {row_offset} leaves no rows of {}", + q.states.seq_len + ); + let seq_len = q.states.seq_len - row_offset; + anyhow::ensure!( + pages_offset < page_indices.len(), + "hd256 paged prep pages_offset {pages_offset} exceeds table len {}", + page_indices.len() + ); + let q_dim = num_q_heads.checked_mul(256).ok_or_else(|| { + anyhow::anyhow!("hd256 paged prep num_q_heads {num_q_heads} * 256 overflows") + })?; + let kv_dim = num_kv_heads.checked_mul(256).ok_or_else(|| { + anyhow::anyhow!("hd256 paged prep num_kv_heads {num_kv_heads} * 256 overflows") + })?; + anyhow::ensure!( + q.width == q_dim, + "hd256 paged prep q width {} != num_q_heads {num_q_heads} * 256", + q.width + ); + anyhow::ensure!( + q_out.hidden_dim == q.width, + "hd256 paged prep q_out.hidden_dim {} != q width {}", + q_out.hidden_dim, + q.width + ); + anyhow::ensure!( + q_out.seq_len == q.states.seq_len, + "hd256 paged prep q_out.seq_len {} != q rows {}", + q_out.seq_len, + q.states.seq_len + ); + anyhow::ensure!( + k.width == kv_dim, + "hd256 paged prep k width {} != num_kv_heads {num_kv_heads} * 256", + k.width + ); + anyhow::ensure!( + v.width == k.width, + "hd256 paged prep v width {} != k width {}", + v.width, + k.width + ); + anyhow::ensure!( + k.states.seq_len == q.states.seq_len, + "hd256 paged prep k rows {} != q rows {}", + k.states.seq_len, + q.states.seq_len + ); + anyhow::ensure!( + v.states.seq_len == q.states.seq_len, + "hd256 paged prep v rows {} != q rows {}", + v.states.seq_len, + q.states.seq_len + ); + let geometry = checked_paged_geometry( + "hd256 paged prep", + layout, + kv_pool.len(), + layer, + 256, + num_kv_heads, + true, + )?; + ensure_vec_backed(q_norm_weight, "hd256 paged prep q_norm_weight")?; + ensure_vec_backed(k_norm_weight, "hd256 paged prep k_norm_weight")?; + ensure_vec_backed(cos_cache, "hd256 paged prep cos_cache")?; + ensure_vec_backed(sin_cache, "hd256 paged prep sin_cache")?; + anyhow::ensure!( + q_norm_weight.len == 256, + "hd256 paged prep q_norm_weight len {} != 256", + q_norm_weight.len + ); + anyhow::ensure!( + k_norm_weight.len == 256, + "hd256 paged prep k_norm_weight len {} != 256", + k_norm_weight.len + ); + let end_pos = start_pos.checked_add(seq_len).ok_or_else(|| { + anyhow::anyhow!("hd256 paged prep start_pos {start_pos} + seq_len {seq_len} overflows") + })?; + let covered = (page_indices.len() - pages_offset) + .checked_mul(layout.page_size) + .ok_or_else(|| { + anyhow::anyhow!( + "hd256 paged prep page window {} * page_size {} overflows", + page_indices.len() - pages_offset, + layout.page_size + ) + })?; + let row_start = page_origin.checked_mul(layout.page_size).ok_or_else(|| { + anyhow::anyhow!( + "hd256 paged prep page_origin {page_origin} * page_size {} overflows", + layout.page_size + ) + })?; + let row_end = row_start.checked_add(covered).ok_or_else(|| { + anyhow::anyhow!("hd256 paged prep row start {row_start} + span {covered} overflows") + })?; + anyhow::ensure!( + row_start <= start_pos, + "hd256 paged prep page_origin {page_origin} starts after start_pos {start_pos}" + ); + anyhow::ensure!( + row_end >= end_pos, + "hd256 paged prep row covers tokens {row_start}..{row_end}, need start_pos \ + {start_pos} + seq_len {seq_len}" + ); + anyhow::ensure!( + end_pos <= cos_max_pos, + "hd256 paged prep start_pos {start_pos} + seq_len {seq_len} > cos_max_pos {cos_max_pos}" + ); + let table_len = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { + anyhow::anyhow!( + "hd256 paged prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" + ) + })?; + anyhow::ensure!( + cos_cache.len >= table_len, + "hd256 paged prep cos_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", + cos_cache.len + ); + anyhow::ensure!( + sin_cache.len >= table_len, + "hd256 paged prep sin_cache len {} < cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim}", + sin_cache.len + ); + let q_elems = q.states.checked_extent("hd256 paged prep q")?; + q_out.checked_extent("hd256 paged prep q_out")?; + let k_elems = k.states.checked_extent("hd256 paged prep k")?; + v.states.checked_extent("hd256 paged prep v")?; + crate::ops::checked_i32(q_elems, "hd256 paged prep q extent")?; + crate::ops::checked_i32(k_elems, "hd256 paged prep k extent")?; + crate::ops::checked_i32(table_len, "hd256 paged prep rope table extent")?; + + let page_origin_i32 = crate::ops::checked_i32(page_origin, "hd256 paged prep page_origin")?; + let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd256 paged prep num_q_heads")?; + let num_kv_heads_i32 = crate::ops::checked_i32(num_kv_heads, "hd256 paged prep num_kv_heads")?; + let seq_len_i32 = crate::ops::checked_i32(seq_len, "hd256 paged prep seq_len")?; + let start_pos_i32 = crate::ops::checked_i32(start_pos, "hd256 paged prep start_pos")?; + let cos_max_pos_i32 = crate::ops::checked_i32(cos_max_pos, "hd256 paged prep cos_max_pos")?; + let rotary_dim_i32 = crate::ops::checked_i32(rotary_dim, "hd256 paged prep rotary_dim")?; + + let page_indices_len = crate::ops::checked_i32( + page_indices.len() - pages_offset, + "hd256 paged prep page window len", + )?; + let q_row_bytes = checked_band_offset(&q, row_offset, seq_len, "hd256 paged prep q")?; + let k_row_bytes = checked_band_offset(&k, row_offset, seq_len, "hd256 paged prep k")?; + let v_row_bytes = checked_band_offset(&v, row_offset, seq_len, "hd256 paged prep v")?; + let qo_row_bytes = checked_row_offset(q_out, row_offset, seq_len, "hd256 paged prep q_out")?; + let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); + let q_ptr = q_ptr + q_row_bytes; + let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); + let k_ptr = k_ptr + k_row_bytes; + let (v_ptr, _gv) = v.states.data.device_ptr(&ctx.stream); + let v_ptr = v_ptr + v_row_bytes; + let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); + let qo_ptr = qo_ptr + qo_row_bytes; + // The KV state owns mutation; this call borrows its shared pool handle. + let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); + let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); + let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); + let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); + let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); + let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); + let pi_ptr = pi_ptr + (pages_offset * std::mem::size_of::()) as u64; + + // Two `if`s, not a tuple-building match: a fn item only coerces to a + // fn pointer across plain `if` arms. + let fp8 = layout.storage == KvStorage::E4m3; + let launch = if fp8 { + ffi::qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda + } else { + ffi::qkv_norm_rope_paged_prefill_hd256_plain_cuda + }; + let entry_point = if fp8 { + "qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda" + } else { + "qkv_norm_rope_paged_prefill_hd256_plain_cuda" + }; + let q_stride = crate::ops::checked_i32(q.states.hidden_dim, "hd256 paged prep q stride")?; + let k_stride = crate::ops::checked_i32(k.states.hidden_dim, "hd256 paged prep k stride")?; + let v_stride = crate::ops::checked_i32(v.states.hidden_dim, "hd256 paged prep v stride")?; + let result = unsafe { + launch( + q_ptr as *const ffi::Half, + k_ptr as *const ffi::Half, + v_ptr as *const ffi::Half, + q_stride, + k_stride, + v_stride, + qn_ptr as *const ffi::Half, + kn_ptr as *const ffi::Half, + cos_ptr as *const ffi::Half, + sin_ptr as *const ffi::Half, + qo_ptr as *mut ffi::Half, + pool_ptr as *mut ffi::Half, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + page_indices_len, + page_origin_i32, + num_q_heads_i32, + num_kv_heads_i32, + seq_len_i32, + start_pos_i32, + cos_max_pos_i32, + rotary_dim_i32, + rms_eps, + geometry.page_size, + geometry.num_pages, + geometry.stride_page, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "{entry_point} failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + Ok(()) +} + +/// Per-request hd256 prep; `page_origins` is the released local-window front. +/// The decode batch is the row suffix `[row_offset..seq_len)` — 0 covers the +/// whole tensor, a mixed step parks its prefill segment in the prefix. +/// Invalid device metadata traps before the first paged-pool access. +#[allow(clippy::too_many_arguments)] +pub fn qkv_norm_rope_paged_decode_hd256_plain_into<'a>( + ctx: &DeviceContext, + q: impl Into>, + k: impl Into>, + v: impl Into>, + q_out: &mut HiddenStates, + row_offset: usize, + kv_pool: &CudaSlice, + layout: &PagedKvLayout, + q_norm_weight: &DeviceVec, + k_norm_weight: &DeviceVec, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + layer: usize, + page_indices: &CudaSlice, + page_indptr: &CudaSlice, + page_origins: &CudaSlice, + positions: &CudaSlice, + cos_max_pos: usize, + num_q_heads: usize, + num_kv_heads: usize, + rotary_dim: usize, + rms_eps: f32, +) -> Result<()> { + let (q, k, v) = (q.into(), k.into(), v.into()); + anyhow::ensure!( + row_offset < q.states.seq_len, + "hd256 paged decode prep row_offset {row_offset} leaves no rows of {}", + q.states.seq_len + ); + let batch = q.states.seq_len - row_offset; + anyhow::ensure!( + Some(q.width) == num_q_heads.checked_mul(256), + "hd256 paged decode prep q width {} != num_q_heads {} * 256", + q.width, + num_q_heads + ); + anyhow::ensure!( + q_out.hidden_dim == q.width && q_out.seq_len == q.states.seq_len, + "hd256 paged decode prep q_out [{} x {}] != q [{} x {}]", + q_out.hidden_dim, + q_out.seq_len, + q.width, + q.states.seq_len + ); + anyhow::ensure!( + Some(k.width) == num_kv_heads.checked_mul(256) && k.states.seq_len == q.states.seq_len, + "hd256 paged decode prep k [{} x {}] != [num_kv_heads {} * 256 x {}]", + k.width, + k.states.seq_len, + num_kv_heads, + q.states.seq_len + ); + anyhow::ensure!( + v.width == k.width && v.states.seq_len == q.states.seq_len, + "hd256 paged decode prep v [{} x {}] != k [{} x {}]", + v.width, + v.states.seq_len, + k.width, + q.states.seq_len + ); + let geometry = checked_paged_geometry( + "hd256 paged decode prep", + layout, + kv_pool.len(), + layer, + 256, + num_kv_heads, + true, + )?; + ensure_vec_backed(q_norm_weight, "hd256 paged decode prep q_norm_weight")?; + ensure_vec_backed(k_norm_weight, "hd256 paged decode prep k_norm_weight")?; + ensure_vec_backed(cos_cache, "hd256 paged decode prep cos_cache")?; + ensure_vec_backed(sin_cache, "hd256 paged decode prep sin_cache")?; + anyhow::ensure!( + q_norm_weight.len == 256 && k_norm_weight.len == 256, + "hd256 paged decode prep norm weight lens {} / {} != 256", + q_norm_weight.len, + k_norm_weight.len + ); + let table_rows = cos_max_pos.checked_mul(rotary_dim).ok_or_else(|| { + anyhow::anyhow!( + "hd256 paged decode prep cos_max_pos {cos_max_pos} * rotary_dim {rotary_dim} overflows" + ) + })?; + crate::ops::checked_i32(table_rows, "hd256 paged decode prep rope table extent")?; + anyhow::ensure!( + cos_cache.len >= table_rows && sin_cache.len >= table_rows, + "hd256 paged decode prep cos/sin lens {} / {} < cos_max_pos {cos_max_pos} * \ + rotary_dim {rotary_dim}", + cos_cache.len, + sin_cache.len + ); + anyhow::ensure!( + positions.len() >= batch && page_origins.len() >= batch && page_indptr.len() > batch, + "hd256 paged decode prep metadata lens (positions {}, origins {}, indptr {}) \ + do not cover batch {batch}", + positions.len(), + page_origins.len(), + page_indptr.len() + ); + let page_indices_len = crate::ops::checked_i32( + page_indices.len(), + "hd256 paged decode prep page_indices len", + )?; + let q_elems = q.states.checked_extent("hd256 paged decode prep q")?; + q_out.checked_extent("hd256 paged decode prep q_out")?; + let k_elems = k.states.checked_extent("hd256 paged decode prep k")?; + v.states.checked_extent("hd256 paged decode prep v")?; + crate::ops::checked_i32(q_elems, "hd256 paged decode prep q extent")?; + crate::ops::checked_i32(k_elems, "hd256 paged decode prep k extent")?; + let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd256 paged decode prep q heads")?; + let num_kv_heads_i32 = + crate::ops::checked_i32(num_kv_heads, "hd256 paged decode prep kv heads")?; + let batch_i32 = crate::ops::checked_i32(batch, "hd256 paged decode prep batch")?; + let cos_max_pos_i32 = + crate::ops::checked_i32(cos_max_pos, "hd256 paged decode prep cos_max_pos")?; + let rotary_dim_i32 = crate::ops::checked_i32(rotary_dim, "hd256 paged decode prep rotary_dim")?; + + let q_row_bytes = checked_band_offset(&q, row_offset, batch, "hd256 paged decode prep q")?; + let k_row_bytes = checked_band_offset(&k, row_offset, batch, "hd256 paged decode prep k")?; + let v_row_bytes = checked_band_offset(&v, row_offset, batch, "hd256 paged decode prep v")?; + let qo_row_bytes = + checked_row_offset(q_out, row_offset, batch, "hd256 paged decode prep q_out")?; + let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); + let q_ptr = q_ptr + q_row_bytes; + let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); + let k_ptr = k_ptr + k_row_bytes; + let (v_ptr, _gv) = v.states.data.device_ptr(&ctx.stream); + let v_ptr = v_ptr + v_row_bytes; + let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); + let qo_ptr = qo_ptr + qo_row_bytes; + // The KV states own mutation; this call borrows the shared pool handle. + let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); + let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); + let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); + let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); + let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); + let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); + let (ip_ptr, _gip) = page_indptr.device_ptr(&ctx.stream); + let (og_ptr, _gog) = page_origins.device_ptr(&ctx.stream); + let (ps_ptr, _gps) = positions.device_ptr(&ctx.stream); + + // Two `if`s, not a tuple-building match: a fn item only coerces to a + // fn pointer across plain `if` arms. + let fp8 = layout.storage == KvStorage::E4m3; + let launch = if fp8 { + ffi::qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda + } else { + ffi::qkv_norm_rope_paged_decode_hd256_plain_cuda + }; + let entry_point = if fp8 { + "qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda" + } else { + "qkv_norm_rope_paged_decode_hd256_plain_cuda" + }; + let q_stride = + crate::ops::checked_i32(q.states.hidden_dim, "hd256 paged decode prep q stride")?; + let k_stride = + crate::ops::checked_i32(k.states.hidden_dim, "hd256 paged decode prep k stride")?; + let v_stride = + crate::ops::checked_i32(v.states.hidden_dim, "hd256 paged decode prep v stride")?; + let result = unsafe { + launch( + q_ptr as *const ffi::Half, + k_ptr as *const ffi::Half, + v_ptr as *const ffi::Half, + q_stride, + k_stride, + v_stride, + qn_ptr as *const ffi::Half, + kn_ptr as *const ffi::Half, + cos_ptr as *const ffi::Half, + sin_ptr as *const ffi::Half, + qo_ptr as *mut ffi::Half, + pool_ptr as *mut ffi::Half, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + page_indices_len, + ip_ptr as *const i32, + og_ptr as *const i32, + ps_ptr as *const i32, + num_q_heads_i32, + num_kv_heads_i32, + batch_i32, + cos_max_pos_i32, + rotary_dim_i32, + rms_eps, + geometry.page_size, + geometry.num_pages, + geometry.stride_page, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "{entry_point} failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + Ok(()) +} + +/// QK RMSNorm + partial RoPE for the hd512 paged-prefill prep (Gemma 4 +/// global layers). Q is normalised + partially rotated into a contiguous +/// `q_out`; K is normalised + partially rotated straight into the paged KV +/// pool at layer `layer` (feeds `batch_prefill_paged`); V — the K=V fork — +/// is the weightless RMS norm of the same raw K, written in the same pass. +/// No gate; plain-w RMSNorm. +/// +/// The rotation is the engine's proportional one over `cos_max_pos` rows of +/// 512: pairs `(d, d + 256)` with the live angles first and the identity +/// past them. How the row is laid out — K and V as two blocks, or the +/// folded row of `KvFormat::Folded` — is the layout's format, and the +/// kernel is told it as bands rather than asked to know the format. `q` and +/// `k` are column bands read at their own row strides, so they may share one +/// fused Q|K projection row. +/// +/// The kernel `__trap()`s on any out-of-range pos or page id as the +/// second layer of the host validation's defence. +#[allow(clippy::too_many_arguments)] +pub fn qk_norm_partial_rope_paged_prefill_hd512_into<'a>( + ctx: &DeviceContext, + q: impl Into>, + k: impl Into>, + q_out: &mut HiddenStates, + row_offset: usize, + kv_pool: &CudaSlice, + layout: &PagedKvLayout, + q_norm_weight: &DeviceVec, + k_norm_weight: &DeviceVec, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + layer: usize, + page_indices: &CudaSlice, + pages_offset: usize, + start_pos: usize, + cos_max_pos: usize, + num_q_heads: usize, + num_kv_heads: usize, + rms_eps: f32, +) -> Result<()> { + let (q, k) = (q.into(), k.into()); + // Same suffix-window contract as the hd256 prefill prep: the segment + // is the rows `[row_offset..seq_len)` with its table at `pages_offset`. + anyhow::ensure!( + row_offset < q.states.seq_len, + "hd512 prefill prep row_offset {row_offset} leaves no rows of {}", + q.states.seq_len + ); + let seq_len = q.states.seq_len - row_offset; + anyhow::ensure!( + pages_offset < page_indices.len(), + "hd512 prefill prep pages_offset {pages_offset} exceeds table len {}", + page_indices.len() + ); + let q_dim = num_q_heads.checked_mul(512).ok_or_else(|| { + anyhow::anyhow!("hd512 prefill prep num_q_heads {num_q_heads} * 512 overflows") + })?; + let kv_dim = num_kv_heads.checked_mul(512).ok_or_else(|| { + anyhow::anyhow!("hd512 prefill prep num_kv_heads {num_kv_heads} * 512 overflows") + })?; + anyhow::ensure!( + q.width == q_dim, + "hd512 prefill prep q width {} != num_q_heads {num_q_heads} * 512", + q.width + ); + anyhow::ensure!( + q_out.hidden_dim == q.width, + "hd512 prefill prep q_out.hidden_dim {} != q width {}", + q_out.hidden_dim, + q.width + ); + anyhow::ensure!( + q_out.seq_len == q.states.seq_len, + "hd512 prefill prep q_out.seq_len {} != q rows {}", + q_out.seq_len, + q.states.seq_len + ); + anyhow::ensure!( + k.width == kv_dim, + "hd512 prefill prep k width {} != num_kv_heads {num_kv_heads} * 512", + k.width + ); + anyhow::ensure!( + k.states.seq_len == q.states.seq_len, + "hd512 prefill prep k rows {} != q rows {}", + k.states.seq_len, + q.states.seq_len + ); + let target = hd512_prep_target( + "hd512 prefill prep", + layout, + kv_pool.len(), + layer, + num_kv_heads, + )?; + anyhow::ensure!( + q_norm_weight.len == 512, + "hd512 prefill prep q_norm_weight len {} != 512", + q_norm_weight.len + ); + anyhow::ensure!( + k_norm_weight.len == 512, + "hd512 prefill prep k_norm_weight len {} != 512", + k_norm_weight.len + ); + let end_pos = start_pos.checked_add(seq_len).ok_or_else(|| { + anyhow::anyhow!("hd512 prefill prep start_pos {start_pos} + seq_len {seq_len} overflows") + })?; + let covered = (page_indices.len() - pages_offset) + .checked_mul(layout.page_size) + .ok_or_else(|| { + anyhow::anyhow!( + "hd512 prefill prep page window {} * page_size {} overflows", + page_indices.len() - pages_offset, + layout.page_size + ) + })?; + anyhow::ensure!( + covered >= end_pos, + "hd512 prefill prep pages cover {covered} tokens, need start_pos {start_pos} + seq_len {seq_len}" + ); + anyhow::ensure!( + end_pos <= cos_max_pos, + "hd512 prefill prep start_pos {start_pos} + seq_len {seq_len} > cos_max_pos {cos_max_pos}" + ); + let table_len = cos_max_pos.checked_mul(512).ok_or_else(|| { + anyhow::anyhow!("hd512 prefill prep cos_max_pos {cos_max_pos} * 512 overflows") + })?; + anyhow::ensure!( + cos_cache.len >= table_len, + "hd512 prefill prep cos_cache len {} < cos_max_pos {cos_max_pos} * 512", + cos_cache.len + ); + anyhow::ensure!( + sin_cache.len >= table_len, + "hd512 prefill prep sin_cache len {} < cos_max_pos {cos_max_pos} * 512", + sin_cache.len + ); + let q_elems = q.states.checked_extent("hd512 prefill prep q")?; + q_out.checked_extent("hd512 prefill prep q_out")?; + let k_elems = k.states.checked_extent("hd512 prefill prep k")?; + crate::ops::checked_i32(q_elems, "hd512 prefill prep q extent")?; + crate::ops::checked_i32(k_elems, "hd512 prefill prep k extent")?; + crate::ops::checked_i32(table_len, "hd512 prefill prep rope table extent")?; + ensure_vec_backed(q_norm_weight, "hd512 prefill prep q_norm_weight")?; + ensure_vec_backed(k_norm_weight, "hd512 prefill prep k_norm_weight")?; + ensure_vec_backed(cos_cache, "hd512 prefill prep cos_cache")?; + ensure_vec_backed(sin_cache, "hd512 prefill prep sin_cache")?; + + let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd512 prefill prep num_q_heads")?; + let num_kv_heads_i32 = + crate::ops::checked_i32(num_kv_heads, "hd512 prefill prep num_kv_heads")?; + let seq_len_i32 = crate::ops::checked_i32(seq_len, "hd512 prefill prep seq_len")?; + let start_pos_i32 = crate::ops::checked_i32(start_pos, "hd512 prefill prep start_pos")?; + let cos_max_pos_i32 = crate::ops::checked_i32(cos_max_pos, "hd512 prefill prep cos_max_pos")?; + + let page_indices_len = crate::ops::checked_i32( + page_indices.len() - pages_offset, + "hd512 paged prep page window len", + )?; + let q_row_bytes = checked_band_offset(&q, row_offset, seq_len, "hd512 prefill prep q")?; + let k_row_bytes = checked_band_offset(&k, row_offset, seq_len, "hd512 prefill prep k")?; + let qo_row_bytes = checked_row_offset(q_out, row_offset, seq_len, "hd512 prefill prep q_out")?; + let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); + let q_ptr = q_ptr + q_row_bytes; + let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); + let k_ptr = k_ptr + k_row_bytes; + let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); + let qo_ptr = qo_ptr + qo_row_bytes; + // The KV state owns mutation; this call borrows its shared pool handle. + let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); + let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); + let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); + let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); + let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); + let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); + let pi_ptr = pi_ptr + (pages_offset * std::mem::size_of::()) as u64; + + let q_stride = crate::ops::checked_i32(q.states.hidden_dim, "hd512 prefill prep q stride")?; + let k_stride = crate::ops::checked_i32(k.states.hidden_dim, "hd512 prefill prep k stride")?; + let result = unsafe { + ffi::qk_norm_partial_rope_paged_prefill_hd512_cuda( + q_ptr as *const ffi::Half, + k_ptr as *const ffi::Half, + q_stride, + k_stride, + qn_ptr as *const ffi::Half, + kn_ptr as *const ffi::Half, + cos_ptr as *const ffi::Half, + sin_ptr as *const ffi::Half, + qo_ptr as *mut ffi::Half, + pool_ptr as *mut ffi::Half, + target.k_offset_elems, + target.v_offset_elems, + pi_ptr as *const i32, + page_indices_len, + num_q_heads_i32, + num_kv_heads_i32, + seq_len_i32, + start_pos_i32, + cos_max_pos_i32, + target.row_width, + target.fold_rotary, + rms_eps, + target.page_size, + target.num_pages, + target.stride_page, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "qk_norm_partial_rope_paged_prefill_hd512_cuda failed with error \ + {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + Ok(()) +} + +/// Per-request hd512 prep for the non-evicting global cache; V is the K fork. +/// Invalid device metadata traps before the first paged-pool access. +#[allow(clippy::too_many_arguments)] +pub fn qk_norm_partial_rope_paged_decode_hd512_into<'a>( + ctx: &DeviceContext, + q: impl Into>, + k: impl Into>, + q_out: &mut HiddenStates, + row_offset: usize, + kv_pool: &CudaSlice, + layout: &PagedKvLayout, + q_norm_weight: &DeviceVec, + k_norm_weight: &DeviceVec, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + layer: usize, + page_indices: &CudaSlice, + page_indptr: &CudaSlice, + page_origins: &CudaSlice, + positions: &CudaSlice, + cos_max_pos: usize, + num_q_heads: usize, + num_kv_heads: usize, + rms_eps: f32, +) -> Result<()> { + let (q, k) = (q.into(), k.into()); + anyhow::ensure!( + row_offset < q.states.seq_len, + "hd512 paged decode prep row_offset {row_offset} leaves no rows of {}", + q.states.seq_len + ); + let batch = q.states.seq_len - row_offset; + anyhow::ensure!( + Some(q.width) == num_q_heads.checked_mul(512), + "hd512 paged decode prep q width {} != num_q_heads {} * 512", + q.width, + num_q_heads + ); + anyhow::ensure!( + q_out.hidden_dim == q.width && q_out.seq_len == q.states.seq_len, + "hd512 paged decode prep q_out [{} x {}] != q [{} x {}]", + q_out.hidden_dim, + q_out.seq_len, + q.width, + q.states.seq_len + ); + anyhow::ensure!( + Some(k.width) == num_kv_heads.checked_mul(512) && k.states.seq_len == q.states.seq_len, + "hd512 paged decode prep k [{} x {}] != [num_kv_heads {} * 512 x {}]", + k.width, + k.states.seq_len, + num_kv_heads, + q.states.seq_len + ); + let target = hd512_prep_target( + "hd512 paged decode prep", + layout, + kv_pool.len(), + layer, + num_kv_heads, + )?; + ensure_vec_backed(q_norm_weight, "hd512 paged decode prep q_norm_weight")?; + ensure_vec_backed(k_norm_weight, "hd512 paged decode prep k_norm_weight")?; + ensure_vec_backed(cos_cache, "hd512 paged decode prep cos_cache")?; + ensure_vec_backed(sin_cache, "hd512 paged decode prep sin_cache")?; + anyhow::ensure!( + q_norm_weight.len == 512 && k_norm_weight.len == 512, + "hd512 paged decode prep norm weight lens {} / {} != 512", + q_norm_weight.len, + k_norm_weight.len + ); + let table_rows = cos_max_pos.checked_mul(512).ok_or_else(|| { + anyhow::anyhow!("hd512 paged decode prep cos_max_pos {cos_max_pos} * 512 overflows") + })?; + crate::ops::checked_i32(table_rows, "hd512 paged decode prep rope table extent")?; + anyhow::ensure!( + cos_cache.len >= table_rows && sin_cache.len >= table_rows, + "hd512 paged decode prep cos/sin lens {} / {} < cos_max_pos {cos_max_pos} * 512", + cos_cache.len, + sin_cache.len + ); + anyhow::ensure!( + positions.len() >= batch && page_indptr.len() > batch && page_origins.len() >= batch, + "hd512 paged decode prep metadata lens (positions {}, indptr {}, origins {}) do \ + not cover batch {batch}", + positions.len(), + page_indptr.len(), + page_origins.len() + ); + let page_indices_len = crate::ops::checked_i32( + page_indices.len(), + "hd512 paged decode prep page_indices len", + )?; + let q_elems = q.states.checked_extent("hd512 paged decode prep q")?; + q_out.checked_extent("hd512 paged decode prep q_out")?; + let k_elems = k.states.checked_extent("hd512 paged decode prep k")?; + crate::ops::checked_i32(q_elems, "hd512 paged decode prep q extent")?; + crate::ops::checked_i32(k_elems, "hd512 paged decode prep k extent")?; + let num_q_heads_i32 = crate::ops::checked_i32(num_q_heads, "hd512 paged decode prep q heads")?; + let num_kv_heads_i32 = + crate::ops::checked_i32(num_kv_heads, "hd512 paged decode prep kv heads")?; + let batch_i32 = crate::ops::checked_i32(batch, "hd512 paged decode prep batch")?; + let cos_max_pos_i32 = + crate::ops::checked_i32(cos_max_pos, "hd512 paged decode prep cos_max_pos")?; + let q_row_bytes = checked_band_offset(&q, row_offset, batch, "hd512 paged decode prep q")?; + let k_row_bytes = checked_band_offset(&k, row_offset, batch, "hd512 paged decode prep k")?; + let qo_row_bytes = + checked_row_offset(q_out, row_offset, batch, "hd512 paged decode prep q_out")?; + let (q_ptr, _gq) = q.states.data.device_ptr(&ctx.stream); + let q_ptr = q_ptr + q_row_bytes; + let (k_ptr, _gk) = k.states.data.device_ptr(&ctx.stream); + let k_ptr = k_ptr + k_row_bytes; + let (qo_ptr, _gqo) = q_out.data.device_ptr_mut(&ctx.stream); + let qo_ptr = qo_ptr + qo_row_bytes; + // The KV states own mutation; this call borrows the shared pool handle. + let (pool_ptr, _gp) = kv_pool.device_ptr(&ctx.stream); + let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); + let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); + let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); + let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); + let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); + let (ip_ptr, _gip) = page_indptr.device_ptr(&ctx.stream); + let (po_ptr, _gpo) = page_origins.device_ptr(&ctx.stream); + let (ps_ptr, _gps) = positions.device_ptr(&ctx.stream); + + let q_stride = + crate::ops::checked_i32(q.states.hidden_dim, "hd512 paged decode prep q stride")?; + let k_stride = + crate::ops::checked_i32(k.states.hidden_dim, "hd512 paged decode prep k stride")?; + let result = unsafe { + ffi::qk_norm_partial_rope_paged_decode_hd512_cuda( + q_ptr as *const ffi::Half, + k_ptr as *const ffi::Half, + q_stride, + k_stride, + qn_ptr as *const ffi::Half, + kn_ptr as *const ffi::Half, + cos_ptr as *const ffi::Half, + sin_ptr as *const ffi::Half, + qo_ptr as *mut ffi::Half, + pool_ptr as *mut ffi::Half, + target.k_offset_elems, + target.v_offset_elems, + pi_ptr as *const i32, + page_indices_len, + ip_ptr as *const i32, + po_ptr as *const i32, + ps_ptr as *const i32, + num_q_heads_i32, + num_kv_heads_i32, + batch_i32, + cos_max_pos_i32, + target.row_width, + target.fold_rotary, + rms_eps, + target.page_size, + target.num_pages, + target.stride_page, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "qk_norm_partial_rope_paged_decode_hd512_cuda failed with error \ + {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + Ok(()) +} diff --git a/pegainfer-kernels/src/ops/kimi_k2/o_proj.rs b/pegainfer-kernels/src/ops/kimi_k2/o_proj.rs index e308d7ab8..f001ff552 100644 --- a/pegainfer-kernels/src/ops/kimi_k2/o_proj.rs +++ b/pegainfer-kernels/src/ops/kimi_k2/o_proj.rs @@ -7,6 +7,7 @@ use cudarc::driver::DevicePtrMut; use super::experts::KIMI_K2_HIDDEN; use super::mla::KIMI_K2_MLA_V_HEAD_DIM; use crate::ffi; +use crate::ops::ffi_status_detail; use crate::tensor::DeviceContext; use crate::tensor::DeviceMatrix; use crate::tensor::GpuTensor; @@ -67,17 +68,9 @@ pub fn kimi_o_proj_cublaslt_into( ctx.stream.cu_stream(), ); if status != 0 { - if status >= 100_000 { - bail!( - "Kimi o_proj cuBLASLt failed: cublas_status={}, batch_size={}", - status - 100_000, - batch_size - ); - } bail!( - "Kimi o_proj cuBLASLt launch failed: cuda_status={}, batch_size={}", - status, - batch_size + "Kimi o_proj cuBLASLt failed: {}, batch_size={batch_size}", + ffi_status_detail(status) ); } } diff --git a/pegainfer-kernels/src/ops/kimi_k2/shared_gate_up.rs b/pegainfer-kernels/src/ops/kimi_k2/shared_gate_up.rs index eea8bd031..c9da9dec7 100644 --- a/pegainfer-kernels/src/ops/kimi_k2/shared_gate_up.rs +++ b/pegainfer-kernels/src/ops/kimi_k2/shared_gate_up.rs @@ -7,6 +7,7 @@ use cudarc::driver::DevicePtrMut; use super::experts::KIMI_K2_HIDDEN; use super::experts::KIMI_K2_SHARED_GATE_UP; use crate::ffi; +use crate::ops::ffi_status_detail; use crate::tensor::DeviceContext; use crate::tensor::DeviceMatrix; use crate::tensor::GpuTensor; @@ -61,17 +62,9 @@ pub fn kimi_shared_gate_up_cublaslt_into( ctx.stream.cu_stream(), ); if status != 0 { - if status >= 100_000 { - bail!( - "Kimi shared_gate_up cuBLASLt failed: cublas_status={}, batch_size={}", - status - 100_000, - batch_size - ); - } bail!( - "Kimi shared_gate_up cuBLASLt launch failed: cuda_status={}, batch_size={}", - status, - batch_size + "Kimi shared_gate_up cuBLASLt failed: {}, batch_size={batch_size}", + ffi_status_detail(status) ); } } diff --git a/pegainfer-kernels/src/ops/linear.rs b/pegainfer-kernels/src/ops/linear.rs index 7f7e26865..feec0c618 100644 --- a/pegainfer-kernels/src/ops/linear.rs +++ b/pegainfer-kernels/src/ops/linear.rs @@ -10,6 +10,7 @@ use cudarc::driver::DevicePtrMut; use half::bf16; use crate::ffi; +use crate::ops::ffi_status_detail; use crate::tensor::DeviceContext; use crate::tensor::DeviceMatrix; use crate::tensor::DeviceVec; @@ -382,16 +383,9 @@ pub fn gemm_lt_pin_into_checked( weight.rows, weight.cols ), - s if s >= 100_000 => bail!( - "cublasLt pin GEMM failed: cublas_status={}, m={}, n={}, k={}", - s - 100_000, - weight.rows, - x.seq_len, - weight.cols - ), s => bail!( - "cublasLt pin GEMM launch failed: cuda_status={}, m={}, n={}, k={}", - s, + "cublasLt pin GEMM failed: {}, m={}, n={}, k={}", + ffi_status_detail(s), weight.rows, x.seq_len, weight.cols @@ -676,18 +670,9 @@ fn gemm_per_token_into_checked( crate::tensor::active_cu_stream(ctx), ); if status != 0 { - if status >= 100_000 { - bail!( - "cuBLAS per-token GEMM failed: cublas_status={}, m={}, batch={}, k={}", - status - 100_000, - weight.rows, - x.seq_len, - weight.cols - ); - } bail!( - "CUDA per-token GEMM launch failed: cuda_status={}, m={}, batch={}, k={}", - status, + "per-token GEMM failed: {}, m={}, batch={}, k={}", + ffi_status_detail(status), weight.rows, x.seq_len, weight.cols @@ -860,13 +845,10 @@ fn launch_gemm_pin( LoRA prefill-delta shape), or the numeric policy was set after executor \ construction. Run without --batch-invariant or report this shape" ), - s if s >= 100_000 => { - bail!( - "cuBLAS pin GEMM failed: cublas_status={}, m={m}, n={n}, k={k}", - s - 100_000 - ) - } - s => bail!("CUDA pin GEMM launch failed: cuda_status={s}, m={m}, n={n}, k={k}"), + s => bail!( + "pin GEMM failed: {}, m={m}, n={n}, k={k}", + ffi_status_detail(s) + ), } } } @@ -892,13 +874,10 @@ fn launch_gemm_pertoken( crate::tensor::active_cu_stream(ctx), ); if status != 0 { - if status >= 100_000 { - bail!( - "cuBLAS per-token GEMM failed: cublas_status={}, m={m}, n={n}, k={k}", - status - 100_000 - ); - } - bail!("CUDA per-token GEMM launch failed: cuda_status={status}, m={m}, n={n}, k={k}"); + bail!( + "per-token GEMM failed: {}, m={m}, n={n}, k={k}", + ffi_status_detail(status) + ); } } PER_TOKEN_SERVED.fetch_add(1, Ordering::Relaxed); @@ -969,19 +948,9 @@ fn launch_gemm( }; } if status != 0 { - if status >= 100_000 { - bail!( - "cuBLAS GEMM failed: cublas_status={}, m={}, n={}, k={}, graphsafe={}", - status - 100_000, - m, - n, - k, - graphsafe - ); - } bail!( - "CUDA GEMM launch failed: cuda_status={}, m={}, n={}, k={}, graphsafe={}", - status, + "GEMM failed: {}, m={}, n={}, k={}, graphsafe={}", + ffi_status_detail(status), m, n, k, diff --git a/pegainfer-kernels/src/ops/norm.rs b/pegainfer-kernels/src/ops/norm.rs index 0ad9630a6..14b52bbee 100644 --- a/pegainfer-kernels/src/ops/norm.rs +++ b/pegainfer-kernels/src/ops/norm.rs @@ -627,9 +627,23 @@ pub fn rms_norm_offset_into( Ok(()) } -/// Batched per-head RMSNorm with F32 weight + SiLU gate multiplication. +/// Which activation multiplies a gated RMSNorm's output. +/// +/// Qwen3.5's Gated DeltaNet uses [`GatedNormActivation::Silu`]. +/// Qwen3.8-Flash-Next declares `output_gate_type: "sigmoid"`, and its reference +/// implementation passes that field straight to the gated norm, so the same +/// kernel has to serve both. Note the *attention* output gate is a hardcoded +/// sigmoid on both lines; this enum is only about the gated norm. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GatedNormActivation { + Silu, + Sigmoid, +} + +/// Batched per-head RMSNorm with F32 weight + a gated activation multiplies. /// HiddenStates are flattened as (seq_len * num_heads) contiguous head slices. -#[allow(clippy::too_many_arguments)] +/// `activation` chooses the gate: Qwen3.5's Gated DeltaNet passes +/// [`GatedNormActivation::Silu`], Qwen3.8-Flash-Next's passes `Sigmoid`. pub fn rms_norm_gated_batch_into( ctx: &DeviceContext, x: &HiddenStates, @@ -639,6 +653,7 @@ pub fn rms_norm_gated_batch_into( num_heads: usize, head_dim: usize, eps: f32, + activation: GatedNormActivation, ) { let total_heads = x.seq_len * num_heads; assert_eq!(x.hidden_dim, num_heads * head_dim); @@ -650,17 +665,30 @@ pub fn rms_norm_gated_batch_into( let (w_ptr, _gw) = weight.device_ptr(&ctx.stream); let (g_ptr, _gg) = gate.data.device_ptr(&ctx.stream); let (o_ptr, _go) = out.data.device_ptr_mut(&ctx.stream); + let stream = crate::tensor::active_cu_stream(ctx); unsafe { - ffi::rms_norm_gated_cuda( - x_ptr as *const ffi::Half, - w_ptr as *const f32, - g_ptr as *const ffi::Half, - o_ptr as *mut ffi::Half, - total_heads as i32, - head_dim as i32, - eps, - crate::tensor::active_cu_stream(ctx), - ); + match activation { + GatedNormActivation::Silu => ffi::rms_norm_gated_cuda( + x_ptr as *const ffi::Half, + w_ptr as *const f32, + g_ptr as *const ffi::Half, + o_ptr as *mut ffi::Half, + total_heads as i32, + head_dim as i32, + eps, + stream, + ), + GatedNormActivation::Sigmoid => ffi::rms_norm_gated_sigmoid_cuda( + x_ptr as *const ffi::Half, + w_ptr as *const f32, + g_ptr as *const ffi::Half, + o_ptr as *mut ffi::Half, + total_heads as i32, + head_dim as i32, + eps, + stream, + ), + } } } @@ -805,4 +833,151 @@ mod parity { .expect("fused combine"); assert_bitwise(&ctx, &split, &fused, "fused MoE combine"); } + + /// CPU reference for the per-head gated RMSNorm, in f64 then rounded to bf16. + /// The kernel reduces in f32 with a warp tree, so this is an oracle with a + /// bf16-sized tolerance, not a bit-exact model. + fn gated_norm_reference( + x: &[bf16], + weight: &[f32], + gate: &[bf16], + head_dim: usize, + eps: f32, + sigmoid: bool, + ) -> Vec { + let mut out = vec![bf16::from_f32(0.0); x.len()]; + for head in 0..(x.len() / head_dim) { + let base = head * head_dim; + let mut sumsq = 0.0f64; + for t in 0..head_dim { + let v = f64::from(x[base + t].to_f32()); + sumsq += v * v; + } + let inv_rms = 1.0 / (sumsq / head_dim as f64 + f64::from(eps)).sqrt(); + for t in 0..head_dim { + let normed = f64::from(x[base + t].to_f32()) * inv_rms * f64::from(weight[t]); + let g = f64::from(gate[base + t].to_f32()); + let sigma = 1.0 / (1.0 + (-g).exp()); + let activated = if sigmoid { sigma } else { g * sigma }; + out[base + t] = bf16::from_f32((normed * activated) as f32); + } + } + out + } + + fn assert_close(got: &[bf16], want: &[bf16], what: &str) { + assert_eq!(got.len(), want.len(), "{what} length"); + assert!(!got.is_empty(), "{what}: nothing to compare"); + let tolerance = |magnitude: f64| 0.02 + 0.03 * magnitude; + let mut worst = (0usize, 0.0f64, 0.0f64); + let mut worst_excess = f64::NEG_INFINITY; + for (index, (got_value, want_value)) in got.iter().zip(want).enumerate() { + let (got_value, want_value) = ( + f64::from(got_value.to_f32()), + f64::from(want_value.to_f32()), + ); + let (abs_diff, magnitude) = ((got_value - want_value).abs(), want_value.abs()); + // A NaN never satisfies the comparison below, so it would be skipped + // rather than reported and an all-NaN output would pass. Non-finite + // values are rejected here instead of being ranked. + assert!( + got_value.is_finite() && want_value.is_finite(), + "{what}: element {index} is not finite (got {got_value}, reference {want_value})" + ); + // Rank by the margin against the actual bound, not by the raw + // deviation: the bound is affine in `magnitude`, so the two orderings + // differ, and ranking by deviation can let a large-magnitude element + // past its own bound while a small-magnitude element, ranked higher, + // passes. Every element is checked for that reason. + let excess = abs_diff - tolerance(magnitude); + if excess > worst_excess { + worst_excess = excess; + worst = (index, abs_diff, magnitude); + } + } + let (index, abs_diff, magnitude) = worst; + assert!( + worst_excess <= 0.0, + "{what}: element {index} differs by {abs_diff} (reference {magnitude}), over the \ + tolerance by {worst_excess}" + ); + } + + /// The gated norm has two activations, because Qwen3.8-Flash-Next's + /// `output_gate_type` selects sigmoid where Qwen3.5 uses silu. Run at the + /// Flash-Next GDN geometry (48 value heads at head_dim 128), each activation + /// against its own CPU reference. The two distinct references are what pin the + /// template parameter: were both entry points instantiated with one activation, + /// the mismatched comparison would fail on nearly every element (silu is + /// `g * sigmoid(g)`, so the two differ wherever `g` is not zero). + #[test] + #[ignore = "requires a GPU"] + fn the_gated_norm_activations_match_their_references() { + let ctx = DeviceContext::new_with_device(0).expect("device"); + let (num_heads, head_dim, seq_len, eps) = (48usize, 128usize, 4usize, 1e-6f32); + let d = num_heads * head_dim; + let x = seeded(&ctx, d, seq_len, 11); + let gate = seeded(&ctx, d, seq_len, 12); + let weight_host: Vec = (0..head_dim) + .map(|i| ((i % 7) as f32 - 3.0) * 0.25 + 1.0) + .collect(); + let weight = ctx.stream.clone_htod(&weight_host).expect("weight up"); + + let mut silu_out = HiddenStates::zeros(&ctx, d, seq_len).expect("silu buf"); + rms_norm_gated_batch_into( + &ctx, + &x, + &weight, + &gate, + &mut silu_out, + num_heads, + head_dim, + eps, + GatedNormActivation::Silu, + ); + let mut sigmoid_out = HiddenStates::zeros(&ctx, d, seq_len).expect("sigmoid buf"); + rms_norm_gated_batch_into( + &ctx, + &x, + &weight, + &gate, + &mut sigmoid_out, + num_heads, + head_dim, + eps, + GatedNormActivation::Sigmoid, + ); + + let x_host = ctx.stream.clone_dtoh(&x.data).expect("x down"); + let gate_host = ctx.stream.clone_dtoh(&gate.data).expect("gate down"); + let silu_host = ctx.stream.clone_dtoh(&silu_out.data).expect("silu down"); + let sigmoid_host = ctx + .stream + .clone_dtoh(&sigmoid_out.data) + .expect("sigmoid down"); + + assert_close( + &silu_host, + &gated_norm_reference(&x_host, &weight_host, &gate_host, head_dim, eps, false), + "silu gated norm", + ); + assert_close( + &sigmoid_host, + &gated_norm_reference(&x_host, &weight_host, &gate_host, head_dim, eps, true), + "sigmoid gated norm", + ); + } + + /// `assert_close` must not treat a non-finite value as "close": a NaN never + /// wins the margin comparison, so without the finiteness check an all-NaN + /// output would leave `worst_excess` at negative infinity and pass. + #[test] + #[should_panic(expected = "not finite")] + fn assert_close_rejects_non_finite_values() { + assert_close( + &[bf16::from_f32(f32::NAN)], + &[bf16::from_f32(1.0)], + "nan output", + ); + } } diff --git a/pegainfer-kernels/src/ops/qwen35.rs b/pegainfer-kernels/src/ops/qwen35.rs new file mode 100644 index 000000000..5807a6a33 --- /dev/null +++ b/pegainfer-kernels/src/ops/qwen35.rs @@ -0,0 +1,360 @@ +//! Qwen3.5 full-attention wrappers at head dim 256: the decode prep and the +//! paged attention reads. + +use anyhow::Result; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use half::bf16; + +use super::PrefillPagedPlan; +use super::attention::PagedGeometry; +use super::attention::checked_paged_geometry; +use crate::ffi; +use crate::paged_kv::PagedKvLayout; +use crate::tensor::DeviceContext; +use crate::tensor::DeviceVec; +use crate::tensor::HiddenStates; + +/// Batched QK RMSNorm + partial RoPE for Qwen3.5 HD256 decode. +/// +/// Reads Q from interleaved `q_full` ([q, gate] per head), writes prepared Q into `q`, +/// and normalizes/applies partial RoPE to `k` in-place using per-request positions. +#[allow(clippy::too_many_arguments)] +pub fn qk_norm_partial_rope_batched_decode_hd256_into( + ctx: &DeviceContext, + q_full: &HiddenStates, + q: &mut HiddenStates, + k: &mut HiddenStates, + q_norm_weight: &DeviceVec, + k_norm_weight: &DeviceVec, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + positions_d: &CudaSlice, + num_q_heads: usize, + num_kv_heads: usize, + rotary_dim: usize, + rms_eps: f32, +) { + let batch_size = q.seq_len; + debug_assert_eq!(q_full.seq_len, batch_size); + debug_assert_eq!(k.seq_len, batch_size); + + let (qf_ptr, _gqf) = q_full.data.device_ptr(&ctx.stream); + let (q_ptr, _gq) = q.data.device_ptr_mut(&ctx.stream); + let (k_ptr, _gk) = k.data.device_ptr_mut(&ctx.stream); + let (qn_ptr, _gqn) = q_norm_weight.data.device_ptr(&ctx.stream); + let (kn_ptr, _gkn) = k_norm_weight.data.device_ptr(&ctx.stream); + let (cos_ptr, _gc) = cos_cache.data.device_ptr(&ctx.stream); + let (sin_ptr, _gs) = sin_cache.data.device_ptr(&ctx.stream); + let (pos_ptr, _gp) = positions_d.device_ptr(&ctx.stream); + + unsafe { + ffi::qk_norm_partial_rope_batched_decode_hd256_cuda( + qf_ptr as *const ffi::Half, + k_ptr as *mut ffi::Half, + qn_ptr as *const ffi::Half, + kn_ptr as *const ffi::Half, + cos_ptr as *const ffi::Half, + sin_ptr as *const ffi::Half, + pos_ptr as *const i32, + q_ptr as *mut ffi::Half, + num_q_heads as i32, + num_kv_heads as i32, + batch_size as i32, + rotary_dim as i32, + rms_eps, + crate::tensor::active_cu_stream(ctx), + ); + } +} + +#[allow(clippy::too_many_arguments)] +fn scatter_decode_kv_into_paged( + ctx: &DeviceContext, + k: &HiddenStates, + v: &HiddenStates, + kv_buffer: &CudaSlice, + layout: &PagedKvLayout, + layer: usize, + page_indices_d: &CudaSlice, + page_indptr_d: &CudaSlice, + last_page_len_d: &CudaSlice, + positions_d: &CudaSlice, + request_indices_d: &CudaSlice, + batch_size: usize, + op_name: &str, +) -> Result<()> { + let num_kv_heads = layout.num_kv_heads; + let head_dim = layout.head_dim; + let page_size = layout.page_size; + + let PagedGeometry { + k_offset_elems: k_offset, + v_offset_elems: v_offset, + stride_page, + .. + } = checked_paged_geometry( + op_name, + layout, + kv_buffer.len(), + layer, + head_dim, + num_kv_heads, + false, + )?; + + let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); + let (k_ptr, _gk) = k.data.device_ptr(&ctx.stream); + let (v_ptr, _gv) = v.data.device_ptr(&ctx.stream); + let (pi_ptr, _gpi) = page_indices_d.device_ptr(&ctx.stream); + let (pip_ptr, _gpip) = page_indptr_d.device_ptr(&ctx.stream); + let (lpl_ptr, _glpl) = last_page_len_d.device_ptr(&ctx.stream); + let (pos_ptr, _gpos) = positions_d.device_ptr(&ctx.stream); + let (ri_ptr, _gri) = request_indices_d.device_ptr(&ctx.stream); + + let src_stride_n = (num_kv_heads * head_dim) as i64; + let src_stride_h = head_dim as i64; + let result = unsafe { + ffi::paged_kv_scatter_cuda( + buf_ptr as *const ffi::Half, + k_offset, + v_offset, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + k_ptr as *const ffi::Half, + v_ptr as *const ffi::Half, + ri_ptr as *const i32, + pos_ptr as *const i32, + batch_size as i32, + num_kv_heads as i32, + head_dim as i32, + page_size as i32, + stride_page, + src_stride_n, + src_stride_h, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "paged_kv_scatter_cuda ({op_name}) failed for layer {layer}, bs={batch_size}, \ + kv_heads={num_kv_heads}, head_dim={head_dim}, page_size={page_size}: {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub fn paged_attention_batch_decode_hd256_into( + ctx: &DeviceContext, + q: &HiddenStates, + k: &HiddenStates, + v: &HiddenStates, + kv_buffer: &CudaSlice, + layout: &PagedKvLayout, + layer: usize, + page_indices_d: &CudaSlice, + page_indptr_d: &CudaSlice, + last_page_len_d: &CudaSlice, + positions_d: &CudaSlice, + request_indices_d: &CudaSlice, + kv_tile_indices_d: &CudaSlice, + kv_chunk_size_d: &CudaSlice, + output: &mut HiddenStates, + num_qo_heads: usize, + batch_size: usize, +) -> Result<()> { + let num_kv_heads = layout.num_kv_heads; + let head_dim = layout.head_dim; + debug_assert_eq!(head_dim, 256); + let page_size = layout.page_size; + + let PagedGeometry { + k_offset_elems: k_offset, + v_offset_elems: v_offset, + stride_page, + .. + } = checked_paged_geometry( + "batch hd256 decode", + layout, + kv_buffer.len(), + layer, + head_dim, + num_kv_heads, + false, + )?; + + let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); + let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); + let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); + let (pi_ptr, _gpi) = page_indices_d.device_ptr(&ctx.stream); + let (pip_ptr, _gpip) = page_indptr_d.device_ptr(&ctx.stream); + let (lpl_ptr, _glpl) = last_page_len_d.device_ptr(&ctx.stream); + let (ri_ptr, _gri) = request_indices_d.device_ptr(&ctx.stream); + let (kti_ptr, _gkti) = kv_tile_indices_d.device_ptr(&ctx.stream); + let (kcs_ptr, _gkcs) = kv_chunk_size_d.device_ptr(&ctx.stream); + + let stream = crate::tensor::active_cu_stream(ctx); + + scatter_decode_kv_into_paged( + ctx, + k, + v, + kv_buffer, + layout, + layer, + page_indices_d, + page_indptr_d, + last_page_len_d, + positions_d, + request_indices_d, + batch_size, + "batch hd256 decode", + )?; + + let sm_scale = 1.0f32 / (head_dim as f32).sqrt(); + let result = unsafe { + ffi::paged_attention_decode_cuda_hd256( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const ffi::Half, + k_offset, + v_offset, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + ri_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + num_qo_heads as i32, + num_kv_heads as i32, + head_dim as i32, + page_size as i32, + batch_size as i32, + stride_page, + sm_scale, + stream, + ) + }; + if result != 0 { + anyhow::bail!( + "paged_attention_decode_cuda_hd256 (batch) failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + } + + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub fn paged_attention_batch_decode_via_prefill_hd256_into( + ctx: &DeviceContext, + q: &HiddenStates, + k: &HiddenStates, + v: &HiddenStates, + kv_buffer: &CudaSlice, + layout: &PagedKvLayout, + layer: usize, + plan: &PrefillPagedPlan, + positions_d: &CudaSlice, + output: &mut HiddenStates, + num_qo_heads: usize, + batch_size: usize, +) -> Result<()> { + let num_kv_heads = layout.num_kv_heads; + let head_dim = layout.head_dim; + debug_assert_eq!(head_dim, 256); + anyhow::ensure!( + batch_size == plan.total_tokens && batch_size == plan.batch_size() as usize, + "decode-via-prefill plan shape mismatch: bs={batch_size}, total_tokens={}, plan_batch={}", + plan.total_tokens, + plan.batch_size() + ); + + scatter_decode_kv_into_paged( + ctx, + k, + v, + kv_buffer, + layout, + layer, + &plan.page_indices_d, + &plan.page_indptr_d, + &plan.last_page_len_d, + positions_d, + &plan.batch_indices_d, + batch_size, + "batch hd256 decode via prefill", + )?; + + let PagedGeometry { + k_offset_elems: k_offset, + v_offset_elems: v_offset, + stride_page, + .. + } = checked_paged_geometry( + "batch hd256 decode via prefill", + layout, + kv_buffer.len(), + layer, + head_dim, + num_kv_heads, + false, + )?; + let sm_scale = 1.0f32 / (head_dim as f32).sqrt(); + + let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); + let (q_ptr, _gq) = q.data.device_ptr(&ctx.stream); + let (out_ptr, _go) = output.data.device_ptr_mut(&ctx.stream); + let (pi_ptr, _gpi) = plan.page_indices_d.device_ptr(&ctx.stream); + let (pip_ptr, _gpip) = plan.page_indptr_d.device_ptr(&ctx.stream); + let (lpl_ptr, _glpl) = plan.last_page_len_d.device_ptr(&ctx.stream); + let (qi_ptr, _gqi) = plan.q_indptr_d.device_ptr(&ctx.stream); + let (ri_ptr, _gri) = plan.request_indices_d.device_ptr(&ctx.stream); + let (qti_ptr, _gqti) = plan.qo_tile_indices_d.device_ptr(&ctx.stream); + let (kti_ptr, _gkti) = plan.kv_tile_indices_d.device_ptr(&ctx.stream); + let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); + let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); + + let result = unsafe { + ffi::batch_prefill_paged_cuda_hd256( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const ffi::Half, + k_offset, + v_offset, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + qi_ptr as *const i32, + ri_ptr as *const i32, + qti_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + tnr_ptr as *const u32, + num_qo_heads as i32, + num_kv_heads as i32, + head_dim as i32, + layout.page_size as i32, + batch_size as i32, + plan.batch_size(), + plan.num_tiles, + stride_page, + sm_scale, + crate::tensor::active_cu_stream(ctx), + ) + }; + if result != 0 { + anyhow::bail!( + "batch_prefill_paged_cuda_hd256 (decode via prefill) failed for layer {layer}, \ + bs={batch_size}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}: {result}{}", + plan.num_tiles, + crate::ops::ffi_exception_message(result) + ); + } + + Ok(()) +} diff --git a/pegainfer-kernels/src/typed_ops.rs b/pegainfer-kernels/src/typed_ops.rs index fc2d1594e..9eff02b9e 100644 --- a/pegainfer-kernels/src/typed_ops.rs +++ b/pegainfer-kernels/src/typed_ops.rs @@ -8,6 +8,7 @@ use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; use crate::ffi; +use crate::ops::ffi_status_detail; use crate::tensor::DeviceContext; use crate::tensor::DeviceMatrix; use crate::tensor::GpuTensor; @@ -495,13 +496,10 @@ fn launch_gemm( ) }; if status != 0 { - if status >= 100_000 { - anyhow::bail!( - "cuBLAS GEMM failed: cublas_status={}, m={m}, n={n}, k={k}", - status - 100_000 - ); - } - anyhow::bail!("CUDA GEMM launch failed: cuda_status={status}, m={m}, n={n}, k={k}"); + anyhow::bail!( + "GEMM failed: {}, m={m}, n={n}, k={k}", + ffi_status_detail(status) + ); } } Ok(()) diff --git a/pegainfer-kernels/tests/hd256_decode_csr_trap.rs b/pegainfer-kernels/tests/hd256_decode_csr_trap.rs index 55383b456..853b7c8c6 100644 --- a/pegainfer-kernels/tests/hd256_decode_csr_trap.rs +++ b/pegainfer-kernels/tests/hd256_decode_csr_trap.rs @@ -3,6 +3,8 @@ //! //! Manual gate — CI compiles this but never runs it. +#![cfg(feature = "gemma4")] + mod common; use cudarc::driver::CudaSlice; diff --git a/pegainfer-kernels/tests/hd256_fp8_pool.rs b/pegainfer-kernels/tests/hd256_fp8_pool.rs index 6f0323785..422847526 100644 --- a/pegainfer-kernels/tests/hd256_fp8_pool.rs +++ b/pegainfer-kernels/tests/hd256_fp8_pool.rs @@ -8,7 +8,7 @@ use cudarc::driver::CudaSlice; use half::bf16; use pegainfer_kernels::ops::PrefillPagedPlan; use pegainfer_kernels::ops::batch_prefill_paged_window_hd256_into; -use pegainfer_kernels::ops::paged_attention_batch_decode_hd256_into; +use pegainfer_kernels::ops::paged_attention_batch_decode_into; use pegainfer_kernels::ops::qkv_norm_rope_paged_decode_hd256_plain_into; use pegainfer_kernels::ops::qkv_norm_rope_paged_prefill_hd256_plain_into; use pegainfer_kernels::paged_kv::KvStorage; @@ -349,7 +349,7 @@ fn decode_wrapper_without_fp8_twin_refuses_e4m3() { let mut output = HiddenStates::zeros(&ctx, HD, 1).expect("output"); let meta = ctx.stream.clone_htod(&[0i32]).expect("metadata"); let indptr = ctx.stream.clone_htod(&[0i32, 1]).expect("indptr"); - let err = paged_attention_batch_decode_hd256_into( + let err = paged_attention_batch_decode_into( &ctx, &state, &state, diff --git a/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs b/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs index a2a28c44d..e0362531f 100644 --- a/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs +++ b/pegainfer-kernels/tests/hd256_qk_rope_plain_smoke.rs @@ -1,14 +1,15 @@ -//! Device gate for csrc/shared/prefill_attention_hd256_plain.cu. +//! Device gate for csrc/gemma4/prefill_attention_hd256_plain.cu. //! //! Manual gate: CI compiles this but never runs it. Run on a GPU box with //! PEGAINFER_REQUIRE_GPU=1, which turns a missing device into a failure //! rather than a skip. +#![cfg(feature = "gemma4")] + mod common; use cudarc::driver::CudaSlice; use half::bf16; -use pegainfer_kernels::ops::qk_norm_rope_prefill_hd256_plain_into; use pegainfer_kernels::ops::qkv_norm_rope_paged_decode_hd256_plain_into; use pegainfer_kernels::ops::qkv_norm_rope_paged_prefill_hd256_plain_into; use pegainfer_kernels::paged_kv::PagedKvLayout; @@ -17,6 +18,7 @@ use pegainfer_kernels::tensor::DeviceVec; use pegainfer_kernels::tensor::HiddenStates; const HD: usize = 256; +const HALF: usize = HD / 2; const EPS: f32 = 1e-6; // The Gemma 4 12B local-layer geometry: 16 query heads over 8 KV heads. const NUM_Q_HEADS: usize = 16; @@ -44,6 +46,15 @@ const SEQ_LEN: usize = 4; const START_POS: usize = 1; const COS_MAX_POS: usize = 8; +// Pool geometry. The prefill cases map positions 1..=4 through PAGE_INDICES +// to pages 3, 7, 7, 5; page id 9 is an out-of-range sentinel the kernel must +// never dereference. The 8-page pool leaves in-range pages unreferenced, so +// stray writes have somewhere visible to land. +const PAGE_SIZE: usize = 2; +const NUM_LAYERS: usize = 2; +const PAGE_INDICES: [i32; 4] = [3, 7, 5, 9]; +const POOL_PAGES: usize = 8; + fn signed(base: f32, head: usize, token: usize) -> f32 { let tok_sign = if token.is_multiple_of(2) { 1.0 } else { -1.0 }; base * HEAD_SIGNS[head] * tok_sign @@ -70,14 +81,14 @@ fn rope_row(row: usize) -> (f32, f32) { } } -/// Laid out as the kernel indexes them: [pos * rotary_dim + d]. -fn cos_sin_tables(ctx: &DeviceContext, rows: usize, rotary_dim: usize) -> (DeviceVec, DeviceVec) { - let mut cos = Vec::with_capacity(rows * rotary_dim); - let mut sin = Vec::with_capacity(rows * rotary_dim); +/// Laid out as the kernel indexes them: [pos * 256 + d]. +fn cos_sin_tables(ctx: &DeviceContext, rows: usize) -> (DeviceVec, DeviceVec) { + let mut cos = Vec::with_capacity(rows * HD); + let mut sin = Vec::with_capacity(rows * HD); for row in 0..rows { let (c, s) = rope_row(row); - cos.extend(vec![bf16::from_f32(c); rotary_dim]); - sin.extend(vec![bf16::from_f32(s); rotary_dim]); + cos.extend(vec![bf16::from_f32(c); HD]); + sin.extend(vec![bf16::from_f32(s); HD]); } ( DeviceVec::from_host(ctx, &cos).expect("cos H2D"), @@ -85,45 +96,66 @@ fn cos_sin_tables(ctx: &DeviceContext, rows: usize, rotary_dim: usize) -> (Devic ) } -fn expected_prep(x: f32, w: &[bf16], inv: f32, d: usize, row: usize, rotary_dim: usize) -> f32 { - let half_rotary = rotary_dim / 2; - if d < half_rotary { +/// rotary_dim = 256, the Gemma 4 local-layer case: the full head rotates +/// as pairs `(d, d + 128)` and the pass-through tail is empty. +fn expected_prep(x: f32, w: &[bf16], inv: f32, d: usize, row: usize) -> f32 { + if d < HALF { let lo = normed(x, w, inv, d); - let hi = normed(x, w, inv, d + half_rotary); + let hi = normed(x, w, inv, d + HALF); match row % 3 { 0 => lo, 1 => -hi, _ => -lo, } - } else if d < rotary_dim { - let lo = normed(x, w, inv, d - half_rotary); + } else { + let lo = normed(x, w, inv, d - HALF); let hi = normed(x, w, inv, d); match row % 3 { 0 => hi, 1 => lo, _ => -hi, } - } else { - normed(x, w, inv, d) } } -fn expected_full(base: f32, w: &[bf16], dim: usize, rotary_dim: usize) -> Vec { - let num_heads = dim / HD; - let mut full = vec![0.0f32; dim * SEQ_LEN]; +fn expected_q(w: &[bf16]) -> Vec { + let mut full = vec![0.0f32; Q_DIM * SEQ_LEN]; for t in 0..SEQ_LEN { - let row = START_POS + t; - for h in 0..num_heads { - let x = signed(base, h, t); - let inv = inv_rms(x); + for h in 0..NUM_Q_HEADS { + let x = signed(Q_BASE, h, t); for d in 0..HD { - full[t * dim + h * HD + d] = expected_prep(x, w, inv, d, row, rotary_dim); + full[t * Q_DIM + h * HD + d] = expected_prep(x, w, inv_rms(x), d, START_POS + t); } } } full } +/// K and V blocks; everything else stays 0.0. The offsets are derived from +/// the layout, so oracle and kernel share no hand-picked raw offset. V is its +/// own source under its own inv_rms, weightless and un-rotated. +fn expected_pool(w: &[bf16], layout: &PagedKvLayout, layer: usize) -> Vec { + let mut exp = vec![0.0f32; layout.page_stride * POOL_PAGES]; + for t in 0..SEQ_LEN { + let pos = START_POS + t; + let page = PAGE_INDICES[pos / PAGE_SIZE] as usize; + for h in 0..NUM_KV_HEADS { + let base = page * layout.page_stride + + layer * layout.layer_stride + + (pos % PAGE_SIZE) * KV_DIM + + h * HD; + let k = signed(K_BASE, h, t); + let v = signed(V_BASE, h, t); + let v_val = bf16::from_f32(v * inv_rms(v)).to_f32(); + for d in 0..HD { + exp[base + d] = expected_prep(k, w, inv_rms(k), d, pos); + exp[base + layout.kv_block_len + d] = v_val; + } + } + } + exp +} + fn assert_close(got: &[f32], expected: &[f32], what: &str) { assert_eq!(got.len(), expected.len()); for (i, (&g, &e)) in got.iter().zip(expected).enumerate() { @@ -134,8 +166,26 @@ fn assert_close(got: &[f32], expected: &[f32], what: &str) { } } -/// Starts at 1: w[0] = 0 would make dim 0 normalise to 0.0, which an -/// uninitialised output slot could imitate. +/// Exact zero is the assertion, not sloppiness: it marks a slot the kernel +/// must never have written — unreferenced pages, the other layer, and the +/// slots outside the request's positions. +#[allow(clippy::float_cmp)] +fn assert_pool(got: &[f32], expected: &[f32]) { + assert_eq!(got.len(), expected.len()); + for (i, (&g, &e)) in got.iter().zip(expected).enumerate() { + if e == 0.0 { + assert_eq!(g, 0.0, "pool[{i}]: expected untouched, got {g}"); + } else { + assert!( + (g - e).abs() < 0.02, + "pool[{i}]: got {g}, expected {e} (tolerance 0.02)" + ); + } + } +} + +/// Starts at 1: w[0] = 0 would make dim 0 normalise to 0.0, which +/// assert_pool cannot tell from an untouched slot. fn q_norm_weights() -> Vec { (1..=HD).map(|d| bf16::from_f32(d as f32)).collect() } @@ -157,76 +207,67 @@ fn hidden_input(ctx: &DeviceContext, base: f32, num_heads: usize, rows: usize) - HiddenStates::from_host(ctx, &host, num_heads * HD, rows).expect("input H2D") } -fn run_prep(ctx: &DeviceContext, rotary_dim: usize) -> (Vec, Vec) { +#[test] +fn prefill_prep_matches_closed_form() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let ctx = &ctx; let qw = q_norm_weights(); let kw = k_norm_weights(); + let layer = 1; + let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); let q = hidden_input(ctx, Q_BASE, NUM_Q_HEADS, SEQ_LEN); let k = hidden_input(ctx, K_BASE, NUM_KV_HEADS, SEQ_LEN); + let v = hidden_input(ctx, V_BASE, NUM_KV_HEADS, SEQ_LEN); let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); - let mut k_out = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("k_out alloc"); - - let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS, rotary_dim); + let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS); let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); - - qk_norm_rope_prefill_hd256_plain_into( + let pool: CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride * POOL_PAGES) + .expect("pool alloc"); + let page_indices: CudaSlice = ctx + .stream + .clone_htod(&PAGE_INDICES) + .expect("page_indices H2D"); + + qkv_norm_rope_paged_prefill_hd256_plain_into( ctx, &q, &k, + &v, &mut q_out, - &mut k_out, + 0, + &pool, + &layout, &qn, &kn, &cos_dev, &sin_dev, + layer, + &page_indices, + 0, + 0, START_POS, COS_MAX_POS, NUM_Q_HEADS, NUM_KV_HEADS, - rotary_dim, + HD, EPS, ) - .expect("prep launch"); + .expect("prefill prep launch"); let qo = q_out.to_host(ctx).expect("q_out D2H"); - let ko = k_out.to_host(ctx).expect("k_out D2H"); - (qo, ko) -} - -// Pool geometry for the row-offset serving cases. Positions map through -// PAGE_INDICES to pages 3, 7, 5; page id 9 is an out-of-range sentinel the -// kernel must never dereference. The 8-page pool leaves in-range pages -// unreferenced, so stray writes have somewhere visible to land. -const PAGE_SIZE: usize = 2; -const NUM_LAYERS: usize = 2; -const PAGE_INDICES: [i32; 4] = [3, 7, 5, 9]; -const POOL_PAGES: usize = 8; - -/// rotary_dim = 256 is the Gemma 4 local-layer case: the full head rotates -/// and the pass-through tail is empty. -#[test] -fn full_rotation_matches_closed_form() { - let Some(ctx) = common::device_or_skip() else { - return; - }; - let ctx = &ctx; - let qw = q_norm_weights(); - let kw = k_norm_weights(); - let (qo, ko) = run_prep(ctx, HD); - assert_close( - &qo, - &expected_full(Q_BASE, &qw, Q_DIM, HD), - "full-rotation Q pairing", - ); - assert_close( - &ko, - &expected_full(K_BASE, &kw, KV_DIM, HD), - "full-rotation K pairing", - ); + assert_close(&qo, &expected_q(&qw), "prefill Q pairing"); + let pool_host: Vec = ctx.stream.clone_dtoh(&pool).expect("pool D2H"); + let pool_f: Vec = pool_host.iter().map(|x| x.to_f32()).collect(); + assert_pool(&pool_f, &expected_pool(&kw, &layout, layer)); } #[test] -fn rejects_bad_rotary_dim() { +fn prefill_rejects_bad_rotary_dim() { let Some(ctx) = common::device_or_skip() else { return; }; @@ -237,65 +278,104 @@ fn rejects_bad_rotary_dim() { // actually reach the launcher. let q = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q alloc"); let k = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("k alloc"); + let v = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("v alloc"); let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); - let mut k_out = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("k_out alloc"); + let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); + let pool: CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride * POOL_PAGES) + .expect("pool alloc"); let cos_dev = DeviceVec::zeros(ctx, COS_MAX_POS * 512).expect("cos alloc"); let sin_dev = DeviceVec::zeros(ctx, COS_MAX_POS * 512).expect("sin alloc"); let qn = DeviceVec::zeros(ctx, HD).expect("qn alloc"); let kn = DeviceVec::zeros(ctx, HD).expect("kn alloc"); + let page_indices: CudaSlice = ctx + .stream + .clone_htod(&PAGE_INDICES) + .expect("page_indices H2D"); // 127: odd — index 126 would never be written. 512: wider than the - // head — smem and the output slices would walk past 256. Both must be - // rejected, not silently accepted. + // head — smem and the output slices would walk past 256. for bad in [127, 512] { - let err = qk_norm_rope_prefill_hd256_plain_into( + let err = qkv_norm_rope_paged_prefill_hd256_plain_into( ctx, &q, &k, + &v, &mut q_out, - &mut k_out, + 0, + &pool, + &layout, &qn, &kn, &cos_dev, &sin_dev, - 0, // start_pos + 0, // layer + &page_indices, + 0, + 0, + START_POS, COS_MAX_POS, NUM_Q_HEADS, NUM_KV_HEADS, bad, EPS, ); - assert!(err.is_err(), "rotary_dim={bad} must be rejected"); + let Err(e) = err else { + panic!("rotary_dim={bad} must be rejected"); + }; + let msg = format!("{e:#}"); + assert!( + msg.contains("rotary_dim must be"), + "rotary_dim={bad} must be rejected by the launcher's check; got: {msg}" + ); } } #[test] -fn rejects_position_beyond_cos_table() { +fn prefill_rejects_position_beyond_cos_table() { let Some(ctx) = common::device_or_skip() else { return; }; let ctx = &ctx; // Rejected on the host, before any launch: with cos_max_pos 4, - // start_pos 1 + seq_len 4 reaches row 5 — past the tables. + // start_pos 1 + seq_len 4 reaches row 5 — past the tables; the page + // window covers 8 rows, so the position check is the one that fires. let q = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q alloc"); let k = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("k alloc"); + let v = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("v alloc"); let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); - let mut k_out = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("k_out alloc"); + let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); + let pool: CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride * POOL_PAGES) + .expect("pool alloc"); let cos_dev = DeviceVec::zeros(ctx, 4 * HD).expect("cos alloc"); let sin_dev = DeviceVec::zeros(ctx, 4 * HD).expect("sin alloc"); let qn = DeviceVec::zeros(ctx, HD).expect("qn alloc"); let kn = DeviceVec::zeros(ctx, HD).expect("kn alloc"); + let page_indices: CudaSlice = ctx + .stream + .clone_htod(&PAGE_INDICES) + .expect("page_indices H2D"); - let err = qk_norm_rope_prefill_hd256_plain_into( + let err = qkv_norm_rope_paged_prefill_hd256_plain_into( ctx, &q, &k, + &v, &mut q_out, - &mut k_out, + 0, + &pool, + &layout, &qn, &kn, &cos_dev, &sin_dev, + 0, // layer + &page_indices, + 0, + 0, 1, // start_pos 4, // cos_max_pos NUM_Q_HEADS, @@ -303,9 +383,13 @@ fn rejects_position_beyond_cos_table() { HD, EPS, ); + let Err(e) = err else { + panic!("start_pos + seq_len beyond cos_max_pos must be rejected"); + }; + let msg = format!("{e:#}"); assert!( - err.is_err(), - "start_pos + seq_len beyond cos_max_pos must be rejected on the host" + msg.contains("hd256 paged prep start_pos 1 + seq_len 4 > cos_max_pos 4"), + "must be rejected on the host, not by the launcher; got: {msg}" ); } @@ -353,7 +437,7 @@ fn decode_prep_row_offset_serves_only_the_suffix() { suffix(K_BASE, NUM_KV_HEADS), suffix(V_BASE, NUM_KV_HEADS), ); - let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS, HD); + let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS); let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); let pages_d: CudaSlice = ctx.stream.clone_htod(&pages_cat).expect("pages H2D"); @@ -449,7 +533,7 @@ fn prefill_prep_row_offset_serves_only_the_suffix() { let kw = k_norm_weights(); let layer = 1; let layout = PagedKvLayout::new(NUM_LAYERS, NUM_KV_HEADS, HD, PAGE_SIZE); - let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS, HD); + let (cos_dev, sin_dev) = cos_sin_tables(ctx, COS_MAX_POS); let qn = DeviceVec::from_host(ctx, &qw).expect("q_norm_weight H2D"); let kn = DeviceVec::from_host(ctx, &kw).expect("k_norm_weight H2D"); diff --git a/pegainfer-kernels/tests/hd256_window_prefill_reject.rs b/pegainfer-kernels/tests/hd256_window_prefill_reject.rs index c453e0f3c..7756ef5bb 100644 --- a/pegainfer-kernels/tests/hd256_window_prefill_reject.rs +++ b/pegainfer-kernels/tests/hd256_window_prefill_reject.rs @@ -2,6 +2,8 @@ //! length with no device-side bounds check, so a plan's host-side bounds are //! the only bounds. This drives them through the public wrapper. +#![cfg(feature = "gemma4")] + mod common; use cudarc::driver::CudaSlice; diff --git a/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs b/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs index daaa762c9..fab546c07 100644 --- a/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs +++ b/pegainfer-kernels/tests/hd512_qk_rope_smoke.rs @@ -1,17 +1,18 @@ -//! Device gate for csrc/shared/prefill_attention_hd512.cu. +//! Device gate for csrc/gemma4/prefill_attention_hd512.cu. //! //! Manual gate: CI compiles this but never runs it. Run on a GPU box with //! PEGAINFER_REQUIRE_GPU=1, which turns a missing device into a failure //! rather than a skip. Traps are in their own binaries — __trap() poisons //! the context for whatever runs next. +#![cfg(feature = "gemma4")] + mod common; use cudarc::driver::CudaSlice; use half::bf16; use pegainfer_kernels::ops::Hd512DecodeMetadata; use pegainfer_kernels::ops::paged_attention_batch_decode_split_kv_hd512_into; -use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd512_into; use pegainfer_kernels::ops::qk_norm_partial_rope_paged_decode_hd512_into; use pegainfer_kernels::ops::qk_norm_partial_rope_paged_prefill_hd512_into; use pegainfer_kernels::paged_kv::KvFormat; @@ -38,8 +39,6 @@ const PAGE_SIZE: usize = 2; const START_POS: usize = 1; // Positions 1..=4 map to pages 3, 7, 7, 5; page 9 is unreferenced. const PAGE_INDICES: [i32; 4] = [3, 7, 5, 9]; -// Non-monotonic rows distinguish request positions from token indices. -const POSITIONS: [i32; 4] = [0, 1, 3, 1]; // Two layers make a wrong K-layer offset observable. const NUM_LAYERS: usize = 2; const PAGE_STRIDE: i64 = 8 * HD as i64; @@ -384,52 +383,6 @@ fn prefill_prep_folded_row_matches_closed_form() { assert_pool(&pool_f, &pool_exp); } -#[test] -fn rejects_bad_rotary_dim() { - let Some(ctx) = common::device_or_skip() else { - return; - }; - let ctx = &ctx; - // Zeroed buffers suffice: the launcher must reject before touching any - // device memory. Tables are sized for the worst case checked here - // (8 x 1024), so the wrapper's table checks pass and both values - // actually reach the launcher. - let q = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q alloc"); - let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); - let mut k_mut = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("k_mut alloc"); - let cos_dev = DeviceVec::zeros(ctx, 8 * 1024).expect("cos alloc"); - let sin_dev = DeviceVec::zeros(ctx, 8 * 1024).expect("sin alloc"); - let qn = DeviceVec::zeros(ctx, HD).expect("qn alloc"); - let kn = DeviceVec::zeros(ctx, HD).expect("kn alloc"); - let positions: CudaSlice = ctx.stream.clone_htod(&POSITIONS).expect("positions H2D"); - - // 127: odd — index 126 would never be written. 1024: wider than the - // head — smem and the output slices would walk past 512. Both must be - // rejected, not silently accepted. - for bad in [127, 1024] { - let decode_err = qk_norm_partial_rope_batched_decode_hd512_into( - ctx, - &q, - &mut q_out, - &mut k_mut, - &qn, - &kn, - &cos_dev, - &sin_dev, - &positions, - 8, // cos_max_pos - NUM_Q_HEADS, - NUM_KV_HEADS, - bad, - EPS, - ); - assert!( - decode_err.is_err(), - "decode: rotary_dim={bad} must be rejected" - ); - } -} - #[test] fn prefill_rejects_position_beyond_cos_table() { let Some(ctx) = common::device_or_skip() else { diff --git a/pegainfer-kernels/tests/hd512_qk_rope_trap.rs b/pegainfer-kernels/tests/hd512_qk_rope_trap.rs deleted file mode 100644 index 9593fc5e6..000000000 --- a/pegainfer-kernels/tests/hd512_qk_rope_trap.rs +++ /dev/null @@ -1,65 +0,0 @@ -//! The hd512 decode position trap. One trap per binary: `__trap()` leaves -//! the context in a sticky error state, so anything sharing the process -//! afterwards would fail for the wrong reason. -//! -//! Manual gate — CI compiles this but never runs it. - -mod common; - -use cudarc::driver::CudaSlice; -use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd512_into; -use pegainfer_kernels::tensor::DeviceVec; -use pegainfer_kernels::tensor::HiddenStates; - -const HD: usize = 512; -const ROTARY_DIM: usize = 128; -const EPS: f32 = 1e-6; -const NUM_Q_HEADS: usize = 2; -const NUM_KV_HEADS: usize = 2; -const Q_DIM: usize = NUM_Q_HEADS * HD; -const KV_DIM: usize = NUM_KV_HEADS * HD; -const SEQ_LEN: usize = 4; - -/// The launch is async, so the trap may surface at the launcher or at the -/// next sync — either counts; a silent `Ok` does not. -#[test] -fn decode_position_trap_is_visible() { - let Some(ctx) = common::device_or_skip() else { - return; - }; - let ctx = &ctx; - // Inputs are zeroed but every wrapper check must PASS so the launch - // actually happens and the kernel traps: weights [512], tables 4 rows - // (cos_max_pos 4), positions [0, 1, 5, 1] — token 1 reads row 5. - let q = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q alloc"); - let mut q_out = HiddenStates::zeros(ctx, Q_DIM, SEQ_LEN).expect("q_out alloc"); - let mut k = HiddenStates::zeros(ctx, KV_DIM, SEQ_LEN).expect("k alloc"); - let qn = DeviceVec::zeros(ctx, HD).expect("qn alloc"); - let kn = DeviceVec::zeros(ctx, HD).expect("kn alloc"); - let cos_dev = DeviceVec::zeros(ctx, 4 * ROTARY_DIM).expect("cos alloc"); - let sin_dev = DeviceVec::zeros(ctx, 4 * ROTARY_DIM).expect("sin alloc"); - let positions: CudaSlice = ctx.stream.clone_htod(&[0, 1, 5, 1]).expect("positions H2D"); - - let res = qk_norm_partial_rope_batched_decode_hd512_into( - ctx, - &q, - &mut q_out, - &mut k, - &qn, - &kn, - &cos_dev, - &sin_dev, - &positions, - 4, // cos_max_pos - NUM_Q_HEADS, - NUM_KV_HEADS, - ROTARY_DIM, - EPS, - ); - if res.is_ok() { - assert!( - ctx.sync().is_err(), - "out-of-range position (5 >= cos_max_pos 4) must fail visibly" - ); - } -} diff --git a/pegainfer-kernels/tests/hd512_qk_rope_trap_page.rs b/pegainfer-kernels/tests/hd512_qk_rope_trap_page.rs index 6b994dcac..6c02f54d7 100644 --- a/pegainfer-kernels/tests/hd512_qk_rope_trap_page.rs +++ b/pegainfer-kernels/tests/hd512_qk_rope_trap_page.rs @@ -1,8 +1,11 @@ -//! The hd512 paged-KV page-id trap. One trap per binary, for the reason -//! given in hd512_qk_rope_trap.rs. +//! The hd512 paged-KV page-id trap. One trap per binary: `__trap()` leaves +//! the context in a sticky error state, so anything sharing the process +//! afterwards would fail for the wrong reason. //! //! Manual gate — CI compiles this but never runs it. +#![cfg(feature = "gemma4")] + mod common; use cudarc::driver::CudaSlice; diff --git a/pegainfer-kimi-k2/src/runner/worker.rs b/pegainfer-kimi-k2/src/runner/worker.rs index 1f1dd2c1f..ab7c5192d 100644 --- a/pegainfer-kimi-k2/src/runner/worker.rs +++ b/pegainfer-kimi-k2/src/runner/worker.rs @@ -35,6 +35,7 @@ use pegainfer_kernels::ops::KIMI_O_PROJ_CUBLASLT_INPUT; use pegainfer_kernels::ops::KimiMarlinRouteWorkspace; use pegainfer_kernels::ops::KimiMarlinWna16Workspace; use pegainfer_kernels::ops::KimiMlaPagedKvLayout; +use pegainfer_kernels::ops::ffi_status_detail; use pegainfer_kernels::ops::flashinfer_top1_row_states_bytes; use pegainfer_kernels::ops::kimi_flashinfer_batch_decode_mla_rt; use pegainfer_kernels::ops::kimi_flashinfer_single_prefill_mla_rt; @@ -680,37 +681,25 @@ fn bind_rank_thread( pegainfer_kernels::ffi::cublas_init(); let status = pegainfer_kernels::ffi::kimi_shared_gate_up_cublaslt_init_cuda(); if status != 0 { - if status >= 100_000 { - anyhow::bail!( - "Kimi shared_gate_up cuBLASLt init failed: cublas_status={}", - status - 100_000 - ); - } anyhow::bail!( - "Kimi shared_gate_up cuBLASLt init failed: cuda_status={}", - status + "Kimi shared_gate_up cuBLASLt init failed: {}", + ffi_status_detail(status) ); } let status = pegainfer_kernels::ffi::kimi_mla_cublaslt_init_cuda(); if status != 0 { - if status >= 100_000 { - anyhow::bail!( - "Kimi MLA cuBLASLt init failed: cublas_status={}", - status - 100_000 - ); - } - anyhow::bail!("Kimi MLA cuBLASLt init failed: cuda_status={}", status); + anyhow::bail!( + "Kimi MLA cuBLASLt init failed: {}", + ffi_status_detail(status) + ); } if local_dims.o_proj_in == KIMI_O_PROJ_CUBLASLT_INPUT { let status = pegainfer_kernels::ffi::kimi_o_proj_cublaslt_init_cuda(); if status != 0 { - if status >= 100_000 { - anyhow::bail!( - "Kimi o_proj cuBLASLt init failed: cublas_status={}", - status - 100_000 - ); - } - anyhow::bail!("Kimi o_proj cuBLASLt init failed: cuda_status={}", status); + anyhow::bail!( + "Kimi o_proj cuBLASLt init failed: {}", + ffi_status_detail(status) + ); } } } diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index d13329642..b93710115 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -149,7 +149,7 @@ impl Qwen35Model { &bufs.k_attn, &bufs.v_attn, kv_buffer, - layout, + &layout.kernel_layout(), layer_idx, &bufs.page_indices_d, &bufs.page_indptr_d, @@ -226,7 +226,7 @@ impl Qwen35Model { &bufs.k_attn, &bufs.v_attn, kv_buffer, - layout, + &layout.kernel_layout(), layer_idx, plan, &bufs.positions_d, @@ -862,6 +862,7 @@ impl Qwen35Model { geom.local_linear_num_value_heads(), self.config.linear_value_head_dim, self.config.rms_norm_eps, + ops::GatedNormActivation::Silu, ); ops::gemm_into( &self.ctx, diff --git a/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index 425dce5f5..7155e9634 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -33,7 +33,7 @@ use std::path::Path; use anyhow::Result; use anyhow::anyhow; pub(crate) use config::probe_config_json; -use pegainfer_frontend::engine::EngineHandle; +use pegainfer_frontend::engine::Engine; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::EpBackend; pub use scheduler::DEFAULT_MAX_PREFILL_TOKENS; @@ -107,7 +107,7 @@ pub fn start_engine( options: EngineLoadOptions, max_batch: usize, max_prefill_tokens: usize, -) -> Result { +) -> Result { start_engine_with_capacity_and_policy( model_path, options, @@ -168,7 +168,7 @@ pub fn launch_with_options_policy_and_overlap( options: Qwen35LaunchOptions, scheduler_policy: Qwen35SchedulerPolicy, decode_overlap: Qwen35DecodeOverlap, -) -> Result { +) -> Result { let device_ordinals = options.device_ordinals()?; start_engine_with_capacity_policy_and_overlap( model_path, @@ -192,7 +192,7 @@ pub fn start_engine_with_capacity( options: EngineLoadOptions, max_batch: usize, max_prefill_tokens: usize, -) -> Result { +) -> Result { start_engine_with_capacity_and_policy( model_path, options, @@ -210,7 +210,7 @@ pub(crate) fn start_engine_with_capacity_and_policy( max_prefill_tokens: usize, scheduler_policy: Qwen35SchedulerPolicy, prefix_cache_mib: usize, -) -> Result { +) -> Result { start_engine_with_capacity_policy_and_overlap( model_path, options, @@ -230,7 +230,7 @@ pub fn start_engine_with_capacity_policy_and_overlap( scheduler_policy: Qwen35SchedulerPolicy, decode_overlap: Qwen35DecodeOverlap, prefix_cache_mib: usize, -) -> Result { +) -> Result { anyhow::ensure!( (1..=MAX_DECODE_BATCH).contains(&max_batch), "Qwen3.5 max_batch must be in 1..={MAX_DECODE_BATCH}, got {max_batch}" diff --git a/pegainfer-qwen35/src/model_line.rs b/pegainfer-qwen35/src/model_line.rs index f6952170a..ca3fa220e 100644 --- a/pegainfer-qwen35/src/model_line.rs +++ b/pegainfer-qwen35/src/model_line.rs @@ -179,7 +179,7 @@ impl ModelLine for Qwen35Line { resolve_decode_overlap(ctx.shared.decode_overlap) .map_err(|error| anyhow::anyhow!(error.to_string()))?, ) - .map(LaunchedEngine::Handle) + .map(LaunchedEngine::Stepped) } } diff --git a/pegainfer-qwen35/src/ops.rs b/pegainfer-qwen35/src/ops.rs index af08deb6b..c922f9c6f 100644 --- a/pegainfer-qwen35/src/ops.rs +++ b/pegainfer-qwen35/src/ops.rs @@ -1,6 +1,7 @@ //! Qwen3.5 GPU operation wrappers. pub(crate) use pegainfer_core::ops::GEMM_LT_MAX_N; +pub(crate) use pegainfer_core::ops::GatedNormActivation; pub(crate) use pegainfer_core::ops::PrefillPagedPlan; pub(crate) use pegainfer_core::ops::SuppressIds; pub(crate) use pegainfer_core::ops::add_batch; @@ -10,14 +11,14 @@ pub(crate) use pegainfer_core::ops::extract_vec; pub(crate) use pegainfer_core::ops::gemm; pub(crate) use pegainfer_core::ops::gemm_into; pub(crate) use pegainfer_core::ops::gemm_lt_tune; -pub(crate) use pegainfer_core::ops::paged_attention_batch_decode_hd256_into; -pub(crate) use pegainfer_core::ops::paged_attention_batch_decode_via_prefill_hd256_into; -pub(crate) use pegainfer_core::ops::qk_norm_partial_rope_batched_decode_hd256_into; pub use pegainfer_core::ops::rms_norm_batch_offset_into; pub(crate) use pegainfer_core::ops::rms_norm_gated_batch_into; pub use pegainfer_core::ops::rms_norm_offset_into; pub(crate) use pegainfer_core::ops::silu_mul_fused_batch_into; pub(crate) use pegainfer_core::ops::write_vec_into; +pub(crate) use pegainfer_kernels::ops::paged_attention_batch_decode_hd256_into; +pub(crate) use pegainfer_kernels::ops::paged_attention_batch_decode_via_prefill_hd256_into; +pub(crate) use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd256_into; pub(crate) use recurrent::conv1d_decode_batch_into; pub(crate) use recurrent::conv1d_prefill_batch_into; pub(crate) use recurrent::gated_delta_rule_decode_batch_into; diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs index 6b95fa34b..8edc7746c 100644 --- a/pegainfer-qwen35/src/prefill.rs +++ b/pegainfer-qwen35/src/prefill.rs @@ -504,6 +504,7 @@ impl Qwen35Model { geom.local_linear_num_value_heads(), c.linear_value_head_dim, c.rms_norm_eps, + ops::GatedNormActivation::Silu, ); *linear_idx += 1; diff --git a/pegainfer-qwen35/src/scheduler/backend.rs b/pegainfer-qwen35/src/scheduler/backend.rs index 9cf74b261..052b7acbf 100644 --- a/pegainfer-qwen35/src/scheduler/backend.rs +++ b/pegainfer-qwen35/src/scheduler/backend.rs @@ -171,7 +171,7 @@ impl SingleGpuBackend { anyhow::bail!("single-GPU commit received TP chunk state") }; for (i, (kv, rec)) in kvs.iter_mut().zip(recs.iter()).enumerate() { - let is_final = chunk.ends[i] == chunk.reqs[i].prompt_tokens.len(); + let is_final = chunk.ends[i] == chunk.reqs[i].request.prompt_tokens.len(); let boundary = self .kv_cache .apply_prefill(kv, is_final.then_some(tokens[i]))?; @@ -287,7 +287,7 @@ impl SingleGpuBackend { self.kv_cache.pool().available_blocks() } - pub(super) fn active_joint_prefix_pages(&self, req: &SchedulerRequest) -> usize { + pub(super) fn active_joint_prefix_pages(&self, req: &Request) -> usize { self.kv_cache .active_joint_prefix_pages(&req.prompt_tokens, req.lora_adapter.as_deref()) } @@ -302,7 +302,7 @@ impl SingleGpuBackend { pub(super) fn alloc_prefill_state( &mut self, - req: &SchedulerRequest, + req: &Request, ) -> Result<(PrefillBackendState, usize), AdmissionError> { let mut rec = self .alloc_recurrent() @@ -492,7 +492,7 @@ impl SingleGpuBackend { pub(super) fn sample_prefill_logits( &mut self, - pending: &[SchedulerRequest], + pending: &[QueuedRequest], logits: &HiddenStates, sample_seed: u64, ) -> Result<(Vec, Vec>)> { @@ -501,10 +501,11 @@ impl SingleGpuBackend { pending.len(), "Qwen3.5 prefill logits rows must preserve pending request order" ); - let requested_logprobs: Vec> = pending.iter().map(|r| r.logprobs).collect(); + let requested_logprobs: Vec> = + pending.iter().map(|r| r.request.logprobs).collect(); let cpu_logits = snapshot_requested_logprobs(self.model.device_ctx(), logits, &requested_logprobs)?; - let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect(); + let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.request.params).collect(); let tokens = self.model.select_tokens_from_logits_varied( logits, &mut self.graph_state.buffers, @@ -604,7 +605,7 @@ impl SingleGpuBackend { impl TpSchedulerBackend { pub(super) fn alloc_prefill_state( &mut self, - req: &SchedulerRequest, + req: &Request, ) -> Result<(PrefillBackendState, usize), AdmissionError> { let request_id = self.alloc_request_id(); let cached_tokens = self @@ -620,7 +621,7 @@ impl TpSchedulerBackend { Ok((PrefillBackendState::Tp { request_id }, cached_tokens)) } - pub(super) fn active_joint_prefix_pages(&self, req: &SchedulerRequest) -> usize { + pub(super) fn active_joint_prefix_pages(&self, req: &Request) -> usize { self.executor .active_joint_prefix_pages(&req.prompt_tokens, req.lora_adapter.as_deref()) } @@ -784,7 +785,7 @@ impl SchedulerBackend { } } - pub(super) fn active_joint_prefix_pages(&self, req: &SchedulerRequest) -> usize { + pub(super) fn active_joint_prefix_pages(&self, req: &Request) -> usize { match self { Self::Single(backend) => backend.active_joint_prefix_pages(req), Self::Tp(backend) => backend.active_joint_prefix_pages(req), @@ -840,7 +841,7 @@ impl SchedulerBackend { pub(super) fn alloc_prefill_state( &mut self, - req: &SchedulerRequest, + req: &Request, ) -> Result<(PrefillBackendState, usize), AdmissionError> { match self { Self::Single(backend) => backend.alloc_prefill_state(req), diff --git a/pegainfer-qwen35/src/scheduler/mod.rs b/pegainfer-qwen35/src/scheduler/mod.rs index c08263fde..98424c93f 100644 --- a/pegainfer-qwen35/src/scheduler/mod.rs +++ b/pegainfer-qwen35/src/scheduler/mod.rs @@ -1,6 +1,6 @@ -//! Scheduler for Qwen3.5: dedicated GPU thread that batches concurrent requests. +//! Qwen3.5 step scheduler: batches requests on a dedicated GPU owner thread. //! -//! Mirrors the Qwen3 scheduler but manages: +//! Retains model-specific state alongside the frontend request ledger: //! - `RecurrentState` alongside content-hashed `RequestKv` (hybrid attention) //! - `BatchDecodeGraphState` for CUDA Graph batch decode (stable-address slots) @@ -9,7 +9,6 @@ mod plan; mod tp; use std::collections::HashMap; use std::collections::HashSet; -use std::collections::VecDeque; use std::sync::Arc; use std::sync::OnceLock; use std::sync::mpsc as std_mpsc; @@ -26,24 +25,29 @@ use log::debug; use log::info; use log::warn; use pegainfer_core::tensor::HiddenStates; -use pegainfer_frontend::engine::EngineHandle as SchedulerHandle; +use pegainfer_frontend::engine::Engine; +use pegainfer_frontend::engine::EngineInfo; use pegainfer_frontend::engine::FinishReason; -use pegainfer_frontend::engine::GenerateRequest as SchedulerRequest; use pegainfer_frontend::engine::KvCapacity; +use pegainfer_frontend::engine::LiveScheduler; +use pegainfer_frontend::engine::QueuedRequest; +use pegainfer_frontend::engine::RejectReason as EngineRejectReason; +use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::RequestId as FrontendRequestId; +use pegainfer_frontend::engine::RequestLedger; +use pegainfer_frontend::engine::Scheduler; use pegainfer_frontend::engine::SchedulerMetrics; -use pegainfer_frontend::engine::SubmittedRequest; -use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenLogprob; -use pegainfer_frontend::engine::TokenSink; +use pegainfer_frontend::engine::drive; use pegainfer_frontend::engine::panic_message; +use pegainfer_frontend::engine::scheduler_pair; +use pegainfer_frontend::engine::spawn_scheduler; use pegainfer_frontend::sampler::SamplingParams; use pegainfer_kv_cache::KvCacheManager; use pegainfer_kv_cache::KvView; use pegainfer_kv_cache::RequestKv; use rand::SeedableRng; use rand::rngs::StdRng; -use tokio::sync::mpsc; -use tokio::sync::watch; use self::backend::*; use self::plan::ActiveDecodeState; @@ -84,11 +88,10 @@ use crate::weights::Qwen35Model; /// An in-flight request being decoded. Recurrent state lives in the /// `BatchDecodeGraphState` at `graph_slot_idx` — NOT owned here. struct ActiveRequest35 { - request_id: Option, - token_tx: TokenSink, + id: FrontendRequestId, + client_label: Option>, backend_state: ActiveBackendState, last_token: u32, - generated_count: usize, max_tokens: usize, prompt_len: usize, params: SamplingParams, @@ -100,7 +103,7 @@ struct ActiveRequest35 { /// It owns its growing KV and recurrent state until the prompt is exhausted, /// at which point it is promoted into the decode batch. struct PrefillingRequest35 { - req: SchedulerRequest, + req: QueuedRequest, backend_state: PrefillBackendState, /// Prompt tokens prefilled so far. cursor: usize, @@ -132,100 +135,6 @@ enum PrefillBackendState { }, } -struct TerminalRequest { - token_tx: TokenSink, - prompt_tokens: usize, - completion_tokens: usize, -} - -impl TerminalRequest { - fn send_error(self, message: &str) { - let _ = self.token_tx.send(TokenEvent::Error { - message: message.to_string(), - prompt_tokens: self.prompt_tokens, - completion_tokens: self.completion_tokens, - }); - } -} - -impl From for TerminalRequest { - fn from(req: SchedulerRequest) -> Self { - Self { - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - token_tx: req.token_tx, - } - } -} - -impl From for TerminalRequest { - fn from(req: ActiveRequest35) -> Self { - Self { - token_tx: req.token_tx, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - } - } -} - -impl From for TerminalRequest { - fn from(req: PrefillingRequest35) -> Self { - req.req.into() - } -} - -struct PrefillCompletionRequest { - req: SchedulerRequest, - backend_state: PrefillBackendState, -} - -trait CompletionRequest { - fn token_tx(&self) -> &TokenSink; - fn into_terminal(self) -> TerminalRequest; -} - -impl CompletionRequest for ActiveRequest35 { - fn token_tx(&self) -> &TokenSink { - &self.token_tx - } - - fn into_terminal(self) -> TerminalRequest { - self.into() - } -} - -impl CompletionRequest for PrefillCompletionRequest { - fn token_tx(&self) -> &TokenSink { - &self.req.token_tx - } - - fn into_terminal(self) -> TerminalRequest { - self.req.into() - } -} - -struct CompletionCandidate { - request: R, - final_events: Vec, -} - -impl CompletionCandidate { - fn commit(self) { - for event in self.final_events { - let _ = self.request.token_tx().send(event); - } - } - - fn into_terminal(self) -> TerminalRequest { - self.request.into_terminal() - } -} - -struct FatalSchedulerError { - message: String, - transient: Vec, -} - #[derive(Clone, Debug, PartialEq)] struct PrefillArtifact { token: u32, @@ -265,29 +174,6 @@ impl PrefillStepArtifacts { } } -impl FatalSchedulerError { - fn new(message: impl Into) -> Self { - Self { - message: message.into(), - transient: Vec::new(), - } - } - - fn with_request(mut self, request: impl Into) -> Self { - self.transient.push(request.into()); - self - } - - fn with_requests(mut self, requests: I) -> Self - where - I: IntoIterator, - R: Into, - { - self.transient.extend(requests.into_iter().map(Into::into)); - self - } -} - pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; /// Env-gated per-step ITL diagnostics (issue #470). When `PEGAINFER_ITL_DEBUG` @@ -343,7 +229,7 @@ pub fn start_with_capacity( seed: u64, max_batch: usize, max_prefill_tokens: usize, -) -> Result { +) -> Result { start_with_capacity_and_policy( model, seed, @@ -361,7 +247,7 @@ pub(crate) fn start_with_capacity_and_policy( max_prefill_tokens: usize, scheduler_policy: Qwen35SchedulerPolicy, decode_overlap: Qwen35DecodeOverlap, -) -> Result { +) -> Result { assert!( max_prefill_tokens > 0, "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" @@ -370,7 +256,6 @@ pub(crate) fn start_with_capacity_and_policy( // still uses the current page budget inside the scheduler loop. let backend = SingleGpuBackend::new(model, max_batch, decode_overlap)?; let total_blocks = backend.capacity_pages_for_requests(); - let kv_total_blocks = total_blocks as u64; let block_size = backend.page_size(); let servable = servable_len( backend.model().config().max_position_embeddings, @@ -378,25 +263,22 @@ pub(crate) fn start_with_capacity_and_policy( block_size, ); - let (submit_tx, submit_rx) = mpsc::unbounded_channel(); + let (handle, wiring) = scheduler_pair(); let (startup_tx, startup_rx) = std_mpsc::channel(); - let (load_tx, load_rx) = watch::channel(SchedulerMetrics { - kv_total_blocks, - ..SchedulerMetrics::default() - }); let join_handle = thread::Builder::new() .name("scheduler-qwen35".into()) .spawn(move || match bind_model_thread(backend.model()) { Ok(_guard) => { let _ = startup_tx.send(Ok(())); - scheduler_loop( - SchedulerBackend::Single(backend), - submit_rx, - seed, - max_prefill_tokens, - scheduler_policy, - load_tx, + drive( + Qwen35Scheduler::new( + SchedulerBackend::Single(backend), + seed, + max_prefill_tokens, + scheduler_policy, + ), + wiring, ); } Err(err) => { @@ -416,15 +298,20 @@ pub(crate) fn start_with_capacity_and_policy( let _ = join_handle.join(); return Err(err); } - Ok( - SchedulerHandle::new_with_join_handle(submit_tx, join_handle) - .with_servable_len(servable) - .with_kv_capacity(KvCapacity { + Ok(Engine { + schedulers: vec![LiveScheduler { + handle, + join: join_handle, + }], + info: EngineInfo { + servable_len: Some(servable), + kv_capacity: Some(KvCapacity { total_blocks, block_size, - }) - .with_metrics_watch(load_rx), - ) + }), + }, + lora: None, + }) } pub(crate) fn start_tp_with_capacity( @@ -435,7 +322,7 @@ pub(crate) fn start_tp_with_capacity( max_prefill_tokens: usize, enable_cuda_graph: bool, prefix_snapshot_bytes: usize, -) -> Result { +) -> Result { assert!( max_prefill_tokens > 0, "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" @@ -458,31 +345,22 @@ pub(crate) fn start_tp_with_capacity( block_size: backend.page_size(), }; - let (submit_tx, submit_rx) = mpsc::unbounded_channel(); - let (load_tx, load_rx) = watch::channel(SchedulerMetrics { - kv_total_blocks: kv_capacity.total_blocks as u64, - ..SchedulerMetrics::default() - }); - let join_handle = thread::Builder::new() - .name("scheduler-qwen35-tp".into()) - .spawn(move || { - scheduler_loop( + Ok(Engine { + schedulers: vec![spawn_scheduler( + "scheduler-qwen35-tp", + Qwen35Scheduler::new( SchedulerBackend::Tp(backend), - submit_rx, seed, max_prefill_tokens, Qwen35SchedulerPolicy::Off, - load_tx, - ); - }) - .expect("failed to spawn Qwen3.5 TP scheduler thread"); - - Ok( - SchedulerHandle::new_with_join_handle(submit_tx, join_handle) - .with_servable_len(servable) - .with_kv_capacity(kv_capacity) - .with_metrics_watch(load_rx), - ) + ), + )], + info: EngineInfo { + servable_len: Some(servable), + kv_capacity: Some(kv_capacity), + }, + lora: None, + }) } fn servable_len(max_context: usize, max_pages: usize, page_size: usize) -> u32 { @@ -520,177 +398,150 @@ fn bind_model_thread(model: &Qwen35Model) -> Result { unsafe { crate::ffi::cublas_init(); } + let guard = CublasThreadGuard; model.tune_decode_gemm_algos()?; - Ok(CublasThreadGuard) + Ok(guard) } // ── Main loop ─────────────────────────────────────────────────────────── -fn publish_load( - load_tx: &watch::Sender, - backend: &SchedulerBackend, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - inflight_prefill_reqs: usize, - num_waiting_reqs: usize, -) { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - let (num_running_reqs, num_waiting_reqs) = - logical_load_counts(active, prefilling, inflight_prefill_reqs, num_waiting_reqs); - load_tx.send_replace(SchedulerMetrics { - kv_used_blocks: kv_total_blocks.saturating_sub(backend.available_pages() as u64), - kv_total_blocks, - num_running_reqs, - num_waiting_reqs, - spec_decode: None, - }); -} - -fn logical_load_counts( - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - inflight_prefill_reqs: usize, - num_waiting_reqs: usize, -) -> (u64, u64) { - ( - (active.len() + prefilling.len() + inflight_prefill_reqs) as u64, - num_waiting_reqs as u64, - ) -} - -fn should_block_on_submit(owned_work_empty: bool, inflight_prefill: bool) -> bool { - owned_work_empty && !inflight_prefill -} - -fn terminal_scheduler_shutdown( - submit_rx: &mut mpsc::UnboundedReceiver, - load_tx: &watch::Sender, - kv_total_blocks: u64, - active: Vec, - prefilling: Vec, - pending: Vec, - deferred: Vec, - inflight_prefill: Option, - failure: FatalSchedulerError, -) { - submit_rx.close(); - - let mut requests = failure.transient; - requests.extend(active.into_iter().map(Into::into)); - requests.extend(prefilling.into_iter().map(Into::into)); - requests.extend(pending.into_iter().map(Into::into)); - requests.extend(deferred.into_iter().map(Into::into)); - if let Some(InflightPrefill { output, chunk, .. }) = inflight_prefill { - // The stream must drain before the chunk's KV/recurrent/conv state is - // released or transferred into terminal request ownership. - drop(output); - requests.extend(chunk.reqs.into_iter().map(Into::into)); - } - while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { - requests.push(req.into()); +fn reject_or_retire(id: FrontendRequestId, reason: EngineRejectReason, ledger: &mut RequestLedger) { + if ledger.is_aborted(id) { + ledger.retire(id); + } else { + ledger.reject(id, reason); } +} - warn!( - "Qwen3.5 TP scheduler terminating after replica failure: {}", - failure.message - ); - for request in requests { - request.send_error(&failure.message); +fn fail_or_retire(id: FrontendRequestId, message: impl Into, ledger: &mut RequestLedger) { + if ledger.is_aborted(id) { + ledger.retire(id); + } else { + ledger.fail(id, message); } - load_tx.send_replace(SchedulerMetrics { - kv_used_blocks: 0, - kv_total_blocks, - num_running_reqs: 0, - num_waiting_reqs: 0, - spec_decode: None, - }); } fn prune_closed_requests( backend: &mut B, active: &mut Vec, prefilling: &mut Vec, - pending: &mut Vec, -) -> std::result::Result<(), FatalSchedulerError> + pending: &mut Vec, + ledger: &mut RequestLedger, +) -> Result<()> where B: DecodeDispatchBackend + PrefillPromoteBackend, { - pending.retain(|req| !req.token_tx.is_closed()); - + pending.retain(|req| { + if ledger.is_aborted(req.id) { + ledger.retire(req.id); + false + } else { + true + } + }); for idx in (0..active.len()).rev() { - if active[idx].token_tx.is_closed() { + if ledger.is_aborted(active[idx].id) { debug!( "request pruned before scheduling: request_id={:?} phase=decode tokens_generated={}", - active[idx].request_id, active[idx].generated_count + active[idx].client_label, + ledger.completion_tokens(active[idx].id) ); let removed = backend.take_active_request(active, idx); - if let Err(err) = backend.drop_active_state(&removed.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); - } + backend.drop_active_state(&removed.backend_state)?; + ledger.retire(removed.id); } } - for idx in (0..prefilling.len()).rev() { - if prefilling[idx].req.token_tx.is_closed() { + if ledger.is_aborted(prefilling[idx].req.id) { let removed = prefilling.remove(idx); debug!( "request pruned before scheduling: request_id={:?} phase=prefill cursor={}", - removed.req.request_id, removed.cursor + removed.req.request.client_label, removed.cursor ); - // Admission now creates rank-local recurrent state, even on a cold miss. + // Admission creates rank-local recurrent state, even on a cold miss. let expectation = DropExpectation::MustExist; - if let Err(err) = backend.drop_prefill_state(&removed.backend_state, expectation) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); - } + backend.drop_prefill_state(&removed.backend_state, expectation)?; + ledger.retire(removed.req.id); } } Ok(()) } -const UNSUPPORTED_PROMPT_LOGPROBS_MESSAGE: &str = - "prompt_logprobs is unsupported by the Qwen3.5 serving contract"; - -fn reject_unsupported_prompt_logprobs(pending: &mut Vec) { +fn reject_unsupported_prompt_logprobs( + pending: &mut Vec, + ledger: &mut RequestLedger, +) { pending.retain(|req| { - if req.prompt_logprobs.is_none() { + if req.request.prompt_logprobs.is_none() { return true; } - let _ = req.token_tx.send(TokenEvent::Rejected { - message: UNSUPPORTED_PROMPT_LOGPROBS_MESSAGE.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); + reject_or_retire( + req.id, + EngineRejectReason::Unsupported { + feature: "prompt_logprobs".to_owned(), + }, + ledger, + ); false }); } -#[allow(clippy::needless_pass_by_value)] -fn scheduler_loop( - mut backend: SchedulerBackend, - mut submit_rx: mpsc::UnboundedReceiver, - seed: u64, +struct Qwen35Scheduler { + // Drain asynchronous work before releasing its request allocations or backend. + inflight_prefill: Option, + active: Vec, + prefilling: Vec, + pending: Vec, + rng: StdRng, prefill_budget: usize, scheduler_policy: Qwen35SchedulerPolicy, - load_tx: watch::Sender, -) { - let mut rng = StdRng::seed_from_u64(seed); - let mut active: Vec = Vec::new(); - let mut deferred: Vec = Vec::new(); - let mut prefilling: Vec = Vec::new(); - let mut inflight_prefill: Option = None; - let max_batch = backend.max_batch(); - let decode_overlap = matches!( - &backend, - SchedulerBackend::Single(single) if single.overlap_enabled() - ); + failed: bool, + backend: SchedulerBackend, +} + +impl Qwen35Scheduler { + fn new( + backend: SchedulerBackend, + seed: u64, + prefill_budget: usize, + scheduler_policy: Qwen35SchedulerPolicy, + ) -> Self { + info!("scheduler ready (max_batch={})", backend.max_batch()); + Self { + inflight_prefill: None, + active: Vec::new(), + prefilling: Vec::new(), + pending: Vec::new(), + rng: StdRng::seed_from_u64(seed), + prefill_budget, + scheduler_policy, + failed: false, + backend, + } + } - info!("scheduler ready (max_batch={})", max_batch); + fn advance(&mut self, ledger: &mut RequestLedger) -> Result<()> { + let Self { + inflight_prefill, + active, + prefilling, + pending, + rng, + prefill_budget, + scheduler_policy, + backend, + .. + } = self; + let decode_overlap = matches!( + backend, + SchedulerBackend::Single(single) if single.overlap_enabled() + ); + let max_batch = backend.max_batch(); - loop { - if inflight_prefill + let prefill_ready = inflight_prefill .as_mut() - .is_some_and(|prefill| prefill.output.is_ready()) - { + .is_some_and(|prefill| prefill.output.is_ready()); + if prefill_ready { let (prefill_tokens, prefill_reqs) = inflight_prefill.as_ref().map_or((0, 0), |prefill| { ( @@ -701,12 +552,13 @@ fn scheduler_loop( let decode_n = active.len(); let step_start = itl_debug_enabled().then(Instant::now); let finish_result = finish_async_prefill( - &mut backend, - &mut active, - &mut prefilling, + backend, + active, + prefilling, inflight_prefill .take() .expect("ready async prefill must still be present"), + ledger, ); log_itl_step( step_start, @@ -715,111 +567,21 @@ fn scheduler_loop( prefill_reqs, decode_n, ); - if let Err(failure) = finish_result { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - Vec::new(), - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - } - - // 1. Merge deferred work with every submission currently available. - let mut pending = std::mem::take(&mut deferred); - while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { - pending.push(req); + finish_result?; } - // 2. Remove closed work before metrics, admission, or planning. Active - // and prefilling cleanup goes through the backend's normal retirement - // paths so graph slots and TP request state are released consistently. - if let Err(failure) = - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) - { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - pending, - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - reject_unsupported_prompt_logprobs(&mut pending); - - // 3. Publish the settled post-prune state. Requests accepted from the - // channel are waiting until admission below; closed requests never - // appear in this snapshot or consume its KV/slot accounting. - publish_load( - &load_tx, - &backend, - &active, - &prefilling, - inflight_prefill - .as_ref() - .map_or(0, |prefill| prefill.chunk.reqs.len()), - pending.len(), - ); + prune_closed_requests(backend, active, prefilling, pending, ledger)?; + reject_unsupported_prompt_logprobs(pending, ledger); - // 4. Nothing in flight and nothing pending: the idle snapshot above is - // already visible, so block until work arrives. Drain and prune again - // after wakeup because the first request may already be closed and more - // submissions may have raced with the blocking receive. - if should_block_on_submit( - active.is_empty() && prefilling.is_empty() && pending.is_empty(), - inflight_prefill.is_some(), - ) { - if let Some((req, _kv_prefix)) = submit_rx.blocking_recv() { - pending.push(req); - } else { - info!("scheduler: all handles dropped, exiting"); - return; - } - while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { - pending.push(req); - } - if let Err(failure) = - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) - { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - pending, - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - reject_unsupported_prompt_logprobs(&mut pending); - publish_load(&load_tx, &backend, &active, &prefilling, 0, pending.len()); - if pending.is_empty() { - continue; - } + if prefill_ready { + // Commit completed prefill updates before starting another prefill or decode. + return Ok(()); } // One async prefill owns its scheduled request state. Do not admit or // launch a second chunk until it resolves. Active decode keeps moving; - // if it retires first, wait on the event instead of blocking on submit. + // if it retires first, wait for the prefill event to complete. if inflight_prefill.is_some() { - deferred = pending; let itl_step_start = itl_debug_enabled().then(Instant::now); let (itl_prefill_tokens, itl_prefill_reqs) = inflight_prefill.as_ref().map_or((0, 0), |prefill| { @@ -831,16 +593,17 @@ fn scheduler_loop( let itl_decode_n = active.len(); let (itl_plan_kind, step_result) = if active.is_empty() { let result = finish_async_prefill( - &mut backend, - &mut active, - &mut prefilling, + backend, + active, + prefilling, inflight_prefill .take() .expect("async prefill must be present before blocking wait"), + ledger, ); ("overlap_wait", result) } else { - let result = decode_step(&mut backend, &mut active, &mut rng); + let result = decode_step(backend, active, rng, ledger); ("overlap_decode", result) }; log_itl_step( @@ -850,43 +613,32 @@ fn scheduler_loop( itl_prefill_reqs, itl_decode_n, ); - if let Err(failure) = step_result { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - Vec::new(), - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - continue; + return step_result; } - // 5. Admit new prompts. In-flight prefills reserve their promotion slot - // and future KV growth, so shrink the slot/page budgets accordingly + if pending.is_empty() && active.is_empty() && prefilling.is_empty() { + return Ok(()); + } + + // Admit new prompts. In-flight prefills reserve their promotion slot + // and future KV growth, so shrink the slot/page budgets accordingly. let active_budget: Vec = active .iter() .map(|req| ActiveKvBudget { prompt_len: req.prompt_len, - generated_count: req.generated_count, + generated_count: ledger.completion_tokens(req.id), max_tokens: req.max_tokens, }) .collect(); let page_size = backend.page_size(); - let mut admission_pending = pending; + let mut admission_pending = std::mem::take(pending); loop { let prefilling_budget: Vec = prefilling .iter() .map(|p| PrefillKvBudget { current_tokens: p.cursor, - prompt_len: p.req.prompt_tokens.len(), - max_tokens: p.req.max_tokens, + prompt_len: p.req.request.prompt_tokens.len(), + max_tokens: p.req.request.max_tokens, }) .collect(); let page_budget = backend @@ -903,107 +655,81 @@ fn scheduler_loop( // construction, so a real request can use at most the remaining pages. backend.capacity_pages_for_requests(), backend.max_position_embeddings(), - |req| req.prompt_tokens.len(), - |req| req.max_tokens, - |req| backend.active_joint_prefix_pages(req), + |req| req.request.prompt_tokens.len(), + |req| req.request.max_tokens, + |req| backend.active_joint_prefix_pages(&req.request), ); for (rejected, reason) in &admission.rejected { - send_rejection(rejected, *reason); + send_rejection(rejected, *reason, ledger); } - // 6. Materialize admitted prompts. A restore can move cached KV out + // Materialize admitted prompts. A restore can move cached KV out // of the available pool, so re-run admission for FIFO followers. let mut restored_prefix = false; - let mut admitted = admission.pending.into_iter(); - while let Some(req) = admitted.next() { + for req in admission.pending { debug!( "request admitted: request_id={:?} prompt_len={} max_tokens={}", - req.request_id, - req.prompt_tokens.len(), - req.max_tokens + req.request.client_label, + req.request.prompt_tokens.len(), + req.request.max_tokens ); - match backend.alloc_prefill_state(&req) { + match backend.alloc_prefill_state(&req.request) { Ok((backend_state, cached_tokens)) => { restored_prefix |= cached_tokens > 0; - let scheduled_at_unix_s = unix_now_s(); - let _ = req.token_tx.send(TokenEvent::Scheduled { - queued_at_unix_s: req.queued_at_unix_s.unwrap_or(scheduled_at_unix_s), - scheduled_at_unix_s, - prompt_tokens: req.prompt_tokens.len(), - cached_tokens, - }); + ledger.admit(req.id); + ledger.set_cached_tokens(req.id, cached_tokens); prefilling.push(PrefillingRequest35 { + req, backend_state, cursor: cached_tokens, step_chunk: 0, - req, }); } Err(AdmissionError::Recoverable(error)) => { warn!("failed to admit new request: {error}"); - let _ = req.token_tx.send(TokenEvent::Error { - message: error.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - } - Err(AdmissionError::Fatal(error)) => { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - let failure = FatalSchedulerError::new(error.to_string()) - .with_request(req) - .with_requests(admitted); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - Vec::new(), - admission.deferred, - inflight_prefill.take(), - failure, - ); - return; + fail_or_retire(req.id, error.to_string(), ledger); } + Err(AdmissionError::Fatal(error)) => return Err(error), } } - if admission.deferred.is_empty() || !restored_prefix { - deferred = admission.deferred; + *pending = admission.deferred; break; } admission_pending = admission.deferred; } - // 7. Choose this tick's prefill budget, take that chunk off the front of - // the queue, then dispatch by plan. Auto can return 0 for a short - // decode-priority tick; the next iteration reconsiders the same FIFO - // prefill without reordering it. + // Choose this tick's prefill budget, take that chunk off the front of + // the queue, then dispatch by plan. Auto can return 0 for a short + // decode-priority tick; the next iteration reconsiders the same FIFO + // prefill without reordering it. let active_decode: Vec = active .iter() .map(|req| ActiveDecodeState { - generated_count: req.generated_count, + generated_count: ledger.completion_tokens(req.id), max_tokens: req.max_tokens, }) .collect(); let prefill_queue: Vec = prefilling .iter() .map(|req| PrefillQueueState { - remaining_tokens: req.req.prompt_tokens.len().saturating_sub(req.cursor), + remaining_tokens: req + .req + .request + .prompt_tokens + .len() + .saturating_sub(req.cursor), }) .collect(); let step_prefill_budget = choose_prefill_budget( - scheduler_policy, - prefill_budget, + *scheduler_policy, + *prefill_budget, &active_decode, &prefill_queue, decode_overlap, ); - let scheduled = take_prefill_chunks( - &mut prefilling, - step_prefill_budget, - backend.snapshot_stride(), - ); + let scheduled = + take_prefill_chunks(prefilling, step_prefill_budget, backend.snapshot_stride()); // ITL diagnostics (#470): capture the *actual* prefill-chunk token count // and the frozen decode width for this step before the plan consumes the // scheduled set. Off unless PEGAINFER_ITL_DEBUG is set. @@ -1023,31 +749,15 @@ fn scheduler_loop( let step_result = match plan { ExecutionPlan::Unified { pending } => { if decode_overlap { - launch_overlap_step( - &mut backend, - &mut active, - pending, - &mut inflight_prefill, - &mut rng, - ) + launch_overlap_step(backend, active, pending, inflight_prefill, rng, ledger) } else { - unified_step_sched( - &mut backend, - &mut active, - pending, - &mut prefilling, - &mut rng, - ) + unified_step_sched(backend, active, pending, prefilling, rng, ledger) } } - ExecutionPlan::Prefill { pending } => prefill_batch( - &mut backend, - &mut active, - pending, - &mut prefilling, - &mut rng, - ), - ExecutionPlan::Decode => decode_step(&mut backend, &mut active, &mut rng), + ExecutionPlan::Prefill { pending } => { + prefill_batch(backend, active, pending, prefilling, rng, ledger) + } + ExecutionPlan::Decode => decode_step(backend, active, rng, ledger), }; log_itl_step( itl_step_start, @@ -1056,48 +766,71 @@ fn scheduler_loop( itl_prefill_reqs, itl_decode_n, ); - if let Err(failure) = step_result { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - Vec::new(), - deferred, - inflight_prefill.take(), - failure, - ); - return; - } + step_result?; } + Ok(()) } } -fn send_rejection(req: &SchedulerRequest, reason: RejectReason) { - let message = match reason { - RejectReason::ZeroMaxTokens => "max_tokens must be greater than zero".to_owned(), - RejectReason::ContextLength { limit } => format!( - "request exceeds this model's maximum context length of {limit} tokens: requested {} (prompt={} + max_tokens={})", - req.prompt_tokens.len().saturating_add(req.max_tokens), - req.prompt_tokens.len(), - req.max_tokens - ), - RejectReason::KvBudget => { - let max_request_tokens = - plan::request_lifetime_tokens(req.prompt_tokens.len(), req.max_tokens); - format!( - "request requires more KV pages than this model instance can provide: prompt_tokens={}, max_request_tokens={max_request_tokens}", - req.prompt_tokens.len() - ) +impl Scheduler for Qwen35Scheduler { + fn submit(&mut self, request: QueuedRequest) { + self.pending.push(request); + } + + fn step(&mut self, ledger: &mut RequestLedger) -> Result<()> { + let result = self.advance(ledger); + if result.is_err() { + // The driver answers every open account after this returns. + drop(self.inflight_prefill.take()); + self.active.clear(); + self.prefilling.clear(); + self.pending.clear(); + self.failed = true; } + result + } + + fn metrics(&self) -> SchedulerMetrics { + let kv_total_blocks = self.backend.capacity_pages_for_requests() as u64; + let inflight_prefill_reqs = self + .inflight_prefill + .as_ref() + .map_or(0, |prefill| prefill.chunk.reqs.len()); + SchedulerMetrics { + kv_used_blocks: if self.failed { + // TP pages stay owned until backend teardown; zero is display-only. + 0 + } else { + kv_total_blocks.saturating_sub(self.backend.available_pages() as u64) + }, + kv_total_blocks, + num_running_reqs: (self.active.len() + self.prefilling.len() + inflight_prefill_reqs) + as u64, + num_waiting_reqs: self.pending.len() as u64, + spec_decode: None, + } + } +} + +fn send_rejection(req: &QueuedRequest, reason: RejectReason, ledger: &mut RequestLedger) { + let reason = match reason { + RejectReason::ZeroMaxTokens => EngineRejectReason::Unsupported { + feature: "max_tokens=0".to_owned(), + }, + RejectReason::ContextLength { limit } => EngineRejectReason::ContextLength { + prompt_tokens: req.request.prompt_tokens.len(), + max_tokens: req.request.max_tokens, + limit, + }, + RejectReason::KvBudget => EngineRejectReason::KvBudget { + prompt_tokens: req.request.prompt_tokens.len(), + worst_case_tokens: plan::request_lifetime_tokens( + req.request.prompt_tokens.len(), + req.request.max_tokens, + ), + }, }; - let _ = req.token_tx.send(TokenEvent::Rejected { - message, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); + reject_or_retire(req.id, reason, ledger); } // ── Batch prefill ─────────────────────────────────────────────────────── @@ -1108,7 +841,8 @@ fn prefill_batch( scheduled: Vec, prefilling: &mut Vec, rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { let mut chunk = ScheduledChunk::from(scheduled); let sample_seed = rand::RngExt::random(rng); let artifacts = match backend { @@ -1119,23 +853,19 @@ fn prefill_batch( Ok(v) => v, Err(e) => { warn!("batch prefill failed: {e}"); - fail_chunk(chunk, &e.to_string()); + fail_chunk(chunk, &e.to_string(), ledger); return Ok(()); } }; let prefill_sample_seed = rand::RngExt::random(rng); match single.sample_prefill_logits(&chunk.reqs, &logits, prefill_sample_seed) { Ok((tokens, logprobs)) => { - if let Err(error) = single.apply_prefill(&mut chunk, &tokens) { - return Err( - FatalSchedulerError::new(error.to_string()).with_requests(chunk.reqs) - ); - } + single.apply_prefill(&mut chunk, &tokens)?; PrefillStepArtifacts::Single { tokens, logprobs } } Err(e) => { warn!("prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); + fail_chunk(chunk, &e.to_string(), ledger); return Ok(()); } } @@ -1144,12 +874,12 @@ fn prefill_batch( Ok(v) => PrefillStepArtifacts::Tp(v), Err(e) => { warn!("TP prefill chunk failed: {e}"); - return Err(FatalSchedulerError::new(e.to_string()).with_requests(chunk.reqs)); + return Err(e); } }, }; - promote_or_requeue(backend, active, prefilling, chunk, &artifacts) + promote_or_requeue(backend, active, prefilling, chunk, &artifacts, ledger) } fn launch_overlap_step( @@ -1158,7 +888,8 @@ fn launch_overlap_step( scheduled: Vec, inflight_prefill: &mut Option, rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { debug_assert!(inflight_prefill.is_none()); let mut chunk = ScheduledChunk::from(scheduled); let decode_seed = rand::RngExt::random(rng); @@ -1177,10 +908,10 @@ fn launch_overlap_step( } Err(err) => { warn!("async prefill launch failed: {err}"); - fail_chunk(chunk, &err.to_string()); + fail_chunk(chunk, &err.to_string(), ledger); } } - decode_step_with_seed(backend, active, decode_seed) + decode_step_with_seed(backend, active, decode_seed, ledger) } fn finish_async_prefill( @@ -1188,7 +919,8 @@ fn finish_async_prefill( active: &mut Vec, prefilling: &mut Vec, inflight: InflightPrefill, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { let InflightPrefill { mut chunk, output, @@ -1202,15 +934,13 @@ fn finish_async_prefill( Ok(result) => result, Err(err) => { warn!("async prefill sampling failed: {err}"); - fail_chunk(chunk, &err.to_string()); + fail_chunk(chunk, &err.to_string(), ledger); return Ok(()); } }; - if let Err(error) = single.apply_prefill(&mut chunk, &tokens) { - return Err(FatalSchedulerError::new(error.to_string()).with_requests(chunk.reqs)); - } + single.apply_prefill(&mut chunk, &tokens)?; let artifacts = PrefillStepArtifacts::Single { tokens, logprobs }; - promote_or_requeue(single, active, prefilling, chunk, &artifacts) + promote_or_requeue(single, active, prefilling, chunk, &artifacts, ledger) } // ── Unified step (prefill chunk + decode in one forward pass) ────────────── @@ -1221,7 +951,8 @@ fn unified_step_sched( scheduled: Vec, prefilling: &mut Vec, rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { let mut chunk = ScheduledChunk::from(scheduled); if matches!(backend, SchedulerBackend::Tp(_)) { // Preserve the established scheduler RNG order: decode seed first, @@ -1239,19 +970,15 @@ fn unified_step_sched( Ok(artifacts) => artifacts, Err(err) => { warn!("TP unified step failed: {err}"); - return Err(FatalSchedulerError::new(err.to_string()).with_requests(chunk.reqs)); + return Err(err); } }; let (decode_tokens, decode_logprobs) = split_decode_artifacts(&artifacts.decode); - if let Err(failure) = - dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs) - { - return Err(failure.with_requests(chunk.reqs)); - } + dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs, ledger)?; let prefill = PrefillStepArtifacts::Tp(artifacts.prefill); - return promote_or_requeue(backend, active, prefilling, chunk, &prefill); + return promote_or_requeue(backend, active, prefilling, chunk, &prefill, ledger); } let SchedulerBackend::Single(backend) = backend else { @@ -1265,14 +992,8 @@ fn unified_step_sched( Err(e) => { warn!("unified step failed: {e}"); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - fail_chunk(chunk, &message); + fail_active(active, &message, ledger); + fail_chunk(chunk, &message, ledger); return Ok(()); } }; @@ -1281,10 +1002,8 @@ fn unified_step_sched( // Process decode results FIRST (it may retire requests and free graph slots // that promotion then fills densely). - if output.decoded - && let Err(failure) = process_decode_logits(backend, active, decode_seed) - { - return Err(failure.with_requests(chunk.reqs)); + if output.decoded { + process_decode_logits(backend, active, decode_seed, ledger)?; } let prefill_logits = output @@ -1296,15 +1015,13 @@ fn unified_step_sched( Ok(v) => v, Err(e) => { warn!("unified prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); + fail_chunk(chunk, &e.to_string(), ledger); return Ok(()); } }; - if let Err(error) = backend.apply_prefill(&mut chunk, &tokens) { - return Err(FatalSchedulerError::new(error.to_string()).with_requests(chunk.reqs)); - } + backend.apply_prefill(&mut chunk, &tokens)?; let prefill = PrefillStepArtifacts::Single { tokens, logprobs }; - promote_or_requeue(backend, active, prefilling, chunk, &prefill) + promote_or_requeue(backend, active, prefilling, chunk, &prefill, ledger) } // ── Decode step (pure decode, CUDA Graph enabled) ────────────────────── @@ -1313,7 +1030,8 @@ fn decode_step( backend: &mut SchedulerBackend, active: &mut Vec, rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { // Preserve the historical scheduler RNG sequence: TP consumes the first // seed, while single-GPU decode consumed a second seed inside sampling. let first_seed = rand::RngExt::random(rng); @@ -1322,26 +1040,21 @@ fn decode_step( } else { first_seed }; - decode_step_with_seed(backend, active, sample_seed) + decode_step_with_seed(backend, active, sample_seed, ledger) } fn decode_step_with_seed( backend: &mut SchedulerBackend, active: &mut Vec, sample_seed: u64, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { let (tokens, logprobs_vec) = match backend { SchedulerBackend::Single(single) => { if let Err(e) = single.decode_graph(active) { warn!("batch_decode_graph error: {e}"); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } + fail_active(active, &message, ledger); return Ok(()); } // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) @@ -1350,13 +1063,7 @@ fn decode_step_with_seed( Err(e) => { warn!("decode sampling/logprobs error: {e}"); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } + fail_active(active, &message, ledger); return Ok(()); } } @@ -1365,17 +1072,15 @@ fn decode_step_with_seed( Ok(v) => split_decode_artifacts(&v), Err(e) => { warn!("TP eager decode error: {e}"); - return Err(FatalSchedulerError::new(e.to_string())); + return Err(e); } }, }; if let SchedulerBackend::Single(single) = backend { - single - .apply_decode(active, &tokens) - .map_err(|e| FatalSchedulerError::new(e.to_string()))?; + single.apply_decode(active, &tokens)?; } - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec, ledger) } /// Process decode logits from unified step: sample, extract logprobs, dispatch. @@ -1383,30 +1088,23 @@ fn process_decode_logits( backend: &mut SingleGpuBackend, active: &mut Vec, sample_seed: u64, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, sample_seed) { Ok(v) => v, Err(e) => { warn!("decode sampling/logprobs error: {e}"); let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } + fail_active(active, &message, ledger); return Ok(()); } }; - backend - .apply_decode(active, &tokens) - .map_err(|e| FatalSchedulerError::new(e.to_string()))?; - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) + backend.apply_decode(active, &tokens)?; + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec, ledger) } -/// Dispatch sampled decode tokens: send events, check EOS/limits, retire finished. +/// Dispatch sampled decode tokens and retire completed backend state. /// /// `tokens` and `logprobs` are indexed by original position in `active`. /// Retirements collected first, then compacted in reverse order. @@ -1415,103 +1113,48 @@ fn dispatch_decode_tokens( active: &mut Vec, tokens: &[u32], logprobs: &[Option], -) -> std::result::Result<(), FatalSchedulerError> { - enum Retirement { - Completion(Vec), - CleanupOnly, - Disconnected, - } - - let n = active.len(); + ledger: &mut RequestLedger, +) -> Result<()> { let mut to_retire = Vec::new(); - - for i in 0..n { + for (i, req) in active.iter_mut().enumerate() { + if ledger.is_aborted(req.id) { + to_retire.push((i, None)); + continue; + } let token = tokens[i]; - let logprob = logprobs[i].clone(); - let req = &mut active[i]; - req.generated_count += 1; - - let is_eos = !req.params.ignore_eos && backend.is_stop_token(token); - let at_limit = req.generated_count >= req.max_tokens; - - if is_eos { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - req.prompt_len, - req.generated_count, - FinishReason::Stop - ); - let event = TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }; - if backend.completion_requires_drop_ack() { - to_retire.push((i, Retirement::Completion(vec![event]))); - } else { - let _ = req.token_tx.send(event); - to_retire.push((i, Retirement::CleanupOnly)); - } - } else if at_limit { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - req.prompt_len, - req.generated_count, - FinishReason::Length - ); - let events = vec![ - TokenEvent::Token { id: token, logprob }, - TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }, - ]; - if backend.completion_requires_drop_ack() { - to_retire.push((i, Retirement::Completion(events))); - } else { - for event in events { - let _ = req.token_tx.send(event); - } - to_retire.push((i, Retirement::CleanupOnly)); - } - } else if req - .token_tx - .send(TokenEvent::Token { id: token, logprob }) - .is_err() - { - debug!( - "request dropped: client disconnected: request_id={:?} tokens_generated={}", - req.request_id, req.generated_count - ); - to_retire.push((i, Retirement::Disconnected)); + let finish = if !req.params.ignore_eos && backend.is_stop_token(token) { + Some(FinishReason::Stop) + } else if ledger.completion_tokens(req.id) + 1 >= req.max_tokens { + Some(FinishReason::Length) + } else { + None + }; + if let Some(reason) = finish { + to_retire.push((i, Some(reason))); } else { + ledger.push_tokens(req.id, &[token], std::slice::from_ref(&logprobs[i])); req.last_token = token; } } - // Remove in reverse order so compact_slot indices stay valid - for (i, retirement) in to_retire.into_iter().rev() { + // Keep slots dense and wait for retirement ACK before publishing completion. + for (i, finish) in to_retire.into_iter().rev() { let request = backend.take_active_request(active, i); - match retirement { - Retirement::Completion(final_events) => { - let candidate = CompletionCandidate { - request, - final_events, - }; - if let Err(err) = backend.drop_active_state(&candidate.request.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()) - .with_request(candidate.into_terminal())); - } - candidate.commit(); - } - Retirement::CleanupOnly | Retirement::Disconnected => { - if let Err(err) = backend.drop_active_state(&request.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(request)); - } + backend.drop_active_state(&request.backend_state)?; + if ledger.is_aborted(request.id) { + ledger.retire(request.id); + } else if let Some(reason) = finish { + if reason == FinishReason::Length { + ledger.push_tokens(request.id, &[tokens[i]], std::slice::from_ref(&logprobs[i])); } + debug!( + "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", + request.client_label, + request.prompt_len, + ledger.completion_tokens(request.id), + reason + ); + ledger.finish(request.id, reason); } } Ok(()) @@ -1519,7 +1162,6 @@ fn dispatch_decode_tokens( trait DecodeDispatchBackend { fn is_stop_token(&self, token: u32) -> bool; - fn completion_requires_drop_ack(&self) -> bool; fn take_active_request( &mut self, active: &mut Vec, @@ -1533,10 +1175,6 @@ impl DecodeDispatchBackend for SingleGpuBackend { self.is_stop_token(token) } - fn completion_requires_drop_ack(&self) -> bool { - false - } - fn take_active_request( &mut self, active: &mut Vec, @@ -1555,10 +1193,6 @@ impl DecodeDispatchBackend for SchedulerBackend { self.is_stop_token(token) } - fn completion_requires_drop_ack(&self) -> bool { - matches!(self, SchedulerBackend::Tp(_)) - } - fn take_active_request( &mut self, active: &mut Vec, @@ -1609,7 +1243,7 @@ fn compact_single_slot( /// Step's scheduled prefill set struct ScheduledChunk { - reqs: Vec, + reqs: Vec, backend_state: ScheduledChunkBackendState, /// Prompt cursor after this step's chunk ends: Vec, @@ -1627,8 +1261,7 @@ struct InflightPrefill { enum ScheduledChunkBackendState { Single { - // Keep request allocations stable while an async prefill owns the chunk; - // failure cleanup can then move each request back through one ownership path. + // Keep request allocations stable while an async prefill owns the chunk. #[allow(clippy::vec_box)] kvs: Vec>, recs: Vec, @@ -1663,7 +1296,7 @@ impl From> for ScheduledChunk { let end = p.cursor + p.step_chunk; chunk .windows - .push(p.req.prompt_tokens[p.cursor..end].to_vec()); + .push(p.req.request.prompt_tokens[p.cursor..end].to_vec()); chunk.ends.push(end); chunk.reqs.push(p.req); match (&mut chunk.backend_state, p.backend_state) { @@ -1694,7 +1327,7 @@ fn take_prefill_chunks( ) -> Vec { let remaining: Vec = prefilling .iter() - .map(|p| p.req.prompt_tokens.len() - p.cursor) + .map(|p| p.req.request.prompt_tokens.len() - p.cursor) .collect(); let chunks = plan_prefill_chunks(&remaining, prefill_budget); let mut scheduled: Vec = prefilling.drain(0..chunks.len()).collect(); @@ -1705,46 +1338,45 @@ fn take_prefill_chunks( } /// Report a forward/sampling failure to every request in the failed chunk. -fn fail_chunk(chunk: ScheduledChunk, message: &str) { +fn fail_chunk(chunk: ScheduledChunk, message: &str, ledger: &mut RequestLedger) { for req in chunk.reqs { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); + fail_or_retire(req.id, message, ledger); + } +} + +fn fail_active(active: &mut Vec, message: &str, ledger: &mut RequestLedger) { + for req in active.drain(..) { + fail_or_retire(req.id, message, ledger); } } -/// For each request in the just-prefilled chunk: if its prompt is now exhausted, -/// sample its first token, emit events, and move it into the decode batch; -/// otherwise re-queue it (with an advanced cursor) at the FRONT of `prefilling`. -/// `artifacts` are indexed by request order in `chunk`. +/// Commit immediate completions after retirement, and put unfinished chunks +/// back at the front of the prefill queue without changing their FIFO order. fn promote_or_requeue( backend: &mut impl PrefillPromoteBackend, active: &mut Vec, prefilling: &mut Vec, chunk: ScheduledChunk, artifacts: &PrefillStepArtifacts, -) -> std::result::Result<(), FatalSchedulerError> { + ledger: &mut RequestLedger, +) -> Result<()> { let ScheduledChunk { reqs, backend_state, ends, .. } = chunk; - let mut still_prefilling: Vec = Vec::new(); + let mut still_prefilling = Vec::new(); let backend_states = split_scheduled_backend_state(backend_state); - let mut entries: VecDeque<_> = reqs - .into_iter() - .zip(backend_states) - .zip(ends) - .enumerate() - .map(|(i, ((req, backend_state), end))| (i, req, backend_state, end)) - .collect(); - - while let Some((i, req, backend_state, end)) = entries.pop_front() { - // Not finished: re-queue with the advanced cursor - if end < req.prompt_tokens.len() { + for (i, ((req, backend_state), end)) in + reqs.into_iter().zip(backend_states).zip(ends).enumerate() + { + if ledger.is_aborted(req.id) { + backend.drop_prefill_state(&backend_state, DropExpectation::MustExist)?; + ledger.retire(req.id); + continue; + } + if end < req.request.prompt_tokens.len() { still_prefilling.push(PrefillingRequest35 { req, backend_state, @@ -1754,133 +1386,45 @@ fn promote_or_requeue( continue; } - let prompt_len = req.prompt_tokens.len(); let artifact = artifacts.final_artifact(i); let first_token = artifact.token; - let logprob = artifact.logprob; - - if !req.params.ignore_eos && backend.is_stop_token(first_token) { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - prompt_len, - 0, - FinishReason::Stop - ); - let candidate = CompletionCandidate { - request: PrefillCompletionRequest { req, backend_state }, - final_events: vec![TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: prompt_len, - completion_tokens: 0, - }], - }; - if let Err(err) = backend - .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - candidate.into_terminal(), - still_prefilling, - entries, - )); - } - candidate.commit(); - continue; - } - - if req.max_tokens <= 1 { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - prompt_len, - 1, - FinishReason::Length - ); - let candidate = CompletionCandidate { - request: PrefillCompletionRequest { req, backend_state }, - final_events: vec![ - TokenEvent::Token { - id: first_token, - logprob, - }, - TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: prompt_len, - completion_tokens: 1, - }, - ], - }; - if let Err(err) = backend - .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - candidate.into_terminal(), - still_prefilling, - entries, - )); - } - candidate.commit(); - continue; - } - - if req - .token_tx - .send(TokenEvent::Token { - id: first_token, - logprob, - }) - .is_err() - { - debug!( - "request dropped: client disconnected: request_id={:?} tokens_generated={}", - req.request_id, 0 - ); - let removed = PrefillCompletionRequest { req, backend_state }; - if let Err(err) = - backend.drop_prefill_state(&removed.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - removed.into_terminal(), - still_prefilling, - entries, - )); + let finish = if !req.request.params.ignore_eos && backend.is_stop_token(first_token) { + Some(FinishReason::Stop) + } else if req.request.max_tokens <= 1 { + Some(FinishReason::Length) + } else { + None + }; + if let Some(reason) = finish { + backend.drop_prefill_state(&backend_state, DropExpectation::MustExist)?; + if ledger.is_aborted(req.id) { + ledger.retire(req.id); + } else { + if reason == FinishReason::Length { + ledger.push_tokens(req.id, &[first_token], &[artifact.logprob]); + } + ledger.finish(req.id, reason); } continue; } let active_backend_state = backend.promote_prefill_state(active.len(), backend_state); + ledger.push_tokens(req.id, &[first_token], &[artifact.logprob]); active.push(ActiveRequest35 { - request_id: req.request_id, - token_tx: req.token_tx, + id: req.id, + client_label: req.request.client_label, backend_state: active_backend_state, last_token: first_token, - generated_count: 1, - max_tokens: req.max_tokens, - prompt_len, - params: req.params, - logprobs: req.logprobs, + max_tokens: req.request.max_tokens, + prompt_len: req.request.prompt_tokens.len(), + params: req.request.params, + logprobs: req.request.logprobs, }); } - prefilling.splice(0..0, still_prefilling); Ok(()) } -fn prefill_lifecycle_failure( - message: String, - current: TerminalRequest, - still_prefilling: Vec, - remaining: VecDeque<(usize, SchedulerRequest, PrefillBackendState, usize)>, -) -> FatalSchedulerError { - FatalSchedulerError::new(message) - .with_request(current) - .with_requests(still_prefilling) - .with_requests(remaining.into_iter().map(|(_, req, _, _)| req)) -} - trait PrefillPromoteBackend { fn is_stop_token(&self, token: u32) -> bool; fn promote_prefill_state( @@ -1993,11 +1537,6 @@ fn active_request_kv(request: &mut ActiveRequest35) -> Option<&mut RequestKv> { ActiveBackendState::Tp { .. } => None, } } -fn unix_now_s() -> f64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_or(0.0, |duration| duration.as_secs_f64()) -} fn clamp_prefill_chunk(cursor: usize, chunk: usize, snapshot_stride: Option) -> usize { snapshot_stride.map_or(chunk, |stride| { debug_assert!(stride > 0); diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 2f48f50ab..c59980720 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -1,61 +1,108 @@ -use std::sync::Arc; -use std::sync::Barrier; use std::time::Duration; -use std::time::Instant; -use super::*; +use pegainfer_frontend::engine::RequestUpdate; +use pegainfer_frontend::engine::StopPolicy; +use pegainfer_frontend::engine::Terminal; -fn test_request(request_id: &str, token_tx: TokenSink) -> SchedulerRequest { - test_request_with_shape(request_id, token_tx, vec![1], 1) -} +use super::*; -fn test_request_with_shape( - request_id: &str, - token_tx: TokenSink, - prompt_tokens: Vec, - max_tokens: usize, -) -> SchedulerRequest { - SchedulerRequest { - trace_parent: None, - request_id: Some(request_id.to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, +fn test_request(label: &str, prompt_tokens: Vec, max_tokens: usize) -> Request { + Request { prompt_tokens, params: SamplingParams { ignore_eos: true, ..SamplingParams::default() }, + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, - token_tx, logprobs: None, prompt_logprobs: None, + trace_parent: None, + client_label: Some(Arc::from(label)), } } -fn active_request(request_id: u64, label: &str, token_tx: TokenSink) -> ActiveRequest35 { +// The production driver owns registration and commit even in these tests. +fn run_step(requests: Vec, aborted: &[usize], step: F) -> Vec +where + F: FnOnce(Vec, &mut RequestLedger) -> Result<()> + Send, +{ + struct TestStep { + pending: Vec, + step: Option, + } + impl Scheduler for TestStep + where + F: FnOnce(Vec, &mut RequestLedger) -> Result<()> + Send, + { + fn submit(&mut self, request: QueuedRequest) { + self.pending.push(request); + } + fn step(&mut self, ledger: &mut RequestLedger) -> Result<()> { + self.step.take().expect("one test step")(std::mem::take(&mut self.pending), ledger) + } + fn metrics(&self) -> SchedulerMetrics { + SchedulerMetrics::default() + } + } + let (mut handle, wiring) = scheduler_pair(); + let mut steps = handle.take_steps().unwrap(); + let controls: Vec<_> = requests + .into_iter() + .map(|request| handle.submit(request)) + .collect(); + for &index in aborted { + controls[index].abort(); + } + drop(handle); + drive( + TestStep { + pending: Vec::new(), + step: Some(step), + }, + wiring, + ); + let mut updates = Vec::new(); + while let Ok(step) = steps.try_recv() { + updates.extend(step.updates); + } + updates +} + +fn active_request( + req: QueuedRequest, + worker_id: u64, + ledger: &mut RequestLedger, +) -> ActiveRequest35 { + ledger.admit(req.id); + ledger.push_tokens(req.id, &[1], &[None]); ActiveRequest35 { - request_id: Some(label.to_string()), - token_tx, + id: req.id, + client_label: req.request.client_label, backend_state: ActiveBackendState::Tp { - request_id: RequestId::new(request_id), + request_id: RequestId::new(worker_id), slot_idx: 0, }, last_token: 1, - generated_count: 1, - max_tokens: 8, - prompt_len: 1, - params: SamplingParams::default(), - logprobs: None, + max_tokens: req.request.max_tokens, + prompt_len: req.request.prompt_tokens.len(), + params: req.request.params, + logprobs: req.request.logprobs, } } -fn prefilling_request(request_id: u64, label: &str, token_tx: TokenSink) -> PrefillingRequest35 { +fn prefilling_request( + req: QueuedRequest, + worker_id: u64, + ledger: &mut RequestLedger, +) -> PrefillingRequest35 { + ledger.admit(req.id); PrefillingRequest35 { - req: test_request(label, token_tx), + req, backend_state: PrefillBackendState::Tp { - request_id: RequestId::new(request_id), + request_id: RequestId::new(worker_id), }, cursor: 0, step_chunk: 0, @@ -63,111 +110,18 @@ fn prefilling_request(request_id: u64, label: &str, token_tx: TokenSink) -> Pref } #[derive(Default)] -struct PruneTestBackend { - retired_active: Vec, - dropped_prefilling: Vec<(RequestId, DropExpectation)>, -} - -impl DecodeDispatchBackend for PruneTestBackend { - fn is_stop_token(&self, _token: u32) -> bool { - false - } - - fn completion_requires_drop_ack(&self) -> bool { - true - } - - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35 { - active.swap_remove(idx) - } - - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - let ActiveBackendState::Tp { request_id, .. } = state else { - panic!("prune test expected TP active state"); - }; - self.retired_active.push(*request_id); - Ok(()) - } -} - -impl PrefillPromoteBackend for PruneTestBackend { - fn is_stop_token(&self, _token: u32) -> bool { - false - } - - fn promote_prefill_state( - &mut self, - _active_len: usize, - _state: PrefillBackendState, - ) -> ActiveBackendState { - panic!("prune test must not promote prefill state") - } - - fn drop_prefill_state( - &mut self, - state: &PrefillBackendState, - expectation: DropExpectation, - ) -> Result<()> { - let PrefillBackendState::Tp { request_id } = state else { - panic!("prune test expected TP prefill state"); - }; - self.dropped_prefilling.push((*request_id, expectation)); - Ok(()) - } -} - struct LifecycleTestBackend { stop_token: Option, - active_completion_requires_drop_ack: bool, fail_active_drop: bool, fail_prefill_drop: bool, active_drops: Vec, - active_events_before_drop: Vec, prefill_drops: Vec<(RequestId, DropExpectation)>, - observer: Option, -} - -impl LifecycleTestBackend { - fn new( - stop_token: Option, - observer: pegainfer_frontend::engine::TokenStreamReceiver, - ) -> Self { - Self { - stop_token, - active_completion_requires_drop_ack: true, - fail_active_drop: false, - fail_prefill_drop: false, - active_drops: Vec::new(), - active_events_before_drop: Vec::new(), - prefill_drops: Vec::new(), - observer: Some(observer), - } - } - - fn assert_no_completion_published(&mut self) { - let Some(observer) = &mut self.observer else { - return; - }; - assert!(matches!( - observer.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - } } impl DecodeDispatchBackend for LifecycleTestBackend { fn is_stop_token(&self, token: u32) -> bool { self.stop_token == Some(token) } - - fn completion_requires_drop_ack(&self) -> bool { - self.active_completion_requires_drop_ack - } - fn take_active_request( &mut self, active: &mut Vec, @@ -175,18 +129,10 @@ impl DecodeDispatchBackend for LifecycleTestBackend { ) -> ActiveRequest35 { active.swap_remove(idx) } - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { let ActiveBackendState::Tp { request_id, .. } = state else { - panic!("lifecycle test expected TP active state"); + panic!("expected TP active state"); }; - if self.active_completion_requires_drop_ack { - self.assert_no_completion_published(); - } else if let Some(observer) = &mut self.observer { - while let Ok((_, event)) = observer.try_recv() { - self.active_events_before_drop.push(event); - } - } self.active_drops.push(*request_id); anyhow::ensure!(!self.fail_active_drop, "injected active drop failure"); Ok(()) @@ -197,603 +143,283 @@ impl PrefillPromoteBackend for LifecycleTestBackend { fn is_stop_token(&self, token: u32) -> bool { self.stop_token == Some(token) } - - fn promote_prefill_state( - &mut self, - _active_len: usize, - _state: PrefillBackendState, - ) -> ActiveBackendState { + fn promote_prefill_state(&mut self, _: usize, _: PrefillBackendState) -> ActiveBackendState { panic!("completion lifecycle test must not promote prefill state") } - fn drop_prefill_state( &mut self, state: &PrefillBackendState, expectation: DropExpectation, ) -> Result<()> { let PrefillBackendState::Tp { request_id } = state else { - panic!("lifecycle test expected TP prefill state"); + panic!("expected TP prefill state"); }; - self.assert_no_completion_published(); self.prefill_drops.push((*request_id, expectation)); anyhow::ensure!(!self.fail_prefill_drop, "injected prefill drop failure"); Ok(()) } } -fn next_event( - rx: &mut pegainfer_frontend::engine::TokenStreamReceiver, - description: &str, -) -> TokenEvent { - rx.blocking_recv() - .unwrap_or_else(|| panic!("{description} channel closed before event")) - .1 -} - -fn assert_no_more_events(rx: &mut pegainfer_frontend::engine::TokenStreamReceiver) { - assert!( - rx.try_recv().is_err(), - "request received more than one terminal event" - ); -} - #[test] -fn closed_pending_work_is_pruned_before_admission() { - let (closed_sink, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let (open_sink, _open_rx) = TokenSink::standalone(); - let mut pending = vec![ - test_request("closed", closed_sink), - test_request("open", open_sink), - ]; - let mut active = Vec::new(); - let mut prefilling = Vec::new(); - let mut backend = PruneTestBackend::default(); - - assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - assert_eq!(pending.len(), 1); - assert_eq!(pending[0].request_id.as_deref(), Some("open")); - let admission = admit_pending_requests( - pending, - &[], - 1, - 16, - 8, - 8, - 128, - |req| req.prompt_tokens.len(), - |req| req.max_tokens, - |_| 0, - ); - assert_eq!(admission.pending.len(), 1); - assert!(admission.deferred.is_empty()); - assert!(admission.rejected.is_empty()); -} - -#[test] -fn closed_resident_work_is_absent_from_post_prune_load() { - let (closed_active_sink, closed_active_rx) = TokenSink::standalone(); - drop(closed_active_rx); - let (open_active_sink, _open_active_rx) = TokenSink::standalone(); - let (closed_prefill_sink, closed_prefill_rx) = TokenSink::standalone(); - drop(closed_prefill_rx); - let (pending_sink, _pending_rx) = TokenSink::standalone(); - let mut active = vec![ - active_request(10, "active-closed", closed_active_sink), - active_request(11, "active-open", open_active_sink), - ]; - let mut prefilling = vec![prefilling_request( - 12, - "prefill-closed", - closed_prefill_sink, - )]; - let mut pending = vec![test_request("pending-open", pending_sink)]; - let mut backend = PruneTestBackend::default(); - - assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - assert_eq!(active.len(), 1); - assert_eq!(active[0].request_id.as_deref(), Some("active-open")); - assert!(prefilling.is_empty()); - assert_eq!( - logical_load_counts(&active, &prefilling, 0, pending.len()), - (1, 1) - ); - assert_eq!(backend.retired_active, vec![RequestId::new(10)]); - assert_eq!( - backend.dropped_prefilling, - vec![(RequestId::new(12), DropExpectation::MustExist)] - ); -} - -#[test] -fn closed_resident_frees_capacity_for_same_tick_admission() { - let (closed_sink, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let (pending_sink, _pending_rx) = TokenSink::standalone(); - let mut active = vec![active_request(20, "resident-closed", closed_sink)]; - let mut prefilling = Vec::new(); - let mut pending = vec![test_request("replacement", pending_sink)]; - let mut backend = PruneTestBackend::default(); - - assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - let active_budget: Vec = active - .iter() - .map(|req| ActiveKvBudget { - prompt_len: req.prompt_len, - generated_count: req.generated_count, - max_tokens: req.max_tokens, - }) - .collect(); - let admission = admit_pending_requests( - pending, - &active_budget, - 1usize.saturating_sub(prefilling.len()), - 16, - 8, - 8, - 128, - |req| req.prompt_tokens.len(), - |req| req.max_tokens, - |_| 0, - ); - - assert!(active.is_empty()); - assert_eq!(backend.retired_active, vec![RequestId::new(20)]); - assert_eq!(admission.pending.len(), 1); - assert_eq!( - admission.pending[0].request_id.as_deref(), - Some("replacement") - ); - assert!(admission.deferred.is_empty()); -} - -#[test] -fn closed_materialized_prefill_requires_existing_worker_state() { - let (closed_sink, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let mut prefilling = vec![PrefillingRequest35 { - cursor: 1, - ..prefilling_request(21, "prefill-materialized", closed_sink) - }]; - let mut active = Vec::new(); - let mut pending = Vec::new(); - let mut backend = PruneTestBackend::default(); - - assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - assert!(prefilling.is_empty()); - assert_eq!( - backend.dropped_prefilling, - vec![(RequestId::new(21), DropExpectation::MustExist)] +fn closed_pending_work_is_pruned() { + let updates = run_step( + vec![ + test_request("closed", vec![1], 1), + test_request("open", vec![1], 1), + ], + &[0], + |mut pending, ledger| { + prune_closed_requests( + &mut LifecycleTestBackend::default(), + &mut Vec::new(), + &mut Vec::new(), + &mut pending, + ledger, + )?; + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].request.client_label.as_deref(), Some("open")); + ledger.retire(pending[0].id); + Ok(()) + }, ); + assert!(updates.is_empty()); } #[test] -fn prune_drop_failure_preserves_pending_for_terminal_fanout() { - let (closed_tx, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let (pending_tx, mut pending_rx) = TokenSink::standalone(); - let (_observer_tx, observer_rx) = TokenSink::standalone(); - let mut active = vec![active_request(22, "closed-active", closed_tx)]; - let mut prefilling = Vec::new(); - let mut pending = vec![test_request("live-pending", pending_tx)]; - let mut backend = LifecycleTestBackend::new(None, observer_rx); - backend.observer = None; - backend.fail_active_drop = true; - - let failure = - match prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) { - Ok(()) => panic!("injected prune drop should fail"), - Err(failure) => failure, - }; - assert!(active.is_empty()); - assert_eq!(failure.transient.len(), 1); - assert_eq!(pending.len(), 1); - - let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - active, - prefilling, - pending, - Vec::new(), - None, - failure, +fn closed_resident_work_is_pruned() { + let updates = run_step( + vec![ + test_request("active-closed", vec![1], 8), + test_request("active-open", vec![1], 8), + test_request("prefill-closed", vec![1], 1), + test_request("pending-open", vec![1], 1), + ], + &[0, 2], + |requests, ledger| { + let mut requests = requests.into_iter(); + let mut active = vec![ + active_request(requests.next().unwrap(), 10, ledger), + active_request(requests.next().unwrap(), 11, ledger), + ]; + let mut prefilling = vec![prefilling_request(requests.next().unwrap(), 12, ledger)]; + let mut pending: Vec<_> = requests.collect(); + let mut backend = LifecycleTestBackend::default(); + prune_closed_requests( + &mut backend, + &mut active, + &mut prefilling, + &mut pending, + ledger, + )?; + assert_eq!(active.len(), 1); + assert_eq!(active[0].client_label.as_deref(), Some("active-open")); + assert!(prefilling.is_empty()); + assert_eq!(backend.active_drops, vec![RequestId::new(10)]); + assert_eq!( + backend.prefill_drops, + vec![(RequestId::new(12), DropExpectation::MustExist)] + ); + ledger.retire(active[0].id); + ledger.retire(pending[0].id); + Ok(()) + }, ); - - assert!(matches!( - next_event(&mut pending_rx, "pending after prune failure"), - TokenEvent::Error { .. } - )); - assert_no_more_events(&mut pending_rx); + assert!(updates.is_empty()); } #[test] fn decode_eos_waits_for_drop_before_finished() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(30, "decode-eos", token_tx); + let mut request = test_request("decode-eos", vec![1], 8); request.params.ignore_eos = false; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(Some(9), token_rx); - - assert!(dispatch_decode_tokens(&mut backend, &mut active, &[9], &[None]).is_ok()); - - assert!(active.is_empty()); - assert_eq!(backend.active_drops, vec![RequestId::new(30)]); - let mut token_rx = backend.observer.take().unwrap(); + let updates = run_step(vec![request], &[], |mut requests, ledger| { + let mut active = vec![active_request(requests.remove(0), 30, ledger)]; + let mut backend = LifecycleTestBackend { + stop_token: Some(9), + ..Default::default() + }; + dispatch_decode_tokens(&mut backend, &mut active, &[9], &[None], ledger)?; + assert!(active.is_empty()); + assert_eq!(backend.active_drops, vec![RequestId::new(30)]); + Ok(()) + }); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].tokens, vec![1]); assert!(matches!( - next_event(&mut token_rx, "decode EOS"), - TokenEvent::Finished { - finish_reason: FinishReason::Stop, + updates[0].terminal, + Some(Terminal::Finished { + reason: FinishReason::Stop, + completion_tokens: 1, + stop_cause: None, .. - } + }) )); - assert_no_more_events(&mut token_rx); } #[test] fn decode_length_waits_for_drop_before_token_and_finished() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(31, "decode-length", token_tx); - request.params.ignore_eos = true; - request.max_tokens = 2; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(None, token_rx); - - assert!(dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None]).is_ok()); - - assert!(active.is_empty()); - assert_eq!(backend.active_drops, vec![RequestId::new(31)]); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - next_event(&mut token_rx, "decode length token"), - TokenEvent::Token { id: 7, .. } - )); - assert!(matches!( - next_event(&mut token_rx, "decode length finish"), - TokenEvent::Finished { - finish_reason: FinishReason::Length, - .. - } - )); - assert_no_more_events(&mut token_rx); -} - -#[test] -fn non_tp_decode_preserves_publish_before_retire_order() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(36, "single-order", token_tx); - request.params.ignore_eos = true; - request.max_tokens = 2; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(None, token_rx); - backend.active_completion_requires_drop_ack = false; - - assert!(dispatch_decode_tokens(&mut backend, &mut active, &[8], &[None]).is_ok()); - - assert!(active.is_empty()); - assert_eq!(backend.active_drops, vec![RequestId::new(36)]); - assert_eq!(backend.active_events_before_drop.len(), 2); - assert!(matches!( - &backend.active_events_before_drop[0], - TokenEvent::Token { id: 8, .. } - )); + let updates = run_step( + vec![test_request("decode-length", vec![1], 2)], + &[], + |mut requests, ledger| { + let mut active = vec![active_request(requests.remove(0), 31, ledger)]; + let mut backend = LifecycleTestBackend::default(); + dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None], ledger)?; + assert!(active.is_empty()); + assert_eq!(backend.active_drops, vec![RequestId::new(31)]); + Ok(()) + }, + ); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].tokens, vec![1, 7]); assert!(matches!( - &backend.active_events_before_drop[1], - TokenEvent::Finished { - finish_reason: FinishReason::Length, + updates[0].terminal, + Some(Terminal::Finished { + reason: FinishReason::Length, + completion_tokens: 2, + stop_cause: None, .. - } + }) )); } #[test] fn decode_completion_drop_failure_publishes_only_terminal_error() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(32, "decode-drop-failure", token_tx); - request.params.ignore_eos = true; - request.max_tokens = 2; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(None, token_rx); - backend.fail_active_drop = true; - - let failure = match dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None]) { - Ok(()) => panic!("injected active drop should fail"), - Err(failure) => failure, - }; - assert!(active.is_empty()); - assert_eq!(failure.transient.len(), 1); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - token_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - - let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - None, - failure, + let updates = run_step( + vec![test_request("decode-drop-failure", vec![1], 2)], + &[], + |mut requests, ledger| { + let mut active = vec![active_request(requests.remove(0), 32, ledger)]; + let id = active[0].id; + let mut backend = LifecycleTestBackend { + fail_active_drop: true, + ..Default::default() + }; + let result = dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None], ledger); + assert!(result.is_err()); + assert!(active.is_empty()); + assert_eq!(backend.active_drops, vec![RequestId::new(32)]); + assert_eq!( + ledger.completion_tokens(id), + 1, + "failed retirement must not commit the final token" + ); + result + }, ); - - assert!(matches!( - next_event(&mut token_rx, "failed decode completion"), - TokenEvent::Error { .. } + assert_eq!(updates.len(), 2); + assert_eq!(updates[0].tokens, vec![1]); + assert!(updates[0].terminal.is_none()); + assert!(updates[1].tokens.is_empty()); + assert!(matches!(&updates[1].terminal, + Some(Terminal::Failed { message, completion_tokens: 1, .. }) + if message.contains("injected active drop failure") )); - assert_no_more_events(&mut token_rx); } #[test] fn immediate_prefill_completion_waits_for_drop() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = prefilling_request(33, "prefill-length", token_tx); - request.req.max_tokens = 1; - request.req.params.ignore_eos = true; - request.step_chunk = 1; - let chunk = ScheduledChunk::from(vec![request]); - let mut active = Vec::new(); - let mut prefilling = Vec::new(); - let mut backend = LifecycleTestBackend::new(None, token_rx); - - assert!( - promote_or_requeue( - &mut backend, - &mut active, - &mut prefilling, - chunk, - &PrefillStepArtifacts::Single { - tokens: vec![11], - logprobs: vec![None], - }, - ) - .is_ok() - ); - - assert!(active.is_empty()); - assert!(prefilling.is_empty()); - assert_eq!( - backend.prefill_drops, - vec![(RequestId::new(33), DropExpectation::MustExist)] + let updates = run_step( + vec![test_request("prefill-length", vec![1], 1)], + &[], + |mut requests, ledger| { + let mut request = prefilling_request(requests.remove(0), 33, ledger); + request.step_chunk = 1; + let chunk = ScheduledChunk::from(vec![request]); + let mut active = Vec::new(); + let mut prefilling = Vec::new(); + let mut backend = LifecycleTestBackend::default(); + promote_or_requeue( + &mut backend, + &mut active, + &mut prefilling, + chunk, + &PrefillStepArtifacts::Single { + tokens: vec![11], + logprobs: vec![None], + }, + ledger, + )?; + assert!(active.is_empty()); + assert!(prefilling.is_empty()); + assert_eq!( + backend.prefill_drops, + vec![(RequestId::new(33), DropExpectation::MustExist)] + ); + Ok(()) + }, ); - let mut token_rx = backend.observer.take().unwrap(); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].tokens, vec![11]); assert!(matches!( - next_event(&mut token_rx, "prefill length token"), - TokenEvent::Token { id: 11, .. } - )); - assert!(matches!( - next_event(&mut token_rx, "prefill length finish"), - TokenEvent::Finished { - finish_reason: FinishReason::Length, + updates[0].terminal, + Some(Terminal::Finished { + reason: FinishReason::Length, + stop_cause: None, + completion_tokens: 1, .. - } + }) )); - assert_no_more_events(&mut token_rx); } #[test] fn immediate_prefill_drop_failure_publishes_only_terminal_error() { - let (token_tx, token_rx) = TokenSink::standalone(); - let (remaining_tx, mut remaining_rx) = TokenSink::standalone(); - let mut request = prefilling_request(34, "prefill-drop-failure", token_tx); - request.req.max_tokens = 1; - request.req.params.ignore_eos = true; - request.step_chunk = 1; - let mut remaining = prefilling_request(35, "remaining-scheduled", remaining_tx); - remaining.req.max_tokens = 1; - remaining.req.params.ignore_eos = true; - remaining.step_chunk = 1; - let chunk = ScheduledChunk::from(vec![request, remaining]); - let mut active = Vec::new(); - let mut prefilling = Vec::new(); - let mut backend = LifecycleTestBackend::new(None, token_rx); - backend.fail_prefill_drop = true; - - let failure = match promote_or_requeue( - &mut backend, - &mut active, - &mut prefilling, - chunk, - &PrefillStepArtifacts::Single { - tokens: vec![12, 13], - logprobs: vec![None, None], + let updates = run_step( + vec![ + test_request("prefill-drop-failure", vec![1], 1), + test_request("remaining-scheduled", vec![1], 1), + ], + &[], + |requests, ledger| { + let mut requests = requests.into_iter(); + let mut first = prefilling_request(requests.next().unwrap(), 34, ledger); + let id = first.req.id; + first.step_chunk = 1; + let mut second = prefilling_request(requests.next().unwrap(), 35, ledger); + second.step_chunk = 1; + let mut active = Vec::new(); + let mut prefilling = Vec::new(); + let mut backend = LifecycleTestBackend { + fail_prefill_drop: true, + ..Default::default() + }; + let result = promote_or_requeue( + &mut backend, + &mut active, + &mut prefilling, + ScheduledChunk::from(vec![first, second]), + &PrefillStepArtifacts::Single { + tokens: vec![12, 13], + logprobs: vec![None, None], + }, + ledger, + ); + assert!(result.is_err()); + assert_eq!(ledger.completion_tokens(id), 0); + assert_eq!( + backend.prefill_drops, + vec![(RequestId::new(34), DropExpectation::MustExist)] + ); + result }, - ) { - Ok(()) => panic!("injected prefill drop should fail"), - Err(failure) => failure, - }; - assert_eq!(failure.transient.len(), 2); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - token_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - assert!(matches!( - remaining_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - - let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - active, - prefilling, - Vec::new(), - Vec::new(), - None, - failure, ); - - assert!(matches!( - next_event(&mut token_rx, "failed prefill completion"), - TokenEvent::Error { .. } - )); - assert_no_more_events(&mut token_rx); - assert!(matches!( - next_event(&mut remaining_rx, "remaining scheduled prefill"), - TokenEvent::Error { .. } - )); - assert_no_more_events(&mut remaining_rx); -} - -#[test] -fn terminal_shutdown_closes_drains_and_errors_every_owner_once() { - let (active_tx, active_rx) = TokenSink::standalone(); - let (prefill_tx, prefill_rx) = TokenSink::standalone(); - let (pending_tx, pending_rx) = TokenSink::standalone(); - let (deferred_tx, deferred_rx) = TokenSink::standalone(); - let (candidate_tx, candidate_rx) = TokenSink::standalone(); - let (scheduled_tx, scheduled_rx) = TokenSink::standalone(); - let (queued_tx, queued_rx) = TokenSink::standalone(); - let (after_close_tx, mut after_close_rx) = TokenSink::standalone(); - let (closed_tx, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - - let active = vec![active_request(40, "active", active_tx)]; - let prefilling = vec![prefilling_request(41, "prefilling", prefill_tx)]; - let pending = vec![ - test_request("duplicate-external-id", pending_tx), - test_request("closed-sink", closed_tx), - ]; - let deferred = vec![test_request("duplicate-external-id", deferred_tx)]; - let candidate = CompletionCandidate { - request: active_request(42, "candidate", candidate_tx), - final_events: vec![TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: 1, - completion_tokens: 2, - }], - }; - let failure = FatalSchedulerError::new("injected TP replica failure") - .with_request(candidate.into_terminal()) - .with_request(test_request("scheduled", scheduled_tx)); - - let queued = test_request("queued-before-close", queued_tx); - let after_close = test_request("queued-after-close", after_close_tx); - let (submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let before_send = Arc::new(Barrier::new(2)); - let sent_before_close = Arc::new(Barrier::new(2)); - let after_receiver_close = Arc::new(Barrier::new(2)); - let sender = { - let before_send = Arc::clone(&before_send); - let sent_before_close = Arc::clone(&sent_before_close); - let after_receiver_close = Arc::clone(&after_receiver_close); - std::thread::spawn(move || { - before_send.wait(); - submit_tx - .send((queued, pegainfer_frontend::engine::KvPrefix::none())) - .expect("close-before request should be accepted"); - sent_before_close.wait(); - after_receiver_close.wait(); - submit_tx - .send((after_close, pegainfer_frontend::engine::KvPrefix::none())) - .is_err() - }) - }; - - before_send.wait(); - sent_before_close.wait(); - let (load_tx, load_rx) = watch::channel(SchedulerMetrics { - kv_used_blocks: 9, - kv_total_blocks: 64, - num_running_reqs: 9, - num_waiting_reqs: 9, - spec_decode: None, - }); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - active, - prefilling, - pending, - deferred, - None, - failure, - ); - after_receiver_close.wait(); - assert!(sender.join().expect("submit race thread panicked")); - - let mut receivers = vec![ - ("active", active_rx), - ("prefilling", prefill_rx), - ("pending", pending_rx), - ("deferred", deferred_rx), - ("candidate", candidate_rx), - ("scheduled", scheduled_rx), - ("queued", queued_rx), - ]; - for (owner, rx) in &mut receivers { - match next_event(rx, owner) { - TokenEvent::Error { message, .. } => { - assert_eq!(message, "injected TP replica failure"); - } - other => panic!("{owner} received non-error terminal event: {other:?}"), - } - assert_no_more_events(rx); - } - assert!(after_close_rx.try_recv().is_err()); - - let snapshot = *load_rx.borrow(); - assert_eq!(snapshot.kv_used_blocks, 0); - assert_eq!(snapshot.kv_total_blocks, 64); - assert_eq!(snapshot.num_running_reqs, 0); - assert_eq!(snapshot.num_waiting_reqs, 0); -} - -fn collect_finished_with_timeout( - token_rx: &mut pegainfer_frontend::engine::TokenStreamReceiver, - description: &str, -) -> (usize, FinishReason) { - let deadline = Instant::now() + Duration::from_secs(30); - let mut token_count = 0; - loop { - match token_rx.try_recv() { - Ok((_, TokenEvent::Token { .. })) => token_count += 1, - Ok((_, TokenEvent::Finished { finish_reason, .. })) => { - return (token_count, finish_reason); - } - Ok((_, TokenEvent::Error { message, .. })) => { - panic!("{description} failed: {message}") - } - Ok((_, TokenEvent::Rejected { message, .. })) => { - panic!("{description} was rejected: {message}") - } - Ok((_, _)) => {} - Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { - assert!( - Instant::now() < deadline, - "timed out waiting for {description}" - ); - std::thread::sleep(Duration::from_millis(1)); - } - Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { - panic!("{description} channel disconnected before Finished") - } - } + for raw_id in 0..2 { + let request_updates: Vec<_> = updates + .iter() + .filter(|update| update.id == FrontendRequestId::new(raw_id)) + .collect(); + let tokens: Vec<_> = request_updates + .iter() + .flat_map(|update| &update.tokens) + .copied() + .collect(); + assert!(tokens.is_empty()); + let terminals: Vec<_> = request_updates + .iter() + .filter_map(|update| update.terminal.as_ref()) + .collect(); + assert_eq!(terminals.len(), 1); + assert!(matches!(terminals[0], Terminal::Failed { message, .. } + if message.contains("injected prefill drop failure"))); } } @@ -809,88 +435,60 @@ fn prefix_cache_chunking_stops_at_snapshot_boundaries() { #[test] fn send_rejection_reports_lifetime_kv_and_context_limits() { - let rejection_message = |reason: RejectReason, max_tokens: usize| { - let (token_tx, mut token_rx) = TokenSink::standalone(); - let req = test_request_with_shape("rejected", token_tx, vec![1; 16], max_tokens); - send_rejection(&req, reason); - match token_rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Rejected { - message, - prompt_tokens, - completion_tokens, - }) => { - assert_eq!((prompt_tokens, completion_tokens), (16, 0)); - message - } - other => panic!("expected rejection event, got {other:?}"), + let rejection_reason = |reason, max_tokens| { + let updates = run_step( + vec![test_request("rejected", vec![1; 16], max_tokens)], + &[], + move |requests, ledger| { + send_rejection(&requests[0], reason, ledger); + Ok(()) + }, + ); + assert_eq!(updates.len(), 1); + match updates.into_iter().next().unwrap().terminal { + Some(Terminal::Rejected { + reason, + prompt_tokens: 16, + }) => reason, + other => panic!("expected rejection, got {other:?}"), } }; - - let kv = rejection_message(RejectReason::KvBudget, 49); - assert!( - kv.contains("max_request_tokens=65"), - "rejection should report the lifetime KV peak used by admission: {kv}" - ); - - let context = rejection_message(RejectReason::ContextLength { limit: 32 }, 17); - assert!( - context.contains("maximum context length of 32 tokens"), - "rejection should report the context-window limit: {context}" + assert_eq!( + rejection_reason(RejectReason::KvBudget, 49), + pegainfer_frontend::engine::RejectReason::KvBudget { + prompt_tokens: 16, + worst_case_tokens: 65, + } ); - assert!( - context.contains("requested 33"), - "rejection should report prompt + max_tokens: {context}" + assert_eq!( + rejection_reason(RejectReason::ContextLength { limit: 32 }, 17), + pegainfer_frontend::engine::RejectReason::ContextLength { + prompt_tokens: 16, + max_tokens: 17, + limit: 32, + } ); } #[test] -fn prompt_logprobs_request_is_rejected_before_backend_admission() { - let (unsupported_tx, mut unsupported_rx) = TokenSink::standalone(); - let (regular_tx, mut regular_rx) = TokenSink::standalone(); - let mut unsupported = test_request_with_shape( - "unsupported-prompt-logprobs", - unsupported_tx, - vec![1, 2, 3], - 4, - ); +fn prompt_logprobs_filter_rejects_unsupported_requests() { + let mut unsupported = test_request("unsupported-prompt-logprobs", vec![1, 2, 3], 4); unsupported.prompt_logprobs = Some(0); - let regular = test_request("regular", regular_tx); - let mut pending = vec![unsupported, regular]; - - reject_unsupported_prompt_logprobs(&mut pending); - - assert_eq!(pending.len(), 1); - assert_eq!(pending[0].request_id.as_deref(), Some("regular")); - assert!( - pending[0].prompt_logprobs.is_none(), - "only requests eligible for backend admission may remain" + let updates = run_step( + vec![unsupported, test_request("regular", vec![1], 1)], + &[], + |mut pending, ledger| { + reject_unsupported_prompt_logprobs(&mut pending, ledger); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].request.client_label.as_deref(), Some("regular")); + ledger.retire(pending[0].id); + Ok(()) + }, ); - match unsupported_rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Rejected { - message, - prompt_tokens, - completion_tokens, - }) => { - assert_eq!(message, UNSUPPORTED_PROMPT_LOGPROBS_MESSAGE); - assert_eq!(prompt_tokens, 3); - assert_eq!(completion_tokens, 0); - } - event => panic!("expected unsupported prompt-logprobs rejection, got {event:?}"), - } - assert!(matches!( - regular_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); -} - -#[test] -fn submit_parking_requires_idle_owned_work_and_no_inflight_prefill() { - // Parking on submit_rx while a prefill is in flight would never observe its - // completion, so only a fully idle scheduler may block there. - assert!(should_block_on_submit(true, false)); - assert!(!should_block_on_submit(true, true)); - assert!(!should_block_on_submit(false, false)); - assert!(!should_block_on_submit(false, true)); + assert_eq!(updates.len(), 1); + assert!(matches!(&updates[0].terminal, Some(Terminal::Rejected { + reason: pegainfer_frontend::engine::RejectReason::Unsupported { feature }, prompt_tokens: 3, + }) if feature == "prompt_logprobs")); } #[test] @@ -901,34 +499,45 @@ fn tp2_scheduler_runs_forced_mixed_steps() { else { return; }; - let handle = start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1, false, 0) + let mut engine = start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1, false, 0) .expect("start TP2 scheduler"); - let (decode_tx, mut decode_rx) = TokenSink::standalone(); - let (prefill_tx, mut prefill_rx) = TokenSink::standalone(); - - handle - .submit(test_request_with_shape( - "mixed-active", - decode_tx, - vec![151_646], - 8, - )) - .expect("submit request that becomes active first"); - handle - .submit(test_request_with_shape( - "mixed-prefill", - prefill_tx, - vec![151_646, 9707], - 2, - )) - .expect("submit request that remains chunk-prefilling"); - - let (decode_tokens, decode_finish) = - collect_finished_with_timeout(&mut decode_rx, "mixed active request"); - let (prefill_tokens, prefill_finish) = - collect_finished_with_timeout(&mut prefill_rx, "mixed prefill request"); - assert_eq!(decode_tokens, 8); - assert_eq!(decode_finish, FinishReason::Length); - assert_eq!(prefill_tokens, 2); - assert_eq!(prefill_finish, FinishReason::Length); + let LiveScheduler { mut handle, join } = engine.schedulers.remove(0); + let mut steps = handle.take_steps().unwrap(); + let decode = handle.submit(test_request("mixed-active", vec![151_646], 8)); + let prefill = handle.submit(test_request("mixed-prefill", vec![151_646, 9707], 2)); + let ids = [decode.id(), prefill.id()]; + let expected = [8, 2]; + let mut finished = [false, false]; + let deadline = Instant::now() + Duration::from_secs(30); + while !finished.iter().all(|done| *done) { + match steps.try_recv() { + Ok(step) => { + for update in step.updates { + let index = ids + .iter() + .position(|id| *id == update.id) + .expect("unknown request id"); + if let Some(terminal) = update.terminal { + assert!( + matches!(terminal, Terminal::Finished { reason: FinishReason::Length, completion_tokens, .. } + if completion_tokens == expected[index]) + ); + finished[index] = true; + } + } + } + Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { + assert!( + Instant::now() < deadline, + "timed out waiting for mixed requests" + ); + std::thread::sleep(Duration::from_millis(1)); + } + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { + panic!("scheduler exited before mixed requests finished") + } + } + } + drop(handle); + join.join().expect("scheduler thread panicked"); } diff --git a/pegainfer-qwen35/src/scheduler/tp.rs b/pegainfer-qwen35/src/scheduler/tp.rs index 9b636256d..702e1c01e 100644 --- a/pegainfer-qwen35/src/scheduler/tp.rs +++ b/pegainfer-qwen35/src/scheduler/tp.rs @@ -23,9 +23,9 @@ pub(super) fn tp_prefill_items(chunk: &ScheduledChunk) -> Result EngineHandle { +fn start_engine(model_path: &str, max_prefill_tokens: usize) -> EngineHarness { pegainfer_qwen35::start_engine( Path::new(model_path), EngineLoadOptions { @@ -35,44 +34,35 @@ fn start_engine(model_path: &str, max_prefill_tokens: usize) -> EngineHandle { MAX_BATCH, max_prefill_tokens, ) + .map(EngineHarness::new) .expect("failed to start Qwen3.5 engine") } -fn generate(handle: &EngineHandle, prompt_tokens: Vec) -> (Vec, FinishReason) { - let (token_tx, mut rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams { - ignore_eos: true, - ..SamplingParams::default() - }, - max_tokens: GENERATED_TOKENS, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: None, - prompt_logprobs: None, - }) - .expect("submit failed"); - +fn generate(handle: &mut EngineHarness, prompt_tokens: Vec) -> (Vec, FinishReason) { + let control = handle.submit(common::request( + prompt_tokens, + SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + GENERATED_TOKENS, + )); let mut tokens = Vec::new(); loop { - match rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Token { id, .. }) => tokens.push(id), - Some( - TokenEvent::Scheduled { .. } - | TokenEvent::PromptTokens { .. } - | TokenEvent::KvTransfer { .. }, - ) => {} - Some(TokenEvent::Finished { finish_reason, .. }) => return (tokens, finish_reason), - Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), - Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), - None => panic!("scheduler channel closed without Finished"), + let update = handle.next(control.id()); + tokens.extend(update.tokens); + if let Some(terminal) = update.terminal { + match terminal { + Terminal::Finished { + reason, + completion_tokens, + .. + } => { + assert_eq!(completion_tokens, GENERATED_TOKENS); + return (tokens, reason); + } + terminal => panic!("generation did not finish: {terminal:?}"), + } } } } @@ -102,8 +92,8 @@ fn chunked_prefill_matches_unchunked_prefill_for_resumed_paged_kv() { ); let (baseline_tokens, baseline_finish) = { - let handle = start_engine(&model_path, BASELINE_PREFILL_BUDGET); - generate(&handle, prompt_tokens.clone()) + let mut handle = start_engine(&model_path, BASELINE_PREFILL_BUDGET); + generate(&mut handle, prompt_tokens.clone()) }; assert_eq!( baseline_finish, @@ -112,8 +102,8 @@ fn chunked_prefill_matches_unchunked_prefill_for_resumed_paged_kv() { ); let (chunked_tokens, chunked_finish) = { - let handle = start_engine(&model_path, CHUNK_BUDGET); - generate(&handle, prompt_tokens) + let mut handle = start_engine(&model_path, CHUNK_BUDGET); + generate(&mut handle, prompt_tokens) }; assert_eq!( chunked_finish, diff --git a/pegainfer-qwen35/tests/common/mod.rs b/pegainfer-qwen35/tests/common/mod.rs index b2bd8fe2c..4ba477fd7 100644 --- a/pegainfer-qwen35/tests/common/mod.rs +++ b/pegainfer-qwen35/tests/common/mod.rs @@ -1,5 +1,19 @@ +use std::collections::HashMap; +use std::collections::VecDeque; use std::sync::Arc; +use pegainfer_frontend::engine::Engine; +use pegainfer_frontend::engine::EosPolicy; +use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::RequestControl; +use pegainfer_frontend::engine::RequestId; +use pegainfer_frontend::engine::RequestUpdate; +use pegainfer_frontend::engine::SchedulerHandle; +use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StepOutputs; +use pegainfer_frontend::engine::StepReceiver; +use pegainfer_frontend::engine::StopPolicy; +use pegainfer_frontend::sampler::SamplingParams; use vllm_text::Error; use vllm_text::Result; use vllm_text::backend::hf::ResolvedModelFiles; @@ -13,6 +27,124 @@ pub(crate) mod model_fixture; pub(crate) use model_fixture::model_path_or_skip; +#[allow(dead_code)] +pub(crate) fn request( + prompt_tokens: Vec, + params: SamplingParams, + max_tokens: usize, +) -> Request { + Request { + prompt_tokens, + stop_policy: StopPolicy::new( + if params.ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + Vec::new(), + ), + params, + max_tokens, + lora_adapter: None, + kv_transfer_params: None, + logprobs: None, + prompt_logprobs: None, + trace_parent: None, + client_label: None, + } +} + +// Cancel on scope exit or assertion failure before the harness joins the scheduler. +pub(crate) struct RequestGuard(RequestControl); + +#[allow(dead_code)] +impl RequestGuard { + pub(crate) fn id(&self) -> RequestId { + self.0.id() + } + + pub(crate) fn abort(&self) { + self.0.abort(); + } +} + +impl Drop for RequestGuard { + fn drop(&mut self) { + self.0.abort(); + } +} + +#[allow(dead_code)] +pub(crate) struct EngineHarness { + handle: Option, + join: Option>, + steps: StepReceiver, + pending: HashMap>, +} + +#[allow(dead_code)] +impl EngineHarness { + pub(crate) fn new(mut engine: Engine) -> Self { + assert_eq!(engine.schedulers.len(), 1); + let mut scheduler = engine.schedulers.pop().unwrap(); + Self { + steps: scheduler + .handle + .take_steps() + .expect("new engine has a step stream"), + handle: Some(scheduler.handle), + join: Some(scheduler.join), + pending: HashMap::new(), + } + } + + pub(crate) fn submit(&self, request: Request) -> RequestGuard { + RequestGuard(self.handle.as_ref().unwrap().submit(request)) + } + + pub(crate) fn metrics(&self) -> SchedulerMetrics { + self.handle.as_ref().unwrap().metrics() + } + + pub(crate) fn next(&mut self, id: RequestId) -> RequestUpdate { + loop { + if let Some(update) = self.try_next(id) { + return update; + } + let step = self + .steps + .blocking_recv() + .expect("scheduler closed before terminal"); + self.buffer(step); + } + } + + pub(crate) fn try_next(&mut self, id: RequestId) -> Option { + while let Ok(step) = self.steps.try_recv() { + self.buffer(step); + } + self.pending.get_mut(&id).and_then(VecDeque::pop_front) + } + + fn buffer(&mut self, step: StepOutputs) { + for update in step.updates { + self.pending.entry(update.id).or_default().push_back(update); + } + } +} + +impl Drop for EngineHarness { + fn drop(&mut self) { + drop(self.handle.take()); + if let Some(join) = self.join.take() { + let result = join.join(); + if !std::thread::panicking() { + result.expect("scheduler thread panicked"); + } + } + } +} + #[allow(dead_code)] pub(crate) fn load_tokenizer(model_path: &str) -> DynTokenizer { try_load_tokenizer(model_path) diff --git a/pegainfer-qwen35/tests/e2e_scheduler.rs b/pegainfer-qwen35/tests/e2e_scheduler.rs index b4b835318..625140813 100644 --- a/pegainfer-qwen35/tests/e2e_scheduler.rs +++ b/pegainfer-qwen35/tests/e2e_scheduler.rs @@ -1,26 +1,26 @@ //! E2E scheduler integration test for Qwen3.5-4B. //! //! Tests the Qwen3.5 reduced-capacity scheduler path (batch prefill + -//! CUDA Graph decode) with sequential, concurrent, and consumer-drop requests. +//! CUDA Graph decode) with sequential, concurrent, and cancelled requests. use std::collections::HashSet; use std::path::Path; use std::time::Instant; use log::info; -use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::FinishReason; -use pegainfer_frontend::engine::GenerateRequest; -use pegainfer_frontend::engine::SchedulerMetrics; -use pegainfer_frontend::engine::TokenEvent; +use pegainfer_frontend::engine::RejectReason; +use pegainfer_frontend::engine::RequestUpdate; +use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::engine::TokenLogprob; -use pegainfer_frontend::engine::TokenSink; -use pegainfer_frontend::engine::TokenStreamReceiver; use pegainfer_frontend::sampler::SamplingParams; use vllm_text::tokenizer::DynTokenizer; mod common; +use common::EngineHarness; +use common::RequestGuard; + const CASES: &[TestCase] = &[ TestCase { name: "tell_story", @@ -101,95 +101,71 @@ struct GenerationResult { } fn generate_tokens( - handle: &EngineHandle, + handle: &mut EngineHarness, tokenizer: &DynTokenizer, prompt: &str, max_tokens: usize, ) -> (Vec, FinishReason) { - let result = generate_tokens_with_logprobs(handle, tokenizer, prompt, max_tokens, 0); + let result = generate_tokens_with_logprobs(handle, tokenizer, prompt, max_tokens, None); (result.tokens, result.finish_reason) } fn generate_tokens_with_logprobs( - handle: &EngineHandle, + handle: &mut EngineHarness, tokenizer: &DynTokenizer, prompt: &str, max_tokens: usize, - logprobs: usize, + logprobs: Option, ) -> GenerationResult { let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); - let (token_tx, mut token_rx) = TokenSink::standalone(); - - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams::default(), - max_tokens, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: (logprobs > 0).then_some(logprobs), - prompt_logprobs: None, - }) - .expect("submit failed"); - - collect_generation(&mut token_rx, prompt, logprobs) + let mut request = common::request(prompt_tokens, SamplingParams::default(), max_tokens); + request.logprobs = logprobs; + let control = handle.submit(request); + collect_generation(handle, &control, prompt, logprobs) } fn submit_repeated_token_request( - handle: &EngineHandle, + handle: &EngineHarness, request_id: &str, token: u32, prompt_len: usize, max_tokens: usize, -) -> TokenStreamReceiver { - let (token_tx, token_rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: Some(request_id.to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens: vec![token; prompt_len], - params: SamplingParams { - ignore_eos: true, - ..SamplingParams::default() - }, - max_tokens, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: None, - prompt_logprobs: None, - }) - .unwrap_or_else(|err| panic!("submit {request_id}: {err}")); - token_rx +) -> RequestGuard { + let mut request = common::request( + vec![token; prompt_len], + SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens, + ); + request.client_label = Some(request_id.into()); + handle.submit(request) } -fn wait_for_first_token(rx: &mut TokenStreamReceiver, request_id: &str) { +fn wait_for_first_token(handle: &mut EngineHarness, control: &RequestGuard, request_id: &str) { let deadline = Instant::now() + std::time::Duration::from_secs(30); loop { - match recv_event_before(rx, request_id, deadline) { - Some(TokenEvent::Token { .. }) => return, - Some(TokenEvent::PromptTokens { .. } | TokenEvent::Scheduled { .. }) => {} - Some(event) => panic!("{request_id} emitted {event:?} before its first token"), - None => panic!("scheduler closed before {request_id} emitted a token"), + let update = recv_event_before(handle, control, request_id, deadline); + assert!( + update.terminal.is_none(), + "{request_id} ended before its first token" + ); + if !update.tokens.is_empty() { + return; } } } fn recv_event_before( - rx: &mut TokenStreamReceiver, + handle: &mut EngineHarness, + control: &RequestGuard, request_id: &str, deadline: Instant, -) -> Option { +) -> RequestUpdate { loop { - if let Ok((_, event)) = rx.try_recv() { - return Some(event); + if let Some(update) = handle.try_next(control.id()) { + return update; } assert!( Instant::now() < deadline, @@ -199,35 +175,31 @@ fn recv_event_before( } } -fn assert_no_generated_event(rx: &mut TokenStreamReceiver, request_id: &str) { - while let Ok((_, event)) = rx.try_recv() { - match event { - TokenEvent::PromptTokens { .. } | TokenEvent::Scheduled { .. } => {} - event => panic!("{request_id} emitted {event:?} before the overlap bound"), - } +fn assert_no_generated_event(handle: &mut EngineHarness, control: &RequestGuard, request_id: &str) { + while let Some(update) = handle.try_next(control.id()) { + assert!( + update.tokens.is_empty() && update.terminal.is_none(), + "{request_id} emitted {update:?} before the overlap bound" + ); } } -fn drain_tokens(rx: &mut TokenStreamReceiver, request_id: &str) -> usize { +fn drain_tokens(handle: &mut EngineHarness, control: &RequestGuard, request_id: &str) -> usize { let mut tokens = 0; - while let Ok((_, event)) = rx.try_recv() { - match event { - TokenEvent::Token { .. } => tokens += 1, - TokenEvent::PromptTokens { .. } | TokenEvent::Scheduled { .. } => {} - event => panic!("{request_id} emitted {event:?} while it must remain active"), - } + while let Some(update) = handle.try_next(control.id()) { + assert!( + update.terminal.is_none(), + "{request_id} ended while it must remain active" + ); + tokens += update.tokens.len(); } tokens } -fn wait_for_running_requests( - load: &mut tokio::sync::watch::Receiver, - expected: u64, - timeout: std::time::Duration, -) { +fn wait_for_running_requests(handle: &EngineHarness, expected: u64, timeout: std::time::Duration) { let deadline = Instant::now() + timeout; loop { - let snapshot = *load.borrow_and_update(); + let snapshot = handle.metrics(); if snapshot.num_running_reqs == expected { return; } @@ -240,83 +212,89 @@ fn wait_for_running_requests( } fn collect_generation( - token_rx: &mut TokenStreamReceiver, + handle: &mut EngineHarness, + control: &RequestGuard, name: &str, - logprobs: usize, + logprobs: Option, ) -> GenerationResult { - collect_generation_until(token_rx, name, logprobs, None) + collect_generation_until(handle, control, name, logprobs, None) } fn collect_generation_with_timeout( - token_rx: &mut TokenStreamReceiver, + handle: &mut EngineHarness, + control: &RequestGuard, name: &str, - logprobs: usize, + logprobs: Option, timeout: std::time::Duration, ) -> GenerationResult { - collect_generation_until(token_rx, name, logprobs, Some(Instant::now() + timeout)) + collect_generation_until( + handle, + control, + name, + logprobs, + Some(Instant::now() + timeout), + ) } fn collect_generation_until( - token_rx: &mut TokenStreamReceiver, + handle: &mut EngineHarness, + control: &RequestGuard, name: &str, - logprobs: usize, + logprobs: Option, deadline: Option, ) -> GenerationResult { let mut tokens = Vec::new(); let mut token_logprobs = Vec::new(); loop { - let event = match deadline { - Some(deadline) => recv_event_before(token_rx, name, deadline), - None => token_rx.blocking_recv().map(|(_, event)| event), + let update = match deadline { + Some(deadline) => recv_event_before(handle, control, name, deadline), + None => handle.next(control.id()), }; - match event { - Some(TokenEvent::Token { id, logprob }) => { - if logprobs == 0 { - assert!( - logprob.is_none(), - "{name}: logprobs=0 should not return token logprobs" - ); - } else { - let lp = logprob - .as_ref() - .unwrap_or_else(|| panic!("{name}: logprobs={logprobs} returned None")); - assert!( - lp.logprob.is_finite(), - "{name}: sampled token logprob must be finite" - ); - assert_eq!( - lp.top_logprobs.len(), - logprobs, - "{name}: top_logprobs length should match the request" - ); - assert!( - lp.top_logprobs.iter().all(|&(_, v)| v.is_finite()), - "{name}: top_logprobs must be finite" - ); + for (&id, logprob) in update.tokens.iter().zip(&update.logprobs) { + if let Some(count) = logprobs { + let lp = logprob + .as_ref() + .unwrap_or_else(|| panic!("{name}: logprobs={count} returned None")); + assert!( + lp.logprob.is_finite(), + "{name}: sampled token logprob must be finite" + ); + assert_eq!( + lp.top_logprobs.len(), + count, + "{name}: top_logprobs length should match the request" + ); + assert!( + lp.top_logprobs.iter().all(|&(_, v)| v.is_finite()), + "{name}: top_logprobs must be finite" + ); + if count > 0 { assert_eq!( lp.top_logprobs.first().map(|&(token, _)| token), Some(id), "{name}: greedy sampled token should match top-1 logprob row" ); } - tokens.push(id); - token_logprobs.push(logprob); + } else { + assert!( + logprob.is_none(), + "{name}: disabled logprobs should return None" + ); } - Some( - TokenEvent::PromptTokens { .. } - | TokenEvent::Scheduled { .. } - | TokenEvent::KvTransfer { .. }, - ) => {} - Some(TokenEvent::Finished { finish_reason, .. }) => { - return GenerationResult { - tokens, - logprobs: token_logprobs, - finish_reason, - }; + } + tokens.extend(update.tokens); + token_logprobs.extend(update.logprobs); + if let Some(terminal) = update.terminal { + match terminal { + Terminal::Finished { reason, .. } => { + return GenerationResult { + tokens, + logprobs: token_logprobs, + finish_reason: reason, + }; + } + terminal => panic!("{name}: generation did not finish: {terminal:?}"), } - Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), - Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), - None => panic!("{name}: scheduler channel closed without Finished"), } } } @@ -334,45 +312,27 @@ fn concurrent_params(case_idx: usize) -> SamplingParams { } } -fn expect_context_window_rejection(handle: &EngineHandle, max_context_tokens: usize) { - let (token_tx, mut token_rx) = TokenSink::standalone(); - - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: Some("over-context-window".to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens: vec![1; max_context_tokens], - params: SamplingParams::default(), - max_tokens: 1, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: None, - prompt_logprobs: None, - }) - .expect("submit over-context request"); - - match token_rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Rejected { - message, +fn expect_context_window_rejection(handle: &mut EngineHarness, max_context_tokens: usize) { + let mut request = common::request(vec![1; max_context_tokens], SamplingParams::default(), 1); + request.client_label = Some("over-context-window".into()); + let control = handle.submit(request); + let update = handle.next(control.id()); + match update.terminal { + Some(Terminal::Rejected { + reason, prompt_tokens, - completion_tokens, }) => { assert_eq!(prompt_tokens, max_context_tokens); - assert_eq!(completion_tokens, 0); - assert!( - message.contains("maximum context length"), - "expected context-window rejection, got: {message}" - ); - assert!( - message.contains(&(max_context_tokens + 1).to_string()), - "rejection should report prompt + max_tokens, got: {message}" + assert_eq!( + reason, + RejectReason::ContextLength { + prompt_tokens: max_context_tokens, + max_tokens: 1, + limit: max_context_tokens, + } ); } - Some(_) => panic!("expected context-window rejection"), - None => panic!("scheduler channel closed without rejection"), + terminal => panic!("expected context-window rejection, got: {terminal:?}"), } } @@ -449,7 +409,7 @@ fn assert_no_model_wide_collapse(collapses: &[(&str, Collapse)]) { } fn run_full_scheduler_e2e( - handle: &EngineHandle, + handle: &mut EngineHarness, tokenizer: &DynTokenizer, max_context_tokens: usize, label: &str, @@ -466,9 +426,9 @@ fn run_full_scheduler_e2e( for case in CASES.iter().take(3) { let max_tokens = case.max_new_tokens.min(16); let no_logprobs = - generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, 0); + generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, None); let with_logprobs = - generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, 1); + generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, Some(1)); assert_eq!(no_logprobs.finish_reason, with_logprobs.finish_reason); assert_eq!( no_logprobs.tokens, with_logprobs.tokens, @@ -477,7 +437,7 @@ fn run_full_scheduler_e2e( ); assert!( no_logprobs.logprobs.iter().all(Option::is_none), - "logprobs=0 should keep the no-host-logprobs path for {:?}", + "logprobs=None should keep the no-host-logprobs path for {:?}", case.name ); assert!( @@ -491,7 +451,7 @@ fn run_full_scheduler_e2e( case.name ); info!( - " PASS: {:?} logprobs=0 and logprobs=1 produced identical greedy tokens", + " PASS: {:?} logprobs=None and logprobs=Some(1) produced identical greedy tokens", case.name ); } @@ -543,35 +503,23 @@ fn run_full_scheduler_e2e( // ── 4. Concurrent requests ────────────────────────────────────────── info!("=== Phase 4: Concurrent requests ==="); { - let mut receivers: Vec<(String, usize, TokenStreamReceiver)> = Vec::new(); + let mut requests = Vec::new(); // Submit all cases concurrently, alternating greedy and sampled rows so // batch decode covers the mixed token-selection path from #284. for (case_idx, case) in CASES.iter().enumerate() { let prompt_tokens = tokenizer.encode(case.prompt, false).expect("encode failed"); - let (token_tx, token_rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: concurrent_params(case_idx), - max_tokens: case.max_new_tokens, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: None, - prompt_logprobs: None, - }) - .expect("submit failed"); - receivers.push((case.name.to_string(), 0, token_rx)); + let control = handle.submit(common::request( + prompt_tokens, + concurrent_params(case_idx), + case.max_new_tokens, + )); + requests.push((case.name, control)); } // Collect all results - for (name, logprobs, mut rx) in receivers { - let result = collect_generation(&mut rx, &name, logprobs); + for (name, control) in requests { + let result = collect_generation(handle, &control, name, None); let text = tokenizer .decode(&result.tokens, true) .expect("decode failed"); @@ -584,82 +532,42 @@ fn run_full_scheduler_e2e( info!("=== Phase 4b: Mixed concurrent logprobs ==="); { let mixed = [ - ("mixed_no_logprobs", CASES[0].prompt, 0usize), - ("mixed_with_logprobs", CASES[1].prompt, 1usize), + ("mixed_no_logprobs", CASES[0].prompt, None), + ("mixed_chosen_logprob", CASES[1].prompt, Some(0)), + ("mixed_top_logprobs", CASES[1].prompt, Some(1)), ]; - let mut receivers: Vec<(&str, usize, TokenStreamReceiver)> = Vec::new(); - + let mut requests = Vec::new(); for (name, prompt, logprobs) in mixed { let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); - let (token_tx, token_rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: Some(name.to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams::default(), - max_tokens: 8, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: (logprobs > 0).then_some(logprobs), - prompt_logprobs: None, - }) - .expect("submit failed"); - receivers.push((name, logprobs, token_rx)); + let mut request = common::request(prompt_tokens, SamplingParams::default(), 8); + request.client_label = Some(name.into()); + request.logprobs = logprobs; + requests.push((name, logprobs, handle.submit(request))); } - - for (name, logprobs, mut rx) in receivers { - let result = collect_generation(&mut rx, name, logprobs); + for (name, logprobs, control) in requests { + let result = collect_generation(handle, &control, name, logprobs); assert!(!result.tokens.is_empty(), "{name}: produced no tokens"); - if logprobs == 0 { - assert!( - result.logprobs.iter().all(Option::is_none), - "{name}: no-logprobs request should stay on the no-copy path" - ); - } else { - assert!( - result.logprobs.iter().all(Option::is_some), - "{name}: requested logprobs should be present" - ); - } info!(" PASS: {name} → {} tokens", result.tokens.len()); } } - // ── 5. Consumer drop safety ───────────────────────────────────────── - info!("=== Phase 5: Consumer drop ==="); + // ── 5. Cancellation safety ───────────────────────────────────────── + info!("=== Phase 5: Request cancellation ==="); { let prompt_tokens = tokenizer.encode("Hello", false).expect("encode failed"); - let (token_tx, rx) = TokenSink::standalone(); - drop(rx); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams::default(), - max_tokens: 10, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: None, - prompt_logprobs: None, - }) - .expect("submit failed"); - std::thread::sleep(std::time::Duration::from_millis(500)); - info!(" PASS: consumer drop handled"); + let control = handle.submit(common::request( + prompt_tokens, + SamplingParams::default(), + 10, + )); + control.abort(); } // Verify scheduler survives let (tokens, _) = generate_tokens(handle, tokenizer, "Hello", 5); let text = tokenizer.decode(&tokens, true).expect("decode failed"); - assert!(!text.is_empty(), "scheduler dead after consumer drop"); - info!(" PASS: scheduler survived consumer drop"); + assert!(!text.is_empty(), "scheduler dead after cancellation"); + info!(" PASS: scheduler survived cancellation"); info!("All Qwen3.5 scheduler tests passed for {label}!"); } @@ -690,10 +598,11 @@ fn test_e2e_qwen35_scheduler() { pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, ) .expect("Failed to start Qwen3.5 scheduler"); + let mut handle = EngineHarness::new(handle); info!("scheduler loaded in {:.2?}", start.elapsed()); let max_context_tokens = max_position_embeddings(&model_path); - run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP1"); + run_full_scheduler_e2e(&mut handle, &tokenizer, max_context_tokens, "TP1"); } #[test] @@ -727,7 +636,8 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { 0, ) .expect("Failed to start Qwen3.5 default-Off scheduler"); - let mut off_rx = submit_repeated_token_request( + let mut off_handle = EngineHarness::new(off_handle); + let off_request = submit_repeated_token_request( &off_handle, "overlap-off-reference", seed_token, @@ -735,9 +645,10 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { 2, ); let off = collect_generation_with_timeout( - &mut off_rx, + &mut off_handle, + &off_request, "overlap-off-reference", - 0, + None, std::time::Duration::from_secs(30), ); assert_eq!( @@ -767,33 +678,44 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { 0, ) .expect("Failed to start Qwen3.5 auto + shared-SM scheduler"); - let mut auto_load = auto_handle - .metrics_watch() - .expect("scheduler must expose metrics"); + let mut auto_handle = EngineHarness::new(auto_handle); - let mut auto_active_rx = submit_repeated_token_request( + let auto_active_request = submit_repeated_token_request( &auto_handle, "overlap-auto-last-decoder", seed_token, 512, 128, ); - wait_for_first_token(&mut auto_active_rx, "overlap-auto-last-decoder"); - let _ = drain_tokens(&mut auto_active_rx, "overlap-auto-last-decoder"); - let mut auto_prefill_rx = submit_repeated_token_request( + wait_for_first_token( + &mut auto_handle, + &auto_active_request, + "overlap-auto-last-decoder", + ); + let _ = drain_tokens( + &mut auto_handle, + &auto_active_request, + "overlap-auto-last-decoder", + ); + let auto_prefill_request = submit_repeated_token_request( &auto_handle, "overlap-auto-inflight-prefill", seed_token, 8192, 2, ); - wait_for_running_requests(&mut auto_load, 2, std::time::Duration::from_secs(10)); - assert_no_generated_event(&mut auto_prefill_rx, "overlap-auto-inflight-prefill"); - drop(auto_active_rx); + wait_for_running_requests(&auto_handle, 2, std::time::Duration::from_secs(10)); + assert_no_generated_event( + &mut auto_handle, + &auto_prefill_request, + "overlap-auto-inflight-prefill", + ); + auto_active_request.abort(); let auto_prefill = collect_generation_with_timeout( - &mut auto_prefill_rx, + &mut auto_handle, + &auto_prefill_request, "overlap-auto-inflight-prefill", - 0, + None, std::time::Duration::from_secs(30), ); assert_eq!( @@ -822,28 +744,27 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { 0, ) .expect("Failed to start Qwen3.5 shared-SM scheduler"); - let mut load = handle - .metrics_watch() - .expect("scheduler must expose metrics"); + let mut handle = EngineHarness::new(handle); - let mut active_rx = + let active_request = submit_repeated_token_request(&handle, "overlap-last-decoder", seed_token, 512, 128); - wait_for_first_token(&mut active_rx, "overlap-last-decoder"); - let _ = drain_tokens(&mut active_rx, "overlap-last-decoder"); - let mut prefill_rx = + wait_for_first_token(&mut handle, &active_request, "overlap-last-decoder"); + let _ = drain_tokens(&mut handle, &active_request, "overlap-last-decoder"); + let prefill_request = submit_repeated_token_request(&handle, "overlap-inflight-prefill", seed_token, 8192, 2); - wait_for_running_requests(&mut load, 2, std::time::Duration::from_secs(10)); - let _ = drain_tokens(&mut active_rx, "overlap-last-decoder"); + wait_for_running_requests(&handle, 2, std::time::Duration::from_secs(10)); + let _ = drain_tokens(&mut handle, &active_request, "overlap-last-decoder"); for _ in 0..2 { - wait_for_first_token(&mut active_rx, "overlap-last-decoder"); - assert_no_generated_event(&mut prefill_rx, "overlap-inflight-prefill"); + wait_for_first_token(&mut handle, &active_request, "overlap-last-decoder"); + assert_no_generated_event(&mut handle, &prefill_request, "overlap-inflight-prefill"); } - drop(active_rx); + active_request.abort(); let prefill = collect_generation_with_timeout( - &mut prefill_rx, + &mut handle, + &prefill_request, "overlap-inflight-prefill", - 0, + None, std::time::Duration::from_secs(30), ); assert_eq!( @@ -856,7 +777,49 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { "Shared-SM overlapped prefill must match the greedy default-Off reference" ); - let (tokens, finish_reason) = generate_tokens(&handle, &tokenizer, "Hello again", 2); + let streaming_decoder = + submit_repeated_token_request(&handle, "overlap-streaming-decoder", seed_token, 512, 1024); + wait_for_first_token(&mut handle, &streaming_decoder, "overlap-streaming-decoder"); + let _ = drain_tokens(&mut handle, &streaming_decoder, "overlap-streaming-decoder"); + let streaming_prefill = + submit_repeated_token_request(&handle, "overlap-streaming-prefill", seed_token, 128, 2); + let deadline = Instant::now() + std::time::Duration::from_secs(30); + let first = loop { + let update = recv_event_before( + &mut handle, + &streaming_prefill, + "overlap-streaming-prefill", + deadline, + ); + if !update.tokens.is_empty() || update.terminal.is_some() { + break update; + } + }; + let decoder_tokens = drain_tokens(&mut handle, &streaming_decoder, "overlap-streaming-decoder"); + streaming_decoder.abort(); + assert_eq!( + first.tokens.len(), + 1, + "completed async prefill must publish its first token before the next decode" + ); + assert!(first.terminal.is_none()); + assert!( + decoder_tokens > 0, + "the existing decoder must remain active" + ); + let remaining = collect_generation_with_timeout( + &mut handle, + &streaming_prefill, + "overlap-streaming-prefill", + None, + std::time::Duration::from_secs(30), + ); + assert_eq!(remaining.tokens.len(), 1); + assert_eq!(remaining.finish_reason, FinishReason::Length); + wait_for_running_requests(&handle, 0, std::time::Duration::from_secs(10)); + assert_eq!(handle.metrics().num_waiting_reqs, 0); + + let (tokens, finish_reason) = generate_tokens(&mut handle, &tokenizer, "Hello again", 2); assert_eq!( tokens.len(), 2, @@ -864,16 +827,28 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { ); assert_eq!(finish_reason, FinishReason::Length); - let mut shutdown_active_rx = + let shutdown_active_request = submit_repeated_token_request(&handle, "overlap-shutdown-decoder", seed_token, 512, 128); - wait_for_first_token(&mut shutdown_active_rx, "overlap-shutdown-decoder"); - let _ = drain_tokens(&mut shutdown_active_rx, "overlap-shutdown-decoder"); - let mut shutdown_prefill_rx = + wait_for_first_token( + &mut handle, + &shutdown_active_request, + "overlap-shutdown-decoder", + ); + let _ = drain_tokens( + &mut handle, + &shutdown_active_request, + "overlap-shutdown-decoder", + ); + let shutdown_prefill_request = submit_repeated_token_request(&handle, "overlap-shutdown-prefill", seed_token, 8192, 2); - wait_for_running_requests(&mut load, 2, std::time::Duration::from_secs(10)); - assert_no_generated_event(&mut shutdown_prefill_rx, "overlap-shutdown-prefill"); - drop(shutdown_active_rx); - drop(shutdown_prefill_rx); + wait_for_running_requests(&handle, 2, std::time::Duration::from_secs(10)); + assert_no_generated_event( + &mut handle, + &shutdown_prefill_request, + "overlap-shutdown-prefill", + ); + shutdown_active_request.abort(); + shutdown_prefill_request.abort(); let (done_tx, done_rx) = std::sync::mpsc::channel(); let shutdown = std::thread::spawn(move || { @@ -908,10 +883,11 @@ fn test_e2e_qwen35_scheduler_tp2() { pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, ) .expect("Failed to start Qwen3.5 TP2 scheduler"); + let mut handle = EngineHarness::new(handle); info!("TP2 scheduler loaded in {:.2?}", start.elapsed()); let max_context_tokens = max_position_embeddings(&model_path); - run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2"); + run_full_scheduler_e2e(&mut handle, &tokenizer, max_context_tokens, "TP2"); } #[test] @@ -939,8 +915,9 @@ fn test_e2e_qwen35_scheduler_tp2_graph() { pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, ) .expect("Failed to start Qwen3.5 TP2 graph scheduler"); + let mut handle = EngineHarness::new(handle); info!("TP2 graph scheduler loaded in {:.2?}", start.elapsed()); let max_context_tokens = max_position_embeddings(&model_path); - run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2 graph"); + run_full_scheduler_e2e(&mut handle, &tokenizer, max_context_tokens, "TP2 graph"); } diff --git a/pegainfer-qwen35/tests/prefix_cache.rs b/pegainfer-qwen35/tests/prefix_cache.rs index b6ab83b02..25574353c 100644 --- a/pegainfer-qwen35/tests/prefix_cache.rs +++ b/pegainfer-qwen35/tests/prefix_cache.rs @@ -2,22 +2,22 @@ //! //! The first request publishes the 256-token boundary. The second identical //! request must restore both state families at that boundary and report the -//! joint hit through `TokenEvent::Scheduled`. +//! joint hit through `RequestUpdate::cached_tokens`. use std::path::Path; -use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::FinishReason; -use pegainfer_frontend::engine::GenerateRequest; -use pegainfer_frontend::engine::TokenEvent; +use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::engine::TokenLogprob; -use pegainfer_frontend::engine::TokenSink; use pegainfer_frontend::sampler::SamplingParams; use pegainfer_qwen35::Qwen35LaunchOptions; use pegainfer_qwen35::Qwen35SchedulerPolicy; mod common; +use common::EngineHarness; +use common::RequestGuard; + const PREFIX_BOUNDARY: usize = 256; const PROMPT_TOKENS: usize = 320; const TRACE_TOKENS: usize = 8; @@ -30,7 +30,7 @@ fn model_path_or_skip() -> Option { common::model_path_or_skip("prefix_cache") } -fn start_engine(model_path: &str, tp_size: usize, prefix_cache_mib: usize) -> EngineHandle { +fn start_engine(model_path: &str, tp_size: usize, prefix_cache_mib: usize) -> EngineHarness { start_engine_with_graph(model_path, tp_size, prefix_cache_mib, tp_size == 1) } @@ -39,14 +39,15 @@ fn start_engine_with_graph( tp_size: usize, prefix_cache_mib: usize, cuda_graph: bool, -) -> EngineHandle { - pegainfer_qwen35::launch_with_options_policy_and_overlap( +) -> EngineHarness { + let handle = pegainfer_qwen35::launch_with_options_policy_and_overlap( Path::new(model_path), Qwen35LaunchOptions::new(0, tp_size, cuda_graph, 2, 1024, prefix_cache_mib), Qwen35SchedulerPolicy::Off, pegainfer_qwen35::Qwen35DecodeOverlap::Off, ) - .unwrap_or_else(|err| panic!("failed to start Qwen3.5 TP{tp_size} prefix-cache engine: {err}")) + .unwrap_or_else(|err| panic!("failed to start Qwen3.5 TP{tp_size} prefix-cache engine: {err}")); + EngineHarness::new(handle) } struct Generation { @@ -56,73 +57,62 @@ struct Generation { } fn submit( - handle: &EngineHandle, + handle: &EngineHarness, prompt_tokens: Vec, max_tokens: usize, logprobs: usize, -) -> pegainfer_frontend::engine::TokenStreamReceiver { - let (token_tx, rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams { - ignore_eos: true, - ..SamplingParams::default() - }, - max_tokens, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: (logprobs > 0).then_some(logprobs), - prompt_logprobs: None, - }) - .expect("submit failed"); - - rx +) -> RequestGuard { + let mut request = common::request( + prompt_tokens, + SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens, + ); + request.logprobs = (logprobs > 0).then_some(logprobs); + handle.submit(request) } fn generate( - handle: &EngineHandle, + handle: &mut EngineHarness, prompt_tokens: Vec, max_tokens: usize, logprobs: usize, ) -> Generation { - let mut rx = submit(handle, prompt_tokens, max_tokens, logprobs); + let control = submit(handle, prompt_tokens, max_tokens, logprobs); let mut cached_tokens = None; let mut generated_tokens = Vec::with_capacity(max_tokens); let mut generated_logprobs = Vec::with_capacity(max_tokens); loop { - match rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Scheduled { - cached_tokens: hit, .. - }) => { - cached_tokens = Some(hit); - } - Some(TokenEvent::Token { id, logprob }) => { - generated_tokens.push(id); - generated_logprobs.push(logprob); - } - Some(TokenEvent::PromptTokens { .. } | TokenEvent::KvTransfer { .. }) => {} - Some(TokenEvent::Finished { finish_reason, .. }) => { - assert_eq!(finish_reason, FinishReason::Length); - return Generation { - cached_tokens: cached_tokens.expect("request did not emit Scheduled"), - tokens: generated_tokens, - logprobs: generated_logprobs, - }; + let update = handle.next(control.id()); + if let Some(hit) = update.cached_tokens { + cached_tokens = Some(hit); + } + generated_tokens.extend(update.tokens); + generated_logprobs.extend(update.logprobs); + if let Some(terminal) = update.terminal { + match terminal { + Terminal::Finished { + reason, + completion_tokens, + .. + } => { + assert_eq!(reason, FinishReason::Length); + assert_eq!(completion_tokens, max_tokens); + return Generation { + cached_tokens: cached_tokens.expect("request did not report cached_tokens"), + tokens: generated_tokens, + logprobs: generated_logprobs, + }; + } + terminal => panic!("generation did not finish: {terminal:?}"), } - Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), - Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), - None => panic!("scheduler channel closed without Finished"), } } } -fn generate_one(handle: &EngineHandle, prompt_tokens: Vec) -> (usize, u32) { +fn generate_one(handle: &mut EngineHarness, prompt_tokens: Vec) -> (usize, u32) { let result = generate(handle, prompt_tokens, 1, 0); ( result.cached_tokens, @@ -209,11 +199,11 @@ fn joint_restore_and_unpinned_lru_eviction_preserve_output() { PROMPT_TOKENS, ); - let handle = start_engine(model_path, 1, PREFIX_CACHE_MIB); - let (cold_cached, cold_token) = generate_one(&handle, prompt_a.clone()); + let mut handle = start_engine(model_path, 1, PREFIX_CACHE_MIB); + let (cold_cached, cold_token) = generate_one(&mut handle, prompt_a.clone()); assert_eq!(cold_cached, 0, "first request must be cold"); - let (warm_cached, warm_token) = generate_one(&handle, prompt_a.clone()); + let (warm_cached, warm_token) = generate_one(&mut handle, prompt_a.clone()); assert_eq!( warm_cached, PREFIX_BOUNDARY, "the longest jointly committed boundary should be restored" @@ -223,22 +213,23 @@ fn joint_restore_and_unpinned_lru_eviction_preserve_output() { "joint restore must preserve greedy output" ); - let (beta_cold_cached, beta_token) = generate_one(&handle, prompt_b.clone()); + let (beta_cold_cached, beta_token) = generate_one(&mut handle, prompt_b.clone()); assert_eq!(beta_cold_cached, 0, "new beta prefix must be cold"); - let (alpha_touched_cached, _) = generate_one(&handle, prompt_a); + let (alpha_touched_cached, _) = generate_one(&mut handle, prompt_a); assert_eq!( alpha_touched_cached, PREFIX_BOUNDARY, "alpha lookup must refresh its snapshot LRU position" ); - let (gamma_cold_cached, _) = generate_one(&handle, prompt_c); + let (gamma_cold_cached, _) = generate_one(&mut handle, prompt_c); assert_eq!( gamma_cold_cached, 0, "new gamma prefix must insert under snapshot pressure" ); - let (beta_after_eviction_cached, beta_after_eviction_token) = generate_one(&handle, prompt_b); + let (beta_after_eviction_cached, beta_after_eviction_token) = + generate_one(&mut handle, prompt_b); assert_eq!( beta_after_eviction_cached, 0, "beta KV may remain resident, but its evicted snapshot must force a joint miss" @@ -260,23 +251,23 @@ fn boundary_selection_and_multitoken_restore_preserve_logits() { "Boundary coverage checks exact alignment, prefix extension, and joint recurrent state restore. ", 576, ); - let handle = start_engine(&model_path, 1, 512); + let mut handle = start_engine(&model_path, 1, 512); - let cold = generate(&handle, long_prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); + let cold = generate(&mut handle, long_prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); assert_eq!(cold.cached_tokens, 0); - let warm = generate(&handle, long_prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); + let warm = generate(&mut handle, long_prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); assert_eq!(warm.cached_tokens, 512); assert_trace_close("tp1 576-token restore", &cold, &warm); - let exact_512 = generate(&handle, long_prompt[..512].to_vec(), 1, 0); + let exact_512 = generate(&mut handle, long_prompt[..512].to_vec(), 1, 0); assert_eq!( exact_512.cached_tokens, 256, "an exactly aligned prompt must retain one token for final prefill" ); - let exact_256 = generate(&handle, long_prompt[..256].to_vec(), 1, 0); + let exact_256 = generate(&mut handle, long_prompt[..256].to_vec(), 1, 0); assert_eq!(exact_256.cached_tokens, 0); - let extended = generate(&handle, long_prompt[..320].to_vec(), 1, 0); + let extended = generate(&mut handle, long_prompt[..320].to_vec(), 1, 0); assert_eq!(extended.cached_tokens, 256); } @@ -294,41 +285,56 @@ fn restore_during_live_decode( "Joint prefix restore must preserve logits while another request decodes. ", 576, ); - let handle = pegainfer_qwen35::launch_with_options_policy_and_overlap( - Path::new(&model_path), - Qwen35LaunchOptions::new(0, tp_size, cuda_graph, 2, 1024, 128), - Qwen35SchedulerPolicy::Off, - overlap, - ) - .unwrap(); - let cold = generate(&handle, prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); - let mut background = submit(&handle, vec![9707], 1024, 0); - loop { - match background.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Token { .. }) => break, - Some(TokenEvent::Scheduled { .. }) => {} - event => panic!("background decode did not start: {event:?}"), - } - } - let warm = generate(&handle, prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); - assert_eq!(warm.cached_tokens, 512); - assert_trace_close("restore during live decode", &cold, &warm); - let mut background_tokens = 1; - while let Ok((_, event)) = background.try_recv() { - match event { - TokenEvent::Token { .. } => background_tokens += 1, - TokenEvent::Finished { .. } => { - panic!("background finished before the mixed-load probe") - } - TokenEvent::Error { message, .. } | TokenEvent::Rejected { message, .. } => { - panic!("{message}") + let start_engine = || { + pegainfer_qwen35::launch_with_options_policy_and_overlap( + Path::new(&model_path), + Qwen35LaunchOptions::new(0, tp_size, cuda_graph, 2, 1024, 128), + Qwen35SchedulerPolicy::Off, + overlap, + ) + .map(EngineHarness::new) + .unwrap() + }; + let generate_with_background = |handle: &mut EngineHarness| { + let background = submit(handle, vec![9707], 1024, 0); + let mut background_tokens = loop { + let update = handle.next(background.id()); + assert!( + update.terminal.is_none(), + "background decode ended before the probe" + ); + if !update.tokens.is_empty() { + break update.tokens.len(); } - _ => {} + }; + let generation = generate(handle, prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); + while let Some(update) = handle.try_next(background.id()) { + background_tokens += update.tokens.len(); + assert!( + update.terminal.is_none(), + "background finished before the mixed-load probe" + ); } - } - assert!(background_tokens < 1024); - drop(background); // Exercise cancellation cleanup followed by another cache hit. - let again = generate(&handle, prompt, TRACE_TOKENS, TOP_LOGPROBS); + assert!(background_tokens < 1024); + background.abort(); + generation + }; + + // Keep a cold reference under the same background decode as the warm request. + let mixed_cold = { + let mut reference = start_engine(); + let generation = generate_with_background(&mut reference); + assert_eq!(generation.cached_tokens, 0); + generation + }; + let mut handle = start_engine(); + let cold = generate(&mut handle, prompt.clone(), TRACE_TOKENS, TOP_LOGPROBS); + assert_eq!(cold.cached_tokens, 0); + assert_eq!(mixed_cold.tokens, cold.tokens); + let warm = generate_with_background(&mut handle); + assert_eq!(warm.cached_tokens, 512); + assert_trace_close("restore during live decode", &mixed_cold, &warm); + let again = generate(&mut handle, prompt, TRACE_TOKENS, TOP_LOGPROBS); assert_eq!(again.cached_tokens, 512); assert_trace_close("restore after cancellation", &cold, &again); } diff --git a/pegainfer-qwen35/tests/sampling_behavior.rs b/pegainfer-qwen35/tests/sampling_behavior.rs index ac5ba155a..9c0e929cf 100644 --- a/pegainfer-qwen35/tests/sampling_behavior.rs +++ b/pegainfer-qwen35/tests/sampling_behavior.rs @@ -10,15 +10,14 @@ use std::path::Path; -use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; -use pegainfer_frontend::engine::GenerateRequest; -use pegainfer_frontend::engine::TokenEvent; -use pegainfer_frontend::engine::TokenSink; +use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; mod common; +use common::EngineHarness; + const GENERATED_TOKENS: usize = 32; fn params(mut params: SamplingParams) -> SamplingParams { @@ -27,38 +26,29 @@ fn params(mut params: SamplingParams) -> SamplingParams { } /// Submit one request and collect the generated token ids until `Finished`. -fn generate(handle: &EngineHandle, prompt_tokens: Vec, params: SamplingParams) -> Vec { - let (token_tx, mut rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params, - max_tokens: GENERATED_TOKENS, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: None, - prompt_logprobs: None, - }) - .expect("submit failed"); - +fn generate( + handle: &mut EngineHarness, + prompt_tokens: Vec, + params: SamplingParams, +) -> Vec { + let control = handle.submit(common::request(prompt_tokens, params, GENERATED_TOKENS)); let mut tokens = Vec::new(); loop { - match rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Token { id, .. }) => tokens.push(id), - Some( - TokenEvent::Scheduled { .. } - | TokenEvent::PromptTokens { .. } - | TokenEvent::KvTransfer { .. }, - ) => {} - Some(TokenEvent::Finished { .. }) => return tokens, - Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), - Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), - None => panic!("scheduler channel closed without Finished"), + let update = handle.next(control.id()); + tokens.extend(update.tokens); + if let Some(terminal) = update.terminal { + match terminal { + Terminal::Finished { + reason, + completion_tokens, + .. + } => { + assert_eq!(reason, pegainfer_frontend::engine::FinishReason::Length); + assert_eq!(completion_tokens, GENERATED_TOKENS); + return tokens; + } + terminal => panic!("generation did not finish: {terminal:?}"), + } } } } @@ -82,6 +72,7 @@ fn sampling_params_steer_the_qwen35_sampler() { pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, ) .expect("failed to start Qwen3.5 engine"); + let mut handle = EngineHarness::new(handle); let tokenizer = common::load_tokenizer(&model_path); let prompt = "Here is a short story about a dragon. Once upon a time"; @@ -89,17 +80,17 @@ fn sampling_params_steer_the_qwen35_sampler() { let greedy_params = params(SamplingParams::default()); - let greedy = generate(&handle, prompt_tokens.clone(), greedy_params); + let greedy = generate(&mut handle, prompt_tokens.clone(), greedy_params); assert_eq!( greedy.len(), GENERATED_TOKENS, "ignore_eos should force a full 32-token generation" ); - let greedy_again = generate(&handle, prompt_tokens.clone(), greedy_params); + let greedy_again = generate(&mut handle, prompt_tokens.clone(), greedy_params); assert_eq!(greedy, greedy_again, "greedy decode must be deterministic"); let top_k_one = generate( - &handle, + &mut handle, prompt_tokens.clone(), params(SamplingParams { temperature: 0.8, @@ -110,7 +101,7 @@ fn sampling_params_steer_the_qwen35_sampler() { assert_eq!(top_k_one, greedy, "top_k=1 must collapse to greedy"); let top_p_tiny = generate( - &handle, + &mut handle, prompt_tokens.clone(), params(SamplingParams { temperature: 1.0, @@ -127,7 +118,7 @@ fn sampling_params_steer_the_qwen35_sampler() { ..SamplingParams::default() }); let runs: Vec> = (0..4) - .map(|_| generate(&handle, prompt_tokens.clone(), hot)) + .map(|_| generate(&mut handle, prompt_tokens.clone(), hot)) .collect(); assert!( runs.iter().any(|run| *run != runs[0]), diff --git a/pegainfer-qwen35/tests/serving_tp2.rs b/pegainfer-qwen35/tests/serving_tp2.rs index c721f9c43..39b161761 100644 --- a/pegainfer-qwen35/tests/serving_tp2.rs +++ b/pegainfer-qwen35/tests/serving_tp2.rs @@ -93,7 +93,7 @@ async fn spawn_ready_server( let server_shutdown = shutdown.clone(); let mut task = tokio::spawn(async move { pegainfer_frontend::vllm::serve( - std::future::ready(Ok(pegainfer_frontend::engine::LaunchedEngine::Handle( + std::future::ready(Ok(pegainfer_frontend::engine::LaunchedEngine::Stepped( handle, ))), &frontend_model_path, diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index 89b50b9ae..5cc5d8396 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -12,7 +12,7 @@ # # PEGAINFER_TEST_MODEL_PATH= \ # PEGAINFER_NVFP4_MODEL= \ -# [PEGAINFER_GATE_GPU=] scripts/gemma4_gates.sh [filter] +# [PEGAINFER_GATE_GPU=] scripts/gemma4_gates.sh [filter] # # A filter runs the subset of manifest gates whose names contain it; the # membership check still covers the whole manifest. @@ -68,6 +68,7 @@ GATES_SERVING_CONTRACT=( "gpu,ckpt engine::lane_gates_logprobs::prompt_scores_bypass_the_prefix_cache_and_match_teacher_forced_decode" "gpu,ckpt engine::lane_gates_logprobs::a_scored_prompt_beside_a_live_batch_is_prefilled_whole" "gpu,ckpt engine::lane_gates_logprobs::a_low_slot_chunked_pool_scores_up_to_what_it_holds" + "gpu,ckpt engine::lane_gates_logprobs::a_failed_prefill_costs_that_request_not_the_engine" ) GATES_KV_AND_LANES=( "gpu,ckpt,fixtures serve::oracle::incremental_serving_matches_recompute" @@ -111,6 +112,12 @@ GATES_KERNELS_HD256_FP8_POOL=( "gpu decode_wrapper_without_fp8_twin_refuses_e4m3" "gpu the_generated_windowed_prefill_refuses_e4m3" ) +GATES_TENSOR_PARALLEL=( + "tp2,ckpt engine::lane_gates_tp::the_two_rank_engine_matches_one_rank" + "tp2,ckpt engine::lane_gates_tp::the_two_rank_engine_scores_prompt_logprobs" + "tp2,ckpt engine::lane_gates_tp::the_two_rank_engine_is_prefix_consistent" + "tp2,ckpt,hf31 engine::lane_gates_tp::the_two_rank_engine_matches_the_hf_reference" +) MANIFEST_LIB=( "${GATES_NUMERIC_PARITY[@]}" "${GATES_ADMISSION[@]}" @@ -121,6 +128,7 @@ MANIFEST_LIB=( "${GATES_DEVICE[@]}" "${GATES_ROUTED[@]}" "${GATES_TILELANG_GLOBAL[@]}" + "${GATES_TENSOR_PARALLEL[@]}" ) GATES_FP8_PROFILE=( "serve::oracle::context_waypoints_match_hf" @@ -212,48 +220,81 @@ moe_ckpt= gpu_uuid= gpu_lock_fd= -require_gpu() { - command -v nvidia-smi >/dev/null 2>&1 || die "nvidia-smi is unavailable, so no device can be claimed" +# Claim `count` devices for this run and export them as one +# `CUDA_VISIBLE_DEVICES` list. The single-device gates and the two-rank gate +# draw from the same reservation, so a full-suite run never tries to lock a +# device it already holds, and a pair may be named in `PEGAINFER_GATE_GPU`. +# `PEGAINFER_GATE_GPU` names the devices ("a" or "a,b"); otherwise they come +# from `CUDA_VISIBLE_DEVICES`, else the first `count` visible devices. +require_devices() { + local count=$1 + command -v nvidia-smi >/dev/null 2>&1 || die "nvidia-smi is unavailable, so no devices can be claimed" command -v flock >/dev/null 2>&1 || die "flock is unavailable, so device ownership cannot be enforced" local selector=${PEGAINFER_GATE_GPU:-} + local explicit=0 if [ -z "$selector" ] && [ "${CUDA_VISIBLE_DEVICES+x}" = x ]; then [ -n "$CUDA_VISIBLE_DEVICES" ] || die \ "CUDA_VISIBLE_DEVICES is empty; set PEGAINFER_GATE_GPU to claim a device" - [[ $CUDA_VISIBLE_DEVICES != *,* ]] || die \ - "CUDA_VISIBLE_DEVICES must name one device; set PEGAINFER_GATE_GPU explicitly" selector=$CUDA_VISIBLE_DEVICES fi - selector=${selector:-0} - [[ $selector != *,* ]] || die "PEGAINFER_GATE_GPU must name exactly one device" - - local rows=() row compute_mode lock_key lock_path - mapfile -t rows < <( - nvidia-smi -i "$selector" --query-gpu=uuid,compute_mode --format=csv,noheader 2>/dev/null - ) - [ ${#rows[@]} -eq 1 ] || die "device selector $selector does not resolve to one GPU" - row=${rows[0]} - gpu_uuid=${row%%,*} - gpu_uuid=${gpu_uuid//[[:space:]]/} - compute_mode=${row#*,} - compute_mode=${compute_mode#"${compute_mode%%[![:space:]]*}"} - compute_mode=${compute_mode%"${compute_mode##*[![:space:]]}"} - [ "$compute_mode" != Prohibited ] || die "GPU $gpu_uuid prohibits compute contexts" - [[ $gpu_uuid =~ ^[A-Za-z0-9._:/-]+$ ]] || die "nvidia-smi returned an unsafe GPU identity" + [ -n "$selector" ] && explicit=1 + local picks=() + [ -z "$selector" ] || IFS=',' read -r -a picks <<<"$selector" + if [ ${#picks[@]} -lt "$count" ]; then + # An explicit selector is never substituted: naming one device (as a SLURM + # job does) means "use this one", so a two-rank run drops its two-rank gates + # rather than take a different device; naming none at all is a mistake. + if [ "$explicit" -eq 1 ]; then + [ ${#picks[@]} -gt 0 ] || die "the device selector names no device" + return 1 + fi + picks=() + mapfile -t picks < <(nvidia-smi --query-gpu=index --format=csv,noheader 2>/dev/null) + fi + # Not enough devices and nobody named them: `return 1` lets a run that can drop + # its two-rank gates carry on rather than fail. + [ ${#picks[@]} -ge "$count" ] || return 1 + + local uuids=() rows=() row compute_mode sel uuid seen + for sel in "${picks[@]:0:$count}"; do + mapfile -t rows < <( + nvidia-smi -i "$sel" --query-gpu=uuid,compute_mode --format=csv,noheader 2>/dev/null + ) + [ ${#rows[@]} -eq 1 ] || die "device selector $sel does not resolve to one GPU" + row=${rows[0]} + uuid=${row%%,*} + uuid=${uuid//[[:space:]]/} + compute_mode=${row#*,} + compute_mode=${compute_mode#"${compute_mode%%[![:space:]]*}"} + compute_mode=${compute_mode%"${compute_mode##*[![:space:]]}"} + [ "$compute_mode" != Prohibited ] || die "GPU $uuid prohibits compute contexts" + [[ $uuid =~ ^[A-Za-z0-9._:/-]+$ ]] || die "nvidia-smi returned an unsafe GPU identity" + for seen in "${uuids[@]:-}"; do + [ "$seen" != "$uuid" ] || die "device selector names $uuid twice" + done + uuids+=("$uuid") + done + gpu_uuid=$(IFS=,; printf '%s' "${uuids[*]}") export CUDA_VISIBLE_DEVICES=$gpu_uuid - lock_key=${gpu_uuid//\//_} - lock_key=${lock_key//:/_} - lock_path=$GPU_LOCK_ROOT/pegainfer-gemma4-gates-$lock_key.lock - if (umask 022; set -o noclobber; : >"$lock_path") 2>/dev/null; then - : - elif [ ! -e "$lock_path" ]; then - die "cannot create device lock $lock_path" - fi - # A read-only descriptor lets separate Unix accounts lock the same inode. - exec {gpu_lock_fd}<"$lock_path" || die "cannot open device lock $lock_path" - flock -n "$gpu_lock_fd" || die "GPU $gpu_uuid is already owned by another Gemma 4 gate runner" - echo "gemma4 gates: claimed GPU $gpu_uuid (selector $selector)" + local lock_fds=() fd uuid2 lock_key lock_path + for uuid2 in "${uuids[@]}"; do + lock_key=${uuid2//\//_} + lock_key=${lock_key//:/_} + lock_path=$GPU_LOCK_ROOT/pegainfer-gemma4-gates-$lock_key.lock + if (umask 022; set -o noclobber; : >"$lock_path") 2>/dev/null; then + : + elif [ ! -e "$lock_path" ]; then + die "cannot create device lock $lock_path" + fi + # A read-only descriptor lets separate Unix accounts lock the same inode. + exec {fd}<"$lock_path" || die "cannot open device lock $lock_path" + flock -n "$fd" || die "GPU $uuid2 is already owned by another Gemma 4 gate runner" + lock_fds+=("$fd") + done + gpu_lock_fd=${lock_fds[0]} + echo "gemma4 gates: claimed $count GPU(s) $gpu_uuid" echo "gemma4 gates: storage profile $gate_storage" if [ -n "${PEGAINFER_KV_FP8:-}" ]; then echo "gemma4 gates: PEGAINFER_KV_FP8=$PEGAINFER_KV_FP8" @@ -552,13 +593,62 @@ if [ -z "$ROUTED_FIXTURE_TAG" ]; then [ ${#selected[@]} -gt 0 ] || die "every selected gate needed the routed fixture set" fi +# The 31B Hugging Face fixture is a separate artifact from its checkpoint, so a +# gate that compares against it leaves the run by name when it is not selected. +# Only that the variable is set is checked here: the fixture's own manifest is +# held against the checkpoint under test by `golden_bytes`, so checking the file +# name on top of that could not catch a mismatched dump — only refuse a valid one +# that happens to be named otherwise. +hf31_mismatch= +require_hf31() { + if [ -z "${PEGAINFER_GEMMA4_GOLDEN:-}" ]; then + hf31_mismatch="PEGAINFER_GEMMA4_GOLDEN is unset (point it at the 31B HF fixture)" + fi +} + # --- prerequisites: the union over what this run selected, and no more ----- needs=" " for entry in "${selected[@]}"; do needs="$needs${entry%%|*} "; done needs=" ${needs//,/ } " demanded="" -for want in gpu ckpt moeckpt prompts fixtures routedfixtures chatgolden tlgeom; do - case "$needs" in *" $want "*) "require_$want"; demanded="$demanded $want" ;; esac +# One reservation covers both the single-device gates and the two-rank gate: +# claim two devices when a tp2 gate is selected, one otherwise. Claiming per +# gate type would try to lock a device this same process already holds. +device_count=0 +case "$needs" in + *" tp2 "*) device_count=2 ;; + *" gpu "*) device_count=1 ;; +esac +if [ "$device_count" -gt 0 ] && ! require_devices "$device_count"; then + # Two devices were wanted and only one is visible; a one-GPU host should still + # run the rest, so drop the two-rank gates by name and claim the single one. + [ "$device_count" -eq 2 ] || die "a device is required but none could be claimed" + kept=(); dropped=0 + for entry in "${selected[@]}"; do + case ",${entry%%|*}," in + *,tp2,*) dropped=$((dropped + 1)) ;; + *) kept+=("$entry") ;; + esac + done + selected=("${kept[@]}") + [ ${#selected[@]} -gt 0 ] || die \ + "only two-rank gates were selected and a second device is unavailable" + echo "gemma4 gates: not selected, a second device is unavailable: $dropped two-rank gate(s)" + needs=" " + for entry in "${selected[@]}"; do needs="$needs${entry%%|*} "; done + needs=" ${needs//,/ } " + require_devices 1 || die "a device is required but none could be claimed" +fi +for want in gpu tp2 ckpt moeckpt prompts fixtures routedfixtures chatgolden tlgeom hf31; do + case "$needs" in + *" $want "*) + case "$want" in + gpu|tp2) : ;; + *) "require_$want" ;; + esac + demanded="$demanded $want" + ;; + esac done echo "gemma4 gates: prerequisites$demanded" @@ -576,6 +666,19 @@ if [ -n "$tlgeom_mismatch" ]; then [ ${#selected[@]} -gt 0 ] || die "every selected gate needed the kernels' own geometry" fi +# And the same for a gate that needs the 31B Hugging Face fixture. +if [ -n "$hf31_mismatch" ]; then + kept=() + for entry in "${selected[@]}"; do + case ",${entry%%|*}," in + *,hf31,*) echo "gemma4 gates: not selected, $hf31_mismatch: ${entry##*|}" ;; + *) kept+=("$entry") ;; + esac + done + selected=("${kept[@]}") + [ ${#selected[@]} -gt 0 ] || die "every selected gate needed the 31B HF fixture" +fi + echo "gemma4 gates: source $(git rev-parse HEAD)$([ -n "$(git status --porcelain)" ] && echo ' (dirty)')" [ -z "$ckpt" ] || echo "gemma4 gates: checkpoint $ckpt" [ -z "$moe_ckpt" ] || echo "gemma4 gates: routed checkpoint $moe_ckpt"