forgo.cloud
Sign in
Repo workspace

forkjoin-ai/gnosis

Prefill attention → GPU: handoff

distributed-inference/PREFILL_ATTENTION_GPU_HANDOFF.md
forkjoin-ai/gnosis
Repo overview Open in workbench

Prefill attention → GPU: handoff

STATUS 2026-06-29 — GPU path IMPLEMENTED + parity-verified on real GPU. The batched GPU attention below is built behind MESH_GPU_PREFILL_ATTN=1 (default OFF) and passes the prefill_chunk_matches_forward_one_multichunk parity test on a physical GTX 1080 Ti, with the in-op FLUX_GPUCHECK probe showing max |gpu−cpu| ≈ 6e-10 (vs the 1e-4 bar). What's done, and what's still open, are in "Implementation status" at the bottom. The blind ~17-min Cloud Build loop is no longer required to ITERATE — the local CUDA box now compiles + GPU-parity-tests this in ~20s (MESH_CUDA_ARCH=sm_61, nvcc via VS BuildTools cl.exe). A cloud build is still needed to DEPLOY + benchmark the real Mistral-7B speedup on the L4.

Long-prompt prefill on the gnosis mesh (Cloud Run L4, gnosis-openai-mesh, Mistral-7B) is bottlenecked by attention, not matmuls. This doc hands off the GPU implementation. A CPU parallelization already shipped as the interim win (see "What already shipped"); the GPU path is the higher ceiling.

Diagnosis (measured)

  • A 1358-token prompt prefilled in ~75s. The matmuls (QKV/O/FFN) are already GPU-batched (cuBLAS, batch-128 after the maxBatchSize fix in distributed-inference-host/src/bin/openai-server.ts) and sum to only ~7s.
  • The other ~68s was attention: an un-GEMM'd, per-token, per-head causal loop in forward_range_chunk that — before the interim fix — ran with a serial chunks_mut over heads inside a serial for i in 0..seq_len, i.e. effectively single-core. At ~242 GFLOP for 1358 tokens that's ~3.6 GFLOP/s — 15–25× under even CPU SIMD, and ~4000× under the L4.
  • Batch-5 vs batch-128 prefill measured the same ~75s → the cost scales with token count (O(context²) attention), not matmul chunking. That is the tell.

What already shipped (interim, CPU)

model.rs forward_range_chunk: attention was split into two phases — (A) the existing per-token rope + save_kv loop, then (B) attention parallelized over the flat seq_len × num_heads work set via rayon par_chunks_mut(head_dim) on batch_xb[..seq_len*qkv_dim] (chunk c → token c/num_heads, head c%num_heads). The per-(i,h) math is unchanged and bit-identical (same attention_scores_f32 / memo path / softmax / accumulate_weighted, causal range 0..=pos), so the prefill_chunk_matches_forward_one_multichunk parity test (1e-4 tol) still passes. Expected ~8× (cores) → prefill ~75s → ~10–18s. This is a wall-clock parallelization, NOT an algorithmic fix — it stays O(n²) on the CPU.

GPU design (the real fix — sub-second prefill)

Replace the Phase-B CPU attention with a batched GPU path, gated by a parity check with clean CPU fallback (mirror the try_* / gpucheck_* discipline in cuda_matmul.rs). Per layer, for the chunk of T = seq_len query tokens at absolute positions start_pos..start_pos+T, with L = start_pos + T context keys:

  • GQA: num_kv_groups = num_heads / num_kv_heads (Mistral: 32/8 = 4). Query head h reads kv head kv_h = h / num_kv_groups. head_dim = 128, scale = 1/sqrt(head_dim).
  • Layout (host buffers already populated by Phase A):
    • Q per head h: batch_q[i*qkv_dim + h*head_dim ..] for i in 0..T — strided by qkv_dim; gather into a contiguous [T × head_dim] (or pass the stride to cuBLAS via lda).
    • K/V per kv head: `kv_cache.keys/values[layer_offset + kv_hmax_seq_lenhead_dim
      • t*head_dim ..]fort in 0..L— already contiguous[L × head_dim]`.
  • Step 1 — QKᵀ: Scores[T × L] = Q[T × head_dim] · K[L × head_dim]ᵀ, batched over the num_heads heads → cuBLAS cublasSgemmStridedBatched (f32; cudarc Gemm/GemmConfig + the raw-pointer bodies already in cuda_matmul.rs, e.g. gemm_with_weight_f16_ptr). Q stride qkv_dim, K stride per kv head; with GQA, 4 q-heads share one kv-head's K (set the K batch stride to 0 within a group, or loop kv heads × 4).
  • Step 2 — scale + causal mask + softmax: one CUDA kernel over [heads × T × L]. Row (h, t) is valid for keys j ≤ start_pos + t; set j > start_pos+t to -inf (→ 0 after softmax, exactly matching the CPU loop that only sums t in 0..=pos). Multiply by scale before the row-max/softmax. New .cu file in src/cuda/ wired through build.rs exactly like q8_gemv.cu (nvcc → OUT_DIR/*.cubin, include_bytes!). Softmax could also run on host if cheap, but keep scores resident to avoid a 2×[heads×T×L] dtoh/htod.
  • Step 3 — AV: Out[T × head_dim] = Scores[T × L] · V[L × head_dim], batched over heads → cublasSgemmStridedBatched. Scatter Out[i, h, :] back to batch_xb[i*qkv_dim + h*head_dim ..].
  • Correctness gate: compare GPU Out to the CPU Phase-B reference to f32 tol (1e-4), like gpucheck_q8. Any mismatch / unsupported shape / GPU error → return false → caller runs the CPU Phase B. Add an env flag (default ON for prefill T>1), decode (T==1) stays on the existing kernels.

Insertion points

  • src/model.rs forward_range_chunk — the Phase-B block (search the "PREFILL ATTENTION (parallel over token × head)" comment added by the interim fix). Wrap: if !try_gpu_prefill_attention(...) { <existing CPU par_chunks_mut block> }.
  • src/model.rs forward_range_attn_chunk (~line 9374) — the split-phase sibling used by layer-sharded paths. Check whether it has the same per-token attention loop; if so, route it through the same GPU entry. (The single-fat-station mesh path that was benchmarked uses forward_range_chunk.)
  • src/cuda_matmul.rs — add try_prefill_attention(...) next to try_mat_vec_q8_0_rows_shared / q8_gemv_kernel, reusing the cuBLAS handle, resident-buffer pools, and op-logging (FLUX_GPU_OP_LOG_EVERY).

Hard constraint (why this wasn't done inline)

No local CUDA toolkit on the dev box (nvcc absent, no /usr/local/cuda) → the cuda feature cannot be compiled or type-checked locally. Every cuBLAS / .cu edit is a blind ~17-minute Cloud Build with no local feedback, so iterating a correctness-critical attention kernel that way is slow and risky on the production daily-driver image. Do this on a CUDA-equipped box with a tight cargo build --features cuda → cargo test loop; verify the parity gate locally on GPU before deploying.

Build / deploy / benchmark (same machinery)

  • Build image (cuda+node, ~17min): gcloud builds submit --config cloudbuild.moonshine-v3.yaml . from the staged context (full workspace minus firefox/regenerables/media). Dockerfile drift already fixed: moonshine/assets COPY + node-builder python3 make g++ (Dockerfile.moonshine).
  • Deploy to a 0%-traffic tagged revision (prod stays 100% on gnosis-openai-mesh-00093-vgs): gcloud run deploy gnosis-openai-mesh --image .../moonshine:<tag> --no-traffic --tag meshfast. URL https://meshfast---gnosis-openai-mesh-6ptd7xm6fq-uc.a.run.app.
  • Benchmark: POST /v1/chat/completions model mistral-7b, a 1358-token prompt, max_tokens:4, read _gnosis.elapsed_ms. Correctness: short-prompt completion must stay coherent and match the CPU path; confirm 0 parity divergence in logs.

Expected ceiling

  • Interim CPU parallelization: ~10–18s for 1358-token prefill (8 vCPU, still O(n²)).
  • GPU batched attention: sub-second for the attention portion (242 GFLOP on an L4's ~30 TFLOP f32 ≈ ~8ms of compute; dominated by the QKᵀ/AV GEMMs + softmax + the per-layer dtoh/htod). Realistic end-to-end long-prompt prefill: low single- digit seconds, bounded by the remaining per-token host glue (norms/rope/residual) and the matmuls (~7s today — itself a candidate for fewer host roundtrips).

Implementation status (2026-06-29)

DONE — built, gated, parity-verified on a physical GPU (GTX 1080 Ti, sm_61):

  • Kernel src/cuda/prefill_attn.cu — prefill_attn_softmax: scale-free per-row causal-mask + softmax (one CUDA block per (head, query row), two shared-memory tree reductions). The 1/sqrt(head_dim) scale is folded into the QKᵀ GEMM alpha, so the kernel only masks + softmaxes; masked tail written 0 so the full-width AV GEMM reproduces the CPU 0..=pos sum. Compiled via build.rs exactly like the q8/q4k/q6k kernels (nvcc --cubin → include_bytes!).
  • build.rs — kernel added to the list; arch now MESH_CUDA_ARCH (default sm_89 for the L4; set sm_61 on the local 1080 Ti box). -Wno-deprecated-gpu- targets added (harmless on sm_89).
  • cuda_matmul.rs — try_prefill_attention() (public gate) → GpuContext::prefill_attention(): per query-tile (Tq=256, bounds the [heads×Tq×L] scores VRAM) it runs QKᵀ and AV as cuBLAS sgemmStridedBatched (f32), batched over the num_kv_groups query heads that share each kv head (GQA via K/V batch stride 0 within a group), with the softmax kernel between. Column-major: QKᵀ = (OP_T, OP_N), AV = (OP_N, OP_N). FLUX_GPUCHECK=1 runs an independent f64 CPU recompute and logs max |gpu−cpu|.
  • model.rs forward_range_chunk Phase B — gated try_prefill_attention call ahead of the CPU par_chunks_mut loop; falls back to CPU on any miss. ONLY eligible when the CPU-only variants are inactive (attention-score memo OFF, bowl q-filter NONE) — the GPU path is plain causal attention.
  • Flag MESH_GPU_PREFILL_ATTN=1, DEFAULT OFF (consistent with every other MESH_GPU_* additive flag; the handoff suggested default-on but a brand-new correctness-critical path ships gated). Decode (T==1) never takes this path.
  • Parity: prefill_chunk_matches_forward_one_multichunk passes with the flag on; GPUCHECK max |gpu−cpu| ≈ 6e-10. Flag-off run shows no PREFILL-ATTN logs (CPU path intact). cargo check (no cuda feature) is clean.

Local loop (this CUDA box): nvcc needs MSVC cl.exe on PATH — VS 2022 BuildTools …/VC/Tools/MSVC/<ver>/bin/Hostx64/x64. Then: MESH_CUDA_ARCH=sm_61 MESH_GPU_PREFILL_ATTN=1 FLUX_GPUCHECK=1 cargo test --features cuda --lib prefill_chunk_matches_forward_one_multichunk -- --nocapture. (The existing q8/q4k/q6k .cu use __builtin_memcpy and don't compile under the MSVC host — they stay CPU-fallback locally; unrelated to this path. Fixing them for Windows is a separate cleanup.)

OPEN — not yet done (needs the L4 cloud build):

  1. Deploy + benchmark on the L4 with the real Mistral-7B knot: build cloudbuild.moonshine-v3.yaml (sm_89 cubin via the default arch), deploy a 0%-traffic meshfast tag, POST the 1358-token prompt with MESH_GPU_PREFILL_ATTN=1, read _gnosis.elapsed_ms, and confirm the attention portion drops from ~68s toward sub-second + 0 GPUCHECK divergence in logs. Only after that does flipping the default (or enabling it in the mesh env) make sense.
  2. forward_range_attn_chunk (~line 9374, the layer-sharded sibling) is NOT yet routed through the GPU entry — only forward_range_chunk (the single-fat-station path that was benchmarked) is. Wire it the same way if/when the sharded path needs it.
  3. Perf follow-ups (correctness-orthogonal): pool the per-tile Q/K/V/scores device buffers instead of per-tile htod/alloc_zeros; keep scores resident across the two GEMMs (already done within a tile); consider fusing the softmax into the AV epilogue. None of these change numerics.