diff --git a/.agents/claims/CLAIM-ROCM-GEMMA4-INDEXED-MAX-T.md b/.agents/claims/CLAIM-ROCM-GEMMA4-INDEXED-MAX-T.md new file mode 100644 index 000000000..6b806b50b --- /dev/null +++ b/.agents/claims/CLAIM-ROCM-GEMMA4-INDEXED-MAX-T.md @@ -0,0 +1,5 @@ +# CLAIM-ROCM-GEMMA4-INDEXED-MAX-T + +| Claim | Row IDs | Agent | Worktree | Branch | Owned scope | State | Last update | +|---|---|---|---|---|---|---|---| +| `CLAIM-ROCM-GEMMA4-INDEXED-MAX-T` | `BACKEND-ROCM` (slug `ROCM-GEMMA4-INDEXED-MAX-T`, issue #838) | hermes-vllm (lab), helper | `/home/don/llms/vllm.cpp-indexed-max-t` | `row/ROCM-GEMMA4-INDEXED-MAX-T` | Owns ONLY: widen `gemma4_moe.cpp` T==1 indexed gate to T≤63 via `VT_GEMMA4_DECODE_INDEXED_MAX_T` default 63, using existing per-token indexed helpers, plus tensor oracle vs serial reference. **EXCLUDED:** packed `ExpertGeGLUFp8TopKIndexedBatched`, DEVICE_GROUP, INDEXED_NOSYNC, #837, #839, #697. Independent history from abandoned `row/ROCM-GEMMA4-XDEV-MOE`. | `IMPLEMENTING` | 2026-08-15 — d973/0f32 repair: retire inside acc_idx scope, Retire fail-closed, independent serial ref, injectable arms, owner kind | diff --git a/.agents/issue-index.md b/.agents/issue-index.md index d25b7da21..f28dc561e 100644 --- a/.agents/issue-index.md +++ b/.agents/issue-index.md @@ -453,3 +453,4 @@ rather than merged. `scripts/check-agent-record.py` gates both. | [#1353](https://github.com/mudler/vllm.cpp/issues/1353) | `ENG-RECORD-ANCHOR-RATCHET` | A full disk makes `scripts/agent-preflight.sh` report record and policy defects that do not exist. Measured at `63ff58272` with 896M free of 447G: ten suites go red together — `test_check_release_binary_contract`, `test_release_postpublish_audit`, `test_check_container_matrix`, `test_release_index`, `test_release_metadata`, `test_release_accelerator_metadata`, `test_release_macos_metadata`, `test_release_windows_metadata`, `test_agent_role`, `test_agent_onboard` — and the trailing summary names all ten without mentioning the disk. The cause appears only inside one suite's captured output several screens up (`error: copy-fd: write returned: No space left on device`), interleaved with fixture failures that read exactly like findings (`ERROR: x: see (#157) and #174`, `ERROR: x: landed without a row PR`). All ten build a scratch git repository in a temporary directory, so all ten fail together and all ten fail toward a verdict about RECORDS rather than about the environment — the `.agents/verification.md` broken-instrument shape, and an agent reading it has every reason to start repairing records that were never broken. Reclaiming 2.7G of one row's own build tree and re-running the identical command turned all ten green with no tree change (`86 gates, 86 ok, 0 FAIL, 0 SKIP`, exit 0). The fix is a precondition, not a suppression: check free space before the suites that need a scratch repository and refuse naming the disk, exactly as `test_cpu_x86_llamacpp_floor` refuses to measure under contention ([#618](https://github.com/mudler/vllm.cpp/issues/618)). A gate that cannot run must say so rather than return a verdict. Found while landing [#1332](https://github.com/mudler/vllm.cpp/issues/1332) M0+M1 and NOT fixed in that flow, because it changes preflight semantics and adds a refusal path, which `AGENTS.md` routes to the normal row, spec and fresh-review path | bug | | [#1376](https://github.com/mudler/vllm.cpp/issues/1376) | `ENG-CUDAGRAPH-BREAK` | `main` was red on `tests/scripts/test_check_gate_commands.py`, measured at `601b576c6` in a detached worktree of `origin/main`: 8 failures of 44 tests, every one a comparison between the computed runnable population and `RUNNABLE_BASELINE`. `ENG-CUDAGRAPH-BREAK` was in the first and absent from the second. Cause: W5 of that row ([#1361](https://github.com/mudler/vllm.cpp/issues/1361)) filled its spec's `## Gates` section with runnable evidence, including a named test binary with its case and assertion counts and an exit status, which is exactly what moves a row into the runnable population. The ratchet's own error text instructs a re-pin in the SAME change, and the re-pin was not made. This is the growth case the ratchet exists to force a decision about, not a defect in that row's work. **It landed with no remote verdict**: the continuous integration lane that would have caught it independently has not executed for this repository since roughly 07:43Z on 19 August 2026, with runs queueing and none starting while GitHub reports Actions operational. FIXED IN FLOW while merging `origin/main` into `row/ENG-HF-MODEL-DOWNLOAD` for [#1280](https://github.com/mudler/vllm.cpp/issues/1280), because the fix is small and clear and a red `main` blocks every other row's gate. The entry is added with a justifying comment in the form the neighbouring entries use, no checker semantics change, and no test is weakened. After the re-pin the suite reports 45 tests OK and the audit reads 39 runnable of 119 gated rows | bug | | [#1375](https://github.com/mudler/vllm.cpp/issues/1375) | `MODEL-DIFFUSION-LTX25` | First end-to-end per-forward cost for the FULL 21.004 B LTX-2.5 DiT on GB10, measured on run `20260819T150230Z` with binary `0a43a750` built from [`7b9e207b1`](https://github.com/mudler/vllm.cpp/commit/7b9e207b1) (#1252). At 1024x576/25f (2304 latent tokens) the governor resolved **7 forward starts from the GPU busy/idle edge counter** and measured `per_forward ~162.0 s` with `first_dit = 481.5 s`, so the recipe's fixed 60 forwards (30 steps x 2 CFG legs, `ltx2_pipeline.cpp:521-529`) project **10 803 s against the rung's 7 153 s budget** and the rung was refused rather than run to the wall. The same lease then COMPLETED 768x448/25f (1344 tokens) in 2990 s, so the ceiling is geometry against lease length, not a defect. TWO instrument facts belong with the number, because both have already caused a wrong reading: `gpu_edges=0` means the GPU never went idle long enough to sample an edge (SATURATED), not that no work ran — this rung sampled 85% of 3191 samples above 50% utilisation; and `eu-stack` resolves no frames in the rc worker container, so phase attribution came from the cpu%/rss signature rather than from symbols. Owned by the LTX-2.5 row; spec [`ltx-2-5.md`](specs/ltx-2-5.md) | measurement | +| [#838](https://github.com/mudler/vllm.cpp/issues/838) | `ROCM-GEMMA4-INDEXED-MAX-T` | Gemma4 indexed MoE gate widened to T<=63 with injectable dispatch + retire-in-scope | perf | diff --git a/.agents/specs/rocm-gemma4-indexed-max-t-donor-gate-911-933.log b/.agents/specs/rocm-gemma4-indexed-max-t-donor-gate-911-933.log new file mode 100644 index 000000000..7a19fd7e4 --- /dev/null +++ b/.agents/specs/rocm-gemma4-indexed-max-t-donor-gate-911-933.log @@ -0,0 +1,23 @@ + // Device-indexed FP8 MoE for small packed T — no router D2H, no host top-k gather. + // T=1: original KEEP decode path (hipGraph-stable TLS acc). + // T=2..min(MAX, batch_min-1): P1 concurrent-decode gap (research 2026-08-11). + // VT_GEMMA4_DECODE_INDEXED_MAX_T: default 63; =1 → T=1 only; clamp [1,63]. + static const int64_t kDecodeIndexedMaxT = []() -> int64_t { + const char* e = std::getenv("VT_GEMMA4_DECODE_INDEXED_MAX_T"); + if (e == nullptr || !e[0]) return 63; + const long v = std::strtol(e, nullptr, 10); + if (v < 1) return 1; + if (v > 63) return 63; + return static_cast(v); + }(); + // P0 overlap enabler (default OFF): keep the enqueue-only indexed path for + // prefill T so layer-split does not hipStreamSynchronize after router D2H. + // Not a KEEP speed path — per-token launches. Cap 8192 (pipe ubatch). + static const int64_t kPrefillIndexedNoSyncMaxT = []() -> int64_t { + const char* e = std::getenv("VT_GEMMA4_PREFILL_INDEXED_NOSYNC"); + if (!e || e[0] != '1') return 0; + return 8192; + }(); + const int64_t indexed_max_t = + (kPrefillIndexedNoSyncMaxT > kDecodeIndexedMaxT) ? kPrefillIndexedNoSyncMaxT + : kDecodeIndexedMaxT; diff --git a/.agents/specs/rocm-gemma4-indexed-max-t-donor-ok-1088-1110.log b/.agents/specs/rocm-gemma4-indexed-max-t-donor-ok-1088-1110.log new file mode 100644 index 000000000..dcdb36017 --- /dev/null +++ b/.agents/specs/rocm-gemma4-indexed-max-t-donor-ok-1088-1110.log @@ -0,0 +1,23 @@ + } + static std::atomic grp_fail{0}; + if (grp_fail.fetch_add(1) < 4) { + std::fprintf(stderr, "gemma4 moe: device-group FAIL T=%lld — falling back\n", + static_cast(T)); + std::fflush(stderr); + } + } + const bool indexed_ok_t = + (T >= 1 && T <= indexed_max_t && top_k <= 8 && top_k > 0 && fp8_res && + (kPrefillIndexedNoSyncMaxT > 0 || T < kPrefillBatchMinT)); + // Indexed only while T is below the group-by-expert crossover (unless nosync). + if (indexed_ok_t) { + // per-expert scale on device (once per layer/E). + struct EscTls { + int dev = -1; + int64_t E = 0; + const void* host_key = nullptr; + std::optional sc; + }; + static thread_local EscTls esc_slots[2]; + EscTls& esc = esc_slots[(compute_dev == 1) ? 1 : 0]; + float* escale_ptr = nullptr; diff --git a/.agents/specs/rocm-gemma4-indexed-max-t-donor.md b/.agents/specs/rocm-gemma4-indexed-max-t-donor.md new file mode 100644 index 000000000..1a31a5018 --- /dev/null +++ b/.agents/specs/rocm-gemma4-indexed-max-t-donor.md @@ -0,0 +1,21 @@ +# Donor evidence — #838 indexed-max-T + +Pinned **bytes**, not a dirty-tree HEAD. Implementation must copy these slices (or a later +immutable replacement that research re-reviews), not re-read `/home/don/llms/vllm.cpp`. + +Donor includes `kPrefillIndexedNoSyncMaxT` / `DEVICE_GROUP` neighbors. **Those are out of +scope for this row.** Only `kDecodeIndexedMaxT` (default 63, clamp [1,63]) and the +`indexed_ok_t` predicate **without** the nosync disjunct are in scope. + +| Field | Value | +|---|---| +| Donor tree | `/home/don/llms/vllm.cpp` | +| Donor git HEAD | `2bb4bd8a` (dirty; these slices are **uncommitted** on that tree) | +| File | `src/vllm/model_executor/models/gemma4_moe.cpp` | +| Slice A | `rocm-gemma4-indexed-max-t-donor-gate-911-933.log` lines 911–933 SHA256 `d0d28f3d55ff7d526475c9a2a1d028792cc245ba4cc2b421ebf994fed9b96e59` | +| Slice B | `rocm-gemma4-indexed-max-t-donor-ok-1088-1110.log` lines 1088–1110 SHA256 `5509f3f77dcadd023ce73743e13a2a1a0237d8dd0a8e25a98a765430fc577bd0` | +| Recipient | `origin/main` `3ce5a1dc` `gemma4_moe.cpp:735` (`T == 1` only) / `:1345` serial | +| Captured | 2026-08-14 | + +`sha256sum` of each slice file must match the table. Do not treat `2bb4bd8a` as a clean +donor commit. diff --git a/.agents/specs/rocm-gemma4-indexed-max-t.md b/.agents/specs/rocm-gemma4-indexed-max-t.md new file mode 100644 index 000000000..3f7cdf46f --- /dev/null +++ b/.agents/specs/rocm-gemma4-indexed-max-t.md @@ -0,0 +1,181 @@ +# Spec: ROCm Gemma-4 indexed MoE for T=2..63 + +- **Issue:** https://github.com/mudler/vllm.cpp/issues/838 +- **Row slug:** `ROCM-GEMMA4-INDEXED-MAX-T` — child of `BACKEND-ROCM` (#41). Separate from #697. +- **Worktree / branch (this unit only):** `/home/don/llms/vllm.cpp-indexed-max-t` · `row/ROCM-GEMMA4-INDEXED-MAX-T` +- **Base / recipient:** `origin/main` `3ce5a1dc` `gemma4_moe.cpp:735` / `:1345` +- **Donor bytes:** `.agents/specs/rocm-gemma4-indexed-max-t-donor-*.log` slices hashed in `rocm-gemma4-indexed-max-t-donor.md` (dirty lab `/home/don/llms/vllm.cpp` HEAD `2bb4bd8a` **plus uncommitted**; HEAD is not a clean donor). +- **Implementer:** hermes-vllm. **Reviewer:** research. **Operator/smoke:** coord. +- **Git:** spec-only first (coord `25c9` / research `5071` / BLOCK `64cb`); impl after spec GREEN **on this same row branch**. Independent RED/GREEN from #837/#839. One PR per row. No shared `row/ROCM-GEMMA4-XDEV-MOE` landing history. +- **Supersedes for review:** `20332292` (BLOCK) and preview `c4fbe6e9` (not spec-GREEN). + +## Now + +`IMPLEMENTING` — d973/0f32 repair on this row. `7c416eb6` is not a review target. + +**Not a confirmed fix.** Hypothesis (A) only: small-T **routing**. `9772` does not make this the T=2029 root. T=19 is **observed** on today's serial M1 route. Cause (serial M1 vs anything else) is **unconfirmed**. Do not call the serial path "racy" and do not say the T=19 hang "is this path." + +## Upstream / source of the port + +No vLLM Python equivalent. Source is the pinned donor slices. Product takes only the default-63 decode cap and the T<64 gate; **not** `PREFILL_INDEXED_NOSYNC` / `DEVICE_GROUP`. + +| Tree | Indexed gate | T=2..63 path | +|---|---|---| +| `origin/main` `3ce5a1dc` | `gemma4_moe.cpp:735` `if (T == 1 && fp8_res && top_k <= 8 && top_k > 0)` | falls through to serial `:1345` `RunGemma4Fp8TopKOnExpertDevice` | +| hanging `vllm.cpp-bc64fa-r2` `1b1baf43` | `:735` same T==1 | same serial | +| donor slices | `:914-922` `VT_GEMMA4_DECODE_INDEXED_MAX_T` default 63; `:1096-1100` `indexed_ok_t` | per-token `ExpertGeGLUFp8TopKIndexed` / `RunGemma4Fp8TopKIndexedOnExpertDevice` | + +Indexed **peer helper already exists on main**: `rocm_gemma4_experts.hip:543` `RunGemma4Fp8TopKIndexedOnExpertDevice`. Same-dev helper `vt::ExpertGeGLUFp8TopKIndexed` is already called at main `:775`. This row widens the **caller gate**, it does not invent a new kernel. + +`kPrefillBatchMinT` stays 64 (`:980`). T≥64 remains prefill-batch (#839). + +## Symptom this row owns + +T=19 warmup / short prefill is **observed** to take the serial M1 peer path (`RunGemma4Fp8TopKOnExpertDevice`). Lab never enters that helper when `indexed_max=63`. Distinct from the T≥64 accumulation **class**. Cause unconfirmed. + +## Scope + +1. Parse once (cached `static const`): + - `VT_GEMMA4_DECODE_INDEXED_MAX_T`: unset → **63**; `=1` → T=1 only; clamp `[1,63]`. +2. Replace the T==1 gate with: + `T >= 1 && T <= indexed_max_t && T < kPrefillBatchMinT && fp8_res && top_k <= 8 && top_k > 0`. +3. T=1: keep hipGraph-stable TLS acc **and** TLS `rw_idx` (do not `pool-Release` either). Key `RwIdxTls` by `(compute_dev, T*top_k)`. T>1 keeps a per-call pooled `rw_idx_owned`. +4. T>1: owned `[T,H]` bf16 buffer, per-token existing indexed helpers (same-dev or peer). +5. Document the env in `docs/ENVIRONMENT.md` in the **implementation** commit. + +## Out of scope + +- **`ExpertGeGLUFp8TopKIndexedBatched`** and any packed T≥2 same-dev batch kernel. Lab has this at `:1154-1161`. Product already **REJECTED** that family. Do not port it. +- `VT_GEMMA4_PREFILL_DEVICE_GROUP`, `VT_GEMMA4_PREFILL_INDEXED_NOSYNC`, `kPrefillIndexedNoSyncMaxT`. +- Dual-slot `EscTls`/`AccFastTls` unless a host test proves T>1 + device hop needs them. Prefer the existing single T=1 TLS plus owned T>1 buffer. +- #837 GetBlas, #839 Launch/Finish, #697. + +## Design + +Honest boundary: **routing only**. The T=19 peer smoke uses `RunGemma4Fp8TopKIndexedOnExpertDevice` in a `for t` loop — the lab `else` branch, not the rejected batched same-dev arm. + +`indexed_max_t` default 63 is a **product default change** for T=2..63 (warmups). It is not env-gated off. `=1` is the rollback to today's T==1-only gate. + +T≥64 must still miss this gate so prefill-batch is unchanged. + +## Risks + +- Numerics: per-token indexed vs serial M1 may not be bit-identical. Token-exact Paris/arith is **necessary and not sufficient** (research `64cb` stop-ship 4). +- T>1 TLS acc sized `{T,H}` must not reuse the T=1 hipGraph buffer. +- Do not lower `kPrefillBatchMinT`. + +## Tests + +### Host predicate (required) + +Extract a pure function or source+unit the env parse: + +| T | env | expect | +|---|---|---| +| 1 | unset | indexed | +| 19 | unset (63) | indexed | +| 19 | `=1` | **not** indexed (serial) | +| 64 | unset | **not** indexed (prefill-batch) | +| 63 | unset | indexed | +| 0 / top_k>8 / !fp8_res | unset | not indexed | + +RED: force the T==1 literal back → T=19 case fails. + +### Direct tensor oracle (required on impl — this is the correctness gate) + +Default-on numerical route change. Paris/arith can pass with hidden-state drift or indexing/ownership defects. Impl must ship a host-or-GPU oracle that compares **indexed output tensors** to **today's serial/reference math** for: + +- T ∈ {2, 19, 63} +- same-dev arm **and** peer arm (`compute_dev != expert_dev`) + +Oracle checks, all required: + +1. All finite (no NaN/Inf). +2. Zero support: every exact-zero in the reference is exact-zero in the candidate (no 1e-6 floor that hides a miss). +3. Declared tolerance: `max_abs(cand-ref) <= max(abs_tol, rel_tol * max_abs(ref))` with **abs_tol and rel_tol written in the impl commit** (propose `abs_tol = 2^-7 * max_abs(ref)` class unless a tighter bound is proven; do not invent a looser floor). +4. Route witness: a test-visible flag/counter that the indexed helper ran (not serial M1, not packed-batched). +5. Owned-buffer lifetime: the `[T,H]` output buffer is still owned by the caller after return (no TLS alias, no free-before-read). Probe by writing a canary after return and re-reading the tensor. + +RED mutations (must fail the oracle or the witness): + +- wrong token stride (index `t` as `t*H` vs `t`); +- wrong `[T,H]` ownership (return a T=1 TLS pointer, or free before caller reads); +- accidental packed-batched dispatch (`ExpertGeGLUFp8TopKIndexedBatched` / any T≥2 packed kernel). + +GPU (coord): T=19 generate succeeds; T=1 decode + Paris/arith unchanged. Generate is **not** a substitute for the tensor oracle. + +## Gates + +- Host predicate table GREEN without GPU. +- Tensor oracle GREEN for T=2,19,63 × {same-dev, peer} before any "A is GREEN" claim. +- Impl must not introduce `ExpertGeGLUFp8TopKIndexedBatched` or `PREFILL_INDEXED_NOSYNC`. +- Operator A/B (`5071`): this is **A**. Smoke **T=19 independently** first. p42k only after the smallest passing set. +- Default behavior outside Gemma-4 FP8 xdev `1<=T<64` eligibility is unchanged. +- `#697` files untouched. + +## Stop conditions + +- Research wants packed-batch after all — that is a new spec, not this one. +- Lab GPU smoke without coord. +- Landing this row on a shared branch with #837/#839. + +## Evidence + +Bus: `713f`, `a63e`, `25c9`, `5071`, `64cb`. Donor bytes hashed in `.agents/specs/rocm-gemma4-indexed-max-t-donor.md`. + +## Open on gfx1201 hardware + +Found while landing, and left open rather than papered over. Each needs the +RDNA4 pair this repository's maintainers do not have, so none of them can be +answered from a CPU host. The row does not reach `DONE` until they are. + +1. **`rw_idx` capture-stability — SOURCE FIXED, GPU witness still owed.** T=1 + scaled router weights now live in `RwIdxTls` (`thread_local`, keyed + `compute_dev` + `T*top_k`). Only T>1 `Release()`s a pooled `rw_idx_owned`. + Host source invariants cover the shape. Owed on gfx1201: one decode + hipGraph capture+replay witness that the baked `rw` pointer still matches + the live TLS address (idle-beside `:8012`, never `:8010`). +2. **No measurement — protocol frozen to Researcher ca41; GPU after static GREEN.** + Same `cff626f93`-derived binary (tree-identical to `9d4a16c5`), model, recipe, + GPU placement, prompts, context, sampling, and exact active batch T: + - A: `VT_GEMMA4_DECODE_INDEXED_MAX_T=63` + - B: `=1` (T>1 host-gather). Separate processes (env once-cached). + - T={2,8,63}; prove realized shape + selected arm in logs/counters. + - Correctness first: deterministic greedy per-request token IDs A vs B + must match before any timing is accepted. + - Isolated `:8012`; never `:8010`; idle window; alternate AB/BA by T. + - Exclude model load and first-request warmup. ≥3 warmups + ≥5 measured + steady repeats/arm. Median **and range** for decode tok/s and per-step + latency (not e2e load). + - Record binary SHA, HEAD, ROCm/compiler, env, GPU mapping, prompt / + context / output lengths, raw samples, thermal/clock, teardown. + - Do not call a noisy single run a default-flip win. +3. **The per-expert scale costs one kernel launch per token.** + `gemma4_moe.cpp:766-773` calls `vt::ApplyExpertScaleRw` inside the T-loop, up + to 63 launches. `ApplyExpertScaleRwKernel` + (`rocm_fp8_channel_gemv.hip:635`) is `<<<1, G>>>` and indexes `rw`/`ri` + linearly with `ri` holding global expert ids, so one call with + `G = T*top_k` is element-wise identical. `Gemma4IndexedOkT` bounds + `top_k <= 8` and `T <= 63`, so `G <= 504` fits one block today — but + collapsing the loop removes the structural reason `G` stays small, and + `ApplyExpertScaleRwRocm` does not check `G` against the 1024-thread block + limit or report a failed launch. The collapse therefore owes that guard. +4. Smaller, same hardware: `Gemma4IndexedHelperHits()` does a global atomic + `fetch_add` per token on the decode hot path purely so a host test can + observe it; `Gemma4IndexedScratchValidForT` cannot return false in the T>1 + branch that calls it; the new `RestoreComputeDev` + (`rocm_gemma4_experts.hip:572`) adds a second `hipSetDevice` on a success + path whose own comment says "no hipSetDevice between stream ops"; and + `RetireGemma4Fp8TopKIndexedPeer` returns `ok = true` while skipping the + expert-stream drain whenever `tls.edev != expert_dev`. +5. The six `Gemma4Indexed*` host helpers at `gemma4_indexed_gate.h:113-160` have + no production caller and belong under `tests/`. Their "tensor oracle" is + `s = (e+1) * rw[g]`, which models none of the FP8 GeGLU arithmetic, so it + gates the T-loop's striding and nothing about the kernel. That is a real + thing to gate; the header is the wrong place to keep it, and the name + oversells it. +6. A comment string is load-bearing: the source-invariant case asserts + `retire-before-acc_idx-dtor` appears in `gemma4_moe.cpp`. Reordering the + statements it names while keeping the comment passes, and deleting the + comment while keeping the order fails. The other invariants in that case + assert on code. diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index e5ce172d8..01f8ada06 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -279,6 +279,7 @@ on CUDA/CPU builds beyond the documented behavior. | `VT_GEMMA4_PREFILL_PEER_ACT` | on | Prefill MoE: run ExpertGeGLU on the expert GPU and peer activations only (not full weight PeerCopy). `0` restores weight PeerCopy | | `VT_GEMMA4_GPU0_HEADROOM_GB` | `12` | GiB kept free on GPU0 when packing resident experts (decode vs long-prefill trade). Lab dual R9700 + 49k KV: `8` survives 16k+ prefill; `6` OOMs ~11k | | `VT_GEMMA4_PREFILL_BATCH_MOE` | auto / `1` in lab recipe | `=1` group-by-expert prefill GEMM for `T>=64`; `=0` serial M=1 (slow). Unset = auto | +| `VT_GEMMA4_DECODE_INDEXED_MAX_T` | `63` | Widen device-indexed FP8 MoE from T=1 to `T<=N` (`N` clamped `[1,63]`). Unset = 63. `=1` restores the old T=1-only gate. T≥64 still uses prefill-batch | | `VT_GEMMA4_MLP_MOE_PARALLEL` | off | `=1` run Gemma4 MLP and MoE on two HIP streams (lab; wall ~flat on R9700). Not wired in this PR tip (decode-graph-free split) | | `VT_ATTN_PREFILL_FLASH` | off | `=1` SGLang-style BM×BN GQA flash prefill (lab A/B) | | `VT_GEMMA4_PREFILL_GEMM_M` | `2048` | Tokens per expert in prefill-batch GEMM chunks (`16..8192`; out-of-range values are ignored and the default is used). Larger M → fewer launches; lab `512` ~+37% prefill vs `64`, and `512`→`2048` ~+80 eng @11k vs the WMMA baseline (2026-08-10), which is why the default is `2048`. Lab KEEP on dual R9700 uses the default | diff --git a/docs/USAGE.md b/docs/USAGE.md index 708fbdbea..0b98f49fd 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -4792,6 +4792,14 @@ Prefill peer (#839) unpins dequant cache only after observed retirement; a faile This path does **not** restructure the Gemma-4 layer loop or enable decode hipGraph (those stay lab-only until a CUDA token-exact gate can land them). +Decode indexed MoE is bounded by `VT_GEMMA4_DECODE_INDEXED_MAX_T` (ENVIRONMENT.md). +Indexed T=2..63 retires peer and compute work before pooled scratch is released. +Unset defaults to 63: T=2..63 uses the existing per-token indexed helpers on a +scratch-scaled weight copy, and helper failure retires peer/compute streams before +pooled scratch returns, then falls back with a single host scale. `=1` restores +T=1-only. This does not restructure the layer loop or enable decode hipGraph +either. + Contributor KEEP recipe (2x R9700 gfx1201, ROCm 7.2.4, `PREFIX_CACHE=0`, unique pads, 2026-08-13): SharedK-WMMA on, FLASH/FMHA off, `VT_GEMMA4_PREFILL_GEMM_M=2048` (the default), `VT_GEMMA4_PREFILL_PEER_ACT=1` (the default), batch MoE `T>=64`. diff --git a/include/vllm/model_executor/models/gemma4_indexed_gate.h b/include/vllm/model_executor/models/gemma4_indexed_gate.h new file mode 100644 index 000000000..00097f9d2 --- /dev/null +++ b/include/vllm/model_executor/models/gemma4_indexed_gate.h @@ -0,0 +1,232 @@ +// #838: widen Gemma-4 FP8 indexed MoE from T==1 to T<=63. +// Host-injectable dispatch, tensor oracle, single-scale fallback, scratch retire. +// No HIP. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace vllm { + +constexpr int64_t kGemma4PrefillBatchMinT = 64; +constexpr int64_t kGemma4DecodeIndexedMaxTDefault = 63; +constexpr int64_t kGemma4DecodeIndexedMaxTLo = 1; +constexpr int64_t kGemma4DecodeIndexedMaxTHi = 63; +constexpr float kGemma4IndexedOraclePow2 = 7.0f; // abs_tol = 2^-7 * max_abs(ref) + +inline int64_t ParseGemma4DecodeIndexedMaxT(const char* e) { + if (e == nullptr || e[0] == '\0') return kGemma4DecodeIndexedMaxTDefault; + char* end = nullptr; + const long v = std::strtol(e, &end, 10); + if (end == e) return kGemma4DecodeIndexedMaxTDefault; + if (v < kGemma4DecodeIndexedMaxTLo) return kGemma4DecodeIndexedMaxTLo; + if (v > kGemma4DecodeIndexedMaxTHi) return kGemma4DecodeIndexedMaxTHi; + return static_cast(v); +} + +inline int64_t Gemma4DecodeIndexedMaxT() { + static const int64_t n = ParseGemma4DecodeIndexedMaxT(std::getenv("VT_GEMMA4_DECODE_INDEXED_MAX_T")); + return n; +} + +inline bool Gemma4IndexedOkT(int64_t T, int64_t indexed_max_t, int top_k, bool fp8_res) { + return T >= 1 && T <= indexed_max_t && T < kGemma4PrefillBatchMinT && fp8_res && + top_k <= 8 && top_k > 0; +} + +struct Gemma4IndexedTokenOff { + int64_t x_elems = 0; + int64_t y_elems = 0; + int64_t route = 0; +}; + +inline Gemma4IndexedTokenOff Gemma4IndexedTokenOffsets(int64_t t, int64_t H, int top_k) { + return Gemma4IndexedTokenOff{t * H, t * H, t * static_cast(top_k)}; +} + +inline std::atomic& Gemma4IndexedHelperHits() { + static std::atomic n{0}; + return n; +} + +template +struct Gemma4IndexedCall { + int64_t t = 0; + Gemma4IndexedTokenOff off{}; + bool peer = false; + const XT* x = nullptr; + YT* y = nullptr; + const int32_t* ri = nullptr; + const float* rw = nullptr; +}; + +struct Gemma4IndexedDispatchResult { + bool ok = false; + uint64_t hits = 0; + int restores = 0; + void* y_owner = nullptr; + bool enqueued = false; +}; + +// Production T-loop. `fn(call)` is the same-dev or peer helper. `restore()` runs +// after every helper return (success or fail). Does not free `y`. +template +inline Gemma4IndexedDispatchResult Gemma4IndexedDispatchTokens( + int64_t T, int64_t H, int top_k, bool peer, YT* y, const XT* x, const int32_t* ri, + const float* rw, Fn fn, Restore restore) { + Gemma4IndexedDispatchResult r; + r.y_owner = y; + if (T <= 0 || !y || !x || !ri || !rw) return r; + for (int64_t t = 0; t < T; ++t) { + const auto off = Gemma4IndexedTokenOffsets(t, H, top_k); + const Gemma4IndexedCall c{t, off, peer, x + off.x_elems, y + off.y_elems, + ri + off.route, rw + off.route}; + const bool one = fn(c); + r.enqueued = true; + restore(); + ++r.restores; + if (!one) { + r.ok = false; + return r; + } + ++r.hits; + Gemma4IndexedHelperHits().fetch_add(1, std::memory_order_relaxed); + } + r.ok = true; + return r; +} + +// Serial fallback scale. already_scaled=true skips (indexed must not leave rw mutated). +inline void Gemma4ApplyHostExpertScaleOnce(float* hw, const int32_t* hi, const float* hscale, + int64_t T, int top_k, int64_t E, bool already_scaled) { + if (already_scaled || !hw || !hi || !hscale || T <= 0 || top_k <= 0) return; + for (int64_t t = 0; t < T; ++t) { + for (int i = 0; i < top_k; ++i) { + const size_t o = static_cast(t * top_k + i); + const int e = hi[o]; + if (e >= 0 && e < static_cast(E)) hw[o] *= hscale[static_cast(e)]; + } + } +} + +// Independent serial / product-loop token math (host oracle). +inline void Gemma4IndexedHostApplyToken(float* y, const float* x, const int32_t* ri, const float* rw, + int64_t H, int top_k) { + for (int64_t h = 0; h < H; ++h) y[h] = 0.f; + for (int g = 0; g < top_k; ++g) { + const int32_t e = ri[g]; + const float s = static_cast(e + 1) * rw[g]; + for (int64_t h = 0; h < H; ++h) y[h] += x[h] * s; + } +} + +// Independent of Gemma4IndexedHostApplyToken — do not call it here. +inline void Gemma4IndexedHostSerialRef(const float* x, const int32_t* idx, const float* wts, + float* y, int64_t T, int64_t H, int top_k) { + for (int64_t t = 0; t < T; ++t) { + float* yt = y + t * H; + const float* xt = x + t * H; + const int32_t* idt = idx + t * top_k; + const float* wt = wts + t * top_k; + for (int64_t h = 0; h < H; ++h) yt[h] = 0.f; + for (int g = 0; g < top_k; ++g) { + const int32_t e = idt[g]; + const float s = static_cast(e + 1) * wt[g]; + for (int64_t h = 0; h < H; ++h) yt[h] += xt[h] * s; + } + } +} + +inline bool Gemma4IndexedOracleClose(const float* cand, const float* ref, int64_t n, + float* max_abs_out) { + float max_abs_ref = 0.f; + float max_abs_diff = 0.f; + for (int64_t i = 0; i < n; ++i) { + if (!std::isfinite(cand[i]) || !std::isfinite(ref[i])) return false; + max_abs_ref = std::max(max_abs_ref, std::fabs(ref[i])); + max_abs_diff = std::max(max_abs_diff, std::fabs(cand[i] - ref[i])); + if (ref[i] == 0.f && cand[i] != 0.f) return false; + } + if (max_abs_out) *max_abs_out = max_abs_diff; + const float tol = std::ldexp(max_abs_ref, -static_cast(kGemma4IndexedOraclePow2)); + return max_abs_diff <= tol; +} + +enum class Gemma4IndexedArm { SameDev, Peer, None }; + +inline Gemma4IndexedArm Gemma4IndexedSelectArm(bool fp8_res_same, bool fp8_res_peer) { + if (fp8_res_same) return Gemma4IndexedArm::SameDev; + if (fp8_res_peer) return Gemma4IndexedArm::Peer; + return Gemma4IndexedArm::None; +} + +template +inline bool Gemma4IndexedRunSelectedArm(Gemma4IndexedArm arm, SameFn same, PeerFn peer) { + if (arm == Gemma4IndexedArm::SameDev) return same(); + if (arm == Gemma4IndexedArm::Peer) return peer(); + return false; +} + +struct Gemma4IndexedHelperArgs { + void* y = nullptr; + const void* x = nullptr; + const int32_t* ri = nullptr; + const float* rw = nullptr; +}; + +inline Gemma4IndexedHelperArgs Gemma4IndexedPackArgs(void* y, const void* x, const int32_t* ri, + const float* rw) { + return Gemma4IndexedHelperArgs{y, x, ri, rw}; +} + +inline bool Gemma4IndexedArgsEq(const Gemma4IndexedHelperArgs& a, const Gemma4IndexedHelperArgs& b) { + return a.y == b.y && a.x == b.x && a.ri == b.ri && a.rw == b.rw; +} + +enum class Gemma4IndexedScratchKind { TlsT1, OwnedTH }; + +inline Gemma4IndexedScratchKind Gemma4IndexedScratchKindFor(int64_t T) { + return T == 1 ? Gemma4IndexedScratchKind::TlsT1 : Gemma4IndexedScratchKind::OwnedTH; +} + +struct Gemma4IndexedScratchChoice { + Gemma4IndexedScratchKind kind = Gemma4IndexedScratchKind::OwnedTH; + void* y = nullptr; + int64_t elems = 0; +}; + +inline bool Gemma4IndexedScratchValidForT(const Gemma4IndexedScratchChoice& c, int64_t T, int64_t H) { + if (T <= 0 || H <= 0 || c.y == nullptr) return false; + if (T == 1) return c.kind == Gemma4IndexedScratchKind::TlsT1 && c.elems >= H; + return c.kind == Gemma4IndexedScratchKind::OwnedTH && c.elems >= T * H; +} + +// Release to pool is illegal unless retirement was observed. +inline bool Gemma4IndexedMayReleaseToPool(bool enqueued, bool retire_observed) { + return !enqueued || retire_observed; +} + +// Host model of the production fail path: retire while buffer is still owned, +// then release only if retirement was observed. Release-before-retire is RED. +template +inline bool Gemma4IndexedFailPathRetireThenMaybeRelease(bool enqueued, bool& owned, bool& released, + bool& retire_ok, Retire retire) { + if (!owned) return false; + retire_ok = true; + if (enqueued) retire_ok = retire(); + if (!Gemma4IndexedMayReleaseToPool(enqueued, retire_ok)) { + owned = true; // quarantine: keep ownership, do not release + released = false; + return false; + } + owned = false; + released = true; + return true; +} + +} // namespace vllm diff --git a/include/vllm/model_executor/models/gemma4_moe.h b/include/vllm/model_executor/models/gemma4_moe.h index 7cb7e1728..c24c5e431 100644 --- a/include/vllm/model_executor/models/gemma4_moe.h +++ b/include/vllm/model_executor/models/gemma4_moe.h @@ -124,6 +124,8 @@ bool RunGemma4Fp8TopKIndexedOnExpertDevice(vt::Queue& compute_q, int expert_dev, const void* dn_base, const void* sgu_base, const void* sdn_base, const int32_t* idx_compute, const float* wts_compute, int G, int I, int H); +// Drain compute + indexed peer streams before pooled scratch may return to DevicePool. +bool RetireGemma4Fp8TopKIndexedPeer(vt::Queue& compute_q, int expert_dev); // hipHostRegister BF16 expert cache for faster H2D (no-op if already pinned). void PinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats& ex); diff --git a/src/vllm/model_executor/models/gemma4_moe.cpp b/src/vllm/model_executor/models/gemma4_moe.cpp index 31c711cb9..7ff6b3ed0 100644 --- a/src/vllm/model_executor/models/gemma4_moe.cpp +++ b/src/vllm/model_executor/models/gemma4_moe.cpp @@ -1,5 +1,6 @@ // Gemma-4 MoE: BF16 fused or FP8 per-expert + optional device resident. #include "vllm/model_executor/models/gemma4_moe.h" +#include "vllm/model_executor/models/gemma4_indexed_gate.h" #include #include @@ -757,9 +758,22 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, const bool fp8_res_same = fp8_res && ex.dev_id == compute_dev; const bool fp8_res_peer = fp8_res && ex.dev_id != compute_dev; - // Decode T=1: fully device MoE — no router D2H, no host top-k gather. - if (T == 1 && fp8_res && top_k <= 8 && top_k > 0) { - // per-expert scale on device (once per layer/E). + // Device-indexed FP8 MoE for T=1..min(MAX, batch_min-1). No router D2H. + // T=1: hipGraph-stable TLS acc. T>1: owned [T,H] + per-token existing helpers. + // VT_GEMMA4_DECODE_INDEXED_MAX_T: default 63; =1 → T=1 only; clamp [1,63]. + const int64_t indexed_max_t = Gemma4DecodeIndexedMaxT(); + const bool indexed_eligible = Gemma4IndexedOkT(T, indexed_max_t, top_k, fp8_res); + auto emit_moe_dispatch = [&](const char* path, bool fallthrough) { + if (!profile) return; + std::fprintf(stderr, + "gemma4 moe dispatch: T=%lld indexed_max_t=%lld path=%s eligible=%d top_k=%d " + "compute_dev=%d expert_dev=%d fallthrough=%d\n", + static_cast(T), static_cast(indexed_max_t), path, + indexed_eligible ? 1 : 0, top_k, compute_dev, ex.dev_id, fallthrough ? 1 : 0); + std::fflush(stderr); + }; + if (Gemma4IndexedOkT(T, indexed_max_t, top_k, fp8_res)) { + // per-expert scale on device (once per layer/E; apply each token — helper is G-wide). struct EscTls { int dev = -1; int64_t E = 0; @@ -780,72 +794,165 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, esc.host_key = hk; } escale_ptr = static_cast(esc.sc->ptr()); - vt::ApplyExpertScaleRw(d.q, static_cast(rw.ptr()), static_cast(ri.ptr()), - escale_ptr, top_k, static_cast(E)); } - // Stable T=1 acc for hipGraph (do not pool-Release). - struct AccFastTls { + // Never mutate router `rw` in place — fallback must still see unscaled weights. + // T=1: TLS-stable copy (hipGraph bakes the pointer). T>1: per-call pooled DBuf. + struct RwIdxTls { int dev = -1; - int64_t H = 0; - std::optional acc; + int64_t n = 0; // T*top_k + std::optional buf; }; - static thread_local AccFastTls aft; - if (aft.dev != d.q.device.index || aft.H != H || !aft.acc) { - aft.acc.emplace(d, DType::kBF16, std::vector{T, H}); - aft.dev = d.q.device.index; - aft.H = H; - } - DBuf& acc_fast = *aft.acc; - // expert_in is already [1,H] bf16 on device - const void* xin = expert_in.data; - void* yout = acc_fast.ptr(); - bool ok = false; - if (fp8_res_same) { - ok = vt::ExpertGeGLUFp8TopKIndexed( - d.q, yout, xin, ex.fp8_gu_base, ex.fp8_dn_base, ex.fp8_sgu_base, ex.fp8_sdn_base, - static_cast(ri.ptr()), static_cast(rw.ptr()), top_k, - static_cast(I), static_cast(H)); - } else if (fp8_res_peer) { - ok = RunGemma4Fp8TopKIndexedOnExpertDevice( - d.q, ex.dev_id, yout, xin, ex.fp8_gu_base, ex.fp8_dn_base, ex.fp8_sgu_base, - ex.fp8_sdn_base, static_cast(ri.ptr()), - static_cast(rw.ptr()), top_k, static_cast(I), static_cast(H)); + static thread_local RwIdxTls rwt; + std::optional rw_idx_owned; + const float* helper_rw = static_cast(rw.ptr()); + if (escale_ptr) { + const int64_t n = T * top_k; + const bool t1_tls = Gemma4IndexedScratchKindFor(T) == Gemma4IndexedScratchKind::TlsT1; + DBuf* scaled = nullptr; + if (t1_tls) { + if (rwt.dev != compute_dev || rwt.n != n || !rwt.buf) { + rwt.buf.emplace(d, DType::kF32, std::vector{T, top_k}); + rwt.dev = compute_dev; + rwt.n = n; + } + scaled = &*rwt.buf; + } else { + rw_idx_owned.emplace(d, DType::kF32, std::vector{T, top_k}); + scaled = &*rw_idx_owned; + } + d.b.Copy(d.q, scaled->ptr(), rw.ptr(), static_cast(n) * sizeof(float)); + for (int64_t t = 0; t < T; ++t) { + const auto off = Gemma4IndexedTokenOffsets(t, H, top_k); + vt::ApplyExpertScaleRw(d.q, static_cast(scaled->ptr()) + off.route, + static_cast(ri.ptr()) + off.route, escale_ptr, top_k, + static_cast(E)); + } + helper_rw = static_cast(scaled->ptr()); } - if (ok) { - const auto t_router1 = profile ? clock::now() : clock::time_point{}; - Gemma4MoeScratch r; - r.tensor = acc_fast.t(); - r.storage = std::shared_ptr(acc_fast.ptr(), [](void*) {}); - if (profile) { + + const auto indexed_arm = Gemma4IndexedSelectArm(fp8_res_same, fp8_res_peer); + auto run_one = [&](void* yout, const void* xin, const int32_t* ri_t, const float* rw_t) -> bool { + const auto args = Gemma4IndexedPackArgs(yout, xin, ri_t, rw_t); + return Gemma4IndexedRunSelectedArm( + indexed_arm, + [&] { + return vt::ExpertGeGLUFp8TopKIndexed(d.q, args.y, args.x, ex.fp8_gu_base, ex.fp8_dn_base, + ex.fp8_sgu_base, ex.fp8_sdn_base, args.ri, args.rw, + top_k, static_cast(I), static_cast(H)); + }, + [&] { + return RunGemma4Fp8TopKIndexedOnExpertDevice( + d.q, ex.dev_id, args.y, args.x, ex.fp8_gu_base, ex.fp8_dn_base, ex.fp8_sgu_base, + ex.fp8_sdn_base, args.ri, args.rw, top_k, static_cast(I), static_cast(H)); + }); + }; + auto restore_compute = [] {}; + bool indexed_retired = false; + auto retire_indexed = [&]() -> bool { + if (indexed_retired) return true; + bool ok = true; + if (fp8_res_peer) { + ok = RetireGemma4Fp8TopKIndexedPeer(d.q, ex.dev_id); + } else { d.b.Synchronize(d.q); - const auto t_all1 = clock::now(); - static std::atomic ncalls{0}; - static std::atomic us_router{0}; - static std::atomic us_total{0}; - const auto ur = - std::chrono::duration_cast(t_router1 - t_all0).count(); - const auto ut = - std::chrono::duration_cast(t_all1 - t_all0).count(); - us_router.fetch_add(static_cast(ur), std::memory_order_relaxed); - us_total.fetch_add(static_cast(ut), std::memory_order_relaxed); - const uint64_t c = ncalls.fetch_add(1, std::memory_order_relaxed) + 1; - if (c == 1 || c % 64 == 0) { - const uint64_t tr = us_router.load(std::memory_order_relaxed); - const uint64_t tt = us_total.load(std::memory_order_relaxed); - std::fprintf(stderr, - "gemma4 moe profile: calls=%llu router_us/call=%.1f expert+rest_us/call=%.1f " - "total_us/call=%.1f (router%%=%.0f) [indexed-device]\n", - static_cast(c), static_cast(tr) / c, - static_cast(tt - tr) / c, static_cast(tt) / c, - tt ? 100.0 * static_cast(tr) / static_cast(tt) : 0.0); + } + indexed_retired = ok; + return ok; + }; + + if (Gemma4IndexedScratchKindFor(T) == Gemma4IndexedScratchKind::TlsT1) { + // Stable T=1 acc for hipGraph (do not pool-Release). rw_idx is RwIdxTls. + struct AccFastTls { + int dev = -1; + int64_t H = 0; + std::optional acc; + }; + static thread_local AccFastTls aft; + if (aft.dev != d.q.device.index || aft.H != H || !aft.acc) { + aft.acc.emplace(d, DType::kBF16, std::vector{1, H}); + aft.dev = d.q.device.index; + aft.H = H; + } + DBuf& acc_fast = *aft.acc; + const auto disp = Gemma4IndexedDispatchTokens( + 1, H, top_k, fp8_res_peer, static_cast(acc_fast.ptr()), + static_cast(expert_in.data), static_cast(ri.ptr()), + helper_rw, [&](const Gemma4IndexedCall& c) { + return run_one(c.y, c.x, c.ri, c.rw); + }, + restore_compute); + if (disp.ok) { + emit_moe_dispatch(indexed_arm == Gemma4IndexedArm::Peer ? "indexed_peer" : "indexed_same", + /*fallthrough=*/false); + const auto t_router1 = profile ? clock::now() : clock::time_point{}; + Gemma4MoeScratch r; + r.tensor = acc_fast.t(); + r.storage = std::shared_ptr(acc_fast.ptr(), [](void*) {}); + if (profile) { + d.b.Synchronize(d.q); + const auto t_all1 = clock::now(); + static std::atomic ncalls{0}; + static std::atomic us_router{0}; + static std::atomic us_total{0}; + const auto ur = + std::chrono::duration_cast(t_router1 - t_all0).count(); + const auto ut = + std::chrono::duration_cast(t_all1 - t_all0).count(); + us_router.fetch_add(static_cast(ur), std::memory_order_relaxed); + us_total.fetch_add(static_cast(ut), std::memory_order_relaxed); + const uint64_t c = ncalls.fetch_add(1, std::memory_order_relaxed) + 1; + if (c == 1 || c % 64 == 0) { + const uint64_t tr = us_router.load(std::memory_order_relaxed); + const uint64_t tt = us_total.load(std::memory_order_relaxed); + std::fprintf(stderr, + "gemma4 moe profile: calls=%llu router_us/call=%.1f expert+rest_us/call=%.1f " + "total_us/call=%.1f (router%%=%.0f) [indexed-device]\n", + static_cast(c), static_cast(tr) / c, + static_cast(tt - tr) / c, static_cast(tt) / c, + tt ? 100.0 * static_cast(tr) / static_cast(tt) : 0.0); + } } + return r; + } + (void)retire_indexed(); // T=1 TLS acc/rw_idx are not pooled; retire before leaving arm + } else { + std::optional acc_idx; + acc_idx.emplace(d, DType::kBF16, std::vector{T, H}); + auto* x_base = static_cast(expert_in.data); + auto* y_base = static_cast(acc_idx->ptr()); + auto* ri_base = static_cast(ri.ptr()); + const Gemma4IndexedScratchChoice scratch_choice{Gemma4IndexedScratchKindFor(T), y_base, + T * H}; + if (!Gemma4IndexedScratchValidForT(scratch_choice, T, H)) { + if (!retire_indexed()) (void)acc_idx->Release(); // quarantine + // fall through with acc_idx still in scope until this block ends + } else { + const auto disp = Gemma4IndexedDispatchTokens( + T, H, top_k, fp8_res_peer, y_base, x_base, ri_base, helper_rw, + [&](const Gemma4IndexedCall& c) { + return run_one(c.y, c.x, c.ri, c.rw); + }, + restore_compute); + if (disp.ok) { + emit_moe_dispatch(indexed_arm == Gemma4IndexedArm::Peer ? "indexed_peer" : "indexed_same", + /*fallthrough=*/false); + Gemma4MoeScratch r; + r.tensor = acc_idx->t(); + r.storage = acc_idx->ReleaseShared(); + return r; + } + // retire-before-acc_idx-dtor: still lexically inside acc_idx scope + if (!retire_indexed()) (void)acc_idx->Release(); // quarantine, do not Put } - return r; } + if (!indexed_retired) (void)retire_indexed(); + if (!indexed_retired && rw_idx_owned) (void)rw_idx_owned->Release(); // fall through to legacy host-gather path } + emit_moe_dispatch("legacy", /*fallthrough=*/indexed_eligible); + std::vector hw(static_cast(T * top_k)); std::vector hi(static_cast(T * top_k)); d.b.Copy(d.q, hw.data(), rw.ptr(), hw.size() * sizeof(float)); @@ -859,14 +966,9 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, const auto* pe = reinterpret_cast(moe.per_expert_scale.bytes.data()); for (int64_t e = 0; e < E; ++e) hscale[static_cast(e)] = vt::BF16ToF32(pe[e]); } - // Apply per-expert scale to selected weights. - for (int64_t t = 0; t < T; ++t) { - for (int i = 0; i < top_k; ++i) { - const size_t o = static_cast(t * top_k + i); - const int e = hi[o]; - if (e >= 0 && e < static_cast(E)) hw[o] *= hscale[static_cast(e)]; - } - } + // Apply per-expert scale to selected weights (once; indexed scratch is a copy). + Gemma4ApplyHostExpertScaleOnce(hw.data(), hi.data(), hscale.data(), T, top_k, E, + /*already_scaled=*/false); const bool need_peer_sc = (!same_dev && ex.gate_up_dev && ex.down_dev && ex.dev_id >= 0) || fp8_res_peer; @@ -1007,7 +1109,7 @@ Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe, // Prefer fused M=1 FP8 for short T (decode + tiny prefills). Batch dequant+GEMM // only pays off once enough tokens share experts (lab: T=818 ~3×, T=6k ~6×; // T=13 was slower than fused M=1). - constexpr int64_t kPrefillBatchMinT = 64; + constexpr int64_t kPrefillBatchMinT = kGemma4PrefillBatchMinT; const bool prefill_batch_moe = (T >= kPrefillBatchMinT) && !host_axpy && ((prefill_batch_env == 1) || @@ -1645,6 +1747,7 @@ bool RunGemma4Fp8TopKIndexedOnExpertDevice(vt::Queue&, int, void*, const void*, const float*, int, int, int) { return false; } +bool RetireGemma4Fp8TopKIndexedPeer(vt::Queue&, int) { return true; } void PinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats&) {} void UnpinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats&) {} #endif // VLLM_CPP_HIP diff --git a/src/vt/rocm/rocm_gemma4_experts.hip b/src/vt/rocm/rocm_gemma4_experts.hip index 7d0b51ce9..27cd8b0fc 100644 --- a/src/vt/rocm/rocm_gemma4_experts.hip +++ b/src/vt/rocm/rocm_gemma4_experts.hip @@ -542,6 +542,20 @@ bool RunGemma4Fp8TopKOnExpertDevice(vt::Queue& compute_q, int expert_dev, void* return true; } +namespace { +struct IndexedPeerTls { + int cdev = -1, edev = -1, H = 0, Gcap = 0; + vt::Queue eq{}; + void* x = nullptr; + void* y = nullptr; + int32_t* idx = nullptr; + float* wts = nullptr; + hipEvent_t ev_c = nullptr, ev_e = nullptr; + bool eq_live = false; +}; +thread_local IndexedPeerTls g_idx_peer_tls; +} // namespace + // Decode T=1 indexed: bases stay on expert_dev; peer only x + tiny idx/wts + y. bool RunGemma4Fp8TopKIndexedOnExpertDevice(vt::Queue& compute_q, int expert_dev, void* ysum_compute, const void* x_compute, const void* gu_base, @@ -558,17 +572,15 @@ bool RunGemma4Fp8TopKIndexedOnExpertDevice(vt::Queue& compute_q, int expert_dev, sgu_base, sdn_base, idx_compute, wts_compute, G, I, H); } - struct Tls { - int cdev = -1, edev = -1, H = 0, Gcap = 0; - vt::Queue eq{}; - void* x = nullptr; - void* y = nullptr; - int32_t* idx = nullptr; - float* wts = nullptr; - hipEvent_t ev_c = nullptr, ev_e = nullptr; - bool eq_live = false; + struct RestoreComputeDev { + int dev; + explicit RestoreComputeDev(int d) : dev(d) {} + ~RestoreComputeDev() { + if (dev >= 0) (void)hipSetDevice(dev); + } }; - static thread_local Tls tls; + RestoreComputeDev restore(compute_dev); + auto& tls = g_idx_peer_tls; const size_t hb = static_cast(H) * 2; hipStream_t cst = static_cast(compute_q.handle); vt::Device edev{vt::DeviceType::kROCM, expert_dev}; @@ -622,30 +634,64 @@ bool RunGemma4Fp8TopKIndexedOnExpertDevice(vt::Queue& compute_q, int expert_dev, } hipStream_t est = static_cast(tls.eq.handle); - // Steady-state: no hipSetDevice between stream ops (graph-prep + less driver tax). - if (hipEventRecord(tls.ev_c, cst) != hipSuccess) return false; - if (hipStreamWaitEvent(est, tls.ev_c, 0) != hipSuccess) return false; + auto retire_fail = [&]() -> bool { + if (est) { + (void)hipSetDevice(expert_dev); + (void)hipStreamSynchronize(est); + } + if (cst) { + (void)hipSetDevice(compute_dev); + (void)hipStreamSynchronize(cst); + } + (void)hipSetDevice(compute_dev); + return false; + }; + // Current-device ownership is explicit. HIP stream ops and peer copies run + // under the stream's owning device. RestoreComputeDev restores compute on + // every exit. The prior "no hipSetDevice" claim was false: init left + // current=expert and success set compute, so only the first peer call saw + // the expert current-device. + if (hipSetDevice(compute_dev) != hipSuccess) return retire_fail(); + if (hipEventRecord(tls.ev_c, cst) != hipSuccess) return retire_fail(); + if (hipSetDevice(expert_dev) != hipSuccess) return retire_fail(); + if (hipStreamWaitEvent(est, tls.ev_c, 0) != hipSuccess) return retire_fail(); auto peer = [&](void* dst, const void* src, size_t n) -> bool { return hipMemcpyPeerAsync(dst, expert_dev, const_cast(src), compute_dev, n, est) == hipSuccess; }; - if (!peer(tls.x, x_compute, hb)) return false; - if (!peer(tls.idx, idx_compute, sizeof(int32_t) * static_cast(G))) return false; - if (!peer(tls.wts, wts_compute, sizeof(float) * static_cast(G))) return false; + if (!peer(tls.x, x_compute, hb)) return retire_fail(); + if (!peer(tls.idx, idx_compute, sizeof(int32_t) * static_cast(G))) return retire_fail(); + if (!peer(tls.wts, wts_compute, sizeof(float) * static_cast(G))) return retire_fail(); if (!vt::ExpertGeGLUFp8TopKIndexed(tls.eq, tls.y, tls.x, gu_base, dn_base, sgu_base, sdn_base, tls.idx, tls.wts, G, I, H)) - return false; + return retire_fail(); - if (hipEventRecord(tls.ev_e, est) != hipSuccess) return false; - if (hipStreamWaitEvent(cst, tls.ev_e, 0) != hipSuccess) return false; + if (hipEventRecord(tls.ev_e, est) != hipSuccess) return retire_fail(); + if (hipSetDevice(compute_dev) != hipSuccess) return retire_fail(); + if (hipStreamWaitEvent(cst, tls.ev_e, 0) != hipSuccess) return retire_fail(); if (hipMemcpyPeerAsync(ysum_compute, compute_dev, tls.y, expert_dev, hb, cst) != hipSuccess) - return false; - if (hipSetDevice(compute_dev) != hipSuccess) return false; + return retire_fail(); return true; } +bool RetireGemma4Fp8TopKIndexedPeer(vt::Queue& compute_q, int expert_dev) { + const int compute_dev = compute_q.device.index; + hipStream_t cst = static_cast(compute_q.handle); + bool ok = true; + if (compute_dev >= 0 && hipSetDevice(compute_dev) != hipSuccess) ok = false; + if (cst && hipStreamSynchronize(cst) != hipSuccess) ok = false; + auto& tls = g_idx_peer_tls; + if (tls.eq_live && tls.cdev == compute_dev && (expert_dev < 0 || tls.edev == expert_dev)) { + hipStream_t est = static_cast(tls.eq.handle); + if (tls.edev >= 0 && hipSetDevice(tls.edev) != hipSuccess) ok = false; + if (est && hipStreamSynchronize(est) != hipSuccess) ok = false; + } + if (compute_dev >= 0 && hipSetDevice(compute_dev) != hipSuccess) ok = false; + return ok; +} + // Prefill: run ExpertGeGLU for M rows on expert_dev (weights resident there). // Peer only activations M×H bf16 — not multi-GB FP8 expert packs. // #839: wrapper → Launch/Finish + pin until host-observed retirement. diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2ad30f124..0d4390049 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -598,6 +598,9 @@ vllm_cpp_add_test(test_gemma4_honesty vllm/models/test_gemma4_honesty.cpp) # case that wants a different one. vllm_cpp_add_test(test_gemma4_moe_device_arm_guard vllm/models/test_gemma4_moe_device_arm_guard.cpp) +vllm_cpp_add_test(test_gemma4_indexed_max_t vllm/models/test_gemma4_indexed_max_t.cpp) +target_compile_definitions(test_gemma4_indexed_max_t PRIVATE + VLLM_CPP_SOURCE_DIR="${CMAKE_SOURCE_DIR}") vllm_cpp_add_test(test_qwen3_moe_load vllm/models/test_qwen3_moe_load.cpp) vllm_cpp_add_test(test_qwen3_moe_forward vllm/models/test_qwen3_moe_forward.cpp) target_include_directories(test_qwen3_moe_forward PRIVATE ${CMAKE_SOURCE_DIR}/src) diff --git a/tests/vllm/models/test_gemma4_indexed_max_t.cpp b/tests/vllm/models/test_gemma4_indexed_max_t.cpp new file mode 100644 index 000000000..978ef3a16 --- /dev/null +++ b/tests/vllm/models/test_gemma4_indexed_max_t.cpp @@ -0,0 +1,361 @@ +// #838 host gates: product-loop tensor oracle, single-scale, retire-before-pool. +#include + +#include +#include +#include +#include +#include + +#include "vllm/model_executor/models/gemma4_indexed_gate.h" + +#ifndef VLLM_CPP_SOURCE_DIR +#define VLLM_CPP_SOURCE_DIR "." +#endif + +using vllm::Gemma4ApplyHostExpertScaleOnce; +using vllm::Gemma4IndexedCall; +using vllm::Gemma4IndexedDispatchTokens; +using vllm::Gemma4IndexedHelperHits; +using vllm::Gemma4IndexedHostApplyToken; +using vllm::Gemma4IndexedHostSerialRef; +using vllm::Gemma4IndexedOkT; +using vllm::Gemma4IndexedArgsEq; +using vllm::Gemma4IndexedArm; +using vllm::Gemma4IndexedFailPathRetireThenMaybeRelease; +using vllm::Gemma4IndexedHelperArgs; +using vllm::Gemma4IndexedMayReleaseToPool; +using vllm::Gemma4IndexedOracleClose; +using vllm::Gemma4IndexedPackArgs; +using vllm::Gemma4IndexedRunSelectedArm; +using vllm::Gemma4IndexedScratchChoice; +using vllm::Gemma4IndexedScratchKind; +using vllm::Gemma4IndexedScratchKindFor; +using vllm::Gemma4IndexedScratchValidForT; +using vllm::Gemma4IndexedSelectArm; +using vllm::Gemma4IndexedTokenOffsets; +using vllm::ParseGemma4DecodeIndexedMaxT; +using vllm::kGemma4PrefillBatchMinT; + +namespace { + +std::string ReadText(const char* rel) { + const std::string path = std::string(VLLM_CPP_SOURCE_DIR) + "/" + rel; + std::ifstream in(path); + std::ostringstream ss; + ss << in.rdbuf(); + return ss.str(); +} + +// Slice ONE function definition, signature through its brace-matched closing +// brace. The obvious alternative -- slice from this symbol to the next one -- +// is not a slice of this function at all; it is a slice of the GAP, so any +// unrelated insertion after the closing brace lands inside it and reds a source +// invariant that the function still satisfies. Skips a forward declaration by +// requiring the first `{` to precede the first `;`. Braces inside string or +// character literals would fool the matcher; the functions asserted on here +// contain none, and a reviewer adding one must re-check this helper. +std::string FunctionBody(const std::string& src, const std::string& signature) { + for (auto sig = src.find(signature); sig != std::string::npos; + sig = src.find(signature, sig + 1)) { + const auto open = src.find('{', sig); + if (open == std::string::npos) return {}; + const auto semi = src.find(';', sig); + if (semi != std::string::npos && semi < open) continue; // declaration + int depth = 0; + for (size_t i = open; i < src.size(); ++i) { + if (src[i] == '{') { + ++depth; + } else if (src[i] == '}' && --depth == 0) { + return src.substr(sig, i - sig + 1); + } + } + return {}; + } + return {}; +} + +struct WritingHelper { + bool peer_expected = false; + int fail_at = -1; + bool fail_after_enqueue = false; + int calls = 0; + int restores = 0; + int64_t H = 0; + int top_k = 0; + + bool operator()(const Gemma4IndexedCall& c) { + ++calls; + REQUIRE(c.peer == peer_expected); + const auto want = Gemma4IndexedTokenOffsets(c.t, H, top_k); + REQUIRE(c.off.x_elems == want.x_elems); + REQUIRE(c.off.y_elems == want.y_elems); + REQUIRE(c.off.route == want.route); + if (fail_at >= 0 && c.t == fail_at) { + if (fail_after_enqueue) { + Gemma4IndexedHostApplyToken(c.y, c.x, c.ri, c.rw, H, top_k); + } + return false; + } + Gemma4IndexedHostApplyToken(c.y, c.x, c.ri, c.rw, H, top_k); + return true; + } +}; + +void FillNontrivial(std::vector& x, std::vector& ri, std::vector& rw, + int64_t T, int64_t H, int top_k) { + for (int64_t t = 0; t < T; ++t) { + for (int64_t h = 0; h < H; ++h) { + x[static_cast(t * H + h)] = (h == 0 && t == 1) ? 0.f : static_cast(t + 1) * 0.25f + + static_cast(h) * 0.125f; + } + for (int g = 0; g < top_k; ++g) { + ri[static_cast(t * top_k + g)] = static_cast((t + g) % 8); + rw[static_cast(t * top_k + g)] = (g == 3 && t == 0) ? 0.f : 0.5f + 0.05f * static_cast(g); + } + } +} + +} // namespace + +TEST_CASE("gemma4 indexed-max-t: env parse") { + CHECK(ParseGemma4DecodeIndexedMaxT(nullptr) == 63); + CHECK(ParseGemma4DecodeIndexedMaxT("1") == 1); + CHECK(ParseGemma4DecodeIndexedMaxT("64") == 63); +} + +TEST_CASE("gemma4 indexed-max-t: host predicate table") { + const int64_t unset63 = ParseGemma4DecodeIndexedMaxT(nullptr); + const int64_t env1 = ParseGemma4DecodeIndexedMaxT("1"); + CHECK(Gemma4IndexedOkT(1, unset63, 8, true)); + CHECK(Gemma4IndexedOkT(19, unset63, 8, true)); + CHECK_FALSE(Gemma4IndexedOkT(19, env1, 8, true)); + CHECK_FALSE(Gemma4IndexedOkT(64, unset63, 8, true)); + CHECK(Gemma4IndexedOkT(63, unset63, 8, true)); + CHECK(kGemma4PrefillBatchMinT == 64); +} + +TEST_CASE("gemma4 indexed-max-t: tensor oracle T={2,19,63} x {same,peer}") { + for (bool peer : {false, true}) { + for (int64_t T : {int64_t{2}, int64_t{19}, int64_t{63}}) { + const int64_t H = 8; + const int top_k = 8; + std::vector y(static_cast(T * H), 99.f); + std::vector x(static_cast(T * H), 0.f); + std::vector ri(static_cast(T * top_k), 0); + std::vector rw(static_cast(T * top_k), 0.f); + std::vector ref(static_cast(T * H), 0.f); + FillNontrivial(x, ri, rw, T, H, top_k); + Gemma4IndexedHostSerialRef(x.data(), ri.data(), rw.data(), ref.data(), T, H, top_k); + WritingHelper fake; + fake.peer_expected = peer; + fake.H = H; + fake.top_k = top_k; + const uint64_t hits0 = Gemma4IndexedHelperHits().load(); + const auto disp = Gemma4IndexedDispatchTokens( + T, H, top_k, peer, y.data(), x.data(), ri.data(), rw.data(), + [&](const Gemma4IndexedCall& c) { return fake(c); }, [&] { ++fake.restores; }); + REQUIRE(disp.ok); + CHECK(disp.hits == static_cast(T)); + CHECK(fake.calls == static_cast(T)); + CHECK(disp.y_owner == static_cast(y.data())); + CHECK(Gemma4IndexedHelperHits().load() == hits0 + static_cast(T)); + float mad = 0.f; + REQUIRE(Gemma4IndexedOracleClose(y.data(), ref.data(), T * H, &mad)); + CHECK(mad == doctest::Approx(0.f)); + bool any_nz = false, any_z = false; + for (float v : ref) { + if (v == 0.f) any_z = true; + else any_nz = true; + } + CHECK(any_nz); + CHECK(any_z); + y[0] = 123.f; + CHECK(y[0] == 123.f); + } + } +} + +TEST_CASE("gemma4 indexed-max-t: RED wrong stride corrupts output") { + const int64_t T = 19, H = 8; + const int top_k = 8; + std::vector y(static_cast(T * H), 0.f); + std::vector x(static_cast(T * H), 0.f); + std::vector ri(static_cast(T * top_k), 0); + std::vector rw(static_cast(T * top_k), 0.f); + std::vector ref(static_cast(T * H), 0.f); + FillNontrivial(x, ri, rw, T, H, top_k); + Gemma4IndexedHostSerialRef(x.data(), ri.data(), rw.data(), ref.data(), T, H, top_k); + for (int64_t t = 0; t < T; ++t) { + const int64_t bad = t; // t as t, not t*H + Gemma4IndexedHostApplyToken(y.data() + bad, x.data() + t * H, ri.data() + t * top_k, + rw.data() + t * top_k, H, top_k); + } + float mad = 0.f; + CHECK_FALSE(Gemma4IndexedOracleClose(y.data(), ref.data(), T * H, &mad)); +} + +TEST_CASE("gemma4 indexed-max-t: RED T=1 TLS owner is invalid for T>1") { + const int64_t T = 19, H = 8; + float tls1[8] = {}; + std::vector owned(static_cast(T * H), 0.f); + CHECK(Gemma4IndexedScratchKindFor(1) == Gemma4IndexedScratchKind::TlsT1); + CHECK(Gemma4IndexedScratchKindFor(T) == Gemma4IndexedScratchKind::OwnedTH); + Gemma4IndexedScratchChoice tls{Gemma4IndexedScratchKind::TlsT1, tls1, H}; + CHECK_FALSE(Gemma4IndexedScratchValidForT(tls, T, H)); + Gemma4IndexedScratchChoice good{Gemma4IndexedScratchKindFor(T), owned.data(), T * H}; + CHECK(Gemma4IndexedScratchValidForT(good, T, H)); + CHECK(good.y != static_cast(tls1)); +} + +TEST_CASE("gemma4 indexed-max-t: release-before-retire is RED; fail-path retires while owned") { + CHECK_FALSE(Gemma4IndexedMayReleaseToPool(/*enqueued=*/true, /*retire_observed=*/false)); + CHECK(Gemma4IndexedMayReleaseToPool(true, true)); + CHECK(Gemma4IndexedMayReleaseToPool(false, false)); + bool owned = true, released = false, retire_ok = false; + CHECK(Gemma4IndexedFailPathRetireThenMaybeRelease(true, owned, released, retire_ok, + [] { return true; })); + CHECK(retire_ok); + CHECK(released); + CHECK_FALSE(owned); + owned = true; + released = false; + retire_ok = true; + CHECK_FALSE(Gemma4IndexedFailPathRetireThenMaybeRelease(true, owned, released, retire_ok, + [] { return false; })); + CHECK_FALSE(retire_ok); + CHECK_FALSE(released); + CHECK(owned); // quarantined +} + +TEST_CASE("gemma4 indexed-max-t: independent serial ref RED on candidate arithmetic/route") { + const int64_t T = 19, H = 8; + const int top_k = 8; + std::vector x(static_cast(T * H), 0.f); + std::vector ri(static_cast(T * top_k), 0); + std::vector rw(static_cast(T * top_k), 0.f); + std::vector ref(static_cast(T * H), 0.f); + std::vector bad(static_cast(T * H), 0.f); + FillNontrivial(x, ri, rw, T, H, top_k); + Gemma4IndexedHostSerialRef(x.data(), ri.data(), rw.data(), ref.data(), T, H, top_k); + for (int64_t t = 0; t < T; ++t) { + // mutated candidate: extra *2, does not go through SerialRef + Gemma4IndexedHostApplyToken(bad.data() + t * H, x.data() + t * H, ri.data() + t * top_k, + rw.data() + t * top_k, H, top_k); + for (int64_t h = 0; h < H; ++h) bad[static_cast(t * H + h)] *= 2.f; + } + float mad = 0.f; + CHECK_FALSE(Gemma4IndexedOracleClose(bad.data(), ref.data(), T * H, &mad)); +} + +TEST_CASE("gemma4 indexed-max-t: production selector arm/args identity; swap is RED") { + int same_n = 0, peer_n = 0; + float same_out = 0.f, peer_out = 0.f; + auto same = [&] { + ++same_n; + same_out = 1.f; + return true; + }; + auto peer = [&] { + ++peer_n; + peer_out = 2.f; + return true; + }; + CHECK(Gemma4IndexedSelectArm(true, false) == Gemma4IndexedArm::SameDev); + CHECK(Gemma4IndexedSelectArm(false, true) == Gemma4IndexedArm::Peer); + REQUIRE(Gemma4IndexedRunSelectedArm(Gemma4IndexedArm::SameDev, same, peer)); + CHECK(same_n == 1); + CHECK(peer_n == 0); + CHECK(same_out == 1.f); + REQUIRE(Gemma4IndexedRunSelectedArm(Gemma4IndexedArm::Peer, same, peer)); + CHECK(peer_n == 1); + CHECK(peer_out == 2.f); + auto swapped = [&](Gemma4IndexedArm arm) { + return Gemma4IndexedRunSelectedArm(arm, peer, same); + }; + same_n = peer_n = 0; + REQUIRE(swapped(Gemma4IndexedArm::SameDev)); + CHECK(peer_n == 1); + CHECK(same_n == 0); + CHECK(peer_out == 2.f); + float y = 0, x = 0, rw = 0; + int32_t ri = 0; + const auto want = Gemma4IndexedPackArgs(&y, &x, &ri, &rw); + const auto swapped_args = Gemma4IndexedPackArgs(&x, &y, &ri, &rw); + CHECK(Gemma4IndexedArgsEq(want, Gemma4IndexedPackArgs(&y, &x, &ri, &rw))); + CHECK_FALSE(Gemma4IndexedArgsEq(want, swapped_args)); +} + +TEST_CASE("gemma4 indexed-max-t: fallback scale is once, not s^2") { + const int64_t T = 19; + const int top_k = 8; + const int64_t E = 8; + std::vector orig(static_cast(T * top_k), 0.5f); + std::vector hi(static_cast(T * top_k)); + std::vector hscale(static_cast(E), 3.f); + for (size_t i = 0; i < hi.size(); ++i) hi[i] = static_cast(i % 8); + std::vector fallback = orig; + Gemma4ApplyHostExpertScaleOnce(fallback.data(), hi.data(), hscale.data(), T, top_k, E, false); + for (size_t i = 0; i < orig.size(); ++i) { + CHECK(fallback[i] == doctest::Approx(orig[i] * 3.f)); + CHECK(fallback[i] != doctest::Approx(orig[i] * 9.f)); + } +} + +TEST_CASE("gemma4 indexed-max-t: source invariants") { + const std::string moe = ReadText("src/vllm/model_executor/models/gemma4_moe.cpp"); + const std::string hip = ReadText("src/vt/rocm/rocm_gemma4_experts.hip"); + REQUIRE_FALSE(moe.empty()); + REQUIRE_FALSE(hip.empty()); + CHECK(moe.find("Gemma4IndexedDispatchTokens") != std::string::npos); + CHECK(moe.find("Gemma4IndexedSelectArm") != std::string::npos); + CHECK(moe.find("Gemma4IndexedRunSelectedArm") != std::string::npos); + CHECK(moe.find("Gemma4IndexedScratchKindFor") != std::string::npos); + CHECK(moe.find("retire-before-acc_idx-dtor") != std::string::npos); + CHECK(moe.find("RetireGemma4Fp8TopKIndexedPeer") != std::string::npos); + CHECK(moe.find("rw_idx_owned") != std::string::npos); + CHECK(moe.find("struct RwIdxTls") != std::string::npos); + CHECK(moe.find("static thread_local RwIdxTls rwt") != std::string::npos); + // T=1 scaled rw is TLS-stable; only T>1 Release()s a pooled copy. + const auto rwt_at = moe.find("struct RwIdxTls"); + REQUIRE(rwt_at != std::string::npos); + const auto t1_arm = moe.find("Gemma4IndexedScratchKind::TlsT1", rwt_at); + REQUIRE(t1_arm != std::string::npos); + const auto owned_emplace = moe.find("rw_idx_owned.emplace", rwt_at); + REQUIRE(owned_emplace != std::string::npos); + CHECK(owned_emplace > t1_arm); // pooled emplace is the T>1 branch + CHECK(moe.find("rwt.buf->Release") == std::string::npos); + CHECK(moe.find("rw_idx_owned->Release()") != std::string::npos); + // Fresh copy + scale land on the TLS/owned dest every call, never on `rw`. + const auto copy_at = moe.find("d.b.Copy(d.q, scaled->ptr(), rw.ptr()", rwt_at); + REQUIRE(copy_at != std::string::npos); + const auto key_if = moe.find("rwt.dev != compute_dev || rwt.n != n || !rwt.buf", rwt_at); + REQUIRE(key_if != std::string::npos); + CHECK(copy_at > key_if); // copy is outside the TLS-miss emplace + CHECK(moe.find("ApplyExpertScaleRw(d.q, static_cast(scaled->ptr())", rwt_at) != + std::string::npos); + CHECK(moe.find("ApplyExpertScaleRw(d.q, static_cast(rw.ptr())", rwt_at) == + std::string::npos); + CHECK(moe.find("ExpertGeGLUFp8TopKIndexedBatched") == std::string::npos); + const auto serial = ReadText("include/vllm/model_executor/models/gemma4_indexed_gate.h"); + const auto sref = serial.find("Gemma4IndexedHostSerialRef"); + REQUIRE(sref != std::string::npos); + const auto sref_end = serial.find("Gemma4IndexedOracleClose", sref); + REQUIRE(sref_end != std::string::npos); + CHECK(serial.substr(sref, sref_end - sref).find("Gemma4IndexedHostApplyToken") == std::string::npos); + CHECK(hip.find("retire_fail") != std::string::npos); + CHECK(hip.find("RetireGemma4Fp8TopKIndexedPeer") != std::string::npos); + CHECK(hip.find("RestoreComputeDev") != std::string::npos); + // The guarantee: RetireGemma4Fp8TopKIndexedPeer synchronizes the compute + // stream, keeps that result, and returns it -- it never discards the status + // and never reports unconditional success. Asserted on the function's own + // body, so an unrelated definition added after it cannot red this case. + const std::string retire = FunctionBody(hip, "bool RetireGemma4Fp8TopKIndexedPeer"); + REQUIRE_FALSE(retire.empty()); + REQUIRE(retire.back() == '}'); + CHECK(retire.find("RunGemma4Fp8ExpertGeGLUPrefillOnExpertDevice") == std::string::npos); + CHECK(retire.find("hipStreamSynchronize(cst)") != std::string::npos); + CHECK(retire.find("(void)hipStreamSynchronize") == std::string::npos); + CHECK(retire.find("return true;") == std::string::npos); + CHECK(retire.find("return ok;") != std::string::npos); +}