Repository navigation
cuda: opt-in quantized-native MMA FlashAttention, incl. decode and small batches (rebased on upstream master) - #85
Open
Piggidragon wants to merge 505 commits into
Conversation
github-actions
Bot
added
documentation
This was referenced Sep 17, 2026
* metal : few-row MMA mat-mul and batched copies for speculative decoding Speculative decoding verifies a few draft tokens per step. Without the tensor API, Metal ran these mat-muls with the mat-vec kernels, whose time grows with every src1 row, so DFlash2 decoding on an M3 Ultra was slower than serial decoding. - add mat-mul kernels for 2..16 src1 rows on 8x8 simdgroup matrices: each weight is dequantized once for all rows, and the simdgroups of a threadgroup split K. Q4_0, Q8_0 and Q5_K have their own kernels, F32, F16, Q4_1, Q5_0, Q5_1, Q4_K and Q6_K use a generic path over the 16-weight dequantizers, and Q4_0 at 2 rows uses a 2-row variant of the mat-vec kernel - use them only on MTLGPUFamilyApple7+ without the tensor API, from the row count at which they beat the mat-vec kernels on an M3 Ultra (F32: 6, F16, Q4_K, Q5_0, Q5_1: 3, other types: 2) - fusion table: MUL_MAT + ADD adds a same-shape residual in the MMA store, and up to 16 adjacent same-layout f32 copies between the same two tensors run as one dispatch - the fusion checks and ggml_graph_optimize take the device props, so the reorder packs MUL_MAT + ADD only on devices that can fuse it, at every src1 row count - views do not count toward GGML_METAL_FUSION_MAX when the reorder packs a group, so 16 recurrent state snapshot copies with views between them stay one group - the encoder checks the inner nodes of a fused group for concurrency, tracks written views by their extent, and does not count the destination of a CPY as a read - CONCAT splits long rows across threadgroups when there are few rows - tests: few-row MUL_MAT, MUL_MAT_ADD, CPY_BATCH and CONCAT cases in test-backend-ops (with a prepare_graph hook for the copy order), test-metal-graph-optimize, test-metal-cpy-batch-alias * metal : remove the CPY_BATCH fusion and the memory range changes Remove the batched copy fusion with its kernel and tests, and revert the memory range changes, as suggested in review. The memory ranges, the graph reorder and the CPY encoder are again the same as on master. * cont : clean-up * cont : drop has_tensor gate * cont : clean-up operand/residual logic * cont : drop Q4_0 ne11=2 special-case * cont : add kernels/mul_mv_mma.metal * cont : consolidate mma pipeline selection logic * cont : decouple fusion logic from device props --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
…ml-org#29901) * cuda: tile the lightning indexer over keys and tokens for 4 heads With too few heads for a wmma tile, a block scores 64 keys against 8 tokens: the keys are staged once in half precision, the queries one head at a time, and each thread owns one key for two tokens, so no dot product needs a cross thread reduction. Batches smaller than a token tile keep the vector kernel. test-backend-ops measures 4 heads. * cuda: multiply the lightning indexer tile in float Address review from am17an: the half2 products overflow once a single q * k exceeds the f16 range. The queries stay in float in shared memory and each half2 of keys is widened once for both tokens, so every product and sum is computed in float. * cuda: widen each lightning indexer key once for all heads The tile kernel stages the queries and weights of every head at once, so each key element is widened from half once and feeds all heads, with a single barrier. F16 keys are copied into the tile without a float round trip. Keeping the keys in float in shared memory measures slower, the occupancy drops. * cuda: stop the lightning indexer tile from spilling registers on ROCm Each thread of the tile kernel now scores two keys for a single token, so a warp shares its token and the query reads are broadcasts: six shared reads per element pair instead of nine for the same products. The inner loop is unrolled by 8, which keeps gfx908 at 63 VGPRs with no spill where the fully unrolled loop needed over a thousand, and makes the kernel 36x faster on an R9700 and slightly faster on CUDA.
…l-org#29483) * add supports for q1/q5/q3_k/q5_k/q6_k/mxfp4 of mmvq path * Add K_QUANTS_HANDLING macro to q1_0 of mmvq path
* ggml-cpu : add Q8_0 IME1 matrix kernel for SpacemiT X60 On the SpacemiT X60, IME matrix acceleration only covered Q4_0/Q4_1/Q4_K. Q8_0 had no IME1 kernel, and since the SpacemiT build sets GGML_CPU_REPACK=OFF there was no repack path compiled in either, so Q8_0 had no accelerated path at all and ran roughly ten times slower than Q4_0 for prefill on the same board. - add make_block_q8_0x16 and the Q8_0 repack entry: interleave the weights into the 16-column layout the IME1 vmadot sequence expects - add ime1::gemm_kernel_i8i8, an int8 x int8 IME1 kernel with a single-row and a 4-row A path; the 4-row path loads each B panel once and reuses it across 4 rows of A - add quantize_a_4row_i8 for the 4-row activation quantization - wire both into forward_mul_mat and the repack factory for Q8_0 - docs: mark Q8_0 as supported on X60 Correctness was checked against a quant-exact integer reference for K = 32 up to 4096, with a max relative error of about 1e-6, and by checking that generation stays coherent across several prompts. Tested on Milk-V Jupiter (SpacemiT X60), Bianbu 2.1.1, gcc 14.2, with Qwen2.5-0.5B-Instruct Q8_0. llama-bench -t 4 under taskset -c 0-3, 5 repetitions on an idle board: pp128 goes from 10.70 to 93.87 t/s. Q4_0 is unchanged at 106.40 -> 107.51 t/s, as expected since this does not touch that path. * ggml-cpu : move q8_0_16x32 decl to IME1 section * ggml-cpu : align q8_0 IME1 kernel assignments
…rotation metadata (ggml-org#28498) * kv-cache: save exact KV rotation metadata, reject restoring mismatched rotation * tests : move the state rotation test to test-save-load-state the test is now part of the save/load test matrix and runs against every model under test, like the rest of the suite it probes the KV cache type combinations supported by the model and treats models that do not use attention rotation as passing vacuously Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : skip unsupported KV caches --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
…org#29958) * llama : fix unexpected graph reallocation in the k-pool models Both k-pool models built a graph shape that depends on state the full-context reserve cannot know: - qwen4exp branched on inp->cache_safe, which turns false as soon as llama_memory_seq_cp shares cells (e.g. batched-bench -pps): the QSA layers swapped scatter+gather for fill+concat and dropped the new_pool_rep leaf, so the decode graph had 12 fewer nodes than the reserved one - glm5-next branched on gather = n_tokens <= 16 && n_kv > n_sel, so the TG decode built the gather shape (7564 nodes) while the last reserve, the PP one, had the dense shape (7762 nodes) Either mismatch forces a decode-time re-reserve that drops the worst-case sizing and bakes in the current state, so the next state growth (n_pool, n_kv, n_new) needs more room at an unchanged graph size and aborts under GGML_SCHED_DEBUG_REALLOC=1. Reproduce with, e.g.: GGML_SCHED_DEBUG_REALLOC=1 ./bin/llama-batched-bench \ -hf ggml-org/GLM-5.3-Flash-GGUF:Q2_K -npp 2500 -ntg 32 -npl 1,2 \ -c 32768 -pps -kvu Always scatter+gather the pooled keys, and pick gather from context constants only: n_ubatch bounds every ubatch, top_k + kpool - 1 bounds n_sel. Every graph of a context then shares one shape, which the reserve covers, and the dense path measured faster than the gather path at 2.5k and 16k context. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * llama : drop the unused k-pool cache_safe graph API The k-pool graphs no longer branch on cache_safe, so nothing reads get_kpool_cache_safe() or the conditional new_pool_rep any more: both models always pass the scatter target, which set_input_kpool now requires instead of merely preferring. Also drop the cache_safe copy in kpool_build_sizes(), a sizes-only helper. The layout and state flag itself stays, it still decides which pools a layout with shared cells must re-pool. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * tests : add a shared-seq graph reserve regression test Decode a prompt into seq 0, share its cells with seq 1 via llama_memory_seq_cp (what llama-batched-bench does for -pps), then keep decoding both sequences. For the k-pool models sharing clears cache_safe, which changes the graph topology while the pools keep growing, so a scheduler that re-reserves with the current state instead of the worst-case one aborts under GGML_SCHED_DEBUG_REALLOC=1. The test registration sets that flag, and the test aborts on both k-pool models before 2220411. kimi-linear and minimax-01 are skipped: they reserve the final pp graph with n_seqs = 1 (see [TAG_RESERVE_DIAG_DECAY] in llama-context.cpp), so every multi-seq graph has a different layout and re-reserves by design. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * cont : add TODOs * cont : fix comment * cuda: match the moe weighted reduction on empty ubatches ggml_cuda_match_moe_weighted_reduction rejected tensors with zero rows. A ubatch without outputs shrinks the last layer to zero rows through inp_out_ids, so graph_optimize dropped its alloc dep there and the scheduler graph lost one node compared to the reserved one. The scheduler then re-reserved at the size of that ubatch, and the next ubatch with the same node count but larger tensors aborted under GGML_SCHED_DEBUG_REALLOC=1. The compute loop already skips empty nodes before trying any fusion, so the guard only made the alloc deps depend on the row count. * tests: build the rollback test only where internal symbols link The shared-seq case calls llm_arch_from_string, which libllama does not export through LLAMA_API, so linking test-recurrent-state-rollback fails on Windows with shared libraries. Its build now sits in the NOT WIN32 OR NOT BUILD_SHARED_LIBS block, next to test-llama-archs and the test registration it already lives under. * tests: skip archs by name in the shared-seq reserve test The skip of kimi-linear and minimax-01 went through llm_arch_from_string, which libllama does not export through LLAMA_API, so the test could not link on Windows with shared libraries. It now compares the general.architecture string directly, and the test builds on every platform again. --------- Co-authored-by: Pascal <admin@serveurperso.com>
* vulkan: sparse flash attention for quantized K/V Assisted-by: Claude * vulkan: single-scan sparse FA index compaction The compaction ran one workgroup per mask row and walked the row in BLOCK_SIZE chunks, with a workgroup scan per chunk. For decode that is one workgroup doing KV/1024 barrier-bound iterations, so at 128k cells it cost more than the sparse attention it feeds. Split the row into contiguous segments instead: one per subgroup with ballot counting over coalesced loads, or one per thread without subgroups. A single scan over the segment counts then gives each segment its output offset. The index list stays ascending.
* server: reject partial media truncation * server: keep only the keep_first fix Drop the mtmd test helper change, which no longer builds since clip_image_f32_batch stores its entries by value, and drop the vision test: no test fixture reaches a cut between two adjacent media chunks with a reused cache (tinygemma3 uses SWA and wraps images in text tokens, tinyopenjev and small-test are recurrent), so the test passed or failed independently of the fix. --------- Co-authored-by: Pascal <admin@serveurperso.com>
* ggml_cuda: optimize accumulation in mmq_vec_dot_fp4_fp4_mma for better performance * remove whitespace * fix: correct indentation in mma_block_scaled_fp4 loop
* server: support vision input for Clef * move input_attn_causal to private * extend old server_batch::embd * server_batch::token::pos to multi dim * nits * fix abort * fix img tokens cap * fix yield_to_queue mutate data
…tatistics (ggml-org#27990) * webui: Use toLocaleString() format consistently across chat message statistics * Add missing semicolon * ui: respect lint
Assisted-by: pi:llama.cpp/Qwen3.8-27B
* skill : add tools/ui review checklist to code-review skill Adds a WebUI (tools/ui) section distilled from a year of review feedback on tools/ui PRs: constants/enums placement, layering (stores/services/utils/hooks), DRY and dead code, Svelte 5 idioms, comment hygiene, streaming perf, frontend security, and PR hygiene (static build, demos, scope discipline). Scoped to tools/ui only. * skill : expand tools/ui constants/enums ruleset Replaces the short constants bullet list with a decision ruleset distilled from review outcomes: when a literal becomes a string enum vs an enum-keyed map vs a declarative config vs a grouped as-const object vs a scalar constant, naming and suffix conventions (_REGEX, _MS, _PX, semantic separators), placement and composition rules, and what legitimately stays inline. * skill : drop stale static-build-commit requirement for tools/ui Static build output is no longer stored in the repo - the UI is built by CI and bundled at release time. Keep the requirement that npm run build succeeds and the UI build job is green. * skill : address copilot review of tools/ui checklist - scope the video-demo requirement to user-facing interactive/visual changes instead of every tools/ui PR - narrow the constants trigger to reused/domain-significant literals, keeping the inline boundary as the escape hatch - carve out independent literals for persisted/wire-contract test assertions so they don't become tautological via shared constants - allow simple presentation-local $derived in components; only shared business logic, orchestration and parsing leave the component - match the existing barrel documentation style instead of requiring per-export JSDoc everywhere - keep non-obvious why/rationale comments; move only examples and edge-case enumerations into tests - point at the actual interpolate-size/grid-template-rows transition pattern in CollapsibleRegion.svelte instead of calc-size() * skill : use spaces for nested list indentation EditorConfig CI requires space indentation for all files; the scalar-constants sub-bullets used tabs. Matches the 2-space nested list convention used in docs/.
* cuda, hip: GDN chunked kernel Ported-from: tetherto#257 Source-commit: 53f0027 Co-authored-by: Aman Gupta <aman.gupta@tether.io> * fix older builds * review fixes * cuda: enable GDN MMA on tested Ampere and Ada GPUs Enable SM86 and SM89 after RTX 3090 and RTX 4090 correctness and prefill benchmarks. Keep other untested architectures on the existing fallback. Assisted-by: Codex * cuda : use refactored swizzle API in GDN MMA * ggml : GDN only guarantees the final state when n_tokens > K * cuda : add short GDN MMA phase comments --------- Co-authored-by: gianni-cor <gianfranco.cordella@tether.io> Co-authored-by: Aman Gupta <aman.gupta@tether.io>
…erand moves Pass a 32-bit shared address to ldmatrix. With a generic pointer the compiler added and subtracted the shared window base for every load. Where K is the B operand of the KQ mma (wide tiles), load its 8x8 matrices in the register order the mma takes them, which removes the register moves that formed the operand pairs. The loaded data is the same. Assisted-by: Claude Opus 5.5
Add ggml_flash_attn_ext_set_native_quants(), a per-node FlashAttention flag that lets a backend read a supported quantized K/V cache in place instead of casting it to F16 first. llama sets it from the new context parameter flash_attn_native_quants, exposed as --flash-attn-native-quants and in llama-bench as -fanq. It is off by default, and nothing reads it yet. Assisted-by: Claude Opus 5.5
Piggidragon
force-pushed
the
llama/fattn-native-quants-flag
branch
from
October 14, 2026 18:39
376e5aa to
5bce922
Compare
Piggidragon
changed the title
Piggidragon
force-pushed
the
llama/fattn-native-quants-flag
branch
from
October 14, 2026 22:48
5bce922 to
c45d165
Compare
Piggidragon
changed the title
Piggidragon
force-pushed
the
llama/fattn-native-quants-flag
branch
from
October 14, 2026 22:54
c45d165 to
2c1d5db
Compare
…kernel With the opt-in flag set, the MMA FlashAttention kernel fills its half2 K/V tiles from the quantized cache instead of a transient F16 copy, for D=256 at an even GQA ratio or one above 4, and D=512 above 4, at every query batch size. q4_0 and q8_0 kernels are always compiled, q4_1, q5_0 and q5_1 when GGML_CUDA_FA_QUANTS selects their pair. Where cp.async exists the raw rows are copied to shared memory while the previous tile computes and converted from there. At D=256 K and V get separate half2 tiles, so converting one overlaps the MMA that reads the other. At D=512 the raw chunks are copied two chunks ahead. Turing reads the blocks directly. An even GQA ratio that is not a multiple of 8 takes 2 heads per tile for large query batches. Small batches take narrower tiles: 4x8 for 2 to 4 tokens at 8 heads per tile, 1x8 for a single token, 4x2 for up to 4 tokens at GQA 2 and 4. The route is checked before the vector kernel, so decode and speculative-decoding steps read K/V once per KV head, without a cast. The conversion matches the F16 cast bit for bit for q4_0, q5_0 and q8_0. q4_1 and q5_1 compute q*d + m in half2, which can differ in the last bit. On sm_86, sm_89 and sm_120 the kernels sum KQ in half over 32 values of D and in float across them, since float accumulators run at half rate on GeForce cards. Assisted-by: Claude Opus 5.5
NATIVE_QUANT_EQUIVALENCE computes the attention over a quantized cache and over an F16 copy of it in one graph, requires the two backend results to agree within an NMSE of 1e-4 (the route may sum in another order and partly in half), and checks which path the CUDA dispatcher took through its native-launch counter. The cases cover every tile shape, decode and small query batches, two sequences and the mask scan. Perf cases time the route against the F16-casting path and the vector kernel on every row it covers. Assisted-by: Claude Opus 5.5
scripts/fattn-native-inventory.py reads the built libggml-cuda back and fails on a missing, unexpected or duplicated native kernel for a given GGML_CUDA_FA_QUANTS. The workflow compiles the default, one-extra-pair, all, Turing and FA-off configurations and runs the equivalence cases on a self-hosted GPU runner where one is registered. Assisted-by: Claude Opus 5.5
Where the route applies, how the cache types are compiled, how the K/V tiles are staged, converted and multiplied, and what it was validated and measured with on an RTX 3060 and an RTX 4070: kernel times for prompt processing, decode and small query batches, llama-bench end to end up to 64k context, KL divergence against an F16 cache, and the compute buffer it saves. Assisted-by: Claude Opus 5.5
Piggidragon
force-pushed
the
llama/fattn-native-quants-flag
branch
from
October 15, 2026 00:45
2c1d5db to
8e959c4
Compare
A speculative-decoding step with one draft token has 2 query tokens. At 8 heads per tile they took the 4x8 tile, which leaves half of it empty; the 2x8 tile loads 64 K/V rows per stage and does half the arithmetic. It is 2-15% faster than 4x8 on an RTX 3060 and RTX 4070. An 8x2 tile for 5-8 tokens at GQA 2 was measured as well and left out: it was slower at short context. Assisted-by: Claude Opus 5.5
Assisted-by: Claude Opus 5.5
…=128 The route only had kernels for head sizes 256 and 512. Most models use D=128 (Llama 3, Qwen3, Mistral) and gpt-oss uses D=64; both now take the D=256 rows. At D=128 the rows of q4_0, q5_0 and q8_0 (72, 88 and 136 bytes) are not a whole number of 16-byte pieces, at D=64 not of 8-byte pieces. Rows are now staged in pieces of 16, 8 or 4 bytes, the largest size that divides both the row and its pitch, and the host checks the K/V alignment against that size. The conversion lane layout no longer assumes 8 or 16 blocks per row: with fewer than 4 blocks per pass the quarter warps take rows 8 apart, and the stage pitch keeps their reads on distinct banks. The D=256 kernels are unchanged within run-to-run noise. On an RTX 3060 and an RTX 4070, decode at D=64 and D=128 is 49-95% faster than the F16-casting path and prompt processing at 512-token batches 3-39% faster. Assisted-by: Claude Opus 5.5
Decode and small query batches are bound by reading K/V. At GQA 4 they took the 4x2 and 16x2 tiles, which read every K/V row twice per KV head, and GQA 3 had no native tile at all. Up to 8 query tokens at GQA 3 and 4 now take the 1x8, 2x8, 4x8 and 8x8 tiles: the unused heads cost arithmetic only. Larger batches take tiles of 2 heads, which now also covers GQA 3. GQA 1 keeps the F16 path, where the vector kernel already reads K/V once per KV head. At D=128 and GQA 4, 1 to 8 query tokens are 3-49% faster on an RTX 4070 and 18-50% faster on an RTX 3060. Assisted-by: Claude Opus 5.5
At D=512 one or two query tokens took the 4x8 tile, which leaves three quarters or half of it empty and made decode compute bound: q4_0 at 64k context read K/V at less than 40% of the DRAM bandwidth on an RTX 4070. The 1x8 and 2x8 tiles make it 7-38% faster on an RTX 3060 and an RTX 4070. Turing batches K of D=512 in 96 half2 below 32 columns, which the synchronous loader cannot split into quant blocks. It keeps the 4x8 tile, and the kernel compiles no device code for these shapes there. Assisted-by: Claude Opus 5.5
The help text said the output is unchanged. The route sums in another order and, on sm_86, sm_89 and sm_120, partly in half precision, so the results differ from the F16 path in rounding. It also applies per GQA ratio now. Assisted-by: Claude Opus 5.5
…oute Update the route table, the kernel inventory and the conversion description, regenerate the kernel timings on an RTX 3060 and an RTX 4070, and add llama-bench and KL divergence results for two D=128 models. The Turing code is now checked with a 75-virtual build on the RTX 4070. Assisted-by: Claude Opus 5.5
The staging modes of the quantized-native route assumed 100 KiB of shared memory per SM, which is what sm_86, sm_89 and sm_120 have. sm_80 and sm_87 have 164 KiB, sm_90 and sm_100 228 KiB, Turing 64 KiB. The plan now takes the value of the architecture the device code is compiled for, on the host and on the device alike. With more of it, q8_0 keeps its mask tile in shared memory instead of registers at D=128 and D=256. Code for sm_80 also runs on sm_86 and sm_89 cards, which have less. The host now computes the shared memory a kernel needs (moved out of the case function so the route can ask for it) and keeps the F16 path where a block would not fit the device. Assisted-by: Claude Opus 5.5
At D=512 the 1x8 tile targets two blocks per SM, and its chunked staging needs 59-75 KiB per block. On 100 KiB parts two blocks do not fit, so it fell back to the synchronous loaders, which leave the global memory latency on the critical path. Chunked mode is now taken whenever one block fits. Decode with one query token at D=512 and 64k context is 9-21% faster on an RTX 3060 and an RTX 4070. On the 4070, q8_0 at 4k and 16k context gets slower in the benchmark, where the whole cache stays in its 36 MB L2; in a model the other layers evict it. Assisted-by: Claude Opus 5.5
GeForce cards run MMAs with float accumulators at half rate, which is why the native route sums KQ partly in half on sm_86, sm_89 and sm_120. The professional and data center cards with the same chips (RTX A6000, A40, L40S, RTX 6000 Ada, ...) run both at full rate: they took the precision cost without the speed gain. The kernels now have a half and a float variant, and the host takes the half one on devices whose name contains "GeForce", where the compiled architecture has it. sm_75 joins the list for the RTX 20 series; the T4, Quadro RTX and TITAN RTX take float. The float variant adds about 8 MB to a build for sm_86 and sm_89 with every type selected. The inventory script now expects both variants per case. Assisted-by: Claude Opus 5.5
D=512 had native kernels for q4_0 and q8_0 only, so the other cache types took the F16-casting path there. The chunked staging of D=512 is generic, so the extra-tier types now get the same four tiles as q4_0 and q8_0, compiled when GGML_CUDA_FA_QUANTS selects their pair. Every cache type with a native loader now has kernels at every head size the route supports. Assisted-by: Claude Opus 5.5
… and D=512 for every type Half KQ accumulation is now a kernel variant the host selects for GeForce cards, the staging plans with the shared memory of the compiled architecture, the D=512 1x8 tile stages in chunks, and q4_1, q5_0 and q5_1 have D=512 kernels. Regenerate the D=512 kernel timings, and record the checks of the sm_80 and Turing code on the RTX 4070. Assisted-by: Claude Opus 5.5
On Ampere and older, launch_fattn splits the work with stream-k only when whole tiles would fill less than 75% of the last wave. A native kernel that fits one block per SM then leaves whole SMs idle for the rest of that wave. Kernels that read a quantized K/V in place now always take stream-k when one block fills an SM; Ada and newer took it already. D=512 prompt processing on an RTX 3060 is 0.4-7% faster with every type. Assisted-by: Claude Opus 5.5
Prompt processing at D=128 and D=256 now takes 64x2 and 16x8 tiles on Ampere and newer: one block of 8 warps per SM instead of two blocks of 4 warps of 64 columns. Every K/V tile is converted once for twice the query columns, which halves the conversion work, and with one block per SM q8_0 keeps its mask tile in shared memory. 17-32 query tokens keep 32x2, up to 8 tokens at 8 heads per tile keep 8x8. D=64 keeps 64 columns, where the wide tile brought nothing. Against the 64-column tiles, prompt processing is on average 4.5% faster on an RTX 4070 (q8_0 about 10%, up to 27% at 64k context) and 2.7% on an RTX 3060 (q8_0 about 6-7%). q4_0 on the 3060 is about even, -5% to +7%. Assisted-by: Claude Opus 5.5
Add the 64x2 and 16x8 tiles to the route table and the implementation, the stream-k rule for kernels with one block per SM, the case count and the library size, and regenerate the D=128, D=256 and D=512 prompt processing timings on both GPUs. Assisted-by: Claude Opus 5.5
Owner
|
Tested head
Example Q8_0 KV, head dimension 256, GQA 6, KV length 16384:
Two prefill regressions were D64 / GQA8 / 4096 KV / 512 queries: Q4_0 took 1.52% longer and Q8_0 took 6.96% longer. These are one-sweep Selected sanitizer checks: seven case executions each passed memcheck and synccheck. Native-MMA-filtered racecheck and the two full wide-prompt racechecks were clean. The five small full-graph racechecks reported four warnings in F16 vector reference kernels. Their source blob is unchanged from the official base, but a pristine-base runtime control was not run, so those warnings remain unresolved reference-path evidence rather than a clean full-graph racecheck claim. Numerical bounds and tests were not weakened. The latest functional changes add 128-column prompt tiles and broaden one-block stream-K admission. SM120 already followed the Ada-and-newer stream-K rule; that admission change is mainly relevant to earlier Ampere devices. Local composition with moe-cache The integrated build passed Matched serving configurations, one off/on pair each, same Q8 target/draft KV, 8194 input / 512 output tokens, generic source execution, cache64, fixed miss split, no profile/adaptation, mmap and 4GiB host pin budget:
All six unprofiled requests completed with exact accounting, zero server/source failure counters and clean teardown. Outputs first diverged at indices 38, 96 and 460 respectively. Dense MTP acceptance changed from 355/418 to 352/424; RCO from 317/387 to 318/382. These are observed rates with different output/effective work, not isolated equal-work gains or repeat-trial confidence bounds. The dense MTP result does not demonstrate an uplift. RCO scope matters: its target trunk/verification uses QSA with Prefill was mixed too: dense no-MTP 1189.65 -> 1114.02 tok/s, dense MTP 1348.64 -> 1080.01, RCO 256.30 -> 455.20. The large RCO prefill difference is not attributable to native target attention, whose route is unchanged; loading/history effects were not isolated. No universal prefill gain is claimed. A separate short Nsight attempt hit the source executor's checked progress timeout after 22 output events. It drained and tore down cleanly. The preserved partial report exported after teardown, but has no GPU kernel activity table, so it does not establish actual model kernel launches or valid GPU timing. No timeout/ownership/numerical checks were weakened and no profiling run retried. Backend native-counter assertions passed; model eligibility is source-derived. Full kernel attribution and semantic quality remain unqualified. Local integration report and raw artifacts: Full raw commands, timings, source/binary hashes and tool caveats are preserved locally at Testing and this report were performed with Codex assistance at the repository owner's request. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Reads a quantized K/V cache in place in the CUDA MMA FlashAttention kernel instead of casting the visible attention window to F16 first. The transient F16 copy goes away. It covers
q4_0,q4_1,q5_0,q5_1andq8_0caches at head sizes 64, 128, 256 and 512 and every GQA ratio from 2. On Ampere and Ada the route is faster than the F16-casting path for prompt processing at the default 512-token batches, and several times faster for decode and the small query batches of speculative decoding. End to end, Qwen3.8-27B at 64k context decodes 31-39% faster, Qwen3-4B at 16k context 26-44%.Opt-in:
--flash-attn-native-quants(LLAMA_ARG_FLASH_ATTN_NATIVE_QUANTS),llama-bench -fanq 0,1to compare. Without the flag every FlashAttention node keeps the established path, and the output is byte identical to upstream.Note
This branch is based on ggml-org
master(62df1815b), not onllama/dev. Againstllama/devGitHub shows every upstream change as part of this PR; the base needs to be moved to a branch that tracks ggml-orgmaster. Earlier heads:c9c79b9f2(before the 128-column tiles),61749ebdd(D=256 and D=512 only,q4_0/q8_0at D=512),2c1d5db67(without the 8-column decode tiles and half KQ at D=512),5bce92203(first rebase onmaster),376e5aa58(before the rebase).Where it applies
The 128-column tiles (16x8, 64x2) exist at D=128 and D=256; D=64 keeps 8x8 and 32x2. Every head size covers all five types. Also required: Turing MMA or newer, no logit softcap, the same type for K and V, the GQA optimizations apply, K/V aligned to the staging copy size (a KV cache always is), no sparse K/V set (
n_kv_max == 0). GQA 1 keeps the F16 path: there the vector kernel already reads K/V once per KV head.q4_0/q8_0kernels are always compiled,q4_1/q5_0/q5_1whenGGML_CUDA_FA_QUANTSselects theirtype-typepair: 58 cases by default, 145 withall, each with a float and a half KQ variant.The route is checked before the vector kernel, so decode takes it too: the vector kernel reads the quantized cache once per Q head, the native kernels once per KV head. A tile of 8 heads reads every K/V row once per KV head; below 8 heads per KV head the unused ones cost arithmetic only, so small batches take 8 heads per tile from GQA 3 on.
How it works
flash_attn_ext_f16gainstype_K/type_Vtemplate parameters (F16 by default, so F16 instantiations are unchanged). Only the filling of the half2 K/V tiles and the KQ accumulation differ; softmax, VKQ and combine are shared.q8_0keeps its mask tile in shared memory. Native kernels that fit one block per SM always split their work with stream-k, also on Ampere, so a partial last wave does not leave whole SMs idle.Rows are copied whole, in pieces of 16, 8 or 4 bytes (the largest that divides the row: 8 for the 72-, 88- and 136-byte rows of
q4_0,q5_0andq8_0at D=128, 4 at D=64). The conversion reads the staged rows as 4-byte words; blocks of 18, 22 and 34 bytes are converted in even and odd passes so the misalignment is a compile-time constant. The 8 lanes of a quarter warp take 8 consecutive rows of one block, and with fewer than 4 blocks per pass the quarter warps take rows 8 apart; the stage pitch is chosen so the reads and the swizzled stores stay free of bank conflicts. A quant becomes half2 with one__byte_perminto0x64XX, a bias subtract and the scale.q4_0,q5_0andq8_0reproduce the F16 cast bit for bit,q4_1andq5_1up to the last bit.Shared memory. The staging plans with the shared memory of the architecture the device code is compiled for: 100 KiB on
sm_86,sm_89,sm_120, 164 KiB onsm_80,sm_87, 228 KiB onsm_90,sm_100, 64 KiB on Turing. Code forsm_80also runs onsm_86/sm_89cards; the host checks every kernel against the device and keeps the F16 path where a block would not fit.KQ accumulation. GeForce cards run MMAs with float accumulators at half rate. There the native kernels sum KQ in half over 32 values of D and in float across them, which saves up to 11% of kernel time at D=256 and 10-15% at D=512 on the 3060 and 4070. Summing all 256 values in half would save another 7-9%, but it moved one sensitive wikitext token from a KL divergence of at most 2.3 to between 5 and 18, so it was dropped. Professional and data center cards with the same chips run both accumulators at full rate, so this is a kernel variant the host selects per device: devices named "GeForce" take it on
sm_75,sm_86,sm_89andsm_120, every other device takes float.-DGGML_FATTN_QUANT_KQ_FLOATkeeps float everywhere for comparison.Shared MMA path (first commit, also speeds up the F16 path with identical results):
ldmatrixgets a 32-bit shared address, and K is loaded in the register order the KQ mma takes it as B operand.Results
RTX 3060 (
sm_86) and RTX 4070 (sm_89), CUDA 13.4.test-backend-ops perf -o FLASH_ATTN_EXT -p expected_native, native against this branch without the flag (cast or vector kernel included), negative = native faster, each value the faster of two runs for both paths.Prompt rows cover all five types at n_kv 1024-16384, the token rows
q4_0andq8_0at n_kv 4096-65536. The 128-column tiles made prompt processing another 4.5% (4070) and 2.7% (3060) faster on average,q8_0about 10% and 6%. The positive values are 2048-token batches at D=64 withq8_0,q5_0andq5_1: every query tile converts the K/V rows it reads again, while the cast converts them once, and small heads have the least arithmetic to hide that behind. At 512-token batches, the default-ub, D=64 is 3-13% faster. At n_kv=65536 theq8_0decode kernel reads K/V at about 90% of the DRAM bandwidth.End to end (measured before the 128-column tiles),
llama-bench, Qwen3.8-27B IQ3_XXS (D=256, GQA 6, 16 of 65 layers are attention), against upstreammaster, flag on:q8_0q4_0q8_0q4_0q8_0q4_0D=128 models on the RTX 4070,
-fanq 0,1, d16384:q8_0q4_0q8_0q4_0pp4andpp2stand in for speculative-decoding steps that verify 4 or 2 tokens. At depth 0 a generation step attends to at most 128 tokens and the native kernel is about 1% slower than the vector kernel there; Qwen3.8 shows no difference at depth 0. Without the flag the branch matches upstream.Memory: the compute buffer of Qwen3.8-27B on 4070 + 3060 at 262144 context drops from 4468 MiB to 2877 MiB with the flag.
All tables are in
docs/quantized-native-flash-attention.md.Correctness
test-backend-ops -o FLASH_ATTN_EXT3993/3993,-o GATED_DELTA_NET43/43,-o NATIVE_QUANT_EQUIVALENCE105/105 on both GPUs. The equivalence cases compute the attention over the quantized cache and over an F16 copy in one graph, require the two backend results to agree within NMSE 1e-4 and assert the route through the backend's native-launch counter. They cover every tile shape at D=64, 128, 256 and 512, all five types, GQA 2, 3, 4, 6, 8 and 16, attention sinks, decode, small batches, two sequences and the mask scan.75-virtualbuild runs the Turing code (synchronous loaders, Turing tiles, half KQ), an80-realbuild thesm_80code (float KQ, 128-column tiles, 164 KiB staging plan, also on the 3060): 105/105 each, andnsysshows the expected kernel variant in every build. This checks results, not speed on those cards.compute-sanitizerracecheck, synccheck, memcheck: no errors or warnings on the native kernels, including the 8- and 4-byte staging at D=128 and D=64.q8_0: mean KLD 0.00082 vs 0.00095, same top token 98.75% vs 98.82%. Llama-3.2-3Bq8_0: 0.00173 vs 0.00176. Qwen3-4Bq8_0: 0.00378 +- 0.00029 vs 0.00360 +- 0.00018, medians equal. Greedy output without the flag is byte identical to upstream; with it the text departs after 25-65 words and stays coherent.scripts/fattn-native-inventory.py: 145 declared cases with both KQ variants, all compiled, nothing else.Build cost: with all five types for
86-real;89-real,libggml-cuda.sogrows from 71.0 MB (upstream) to 86.4 MB, about 8 MB of it for the float variants that professional cards need and 3.6 MB for the 128-column tiles.Testing on a card not measured here (e.g. RTX 50,
sm_120)On a GeForce card, then rebuild with
-DCMAKE_CUDA_FLAGS=-DGGML_FATTN_QUANT_KQ_FLOATand repeat theperfrun: if the native rows are not slower with float accumulation, the half accumulation brings nothing on that architecture and it should come out ofggml_cuda_fattn_quant_kq_half_arch()andGGML_FATTN_QUANT_KQ_HALFinfattn-mma-quant.cuh.Commits
cuda : load swizzled ldmatrix tiles with shared addresses, without operand movesggml, llama : opt in to the quantized-native FlashAttention routecuda : read a quantized K/V cache in place in the MMA FlashAttention kerneltests : assert the quantized-native FlashAttention routeci, scripts : check the native FlashAttention kernel inventorydocs : describe the quantized-native FlashAttention routecuda : give two-token batches a 2x8 native FlashAttention tiledocs : add the 2x8 tile and two-token measurementscuda : extend the quantized-native FlashAttention route to D=64 and D=128cuda : give small query batches at GQA 3 and 4 tiles of 8 headscuda : give D=512 decode the 1x8 and 2x8 native tilescommon : say that --flash-attn-native-quants changes roundingdocs : describe D=64, D=128, GQA 3/4 and D=512 decode on the native routecuda : plan native FA staging with the shared memory of the architecturecuda : stage the D=512 1x8 native tile in chunks at one block per SMcuda : select half KQ accumulation per device, not per architecturecuda : give q4_1, q5_0 and q5_1 the D=512 native kernelsdocs : describe per-device half KQ, the staging plan per architecture and D=512 for every typecuda : take stream-k for native FA kernels with one block per SMcuda : give native FA prompt processing tiles of 128 columnsdocs : describe the 128-column prompt tiles and stream-k for native FAAI assistance: the rebase, the kernels' staging and routing, the generalization to D=64/128/512 and the profiling with ncu/nsys and the measurements were done with Claude Opus 5.5. Each commit carries
Assisted-by:.🤖 Generated with Claude Code