diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b5357e97..6a91f52c2 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 scratch construction retains owning tensors until publication, reclaiming partial allocations on failure so an immediate retry can reuse them. - 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..7c16140db 100644 --- a/crates/hipfire-arch-qwen35/map.md +++ b/crates/hipfire-arch-qwen35/map.md @@ -37,7 +37,7 @@ _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,833 | 16 | 1 | | [`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 | @@ -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 · 58,769 lines · 440 public items · 193 tests · 11 examples diff --git a/crates/hipfire-arch-qwen35/src/qwen35/batch.rs b/crates/hipfire-arch-qwen35/src/qwen35/batch.rs index 67e77b16d..d776be320 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)), }) } @@ -1696,3 +1748,86 @@ 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(); + } +}