perf(KV-FP8): serve fp8 KV prefill through the bf16 dense dequant

[iiLaurens]

Row

KV-FP8 — one row per PR. Spec: .agents/specs/fp8-kv-prefill-dense-dequant.md.

Before starting

• Issue/PR search: KV-FP8 W2 (#1593)
  landed the fp8 store and read-dequant for correctness; no issue covered the
  PREFILL performance half. Filed ISSUE-LOCAL-01M3QG4WWC3X0C84M9PWQZAJ10.
• Pull request shape selected at row claim: one PR for the spec and the
  implementation.
• Anchors inspected: cuda_paged_attn.cu:148-197 (the W2 fp8 read and its
  per-element dequant), :2909, :2954 (the dispatch split that names the
  scalar arm), kv_cache_route.h:40,73-78, and the 33-case suite.

What changed

The fp8 KV-cache read sent every prefill to the scalar CUDA-core flash kernel at
f32 q/out, because the model presents f32 for a non-bf16 store and the vendored
FA-2 admission requires bf16 q, bf16 KV and bf16 out. nsys attributes the whole
gap: 25.1 s of 36 s in PagedFlashKernel<float, unsigned char, …> at 8k, against
2.9 ms/layer for the bf16 store on the FA-2 split-KV kernel.

The fp8 cache is now dequantized ONCE per layer into a dense bf16 scratch, and
the normal bf16 dispatch runs on it. The scratch is one block per request with
block_size = max_seq and an identity block table, so the paged address IS the
dense address and no attention kernel changes. A new shared
KvCachePresentsBf16 presents bf16 for an fp8 store, which admits FA-2 with zero
cast kernels; the CUDA path is model-agnostic, so any model that gains an fp8
store inherits it. VT_ATTN_FP8_DENSE=0 restores the per-read dequant for a
same-binary A/B. benchmarks/paged_attn_prefill_ab.cpp (new target) is the
isolated sweep: one vt::PagedAttention call per arm, so kernel cost is
separated from engine cost.

Two earlier shapes of the same lever are recorded as rejected rather than
deleted, because both are plausible and neither survives the numbers:

1. routing fp8 through the bf16 WMMA ladder with the cast inside the K/V staging
   (8k 28.0 -> 9.1 s) — the per-element dequant stays on every re-stream;
2. the hardware __nv_cvt_fp8_to_halfraw staging conversion — bit-identical and
   cheap, but still per re-stream.

Evidence

$ ./build-cuda/tests/test_ops_paged_attn
test cases:     33 |     33 passed | 0 failed | 0 skipped
assertions: 455286 | 455286 passed | 0 failed

$ ./build-cuda/tests/test_ops_paged_attn \
    -tc='paged_attention CUDA fp8-KV prefill (one-time dense dequant) matches f32 ref at head_dim 256'
MESSAGE: fp8-KV dense prefill max_abs_err vs f32 ref = 1.87755e-06
test cases: 1 | 1 passed | 0 failed | 32 skipped

Build: -DVLLM_CPP_CUDA=ON -DVLLM_CPP_CUTLASS_DIR=<local CUTLASS 4.x> -DVLLM_CPP_CUDA_ARCHITECTURES=120a on the local sm_120a card (RTX PRO 4000
Blackwell, 24 GiB). The 1.87755e-06 is the 1.9e-6 the commit body records — the
exact dequantized values, tighter than the bf16-compute envelope.

The new case pins the fp8 cache against the f32 reference at 1.9e-6 max abs err —
the exact dequantized values, tighter than the bf16-compute envelope — with the
other 33 cases unchanged.

The e2e numbers in the commit body are the AUTHOR's run on the local sm_120a
card, 27B NVFP4 arm, 2048-token chunks, no spec, VT_HOST_EMBEDDING=1, with the
bf16 store and llama.cpp's q8_0 KV as denominators on the same box:

32k prefill  34.9 s / 938 tok/s   (bf16 store 36.4 s / 901; llama.cpp q8_0 39.9 s / 822)
             peak 20.5 GiB        (bf16 store 22.5 GiB)
8k prefill   7.6 s / 1000 tok/s
decode MTP n=3  38.6 tok/s        (37.4 before)

They are recorded as the shape's shape evidence, NOT as a gate: the operator's
same-binary A/B under lease is owed.

• commit style + trailer contracts OK; full scripts/agent-preflight.sh is
  RED AT BASE on this box (35 gates) and each failing gate was re-run on a
  clean origin/main worktree and fails there too
• tests that cover this change: tests/vt/test_ops_paged_attn.cpp
• no public doc changes: VT_ATTN_FP8_DENSE is a kernel-internal A/B knob,
  not a user-facing surface

Integration run on the official checkpoint

This branch alone cannot run the 27B artifact: without #3355's NVFP4 ownership
fix the first forward dies in marlin_repack: malloc bqt: out of memory (the
leak is characterized in PR #3355). The runs below therefore use a scratch tree
with #3355, #3356 and this change merged (not pushed), which is how the arm was
exercised end to end on the local sm_120a 24 GiB card:

$ ./build-cuda/examples/vllm-cli --model <Swift-1.5-Qwen3.8-27b-NVFP4 shards> --device cuda \
    --max-num-seqs 1 --kv-cache-memory 2000000000 --max-tokens 8 \
    --prompt "The capital of France is" --temperature 0 [--kv-cache-dtype fp8|bfloat16]

• fp8 KV and bf16 KV produce the same 8 tokens ("The capital of Germany is")
  with the same prompt, both exit 0.
• A 1440-token prefill with --max-tokens 2: fp8 5.063 s vs bf16 4.996 s
  (284.4 vs 288.2 prefill tok/s) — equal within noise at this length. The
  commit body's 938-vs-901 tok/s is a 32k-chunked-prefill measurement; this short
  run does not reproduce it and is NOT offered as a speed result.

This is a smoke test of the arm, not this PR's gate; the gate is the op parity
above.

Review follow-up, verified on Linux CUDA

The review's findings were fixed in b2f195e19 and re-verified here. Linux,
-DVLLM_CPP_CUDA=ON -DVLLM_CPP_CUTLASS_DIR=<local> -DVLLM_CPP_CUDA_ARCHITECTURES=120a:

$ ./build-cuda/tests/test_ops_paged_attn
test cases:     34 |     34 passed | 0 failed | 0 skipped
assertions: 455290 | 455290 passed | 0 failed

$ ./build-cuda/tests/test_ops_paged_attn \
    -tc='paged_attention CUDA fp8-KV prefill returns both scratch buffers to the pool when a later step throws'
test cases: 1 | 1 passed | 0 failed | 33 skipped
assertions: 4 | 4 passed | 0 failed

MUTATION of the guarantee the new case exists for — disable the scope guard's
free and rebuild cuda_paged_attn.cu:

ERROR: CHECK( Fp8DensePoolUsedBytesForTesting() == used_before ) is NOT correct!
  values: CHECK( 24584 == 12288 )      # 12,296 bytes retained
test cases: 1 | 0 passed | 1 failed;  assertions: 4 | 2 passed | 2 failed

Both injected failures (the identity allocation, and the identity copy after
both allocations succeeded) are covered, the original exception propagates, and
the pool-used-bytes assertion is what catches the leak when the guards are
removed.

The happy path is confirmed end to end on the official checkpoint in the
integration run above: fp8 KV generates the same 8 tokens as bf16 with host
embedding on (device_upload=9.676 GiB, exit 0).

Speed claims

• This PR makes NO speed claim — checked as a helper PR: the numbers above
  are the author's, on a local card. The operator A/B under ${GPU_LOCK} and
  the docs/BENCHMARKS.md entry are owed and named in the spec.

Honest gaps

• No operator-gated A/B, no docs/BENCHMARKS.md entry, and no dgx:gpu0
  reproduction. The numbers are one card, one artifact.
• The dense scratch is allocated per layer with cudaMallocAsync under the
  stream; its own cost is inside the measured prefill time, not separated from
  it (the bench exists so a reviewer can separate them).
• scripts/agent-preflight.sh is red at base (record drift, the
  e126687a9a vs a7c23ac96d oracle pin mismatch, check-test-registration's
  CMake probe, the missing file tool); every failing gate reproduces on clean
  main.

Refs ISSUE-LOCAL-01M3QG4WWC3X0C84M9PWQZAJ10

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:opencode-go/deepseek-v4.1-flash [pi]

[localai-org-maint-bot]

The new dense-scratch path leaks device memory on exceptions. In LaunchPagedFp8Out, the first cudaMallocAsync allocates dense, then the identity-table allocation, copy, launch checks, and LaunchPaged can all throw. Both frees occur only at the end of the success path. For example, if the identity-table allocation fails after the much larger dense allocation succeeds, Check throws and no owner frees dense. A failed request can therefore retain its scratch for the process lifetime and leave less memory for later requests.

Please use stream-aware scope guards/RAII for both allocations immediately after each successful allocation. Cover an injected failure after the first allocation, plus one after both allocations, and verify both buffers are released without masking the original exception.

The new benchmark also says every arm is verified against a bf16 reference with maximum absolute error, but it runs only the selected arm, prints a checksum, and returns zero without a comparison. Please implement the stated comparison or correct that claim; a checksum alone does not establish attention parity.

I read the complete diff at a40ba50. These are static findings; CUDA parity, failure injection, and the model speed/memory gate remain unexecuted here. @mudler @richiejp

[iiLaurens]

Repairs pushed at b2f195e19.

Scratch ownership. The dense bf16 scratch and the identity table are now each
owned by a stream-aware scope guard from the moment the allocation succeeds. The
identity allocation and copy, the dequant launch, the bf16 dispatch it feeds, and
the vectors between them all unwind through the guards, which free on the same
stream (ordered behind the work that reads the buffer) and never throw, so the
original exception is not replaced. The success path keeps its explicit checked
frees through release().

Coverage. A new case in test_ops_paged_attn injects both failures the
review asked for — the identity allocation (after the dense scratch exists) and
the identity copy (after both allocations succeeded) — through a test-only seam,
then asserts the CUDA memory pool's used bytes return to their pre-call value and
that the propagated message is the failing Check. The case skips when no CUDA
device is registered; this host has none, so its execution is owed to the CUDA
lane.

Benchmark claim. The header no longer claims a bf16 parity comparison it does
not perform. It now states that the binary runs one arm and prints a checksum,
and points at test_ops_paged_attn.cpp for the f32-reference parity gate.

Evidence: MSVC 17.14 x64 Release CPU-only — test_ops_paged_attn and
vllm_paged_attn_prefill_ab compile and link.

[iiLaurens]

Addressed in b2f195e19.

• Both allocations are now owned from the moment each succeeds. LaunchPagedFp8Out wraps the dense scratch and the identity table in a stream-aware scope guard: the free runs on the same stream, so it stays ordered behind the dequant and the bf16 dispatch that read the buffer, and the destructor never throws, so the original Check exception is not replaced. The success path keeps its explicit checked frees through release().
• The benchmark claim is corrected. The header now states what the binary does — one selected arm, a checksum, no comparison — and points at test_ops_paged_attn.cpp, which holds the f32-reference parity gate.

Both injected failure points you asked for are covered by paged_attention CUDA fp8-KV prefill returns both scratch buffers to the pool when a later step throws: the identity allocation failure, and the identity copy failure after both allocations succeeded. It asserts the pool's used bytes return to the pre-call value and that the propagated message is the failing Check.

Verified on Linux with CUDA + CUTLASS at arch 120a on the local sm_120a card:

$ ./build-cuda/tests/test_ops_paged_attn
test cases:     34 |     34 passed | 0 failed | 0 skipped
assertions: 455290 | 455290 passed | 0 failed

$ ./build-cuda/tests/test_ops_paged_attn \
    -tc='paged_attention CUDA fp8-KV prefill returns both scratch buffers to the pool when a later step throws'
test cases: 1 | 1 passed | 0 failed | 33 skipped
assertions: 4 | 4 passed | 0 failed

MUTATION of the guarantee that case exists for — disable the scope guard's free and rebuild cuda_paged_attn.cu:

ERROR: CHECK( Fp8DensePoolUsedBytesForTesting() == used_before ) is NOT correct!
  values: CHECK( 24584 == 12288 )      # 12,296 bytes retained
test cases: 1 | 0 passed | 1 failed;  assertions: 4 | 2 passed | 2 failed

Restored, the case is 4/4 green again. The happy path is confirmed end to end too: with the scratch guards in place, fp8 KV generates the same 8 tokens as bf16 on the official ukisai/Swift-1.5-Qwen3.8-27b-NVFP4 checkpoint with host embedding on (device_upload=9.676 GiB, exit 0).