Skip to content

Commit 16da4e2

Browse files
committed
Clear released snapshot-alias pages before their first KV use
Snapshot rows virtually alias KV pages. When the last snapshot owning a page is released, the page returns to the free list still holding row bytes. A request that later receives it as a partially filled tail page reads those bytes past its last token; in FP8 KV some decode as NaN, masked attention turns 0 x NaN into non-finite logits, and the commit reports INVALID_SPECULATION (device error 6) under KV pressure. The scheduler now marks a page when its last alias is released. When commit first moves a marked page from a request's reservation into its KV pages, the submission carries that page id and the mark is cleared. Every executor zeroes the listed pages in each local cache plane with cuMemsetD8Async on the model stream, after the input copies and before the step's graph. Tensor-parallel PREPARE frames carry the list to every rank (handshake layout 7), and peers validate it before READY. Pages are not cleared at release, where inbound host restores could still be writing, or at reservation, which may be returned unused. Inbound restore pages are overwritten whole and never enter a clear list; reservations, cancellations, and snapshot reuse keep the mark. The change is independent of model and cache scheme. Validation: - Scheduler unit test: released alias pages are listed exactly once, when a request first commits them, and survive cancellation and snapshot reuse. - TP wire tests: cleared pages round-trip through PREPARE, and ids outside the cache are rejected. - CUDA test on an RTX PRO 6000 (SM120) using real VMM aliases across three planes: FP8 NaN bytes written through released rows land in freed KV pages and stay inside their alias ranges. clear_pages zeroes exactly the freed pages and leaves other pages and the live row byte-identical. - Helios lib tests: the same 9 frontend, engram, and package-cache failures occur on unmodified cc696a6 in the test environment. - flare release build and heliopause check pass. - Stock Qwen3.8 Flash-Next TP2 under KV pressure (--memory-limit 49.5G, 143 alternating 16k/32k prompts): stock passes 143/143. This model's QSA attention does not read partially filled FP8 tail pages, so the run does not reproduce #116. Limitations: the #116 end-to-end reproduction on dense Qwen3.8-27B has not been run. Exporting qwen3.8-27b at 6fa01a4 failed in "Warm compiled forwards" (compiled forward has unbound static Inductor kernels). No throughput comparison has been measured. Gemma 4 (#61) is untested. Addresses #116. Design-Signed-Off-By: Cursor Agent (Auto); design-review=1c01
1 parent cc696a6 commit 16da4e2

17 files changed

Lines changed: 350 additions & 53 deletions

File tree

‎docs/runtime-package-design.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -598,6 +598,9 @@ and page capacities that all ranks can back.
598598
An owned snapshot excludes all conflicting page ids; overlapping exclusions
599599
are reference-counted. A snapshot fitting entirely in execution-allocation
600600
alignment padding excludes no pages. Execution rows never alias KV pages.
601+
A page whose last exclusion is released still holds snapshot row bytes. Before
602+
a request first writes it as KV, every rank clears that page's local cache
603+
planes on the model stream ahead of the step's graph.
601604

602605
Admission reserves the request's KV pages. System, prompt, and terminal snapshots
603606
acquire backing at their capture boundaries, evicting unpinned checkpoints when

‎helios/src/cuda.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ type HostFree = unsafe extern "C" fn(Handle) -> c_int;
2424
type Upload = unsafe extern "C" fn(u64, *const c_void, usize, Handle) -> c_int;
2525
type Download = unsafe extern "C" fn(Handle, u64, usize, Handle) -> c_int;
2626
type Download2D = unsafe extern "C" fn(*const Copy2D, Handle) -> c_int;
27+
type Fill = unsafe extern "C" fn(u64, u8, usize, Handle) -> c_int;
2728
type EventCreate = unsafe extern "C" fn(*mut Handle, c_uint) -> c_int;
2829
type EventRecord = unsafe extern "C" fn(Handle, Handle) -> c_int;
2930
type EventElapsed = unsafe extern "C" fn(*mut f32, Handle, Handle) -> c_int;
@@ -158,6 +159,7 @@ struct Driver {
158159
upload: Upload,
159160
download: Download,
160161
download_2d: Download2D,
162+
fill: Fill,
161163
event_create: EventCreate,
162164
event_record: EventRecord,
163165
event_elapsed: EventElapsed,
@@ -222,6 +224,7 @@ impl Driver {
222224
upload: symbol!("cuMemcpyHtoDAsync_v2"),
223225
download: symbol!("cuMemcpyDtoHAsync_v2"),
224226
download_2d: symbol!("cuMemcpy2DAsync_v2"),
227+
fill: symbol!("cuMemsetD8Async"),
225228
event_create: symbol!("cuEventCreate"),
226229
event_record: symbol!("cuEventRecord"),
227230
event_elapsed: symbol!("cuEventElapsedTime"),
@@ -910,6 +913,13 @@ impl Stream {
910913
})
911914
}
912915

916+
/// No queued or concurrent work may access the range until this stream passes the fill.
917+
pub(crate) unsafe fn zero(&self, device: u64, nbytes: usize) -> Result<()> {
918+
self.driver.check("cuMemsetD8Async", unsafe {
919+
(self.driver.fill)(device, 0, nbytes, self.handle as Handle)
920+
})
921+
}
922+
913923
pub(crate) fn synchronize(&self) -> Result<()> {
914924
self.driver.check("cuStreamSynchronize", unsafe {
915925
(self.driver.stream_synchronize)(self.handle as Handle)

‎helios/src/execution.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ mod threaded;
2020
pub use rank_stream::RankStream;
2121
pub use ranks::RankPeers;
2222
mod transport;
23-
pub(crate) use transport::validate_speculative_survival;
23+
pub(crate) use transport::{clear_pages, validate_speculative_survival};
2424
mod worker;
2525
pub use parallel::{TensorParallelExecutor, serve_rank};
2626
pub(crate) use ranks::RankExecutor;

‎helios/src/execution/parallel.rs‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ mod wire;
1313
use std::{time::Instant, collections::VecDeque};
1414

1515
use super::{ExecutionManifest, FrameCompletion, GraphExecutor, OwnedUpload, RankStream, Upload};
16-
use crate::{Error, Result, Shutdown, cache::Fence};
16+
use crate::{Error, Result, Shutdown, cache::{Fence, PageId}};
1717
use wire::*;
1818

1919
struct Peer {
@@ -25,6 +25,7 @@ struct Pending {
2525
graph: usize,
2626
grammar: bool,
2727
uploads: Vec<OwnedUpload>,
28+
cleared_pages: Vec<PageId>,
2829
dispatched: Instant,
2930
committed: bool,
3031
queued: bool,
@@ -221,6 +222,7 @@ impl TensorParallelExecutor {
221222
&mut self,
222223
graph: usize,
223224
uploads: &[Upload<'_>],
225+
cleared_pages: &[PageId],
224226
grammar: bool,
225227
) -> Result<Option<Fence>> {
226228
if let Some(error) = &self.poison {
@@ -232,7 +234,7 @@ impl TensorParallelExecutor {
232234
let barrier = grammar || !matches!(self.local.manifest().graphs.get(graph).map(|g| g.kind),
233235
Some(super::GraphKind::Decode | super::GraphKind::Speculative));
234236
if !self.can_submit(barrier) { return Ok(None); }
235-
let Some(fence) = self.local.validate_submission(graph, uploads, grammar)? else {
237+
let Some(fence) = self.local.validate_submission(graph, uploads, cleared_pages, grammar)? else {
236238
return Ok(None);
237239
};
238240
let queued = !self.pending.is_empty();
@@ -241,10 +243,11 @@ impl TensorParallelExecutor {
241243
uploads: uploads.iter().map(|u| OwnedUpload {
242244
buffer: u.buffer.to_owned(), offset: u.offset, bytes: u.bytes.to_vec(),
243245
}).collect(),
246+
cleared_pages: cleared_pages.to_vec(),
244247
ready: vec![false; self.peers.len()],
245248
completions: (0..=self.peers.len()).map(|_| None).collect(),
246249
});
247-
let result = wire::prepare(self.local.manifest(), fence, graph, grammar, uploads, |tag, bytes| {
250+
let result = wire::prepare(self.local.manifest(), fence, graph, grammar, uploads, cleared_pages, |tag, bytes| {
248251
for peer in &mut self.peers { peer.channel.send(tag, bytes)?; }
249252
Ok(())
250253
});
@@ -318,13 +321,14 @@ impl TensorParallelExecutor {
318321
if frame.committed { continue; }
319322
if !frame.ready.iter().all(|ready| *ready) { break; }
320323
let uploads = frame.uploads();
321-
if self.local.validate_submission(frame.graph, &uploads, frame.grammar)? != Some(frame.fence) {
324+
if self.local.validate_submission(frame.graph, &uploads, &frame.cleared_pages, frame.grammar)?
325+
!= Some(frame.fence) {
322326
return Err(invalid("rank-zero reservation changed before commit"));
323327
}
324328
for peer in &mut self.peers {
325329
peer.channel.send(COMMIT, &identity(frame.fence, frame.graph))?;
326330
}
327-
if self.local.submit(frame.graph, &uploads, frame.grammar)? != Some(frame.fence) {
331+
if self.local.submit(frame.graph, &uploads, &frame.cleared_pages, frame.grammar)? != Some(frame.fence) {
328332
return Err(invalid("rank-zero submission changed after coordinated validation"));
329333
}
330334
frame.committed = true;
@@ -534,7 +538,7 @@ pub unsafe fn serve_rank(
534538
}
535539
executor.finish_startup()?;
536540
channel.send(STARTED, &[])?;
537-
let mut prepared: Option<(Fence, usize, bool, Vec<OwnedUpload>)> = None;
541+
let mut prepared: Option<(Fence, usize, bool, Vec<OwnedUpload>, Vec<PageId>)> = None;
538542
let mut prepared_at: Option<Instant> = None;
539543
let mut started = VecDeque::new();
540544
loop {
@@ -558,30 +562,30 @@ pub unsafe fn serve_rank(
558562
}
559563
PREPARE | PREPARE_MORE => {
560564
if prepared.is_some() { return Err(invalid("rank already has a prepared frame")); }
561-
let (fence, graph, grammar, owned) =
565+
let (fence, graph, grammar, owned, cleared) =
562566
receive_prepare(&mut channel, executor.manifest(), tag, bytes)?;
563567
let uploads: Vec<_> = owned.iter().map(|upload| Upload {
564568
buffer: &upload.buffer, offset: upload.offset, bytes: &upload.bytes,
565569
}).collect();
566-
if executor.validate_submission(graph, &uploads, grammar)? != Some(fence) {
570+
if executor.validate_submission(graph, &uploads, &cleared, grammar)? != Some(fence) {
567571
return Err(invalid("rank submission sequence or capacity differs"));
568572
}
569573
channel.send(READY, &identity(fence, graph))?;
570-
prepared = Some((fence, graph, grammar, owned));
574+
prepared = Some((fence, graph, grammar, owned, cleared));
571575
prepared_at = Some(Instant::now());
572576
}
573577
COMMIT => {
574578
let (fence, graph) = read_identity(&bytes)?;
575579
prepared_at = None;
576-
let (expected, expected_graph, grammar, owned) = prepared.take()
580+
let (expected, expected_graph, grammar, owned, cleared) = prepared.take()
577581
.ok_or_else(|| invalid("rank commit arrived without preparation"))?;
578582
if fence != expected || graph != expected_graph {
579583
return Err(invalid("rank commit identity differs"));
580584
}
581585
let uploads: Vec<_> = owned.iter().map(|upload| Upload {
582586
buffer: &upload.buffer, offset: upload.offset, bytes: &upload.bytes,
583587
}).collect();
584-
if executor.submit(graph, &uploads, grammar)? != Some(fence) {
588+
if executor.submit(graph, &uploads, &cleared, grammar)? != Some(fence) {
585589
return Err(invalid("rank lost its prepared submission"));
586590
}
587591
started.push_back(Instant::now());

‎helios/src/execution/parallel/tests.rs‎

Lines changed: 45 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,7 @@ fn prepare_packets_preserve_large_and_partial_uploads_on_both_transports() {
217217
manifest.graphs[0].input_bindings[2].nbytes = buffer.nbytes;
218218
manifest.input_frame_bytes = 16 + buffer.nbytes;
219219
manifest.validate().unwrap();
220-
assert_eq!(u64::from_le_bytes(layout(&manifest)[..8].try_into().unwrap()), 6);
220+
assert_eq!(u64::from_le_bytes(layout(&manifest)[..8].try_into().unwrap()), 7);
221221
for (sender, receiver) in [crate::execution::rank_stream::test_pair().unwrap(),
222222
(client.into(), server.into())] {
223223
let mut sender = Channel::new(sender).unwrap();
@@ -230,7 +230,7 @@ fn prepare_packets_preserve_large_and_partial_uploads_on_both_transports() {
230230
buffer: "tokens", offset: 8 + index * 80 * 1024 * 1024 / parts, bytes,
231231
}).collect();
232232
let mut packets = 0;
233-
wire::prepare(&manifest, Fence(7), 0, true, &uploads, |tag, data| {
233+
wire::prepare(&manifest, Fence(7), 0, true, &uploads, &[], |tag, data| {
234234
assert!(data.len() <= 64 * 1024 * 1024);
235235
assert_eq!(tag, if packets == 0 { PREPARE_MORE } else { PREPARE });
236236
packets += 1;
@@ -240,9 +240,10 @@ fn prepare_packets_preserve_large_and_partial_uploads_on_both_transports() {
240240
sender.send(DRAIN, &[]).unwrap();
241241
});
242242
let (tag, packet) = receiver.wait(false).unwrap();
243-
let (fence, graph, grammar, owned) =
243+
let (fence, graph, grammar, owned, cleared) =
244244
receive_prepare(&mut receiver, &manifest, tag, packet).unwrap();
245245
assert_eq!((fence, graph, grammar), (Fence(7), 0, true));
246+
assert!(cleared.is_empty());
246247
assert!(owned.len() > parts);
247248
let mut offset = 8;
248249
for upload in &owned {
@@ -272,16 +273,53 @@ fn small_prepare_preserves_one_packet_and_upload_extents() {
272273
let (sender, receiver) = crate::execution::rank_stream::test_pair().unwrap();
273274
let mut sender = Channel::new(sender).unwrap();
274275
let mut receiver = Channel::new(receiver).unwrap();
275-
wire::prepare(&manifest, Fence(3), 0, false, &input, |tag, data| {
276+
wire::prepare(&manifest, Fence(3), 0, false, &input, &[], |tag, data| {
276277
assert_eq!(tag, PREPARE);
277278
sender.send(tag, data)
278279
}).unwrap();
279280
let (tag, packet) = receiver.wait(false).unwrap();
280-
let (fence, graph, grammar, owned) = receive_prepare(&mut receiver, &manifest, tag, packet).unwrap();
281+
let (fence, graph, grammar, owned, cleared) = receive_prepare(&mut receiver, &manifest, tag, packet).unwrap();
281282
assert_eq!((fence, graph, grammar), (Fence(3), 0, false));
282283
assert_eq!(owned.len(), 1);
283284
assert_eq!(owned[0].offset, 8);
284285
assert_eq!(owned[0].bytes, bytes);
286+
assert!(cleared.is_empty());
287+
}
288+
289+
#[test]
290+
fn prepare_carries_cleared_pages_to_peers_within_the_cache() {
291+
use crate::cache::PageId;
292+
let mut manifest = super::super::tests::manifest();
293+
let plane = |name: &str| flare_selene::Plane { name: name.into(), address: 0, unit_bytes: 256 };
294+
manifest.cache = Some(flare_selene::CacheGeometry {
295+
page_tokens: 64,
296+
page_capacity: 16,
297+
row_capacity: 2,
298+
row: plane("rows"),
299+
planes: vec![plane("keys"), plane("values")],
300+
snapshot_pages: vec![8..16],
301+
});
302+
let bytes = 7_i64.to_le_bytes();
303+
let input = [Upload { buffer: "tokens", offset: 8, bytes: &bytes }];
304+
let (sender, receiver) = crate::execution::rank_stream::test_pair().unwrap();
305+
let mut sender = Channel::new(sender).unwrap();
306+
let mut receiver = Channel::new(receiver).unwrap();
307+
let pages = [PageId(9), PageId(10), PageId(15)];
308+
wire::prepare(&manifest, Fence(4), 0, false, &input, &pages, |tag, data| {
309+
assert_eq!(tag, PREPARE);
310+
sender.send(tag, data)
311+
}).unwrap();
312+
let (tag, packet) = receiver.wait(false).unwrap();
313+
let (fence, _, _, owned, cleared) = receive_prepare(&mut receiver, &manifest, tag, packet).unwrap();
314+
assert_eq!(fence, Fence(4));
315+
assert_eq!(owned.len(), 1);
316+
assert_eq!(cleared, pages);
317+
wire::prepare(&manifest, Fence(5), 0, false, &[], &[PageId(16)], |tag, data| {
318+
sender.send(tag, data)
319+
}).unwrap();
320+
let (tag, packet) = receiver.wait(false).unwrap();
321+
let error = receive_prepare(&mut receiver, &manifest, tag, packet).err().unwrap().to_string();
322+
assert!(error.contains("outside the cache"), "{error}");
285323
}
286324

287325
#[test]
@@ -533,7 +571,7 @@ fn logprob_uploads_budget_count_and_prefix_wire_bytes_separately() {
533571
let manifest = super::super::tests::with_logprobs(crate::execution::MAX_LOGPROB_TOKEN_IDS);
534572
let measure = |uploads: &[Upload<'_>]| {
535573
let mut size = 0;
536-
wire::prepare(&manifest, Fence(1), 0, false, uploads, |_, bytes| {
574+
wire::prepare(&manifest, Fence(1), 0, false, uploads, &[], |_, bytes| {
537575
size += bytes.len();
538576
Ok(())
539577
})
@@ -577,7 +615,7 @@ fn two_frame_window_retains_preparing_and_unconsumed_completions() {
577615
let mut manifest = crate::execution::tests::manifest();
578616
manifest.graphs[0].kind = crate::execution::GraphKind::Decode;
579617
let frame = |sequence| Pending {
580-
fence: Fence(sequence), graph: 0, grammar: false, uploads: Vec::new(),
618+
fence: Fence(sequence), graph: 0, grammar: false, uploads: Vec::new(), cleared_pages: Vec::new(),
581619
dispatched: Instant::now(), committed: false, queued: sequence > 1,
582620
ready: vec![false; 3], completions: (0..4).map(|_| None).collect(),
583621
};

‎helios/src/execution/parallel/wire.rs‎

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
use std::io::{ErrorKind, Read, Write};
88
use std::time::Duration;
99

10-
use crate::cache::Fence;
10+
use crate::cache::{Fence, PageId};
1111
use crate::execution::{ExecutionManifest, FrameCompletion, OwnedUpload, RankStream, Upload};
1212
use crate::{Error, ExecutionError, Failure, Result};
1313

@@ -167,7 +167,7 @@ impl<'a> Reader<'a> {
167167
pub(super) fn layout(manifest: &ExecutionManifest) -> Vec<u8> {
168168
let mut bytes = Vec::new();
169169
for value in [
170-
6,
170+
7,
171171
manifest.abi_version as usize,
172172
manifest.max_active,
173173
manifest.input_frame_bytes,
@@ -237,15 +237,21 @@ pub(super) fn layout(manifest: &ExecutionManifest) -> Vec<u8> {
237237
}
238238

239239
/// Split validated uploads into bounded packets while preserving element alignment.
240+
/// The terminal packet carries the cleared page ids.
240241
pub(super) fn prepare(
241242
manifest: &ExecutionManifest,
242243
fence: Fence,
243244
graph: usize,
244245
grammar: bool,
245246
uploads: &[Upload<'_>],
247+
cleared_pages: &[PageId],
246248
mut send: impl FnMut(u8, &[u8]) -> Result<()>,
247249
) -> Result<()> {
248-
let capacity = uploads.iter().fold(32usize, |size, upload| {
250+
let clears = 8 * (1 + cleared_pages.len());
251+
if 32 + clears > MAX_PACKET {
252+
return Err(invalid("cleared pages exceed the rank control packet"));
253+
}
254+
let capacity = uploads.iter().fold(32 + clears, |size, upload| {
249255
size.saturating_add(24).saturating_add(upload.bytes.len())
250256
}).min(MAX_PACKET);
251257
let mut packet = Vec::with_capacity(capacity);
@@ -283,6 +289,13 @@ pub(super) fn prepare(
283289
bytes = &bytes[size..];
284290
}
285291
}
292+
if packet.len() + clears > MAX_PACKET {
293+
flush(&mut packet, &mut count, PREPARE_MORE)?;
294+
}
295+
number(&mut packet, cleared_pages.len());
296+
for page in cleared_pages {
297+
number(&mut packet, page.0 as usize);
298+
}
286299
flush(&mut packet, &mut count, PREPARE)
287300
}
288301

@@ -291,7 +304,7 @@ pub(super) fn receive_prepare(
291304
manifest: &ExecutionManifest,
292305
mut tag: u8,
293306
mut bytes: Vec<u8>,
294-
) -> Result<(Fence, usize, bool, Vec<OwnedUpload>)> {
307+
) -> Result<(Fence, usize, bool, Vec<OwnedUpload>, Vec<PageId>)> {
295308
let mut identity = None;
296309
let mut remaining = manifest.input_frame_bytes;
297310
let mut owned = Vec::new();
@@ -313,14 +326,28 @@ pub(super) fn receive_prepare(
313326
}
314327
identity = Some(received);
315328
let part = uploads(manifest, &mut reader, &mut remaining)?;
316-
reader.end()?;
317329
if tag == PREPARE_MORE && part.is_empty() {
318330
return Err(invalid("empty rank prepare continuation"));
319331
}
320332
owned.extend(part);
321333
if tag == PREPARE {
322-
return Ok((fence, graph, grammar, owned));
334+
let count = reader.number()?;
335+
let capacity = manifest.cache.as_ref().map_or(0, |cache| cache.page_capacity);
336+
if count > reader.0.len() / 8 {
337+
return Err(invalid("invalid rank cleared page count"));
338+
}
339+
let mut cleared = Vec::with_capacity(count);
340+
for _ in 0..count {
341+
let page = reader.number()?;
342+
if page as u64 >= capacity {
343+
return Err(invalid("rank cleared page id is outside the cache"));
344+
}
345+
cleared.push(PageId(page as u32));
346+
}
347+
reader.end()?;
348+
return Ok((fence, graph, grammar, owned, cleared));
323349
}
350+
reader.end()?;
324351
(tag, bytes) = channel.wait(false)?;
325352
}
326353
}

‎helios/src/execution/ranks.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
use super::{
1010
ExecutionManifest, FrameCompletion, GraphExecutor, RankStream, TensorParallelExecutor, Upload,
1111
};
12-
use crate::{Result, Shutdown, cache::Fence};
12+
use crate::{Result, Shutdown, cache::{Fence, PageId}};
1313

1414
pub enum RankPeers {
1515
Processes(Vec<RankStream>),
@@ -48,12 +48,13 @@ impl RankExecutor {
4848
&mut self,
4949
graph: usize,
5050
uploads: &[Upload<'_>],
51+
cleared_pages: &[PageId],
5152
grammar: bool,
5253
) -> Result<Option<Fence>> {
5354
match self {
54-
Self::Local(executor) => executor.submit(graph, uploads, grammar),
55-
Self::Parallel(executor) => executor.submit(graph, uploads, grammar),
56-
Self::Threaded(executor) => executor.submit(graph, uploads, grammar),
55+
Self::Local(executor) => executor.submit(graph, uploads, cleared_pages, grammar),
56+
Self::Parallel(executor) => executor.submit(graph, uploads, cleared_pages, grammar),
57+
Self::Threaded(executor) => executor.submit(graph, uploads, cleared_pages, grammar),
5758
}
5859
}
5960

0 commit comments

Comments
 (0)