Follow-up to #121 (PR #128). BatchedGpuMoeFfn (the GPU-SLRU routed-MoE batched prefill, reached only with SHARPI_CPU_MOE=0) batches the routed-expert gate/up/down GEMMs grouped-by-expert, but its Phase 3 per-token reduce and Phase 1 per-token shared-expert still loop on the host:
for (int i = 0; i < N; i++) { // per token
_gpu.Clear(_gpuHidden);
for (int k = 0; k < na; k++) { // per active expert
var partV = _gpu.View(downPartial, ...);
_gpu.AddScaledInPlace(_gpuHidden, partV, weights[...]); // host-issued
}
_gpu.AddInPlace(_gpuHidden, sharedV);
_gpu.CopyDeviceRegion(hiddenAll, ...);
}At N = 4096, na = 8 this is ~120k sequential host-issued stream ops, whose P/Invoke + driver-launch overhead undercuts the grouped-GEMM batching win on this stage (flagged by Gemini review on PR #128).
Scope
- A single NVRTC batched weighted scatter-reduce kernel: for each token row, sum the
na unweighted down-partials in top-k slot order with the per-(token,slot) weights, then add the shared-expert row — in one launch over N × embDim.
- Must stay bit-identical to the current host reduce (which mirrors the sequential
GpuMoeFfn Clear → per-k AddScaledInPlace (top-k order) → AddInPlace(shared) order). The per-token FP accumulation order is load-bearing (cf. the K/V-MatVecDual MTP-parity note) — the kernel must reduce in the same k=0..na-1 order per token.
- Optionally batch the Phase 1 shared-expert (gate/up/down) over N as well.
Notes
- Only affects the GPU-SLRU config (
SHARPI_CPU_MOE=0); the benchmarked 35B-A3B default routes MoE to CPU (BatchedRoutedExperts, already batched). So this is a correctness-neutral throughput optimization for the on-GPU-experts path, not a default-path regression.
Tests
- The existing
BatchedTrunkGpuFfn_BitwiseMatchesSequential_GpuSlruMoe oracle (with the batched flags pinned on) must stay bit-identical; add an A/B prefill measurement on the GPU-SLRU config to quantify the win.
Follow-up to #121 (PR #128).
BatchedGpuMoeFfn(the GPU-SLRU routed-MoE batched prefill, reached only withSHARPI_CPU_MOE=0) batches the routed-expert gate/up/down GEMMs grouped-by-expert, but its Phase 3 per-token reduce and Phase 1 per-token shared-expert still loop on the host:At
N = 4096,na = 8this is ~120k sequential host-issued stream ops, whose P/Invoke + driver-launch overhead undercuts the grouped-GEMM batching win on this stage (flagged by Gemini review on PR #128).Scope
naunweighted down-partials in top-k slot order with the per-(token,slot) weights, then add the shared-expert row — in one launch overN × embDim.GpuMoeFfnClear→ per-kAddScaledInPlace(top-k order) →AddInPlace(shared)order). The per-token FP accumulation order is load-bearing (cf. the K/V-MatVecDual MTP-parity note) — the kernel must reduce in the same k=0..na-1 order per token.Notes
SHARPI_CPU_MOE=0); the benchmarked 35B-A3B default routes MoE to CPU (BatchedRoutedExperts, already batched). So this is a correctness-neutral throughput optimization for the on-GPU-experts path, not a default-path regression.Tests
BatchedTrunkGpuFfn_BitwiseMatchesSequential_GpuSlruMoeoracle (with the batched flags pinned on) must stay bit-identical; add an A/B prefill measurement on the GPU-SLRU config to quantify the win.