Repository navigation
Conversation
Port the E4M3 row-scaled KV path from FlashML-org#354 to current main, add a shared BF16 prefill workspace, and tune gfx1201 decode. Validate BF16 versus FP8 KV with Qwen3-30B-A3B-FP8 on R9700, plus shared-path tests on RTX 5090.
zihaomu
marked this pull request as draft
November 3, 2026 13:34
This branch has not been deployed
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.
Relates to #280 and #354.
Source commit:
cbb9eff56e4459ece8e96b8ed695872c5f46642aonzihaomu:feat/fp8-kv-cache-r9700-20261103(based on upstream9cacd0896c9beb67d88c22a38be54236023b4605).What changes
Add opt-in
--kv-cache-dtype fp8for BF16 compute with the plain MHA/GQA pool, TP=1, and Triton attention. Each physical K/V row stores E4M3 codes as bytes plus its own FP32 scale. The prefill path reconstructs historical FP8 rows once per layer into one shared BF16 workspace; current-chunk K/V remains BF16. R9700 decode selects a one-stage pipeline for D=128/GQA=8 and two warps for the 32-key tile. Unsupported pools, attention backends, dtypes, TP>1, and speculative decoding fail during configuration.This is based on #354's FP8 KV design and row format. #354 covers additional pool families but its current head conflicts with
main. This branch is rebased on9cacd08and deliberately scopes the validated implementation to R9700 plain MHA/GQA. It does not depend on the FP8 weight kernels in #779. I am happy for the #354 author/maintainers to take the R9700-specific changes if that is the preferred integration route.Hardware and reproducibility
0000:23:00.0, 32 GiB; driver7.1.9.31600000, kernel7.0.0-38-generic. Imagefreetoken:rocm10.1.0-transformers518-20261026(SHA-2566cf7ba2162685f0282ebbe3b3f98dfbb9c06af0663809bef5d72b508ef26c962), PyTorch2.14.0+rocm10.1.0, Triton3.8.0.freetoken_weight.jsonSHA-2565df75001dcfb0f33d539560f025032811db1942ab275bdbc4508b04176535fb3.FREETOKEN_USE_ARK=0,ROCR_VISIBLE_DEVICES=3(this host maps it to PCI23:00.0),--model-path /model-ftw --dtype bfloat16 --attention-backend triton --moe-strategy offload --moe-cpu-threads 17 --moe-cache-mem 20GiB --moe-cache-scope unified --moe-cache-split equal --num-tokens 33058 --page-size 1 --max-seq-len-override 33024 --max-prefill-length 8192 --max-running-requests 1 --cuda-graph-max-bs 1 --max-output-tokens 256 --memory-ratio 0.9 --expert-load serial --cache-type naive --enable-cache-report --sampling-defaults none --reasoning-parser offplus--kv-cache-dtype bf16orfp8. CPU affinity0-31,64-95, memory nodes1-2for both variants.Validation
104 passed, 5 deselected; engine/KV/FTW/FP8-MoE regression434 passed, 43 skipped, 2 deselected. The unrelated ROCm slow testtest_forced_emu_matches_nativealso fails on cleanmainbecause it interprets gfx1201 capability as NVIDIA sm; standard-m 'not slow'excludes it.103 passed, 1 skipped, 5 deselectedwith Torch2.14.1+cu130and Triton3.8.0.33058 -> 33400 -> 32900 -> 33058recaptured and served; an oversized rebuild returned 503 without disturbing subsequent serving; and two concurrent requests each generated 64 tokens. With fixed 20 GiB capacity, BF16 and FP8 each passed 7/8 bounded tasks (the same arithmetic failure) and all three exact-32K retrievals. Repeated 8K radix prompts hit 8,191 cached tokens in both modes.0.54525000, FP80.54482556, FP8 minus BF16-0.00042445nats/token (gate<= +0.01).3,249,831,936bytes, FP8 codes + scales1,675,694,592bytes, and one shared BF16 workspace67,704,832bytes. FP8 persistent KV plus workspace saves46.35%. This excludes weights, MoE cache, activations and allocator reserve.Paired end-to-end A/B
Three alternating fresh-process rounds per variant on the same branch, 20 GiB MoE cache, Graph batch 1, zero prefix reuse. Every measured request had 256 generated tokens and 256 nonempty SSE text events. Values are medians; the slow first-round FP8 8K sample (
2,470prefill token/s) is retained in the raw results.All planned median gates pass: FP8 prefill and decode >=95% of BF16, TTFT <=105% of BF16. No other process allocated VRAM on the selected GPU during the six measured services. The 8K historical-prefix + 8K new-token operator check took
17.946 mswith reconstruction versus17.886 msfor BF16; 32K FP8 decode stage 1 took88.0 usunder Graph and compiled with zero spills.Follow-up model/workload search (2026-11-03)
The same commit and R9700 were tested with a longer native-context Qwen3-30B-A3B-FP8 workload and two dense-model pilots. This is an update to the performance evidence, not a source change. The 30B run kept the checkpoint, 20 GiB MoE cache, CPU affinity, Graph batch 1, and other runtime settings above; it changed
--num-tokensto40290,--max-seq-len-overrideto40256, and used exactly 40,000 input / 256 output tokens (within the model's 40,960-token context). Three fresh-process BF16/FP8 pairs alternated launch order, with a warmup before each measurement and zero prefix reuse. All six measured outputs had the same text SHA-256, exact token counts and 256 nonempty SSE events; GPU process sampling found no competing allocation on PCI23:00.0.Decode pairs were
60.66 -> 64.88,60.71 -> 64.80, and54.86 -> 61.92tok/s. The slower third BF16 sample is retained. This supports a modest long-context decode gain; the full-request gain is smaller because TTFT is about 15 s for a 40K prefill and only 256 tokens are decoded. At 40K the reported KV allocation is 98,304 bytes/token for BF16 versus 50,688 for FP8, about 48.4% less. The FP8 path also quantizes each new K/V row and loads per-row FP32 scales during attention, so less KV traffic does not imply faster serving on every model. The exact time breakdown of these costs has not been profiled in this search.Exploratory single-pair 32K/256-output pilots on pinned dense checkpoints were negative:
Qwen/Qwen3-4B(revision1cfa9a7208912126459214e8b04321603b3df60c) decoded at37.14 -> 33.29tok/s (BF16 -> FP8), andQwen/Qwen3-0.6B(revisionc1899de289a04d12100db370d81485cdf75e47ca) at78.47 -> 60.56tok/s. Four simultaneous 32K requests on 0.6B also reduced total output throughput, including prefill and queueing, from61.72to58.91tok/s. These pilots are not quality qualifications or multi-round performance claims. The 30B gain should not be generalized to dense models. Exact scripts, prompts, service logs and per-request JSON remain on the test host under.freetoken-zihao/fp8-kv-30b-long-benchmark-20261103/,.freetoken-zihao/fp8-kv-dense-benchmark-20261103/, and.freetoken-zihao/fp8-kv-06b-benchmark-20261103/.Scope remains BF16 compute, plain MHA/GQA, TP=1, Triton, R9700 as the primary performance target. RTX 5090 validation covers the shared numerical path; performance is not claimed there. Speculative decoding and other KV pool families are rejected at startup.
ARK compatibility is tracked separately from this main-based PR: upstream
mainhas noFREETOKEN_USE_ARKintegration. On the rebased #779 experimental stack, an ARK-native short request (11 input / 16 output) and two-chunk long request (8,507 input / 8 output) returned the same text as its FT fallback, with Graph disabled. That stack result is not a validation claim for this commit or for ARK performance.