diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b5357e97..707bc9bfc 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). +- Qwen35 prefill/decode scratch and per-layer weight construction retain actual owners until publication, reclaiming every staged allocation on failure so immediate retries reuse the pool. - 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-arch-qwen35/map.md b/crates/hipfire-arch-qwen35/map.md index 771f38d66..5f5812333 100644 --- a/crates/hipfire-arch-qwen35/map.md +++ b/crates/hipfire-arch-qwen35/map.md @@ -29,7 +29,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/dflash_verify_pm4.rs`](src/dflash_verify_pm4.rs) | 739 | 35 | 9 | | [`src/forward_slots.rs`](src/forward_slots.rs) | 3,154 | 14 | 3 | | [`src/grammar_config.rs`](src/grammar_config.rs) | 143 | 2 | 4 | -| [`src/layer_driver.rs`](src/layer_driver.rs) | 112 | 0 | 0 | +| [`src/layer_driver.rs`](src/layer_driver.rs) | 629 | 0 | 2 | | [`src/lib.rs`](src/lib.rs) | 121 | 19 | 0 | | [`src/mtp_compose.rs`](src/mtp_compose.rs) | 1,374 | 8 | 0 | | [`src/mtp_head.rs`](src/mtp_head.rs) | 2,616 | 32 | 2 | @@ -37,13 +37,13 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/mtp_spec.rs`](src/mtp_spec.rs) | 3,827 | 33 | 12 | | [`src/mtp_speculator.rs`](src/mtp_speculator.rs) | 522 | 3 | 0 | | [`src/paro_moe.rs`](src/paro_moe.rs) | 222 | 0 | 0 | -| [`src/qwen35/batch.rs`](src/qwen35/batch.rs) | 1,698 | 16 | 0 | +| [`src/qwen35/batch.rs`](src/qwen35/batch.rs) | 1,949 | 16 | 2 | | [`src/qwen35/config.rs`](src/qwen35/config.rs) | 1,643 | 41 | 21 | | [`src/qwen35/ep_batch.rs`](src/qwen35/ep_batch.rs) | 4,805 | 20 | 7 | | [`src/qwen35/forward.rs`](src/qwen35/forward.rs) | 6,251 | 31 | 12 | | [`src/qwen35/load.rs`](src/qwen35/load.rs) | 4,906 | 10 | 0 | | [`src/qwen35/prefill.rs`](src/qwen35/prefill.rs) | 10,126 | 12 | 49 | -| [`src/qwen35/weights.rs`](src/qwen35/weights.rs) | 1,978 | 43 | 10 | +| [`src/qwen35/weights.rs`](src/qwen35/weights.rs) | 2,020 | 43 | 10 | | [`src/qwen35.rs`](src/qwen35.rs) | 63 | 7 | 0 | | [`src/scheduler.rs`](src/scheduler.rs) | 142 | 3 | 4 | | [`src/serve_engine.rs`](src/serve_engine.rs) | 1,273 | 8 | 2 | @@ -97,6 +97,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 29 modules · 58,634 lines · 440 public items · 192 tests · 11 examples +- 29 modules · 59,444 lines · 440 public items · 196 tests · 11 examples diff --git a/crates/hipfire-arch-qwen35/src/layer_driver.rs b/crates/hipfire-arch-qwen35/src/layer_driver.rs index 41cd83619..cba8e6806 100644 --- a/crates/hipfire-arch-qwen35/src/layer_driver.rs +++ b/crates/hipfire-arch-qwen35/src/layer_driver.rs @@ -6,12 +6,85 @@ //! `WeightBackend`. `load_weights` (HFQ), `load_weights_paroquant` (PaRo), and //! `load_layer_into` (multi-GPU HFQ) all funnel through `load_layer`. +use crate::qwen35::weights::{free_moe_ffn_with, free_weight_with}; use crate::qwen35::{ DeltaNetLayerWeights, DeltaNetMoeLayerWeights, FullAttnLayerWeights, FullAttnMoeLayerWeights, LayerType, LayerWeights, MoeFfnWeights, Qwen35Config, }; use hip_bridge::HipResult; +use hipfire_runtime::llama::WeightTensor; use hipfire_runtime::weight_backend::WeightBackend; +use rdna_compute::GpuTensor; + +/// All owners allocated while one layer is being assembled. Every field stays +/// optional until publication so an error can drain only the owners that +/// actually exist. +#[derive(Default)] +struct PendingLayer { + attn_norm: Option, + wqkv: Option, + wz: Option, + w_alpha: Option, + w_beta: Option, + a_log: Option, + dt_bias: Option, + conv_weight: Option, + norm_weight: Option, + wo: Option, + wq: Option, + wk: Option, + wv: Option, + q_norm: Option, + k_norm: Option, + ffn_norm: Option, + w_gate: Option, + w_up: Option, + w_down: Option, + ffn: Option, +} + +impl PendingLayer { + fn cleanup(&mut self, b: &mut B) { + if let Some(ffn) = self.ffn.take() { + let mut free = |tensor: GpuTensor| b.free_tensor(tensor); + free_moe_ffn_with(ffn, &mut free); + } + for weight in [ + self.wqkv.take(), + self.wz.take(), + self.w_alpha.take(), + self.w_beta.take(), + self.wo.take(), + self.wq.take(), + self.wk.take(), + self.wv.take(), + self.w_gate.take(), + self.w_up.take(), + self.w_down.take(), + ] + .into_iter() + .flatten() + { + let mut free = |tensor: GpuTensor| b.free_tensor(tensor); + free_weight_with(weight, &mut free); + } + for tensor in [ + self.attn_norm.take(), + self.a_log.take(), + self.dt_bias.take(), + self.conv_weight.take(), + self.norm_weight.take(), + self.q_norm.take(), + self.k_norm.take(), + self.ffn_norm.take(), + ] + .into_iter() + .flatten() + { + b.free_tensor(tensor); + } + } +} /// Load one layer's weights. `load_moe` builds the MoE FFN block for MoE layers /// (format-specific: HFQ `load_moe_ffn` vs PaRo `paro_load_moe_ffn`), supplied by @@ -30,83 +103,527 @@ pub(crate) fn load_layer( let q_out_dim = config.n_heads * config.head_dim * 2; let kv_dim = config.n_kv_heads * config.head_dim; let o_in = config.n_heads * config.head_dim; + let mut pending = PendingLayer::default(); + + macro_rules! stage { + ($slot:ident, $load:expr) => { + match $load { + Ok(owner) => pending.$slot = Some(owner), + Err(err) => { + pending.cleanup(b); + return Err(err); + } + } + }; + } + macro_rules! take { + ($slot:ident) => { + pending.$slot.take().expect(concat!( + "load_layer: missing staged owner ", + stringify!($slot) + )) + }; + } + + let layer = match (config.layer_types[layer_idx], is_moe) { + (LayerType::LinearAttention, false) => { + stage!(attn_norm, b.norm("input_layernorm.weight", &[config.dim])); + stage!(wqkv, b.proj("linear_attn.in_proj_qkv", qkv_dim, config.dim)); + stage!(wz, b.proj("linear_attn.in_proj_z", d_inner, config.dim)); + stage!( + w_alpha, + b.proj( + "linear_attn.in_proj_a", + config.linear_num_value_heads, + config.dim, + ) + ); + stage!( + w_beta, + b.proj( + "linear_attn.in_proj_b", + config.linear_num_value_heads, + config.dim, + ) + ); + stage!( + a_log, + b.raw_f32("linear_attn.A_log", config.linear_num_value_heads) + ); + stage!( + dt_bias, + b.raw_f32("linear_attn.dt_bias", config.linear_num_value_heads) + ); + stage!( + conv_weight, + b.raw_f32( + "linear_attn.conv1d.weight", + qkv_dim * config.conv_kernel_dim, + ) + ); + stage!( + norm_weight, + b.raw_f32("linear_attn.norm.weight", config.linear_value_head_dim) + ); + stage!(wo, b.proj("linear_attn.out_proj", config.dim, d_inner)); + stage!( + ffn_norm, + b.norm("post_attention_layernorm.weight", &[config.dim]) + ); + stage!( + w_gate, + b.proj("mlp.gate_proj", config.hidden_dim, config.dim) + ); + stage!(w_up, b.proj("mlp.up_proj", config.hidden_dim, config.dim)); + stage!( + w_down, + b.proj("mlp.down_proj", config.dim, config.hidden_dim) + ); + LayerWeights::DeltaNet(DeltaNetLayerWeights { + attn_norm: take!(attn_norm), + wqkv: take!(wqkv), + wz: take!(wz), + w_alpha: take!(w_alpha), + w_beta: take!(w_beta), + a_log: take!(a_log), + dt_bias: take!(dt_bias), + conv_weight: take!(conv_weight), + norm_weight: take!(norm_weight), + wo: take!(wo), + ffn_norm: take!(ffn_norm), + w_gate: take!(w_gate), + w_up: take!(w_up), + w_down: take!(w_down), + }) + } + (LayerType::FullAttention, false) => { + stage!(attn_norm, b.norm("input_layernorm.weight", &[config.dim])); + stage!(wq, b.proj("self_attn.q_proj", q_out_dim, config.dim)); + stage!(wk, b.proj("self_attn.k_proj", kv_dim, config.dim)); + stage!(wv, b.proj("self_attn.v_proj", kv_dim, config.dim)); + stage!(wo, b.proj("self_attn.o_proj", config.dim, o_in)); + stage!( + q_norm, + b.norm("self_attn.q_norm.weight", &[config.head_dim]) + ); + stage!( + k_norm, + b.norm("self_attn.k_norm.weight", &[config.head_dim]) + ); + stage!( + ffn_norm, + b.norm("post_attention_layernorm.weight", &[config.dim]) + ); + stage!( + w_gate, + b.proj("mlp.gate_proj", config.hidden_dim, config.dim) + ); + stage!(w_up, b.proj("mlp.up_proj", config.hidden_dim, config.dim)); + stage!( + w_down, + b.proj("mlp.down_proj", config.dim, config.hidden_dim) + ); + LayerWeights::FullAttn(FullAttnLayerWeights { + attn_norm: take!(attn_norm), + wq: take!(wq), + wk: take!(wk), + wv: take!(wv), + wo: take!(wo), + q_norm: take!(q_norm), + k_norm: take!(k_norm), + ffn_norm: take!(ffn_norm), + w_gate: take!(w_gate), + w_up: take!(w_up), + w_down: take!(w_down), + }) + } + (LayerType::LinearAttention, true) => { + stage!(attn_norm, b.norm("input_layernorm.weight", &[config.dim])); + stage!(wqkv, b.proj("linear_attn.in_proj_qkv", qkv_dim, config.dim)); + stage!(wz, b.proj("linear_attn.in_proj_z", d_inner, config.dim)); + stage!( + w_alpha, + b.proj( + "linear_attn.in_proj_a", + config.linear_num_value_heads, + config.dim, + ) + ); + stage!( + w_beta, + b.proj( + "linear_attn.in_proj_b", + config.linear_num_value_heads, + config.dim, + ) + ); + stage!( + a_log, + b.raw_f32("linear_attn.A_log", config.linear_num_value_heads) + ); + stage!( + dt_bias, + b.raw_f32("linear_attn.dt_bias", config.linear_num_value_heads) + ); + stage!( + conv_weight, + b.raw_f32( + "linear_attn.conv1d.weight", + qkv_dim * config.conv_kernel_dim, + ) + ); + stage!( + norm_weight, + b.raw_f32("linear_attn.norm.weight", config.linear_value_head_dim) + ); + stage!(wo, b.proj("linear_attn.out_proj", config.dim, d_inner)); + stage!( + ffn_norm, + b.norm("post_attention_layernorm.weight", &[config.dim]) + ); + stage!(ffn, load_moe(b, config, layer_idx)); + LayerWeights::DeltaNetMoe(DeltaNetMoeLayerWeights { + attn_norm: take!(attn_norm), + wqkv: take!(wqkv), + wz: take!(wz), + w_alpha: take!(w_alpha), + w_beta: take!(w_beta), + a_log: take!(a_log), + dt_bias: take!(dt_bias), + conv_weight: take!(conv_weight), + norm_weight: take!(norm_weight), + wo: take!(wo), + ffn_norm: take!(ffn_norm), + ffn: take!(ffn), + }) + } + (LayerType::FullAttention, true) => { + stage!(attn_norm, b.norm("input_layernorm.weight", &[config.dim])); + stage!(wq, b.proj("self_attn.q_proj", q_out_dim, config.dim)); + stage!(wk, b.proj("self_attn.k_proj", kv_dim, config.dim)); + stage!(wv, b.proj("self_attn.v_proj", kv_dim, config.dim)); + stage!(wo, b.proj("self_attn.o_proj", config.dim, o_in)); + stage!( + q_norm, + b.norm("self_attn.q_norm.weight", &[config.head_dim]) + ); + stage!( + k_norm, + b.norm("self_attn.k_norm.weight", &[config.head_dim]) + ); + stage!( + ffn_norm, + b.norm("post_attention_layernorm.weight", &[config.dim]) + ); + stage!(ffn, load_moe(b, config, layer_idx)); + LayerWeights::FullAttnMoe(FullAttnMoeLayerWeights { + attn_norm: take!(attn_norm), + wq: take!(wq), + wk: take!(wk), + wv: take!(wv), + wo: take!(wo), + q_norm: take!(q_norm), + k_norm: take!(k_norm), + ffn_norm: take!(ffn_norm), + ffn: take!(ffn), + }) + } + }; + Ok(layer) +} + +#[cfg(test)] +mod tests { + use super::*; + use hip_bridge::{HipError, HipResult}; + use rdna_compute::{DType, Gpu}; + + struct FaultBackend { + gpu: Gpu, + calls: usize, + fail_at: Option, + freed: usize, + live: usize, + } + + impl FaultBackend { + fn new(gpu: Gpu) -> Self { + Self { + gpu, + calls: 0, + fail_at: None, + freed: 0, + live: 0, + } + } + + fn next(&mut self) -> HipResult<()> { + self.calls += 1; + if self.fail_at == Some(self.calls) { + Err(HipError::new( + 0, + &format!("fault at layer operation {}", self.calls), + )) + } else { + Ok(()) + } + } + + fn alloc_tensor(&mut self) -> HipResult { + let tensor = self.gpu.alloc_tensor(&[1], DType::F32)?; + self.live += 1; + Ok(tensor) + } + + fn alloc_weight(&mut self, m: usize, k: usize) -> HipResult { + Ok(WeightTensor { + buf: self.alloc_tensor()?, + gpu_dtype: DType::F32, + m, + k, + row_stride: 0, + paro: None, + awq_scale: None, + }) + } + + fn alloc_moe(&mut self) -> HipResult { + let mut buffers = Vec::with_capacity(7); + for _ in 0..7 { + match self.alloc_tensor() { + Ok(tensor) => buffers.push(tensor), + Err(err) => { + for tensor in buffers.drain(..) { + self.free_tensor(tensor); + } + return Err(err); + } + } + } + fn weight_from(buffers: &mut Vec) -> WeightTensor { + WeightTensor { + buf: buffers.pop().expect("MoE test buffer"), + gpu_dtype: DType::F32, + m: 1, + k: 1, + row_stride: 0, + paro: None, + awq_scale: None, + } + } + Ok(MoeFfnWeights { + router: weight_from(&mut buffers), + experts: Vec::new(), + packed_expert_owners: None, + shared_expert: crate::qwen35::SharedExpertWeights { + gate: weight_from(&mut buffers), + up: weight_from(&mut buffers), + down: weight_from(&mut buffers), + }, + shared_expert_gate: weight_from(&mut buffers), + expert_gate_up_ptrs: buffers.pop().expect("MoE test pointer buffer"), + expert_down_ptrs: buffers.pop().expect("MoE test pointer buffer"), + expert_down_awq_ptrs: None, + expert_dtype_tags: None, + layer_idx: 0, + expert_shape: None, + paro_shared: None, + global_expert_dtypes: None, + ep_dummy_buffers: Vec::new(), + }) + } + + fn assert_drained(&self) { + assert_eq!(self.live, 0, "all GPU owners must be reclaimed"); + } + } + + impl WeightBackend for FaultBackend { + fn set_layer(&mut self, _layer: usize) {} + + fn proj(&mut self, _rel: &str, m: usize, k: usize) -> HipResult { + self.next()?; + self.alloc_weight(m, k) + } + + fn norm(&mut self, _rel: &str, _shape: &[usize]) -> HipResult { + self.next()?; + self.alloc_tensor() + } + + fn raw_f32(&mut self, _rel: &str, _n: usize) -> HipResult { + self.next()?; + self.alloc_tensor() + } + + fn bias(&mut self, _rel: &str, _n: usize) -> HipResult { + Err(HipError::new(0, "test backend does not load biases")) + } + + fn free_tensor(&mut self, tensor: GpuTensor) { + self.gpu + .free_tensor(tensor) + .expect("test owner free must succeed"); + self.live = self.live.checked_sub(1).expect("owner freed twice"); + self.freed += 1; + } + } + + fn test_config(moe: bool) -> Qwen35Config { + Qwen35Config { + dim: 1, + n_layers: 1, + vocab_size: 1, + norm_eps: 1e-5, + eos_token: 0, + n_heads: 1, + n_kv_heads: 1, + head_dim: 1, + rope_theta: 1.0, + partial_rotary_factor: 1.0, + is_vl_text: false, + mrope_interleaved: false, + mrope_section: [0; 3], + linear_num_key_heads: 1, + linear_num_value_heads: 1, + linear_key_head_dim: 1, + linear_value_head_dim: 1, + conv_kernel_dim: 1, + hidden_dim: 1, + num_experts: if moe { 1 } else { 0 }, + num_experts_per_tok: if moe { 1 } else { 0 }, + moe_intermediate_size: 1, + shared_expert_intermediate_size: 1, + has_shared_expert: moe, + norm_topk_prob: false, + layer_types: vec![LayerType::FullAttention], + paged_experts: false, + vram_budget_bytes: u64::MAX, + reap_keep: None, + } + } + + fn no_moe( + _backend: &mut FaultBackend, + _config: &Qwen35Config, + _layer: usize, + ) -> HipResult { + Err(HipError::new(0, "dense test must not load MoE")) + } + + fn free_weight(backend: &mut B, weight: WeightTensor) { + let mut free = |tensor: GpuTensor| backend.free_tensor(tensor); + free_weight_with(weight, &mut free); + } + + fn free_test_layer(backend: &mut B, layer: LayerWeights) { + match layer { + LayerWeights::FullAttn(layer) => { + let FullAttnLayerWeights { + attn_norm, + wq, + wk, + wv, + wo, + q_norm, + k_norm, + ffn_norm, + w_gate, + w_up, + w_down, + } = layer; + for tensor in [attn_norm, q_norm, k_norm, ffn_norm] { + backend.free_tensor(tensor); + } + for weight in [wq, wk, wv, wo, w_gate, w_up, w_down] { + free_weight(backend, weight); + } + } + LayerWeights::FullAttnMoe(layer) => { + let FullAttnMoeLayerWeights { + attn_norm, + wq, + wk, + wv, + wo, + q_norm, + k_norm, + ffn_norm, + ffn, + } = layer; + for tensor in [attn_norm, q_norm, k_norm, ffn_norm] { + backend.free_tensor(tensor); + } + for weight in [wq, wk, wv, wo] { + free_weight(backend, weight); + } + let mut free = |tensor: GpuTensor| backend.free_tensor(tensor); + free_moe_ffn_with(ffn, &mut free); + } + _ => panic!("test helper only handles full-attention variants"), + } + } + + #[test] + #[ignore = "requires a real HIP GPU"] + fn dense_layer_failure_reclaims_owners_and_retry_succeeds() { + let Some(gpu) = Gpu::init().ok() else { + eprintln!("skip: no GPU"); + return; + }; + let config = test_config(false); + let mut backend = FaultBackend::new(gpu); + backend.fail_at = Some(11); + + let failed = load_layer(&mut backend, &config, 0, no_moe); + assert!(failed.is_err()); + assert_eq!(backend.calls, 11); + assert_eq!(backend.freed, 10); + backend.assert_drained(); + + backend.calls = 0; + backend.fail_at = None; + let layer = load_layer(&mut backend, &config, 0, no_moe).expect("retry"); + assert_eq!(backend.live, 11); + free_test_layer(&mut backend, layer); + assert_eq!(backend.freed, 21); + backend.assert_drained(); + } + + #[test] + #[ignore = "requires a real HIP GPU"] + fn moe_boundary_failure_reclaims_attention_owners_and_retry_succeeds() { + let Some(gpu) = Gpu::init().ok() else { + eprintln!("skip: no GPU"); + return; + }; + let config = test_config(true); + let mut backend = FaultBackend::new(gpu); + + let failed = load_layer( + &mut backend, + &config, + 0, + |_backend: &mut FaultBackend, + _config: &Qwen35Config, + _layer: usize| + -> HipResult { + Err(HipError::new(0, "injected late MoE failure")) + }, + ); + assert!(failed.is_err()); + assert_eq!(backend.calls, 8); + assert_eq!(backend.freed, 8); + backend.assert_drained(); - Ok(match (config.layer_types[layer_idx], is_moe) { - (LayerType::LinearAttention, false) => LayerWeights::DeltaNet(DeltaNetLayerWeights { - attn_norm: b.norm("input_layernorm.weight", &[config.dim])?, - wqkv: b.proj("linear_attn.in_proj_qkv", qkv_dim, config.dim)?, - wz: b.proj("linear_attn.in_proj_z", d_inner, config.dim)?, - w_alpha: b.proj( - "linear_attn.in_proj_a", - config.linear_num_value_heads, - config.dim, - )?, - w_beta: b.proj( - "linear_attn.in_proj_b", - config.linear_num_value_heads, - config.dim, - )?, - a_log: b.raw_f32("linear_attn.A_log", config.linear_num_value_heads)?, - dt_bias: b.raw_f32("linear_attn.dt_bias", config.linear_num_value_heads)?, - conv_weight: b.raw_f32( - "linear_attn.conv1d.weight", - qkv_dim * config.conv_kernel_dim, - )?, - norm_weight: b.raw_f32("linear_attn.norm.weight", config.linear_value_head_dim)?, - wo: b.proj("linear_attn.out_proj", config.dim, d_inner)?, - ffn_norm: b.norm("post_attention_layernorm.weight", &[config.dim])?, - w_gate: b.proj("mlp.gate_proj", config.hidden_dim, config.dim)?, - w_up: b.proj("mlp.up_proj", config.hidden_dim, config.dim)?, - w_down: b.proj("mlp.down_proj", config.dim, config.hidden_dim)?, - }), - (LayerType::FullAttention, false) => LayerWeights::FullAttn(FullAttnLayerWeights { - attn_norm: b.norm("input_layernorm.weight", &[config.dim])?, - wq: b.proj("self_attn.q_proj", q_out_dim, config.dim)?, - wk: b.proj("self_attn.k_proj", kv_dim, config.dim)?, - wv: b.proj("self_attn.v_proj", kv_dim, config.dim)?, - wo: b.proj("self_attn.o_proj", config.dim, o_in)?, - q_norm: b.norm("self_attn.q_norm.weight", &[config.head_dim])?, - k_norm: b.norm("self_attn.k_norm.weight", &[config.head_dim])?, - ffn_norm: b.norm("post_attention_layernorm.weight", &[config.dim])?, - w_gate: b.proj("mlp.gate_proj", config.hidden_dim, config.dim)?, - w_up: b.proj("mlp.up_proj", config.hidden_dim, config.dim)?, - w_down: b.proj("mlp.down_proj", config.dim, config.hidden_dim)?, - }), - (LayerType::LinearAttention, true) => LayerWeights::DeltaNetMoe(DeltaNetMoeLayerWeights { - attn_norm: b.norm("input_layernorm.weight", &[config.dim])?, - wqkv: b.proj("linear_attn.in_proj_qkv", qkv_dim, config.dim)?, - wz: b.proj("linear_attn.in_proj_z", d_inner, config.dim)?, - w_alpha: b.proj( - "linear_attn.in_proj_a", - config.linear_num_value_heads, - config.dim, - )?, - w_beta: b.proj( - "linear_attn.in_proj_b", - config.linear_num_value_heads, - config.dim, - )?, - a_log: b.raw_f32("linear_attn.A_log", config.linear_num_value_heads)?, - dt_bias: b.raw_f32("linear_attn.dt_bias", config.linear_num_value_heads)?, - conv_weight: b.raw_f32( - "linear_attn.conv1d.weight", - qkv_dim * config.conv_kernel_dim, - )?, - norm_weight: b.raw_f32("linear_attn.norm.weight", config.linear_value_head_dim)?, - wo: b.proj("linear_attn.out_proj", config.dim, d_inner)?, - ffn_norm: b.norm("post_attention_layernorm.weight", &[config.dim])?, - ffn: load_moe(b, config, layer_idx)?, - }), - (LayerType::FullAttention, true) => LayerWeights::FullAttnMoe(FullAttnMoeLayerWeights { - attn_norm: b.norm("input_layernorm.weight", &[config.dim])?, - wq: b.proj("self_attn.q_proj", q_out_dim, config.dim)?, - wk: b.proj("self_attn.k_proj", kv_dim, config.dim)?, - wv: b.proj("self_attn.v_proj", kv_dim, config.dim)?, - wo: b.proj("self_attn.o_proj", config.dim, o_in)?, - q_norm: b.norm("self_attn.q_norm.weight", &[config.head_dim])?, - k_norm: b.norm("self_attn.k_norm.weight", &[config.head_dim])?, - ffn_norm: b.norm("post_attention_layernorm.weight", &[config.dim])?, - ffn: load_moe(b, config, layer_idx)?, - }), - }) + backend.calls = 0; + let layer = load_layer(&mut backend, &config, 0, |backend, _, _| { + backend.alloc_moe() + }) + .expect("MoE retry"); + assert_eq!(backend.live, 15); + free_test_layer(&mut backend, layer); + assert_eq!(backend.freed, 23); + backend.assert_drained(); + } } diff --git a/crates/hipfire-arch-qwen35/src/qwen35/batch.rs b/crates/hipfire-arch-qwen35/src/qwen35/batch.rs index 67e77b16d..152a7aa61 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/batch.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/batch.rs @@ -190,6 +190,16 @@ impl PrefillBatchScratch { config: &Qwen35Config, max_batch: usize, cap_gdn_tape: bool, + ) -> HipResult { + Self::new_opt_with_alloc(gpu, config, max_batch, cap_gdn_tape, Gpu::alloc_tensor) + } + + fn new_opt_with_alloc( + gpu: &mut Gpu, + config: &Qwen35Config, + max_batch: usize, + cap_gdn_tape: bool, + mut allocate: impl FnMut(&mut Gpu, &[usize], DType) -> HipResult, ) -> HipResult { let dim = config.dim; let hidden_dim = config.hidden_dim; @@ -200,47 +210,28 @@ impl PrefillBatchScratch { let q_dim = config.n_heads * config.head_dim; let kv_dim = config.n_kv_heads * config.head_dim; - // hunt3 H-E residual: this struct literal allocates ~40 GpuTensors via - // `?` early-returns. PrefillBatchScratch has no Drop impl (GpuTensor - // carries no Gpu handle; free_tensor needs &mut Gpu), so a `?` failure - // partway through would drop the already-allocated tensors WITHOUT - // freeing them on the device — the exact intra-`new` leak the - // cross-band H-E recovery can't reach. OOM during new() is precisely - // when a mid-literal failure is most likely. Fix: route every alloc - // through a ledger and, on the first error, free everything allocated - // so far before propagating. `alloc!` records mandatory tensors; - // `alloc_opt!` records the inner tensor of an `if cond { Some(..) }`. - // - // The ledger stores non-owning aliases (DeviceBuffer has no Drop and - // GpuTensor is not Clone), so on success the aliases drop as no-ops and - // the real tensors live on in the struct (no double-free); on error we - // free each alias once, which releases the same pool buffer the - // partially-built (and about-to-be-dropped, never-freed) field held. - let mut ledger: Vec = Vec::with_capacity(48); + // Transactional construction: slots own every successful allocation + // until the struct is built. On an allocation error, reverse-drain the + // slots and free each actual owner before returning; GpuTensor has no + // Drop implementation that could release device memory for us. + let mut slots: Vec> = Vec::with_capacity(54); macro_rules! alloc { - ($shape:expr, $dt:expr) => { - match gpu.alloc_tensor($shape, $dt) { + ($shape:expr, $dt:expr) => {{ + match allocate(gpu, $shape, $dt) { Ok(t) => { - // SAFETY: alias lives only inside `new`; if used it is - // freed in the error arm below (the original field is - // dropped without freeing, no Drop on GpuTensor), and - // on success it is dropped untouched (no Drop on - // DeviceBuffer) while the original is moved into Self. - ledger.push(GpuTensor { - buf: unsafe { t.buf.alias() }, - shape: t.shape.clone(), - dtype: t.dtype, - }); - t + slots.push(Some(t)); + slots.len() - 1 } Err(e) => { - for prev in ledger.drain(..) { - let _ = gpu.free_tensor(prev); + while let Some(slot) = slots.pop() { + if let Some(t) = slot { + let _ = gpu.free_tensor(t); + } } return Err(e); } } - }; + }}; } macro_rules! alloc_opt { ($cond:expr, $shape:expr, $dt:expr) => { @@ -251,170 +242,231 @@ impl PrefillBatchScratch { } }; } + macro_rules! take { + ($i:expr) => {{ + slots[$i].take().expect("prefill scratch slot taken twice") + }}; + } // Hoisted grouped-GEMM sizing (same value across the Path-2 fields). let grouped_m_total_max = moe_grouped_m_total_max(max_batch, config.num_experts_per_tok, config.num_experts); let grouped_total_slots_max = max_batch * config.num_experts_per_tok; + let i_x_batch = alloc!(&[max_batch * dim], DType::F32); + let i_x_rot_batch = alloc!(&[max_batch * dim], DType::F32); + let i_x_norm_batch = alloc!(&[max_batch * dim], DType::F32); + let i_dn_qkv_batch = alloc!(&[max_batch * qkv_dim], DType::F32); + let i_dn_z_batch = alloc!(&[max_batch * v_dim], DType::F32); + let i_dn_alpha_batch = alloc!(&[max_batch * n_v_heads], DType::F32); + let i_dn_beta_batch = alloc!(&[max_batch * n_v_heads], DType::F32); + let i_dn_q_raw_batch = alloc!(&[max_batch * k_dim], DType::F32); + let i_dn_k_raw_batch = alloc!(&[max_batch * k_dim], DType::F32); + let i_dn_v_batch = alloc!(&[max_batch * v_dim], DType::F32); + let i_dn_q_batch = alloc!(&[max_batch * v_dim], DType::F32); + let i_dn_k_batch = alloc!(&[max_batch * v_dim], DType::F32); + let i_dn_attn_out_batch = alloc!(&[max_batch * v_dim], DType::F32); + let i_dn_normed_batch = alloc!(&[max_batch * v_dim], DType::F32); + let i_gate_ffn_batch = alloc!(&[max_batch * hidden_dim], DType::F32); + let i_up_batch = alloc!(&[max_batch * hidden_dim], DType::F32); + let i_ffn_hidden_batch = alloc!(&[max_batch * hidden_dim], DType::F32); + let i_dn_normed_rot_batch = alloc!(&[max_batch * v_dim], DType::F32); + // F32 dtype = 4 bytes/element, same layout as i32. The rope / + // attention / kv_write kernels cast the pointer to `const int*`, + // so dtype is cosmetic. Upload i32 bits via memcpy_htod. + let i_positions = alloc!(&[max_batch], DType::F32); + // Depth-based RoPE angles for DDTree verify (39aa358 fix): + // `positions` stays the flat linear KV slot index; this buffer + // carries `base_pos + depth(node)` so FA-layer RoPE rotates Q/K + // at the logically-correct phase while KV writes stay on + // distinct linear slots. Uploaded per cycle in tree-verify mode + // from `TreeVerifyCtx.positions`; FA RoPE kernels read it ONLY + // when `tree_verify.is_some()`. Same i32-in-F32 cosmetic dtype + // pattern as `positions`. + let i_rope_positions = alloc!(&[max_batch], DType::F32); + let i_tokens = alloc!(&[max_batch], DType::F32); + let i_fa_q_full_batch = alloc!(&[max_batch * q_dim * 2], DType::F32); + let i_fa_q_batch = alloc!(&[max_batch * q_dim], DType::F32); + let i_fa_gate_batch = alloc!(&[max_batch * q_dim], DType::F32); + let i_fa_k_batch = alloc!(&[max_batch * kv_dim], DType::F32); + let i_fa_v_batch = alloc!(&[max_batch * kv_dim], DType::F32); + let i_fa_attn_out_batch = alloc!(&[max_batch * q_dim], DType::F32); + let i_fa_attn_out_rot_batch = alloc!(&[max_batch * q_dim], DType::F32); + let i_x_rot_f16_batch = alloc!(&[max_batch * dim], DType::F16); + let i_dn_normed_rot_f16_batch = alloc!(&[max_batch * v_dim], DType::F16); + let i_ffn_hidden_f16_batch = alloc!(&[max_batch * hidden_dim], DType::F16); + let i_fa_attn_out_rot_f16_batch = alloc!(&[max_batch * q_dim], DType::F16); + // S9 prologue control plane: 256 bytes of device-resident + // counters/generations. Raw dtype counts bytes. + let i_mq_prologue_ctrl = alloc!(&[256], DType::Raw); + let i_moe_router_logits_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.num_experts], + DType::F32 + ); + let i_moe_shared_scalar_batch = + alloc_opt!(config.num_experts > 0, &[max_batch], DType::F32); + let i_moe_shared_gate_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.shared_expert_intermediate_size], + DType::F32 + ); + let i_moe_shared_up_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.shared_expert_intermediate_size], + DType::F32 + ); + let i_moe_shared_rot_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.shared_expert_intermediate_size], + DType::F32 + ); + let i_moe_topk_indices_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.num_experts_per_tok], + DType::F32 + ); + let i_moe_topk_weights_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.num_experts_per_tok], + DType::F32 + ); + let i_moe_gate_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.num_experts_per_tok * config.moe_intermediate_size], + DType::F32 + ); + let i_moe_up_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.num_experts_per_tok * config.moe_intermediate_size], + DType::F32 + ); + let i_moe_rot_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.num_experts_per_tok * config.moe_intermediate_size], + DType::F32 + ); + let i_moe_down_expanded_batch = alloc_opt!( + config.num_experts > 0, + &[max_batch * config.num_experts_per_tok * config.dim], + DType::F32 + ); + // Path 2 scatter + grouped-WMMA-GEMM scratch (gated at runtime by + // HIPFIRE_MOE_GROUPED_GEMM=1). m_total_max = N*K_TOP + E*(BLOCK_M-1). + // i32 buffers stored as Raw (4 bytes/elem matches; no DType::I32 yet). + let i_moe_expert_token_counts = alloc_opt!( + config.num_experts > 0, + &[config.num_experts * 4], + DType::Raw + ); + let i_moe_expert_offsets = alloc_opt!( + config.num_experts > 0, + &[(config.num_experts + 1) * 4], + DType::Raw + ); + let i_moe_sorted_slot_index = alloc_opt!( + config.num_experts > 0, + &[grouped_m_total_max * 4], + DType::Raw + ); + let i_moe_inverse_perm = alloc_opt!( + config.num_experts > 0, + &[grouped_total_slots_max * 4], + DType::Raw + ); + let i_moe_expert_tile_ids = alloc_opt!( + config.num_experts > 0, + &[(grouped_m_total_max / MOE_GROUPED_BLOCK_M) * 4], + DType::Raw + ); + let i_moe_y_gate_up_grouped = alloc_opt!( + config.num_experts > 0, + &[grouped_m_total_max * 2 * config.moe_intermediate_size], + DType::F32 + ); + let i_moe_y_down_grouped = alloc_opt!( + config.num_experts > 0, + &[grouped_m_total_max * config.dim], + DType::F32 + ); + let i_dn_s_tape_q8 = alloc_opt!( + cap_gdn_tape && config.linear_num_value_heads > 0, + &[max_batch + * config.linear_num_value_heads + * config.linear_value_head_dim + * config.linear_value_head_dim], + DType::Raw + ); + let i_dn_s_tape_scales = alloc_opt!( + cap_gdn_tape && config.linear_num_value_heads > 0, + &[max_batch * config.linear_num_value_heads * config.linear_value_head_dim], + DType::F32 + ); + let i_dn_s_tape_f32 = alloc_opt!( + cap_gdn_tape && config.linear_num_value_heads > 0, + &[max_batch + * config.linear_num_value_heads + * config.linear_value_head_dim + * config.linear_value_head_dim], + DType::F32 + ); + Ok(Self { max_batch, - x_batch: alloc!(&[max_batch * dim], DType::F32), - x_rot_batch: alloc!(&[max_batch * dim], DType::F32), - x_norm_batch: alloc!(&[max_batch * dim], DType::F32), - dn_qkv_batch: alloc!(&[max_batch * qkv_dim], DType::F32), - dn_z_batch: alloc!(&[max_batch * v_dim], DType::F32), - dn_alpha_batch: alloc!(&[max_batch * n_v_heads], DType::F32), - dn_beta_batch: alloc!(&[max_batch * n_v_heads], DType::F32), - dn_q_raw_batch: alloc!(&[max_batch * k_dim], DType::F32), - dn_k_raw_batch: alloc!(&[max_batch * k_dim], DType::F32), - dn_v_batch: alloc!(&[max_batch * v_dim], DType::F32), - dn_q_batch: alloc!(&[max_batch * v_dim], DType::F32), - dn_k_batch: alloc!(&[max_batch * v_dim], DType::F32), - dn_attn_out_batch: alloc!(&[max_batch * v_dim], DType::F32), - dn_normed_batch: alloc!(&[max_batch * v_dim], DType::F32), - gate_ffn_batch: alloc!(&[max_batch * hidden_dim], DType::F32), - up_batch: alloc!(&[max_batch * hidden_dim], DType::F32), - ffn_hidden_batch: alloc!(&[max_batch * hidden_dim], DType::F32), - dn_normed_rot_batch: alloc!(&[max_batch * v_dim], DType::F32), - // F32 dtype = 4 bytes/element, same layout as i32. The rope / - // attention / kv_write kernels cast the pointer to `const int*`, - // so dtype is cosmetic. Upload i32 bits via memcpy_htod. - positions: alloc!(&[max_batch], DType::F32), - // Depth-based RoPE angles for DDTree verify (39aa358 fix): - // `positions` stays the flat linear KV slot index; this buffer - // carries `base_pos + depth(node)` so FA-layer RoPE rotates Q/K - // at the logically-correct phase while KV writes stay on - // distinct linear slots. Uploaded per cycle in tree-verify mode - // from `TreeVerifyCtx.positions`; FA RoPE kernels read it ONLY - // when `tree_verify.is_some()`. Same i32-in-F32 cosmetic dtype - // pattern as `positions`. - rope_positions: alloc!(&[max_batch], DType::F32), - tokens: alloc!(&[max_batch], DType::F32), - fa_q_full_batch: alloc!(&[max_batch * q_dim * 2], DType::F32), - fa_q_batch: alloc!(&[max_batch * q_dim], DType::F32), - fa_gate_batch: alloc!(&[max_batch * q_dim], DType::F32), - fa_k_batch: alloc!(&[max_batch * kv_dim], DType::F32), - fa_v_batch: alloc!(&[max_batch * kv_dim], DType::F32), - fa_attn_out_batch: alloc!(&[max_batch * q_dim], DType::F32), - fa_attn_out_rot_batch: alloc!(&[max_batch * q_dim], DType::F32), - x_rot_f16_batch: alloc!(&[max_batch * dim], DType::F16), - dn_normed_rot_f16_batch: alloc!(&[max_batch * v_dim], DType::F16), - ffn_hidden_f16_batch: alloc!(&[max_batch * hidden_dim], DType::F16), - fa_attn_out_rot_f16_batch: alloc!(&[max_batch * q_dim], DType::F16), - // S9 prologue control plane: 256 bytes of device-resident - // counters/generations. Raw dtype counts bytes. - mq_prologue_ctrl: alloc!(&[256], DType::Raw), - moe_router_logits_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.num_experts], - DType::F32 - ), - moe_shared_scalar_batch: alloc_opt!(config.num_experts > 0, &[max_batch], DType::F32), - moe_shared_gate_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.shared_expert_intermediate_size], - DType::F32 - ), - moe_shared_up_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.shared_expert_intermediate_size], - DType::F32 - ), - moe_shared_rot_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.shared_expert_intermediate_size], - DType::F32 - ), - moe_topk_indices_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.num_experts_per_tok], - DType::F32 - ), - moe_topk_weights_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.num_experts_per_tok], - DType::F32 - ), - moe_gate_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.num_experts_per_tok * config.moe_intermediate_size], - DType::F32 - ), - moe_up_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.num_experts_per_tok * config.moe_intermediate_size], - DType::F32 - ), - moe_rot_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.num_experts_per_tok * config.moe_intermediate_size], - DType::F32 - ), - moe_down_expanded_batch: alloc_opt!( - config.num_experts > 0, - &[max_batch * config.num_experts_per_tok * config.dim], - DType::F32 - ), - // Path 2 scatter + grouped-WMMA-GEMM scratch (gated at runtime by - // HIPFIRE_MOE_GROUPED_GEMM=1). m_total_max = N*K_TOP + E*(BLOCK_M-1). - // i32 buffers stored as Raw (4 bytes/elem matches; no DType::I32 yet). - moe_expert_token_counts: alloc_opt!( - config.num_experts > 0, - &[config.num_experts * 4], - DType::Raw - ), - moe_expert_offsets: alloc_opt!( - config.num_experts > 0, - &[(config.num_experts + 1) * 4], - DType::Raw - ), - moe_sorted_slot_index: alloc_opt!( - config.num_experts > 0, - &[grouped_m_total_max * 4], - DType::Raw - ), - moe_inverse_perm: alloc_opt!( - config.num_experts > 0, - &[grouped_total_slots_max * 4], - DType::Raw - ), - moe_expert_tile_ids: alloc_opt!( - config.num_experts > 0, - &[(grouped_m_total_max / MOE_GROUPED_BLOCK_M) * 4], - DType::Raw - ), - moe_y_gate_up_grouped: alloc_opt!( - config.num_experts > 0, - &[grouped_m_total_max * 2 * config.moe_intermediate_size], - DType::F32 - ), - moe_y_down_grouped: alloc_opt!( - config.num_experts > 0, - &[grouped_m_total_max * config.dim], - DType::F32 - ), - dn_s_tape_q8: alloc_opt!( - cap_gdn_tape && config.linear_num_value_heads > 0, - &[max_batch - * config.linear_num_value_heads - * config.linear_value_head_dim - * config.linear_value_head_dim], - DType::Raw - ), - dn_s_tape_scales: alloc_opt!( - cap_gdn_tape && config.linear_num_value_heads > 0, - &[max_batch * config.linear_num_value_heads * config.linear_value_head_dim], - DType::F32 - ), - dn_s_tape_f32: alloc_opt!( - cap_gdn_tape && config.linear_num_value_heads > 0, - &[max_batch - * config.linear_num_value_heads - * config.linear_value_head_dim - * config.linear_value_head_dim], - DType::F32 - ), + x_batch: take!(i_x_batch), + x_rot_batch: take!(i_x_rot_batch), + x_norm_batch: take!(i_x_norm_batch), + dn_qkv_batch: take!(i_dn_qkv_batch), + dn_z_batch: take!(i_dn_z_batch), + dn_alpha_batch: take!(i_dn_alpha_batch), + dn_beta_batch: take!(i_dn_beta_batch), + dn_q_raw_batch: take!(i_dn_q_raw_batch), + dn_k_raw_batch: take!(i_dn_k_raw_batch), + dn_v_batch: take!(i_dn_v_batch), + dn_q_batch: take!(i_dn_q_batch), + dn_k_batch: take!(i_dn_k_batch), + dn_attn_out_batch: take!(i_dn_attn_out_batch), + dn_normed_batch: take!(i_dn_normed_batch), + gate_ffn_batch: take!(i_gate_ffn_batch), + up_batch: take!(i_up_batch), + ffn_hidden_batch: take!(i_ffn_hidden_batch), + dn_normed_rot_batch: take!(i_dn_normed_rot_batch), + positions: take!(i_positions), + rope_positions: take!(i_rope_positions), + tokens: take!(i_tokens), + fa_q_full_batch: take!(i_fa_q_full_batch), + fa_q_batch: take!(i_fa_q_batch), + fa_gate_batch: take!(i_fa_gate_batch), + fa_k_batch: take!(i_fa_k_batch), + fa_v_batch: take!(i_fa_v_batch), + fa_attn_out_batch: take!(i_fa_attn_out_batch), + fa_attn_out_rot_batch: take!(i_fa_attn_out_rot_batch), + x_rot_f16_batch: take!(i_x_rot_f16_batch), + dn_normed_rot_f16_batch: take!(i_dn_normed_rot_f16_batch), + ffn_hidden_f16_batch: take!(i_ffn_hidden_f16_batch), + fa_attn_out_rot_f16_batch: take!(i_fa_attn_out_rot_f16_batch), + mq_prologue_ctrl: take!(i_mq_prologue_ctrl), + moe_router_logits_batch: i_moe_router_logits_batch.map(|i| take!(i)), + moe_shared_scalar_batch: i_moe_shared_scalar_batch.map(|i| take!(i)), + moe_shared_gate_batch: i_moe_shared_gate_batch.map(|i| take!(i)), + moe_shared_up_batch: i_moe_shared_up_batch.map(|i| take!(i)), + moe_shared_rot_batch: i_moe_shared_rot_batch.map(|i| take!(i)), + moe_topk_indices_batch: i_moe_topk_indices_batch.map(|i| take!(i)), + moe_topk_weights_batch: i_moe_topk_weights_batch.map(|i| take!(i)), + moe_gate_batch: i_moe_gate_batch.map(|i| take!(i)), + moe_up_batch: i_moe_up_batch.map(|i| take!(i)), + moe_rot_batch: i_moe_rot_batch.map(|i| take!(i)), + moe_down_expanded_batch: i_moe_down_expanded_batch.map(|i| take!(i)), + moe_expert_token_counts: i_moe_expert_token_counts.map(|i| take!(i)), + moe_expert_offsets: i_moe_expert_offsets.map(|i| take!(i)), + moe_sorted_slot_index: i_moe_sorted_slot_index.map(|i| take!(i)), + moe_inverse_perm: i_moe_inverse_perm.map(|i| take!(i)), + moe_expert_tile_ids: i_moe_expert_tile_ids.map(|i| take!(i)), + moe_y_gate_up_grouped: i_moe_y_gate_up_grouped.map(|i| take!(i)), + moe_y_down_grouped: i_moe_y_down_grouped.map(|i| take!(i)), + dn_s_tape_q8: i_dn_s_tape_q8.map(|i| take!(i)), + dn_s_tape_scales: i_dn_s_tape_scales.map(|i| take!(i)), + dn_s_tape_f32: i_dn_s_tape_f32.map(|i| take!(i)), }) } @@ -527,6 +579,24 @@ impl Qwen35DecodeBatchState { max_batch: usize, lane_capacity: usize, sample_repeat_capacity: usize, + ) -> HipResult { + Self::new_with_output_alloc( + gpu, + config, + max_batch, + lane_capacity, + sample_repeat_capacity, + Gpu::zeros, + ) + } + + fn new_with_output_alloc( + gpu: &mut Gpu, + config: &Qwen35Config, + max_batch: usize, + lane_capacity: usize, + sample_repeat_capacity: usize, + mut allocate_output: impl FnMut(&mut Gpu, &[usize], DType) -> HipResult, ) -> HipResult { if max_batch == 0 || lane_capacity == 0 || sample_repeat_capacity == 0 { return Err(HipError::new( @@ -555,10 +625,8 @@ impl Qwen35DecodeBatchState { // GpuTensor / KvCache / DeltaNetState / PrefillBatchScratch have no // freeing Drop (free needs &mut Gpu). A mid-`new` `?` would leak every // prior stage while the daemon falls back to sequential with the leak - // still resident. Stage each compound owner, then ordinary tensors - // through a ledger of non-owning aliases (same pattern as - // PrefillBatchScratch::new_opt): on error free aliases + compound - // owners before propagating; on success aliases drop as no-ops. + // still resident. Stage each compound owner, then stage ordinary + // tensors as actual owners until the struct is published. let kv_cache = llama::KvCache::new_gpu_q8_filtered( gpu, &is_kv_layer, @@ -583,42 +651,53 @@ impl Qwen35DecodeBatchState { } }; - let mut ledger: Vec = Vec::with_capacity(7); - macro_rules! zeros { - ($shape:expr) => { - match gpu.zeros($shape, DType::F32) { + // Keep the actual output owners in the ledger. Borrowed aliases cannot + // be passed to free_tensor, so they are not useful for rollback. + let mut outputs: Vec> = Vec::with_capacity(7); + macro_rules! output { + ($shape:expr) => {{ + match allocate_output(gpu, $shape, DType::F32) { Ok(t) => { - // SAFETY: alias lives only inside `new`. On error it is - // freed below (original field drops without freeing); - // on success it drops untouched while the original - // moves into Self. - ledger.push(GpuTensor { - buf: unsafe { t.buf.alias() }, - shape: t.shape.clone(), - dtype: t.dtype, - }); - t + outputs.push(Some(t)); + outputs.len() - 1 } Err(e) => { - for prev in ledger.drain(..) { - let _ = gpu.free_tensor(prev); + while let Some(slot) = outputs.pop() { + if let Some(t) = slot { + let _ = gpu.free_tensor(t); + } } - pbs.free_gpu(gpu); + let _ = pbs.free_gpu(gpu); dn_state.free_gpu(gpu); let _ = kv_cache.free_gpu(gpu); return Err(e); } } + }}; + } + macro_rules! take_output { + ($index:expr) => { + outputs[$index] + .take() + .expect("decode batch output staged twice or missing") }; } - let final_hidden = zeros!(&[max_batch * config.dim]); - let logits = zeros!(&[max_batch * config.vocab_size]); - let lm_rot = zeros!(&[max_batch * config.dim]); - let sample_out = zeros!(&[max_batch * 2]); - let sample_repeat_tokens = zeros!(&[repeat_tokens_len]); - let sample_repeat_lengths = zeros!(&[max_batch]); - let sample_rng_states = zeros!(&[max_batch]); + let i_final_hidden = output!(&[max_batch * config.dim]); + let i_logits = output!(&[max_batch * config.vocab_size]); + let i_lm_rot = output!(&[max_batch * config.dim]); + let i_sample_out = output!(&[max_batch * 2]); + let i_sample_repeat_tokens = output!(&[repeat_tokens_len]); + let i_sample_repeat_lengths = output!(&[max_batch]); + let i_sample_rng_states = output!(&[max_batch]); + + let final_hidden = take_output!(i_final_hidden); + let logits = take_output!(i_logits); + let lm_rot = take_output!(i_lm_rot); + let sample_out = take_output!(i_sample_out); + let sample_repeat_tokens = take_output!(i_sample_repeat_tokens); + let sample_repeat_lengths = take_output!(i_sample_repeat_lengths); + let sample_rng_states = take_output!(i_sample_rng_states); Ok(Self { max_batch, lane_capacity, @@ -1696,3 +1775,175 @@ pub fn forward_decode_batch_prepared( let lm_rot = state.lm_rot.sub_offset(0, n * config.dim); lm_head_batched(gpu, &weights.output, &final_hidden, &lm_rot, &logits, n) } + +#[cfg(test)] +mod allocation_tests { + use super::*; + + #[test] + #[ignore = "requires an AMD GPU; exercises real allocation rollback and retry"] + fn prefill_scratch_failure_preserves_reusable_allocations() { + let mut gpu = Gpu::init().expect("GPU required for allocation rollback"); + let config = super::super::config::config_from_metadata_json( + &serde_json::json!({"config": { + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 16, + "vocab_size": 64, + "linear_num_key_heads": 1, + "linear_num_value_heads": 2, + "linear_key_head_dim": 16, + "linear_value_head_dim": 16, + "num_experts": 4, + "num_experts_per_tok": 2, + "moe_intermediate_size": 32, + "shared_expert_intermediate_size": 32 + }}) + .to_string(), + ) + .expect("scratch fixture config"); + let mut allocations = 0; + let warm = PrefillBatchScratch::new_opt_with_alloc( + &mut gpu, + &config, + 2, + true, + |gpu, shape, dtype| { + allocations += 1; + gpu.alloc_tensor(shape, dtype) + }, + ) + .expect("warm scratch"); + warm.free_gpu(&mut gpu).expect("release warm scratch"); + let fresh_allocations = gpu.pool_stats().0; + let mut attempted = 0; + let failure = PrefillBatchScratch::new_opt_with_alloc( + &mut gpu, + &config, + 2, + true, + |gpu, shape, dtype| { + attempted += 1; + if attempted == allocations { + Err(HipError::new( + 2, + "injected final scratch allocation failure", + )) + } else { + gpu.alloc_tensor(shape, dtype) + } + }, + ); + match failure { + Err(error) => assert_eq!(error.code, 2), + Ok(scratch) => { + scratch + .free_gpu(&mut gpu) + .expect("release unexpected success"); + panic!("allocation fault did not trigger"); + } + } + let retry = PrefillBatchScratch::new_opt(&mut gpu, &config, 2, true) + .expect("immediate retry after allocation failure"); + retry.free_gpu(&mut gpu).expect("release retried scratch"); + assert_eq!( + gpu.pool_stats().0, + fresh_allocations, + "failed construction lost reusable allocations instead of rolling them back", + ); + eprintln!("late failure at allocation {allocations}: retry reused the complete warm pool"); + gpu.drain_pool(); + } + #[test] + #[ignore = "requires an AMD GPU; exercises decode batch final-output rollback and retry"] + fn decode_batch_final_output_failure_preserves_reusable_allocations() { + let mut gpu = Gpu::init().expect("GPU required for allocation rollback"); + let config = super::super::config::config_from_metadata_json( + &serde_json::json!({"config": { + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_dim": 32, + "vocab_size": 64, + "linear_num_key_heads": 1, + "linear_num_value_heads": 1, + "linear_key_head_dim": 16, + "linear_value_head_dim": 16, + "linear_conv_kernel_dim": 2, + "layer_types": ["full_attention", "full_attention"] + }}) + .to_string(), + ) + .expect("decode batch fixture config"); + + let mut output_allocations = 0; + let warm = Qwen35DecodeBatchState::new_with_output_alloc( + &mut gpu, + &config, + 1, + 2, + 2, + |gpu, shape, dtype| { + output_allocations += 1; + gpu.zeros(shape, dtype) + }, + ) + .expect("warm decode batch"); + warm.free_gpu(&mut gpu).expect("release warm decode batch"); + assert_eq!( + output_allocations, 7, + "constructor must stage seven outputs" + ); + let fresh_allocations = gpu.pool_stats().0; + + let mut attempted = 0; + let failure = Qwen35DecodeBatchState::new_with_output_alloc( + &mut gpu, + &config, + 1, + 2, + 2, + |gpu, shape, dtype| { + attempted += 1; + if attempted == 7 { + Err(HipError::new( + 2, + "injected final decode batch output allocation failure", + )) + } else { + gpu.zeros(shape, dtype) + } + }, + ); + match failure { + Err(error) => assert_eq!(error.code, 2), + Ok(state) => { + state + .free_gpu(&mut gpu) + .expect("release unexpected decode batch success"); + panic!("allocation fault did not trigger"); + } + } + assert_eq!(attempted, 7, "failure must occur on the final output"); + + let retry = Qwen35DecodeBatchState::new(&mut gpu, &config, 1, 2, 2) + .expect("immediate retry after allocation failure"); + retry + .free_gpu(&mut gpu) + .expect("release retried decode batch"); + assert_eq!( + gpu.pool_stats().0, + fresh_allocations, + "failed construction lost reusable allocations instead of rolling them back", + ); + eprintln!( + "late failure at output allocation {attempted}: retry reused the complete warm pool" + ); + gpu.drain_pool(); + } +} diff --git a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs index fc6089ea9..a5b06960b 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs @@ -1339,67 +1339,109 @@ impl MmqScreenable for Qwen35Weights { } } -fn free_moe_ffn(gpu: &mut Gpu, ffn: MoeFfnWeights) { - ffn.router.free_all(gpu); - ffn.shared_expert_gate.free_all(gpu); - ffn.shared_expert.gate.free_all(gpu); - ffn.shared_expert.up.free_all(gpu); - ffn.shared_expert.down.free_all(gpu); - let _ = gpu.free_tensor(ffn.expert_gate_up_ptrs); - let _ = gpu.free_tensor(ffn.expert_down_ptrs); +/// Free a [`WeightTensor`] through a caller-supplied GPU-tensor cleanup seam. +/// The callback receives every owned sidecar and the weight buffer exactly once. +pub(crate) fn free_weight_with(weight: WeightTensor, free: &mut F) +where + F: FnMut(GpuTensor), +{ + if let Some(paro) = weight.paro { + if !paro.is_alias { + free(paro.pairs); + free(paro.theta); + free(paro.channel_scales); + } + } + if let Some(awq) = weight.awq_scale { + free(awq); + } + free(weight.buf); +} + +/// Free a [`WeightTensor`]'s owning sidecars without freeing its weight buffer. +/// Used only for non-owning views into [`PackedExpertOwners`]. +fn free_weight_metadata_with(weight: WeightTensor, free: &mut F) +where + F: FnMut(GpuTensor), +{ + if let Some(paro) = weight.paro { + if !paro.is_alias { + free(paro.pairs); + free(paro.theta); + free(paro.channel_scales); + } + } + if let Some(awq) = weight.awq_scale { + free(awq); + } +} + +/// Free a staged MoE owner through a caller-supplied GPU-tensor cleanup seam. +/// +/// The ownership branches here are authoritative for all current routed-expert +/// layouts: ordinary per-expert weights, packed uniform-MQ4 owners, ParoQuant +/// shared sidecars, EP dummy buffers, and paged-mode's empty expert vector. +/// Each callback invocation consumes one actual owning buffer exactly once. +pub(crate) fn free_moe_ffn_with(ffn: MoeFfnWeights, free: &mut impl FnMut(GpuTensor)) { + free_weight_with(ffn.router, free); + free_weight_with(ffn.shared_expert_gate, free); + free_weight_with(ffn.shared_expert.gate, free); + free_weight_with(ffn.shared_expert.up, free); + free_weight_with(ffn.shared_expert.down, free); + free(ffn.expert_gate_up_ptrs); + free(ffn.expert_down_ptrs); // Non-owning pointer table — free the buffer only; the per-expert scales it // points into are owned by `experts[i].down.awq_scale` and freed below via - // `e.down.free_all`. + // `free_weight_with`. if let Some(t) = ffn.expert_down_awq_ptrs { - let _ = gpu.free_tensor(t); + free(t); } // Owned device buffer (built from per-expert gpu_dtype). Free it. if let Some(t) = ffn.expert_dtype_tags { - let _ = gpu.free_tensor(t); + free(t); } if let Some(owners) = ffn.packed_expert_owners { // Packed expert WeightTensors are non-owning views. Free only metadata // that remains individually owned, then return each layer blob once. for e in ffn.experts { - free_weight_metadata_only(gpu, e.gate_up); - free_weight_metadata_only(gpu, e.down); + free_weight_metadata_with(e.gate_up, free); + free_weight_metadata_with(e.down, free); } - let _ = gpu.free_tensor(owners.gate_up); - let _ = gpu.free_tensor(owners.down); + free(owners.gate_up); + free(owners.down); } else { for e in ffn.experts { - e.gate_up.free_all(gpu); - e.down.free_all(gpu); + free_weight_with(e.gate_up, free); + free_weight_with(e.down, free); } } // ParoQuant MoE: free the owning shared sidecars (per-expert `paro` fields // alias these and must NOT be freed separately — they're non-owning views). if let Some(s) = ffn.paro_shared { - let _ = gpu.free_tensor(s.gate_up_pairs); - let _ = gpu.free_tensor(s.gate_up_theta); - let _ = gpu.free_tensor(s.gate_up_channel_scales); - let _ = gpu.free_tensor(s.down_pairs); - let _ = gpu.free_tensor(s.down_theta); - let _ = gpu.free_tensor(s.down_channel_scales); + free(s.gate_up_pairs); + free(s.gate_up_theta); + free(s.gate_up_channel_scales); + free(s.down_pairs); + free(s.down_theta); + free(s.down_channel_scales); } for d in ffn.ep_dummy_buffers { - let _ = gpu.free_tensor(d); + free(d); } } -/// Free a [`WeightTensor`]'s owning sidecars without freeing its weight buffer. -/// Used only for non-owning views into [`PackedExpertOwners`]. +fn free_moe_ffn(gpu: &mut Gpu, ffn: MoeFfnWeights) { + let mut free = |tensor| { + let _ = gpu.free_tensor(tensor); + }; + free_moe_ffn_with(ffn, &mut free); +} + fn free_weight_metadata_only(gpu: &mut Gpu, weight: WeightTensor) { - if let Some(paro) = weight.paro { - if !paro.is_alias { - let _ = gpu.free_tensor(paro.pairs); - let _ = gpu.free_tensor(paro.theta); - let _ = gpu.free_tensor(paro.channel_scales); - } - } - if let Some(awq) = weight.awq_scale { - let _ = gpu.free_tensor(awq); - } + let mut free = |tensor| { + let _ = gpu.free_tensor(tensor); + }; + free_weight_metadata_with(weight, &mut free); } // ─── State ────────────────────────────────────────────────────────────── diff --git a/crates/hipfire-runtime/map.md b/crates/hipfire-runtime/map.md index 237e40a68..698a27638 100644 --- a/crates/hipfire-runtime/map.md +++ b/crates/hipfire-runtime/map.md @@ -81,7 +81,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/tool_call.rs`](src/tool_call.rs) | 716 | 7 | 15 | | [`src/tp_shard.rs`](src/tp_shard.rs) | 731 | 25 | 20 | | [`src/triattn.rs`](src/triattn.rs) | 1,355 | 45 | 9 | -| [`src/weight_backend.rs`](src/weight_backend.rs) | 2,038 | 22 | 37 | +| [`src/weight_backend.rs`](src/weight_backend.rs) | 2,049 | 22 | 37 | | [`src/weight_pager.rs`](src/weight_pager.rs) | 850 | 31 | 6 | ### Public API surface @@ -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 · 54,713 lines · 926 public items · 629 tests · 132 examples diff --git a/crates/hipfire-runtime/src/weight_backend.rs b/crates/hipfire-runtime/src/weight_backend.rs index 43a7d09dc..482a6aa61 100644 --- a/crates/hipfire-runtime/src/weight_backend.rs +++ b/crates/hipfire-runtime/src/weight_backend.rs @@ -1265,6 +1265,11 @@ pub trait WeightBackend { fn raw_f32(&mut self, rel: &str, n: usize) -> HipResult; /// Load a bias vector (f32). Only qwen2 attention biases use this today. fn bias(&mut self, rel: &str, n: usize) -> HipResult; + /// Return an already allocated tensor to this backend's GPU pool. + /// + /// Layer loading uses this narrow seam to roll back staged owners without + /// exposing the backend's device handle to arch crates. + fn free_tensor(&mut self, tensor: GpuTensor); } /// HFQ backend. `norm_bias`: `1.0` (qwen3.5/gemma) or `0.0` (qwen2/llama). @@ -1320,6 +1325,9 @@ impl<'a> WeightBackend for HfqBackend<'a> { ); Ok(t) } + fn free_tensor(&mut self, tensor: GpuTensor) { + let _ = self.gpu.free_tensor(tensor); + } } /// Resolve `name` via `candidates` and return the first tensor's `(info, bytes)`. @@ -1373,6 +1381,9 @@ impl<'a> WeightBackend for ParoBackend<'a> { fn raw_f32(&mut self, rel: &str, n: usize) -> HipResult { paro_load_f32(self.source, self.gpu, &paro_plain_name(self.layer, rel), n) } + fn free_tensor(&mut self, tensor: GpuTensor) { + let _ = self.gpu.free_tensor(tensor); + } fn bias(&mut self, _rel: &str, _n: usize) -> HipResult { Err(hip_bridge::HipError::new( 0,