Skip to content

[Engine] gfx906 prefill: native MMVQ for small-chunk dense projections (opt-in STRATA_PF_MMVQ) - #4158

Draft
KevinX8 wants to merge 4 commits into
Niko1221:mainfrom
KevinX8:gfx906-pf-mmvq
Draft

KevinX8 wants to merge 4 commits into
Niko1221:mainfrom
KevinX8:gfx906-pf-mmvq

Conversation

@KevinX8

@KevinX8 KevinX8 commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Problem and result: see #4157. Agent turns of 100-256 new tokens on gfx906 spend most of the batched pass dequantizing the dense projections' GGUF blocks and running small rocBLAS GEMMs. With STRATA_PF_MMVQ=1 those chunks use decode's native MMVQ on the blocks instead: a 165-token turn on an MI50 goes from 1.83 to 1.57 s (two passes each); larger chunks are unchanged.

Related discussion: #4157 (scope question open; draft until it is answered).

Approach: Gemm::native_mmvq (src/prefill/gemm.cu, declared in include/strata/prefill/gemm.hpp), called first in Gemm::native, reusing native_mmvq_supported() / native_quantize_q8_1() / native_mmvq() from include/strata/kernels/native_mmvq.hpp. The quantizer and its byte count take 1..8 columns; Q8_1 blocks are per 32 values with column j at j * K / 32, so a slice is quantized as one column of n * K values (byte-identical; the quantizer itself runs one flat kernel over n_in * ncols). One 16 MiB workspace per Gemm, allocated at first use, never grown (no mid-prompt free/sync); a failed allocation turns the path off for that Gemm. Off unless STRATA_PF_MMVQ=1; with it unset the cost is one static bool test per native() call. Based on main v0.2.7.1; the gfx906 build needs #3986 to compile at all (dflash2.cpp), so this branch was built with #3986 applied under it (not part of this PR).

Validation:

  • Passed: tests/hip/pf_mmvq_parity.cpp (registered in the STRATA_HIP_GFX906 test block), this branch on v0.2.7.1 + [Engine] gfx906: fix the build of src/core/dflash2.cpp #3986, MI50, ROCm 10.1: Q4_K, Q6_K, IQ4_XS, Q8_0 at T = 1 / 7 / 8 / 129 / 512 against a host double product on the dequantized weights, one row x1000 — rel RMS 5.3e-3 to 6.3e-3, max 4.7e-3 to 9.0e-3 per token row; ldy = N + 16, beta = 1, T above the threshold, K = 640 and BF16 are declined and leave Y untouched; device memory does not change between a 1-token and a 512-token call.
  • Passed: ctest on the gfx906 build, 150/157. The 7 failures are the same with the flag unset and do not touch this path: mmvq_il_parity, verify_parity (CUDA-only kernels), pf_wmma_parity (RDNA WMMA), ple_parity (Q2_0 fixture not present), split_lend_complement_test, platform_memory_test, isa_guard_check.
  • Passed: through serve/server.py with the switch on: a 30-minute soak of one growing coding conversation (276 turns, context to 76k), two turn benchmarks and a 118k-token fresh read — no verify timeouts, 910 MiB VRAM free after load. Decode unchanged (44.1 -> 45.5 tok/s short, 31.1 -> 30.7 at 118k depth).
  • Passed: quality harness numbers in [Engine] gfx906 prefill: native MMVQ for the dense projections of small prompt chunks #4157.
  • Not run: CUDA and RDNA HIP builds (no hardware or toolchain here). The code compiles into them; a build check there would be welcome. No model other than Qwen3.8-Flash-Next quality-checked.

Limits: measured on one gfx906 card; the 256 default rests on the 165- and 480-token points of one sweep (break-even ~290 by interpolation). Not covered by the test yet: Inf/NaN rows, a K large enough to hit the 16 MiB slice cap, an ldx = K + 64 decline. Below the threshold the result depends on the chunk size (Q8_1 activations).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant