Repository navigation
perf(attention): add packed fp8 kv reads on gfx1201 - #897
Conversation
|
@zihaomu 1. Kernel A/B against main. The only main number is one service at 33.65 tok/s, while the same HSACO measured 37.6 and 43.7 tok/s in two services, so it cannot show the gain. Please add a stage-1 A/B on the R9700 against main (D128, GQA 4/8, 8K and 32K keys), plus the generic reader with the same 64/1/2 launch, so we can see what the packed load buys apart from the launch change. Please also give the VGPR and spill count of the packed kernel; we keep spills at 0. 2. Launch in attention.py, arch from the existing helper. def _use_gfx12_fp8_reader(kv_format, q, k_cache, v_cache) -> bool:
return (
kv_format.name == "fp8_per_token_head" and not FORCE_EMU
and get_rocm_gfx_arch() == "gfx1201"
and ... # the same dtype, shape and alignment checks
)Is there a reason to leave out gfx1200 (RX 9060 series)? LLVM gives gfx1200 and gfx1201 the same feature set (both 3. Drop the test-only arm. The kernel calls 4. Comments. "Preserve the merged reader's address expressions on non-packed targets." (attention.py:320) describes the PR history; please state the reason instead, e.g. |
Merge main through 66190a9, retaining the multi-row decode changes from FlashML-org#900. Move architecture and launch selection into attention.py, reuse the shared architecture helper, and remove the packed reader's test-only generic arm. Use four warps and require 16-byte aligned cache bases and strides to avoid packed-reader register spills. Keep CUDA address expressions inline and document the gfx12 VOP3 conversion workaround. Validate D128 GQA4/8 against main and same-launch generic controls. Confirm zero packed spills, 14 layout/Graph cases, 58 final operator cases and 14 forced-emulation cases. All 81 RTX 5090 main/head codegen comparisons match. Preserve four unrelated AMD suite failures reproduced on pinned main.
|
I reproduced and narrowed the pre-existing gfx1201 BF16 D256 multi-row LDS failure mentioned in the PR description, using current main Minimal stage-1 specialization that reproduces the exact reported number:
Triton metadata: The 65,792-byte result is therefore independently reproducible from the kernel compiler alone and exactly matches the failure reported here. I also isolated the stage effect: One useful detail: The lowest-risk fix looks like either:
Option 2 preserves the 3-stage path and, in this exact specialization, cuts compiled LDS from 65,792 to 32,768 bytes. I have the compile-only reproducer saved locally if any other shapes are useful to sweep. Hardware context for the host doing the work: 2x 32 GiB gfx1201 R9700 + one gfx1201 RX 9070 XT; host ROCm 7.2.3. The compile result above is target-pinned and does not claim a runtime benchmark. |
Adds a packed FP8 KV reader for the qualified gfx1201, BF16 query, D128, GQA4/8, batch-one single-row decode path. Four contiguous E4M3FN bytes are loaded as one uint32, converted with native packed instructions, and reconstructed as BF16 for the existing attention math. KV storage, per-token/head scales and BF16 WMMA remain unchanged; this is not native FP8 WMMA attention.
Merged current main 66190a9 (including #900) into the PR. This revision addresses the review with 64/1/4 (
BLOCK_N/num_stages/num_warps), replacing the previous spilling 64/1/2 launch. Current-revision kernel and regression validation applies to fc3cb03 and the pinned main above. The model-context matrix below is explicitly labelled as historical.attention.py, using the existingget_rocm_gfx_arch()helper. Packed dispatch requires 16-byte aligned bases and row/head strides: a merely 4-byte aligned padded layout produced 33 SGPR spills in validation, so it now uses the original generic reader.kv_rows.pycontains only the packed addressing/conversion; its converter is private and the test-onlyPACK=1arm is removed.Model speed across context lengths: BF16 KV vs FP8 KV
Historical measurements from 2026-11-12, restored here for visibility; these models were not rerun for the current review revision.
These measurements compare BF16 KV versus FP8 KV on the same branch, not main versus the new native reader. Neither model enters this PR's D128 native dispatch: Qwen3.6 uses D256 full attention among its GDN layers; Qwen3.8 uses QSA sparse attention. These results qualify the exercised generic paths and document remaining performance work, not gains attributable to this patch.
Both checkpoints were already present locally; no model was downloaded or requantized:
995ad96eacd98c81ed38be0c5b274b04031597b0.fc694b54fb0174e0913e6adf86691ef85a4ead47.Controls: same physical GPU0 / PCI
0000:03:00.0, TP=1, batch=1, BF16 compute, Graph replay, MoE offload, fixed 16 GiB expert cache, 35,072 KV slots, naive cache, prefill chunks of 2,048, greedy and exactly 128 output tokens. Qwen3.6 usestriton; Qwen3.8 usesqsa_sparse,moe.nvfp4=triton, and disk PLE. CPU affinity is16-31,80-95; memory NUMA policy remains default because the container deniesset_mempolicy.Each format has one service process; each input length is warmed once, then measured three times. All 36 measured requests and 12 warmups completed, and usage/admission counters verified exact prompt/output lengths and no prefix reuse. Medians below include every completed measurement, including slow requests.
Decode timing for this matrix is
(completion_tokens - 1) / (terminal finish_reason SSE timestamp - first nonempty text timestamp). It includes the terminal event even when its text is empty. Raw runner logs retain the older first/last-nonempty estimate; this table is consistently recomputed from the saved SSE events. TTFT includes full prefill.Known regression and uncertainty
FP8 KV is slower in the measured Qwen3.6 512-token case. Its BF16 samples are 31.93 / 28.18 / 29.07 tok/s; FP8 samples are 18.53 / 16.88 / 20.33. This D256 case is not fixed by the D128 dispatch in this PR.
Short contexts offer less KV bandwidth to save while conversion and row-scale work remain. The model's D256 decode specialization reports 73 VGPR spills with FP8 versus 2 with BF16. This is a profiling lead, not an isolated root-cause result: kernel timing and controlled ablations have not established how much each factor contributes.
There is substantial runtime variation. Qwen3.6 FP8 at 32K has samples 9.03 / 26.25 / 28.63 tok/s; Qwen3.8 BF16 at 32K has 24.78 / 24.93 / 31.80. No slow sample was dropped. These are within-service repeats in sequential format runs, not independent alternating service A/B rounds. The percentages are descriptive and do not establish stable gains or a universal -36.24% penalty.
At the same KV capacity, logged K/V allocations are 0.67 -> 0.34 GiB for Qwen3.6 and 0.83 -> 0.43 GiB for Qwen3.8. GDN state and expert weights remain unchanged. Saved VRAM is not reassigned to the expert cache in these comparisons.
Why the handwritten conversion remains
The pinned Triton 3.8 ROCm build lowers plain OCP FP8-to-BF16 casts in software on RDNA4. The handwritten
v_cvt_pk_f32_fp8_e64forces VOP3: LLVM's word-0 packed intrinsic can select a VOP1 form rejected by the gfx12 assembler for VGPR inputs. Triton #11809, merged as7e239bb, adds hardware upcasts using one byte-select FP8-to-FP32 conversion per byte, then narrowing. Its merge does not change our pinned compiler. Recheck plain casts for correctness, generated instructions and cold/Graph performance on a compiler containing that change before removing this workaround. FP8-to-BF16 casts were already numerically possible before that PR; efficient lowering is the issue.Fresh R9700 stage-1 A/B
Radeon AI PRO R9700/gfx1201, PCI
0000:03:00.0; dual EPYC 9334, CPU affinity16-31,80-95; Linux7.0.0-42-generic, amdgpu7.1.9.31600000. Locked ROCm 10.1, HIP 7.16.26385, Torch 2.14.0+rocm10.1.0, Triton 3.8.0 (669b31acROCm build); imagesha256:a79c374ab2bc93c0e4358da6c7bda5cb84cf1308e8aff74fc99fc95c511cbd03. Physical GPU0 was reserved and monitored; no foreign GPU process was observed.Same randomized physical slots, Q/K/V/scales, D128, 32 query heads, batch 1, 64 scratch/split cap. Independent CPU FP8 dequantization plus FP32 PyTorch attention reference, unchanged
atol=0.002, rtol=0.02. Five seeded shuffled rounds per case; every sample retained. Stage 1 is timed separately from the merge kernel. Each cell is cache-evicted / warm Graph latency in microseconds, median across five rounds. Cache eviction usestriton.testing.do_bench; Graph usesdo_bench_cudagraph. Main selects 32/3/4 in this matrix.The same-launch controls isolate the packed reader from launch tuning; the reader delta includes uint32 loads,
.cgpolicy and native conversion. The final packed reader is faster than every main and generic-64/1/4 sample at each tested length, in both cache conditions. One main GQA4/8K cold round is a long tail (max/min 1.363); it is retained, not discarded or corrected. These are kernel measurements, not model throughput multipliers.Register/spill correction: the previous packed 2-warps kernel had 61/65 native VGPR spills and 248/264 bytes scratch for GQA4/8. Four warps remove them. At GQA8/32K, this costs about 4% cold / 5% Graph latency relative to the spilling 2-warps version; the final configuration follows the zero-spill requirement.
Counts above come from native AMDGPU metadata (
.vgpr_count,.vgpr_spill_count,.sgpr_spill_count,.private_segment_fixed_size), not conflated with Triton'sn_spills. Final packed kernels also reportn_spills=0. Zero-spill qualification covers the new packed path; it does not claim to fix existing spills in main's generic AMD kernels.Regression and CUDA isolation
All four AMD failures reproduce on pinned main with the same failure reasons. The BF16 D256 two-query backend case routed through #900's decode requires 65,792 bytes of shared memory on this compiler, exceeding the 65,536-byte device limit. The three e4m3 suite failures are the native-vs-emulation assertion (
e4m3_native()is false on ROCm) and the sm89/sm120 foreign-target compile gates (bf16/uint8operand mismatch). Their tests and production paths are unchanged by this PR; they are retained as failures, not skipped or made green. Main's multi-row AMD shared-memory launch needs a separate fix. The new packed reader is disabled for that path.Fresh RTX 5090 GPU1 / PCI
0000:e1:00.0, EPYC 9554, driver 615.71.09, immutable NGC imagesha256:42263b2424fc237b34c4fc4a91c30d603c57eed36e37d31ff6d9a4f1f801edee, existing Torch 2.14.1+cu130, CUDA 13.0, Triton 3.8.0. 81 main/head comparisons: BF16/FP8/INT8 KV; D128 GQA4/8 and D256 GQA8; 17/8192/32768 keys; 1/3/8 query rows. All independently checked outputs and Graph replays pass; output bits, CUBIN executable sections, normalized SASS, launch parameters and register/shared/spill counts match main. All NVIDIA cases report zero spills. This verifies the tested SM120 stack, not every NVIDIA architecture.Fresh model A/B against current main
Local Qwen/Qwen3-4B, BF16 weights, revision
1cfa9a7208912126459214e8b04321603b3df60c; no download. Same reserved R9700, TP=1, batch=1, BF16 compute, FP8 KV, Triton attention, Graph, greedy, 32,768 prompt tokens and 128 outputs. Fixed 35,072 KV slots, page size 1, prefill chunks of 2048, memory ratio 0.7, radix cache. Four fresh services in main/head/head/main order, each warmed once then measured three times. TTFT is for the warm prefix, not full prefill. Decode timing includes the terminalfinish_reasonevent:127 / (terminal - first nonempty text).The model's compiled packed specialization uses 245 VGPRs/77 SGPRs and also has zero VGPR/SGPR spills and zero scratch; its real cache/scale strides differ from the microbenchmark's separate allocations.
All 12 measurements and four warmups completed; no slow sample was dropped. Within-source service-median max/min: main 1.1290, head 1.0022.
The predeclared 1.10 startup-drift gate failed; these service rates are descriptive only and do not establish a stable end-to-end gain.
Output text SHA256 sets: main
00a5408efb8f9b76a4b92a0536ec84b6b1635e6ac6bb1140216cbf8a4064e78d; head9bc101aaf7a671cc0bda7ac257ffda622359b8fd69fd24ce848f92cd76dd9901.Main/head synthetic continuations differ; do not treat this as a model-quality or token-identity gate. Independent operator numerical checks above remain the correctness evidence.
These synthetic continuations are not a task-quality evaluation. Previous service rates of 33.65/37.6/43.7 tok/s were from earlier revisions and inadequately paired cohorts; they are superseded as evidence for this revision's performance.
Reproduction and retained evidence
Run the same commands at pinned main and this head inside the recorded hardware/runtime, after selecting an idle device by PCI identity:
PYTHONPATH=python python -m pytest tests/kernels/test_triton_attention.py \ tests/kernels/test_e4m3_compat.py tests/kvcache/test_kv_cache_rebuild.py \ tests/kvcache/test_pool_sizing_surface.py tests/engine/test_attention_backend_arch.py -q FREETOKEN_FORCE_E4M3_EMU=1 TRITON_CACHE_DIR=/tmp/ft-pr897-emu PYTHONPATH=python \ python -m pytest tests/kernels/test_triton_attention.py \ -k 'matches_dequantized_reference and 128' -q PYTHONPATH=python python -m freetoken.cli serve \ --model /models/qwen3-4b --served-model-name cast-4b \ --dtype bfloat16 --attention-backend triton --kv-cache-dtype fp8_per_token_head \ --moe-strategy offload --num-tokens 35072 --page-size 1 \ --max-seq-len-override 34816 --max-prefill-length 2048 \ --max-running-requests 1 --cuda-graph-max-bs 1 --memory-ratio 0.7 \ --text-model-only --enable-cache-report --sampling-defaults none \ --reasoning-parser off --host 127.0.0.1 --port 19854Send the saved exact 32K text fixture to
/v1/completionswithmodel=cast-4b,max_tokens=128,temperature=0,ignore_eos=true, streaming with usage. Prompt SHA256:28e8514172da8e98fd27b06d31c4650eb048c3a61faf82de233957f566ffed9e. Warm once and measure three times in each ABBA service; verify 32,768 input/128 output and 32,767 cached prompt tokens per measured request.New host artifacts:
.freetoken-zihao/pr897-review-20261113/:bench.py(main, same-launch generic, packed controls; captures stage1), all individual timing samples,asm/HSACO/native assembly,*-sources.json, runtime/ownership records, test XML, padded/unaligned/window integration checks, exact prompt/request/SSE files, andsummary.json. NVIDIA artifacts: same root'snvidia/, remotetest_lap/pr897-review-20261113/. Historical primitive casts, mutation checks and hybrid runs remain under.freetoken-zihao/amd-fp8-kv-pr894-20261112/. These paths are retained local artifacts, not public download URLs.