From 53bba32a52aae92ee162d17e0891e5e91ea8aa64 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 8 Sep 2026 18:52:00 +0200 Subject: [PATCH] fix: make DFlash construction transactional --- CHANGELOG.md | 1 + crates/hipfire-runtime/map.md | 4 +- crates/hipfire-runtime/src/dflash.rs | 469 ++++++++++++++++++++++++--- 3 files changed, 420 insertions(+), 54 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b5357e97..180c75f02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ - DFlash prompt-cache repair on terminal overshoot (`RepairForTerminal`) (#695). - Template-aware primer splice (#692). - Transactional DFlash constructors with emitter rollback (#691). +- DFlash weight and scratch constructors now roll back late failures for immediate retry. - MQ-V2 prefill admit rule (#690). - gfx1100 DFlash launch fusion and split-K residual tiers (#702 body, S1–S8). - Dense-TP prefill chunking equals arch batch × tp (#725). diff --git a/crates/hipfire-runtime/map.md b/crates/hipfire-runtime/map.md index 237e40a68..2afe13d93 100644 --- a/crates/hipfire-runtime/map.md +++ b/crates/hipfire-runtime/map.md @@ -39,7 +39,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/cpu_router.rs`](src/cpu_router.rs) | 200 | 4 | 3 | | [`src/ddtree.rs`](src/ddtree.rs) | 2,046 | 17 | 24 | | [`src/device_mesh.rs`](src/device_mesh.rs) | 580 | 21 | 8 | -| [`src/dflash.rs`](src/dflash.rs) | 3,979 | 47 | 5 | +| [`src/dflash.rs`](src/dflash.rs) | 4,344 | 47 | 9 | | [`src/dflash_generic.rs`](src/dflash_generic.rs) | 1,382 | 3 | 13 | | [`src/dspark_block_controller.rs`](src/dspark_block_controller.rs) | 442 | 0 | 10 | | [`src/dspark_core.rs`](src/dspark_core.rs) | 1,773 | 11 | 0 | @@ -160,6 +160,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 60 modules · 54,702 lines · 926 public items · 629 tests · 132 examples +- 60 modules · 55,067 lines · 926 public items · 633 tests · 132 examples diff --git a/crates/hipfire-runtime/src/dflash.rs b/crates/hipfire-runtime/src/dflash.rs index 16c633e5e..35b3da368 100644 --- a/crates/hipfire-runtime/src/dflash.rs +++ b/crates/hipfire-runtime/src/dflash.rs @@ -351,6 +351,20 @@ fn hfq_tensor_f32( gpu.upload_f32(&f32_data, &shape) } +/// Upload a DFlash raw weight through the reusable pool. The global raw-upload +/// helper allocates directly from HIP, so an immediate constructor retry cannot +/// reclaim a successfully staged predecessor. This local path also returns the +/// allocation if the host-to-device copy itself fails. +fn upload_raw_weight(gpu: &mut Gpu, data: &[u8], shape: &[usize]) -> HipResult { + let mut tensor = gpu.alloc_tensor(&[data.len()], DType::Raw)?; + if let Err(error) = gpu.hip.memcpy_htod(&tensor.buf, data) { + let _ = gpu.free_tensor(tensor); + return Err(error); + } + tensor.shape = shape.to_vec(); + Ok(tensor) +} + /// Load a matrix tensor as a `WeightTensor` carrying its native dtype. /// Supported quant_types: /// 1 (F16) → lifted to F32 on GPU (legacy path). @@ -395,7 +409,7 @@ fn hfq_weight( m * k * 2, "dflash {name} F16 byte-size mismatch" ); - let buf = gpu.upload_raw(data, &[m * k])?; + let buf = upload_raw_weight(gpu, data, &[m * k])?; Ok::(WeightTensor { buf, gpu_dtype: DType::F16, @@ -443,7 +457,7 @@ fn hfq_weight( 13 => { // MQ4-G256: 136 bytes per 256 weights. The buffer is opaque to // the engine; the gemm_hfq4g256 kernel reads it directly. - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ4G256, @@ -457,7 +471,7 @@ fn hfq_weight( 15 => { // MQ6-G256: 200 bytes per 256 weights. Same opaque-buffer pattern // as MQ4/MQ3; dispatch rotates activations and calls HFQ6 GEMM. - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ6G256, @@ -472,7 +486,7 @@ fn hfq_weight( // MQ3-G256: 104 bytes per 256 weights. Same opaque-buffer pattern // as MQ4. Dispatch path (`gemm_dispatch`) routes through // `rotate_x_mq_batched` + `gemm_hfq3g256_batched_lmhead`. - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ3G256, @@ -498,7 +512,7 @@ fn hfq_weight( data.len() ); } - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ4G256V2, @@ -522,7 +536,7 @@ fn hfq_weight( data.len() ); } - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ6G256V2, @@ -546,7 +560,7 @@ fn hfq_weight( data.len() ); } - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ5G256V2, @@ -570,7 +584,7 @@ fn hfq_weight( data.len() ); } - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ3G256V2, @@ -594,7 +608,7 @@ fn hfq_weight( data.len() ); } - let buf = gpu.upload_raw(data, &[data.len()])?; + let buf = upload_raw_weight(gpu, data, &[data.len()])?; Ok(WeightTensor { buf, gpu_dtype: DType::MQ2G256V2, @@ -658,39 +672,59 @@ impl DflashWeights { && self.successor_codebook.is_some() } pub fn load(gpu: &mut Gpu, hfq: &HfqFile, cfg: &DflashConfig) -> HipResult { - // Transactional construction. Every GPU tensor allocated below is - // recorded in `live_t` (plain F32) / `live_w` (weight + sidecars) and - // taken exactly once into its final owner; completed layers accumulate - // in `layers`. Any failure frees the completed layers plus every + Self::load_with_boundaries(gpu, hfq, cfg, || Ok(()), || Ok(())) + } + + fn load_with_boundaries( + gpu: &mut Gpu, + hfq: &HfqFile, + cfg: &DflashConfig, + mut after_owner: impl FnMut() -> HipResult<()>, + mut after_layer: impl FnMut() -> HipResult<()>, + ) -> HipResult { + // Transactional construction. Every GPU owner returned by the leaf + // loaders is recorded in `live_t` (plain F32) / `live_w` (weight + + // sidecars) and taken exactly once into its final owner; completed + // layers accumulate in `layers`. Any ordinary loader or injected + // boundary failure frees the completed layers plus every // recorded-but-unplaced tensor before returning Err — a bare `?` - // would leak them (`GpuTensor`/`DeviceBuffer` have no `Drop`). The - // leaf loaders (`hfq_weight`, `hfq_tensor_f32`) are single-alloc and - // need no cover. Mirrors the `or_free!` style in - // `hipfire-arch-qwen35`'s `load_dflash_state`. + // would leak them (`GpuTensor`/`DeviceBuffer` have no `Drop`). Mirrors + // the `or_free!` style in `hipfire-arch-qwen35`'s + // `load_dflash_state`. let mut live_t: Vec> = Vec::new(); let mut live_w: Vec> = Vec::new(); let mut layers: Vec = Vec::with_capacity(cfg.n_layers); + macro_rules! cleanup { + () => {{ + for l in layers.drain(..).rev() { + l.free_gpu(gpu); + } + for slot in live_w.iter_mut().rev() { + if let Some(w) = slot.take() { + w.free_all(gpu); + } + } + for slot in live_t.iter_mut().rev() { + if let Some(t) = slot.take() { + let _ = gpu.free_tensor(t); + } + } + }}; + } macro_rules! gt { ($e:expr) => {{ match $e { Ok(t) => { live_t.push(Some(t)); - live_t.len() - 1 + let index = live_t.len() - 1; + if let Err(e) = after_owner() { + cleanup!(); + return Err(e); + } + index } Err(e) => { - for l in layers.drain(..) { - l.free_gpu(gpu); - } - for slot in live_w.iter_mut() { - if let Some(w) = slot.take() { - w.free_all(gpu); - } - } - for slot in live_t.iter_mut() { - if let Some(t) = slot.take() { - let _ = gpu.free_tensor(t); - } - } + cleanup!(); return Err(e); } } @@ -701,22 +735,15 @@ impl DflashWeights { match $e { Ok(w) => { live_w.push(Some(w)); - live_w.len() - 1 + let index = live_w.len() - 1; + if let Err(e) = after_owner() { + cleanup!(); + return Err(e); + } + index } Err(e) => { - for l in layers.drain(..) { - l.free_gpu(gpu); - } - for slot in live_w.iter_mut() { - if let Some(w) = slot.take() { - w.free_all(gpu); - } - } - for slot in live_t.iter_mut() { - if let Some(t) = slot.take() { - let _ = gpu.free_tensor(t); - } - } + cleanup!(); return Err(e); } } @@ -930,6 +957,10 @@ impl DflashWeights { mlp_conv_base: mlp_conv_base.map(|j| take_t!(j)), mlp_conv_proj: mlp_conv_proj.map(|j| take_w!(j)), }); + if let Err(e) = after_layer() { + cleanup!(); + return Err(e); + } } // Selector: hidden_projection [rank, hidden] + two codebooks [vocab, rank] host-side @@ -1489,12 +1520,36 @@ impl DflashScratch { w_full: usize, max_ctx: usize, with_mq: bool, + ) -> HipResult { + Self::new_windowed_with_boundary( + gpu, + cfg, + max_block_size, + w, + w_full, + max_ctx, + with_mq, + &mut || Ok(()), + ) + } + + #[allow(clippy::too_many_arguments)] // Mirrors the public constructor plus a test-only boundary. + fn new_windowed_with_boundary( + gpu: &mut Gpu, + cfg: &DflashConfig, + max_block_size: usize, + w: usize, + w_full: usize, + max_ctx: usize, + with_mq: bool, + after_owner: &mut impl FnMut() -> HipResult<()>, ) -> HipResult { // DFlash2 all-sliding: every layer shares the same W ring. Skip the // last-layer full replacement/backfill path and keep the footprint at // Legacy-at-ctx=w for all layers. if cfg.all_layers_sliding { - let mut s = Self::new_with_mq(gpu, cfg, max_block_size, w, with_mq)?; + let mut s = + Self::new_with_mq_with_boundary(gpu, cfg, max_block_size, w, with_mq, after_owner)?; s.max_ctx_len = max_ctx; s.ctx_mode = DraftCtxMode::Windowed { w, w_full: w }; return Ok(s); @@ -1506,7 +1561,7 @@ impl DflashScratch { let w_full = w_full.max(w); // Base at ctx=w: SWA rings, (w+B) concat buffers, w-row target_hidden // ring — the entire draft footprint except the one long-reach layer. - let mut s = Self::new_with_mq(gpu, cfg, b, w, with_mq)?; + let mut s = Self::new_with_mq_with_boundary(gpu, cfg, b, w, with_mq, after_owner)?; // The last (full-attention) layer gets w_full-row rings + its own // concat pair; its w-sized base caches are freed. if let Some(k) = s.k_ctx_cached.pop() { @@ -1536,20 +1591,33 @@ impl DflashScratch { } }; } + macro_rules! boundary_or_free { + () => { + if let Err(e) = after_owner() { + s.free_gpu(gpu); + return Err(e); + } + }; + } s.k_full_cached = Some(alloc_or_free!(gpu.alloc_tensor(&[w_full * kvd], DType::F32))); + boundary_or_free!(); s.v_full_cached = Some(alloc_or_free!(gpu.alloc_tensor(&[w_full * kvd], DType::F32))); + boundary_or_free!(); s.k_cat_full = Some(alloc_or_free!( gpu.alloc_tensor(&[(w_full + b) * kvd], DType::F32) )); + boundary_or_free!(); s.v_cat_full = Some(alloc_or_free!( gpu.alloc_tensor(&[(w_full + b) * kvd], DType::F32) )); + boundary_or_free!(); // positions_k holds the last w_full context rows + the B noise rows // (the forward uploads only that suffix; every layer's span is one). // Allocate before freeing the old buffer so a failure still leaves // `s` intact for the error arm above. let new_positions_k = alloc_or_free!(gpu.alloc_tensor(&[w_full + b], DType::F32)); let _ = gpu.free_tensor(std::mem::replace(&mut s.positions_k, new_positions_k)); + boundary_or_free!(); // The ctx bound is the target's physical capacity, not the window — // l may cross w_full (the last layer's span just slides). s.max_ctx_len = max_ctx; @@ -1566,6 +1634,19 @@ impl DflashScratch { max_block_size: usize, max_ctx_len: usize, with_mq: bool, + ) -> HipResult { + Self::new_with_mq_with_boundary(gpu, cfg, max_block_size, max_ctx_len, with_mq, &mut || { + Ok(()) + }) + } + + fn new_with_mq_with_boundary( + gpu: &mut Gpu, + cfg: &DflashConfig, + max_block_size: usize, + max_ctx_len: usize, + with_mq: bool, + after_owner: &mut impl FnMut() -> HipResult<()>, ) -> HipResult { let b = max_block_size; let l = max_ctx_len; @@ -1591,10 +1672,19 @@ impl DflashScratch { match gpu.alloc_tensor($shape, $dtype) { Ok(t) => { live.push(Some(t)); - live.len() - 1 + let index = live.len() - 1; + if let Err(e) = after_owner() { + for slot in live.iter_mut().rev() { + if let Some(t) = slot.take() { + let _ = gpu.free_tensor(t); + } + } + return Err(e); + } + index } Err(e) => { - for slot in live.iter_mut() { + for slot in live.iter_mut().rev() { if let Some(t) = slot.take() { let _ = gpu.free_tensor(t); } @@ -1718,8 +1808,7 @@ impl DflashScratch { k_ctx_cached.push(take!(ik)); v_ctx_cached.push(take!(iv)); } - debug_assert!(live.iter().all(|s| s.is_none())); - Ok(DflashScratch { + let scratch = DflashScratch { max_block_size: b, max_ctx_len: l, @@ -1761,7 +1850,9 @@ impl DflashScratch { v_cat_full: None, draft_ffn_graphs, draft_ffn_warmed_up, - }) + }; + debug_assert!(live.iter().all(|s| s.is_none())); + Ok(scratch) } /// Reset the incremental-upload tracker for target_hidden. Call this @@ -3594,6 +3685,280 @@ mod ring_tests { } } +#[cfg(test)] +mod construction_tests { + use super::*; + use std::path::PathBuf; + + fn tiny_config() -> DflashConfig { + DflashConfig { + n_layers: 2, + hidden: 32, + intermediate: 64, + n_heads: 2, + n_kv_heads: 1, + head_dim: 16, + vocab_size: 64, + norm_eps: 1e-6, + rope_theta: 1_000_000.0, + block_size: 2, + mask_token_id: 0, + target_layer_ids: vec![0], + num_target_layers: 1, + declared_window: Some(2), + all_layers_sliding: false, + conv_group_size: Some(16), + conv_kernel_size: Some(2), + selector_rank: Some(8), + selector_top_k: Some(2), + } + } + + fn draft_path() -> PathBuf { + PathBuf::from(std::env::var_os("HOME").expect("HOME is required")) + .join(".hipfire/models/qwen35-27b-dflash-mq4.hfq") + } + + #[test] + #[ignore = "requires an AMD GPU and the canonical Qwen3.5-27B DFlash draft"] + fn late_weight_failure_reuses_every_staged_owner() { + let path = draft_path(); + assert!(path.is_file(), "missing DFlash fixture: {}", path.display()); + let hfq = HfqFile::open(&path).expect("open DFlash fixture"); + let cfg = DflashConfig::from_hfq(&hfq).expect("parse DFlash config"); + let mut gpu = Gpu::init().expect("GPU required for DFlash rollback"); + + let mut allocations = 0usize; + let warm = DflashWeights::load_with_boundaries( + &mut gpu, + &hfq, + &cfg, + || { + allocations += 1; + Ok(()) + }, + || Ok(()), + ) + .expect("warm DFlash weights"); + warm.free_gpu(&mut gpu); + let fresh_allocations = gpu.pool_stats().0; + + let mut attempted = 0usize; + let failure = DflashWeights::load_with_boundaries( + &mut gpu, + &hfq, + &cfg, + || { + attempted += 1; + if attempted == allocations { + Err(hip_bridge::HipError::new( + 2, + "injected failure after final DFlash weight owner", + )) + } else { + Ok(()) + } + }, + || Ok(()), + ); + match failure { + Err(error) => assert_eq!(error.code, 2), + Ok(weights) => { + weights.free_gpu(&mut gpu); + panic!("weight fault did not trigger"); + } + } + assert_eq!( + attempted, allocations, + "failure must follow the final owner" + ); + + let retry = DflashWeights::load(&mut gpu, &hfq, &cfg) + .expect("immediate weight retry after late failure"); + retry.free_gpu(&mut gpu); + assert_eq!( + gpu.pool_stats().0, + fresh_allocations, + "failed weight construction lost reusable allocations", + ); + gpu.drain_pool(); + } + + #[test] + #[ignore = "requires an AMD GPU and the canonical Qwen3.5-27B DFlash draft"] + fn late_layer_failure_reuses_every_completed_layer() { + let path = draft_path(); + assert!(path.is_file(), "missing DFlash fixture: {}", path.display()); + let hfq = HfqFile::open(&path).expect("open DFlash fixture"); + let cfg = DflashConfig::from_hfq(&hfq).expect("parse DFlash config"); + let mut gpu = Gpu::init().expect("GPU required for DFlash rollback"); + + let warm = DflashWeights::load(&mut gpu, &hfq, &cfg).expect("warm DFlash weights"); + warm.free_gpu(&mut gpu); + let fresh_allocations = gpu.pool_stats().0; + + let mut completed_layers = 0usize; + let failure = DflashWeights::load_with_boundaries( + &mut gpu, + &hfq, + &cfg, + || Ok(()), + || { + completed_layers += 1; + if completed_layers == cfg.n_layers { + Err(hip_bridge::HipError::new( + 2, + "injected failure after final DFlash layer", + )) + } else { + Ok(()) + } + }, + ); + match failure { + Err(error) => assert_eq!(error.code, 2), + Ok(weights) => { + weights.free_gpu(&mut gpu); + panic!("layer fault did not trigger"); + } + } + assert_eq!( + completed_layers, cfg.n_layers, + "failure must follow the final completed layer" + ); + + let retry = + DflashWeights::load(&mut gpu, &hfq, &cfg).expect("immediate retry after layer failure"); + retry.free_gpu(&mut gpu); + assert_eq!( + gpu.pool_stats().0, + fresh_allocations, + "failed layer construction lost reusable allocations", + ); + gpu.drain_pool(); + } + + #[test] + #[ignore = "requires an AMD GPU; exercises real base-scratch rollback and retry"] + fn late_base_scratch_failure_reuses_every_staged_owner() { + let cfg = tiny_config(); + let mut gpu = Gpu::init().expect("GPU required for DFlash rollback"); + + let mut allocations = 0usize; + let warm = + DflashScratch::new_with_mq_with_boundary(&mut gpu, &cfg, 2, 4, true, &mut || { + allocations += 1; + Ok(()) + }) + .expect("warm DFlash scratch"); + warm.free_gpu(&mut gpu); + + let mut attempted = 0usize; + let failure = + DflashScratch::new_with_mq_with_boundary(&mut gpu, &cfg, 2, 4, true, &mut || { + attempted += 1; + if attempted == allocations { + Err(hip_bridge::HipError::new( + 2, + "injected failure after final DFlash scratch owner", + )) + } else { + Ok(()) + } + }); + match failure { + Err(error) => assert_eq!(error.code, 2), + Ok(scratch) => { + scratch.free_gpu(&mut gpu); + panic!("base scratch fault did not trigger"); + } + } + assert_eq!( + attempted, allocations, + "failure must follow the final owner" + ); + let fresh_allocations = gpu.pool_stats().0; + + let retry = DflashScratch::new_with_mq(&mut gpu, &cfg, 2, 4, true) + .expect("immediate base-scratch retry"); + retry.free_gpu(&mut gpu); + assert_eq!( + gpu.pool_stats().0, + fresh_allocations, + "failed base-scratch construction lost reusable allocations", + ); + gpu.drain_pool(); + } + + #[test] + #[ignore = "requires an AMD GPU; exercises real window-extension rollback and retry"] + fn late_window_extension_failure_reuses_base_and_extensions() { + let cfg = tiny_config(); + let mut gpu = Gpu::init().expect("GPU required for DFlash rollback"); + + let mut allocations = 0usize; + let warm = DflashScratch::new_windowed_with_boundary( + &mut gpu, + &cfg, + 2, + 2, + 8, + 32, + true, + &mut || { + allocations += 1; + Ok(()) + }, + ) + .expect("warm windowed DFlash scratch"); + warm.free_gpu(&mut gpu); + + let mut attempted = 0usize; + let failure = DflashScratch::new_windowed_with_boundary( + &mut gpu, + &cfg, + 2, + 2, + 8, + 32, + true, + &mut || { + attempted += 1; + if attempted == allocations { + Err(hip_bridge::HipError::new( + 2, + "injected failure after final DFlash window owner", + )) + } else { + Ok(()) + } + }, + ); + match failure { + Err(error) => assert_eq!(error.code, 2), + Ok(scratch) => { + scratch.free_gpu(&mut gpu); + panic!("window fault did not trigger"); + } + } + assert_eq!( + attempted, allocations, + "failure must follow the final owner" + ); + let fresh_allocations = gpu.pool_stats().0; + + let retry = DflashScratch::new_windowed(&mut gpu, &cfg, 2, 2, 8, 32, true) + .expect("immediate windowed-scratch retry"); + retry.free_gpu(&mut gpu); + assert_eq!( + gpu.pool_stats().0, + fresh_allocations, + "failed window construction lost reusable allocations", + ); + gpu.drain_pool(); + } +} + // ─── Candidate selector (DFlash2 chain-only) ─────────────────────────────── /// Proposal returned by the DFlash2 selector. All rows flattened row-major.