Pull requests / #439

#439 prefill: a group's expert gathers in one launch (bit-identical, long prompts +4%)

closed · @architectds · 0 comentarios · En GitHub

BenchmarksSetup & installAMD / HIPNVIDIA / CUDAModels & quantsWindowsLinux

Descripción

## What

The MMQ prompt path gathered every routed expert into its MMQ group buffer with its own launch: one
`mmq::gather_native` per expert, ~22,000 per chunk on IQ3_XXS. This change launches a group's gathers together.

- **One launch per group.** `mmq::gather_native_batch` copies up to 16 experts (`MMQ_GROUP`) in one kernel
  (`blockIdx.y` = expert). The same bytes land in the same places, so the products are bit-identical.
- **Gathers wait in a batch.** A group's gathers are launched when the group's products need them.
- **Ring-slot safety.**
  - A slot is given back only after its gather has been launched: its `used` event, and on the streamed walk the
    issuer's `consumed`.
  - The 8-slot staging ring (chunks below 1,024 tokens) launches the waiting gathers before it reuses a slot.
  - The copy engine therefore never refills a slot that is still to be read.
- **Native packs only.** The Strata Q2 blob path and the FP16 expert path are unchanged.
- **A/B switch.** `STRATA_PREFILL_GATHER_ONE=1` restores one launch per expert.
- **New test.** `tests/cuda/prefill_gather_batch_test` (ctest, `STRATA_BUILD_TESTS`) compares the batch with one
  `gather_native` per expert, byte for byte. It covers batches of 1-16 experts, group positions starting mid-group,
  and the unaligned / odd-size fallback.

## What it speeds up, and what it does not

- **Long prompts (8K chunks): +3.6% at 8K tokens, +4.3% at 30K.** At 8K chunks the prompt path is bound by the GPU's
  own work, and the per-expert launches were part of it.
- **Prompts below ~4K tokens: unchanged (±0.2%).** These are bound by the expert copies over PCIe. From 1,024-token
  chunks every non-resident expert streams, so a 1K-token prompt copies about as much as an 8K one. The launches were
  hidden behind those copies.
- **Decode: untouched.**
- **About `STRATA_PREFILL_TIMING`:** its "dequant" phase (22-42% of the GPU timeline here) overstates the launch
  cost, because the timer adds events of its own per expert. The table below is uninstrumented.

## Measured

**Setup:**
- RTX 5070 Ti 16 GB on PCIe 3.0 x16 (X370), Ryzen 9 5900XT, DDR4-2133, Windows 11, CUDA 13.0.
- IQ3_XXS native pack, text-only args (`--vram-reserve-mib 1900`: 4,560 expert slots), `--prefill auto`
  (8,192-token chunks).
- One binary, runs in the order one, batch, one, batch. Every prompt starts with its own first line, so nothing is
  reused.

| Prompt tokens | One launch per expert (ms) | Batched (ms) | Change |
| ---: | ---: | ---: | ---: |
| 254 | 1,812 / 1,817 | 1,813 / 1,818 | ±0 |
| 1,030 | 2,204 / 2,208 | 2,204 / 2,205 | ±0 |
| 3,866 | 3,430 / 3,425 | 3,420 / 3,419 | +0.2% |
| 8,179 | 4,881 / 4,881 | 4,717 / 4,706 | **+3.6%** |
| 30,118 | 18,452 / 18,439 | 17,704 / 17,679 | **+4.3%** (1,633 → 1,702 tok/s) |

**Exactness:**
- `STRATA_PREFILL_DUMP_R` (the prompt path's final residuals at every 64th position) is byte-identical between the
  two modes for prompts of 128, 256, 512, 1,024, 2,048, 4,096, 8,192 and 9,200 tokens. That covers the staging
  path, the streamed walk, and a two-chunk prompt.
- The 24 greedy tokens after each of those prompts are identical too.
- `prefill_gather_batch_test` passes.

**Not tested here:**
- **HIP:** the file builds as HIP, and the new kernel uses only what it already uses (`uint4`, `dim3` launches).
- A layer split.
- Quants other than IQ3_XXS.
- Linux.

Developed with an AI coding assistant; all numbers measured on the machine above.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

En el sitio

Enlaces a install, modelos, releases.