Skip to content

feat(kvcache): add opt-in fp8 kv cache for r9700 - #853

Draft
zihaomu wants to merge 1 commit into
FlashML-org:mainfrom
zihaomu:feat/fp8-kv-cache-r9700-20261103
Draft

zihaomu wants to merge 1 commit into
FlashML-org:mainfrom
zihaomu:feat/fp8-kv-cache-r9700-20261103

Conversation

@zihaomu

@zihaomu zihaomu commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

Relates to #280 and #354.

Source commit: cbb9eff56e4459ece8e96b8ed695872c5f46642a on zihaomu:feat/fp8-kv-cache-r9700-20261103 (based on upstream 9cacd0896c9beb67d88c22a38be54236023b4605).

What changes

Add opt-in --kv-cache-dtype fp8 for 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 on 9cacd08 and 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

  • GPU: one AMD Radeon AI PRO R9700, gfx1201, PCI 0000:23:00.0, 32 GiB; driver 7.1.9.31600000, kernel 7.0.0-38-generic. Image freetoken:rocm10.1.0-transformers518-20261026 (SHA-256 6cf7ba2162685f0282ebbe3b3f98dfbb9c06af0663809bef5d72b508ef26c962), PyTorch 2.14.0+rocm10.1.0, Triton 3.8.0.
  • Checkpoint: Qwen/Qwen3-30B-A3B-FP8, converted to FTW; freetoken_weight.json SHA-256 5df75001dcfb0f33d539560f025032811db1942ab275bdbc4508b04176535fb3.
  • Runtime: FREETOKEN_USE_ARK=0, ROCR_VISIBLE_DEVICES=3 (this host maps it to PCI 23: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 off plus --kv-cache-dtype bf16 or fp8. CPU affinity 0-31,64-95, memory nodes 1-2 for both variants.

Validation

  • R9700: focused FP8 KV/attention/E4M3 suite 104 passed, 5 deselected; engine/KV/FTW/FP8-MoE regression 434 passed, 43 skipped, 2 deselected. The unrelated ROCm slow test test_forced_emu_matches_native also fails on clean main because it interprets gfx1201 capability as NVIDIA sm; standard -m 'not slow' excludes it.
  • RTX 5090: same 19 changed-file SHA-256 hashes, focused suite 103 passed, 1 skipped, 5 deselected with Torch 2.14.1+cu130 and Triton 3.8.0.
  • Full Qwen3 model: Graph batch 1 captured at both 2 GiB smoke and 20 GiB formal MoE cache capacities. With the 2 GiB smoke capacity, short and two-chunk 8,507-token requests succeeded; pool rebuild 33058 -> 33400 -> 32900 -> 33058 recaptured 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.
  • Teacher-forced NLL on 32 prompts / 1,024 target tokens: BF16 0.54525000, FP8 0.54482556, FP8 minus BF16 -0.00042445 nats/token (gate <= +0.01).
  • At 33,059 physical rows: BF16 persistent KV 3,249,831,936 bytes, FP8 codes + scales 1,675,694,592 bytes, and one shared BF16 workspace 67,704,832 bytes. FP8 persistent KV plus workspace saves 46.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,470 prefill token/s) is retained in the raw results.

Input tokens BF16 prefill tok/s FP8 prefill tok/s BF16 decode tok/s FP8 decode tok/s BF16 TTFT s FP8 TTFT s
512 632.32 633.76 101.19 100.47 0.810 0.808
8,192 4,479.52 4,463.12 88.85 90.95 1.829 1.835
32,768 2,953.91 3,114.11 65.13 69.29 11.093 10.522

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 ms with reconstruction versus 17.886 ms for BF16; 32K FP8 decode stage 1 took 88.0 us under 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-tokens to 40290, --max-seq-len-override to 40256, 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 PCI 23:00.0.

40K Qwen3-30B-A3B-FP8 metric (3-pair median) BF16 KV FP8 KV Change
Decode 60.66 tok/s 64.80 tok/s +6.8%
Prefill 2,687.67 tok/s 2,752.05 tok/s +2.4%
Full request 19.08 s 18.65 s 2.3% less time

Decode pairs were 60.66 -> 64.88, 60.71 -> 64.80, and 54.86 -> 61.92 tok/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 (revision 1cfa9a7208912126459214e8b04321603b3df60c) decoded at 37.14 -> 33.29 tok/s (BF16 -> FP8), and Qwen/Qwen3-0.6B (revision c1899de289a04d12100db370d81485cdf75e47ca) at 78.47 -> 60.56 tok/s. Four simultaneous 32K requests on 0.6B also reduced total output throughput, including prefill and queueing, from 61.72 to 58.91 tok/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 main has no FREETOKEN_USE_ARK integration. 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.

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
zihaomu marked this pull request as draft November 3, 2026 13:34

This branch has not been deployed

No deployments
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.

1 participant