Skip to content

cuda: optimize native quant FlashAttention integration (follow-up to #50) - #55

Open
GenerelSchwerz wants to merge 12 commits into
llama/devfrom
codex/pr50-native-quant-integration
Open

cuda: optimize native quant FlashAttention integration (follow-up to #50)#55
GenerelSchwerz wants to merge 12 commits into
llama/devfrom
codex/pr50-native-quant-integration

Conversation

@GenerelSchwerz

Copy link
Copy Markdown
Owner

No description provided.

Piggidragon and others added 12 commits August 28, 2026 12:38
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
@Piggidragon

Copy link
Copy Markdown

Follow-up plan: FlashInfer-inspired quantized KV-cache acceleration

This 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. Objective

Improve quantized KV-cache performance throughout llama.cpp by:

  • Reading quantized K/V directly inside optimized attention kernels.
  • Avoiding complete F16 cache-window materialization where a native path exists.
  • Selecting kernels according to head dimension, GQA ratio, query length, KV length, cache type, layout, and device architecture.
  • Treating decode and prompt processing as different workloads.
  • Adding split-KV execution for long-context decode where it improves occupancy.
  • Preserving portable and conservative fallbacks for every unsupported combination.
  • Supporting incremental implementation across CUDA, HIP, SYCL, Metal, Vulkan, and CPU.

Replacing llama.cpp or introducing FlashInfer, PyTorch, or another runtime dependency is out of scope.

2. Current starting point

PR 55 already establishes useful infrastructure:

  • Native quantized K/V loaders feeding the CUDA MMA FlashAttention body.
  • Q4_0 and Q8_0 in the default build.
  • Optional Q4_1, Q5_0, and Q5_1 support.
  • Generated cache-type inventory and instantiation checks.
  • D64, D128, and D256 kernel families, with limited D512 coverage.
  • Conservative fallback through an F16 materialized window.
  • Ampere/Ada capability gating.

The main remaining limitations are:

  • Small-query decode commonly selects the quantized vector kernel before the native MMA path is considered.
  • Profitability rules cover only a small set of geometries.
  • Support and profitability are currently too closely coupled.
  • Some shape/type combinations still require safety validation.
  • D512 decode and MHA/GQA1 do not have general native coverage.
  • Replaying several changing geometries can expose workspace or CUDA Graph reuse problems.
  • Other backends do not share a common problem description or planning model.

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 architecture

FlashAttention problem description
              |
              v
    Common capability query
              |
              v
       Backend planner
       /      |       \
 vector   native-tile  materialized
 kernel      kernel      fallback
              |
              v
      optional split-KV
              |
              v
       softmax-state merge

The common layer should describe the attention problem. Each backend should remain responsible for selecting and launching its implementation.

Common problem description

The description should contain enough information for backend selection without encoding backend-specific tile details:

  • Q, K, and V data types.
  • QK and V head dimensions.
  • Query length and KV length.
  • Query-head and KV-head counts.
  • GQA ratio.
  • Batch and sequence counts.
  • KV layout, views, offsets, and strides.
  • Host or device cache residency.
  • Mask representation.
  • Sliding-window parameters.
  • Attention sinks.
  • ALiBi and logit soft cap.
  • Requested precision.
  • Backend and device identity.
  • Graph-capture and graph-replay constraints.

Backend plan

A backend-owned plan should select:

  • Kernel family.
  • Native or materialized KV input.
  • Tile geometry.
  • GQA grouping strategy.
  • Number and size of KV splits.
  • Temporary workspace layout and size.
  • Reduction or merge kernel.
  • Graph compatibility.
  • Alignment and layout requirements.
  • Safe fallback.

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 matrix

Use the existing test-backend-ops infrastructure for development. Avoid adding a new permanent test file unless maintainers agree that it is necessary.

At minimum, cover:

  • Head dimensions: 64, 80, 96, 128, 192, 256, and 512.
  • GQA ratios: 1, 2, 4, 6, 8, 16, and 32.
  • Query lengths: 1, 2, 4, 16, 128, and 512.
  • KV lengths: 512, 1024, 4096, 8192, 16384, and 32768.
  • KV types: F16, BF16 where supported, Q8_0, Q5_1, Q5_0, Q4_1, and Q4_0.
  • Equal and mixed K/V types.
  • Contiguous tensors and cache views.
  • Single-sequence and batched execution.
  • MHA, GQA, and MQA.
  • Full and sliding-window attention.
  • Normal and non-trivial strides.
  • KV lengths immediately below and above quant-block, tile, page, and split boundaries.

Record:

  • Kernel latency.
  • Effective KV bandwidth.
  • Temporary memory.
  • Prompt and generation throughput.
  • Peak device memory.
  • Single-GPU and multi-GPU behavior.
  • Binary-size impact.
  • Graph capture and replay behavior.

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 path

The first implementation phase should reuse the native loader design already present in this PR.

Required work:

  1. Separate supported from profitable completely.
  2. Audit Q4/Q5/Q8 loaders at every compiled head dimension and GQA grouping.
  3. Validate non-power-of-two GQA ratios such as 3 and 6.
  4. Investigate D256/GQA8 Q8_0 memory safety before enabling that route.
  5. Validate changing context sizes and shapes in one process.
  6. Make workspace and CUDA Graph reuse safe for every plan transition.
  7. Allow native MMA consideration before the vector fallback when the selected plan says it is profitable.
  8. Add measured D64 and D128 decode routing.
  9. Add measured D256 routing by type, GQA ratio, query length, and KV-length range.
  10. Add a dedicated D512 decode implementation or explicitly retain the materialized fallback.
  11. Keep every unknown or unsupported combination on the current conservative path.

The initial dispatch table should be keyed by:

  • Compute-capability family.
  • Head dimension.
  • GQA ratio.
  • Query-length range.
  • KV-length range.
  • K and V types.
  • Host or device cache residency.
  • Layout and alignment class.

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 decode

Long-context decode can expose too little parallel work when the grid is based only on sequence and KV head. Add optional KV partitioning:

  1. Divide the active KV range into independently processed chunks.
  2. Produce partial output, maximum logit, and softmax sum for each chunk.
  3. Merge the partial states with the numerically stable online-softmax formula.
  4. Select the split count from KV length, head count, GQA ratio, batch, device occupancy, and workspace budget.
  5. Allocate all temporary state through the backend plan.
  6. Keep workspace addresses and maximum launch dimensions stable during graph replay.
  7. Disable splitting when launch and merge overhead exceed the expected benefit.

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 append

Where graph structure permits it, avoid separate operations for producing K/V, applying RoPE, quantizing, and writing the cache.

Potential backend operations:

  • RoPE plus K quantization plus cache append.
  • V quantization plus cache append.
  • Combined K/V append.
  • Block-scale calculation plus quantized store.
  • Paged-cache address calculation plus append.

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 layout

Extend existing llama.cpp KV-cache infrastructure rather than replacing it.

Goals:

  • Expose page or block metadata to optimized attention kernels.
  • Allow direct reads from quantized cache pages.
  • Avoid copying the active region into a contiguous F16 window.
  • Preserve context shifting, defragmentation, cache views, and sequence operations.
  • Make page size and layout part of the backend plan.
  • Support batched sequences without per-token cache repacking.

Paged access should only be selected when its address-calculation cost is lower than materialization or copying for the current geometry.

9. Backend rollout

Only 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

  • Continue using the PR 55 MMA body and native quantized loaders.
  • Stabilize Q4_0 and Q8_0 first.
  • Add optional Q4_1, Q5_0, and Q5_1 after correctness and binary-size review.
  • Add split-KV and graph-stable planning.

HIP

  • Reuse common CUDA/HIP source structure where practical.
  • Implement native quantized loaders feeding MFMA or WMMA tiles.
  • Maintain separate RDNA and CDNA profitability tables.
  • Validate wave32 and wave64 independently.
  • Do not reuse NVIDIA thresholds.

SYCL

  • Preserve existing direct quantized vector paths.
  • Remove full F16 materialization from fast prompt paths where possible.
  • Add XMX or subgroup-matrix loaders on capable devices.
  • Keep oneDNN/MKL or existing kernels as fallbacks.

Metal

  • Prefer existing direct quantized specializations.
  • Add simdgroup-matrix paths where they win.
  • Remove F16 prepasses only after validating each relevant Apple GPU family.
  • Account for unified-memory behavior in profitability decisions.

Vulkan

  • Build on the existing fused quantized dequantization shaders.
  • Add cooperative-matrix paths where supported.
  • Tune subgroup and workgroup geometry per capability class.
  • Preserve subgroup-only fallbacks.

CPU

  • Extend split-KV to quantized K/V.
  • Vectorize V-cache dequantization.
  • Reuse architecture-specific dot-product implementations for K.
  • Add SIMD-specialized Q4/Q8 tile loaders where profitable.
  • Dispatch by ISA, including AVX2, AVX-512, AVX-VNNI, AMX, NEON, SVE, and RVV where available.

10. Correctness and safety requirements

Every optimized route must be compared against the current reference path.

Required coverage:

  • Numerical comparison with F32 or the current F16-materialized implementation.
  • Repeated graph capture and replay.
  • Alternating context sizes in one process.
  • Context shifting and cache defragmentation.
  • Multiple sequences and mixed sequence lengths.
  • Cache views, offsets, and non-trivial strides.
  • Odd KV lengths and all relevant block boundaries.
  • Non-power-of-two GQA ratios.
  • Masked and unmasked attention.
  • Sliding-window layers.
  • Attention sinks, ALiBi, and logit soft cap.
  • Single-GPU and multi-GPU execution.
  • CUDA Compute Sanitizer for CUDA kernels.
  • Address and undefined-behavior sanitizers for host planning code.

Performance tests alone must not be treated as correctness evidence.

11. Binary-size control

Template specialization must remain bounded.

Controls:

  • Keep the cache-type manifest as the single source of truth.
  • Share runtime V-type dispatch where practical.
  • Group types with compatible loader structure.
  • Compile common same-type Q4_0 and Q8_0 kernels by default initially.
  • Keep uncommon mixed-type combinations optional until their value is demonstrated.
  • Track shared-library size in CI.
  • Reject duplicate or accidental instantiations.
  • Consider runtime-selected loader traits if they reduce size without material performance loss.

The optional all-quant build should not dictate the size of the default build.

12. Tuning policy

Use three layers:

  1. Hard capability checks for correctness and safety.
  2. Shipped deterministic profitability tables derived from a representative benchmark matrix.
  3. Optional developer-only tools that regenerate or validate those tables.

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 behavior

Attention optimization does not remove layer-split communication costs.

The implementation should:

  • Select plans independently for each GPU architecture.
  • Allow different plans on heterogeneous GPUs.
  • Avoid assumptions about equal SM count or memory bandwidth.
  • Keep KV pages local to the GPU that owns the layer.
  • Avoid cross-GPU KV materialization.
  • Measure layer split, row split, and pipeline effects separately where applicable.
  • Preserve multi-GPU as a capacity feature even when adding a slower PCIe GPU does not improve throughput.

14. Proposed change sequence

Keep each step independently testable and reviewable:

  1. Fix Q8/GQA and graph/workspace safety issues in PR 55.
  2. Add measured CUDA decode dispatch for Q4_0.
  3. Add validated Q8_0 decode dispatch.
  4. Extend CUDA coverage across D64, D128, D256, and then D512.
  5. Introduce a common attention problem description and backend planning interface.
  6. Implement CUDA split-KV and stable attention-state merge.
  7. Add fused quantized cache append where graph structure permits it.
  8. Integrate direct paged-cache access.
  9. Port native tile loading to HIP.
  10. Improve Vulkan and Metal direct-quantized paths.
  11. Replace avoidable SYCL materialization paths.
  12. Add quantized CPU split-KV and SIMD loaders.
  13. Expand default cache-format coverage after performance, quality, and binary-size validation.
  14. Enable optimized routes by default only after the complete correctness matrix passes.

15. Acceptance criteria

The project is successful when:

  • No model-name-specific dispatch exists.
  • Supported quantized KV formats avoid full F16 materialization.
  • An optimized route is selected only when it is measurably competitive with the fallback.
  • Long-context GQA decode improves substantially on supported devices.
  • MHA and D512 have either optimized paths or explicit safe fallbacks.
  • Graph replay remains safe across supported shape and context transitions.
  • Single-GPU and multi-GPU results remain correct.
  • Unsupported devices and backends preserve existing behavior.
  • Binary-size growth is controlled and measured.
  • llama.cpp does not acquire a FlashInfer, PyTorch, or CUDA-only architectural dependency.
  • Every backend can adopt the planning interface incrementally.

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 GenerelSchwerz left a comment

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

  1. The change breaks HIP source selection. The existing HIP build glob includes every fattn-mma*.cu file, so it compiles the new CUDA-native quant instances (HIP source selection). The resulting build fails in the new Q4 loader's raw half2 initialization (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.

  2. 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.

  3. The focused native-route test currently selects zero native cases and can report success. The workflow filters for native_quants=1 and sets a removed verbose environment variable (workflow), while the final test identity contains expected_native= instead (test identity). The filter removes nonmatching cases (filter), and the harness returns success when both n_ok and tests_run are 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.

  4. 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.

  5. 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 from supported && profitable (selection). Decide on one contract and make code, tests, workflow, benchmarks, and documentation agree before release.

  6. 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.

  7. The inventory checker cannot prove that exactness. It stores geometries in sets, which discards duplicates, and its positive expectation is only some rather 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.

  8. 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.

  9. 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.md and applicable docs/ops/*.csv after 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]

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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) };

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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;

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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);

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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);

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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% |

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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.

3 participants