spark_model/layers/fp8_calibration/staging.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! BF16 staging for the FP8 KV calibration window (Atlas #919).
4//!
5//! ## Why this exists
6//!
7//! The FP8 KV round-trip is only correct if the scale that quantized a cache
8//! entry is the scale that dequantizes it: paged attention reads a sequence's
9//! whole history in one pass with ONE `k_scale`/`v_scale`. That is why the
10//! 2026-07-25 hardening froze the scale on the FIRST observe — but it is also
11//! why #919 happened, because "the first observe" is the readiness probe and a
12//! 24k serve ended up calibrated on 13 tokens.
13//!
14//! ## The fix, and its tradeoff
15//!
16//! Accumulate the amax over the requested window and, when the freeze finally
17//! fires, make the already-written entries agree with the new scale by
18//! REWRITING them. Of the three options considered (#919):
19//!
20//! * (a) hold pre-freeze tokens in BF16 — impossible without a second pool:
21//! the FP8 pools are sized for 1 byte/element and attention must be able to
22//! read those tokens back DURING the window.
23//! * (b) rescale the stored FP8 codes in place — needs a new dequant/requant
24//! kernel in all five arch trees (`kernels/{hopper,b200,gb10,strix,strix-hip}`)
25//! and loses precision twice.
26//! * (c) **chosen**: keep the original BF16 K/V of the window's batches, plus
27//! their slot mappings, and replay them through the EXISTING
28//! `reshape_and_cache_fp8` kernel at the frozen scale. No new kernel, and the
29//! rewritten entries are quantized once, from the original BF16, at the final
30//! scale — strictly better than a code-domain rescale.
31//!
32//! Replay is chronological, so "last writer wins" per slot is identical to the
33//! original write order even if a block was freed and re-allocated inside the
34//! window. The crossing batch is written by the caller AFTER the replay, at the
35//! frozen scale, so it needs no staging.
36//!
37//! Cost: `window_tokens * num_kv_heads * head_dim * 2 B * 2` device bytes per
38//! attention layer, freed at the freeze. 256 tokens x 4 KV heads x 128 dims is
39//! 512 KiB/layer. `MAX_STAGED_TOKENS` caps a pathological `--fp8-kv-calibration-
40//! tokens`.
41//!
42//! Known gap, documented rather than fixed: KV spilled to host
43//! (`--high-speed-swap`) inside the window is not replayed. The window is a few
44//! hundred tokens, long before spill pressure.
45
46use anyhow::Result;
47use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
48
49use crate::layers::ops;
50
51/// Hard cap on staged calibration tokens, so an absurd
52/// `--fp8-kv-calibration-tokens` cannot reserve gigabytes of device memory per
53/// layer. The effective window is clamped to this in `Fp8KvCalibration::new`.
54pub(super) const MAX_STAGED_TOKENS: usize = 4096;
55
56/// Bytes per BF16 element.
57const BF16: usize = 2;
58/// Bytes per `slot_mapping` entry (int64, see `reshape_and_cache_fp8.cu`).
59const SLOT: usize = 8;
60
61/// Everything the replay needs to re-run `reshape_and_cache_fp8` for one layer.
62#[derive(Debug, Clone, Copy)]
63pub struct Fp8KvWriteTarget {
64 /// `reshape_and_cache_fp8` kernel handle (the layer's `reshape_cache_k`).
65 pub kernel: KernelHandle,
66 /// Base pointer of this layer's K pool.
67 pub k_pool: DevicePtr,
68 /// Base pointer of this layer's V pool.
69 pub v_pool: DevicePtr,
70 /// Tokens per KV block.
71 pub block_size: u32,
72 /// Pool stride in ELEMENTS (`block_size * num_kv_heads * head_dim`).
73 pub cache_stride: u64,
74 /// Source K row stride in elements, as passed to the live write.
75 pub key_stride: u32,
76 /// Source V row stride in elements, as passed to the live write.
77 pub value_stride: u32,
78 /// This batch's `slot_mapping` (int64 per token), staged alongside the
79 /// BF16 K/V so the replay lands on exactly the entries the window wrote.
80 pub slot: DevicePtr,
81}
82
83#[derive(Debug, Clone, Copy)]
84struct StagedBatch {
85 offset_tokens: usize,
86 num_tokens: u32,
87}
88
89/// Device-side copies of the calibration window's BF16 K/V and slot mappings.
90#[derive(Debug)]
91pub(super) struct KvStaging {
92 k: DevicePtr,
93 v: DevicePtr,
94 slots: DevicePtr,
95 capacity_tokens: usize,
96 elems_per_token: usize,
97 used_tokens: usize,
98 batches: Vec<StagedBatch>,
99}
100
101impl Default for KvStaging {
102 fn default() -> Self {
103 Self {
104 k: DevicePtr(0),
105 v: DevicePtr(0),
106 slots: DevicePtr(0),
107 capacity_tokens: 0,
108 elems_per_token: 0,
109 used_tokens: 0,
110 batches: Vec::new(),
111 }
112 }
113}
114
115impl KvStaging {
116 /// Tokens staged so far.
117 pub(super) fn used_tokens(&self) -> usize {
118 self.used_tokens
119 }
120
121 /// Copy one pre-freeze batch aside. Silently declines (returns `false`) if
122 /// the window no longer fits — the caller then writes the batch at the
123 /// provisional scale and the freeze simply leaves it alone, which is the
124 /// pre-#919 behaviour for that batch rather than a correctness cliff.
125 #[allow(clippy::too_many_arguments)]
126 pub(super) fn stage(
127 &mut self,
128 gpu: &dyn GpuBackend,
129 k: DevicePtr,
130 v: DevicePtr,
131 num_tokens: u32,
132 elems_per_token: usize,
133 target: &Fp8KvWriteTarget,
134 capacity_tokens: usize,
135 stream: u64,
136 ) -> Result<bool> {
137 if self.capacity_tokens == 0 {
138 self.k = gpu.alloc(capacity_tokens * elems_per_token * BF16)?;
139 self.v = gpu.alloc(capacity_tokens * elems_per_token * BF16)?;
140 self.slots = gpu.alloc(capacity_tokens * SLOT)?;
141 self.capacity_tokens = capacity_tokens;
142 self.elems_per_token = elems_per_token;
143 }
144 let n = num_tokens as usize;
145 if elems_per_token != self.elems_per_token || self.used_tokens + n > self.capacity_tokens {
146 return Ok(false);
147 }
148
149 let row = elems_per_token * BF16;
150 let dst_off = self.used_tokens * row;
151 copy_rows(
152 gpu,
153 k,
154 target.key_stride,
155 self.k.offset(dst_off),
156 row,
157 n,
158 stream,
159 )?;
160 copy_rows(
161 gpu,
162 v,
163 target.value_stride,
164 self.v.offset(dst_off),
165 row,
166 n,
167 stream,
168 )?;
169 gpu.copy_d2d_async(
170 target.slot,
171 self.slots.offset(self.used_tokens * SLOT),
172 n * SLOT,
173 stream,
174 )?;
175
176 self.batches.push(StagedBatch {
177 offset_tokens: self.used_tokens,
178 num_tokens,
179 });
180 self.used_tokens += n;
181 Ok(true)
182 }
183
184 /// Requantize every staged batch at the frozen scale, in write order, then
185 /// release the staging buffers. Returns the number of tokens rewritten.
186 #[allow(clippy::too_many_arguments)]
187 pub(super) fn replay_and_release(
188 &mut self,
189 gpu: &dyn GpuBackend,
190 target: &Fp8KvWriteTarget,
191 num_kv_heads: u32,
192 head_dim: u32,
193 k_scale: f32,
194 v_scale: f32,
195 stream: u64,
196 ) -> Result<usize> {
197 let rewritten = self.used_tokens;
198 let row = self.elems_per_token * BF16;
199 for batch in std::mem::take(&mut self.batches) {
200 let off = batch.offset_tokens;
201 ops::reshape_and_cache_fp8(
202 gpu,
203 target.kernel,
204 self.k.offset(off * row),
205 self.v.offset(off * row),
206 target.k_pool,
207 target.v_pool,
208 self.slots.offset(off * SLOT),
209 batch.num_tokens,
210 num_kv_heads,
211 head_dim,
212 target.block_size,
213 k_scale,
214 v_scale,
215 // The staging buffers are packed, whatever the live source
216 // strides were.
217 self.elems_per_token as u32,
218 self.elems_per_token as u32,
219 target.cache_stride,
220 stream,
221 )?;
222 }
223 self.release(gpu)?;
224 Ok(rewritten)
225 }
226
227 /// Free the staging buffers. Idempotent.
228 pub(super) fn release(&mut self, gpu: &dyn GpuBackend) -> Result<()> {
229 if self.capacity_tokens == 0 {
230 return Ok(());
231 }
232 gpu.free(self.k)?;
233 gpu.free(self.v)?;
234 gpu.free(self.slots)?;
235 *self = Self::default();
236 Ok(())
237 }
238}
239
240/// Copy `rows` rows of `row_bytes` from a source with `src_stride` ELEMENTS
241/// between rows into a packed destination.
242fn copy_rows(
243 gpu: &dyn GpuBackend,
244 src: DevicePtr,
245 src_stride: u32,
246 dst: DevicePtr,
247 row_bytes: usize,
248 rows: usize,
249 stream: u64,
250) -> Result<()> {
251 let src_pitch = src_stride as usize * BF16;
252 if src_pitch == row_bytes {
253 return gpu.copy_d2d_async(src, dst, row_bytes * rows, stream);
254 }
255 gpu.copy_d2d_2d_async(src, src_pitch, dst, row_bytes, row_bytes, rows, stream)
256}