# PREFILL_GEMM_SCOPE - large-M NVFP4 expert/dense GEMM (design only)

**Status: DESIGN + PLAN ONLY. No kernel written, no GPU run in this pass.**
This scopes the #1 prefill lever for `llama-cpp-localai-paged`: the NVFP4 weight
GEMM at large M (prefill), where llama.cpp's `mul_mat_q` (MMQ) NVFP4 path is far
slower than vLLM's `marlin_moe_wna16` (MoE) + cutlass/nvjet (dense). Per the
prefill ground-truth that motivated this scope, the GEMM bucket is ~232 us/tok
(paged) vs ~68 us/tok (vLLM) - 3.4x slower, ~51% of the paged-vs-vLLM prefill
gap (164 us/tok).

> **Regime warning (read first).** Every "GEMM is at the BW floor / ties vLLM"
> conclusion in `README.md` section 5 is a **DECODE** finding (M<=128,
> bandwidth-bound). This document is about **PREFILL** (large M, compute /
> tensor-core-throughput bound) - a different regime, which is exactly why the
> rejected "W4A16-Marlin MoE GEMM" lever is revisited here **for prefill only**.
> The 232/164/68 us/tok prefill bucket came from the prefill ground-truth that
> commissioned this scope and is **not** in a committed in-repo profile (the
> committed profiling - `GAP_PROGRESS.md` etc. - is decode-focused). Per the
> "profile-don't-assume" rule in `.agents/vllm-parity-methodology.md`, **step 0 of
> any build is to re-confirm the prefill GEMM bucket on GPU** (nsys, prefill-only
> window) before touching code.

---

## 1. Why `mul_mat_q` is slow at large M (confirmed from source)

Source: `ggml/src/ggml-cuda/mmq.cu`, `mmq.cuh` at this backend's pin (`9d5d882d`).

MMQ is built for the **M<=128 decode tile**. Three structural facts from the code:

1. **The M (column/token) tile is capped at 128.**
   `get_mmq_x_max_host()` / `get_mmq_x_max_device()` (mmq.cuh ~108-140) return
   `128` on Blackwell (`turing_mma_available(cc)`), and the host launch loop
   (mmq.cuh ~4237) picks `mmq_x_best` only to *minimise the column-tile count for
   `ncols_max`, never exceeding `mmq_x_max`*. So a prefill ubatch of M=512 (or
   4096) tokens is processed as many `mmq_x<=128` column-tiles. The compile-time
   accumulator tile is `mmq_x`-wide; there is no large-M (e.g. 256-wide) tile
   variant. The whole tile-selection machinery exists to pick a *small* tile for
   *small* batches, not to grow for large ones.

2. **The FP4-MMA kernel is register-bound to 1 CTA/SM.**
   `mul_mat_q` for FP4 is `__launch_bounds__(warp_size*nwarps, min_blocks=1)`
   (mmq.cuh ~3579-3585), i.e. 256 threads, 1 resident block/SM (~255 regs/thread).
   The patch-0017 comment in-tree states this plainly: the kernel is
   "REGISTER-bound to 1 CTA/SM ... the under-occupancy that strands the kernel at
   ~3% of FP4 peak at M=128." At large M the work per tile is bigger, but with one
   CTA/SM the tensor cores still stall on LPDDR5x / shared-memory weight loads
   with no CTA-level latency hiding - the design has no async multi-stage global->
   shared pipeline (cp.async double-buffering) that large-M GEMMs need.

3. **Per-tile fixed overheads amortise poorly only because the tile stays small.**
   Each tile re-stages weights into shared memory, runs the `MMQ_ITER_K_FP4=512`
   K-loop, and the activations are quantized to Q8_1 (`quantize_mmq_fp4_cuda`,
   block_fp4_mmq = FP4 weights x int8 activations). For decode this is the right
   trade (FP4 weight traffic is the bottleneck). For large-M prefill the GEMM is
   compute-bound, so the right structure is big tensor-core output tiles (e.g.
   128x256), a deep async load pipeline, and full SM occupancy - exactly what
   cutlass 3.x / nvjet (cuBLAS) and marlin implement and MMQ does not.

Patch 0017 already proved every *cheap* large-tile/occupancy lever inside MMQ
(`GGML_CUDA_FP4_MMQ_Y`, `GGML_CUDA_FP4_MINBLOCKS`) is a no-win on GB10 - because
the limit is the small-tile kernel *structure*, not a tunable. To win at large M
you must leave MMQ for a large-M kernel.

---

## 2. Options (feasibility / bit-exactness / effort)

### Key enabling facts already in the tree

- **NVFP4 -> bf16/f16 dequant kernels already exist.** `convert.cu` defines
  `dequantize_row_nvfp4_cuda`; `ggml_get_to_bf16_cuda` / `ggml_get_to_fp16_cuda`
  / `ggml_get_to_fp16_nc_cuda` all return it for `GGML_TYPE_NVFP4`. The
  non-Blackwell fallback ("falls back to dequant", README s2) already uses this.
- **cuBLAS on GB10 dispatches to nvjet** (NVIDIA's JIT tensor-core GEMM) - the
  committed profiles already show `nvjet lm_head` and `nvjet non-FP4 cublas GEMM`
  rows. So a dequant->cuBLAS bf16 GEMM lands on a vendor-tuned large-M kernel for
  free.
- **BUT NVFP4 is explicitly excluded from the tensor-core cuBLAS path.** In
  `ggml_cuda_op_mul_mat_cublas` (ggml-cuda.cu ~1659) the `use_fp16` predicate
  begins `src0->type != GGML_TYPE_NVFP4 && ...`. So if NVFP4 reaches cuBLAS today
  it falls to the `else` branch: dequant to **F32** + `cublasSgemm` (**no tensor
  cores**) - useless for prefill. Relaxing this one exclusion (route NVFP4 to the
  bf16/f16 tensor-core branch, where `to_*_cuda(NVFP4)` already exists) is the
  pivot that makes option (a) a few-line change rather than a kernel.

### (a) Dequant -> cuBLAS/cutlass bf16 GEMM for large M  -- RECOMMENDED

Dequant the NVFP4 weights to bf16 (transient pool buffer) once per prefill step,
then a large-M tensor-core `cublasGemmEx` (CUBLAS_COMPUTE_32F accumulate, bf16
inputs). Activations stay bf16 (not Q8_1-quantized).

- **Feasibility: HIGH.** All pieces exist (dequant kernels, cuBLAS bf16 path,
  pool allocator). The only code change for the dense path is (i) make
  `ggml_cuda_should_use_mmq` return false for NVFP4 dense above an M threshold so
  the dispatch falls through to `ggml_cuda_op_mul_mat_cublas`, and (ii) relax the
  `src0->type != GGML_TYPE_NVFP4` exclusion so it dequants to bf16 and uses
  `cublasGemmEx` tensor-core, not f32 Sgemm.
- **Cost model (the crux - why it wins ONLY at large M).** Dequant is one extra
  weight-sized memory pass (read ~0.5B/elt FP4 + scales, write 2B/elt bf16). The
  bf16 GEMM then reads weights as bf16 = **4x the byte traffic of the FP4-MMQ
  read**. At small M (decode) this 4x weight traffic dominates -> bf16-cuBLAS
  loses -> keep MMQ (this is why decode stays FP4-MMQ; consistent with the
  README decode verdict). At large M the GEMM is compute-bound and weight traffic
  is amortised over hundreds of columns, so the 4x is cheap and cuBLAS's mature
  large tiles + async pipeline + full occupancy dominate MMQ's 3%-of-peak small
  tile. The dequant pass itself is ~one weight-read amortised over the whole
  prefill step - negligible at large M.
- **Honest ceiling.** GB10 bf16 tensor-core peak is ~**half** the FP4 tensor-core
  peak. A bf16 cuBLAS GEMM at ~70-80% of bf16 peak is ~35-40% of FP4 peak. That
  is a huge jump from MMQ's ~3% large-M utilisation, but it is **not** automatic
  full vLLM parity (vLLM prefill uses 4-bit weight tiles, staying near FP4-class
  throughput). Expect this to recover most, not all, of the 232->68 gap. See s4.
- **Bit-exactness: NEW FP path** (NVFP4->bf16 round, bf16 TC, f32 accumulate) vs
  fused FP4xQ8_1 MMQ. **Not byte-identical** - gate per-path via KLD exactly like
  the paged-MoE `8cb0ce23` precedent (README s5 / `PAGED_BITEXACT_NOTE.md`). It
  should pass *easily and favourably*: keeping activations in bf16 instead of
  Q8_1 is strictly more precise than the MMQ path, so KLD(dequant-bf16 || f16)
  should be <= KLD(FP4-MMQ || f16). This is a precision-neutral-to-better change,
  not a precision regression like the rejected lever 4.
- **Effort: LOW-MEDIUM (a few days).** Dispatch flip + exclusion relax + an M
  threshold + the KL gate + a prefill bench. No new kernel. Dense first; MoE is
  the harder follow-on (see (c)/plan).
- **Memory note.** Dequant into a *transient* pool scratch per step (do **not**
  cache bf16 weights - a persistent bf16 copy is 4x VRAM for those tensors and
  would erase the backend's "1.5-3x less memory" property). The per-step dequant
  pass is the price of keeping the model FP4-resident.

### (b) Marlin-style fused NVFP4 large-M MoE GEMM (port `marlin_moe_wna16`)

Port vLLM's marlin grouped MoE kernel (4-bit weights, f16 activations, dequant-
in-register, async cp.async pipelines, swizzled layouts).

- **Feasibility: LOW (hardest).** Marlin is a hand-tuned CUTLASS-class kernel and
  is **not NVFP4-aware** (it targets wna16 group-quant, not NVFP4's 16-elt blocks
  with ue4m3 micro-scales). You would either (i) adapt marlin to dequant NVFP4
  in-register and accumulate in f16 (abandoning native Blackwell FP4-MMA), or
  (ii) write a brand-new Blackwell sm_121 FP4-MMA large-M kernel - which is
  essentially re-implementing what cutlass 3.x / nvjet already give you via (a).
- **Bit-exactness:** new FP path, KL-gate (same as (a)).
- **Effort: HIGH (multi-week, high risk),** kernel + layout + Blackwell MMA
  scheduling + graph-safety + the bit-exact gate.
- **Verdict: do NOT start here.** Its only structural advantage over (a) is 4-bit
  weight traffic, which matters only when BW-bound = small M = **decode**, the
  regime already rejected. At large M (a) reaches the same vendor large-M kernels
  for ~1% of the effort. Keep (b) on the shelf as the *only* route to true 68
  us/tok parity if (a)'s bf16 ceiling proves insufficient and the win justifies a
  multi-week kernel.

### (c) M-threshold routing (the integration mechanism for (a))

Not an alternative to (a) - it is *how* (a) is wired. Keep FP4-MMQ for decode
(M<=threshold), switch to the large-M path for prefill.

- **Cleanest hook:** `ggml_cuda_should_use_mmq(type, cc, ne11_or_ne12, n_experts)`
  already receives M (`ne11` dense / `ne12` MoE tokens). Add an NVFP4+Blackwell
  branch: return false when M > `LLAMA_FP4_PREFILL_M` (default e.g. 256-512,
  env/`-D` tunable, default value chosen so default == today's behaviour until
  validated). It is called from both `ggml_cuda_mul_mat` (~2573/2582) and
  `ggml_cuda_mul_mat_id` (~2664), so one edit covers dense + MoE routing.
- **Dense fallthrough is clean:** `ggml_cuda_mul_mat` final `else` ->
  `ggml_cuda_op_mul_mat(..., ggml_cuda_op_mul_mat_cublas, ...)` -> with the
  exclusion relaxed, dequant->bf16->`cublasGemmEx`. Works.
- **MoE fallthrough is NOT clean (the catch):** in `ggml_cuda_mul_mat_id`, a
  false `should_use_mmq` falls to `should_use_mmf` (no NVFP4 support) then to the
  **host-side sorted per-expert loop** with a `cudaStreamSynchronize` (ggml-cuda.cu
  ~2700) - slow and **not CUDA-graph-safe** (it would break the MoE re-graph,
  patch 0025). So MoE large-M needs a *dedicated graph-safe grouped GEMM* (dequant
  the expert-gathered weights to bf16 + `cublasGemmGroupedBatchedEx`, CUDA 12.5+,
  over the existing `expert_bounds`/`ids_dst` sorted layout), not a bare
  fallthrough. This is why the plan ships **dense first, MoE second**.

---

## 3. Recommended approach + implementation plan

**Recommendation: (a) dequant->bf16 cuBLAS, wired via (c) M-threshold routing,
dense-path first, MoE grouped-cuBLAS second. Reject (b).**

### Phase 0 - confirm the bucket on GPU (no code)
- nsys prefill-only window (`-npp <large> -ntg 0/1`, exclude the graph-capture
  step) on q36-27b dense and q36-35b-a3b MoE at the backend pin. Confirm the
  NVFP4 `mul_mat_q` / `mul_mat_id` bucket is ~232 us/tok and that it is
  compute-bound at prefill M (check tensor-core active % low, not BW-saturated).
  If the bucket is not what the ground-truth claims, stop and re-scope.

### Phase 1 - dense large-M NVFP4 -> bf16 cuBLAS (the bankable win)
Files / edits:
1. `ggml/src/ggml-cuda/mmq.cu` - `ggml_cuda_should_use_mmq`: add
   `if (type==GGML_TYPE_NVFP4 && blackwell_mma_available(cc) && ne11 > LLAMA_FP4_PREFILL_M && n_experts==0) return false;`
   (n_experts==0 = dense only in Phase 1). Default threshold == effectively
   disabled until A/B-validated, env/`-D` overridable (mirror the 0017
   `GGML_CUDA_FP4_*` knob style + in-tree comment).
2. `ggml/src/ggml-cuda/ggml-cuda.cu` - `ggml_cuda_op_mul_mat_cublas`: relax the
   `src0->type != GGML_TYPE_NVFP4` guard in `use_fp16` (prefer a dedicated bf16
   branch: NVFP4 -> `ggml_get_to_bf16_cuda` -> `cublasGemmEx` CUDA_R_16BF /
   COMPUTE_32F, matching the existing BF16 src0 branch for best accuracy).
3. Transient pool scratch for the dequanted weights (reuse `ggml_cuda_pool_alloc`
   as the existing branch does; no persistent allocation).

### Phase 2 - MoE grouped large-M (the harder, higher-value follow-on)
1. New grouped path reached from `ggml_cuda_mul_mat_id` when
   `should_use_mmq`==false for NVFP4+large-M+`n_experts>0`: dequant the
   expert-gathered weights to bf16 and run `cublasGemmGroupedBatchedEx` over the
   existing `expert_bounds` / `ids_dst` sorted layout that `mul_mat_q` already
   builds. Reuse the patch-0023 de-dup'd activation gather where applicable.
2. **Must stay CUDA-graph-safe** - no host sync (do not fall into the legacy
   sorted loop). Validate the MoE re-graph (patch 0025 / `LLAMA_MOE_FORCE_GRAPHS`)
   still captures.

### The bit-exact / KL gate (both phases)
- Greedy md5 on the standard prompt (README s5) to detect *unexpected* divergence
  on the non-prefill paths (must stay == the per-path reference: dense
  `5951a5b4`, paged-MoE `8cb0ce23`). The large-M path itself will differ -> gate
  it by KLD vs the f16 reference, requiring `KLD(new||f16) <= KLD(FP4-MMQ||f16)`
  and PPL within the established band, recorded in `PAGED_BITEXACT_NOTE.md`.
- `test-backend-ops` MUL_MAT / MUL_MAT_ID at NVFP4 **prefill shapes** (large M)
  CUDA0-vs-CPU, plus the existing decode shapes to prove decode is byte-untouched
  (default threshold keeps decode on MMQ).

### The bench
- `llama-batched-bench -fa on -ngl 99` reporting **S_PP** (prefill t/s), swept
  over prefill length and `npl`, A/B with `LLAMA_FP4_PREFILL_M` off vs on, dense
  and MoE, vs stock and vs the vLLM prefill reference. Per-lever A/B discipline
  (`.agents/vllm-parity-methodology.md`): one knob at a time, record the rejected
  threshold values too.

---

## 4. Honest risk + expected speedup

- **Phase 1 (dense) is a tractable routing change, not a kernel project** - days,
  low risk. It reuses existing dequant kernels and the existing nvjet/cuBLAS
  large-M path; the net new code is a threshold + a one-line exclusion relax + a
  KL gate.
- **Phase 2 (MoE) is medium risk** - the grouped-batched cuBLAS wiring +
  CUDA-graph-safety is real work (the bare fallthrough is a slow, graph-breaking
  host loop), but still far short of a from-scratch kernel.
- **Will the GEMM bucket hit 232 -> ~68 us/tok (full vLLM parity)? Honestly, no -
  not from bf16-cuBLAS alone.** bf16 tensor-core peak on GB10 is ~half FP4 peak,
  so the realistic floor for a dequant->bf16 GEMM is ~**90-130 us/tok** (roughly
  35-45% of FP4 peak at ~70-80% of bf16 peak). That recovers ~**60-75%** of the
  232->68 bucket gap = a large prefill win (the GEMM is ~51% of the total prefill
  gap, so closing ~two-thirds of it is a meaningful S_PP improvement), but it
  leaves a residual. **True 68 us/tok parity requires a native FP4-MMA large-M
  kernel (option (b)) - the multi-week project** to greenlight only if Phase 1's
  measured win proves the prefill regime matters enough to fund it.
- **Recommendation:** build Phase 1, measure, and let the measured dense S_PP
  gain decide whether Phase 2 (MoE grouped cuBLAS) and ultimately (b) (native FP4
  large-M kernel) are worth funding. Bank the cheap two-thirds before paying for
  the kernel.

---

## 5. Summary table

| Option | Feasibility | Bit-exact | Effort | Verdict |
|---|---|---|---|---|
| (a) dequant->bf16 cuBLAS large-M | HIGH (parts exist) | new FP path, KL-gate (likely better PPL) | LOW-MED (days) | **RECOMMENDED** (dense first) |
| (b) Marlin/native FP4 large-M kernel | LOW | new FP path, KL-gate | HIGH (multi-week) | shelf - only route to true 68 us/tok |
| (c) M-threshold routing | HIGH | n/a (mechanism) | LOW | **the wiring for (a)** |

Decode is untouched by all of the above (threshold keeps M<=128 on FP4-MMQ); this
is a **prefill-only** lever.