cuda: optimize native quant FlashAttention integration (follow-up to #50) - #55
cuda: optimize native quant FlashAttention integration (follow-up to #50)#55GenerelSchwerz wants to merge 12 commits into
Conversation
Lets a backend read a quantized K/V cache in place instead of materializing an F16 copy. Plumbed through as the runtime option --flash-attn-native-quants; no backend acts on it yet. Assisted-by: Claude Opus 5
The native loaders dequantize the current tile straight into the shared half2 tiles the existing MMA body consumes, so the F16 window copy is never allocated. Each loader reproduces its type's F16 cast path bit for bit. K and V become template parameters of the existing kernel; a sentinel V type selects the loader from a runtime argument, so the mixed pairs cost one kernel per K type instead of one per pair. Mixed pairs and the extra cache types are compiled only with GGML_CUDA_FA_ALL_QUANTS, which is the build that can reach them. The cache-type inventory lives once, in fattn-mma-quant-types.h. The route predicate, the extern declarations, the host dispatch, the device loader selection, the generator and CMake's source filter all read it. Assisted-by: Claude Opus 5
The type list is expanded from the same manifest the kernels are generated from, so a new cache type cannot become routable untested. In a build without GGML_CUDA_FA_ALL_QUANTS the extra-tier cases check the materializing fallback instead. Assisted-by: Claude Opus 5
The failure mode here is a build/inventory failure, not a wrong result: an extern-macro mistake can instantiate hundreds of kernels in one translation unit while every runtime test passes. The checker reads the built library back and compares it against the type manifest. Assisted-by: Claude Opus 5
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Follow-up plan: FlashInfer-inspired quantized KV-cache accelerationThis is a proposed direction for extending the native quantized FlashAttention work beyond the current CUDA integration. The goal is not to import FlashInfer or copy its kernels. The goal is to adopt the general design principles that make its attention implementation effective while preserving llama.cpp's backend independence and existing fallbacks. The plan is intentionally model-neutral and device-neutral. Dispatch must be based on tensor geometry, data type, cache layout, runtime state, and backend capabilities. It must never depend on model names. 1. ObjectiveImprove quantized KV-cache performance throughout llama.cpp by:
Replacing llama.cpp or introducing FlashInfer, PyTorch, or another runtime dependency is out of scope. 2. Current starting pointPR 55 already establishes useful infrastructure:
The main remaining limitations are:
Exploratory Ampere/Ada measurements show that the existing native loaders can provide large kernel-level gains for several unrelated geometries, including D64/GQA8, D128/GQA16, D256/GQA2, and D256/GQA6-8. The measurements also show that a native path is not universally faster, and at least one D256/GQA8 Q8_0 case requires further correctness and memory-safety investigation. These observations should motivate a planner, not become universal thresholds. 3. Target architectureThe common layer should describe the attention problem. Each backend should remain responsible for selecting and launching its implementation. Common problem descriptionThe description should contain enough information for backend selection without encoding backend-specific tile details:
Backend planA backend-owned plan should select:
Plans should be cached by device and problem geometry. Normal inference should use deterministic shipped rules. It should not benchmark kernels in the user hot path. 4. Benchmark and validation matrixUse the existing At minimum, cover:
Record:
Benchmark geometries should be derived from model architectures, but results and routing must be keyed by geometry rather than model identity. 5. Stabilize the CUDA native pathThe first implementation phase should reuse the native loader design already present in this PR. Required work:
The initial dispatch table should be keyed by:
Thresholds from one GPU or model must not be applied globally. They should be generated from a representative hardware matrix and reviewed before being shipped. 6. Split-KV decodeLong-context decode can expose too little parallel work when the grid is based only on sequence and KV head. Add optional KV partitioning:
The split and merge algorithm can be common conceptually, while launch geometry and kernels remain backend-specific. Reference: FlashInfer decode scheduler. 7. Fused KV-cache appendWhere graph structure permits it, avoid separate operations for producing K/V, applying RoPE, quantizing, and writing the cache. Potential backend operations:
The existing unfused path should remain the correctness reference and portable fallback. Reference: FlashInfer fused RoPE, quantize, and append implementation. 8. Paged KV and cache layoutExtend existing llama.cpp KV-cache infrastructure rather than replacing it. Goals:
Paged access should only be selected when its address-calculation cost is lower than materialization or copying for the current geometry. 9. Backend rolloutOnly the problem description, planning semantics, split-state representation, and test matrix should be shared. Device kernels must follow each backend's native execution model. CUDA
HIP
SYCL
Metal
Vulkan
CPU
10. Correctness and safety requirementsEvery optimized route must be compared against the current reference path. Required coverage:
Performance tests alone must not be treated as correctness evidence. 11. Binary-size controlTemplate specialization must remain bounded. Controls:
The optional all-quant build should not dictate the size of the default build. 12. Tuning policyUse three layers:
Runtime inference should not perform surprise benchmarking. A local opt-in autotuner can be considered later, but it should not be required for correct or fast default behavior. 13. Multi-GPU behaviorAttention optimization does not remove layer-split communication costs. The implementation should:
14. Proposed change sequenceKeep each step independently testable and reviewable:
15. Acceptance criteriaThe project is successful when:
This plan uses FlashInfer as an architectural reference, especially its decode planning, KV splitting, paged-cache handling, and fused cache append. It does not propose source-level kernel reuse. AI-assisted draft prepared with OpenAI Codex. The implementation details and thresholds require maintainer review and independent validation on representative hardware. |
GenerelSchwerz
left a comment
There was a problem hiding this comment.
Review outcome: changes required before merge
I do not think this branch should be merged or used as the initial release in its current form. The native-quant idea is valuable, and the direct-load approach is worth preserving, but this revision combines an unvalidated automatic route, a large mostly unreachable specialization matrix, stale user documentation, a broken focused test, and a confirmed HIP regression. The branch is also now far enough behind llama/dev that a small replacement PR will be easier to validate than repairing all 102 changed files in place.
Review scope: head b4c43c047f26674a5d7ae8b854cd82f4c0e34ad0, original base 01b141fc8969fbcb876c039dcfa268981918da81, 12 commits, 2,670 additions, 99 deletions, and 102 changed files. I reviewed the CUDA dispatch and loaders, generated instances, build selection, inventory checker, backend tests, workflow, documentation, CI results, and relevant upstream FlashAttention reviews. I did not run GPU workloads locally, so the runtime conclusions below are based on the submitted measurements and CI artifacts.
Blocking issues
-
The change breaks HIP source selection. The existing HIP build glob includes every
fattn-mma*.cufile, so it compiles the new CUDA-native quant instances (HIP source selection). The resulting build fails in the new Q4 loader's rawhalf2initialization (loader; failed HIP job). MUSA has the same broad glob and unnecessarily compiles the native CUDA instances (MUSA source selection). The new sources must be explicitly CUDA-only. -
The new workflow can execute pull-request payloads on a persistent self-hosted runner. It triggers on
pull_request(trigger), selects a self-hosted CUDA runner using only a repository variable, and checks out the PR ref before building it (job). A variable gate does not make untrusted PR code safe. Use an ephemeral isolated runner, restrict execution to trusted refs/manual dispatch, or remove the job. -
The focused native-route test currently selects zero native cases and can report success. The workflow filters for
native_quants=1and sets a removed verbose environment variable (workflow), while the final test identity containsexpected_native=instead (test identity). The filter removes nonmatching cases (filter), and the harness returns success when bothn_okandtests_runare zero (exit condition). The focused job was skipped in the submitted run, so there is no CI runtime coverage either. Make the job assert a nonzero expected count for both native and fallback routes. -
Automatic dispatch enables a geometry that this PR already identifies as requiring memory-safety investigation. The roadmap records an unresolved D256/GQA8 Q8_0 correctness and memory-safety concern, including graph/workspace reuse across changing geometries (roadmap). The current profitability function still selects device-resident D256/GQA8 Q8_0 when the query batch is greater than four and KV length is at most 512 (gate). The explicit equivalence matrix covers GQA6, GQA2, and D512/GQA16, but not D256/GQA8 (matrix). Disable that route until it has dedicated numerical, graph-replay, workspace-reuse, and Compute Sanitizer coverage.
-
The documentation describes an interface and fallback contract that no longer exist. It advertises
--flash-attn-native-quants,-fanq,LLAMA_ARG_FLASH_ATTN_NATIVE_QUANTS,GGML_CUDA_FATTN_NATIVE_VERBOSE, an opt-in graph flag, and a fallback warning (documentation; build documentation). The final code removed that plumbing and now selects the route automatically fromsupported && profitable(selection). Decide on one contract and make code, tests, workflow, benchmarks, and documentation agree before release. -
The generated matrix is much larger than the reachable route table. The generator emits D64, D128, and D256 kernels for every enabled type, runtime-V variants, and D512 for all five types (generator). The profitability gate rejects every mixed pair, every D64/D128 case, and D512 types other than Q4_0/Q8_0 (gate). At least 403 of 485 all-quants symbols and 64 of 98 default symbols are therefore unreachable through this dispatcher. The documented result is +22.24% for the default CUDA library and +104.74% for all-quants (size table). Compile only the exact release route table.
-
The inventory checker cannot prove that exactness. It stores geometries in sets, which discards duplicates, and its positive expectation is only
somerather than an exact tuple set (checker). It also does not reject every unexpected symmetric or mixed symbol. Generate an exact expected multiset from the release manifest and fail on missing, unexpected, or duplicate symbols. -
The performance evidence does not validate the final automatic route. The table reports manual Off/On results for D256/GQA6 Q8_0 at depths 4096 and 16384 (results). The final revision no longer has that switch, and the submitted profitability gate falls back for this device-resident Q8 geometry above KV 512. Those measurements therefore do not show the code path selected by the final dispatcher. Repeat them with explicit route assertion after the route table is final.
-
The hardware gate is broader than the validation.
ampere_mma_available()accepts NVIDIA architectures from Ampere onward, including Hopper and Blackwell builds (definition), but the submitted validation covers only sm86 and sm89 (hardware). Restrict v1 to validated families or provide correctness and profiling results for every enabled family.
Recommended path
Start a replacement branch from current llama/dev and port only the smallest proven slice. This branch is conflicting, and the target has moved substantially through the same FlashAttention files, including merged XOR-swizzle and sparse-FA changes. A clean replacement will make it possible to reason about the final shared-memory layout and route selection instead of preserving obsolete integration work.
For the initial release, define the goal narrowly: remove transient F16 KV materialization for one proven long-context CUDA prefill configuration without changing numerical behavior or slowing fallback workloads.
Suggested v1 boundary:
- CUDA only.
- NVIDIA Ampere and Ada only, until newer families are tested.
- Device-resident K and V.
- Symmetric Q8_0 K/V only.
- One measured head dimension and GQA configuration, most plausibly D256/GQA6.
- Prefill only, with explicit query-batch and KV-length bounds.
- Existing materialized implementation as the fallback.
- No new public option until the route is stable; a developer-only build gate is sufficient for the first experimental revision.
Leave mixed K/V, Q4_1/Q5_0/Q5_1, D64/D128/D512, host-resident KV, decode, split-KV, paged cache, fused append, a common planner, multi-GPU tuning, and other backends for separately measured follow-ups. Q4_0 is a reasonable second PR once it meets the same evidence bar.
Before enabling the first route, require:
- Exact native-versus-materialized equivalence for the supported geometry and explicit fallback assertions for unsupported cases.
- KV lengths around loader/tile boundaries and long-context values; query-batch boundaries; masks, sinks, views, offsets, strides, cache shift/defrag, and changing shapes in one process.
- CUDA graph capture/replay and workspace reuse across shape transitions.
- Compute Sanitizer memcheck, initcheck, racecheck, and synccheck where applicable.
- NCU kernel metrics and NSYS end-to-end traces with the actual selected route recorded in every run.
- Repeated base-versus-branch measurements with dispersion, plus binary-size and compile-time budgets.
- A focused CI test that fails if it executes zero native cases or zero fallback cases.
- Updated
docs/ops.mdand applicabledocs/ops/*.csvafter the final operation test matrix is established. - A complete PR description covering the problem, exact support table, fallback behavior, correctness, performance, limitations, and AI disclosure.
This narrower direction matches upstream review practice. Reviewers asked PR 25635 to use the simplest possible implementation and limit support to one exact tile where appropriate (design review, scope review). PR 20998 was asked to compile only the specialization values actually used (discussion). The broader quantized-FA PR 25312 was replaced with a cleaner, split implementation after the author recognized the need for per-type and per-hardware evidence (discussion, closure). PR 24526 also shows why exact route tests matter: an apparently plausible D512/GQA specialization produced incorrect attention because the selected configuration was wrong (result).
The direct quantized tile loaders can be retained, but the first mergeable unit should contain only the route that has complete evidence. That gives later extensions a trustworthy baseline instead of making the initial release responsible for an entire speculative framework.
Review performed by OpenAI Codex at the repository owner's request.
| # repository variable CUDA_SELF_HOSTED_RUNNER is set to 'true'; otherwise it | ||
| # is skipped, which is reported as skipped rather than as a pass. | ||
| if: ${{ vars.CUDA_SELF_HOSTED_RUNNER == 'true' }} | ||
| runs-on: [self-hosted, linux, x64, cuda] |
There was a problem hiding this comment.
Blocking: this job is triggered by pull_request, checks out the PR payload, and then builds and runs it on a persistent self-hosted CUDA runner. CUDA_SELF_HOSTED_RUNNER is only a repository variable; it does not establish that the checked-out code is trusted. Please move this to an ephemeral isolated runner or restrict it to a trusted ref/manual workflow that never executes arbitrary fork payloads.
| env: | ||
| GGML_CUDA_FATTN_NATIVE_VERBOSE: '1' | ||
| run: | | ||
| ./build-test/bin/test-backend-ops -b CUDA0 -o FLASH_ATTN_EXT -p "native_quants=1" |
There was a problem hiding this comment.
Blocking: native_quants=1 is no longer present in test_flash_attn_ext::vars(); route cases now expose expected_native=. This filter therefore removes every intended case. The harness also returns success for zero tests because n_ok == tests_run == 0. Please filter on the current identity and assert the exact nonzero native and fallback case counts so this job cannot pass vacuously. GGML_CUDA_FATTN_NATIVE_VERBOSE above is stale as well.
| #if defined(__CUDA_ARCH__) | ||
| static __device__ __forceinline__ half2 fattn_mma_q4_centered_half2(const uint32_t q, const uint32_t selector) { | ||
| const uint32_t bits = __byte_perm(q, 0x64646464, selector); | ||
| const __half2_raw raw = { (uint16_t) bits, (uint16_t) (bits >> 16) }; |
There was a problem hiding this comment.
Blocking cross-backend regression: the existing HIP CMake glob (template-instances/fattn-mma*.cu) picks up the new quant instance files and HIP rejects this raw __half2 initialization with narrowing diagnostics. The HIP and HIP-quality jobs fail on this line. MUSA has the same broad glob and compiles these CUDA-only units unnecessarily. Please exclude the native-quant instance family from HIP/MUSA source discovery rather than adapting this CUDA loader to backends that do not use the route.
| @@ -0,0 +1,192 @@ | |||
| # Quantized-native CUDA FlashAttention | |||
|
|
|||
| `--flash-attn-native-quants` lets the CUDA MMA FlashAttention kernel read a | |||
There was a problem hiding this comment.
Blocking documentation/behavior mismatch: the final revision removed --flash-attn-native-quants, -fanq, LLAMA_ARG_FLASH_ATTN_NATIVE_QUANTS, the graph opt-in, verbose route logging, and the fallback warning. ggml_cuda_fattn_native_applies() now selects automatically from supported && profitable. Please decide on the intended release contract and update the implementation, tests, workflow, benchmark method, and both documentation files together.
| if (K->type == GGML_TYPE_Q5_0) { | ||
| return K->ne[1] >= 16384; | ||
| } | ||
| return K->ne[1] <= 512; |
There was a problem hiding this comment.
Blocking: this fallthrough selects device-resident D256/GQA8 Q8_0 for Q->ne[1] > 4 and KV length <= 512, but the PR roadmap explicitly records a D256/GQA8 Q8_0 correctness/memory-safety issue under graph/workspace reuse. The explicit equivalence cases cover GQA6, GQA2, and D512/GQA16, not this route. Please reject GQA8 until dedicated numerical, graph-replay, shape-transition, and Compute Sanitizer tests pass.
| extra_k = type_k in FATTN_MMA_QUANT_TYPES_EXTRA | ||
| if extra_k: | ||
| f.write("#ifdef GGML_CUDA_FA_ALL_QUANTS\n") | ||
| for head_size in FATTN_MMA_QUANT_HEAD_SIZES: |
There was a problem hiding this comment.
This generates D64, D128, and D256 specializations for every enabled K type even though the profitability gate rejects D64/D128 entirely. The runtime-V matrix is also emitted although the gate rejects all mixed K/V pairs, and three extra D512 types are unreachable. By the current symbol inventory, at least 403/485 all-quants symbols and 64/98 default symbols cannot be selected. Please derive generation from the exact shipped route table; the current cost is documented as +22.24% default and +104.74% all-quants library size.
| if name is None: | ||
| foreign.add(k) # a runtime V loader for a type the manifest does not name | ||
| else: | ||
| runtime_v_geometries.setdefault(name, set()).add(geometry) |
There was a problem hiding this comment.
Using a set here discards duplicate instantiations, and the later expectation only checks some versus zero. That cannot prove the documented exact inventory or detect unexpected head/tile tuples. Please compare an exact expected multiset against the demangled symbols and fail on missing, unexpected, or duplicate symmetric and runtime-V entries, including foreign symmetric types.
| return test_case::max_err(backend); | ||
| } | ||
| GGML_ASSERT(route_out && route_out->buffer); | ||
| const bool native = ggml_backend_buffer_get_alloc_size(route_out->buffer, route_out) == ggml_nbytes(route_out); |
There was a problem hiding this comment.
This infers route selection from output allocation size, which is an incidental allocator property rather than a backend contract. A later workspace or allocation change can make the test report the wrong route without changing dispatch. Please expose a test-only CUDA route counter/query or another explicit backend signal, then require the focused run to observe the expected native and fallback counts.
| add_native_equivalence(type, 256, 4, {6, 1}, type == GGML_TYPE_Q5_0 ? 16384 : 512, 65, 1); | ||
| } | ||
| for (ggml_type type : { GGML_TYPE_Q4_0, GGML_TYPE_Q8_0 }) { | ||
| add_native_equivalence(type, 256, 8, { 2, 1}, 1024, 65, 1); |
There was a problem hiding this comment.
The explicit matrix covers ratio 6 here and ratio 2 on this line, but not D256/GQA8 even though the profitability function can select that geometry and the roadmap identifies it as the outstanding memory-safety case. Add the exact boundary cases only after fixing the issue: KV lengths around the tile boundary, changing shapes in one process, workspace reuse, and CUDA graph capture/replay. Until then, the dispatcher must reject it.
| |---|---:|---:|---:| | ||
| | 1 GPU, reserve compute buffer @ 16K ctx | 160.28 MiB | 138.28 MiB | -13.7% | | ||
| | 2 GPU, reserve compute buffer @ 64K ctx, per device | 633.13 MiB | 419.13 MiB | -33.8% | | ||
| | 1 GPU, pp2048 @ depth 4096 | 1194.31 t/s | 1195.39 t/s | +0.09% | |
There was a problem hiding this comment.
These Off/On measurements do not validate the final revision. The option was removed, and the final Q8 D256/GQA6 device-resident profitability gate returns fallback above KV 512, so the reported depth-4096 and depth-16384 results are not evidence for the route now selected by the code. Please rerun after finalizing dispatch and record an explicit route assertion with each benchmark, plus repeated measurements and dispersion.
No description provided.