Pull requests / #2995
#2995 prefill: per-expert NVFP4 scales on the peer-prefill route (opt-in)
open · @florian862-art · 0 comentarios · En GitHub
BenchmarksSetup & installMulti-GPUNVIDIA / CUDAModels & quantsDocumentation
Descripción
## What The peer GPU in a `--peer-device` setup computes its experts' MoE rows for every prompt chunk, but on NVFP4 packs its MMQ products ran **without any per-expert scales** (`grp_tail` null, no `s_gate/s_up`, no `s_down`, no `y_scale`) - so its rows were numerically wrong and `set_peer` refused the combination. This PR wires the scales through and keeps the whole route behind `STRATA_PF_PEER_NVFP4=1`. **Commit 1 (`70d6f941`) - the core patch, `prefill.cpp` + the `generate.cpp` gate only:** - `gather_native` carries the expert's scale tail into `P.grp_tail` (peer group gather) - `swiglu_scaled` applies `s_gate/s_up` on the native path (peer loops + primary loop) - `scale_down_rows` applies `s_down` **in place** on both GPUs' `Dm` rows right after the down-proj - instead of the `row_sd` table the combine reads - `y_scale` feeds the w4a4 dequant on the gu/dn products - `moe_combine` gets `row_sd = nullptr` when the peer is active Two existing bugs fall out of the in-place design: the sums mode never applied `s_down` to GPU 0's own rows (`moe_combine_peer[16]` has no `row_sd` param), and with the peer active `moe_combine` read `m.row_sd[rows_local..T*K)`, which is never written (the 31 non-finite logits behind every cold >=20k prompt on the peer route). `kernels.cu`/`kernels.hpp` are untouched - no kernel ABI change. **Commit 2 (`0b69db53`) - separate, droppable:** sums mode on the **P2P** route (`STRATA_PF_PEER_SUMS_P2P=1`): forward keeps P2P, the return carries one weighted FP16 sum per token through mapped host memory instead of FP32 rows. Correct but slower than the host route on this box (P2P link here runs at 5.5 GB/s, below PCIe-to-RAM); it should pay off on boxes where P2P beats PCIe (NVLink etc.), hence opt-in + documented. ## Measured (2x RTX 5070 Ti 16 GB, NVFP4 pack, 4k cold prompt, median of 3) | route | prompt tok/s | correctness | |---|---|---| | P2P route, FP32 rows back | 1383 | (pre-existing, scales n/a on this route) | | host route + sums, **no scales** (old C reference) | 2683 | refused for NVFP4 before this PR | | host route + sums, **with scales (this PR)** | **2718** (2717/2718/2729; 84k prompt: 2638) | recall gate: both needles exact at 24,933 prompt tokens, **0 non-finite logits** (was: every cold >=20k prompt died) | | P2P forward + sums back (commit 2) | 2259 | recall gate passes, 0 non-finite | The scale math (two tiny row-scale kernels + `swiglu_scaled` per group, all on `ps`) costs nothing measurable - the patched route is slightly *above* the unscaled reference. ## How to test ``` STRATA_PF_PEER_NVFP4=1 STRATA_PEER_NO_P2P=1 # config C (host route, this PR's core) STRATA_PF_PEER_NVFP4=1 STRATA_PF_PEER_SUMS_P2P=1 # + commit 2, only interesting with fast P2P ``` The route banner in the log names the live route. Without the env flags nothing changes: `set_peer` keeps refusing NVFP4 experts and `generate.cpp` keeps the `rows == 0` skip. Design notes, the review history and the full measurement log: `src/prefill/README-nvfp4-peer.md`.
En el sitio
Enlaces a install, modelos, releases.