From c1299d5fe2ac45059bfa5df8a79800a273e5e0d8 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 17:03:04 +0200 Subject: [PATCH 01/25] feat(hardware): move multi-GPU owner and add device mesh --- Cargo.lock | 15 + Cargo.toml | 1 + crates/hipfire-arch-deepseek4/Cargo.toml | 1 + .../examples/ep_deepseek4.rs | 2 +- .../examples/ep_dspark_topology_probe.rs | 2 +- crates/hipfire-arch-deepseek4/src/ep.rs | 12 +- crates/hipfire-arch-deepseek4/src/forward.rs | 2 +- crates/hipfire-arch-deepseek4/src/mtp.rs | 2 +- crates/hipfire-arch-minimax/Cargo.toml | 1 + .../examples/ep_minimax.rs | 2 +- crates/hipfire-arch-minimax/src/forward.rs | 2 +- crates/hipfire-arch-qwen35/Cargo.toml | 1 + .../examples/qwen_dense_tp2_parity.rs | 2 +- .../examples/test_qwen35_load_multi.rs | 2 +- .../examples/test_qwen35_state_multi.rs | 2 +- .../src/qwen35/ep_batch.rs | 25 +- .../hipfire-arch-qwen35/src/qwen35/forward.rs | 2 +- .../hipfire-arch-qwen35/src/qwen35/weights.rs | 2 +- crates/hipfire-arch-qwen35/tests/pp_parity.rs | 2 +- crates/hipfire-generate/Cargo.toml | 1 + crates/hipfire-generate/src/batch.rs | 6 +- crates/hipfire-hardware/Cargo.toml | 14 + .../src/lib.rs} | 270 ++++++++----- crates/hipfire-hardware/src/mesh.rs | 368 ++++++++++++++++++ crates/hipfire-hardware/tests/ownership.rs | 28 ++ crates/hipfire-loader/Cargo.toml | 1 + crates/hipfire-loader/src/carriers.rs | 4 +- crates/hipfire-loader/src/lib.rs | 2 +- crates/hipfire-runtime/Cargo.toml | 1 + .../bench_tp_graph_host_overhead_gfx1201.rs | 2 +- .../ds4_gfx1201_owner_worker_transport.rs | 2 +- .../examples/ep_decode_parity.rs | 2 +- .../examples/pp_parity_chatml.rs | 2 +- .../tp4_cross_device_graph_barrier.rs | 2 +- .../examples/tp_allreduce_smoke.rs | 12 +- crates/hipfire-runtime/src/config.rs | 2 +- crates/hipfire-runtime/src/ep.rs | 9 +- crates/hipfire-runtime/src/lib.rs | 1 - crates/hipfire-runtime/src/llama.rs | 6 +- crates/hipfire-runtime/src/model_load.rs | 2 +- crates/saddle-lab/Cargo.toml | 1 + crates/saddle-lab/examples/gpus_smoke.rs | 2 +- crates/saddle-lab/examples/pp2_vram_probe.rs | 2 +- crates/saddle-lab/examples/pp_parity.rs | 2 +- 44 files changed, 674 insertions(+), 150 deletions(-) create mode 100644 crates/hipfire-hardware/Cargo.toml rename crates/{hipfire-runtime/src/multi_gpu.rs => hipfire-hardware/src/lib.rs} (91%) create mode 100644 crates/hipfire-hardware/src/mesh.rs create mode 100644 crates/hipfire-hardware/tests/ownership.rs diff --git a/Cargo.lock b/Cargo.lock index 4ed89d3d8d..0f90443e0c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1126,6 +1126,7 @@ dependencies = [ "hipfire-config", "hipfire-dispatch", "hipfire-ds4-parent", + "hipfire-hardware", "hipfire-reap", "hipfire-runtime", "libloading", @@ -1219,6 +1220,7 @@ dependencies = [ "hip-bridge", "hipfire-config", "hipfire-dispatch", + "hipfire-hardware", "hipfire-reap", "hipfire-runtime", "rdna-compute", @@ -1259,6 +1261,7 @@ dependencies = [ "hipfire-arch-qwen35-vl", "hipfire-config", "hipfire-dispatch", + "hipfire-hardware", "hipfire-reap", "hipfire-runtime", "rdna-compute", @@ -1446,6 +1449,7 @@ dependencies = [ "hipfire-config", "hipfire-dispatch", "hipfire-engine", + "hipfire-hardware", "hipfire-loader", "hipfire-pflash", "hipfire-runtime", @@ -1456,6 +1460,14 @@ dependencies = [ "tracing", ] +[[package]] +name = "hipfire-hardware" +version = "0.3.0" +dependencies = [ + "hip-bridge", + "rdna-compute", +] + [[package]] name = "hipfire-loader" version = "0.3.0" @@ -1475,6 +1487,7 @@ dependencies = [ "hipfire-arch-qwen35", "hipfire-arch-qwen35-vl", "hipfire-config", + "hipfire-hardware", "hipfire-runtime", "rdna-compute", "saddle-core", @@ -1563,6 +1576,7 @@ dependencies = [ "hipfire-detect", "hipfire-dispatch", "hipfire-engine", + "hipfire-hardware", "hipfire-loader", "hipfire-pflash", "libc", @@ -2902,6 +2916,7 @@ dependencies = [ "hipfire-detect", "hipfire-dispatch", "hipfire-engine", + "hipfire-hardware", "hipfire-loader", "hipfire-pflash", "hipfire-runtime", diff --git a/Cargo.toml b/Cargo.toml index 1f764a9b7b..ee69d0ff56 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ members = [ "crates/hipfire-generate", "crates/radiowave", "crates/hipfire-runtime", + "crates/hipfire-hardware", "crates/hipfire-arch-qwen35", "crates/hipfire-pflash", "crates/hipfire-arch-qwen35-vl", diff --git a/crates/hipfire-arch-deepseek4/Cargo.toml b/crates/hipfire-arch-deepseek4/Cargo.toml index 338dd916b5..75ee4df609 100644 --- a/crates/hipfire-arch-deepseek4/Cargo.toml +++ b/crates/hipfire-arch-deepseek4/Cargo.toml @@ -13,6 +13,7 @@ description = "DeepSeek V4 Flash architecture for hipfire (Hyper-Connections + c [dependencies] hipfire-config = { path = "../hipfire-config" } hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-ds4-parent = { path = "../hipfire-ds4-parent" } hip-bridge = { path = "../hip-bridge" } rdna-compute = { path = "../rdna-compute" } diff --git a/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs b/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs index 22e16f8a31..8e57147221 100644 --- a/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs +++ b/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs @@ -32,7 +32,7 @@ fn main() { use hipfire_arch_deepseek4::{DeepseekV4, DeepseekV4State}; use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::HfqFile; - use hipfire_runtime::multi_gpu::Gpus; + use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; diff --git a/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs b/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs index e74f694b10..56a66c1f68 100644 --- a/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs +++ b/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs @@ -17,7 +17,7 @@ use hipfire_arch_deepseek4::forward::{ use hipfire_arch_deepseek4::{DeepseekV4, DeepseekV4State}; use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::HfqFile; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::path::{Path, PathBuf}; diff --git a/crates/hipfire-arch-deepseek4/src/ep.rs b/crates/hipfire-arch-deepseek4/src/ep.rs index e0104087d2..6f2a97a882 100644 --- a/crates/hipfire-arch-deepseek4/src/ep.rs +++ b/crates/hipfire-arch-deepseek4/src/ep.rs @@ -11,7 +11,7 @@ use crate::forward::{ use crate::forward::{precompute_positions, precompute_token_id, update_attn_state_host, update_pos_array_host, update_token_id_host, ensure_compressor_capacity}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::superop::{self, SuperOpKind}; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use rdna_compute::{Gpu, GpuTensor}; // ───────────────────────── Ship 6 substrate-EP (DeepSeek-V4) ───────────────── @@ -41,7 +41,7 @@ use rdna_compute::{Gpu, GpuTensor}; /// access enabled for the fast peer-direct all-reduce. #[allow(clippy::too_many_arguments)] pub fn forward_ep( - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, weights_per_rank: &[DeepseekV4Weights], cfg: &DeepseekV4Config, state_per_rank: &mut [DeepseekV4State], @@ -116,7 +116,7 @@ pub fn forward_ep( #[allow(clippy::too_many_arguments)] fn forward_ep_tp_graph_body( - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, weights_per_rank: &[DeepseekV4Weights], cfg: &DeepseekV4Config, state_per_rank: &mut [DeepseekV4State], @@ -161,7 +161,7 @@ fn forward_ep_tp_graph_body( ) } -fn sync_ep_ranks(gpus: &mut hipfire_runtime::multi_gpu::Gpus, label: &str) -> Result<(), String> { +fn sync_ep_ranks(gpus: &mut hipfire_hardware::Gpus, label: &str) -> Result<(), String> { for rank in 0..gpus.devices.len() { gpus.devices[rank] .bind_thread() @@ -176,7 +176,7 @@ fn sync_ep_ranks(gpus: &mut hipfire_runtime::multi_gpu::Gpus, label: &str) -> Re #[allow(clippy::too_many_arguments)] fn forward_ep_tp_graph( - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, weights_per_rank: &[DeepseekV4Weights], cfg: &DeepseekV4Config, state_per_rank: &mut [DeepseekV4State], @@ -324,7 +324,7 @@ fn forward_ep_tp_graph( #[allow(clippy::too_many_arguments)] fn forward_ep_direct( - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, weights_per_rank: &[DeepseekV4Weights], cfg: &DeepseekV4Config, state_per_rank: &mut [DeepseekV4State], diff --git a/crates/hipfire-arch-deepseek4/src/forward.rs b/crates/hipfire-arch-deepseek4/src/forward.rs index a646f869c0..5a0c29c2f2 100644 --- a/crates/hipfire-arch-deepseek4/src/forward.rs +++ b/crates/hipfire-arch-deepseek4/src/forward.rs @@ -12875,7 +12875,7 @@ pub fn forward_prefill_batch_chunked( /// rank enters the next stage with bit-identical residual streams. #[allow(clippy::too_many_arguments)] pub fn forward_ep_prefill_batch_chunked( - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, weights_per_rank: &[DeepseekV4Weights], cfg: &DeepseekV4Config, state_per_rank: &mut [DeepseekV4State], diff --git a/crates/hipfire-arch-deepseek4/src/mtp.rs b/crates/hipfire-arch-deepseek4/src/mtp.rs index 5606d89b7c..fc3cb3fb3a 100644 --- a/crates/hipfire-arch-deepseek4/src/mtp.rs +++ b/crates/hipfire-arch-deepseek4/src/mtp.rs @@ -526,7 +526,7 @@ fn mtp_head( /// all-reduce. #[allow(clippy::too_many_arguments)] pub fn mtp_forward_ep( - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, weights_per_rank: &[DeepseekV4Weights], cfg: &DeepseekV4Config, state_per_rank: &mut [DeepseekV4State], diff --git a/crates/hipfire-arch-minimax/Cargo.toml b/crates/hipfire-arch-minimax/Cargo.toml index 23f4e3e03b..f6f292446c 100644 --- a/crates/hipfire-arch-minimax/Cargo.toml +++ b/crates/hipfire-arch-minimax/Cargo.toml @@ -18,6 +18,7 @@ deltanet = ["hipfire-runtime/deltanet", "rdna-compute/deltanet"] [dependencies] hipfire-config = { path = "../hipfire-config" } hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-dispatch = { path = "../hipfire-dispatch", features = ["from-hip-error"] } hip-bridge = { path = "../hip-bridge" } rdna-compute = { path = "../rdna-compute" } diff --git a/crates/hipfire-arch-minimax/examples/ep_minimax.rs b/crates/hipfire-arch-minimax/examples/ep_minimax.rs index 4a2c90931e..24c88d6745 100644 --- a/crates/hipfire-arch-minimax/examples/ep_minimax.rs +++ b/crates/hipfire-arch-minimax/examples/ep_minimax.rs @@ -36,7 +36,7 @@ fn main() { use hipfire_arch_minimax::forward; use hipfire_arch_minimax::minimax::{MiniMaxConfig, MiniMaxState, MiniMaxWeights}; use hipfire_runtime::hfq::HfqFile; - use hipfire_runtime::multi_gpu::Gpus; + use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; diff --git a/crates/hipfire-arch-minimax/src/forward.rs b/crates/hipfire-arch-minimax/src/forward.rs index 6c9830de63..aa42b84f97 100644 --- a/crates/hipfire-arch-minimax/src/forward.rs +++ b/crates/hipfire-arch-minimax/src/forward.rs @@ -1635,7 +1635,7 @@ pub fn forward_batch( /// enabled for the fast peer-direct all-reduce. #[allow(clippy::too_many_arguments)] pub fn forward_ep( - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, weights_per_rank: &[MiniMaxWeights], cfg: &MiniMaxConfig, state_per_rank: &mut [MiniMaxState], diff --git a/crates/hipfire-arch-qwen35/Cargo.toml b/crates/hipfire-arch-qwen35/Cargo.toml index 2671c02b62..abe696730d 100644 --- a/crates/hipfire-arch-qwen35/Cargo.toml +++ b/crates/hipfire-arch-qwen35/Cargo.toml @@ -14,6 +14,7 @@ deltanet = ["hipfire-runtime/deltanet", "rdna-compute/deltanet"] [dependencies] hipfire-config = { path = "../hipfire-config" } hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-arch-qwen35-vl = { path = "../hipfire-arch-qwen35-vl" } hipfire-dispatch = { path = "../hipfire-dispatch", features = ["from-hip-error"] } hip-bridge = { path = "../hip-bridge" } diff --git a/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs b/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs index 8c3610e729..93dea57513 100644 --- a/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs +++ b/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs @@ -6,7 +6,7 @@ use hipfire_arch_qwen35::qwen35::{ use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::kv_mode::KvMode; use hipfire_runtime::llama::{KvCache, KvCacheExt, KvDims, KvLayers, KvTarget}; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::Gpu; diff --git a/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs b/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs index fa06e8e313..a5343e0fe4 100644 --- a/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs +++ b/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs @@ -15,7 +15,7 @@ use hipfire_arch_qwen35::qwen35; use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::hfq::HfqFile; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use std::path::Path; fn main() { diff --git a/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs b/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs index ca9c2f18e4..1842e45478 100644 --- a/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs +++ b/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs @@ -19,7 +19,7 @@ use hipfire_arch_qwen35::qwen35::{self, DeltaNetState, LayerType, Qwen35ScratchS use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use std::path::Path; fn main() { diff --git a/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs b/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs index 39a642d05e..561d84d74f 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs @@ -56,7 +56,7 @@ use hipfire_runtime::llama::weight_gemv_prerotated; use hipfire_runtime::llama::weight_gemv_swiglu_residual; use hipfire_runtime::llama::EmbeddingFormat; use hipfire_runtime::llama::WeightTensor; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use rdna_compute::DType; use rdna_compute::GpuTensor; @@ -836,7 +836,7 @@ pub fn validate_ep_batch_compatibility( let requested_bytes = usize::try_from(requested_bytes_u64) .map_err(|_| HipError::new(0, "peer requested bytes overflow usize"))?; let peer_bytes_per_rank = - hipfire_runtime::multi_gpu::peer_reduce_scratch_bytes_per_rank(4, requested_bytes) + hipfire_hardware::peer_reduce_scratch_bytes_per_rank(4, requested_bytes) .ok_or_else(|| HipError::new(0, "peer per-rank projection overflow"))?; let per_rank_total = per_rank_decode .checked_add(per_rank_seed_pbs) @@ -926,7 +926,7 @@ pub struct Qwen35DecodeBatchEpState { dim: usize, norm_eps: f32, expert_to_rank: Box<[u8]>, - peer_lease: Option, + peer_lease: Option, } /// Transactional ownership guard for `Qwen35DecodeBatchEpState::new`. @@ -938,7 +938,7 @@ struct EpBatchBuildGuard { seed_pbs: Vec>, seed_partials: Vec>, scratches: Vec>, - lease: Option, + lease: Option, } impl EpBatchBuildGuard { @@ -967,7 +967,7 @@ impl EpBatchBuildGuard { fn set_scratch(&mut self, idx: usize, v: Qwen35Scratch) { self.scratches[idx] = Some(v); } - fn set_lease(&mut self, lease: hipfire_runtime::multi_gpu::PeerReduceScratchLease) { + fn set_lease(&mut self, lease: hipfire_hardware::PeerReduceScratchLease) { self.lease = Some(lease); } /// Rollback on owning devices, attempting every free, preserving init error plus first cleanup error. @@ -1044,7 +1044,7 @@ impl EpBatchBuildGuard { Vec, Vec, Vec, - Option, + Option, ) { let ranks = self .ranks @@ -2452,15 +2452,20 @@ pub fn forward_prefill_batch_ep( t_chunk += t_c.elapsed().as_secs_f64() * 1000.0; } - // 3. All-reduce the routed partials, add into each rank's residual. + // 3. All-reduce the routed partials over the full EP group, then add + // the reduced result into each rank's residual. if is_moe && !ep_skip_ar { let t_a = std::time::Instant::now(); - let refs: Vec<&hip_bridge::DeviceBuffer> = partials.iter().map(|p| &p.buf).collect(); + let refs: Vec<&hip_bridge::DeviceBuffer> = + partials.iter().map(|partial| &partial.buf).collect(); + let group: Vec = (0..refs.len()).collect(); if ep_peer_ar { - gpus.all_reduce_sum_f32_peer(&refs, n * dim) + gpus + .all_reduce_sum_f32_peer(&group, &refs, n * dim) .map_err(|e| HipError::new(0, &e.to_string()))?; } else { - gpus.all_reduce_sum_f32(&refs, n * dim) + gpus + .all_reduce_sum_f32(&group, &refs, n * dim) .map_err(|e| HipError::new(0, &e.to_string()))?; } if ep_timing { diff --git a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs index 00c86bd151..f45076f570 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs @@ -51,7 +51,7 @@ use hipfire_runtime::llama::EmbeddingFormat; use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::llama::ParoRotation; use hipfire_runtime::llama::WeightTensor; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hipfire_runtime::tp_shard::ShardConfig; use rdna_compute::DType; use rdna_compute::Gpu; diff --git a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs index cf25aa7ead..fc74135817 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs @@ -12,7 +12,7 @@ use hip_bridge::HipResult; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::EmbeddingFormat; use hipfire_runtime::llama::WeightTensor; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hipfire_runtime::screen_weight_tensor; use hipfire_runtime::MmqScreenable; use rdna_compute::DType; diff --git a/crates/hipfire-arch-qwen35/tests/pp_parity.rs b/crates/hipfire-arch-qwen35/tests/pp_parity.rs index a6b67c53db..46b4d59a04 100644 --- a/crates/hipfire-arch-qwen35/tests/pp_parity.rs +++ b/crates/hipfire-arch-qwen35/tests/pp_parity.rs @@ -24,7 +24,7 @@ use hipfire_arch_qwen35::qwen35::{ use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use rdna_compute::Gpu; use std::path::Path; diff --git a/crates/hipfire-generate/Cargo.toml b/crates/hipfire-generate/Cargo.toml index b23a5b9f76..081c8929d7 100644 --- a/crates/hipfire-generate/Cargo.toml +++ b/crates/hipfire-generate/Cargo.toml @@ -28,6 +28,7 @@ rdna-compute = { path = "../rdna-compute" } saddle-core = { path = "../saddle-core" } hipfire-dispatch = { path = "../hipfire-dispatch" } hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-loader = { path = "../hipfire-loader" } hipfire-engine = { path = "../hipfire-engine" } hipfire-pflash = { path = "../hipfire-pflash" } diff --git a/crates/hipfire-generate/src/batch.rs b/crates/hipfire-generate/src/batch.rs index 685791d3d2..3b0f686c42 100644 --- a/crates/hipfire-generate/src/batch.rs +++ b/crates/hipfire-generate/src/batch.rs @@ -2556,7 +2556,7 @@ pub fn drive_qwen35_ep_continuous_batch( } }; ( - &mut ep.gpus as *mut hipfire_runtime::multi_gpu::Gpus, + &mut ep.gpus as *mut hipfire_hardware::Gpus, config as *const qwen35::Qwen35Config, weights as *const Vec, b, @@ -2571,7 +2571,7 @@ pub fn drive_qwen35_ep_continuous_batch( _ => return Err(BatchDriveError::Gpu("EP batch: not Qwen35 EP".to_string())), } }; - let gpus: &mut hipfire_runtime::multi_gpu::Gpus = unsafe { &mut *gpus_ptr }; + let gpus: &mut hipfire_hardware::Gpus = unsafe { &mut *gpus_ptr }; let config: &qwen35::Qwen35Config = unsafe { &*config_ptr }; let weights: &Vec = unsafe { &*weights_ptr }; let batch_state: &mut qwen35::Qwen35DecodeBatchEpState = unsafe { &mut *batch_ptr }; @@ -2591,7 +2591,7 @@ pub fn drive_qwen35_ep_continuous_batch( // Track last attested receipt for evidence; must be from runtime, never load logs. let mut last_receipt: Option = None; let fail_all = |sched: &mut ContinuousBatchScheduler, - gpus: &mut hipfire_runtime::multi_gpu::Gpus, + gpus: &mut hipfire_hardware::Gpus, batch_state: &mut qwen35::Qwen35DecodeBatchEpState, stdout: &mut std::io::Stdout, reason: String| diff --git a/crates/hipfire-hardware/Cargo.toml b/crates/hipfire-hardware/Cargo.toml new file mode 100644 index 0000000000..89c6967191 --- /dev/null +++ b/crates/hipfire-hardware/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "hipfire-hardware" +version.workspace = true +edition.workspace = true +license.workspace = true +description = "Hardware ownership and multi-device topology for hipfire" + +[dependencies] +hip-bridge = { path = "../hip-bridge" } +rdna-compute = { path = "../rdna-compute" } + +[features] +default = [] +deltanet = ["rdna-compute/deltanet"] diff --git a/crates/hipfire-runtime/src/multi_gpu.rs b/crates/hipfire-hardware/src/lib.rs similarity index 91% rename from crates/hipfire-runtime/src/multi_gpu.rs rename to crates/hipfire-hardware/src/lib.rs index d88214d294..8363fd7ceb 100644 --- a/crates/hipfire-runtime/src/multi_gpu.rs +++ b/crates/hipfire-hardware/src/lib.rs @@ -21,11 +21,55 @@ //! 3. Pass the multi-GPU coherence gate. use hip_bridge::{ - DeviceBuffer, Event, HipError, HipResult, RcclComms, HIP_ERROR_PEER_ACCESS_ALREADY_ENABLED, - HIP_ERROR_PEER_ACCESS_UNSUPPORTED, HIP_EVENT_DISABLE_TIMING, HIP_EVENT_RELEASE_TO_SYSTEM, + DeviceBuffer, Event, HipError, HipResult, HipRuntime, RcclComms, + HIP_ERROR_PEER_ACCESS_ALREADY_ENABLED, HIP_ERROR_PEER_ACCESS_UNSUPPORTED, + HIP_EVENT_DISABLE_TIMING, HIP_EVENT_RELEASE_TO_SYSTEM, }; use rdna_compute::{DType, Gpu, GpuTensor}; +mod mesh; +pub use mesh::{Axis, CollectiveHint, DeviceMesh, DimKind, MeshEpoch}; + +/// Device-resolution knobs used when constructing a [`Gpus`] owner. +/// +/// The hardware leaf reads the same legacy `HIPFIRE_*` environment variables +/// as the runtime did, without depending on runtime configuration. Higher +/// layers may eventually supply an explicit option set; the constructors +/// below intentionally keep the existing process-environment behavior. +#[derive(Clone, Debug, Default)] +pub struct DeviceResolveOpts { + pub tp_use_rccl: Option, + pub devices: Option, + pub emulate_gpus: Option, + pub allow_mixed_arch: bool, + pub uniform_vram_tolerance_gb: Option, +} + +impl DeviceResolveOpts { + pub fn from_env() -> Self { + Self { + tp_use_rccl: std::env::var("HIPFIRE_TP_USE_RCCL") + .ok() + .as_deref() + .map(|value| value != "0" && !value.eq_ignore_ascii_case("false")), + devices: std::env::var("HIPFIRE_DEVICES") + .ok() + .filter(|value| !value.is_empty()), + emulate_gpus: std::env::var("HIPFIRE_EMULATE_GPUS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|&count| count >= 2), + allow_mixed_arch: std::env::var("HIPFIRE_ALLOW_MIXED_ARCH") + .ok() + .as_deref() + == Some("1"), + uniform_vram_tolerance_gb: std::env::var("HIPFIRE_UNIFORM_VRAM_TOLERANCE_GB") + .ok() + .and_then(|value| value.parse().ok()), + } + } +} + /// Stream-event handoff returned by `Gpus::boundary_copy`. When the src /// device has an active stream, `completion` holds a HIP event recorded /// after the async peer copy; `Gpus::wait_boundary` makes the dst stream @@ -871,7 +915,7 @@ impl Gpus { if self.rccl_comms.is_some() { return Ok(()); } - if matches!(crate::config::get().tp_use_rccl, Some(false)) { + if matches!(DeviceResolveOpts::from_env().tp_use_rccl, Some(false)) { return Err(HipError::new( 0, "ensure_rccl: HIPFIRE_TP_USE_RCCL=0 — RCCL path opted out. \ @@ -896,71 +940,73 @@ impl Gpus { Ok(()) } - /// All-reduce-sum of f32 buffers across all ranks. `buffers[r]` must - /// be a device pointer on `devices[r]` holding `count` f32 elements; - /// after this call, each buffer holds the element-wise sum across - /// all ranks. In-place (send == recv) — saves a memcpy and matches - /// how the TP forward path uses the result. + /// All-reduce-sum of f32 buffers across a participating device group. + /// `buffers[k]` must be a device pointer on `self.devices[group[k]]` + /// holding `count` f32 elements; after this call, every group buffer + /// holds the element-wise sum. In-place (send == recv). /// - /// Requires each device to have an `active_stream` set (the stream - /// the collective runs on). Synchronization is the caller's - /// responsibility: this call enqueues the collective and returns - /// immediately; the buffers are valid only after a subsequent - /// `stream_synchronize` (or a downstream dispatch that's already - /// ordered behind the same stream). - pub fn all_reduce_sum_f32(&mut self, buffers: &[&DeviceBuffer], count: usize) -> HipResult<()> { - if buffers.len() != self.devices.len() { + /// The RCCL communicator is owned by this entire `Gpus` set, so this + /// implementation accepts only the full ordered group `0..n`. Use + /// [`Self::all_reduce_sum_f32_peer`] for genuine subgroup reductions. + /// Synchronization remains the caller's responsibility. + pub fn all_reduce_sum_f32( + &mut self, + group: &[usize], + buffers: &[&DeviceBuffer], + count: usize, + ) -> HipResult<()> { + if buffers.len() != group.len() { return Err(HipError::new( 0, &format!( - "all_reduce_sum_f32: buffers.len()={} != n_devices={}", + "all_reduce_sum_f32: buffers.len()={} != group.len()={}", buffers.len(), - self.devices.len() + group.len() ), )); } - // Single-rank (TP=1) degenerate case: the all-reduce-sum over one - // buffer is the identity — the buffer already holds the only rank's - // partial. Short-circuit so the TP=1 EP path is a pure single-GPU - // reference that exercises the full EP executor WITHOUT requiring - // librccl (a 1-rank communicator would also work, but skipping it - // keeps TP=1 dependency-free and the parity baseline trivially exact). - if self.devices.len() == 1 { + let n = self.devices.len(); + if group.len() != n || !group.iter().copied().eq(0..n) { + return Err(HipError::new( + 0, + "all_reduce_sum_f32 (RCCL): sub-group reduction needs ncclCommSplit \ + (Phase 5b); use all_reduce_sum_f32_peer for sub-groups.", + )); + } + // Single-rank (TP=1) all-reduce is the identity and does not require + // librccl. + if n == 1 { return Ok(()); } self.ensure_rccl()?; - // Borrow-check note: `self.rccl_comms.as_ref()` projects through - // a single field, leaving `self.devices` independently - // borrow-able for the per-rank stream lookup below. let rccl = self.rccl_comms.as_ref().expect("ensure_rccl populated it"); - rccl.group_start() .map_err(|e| HipError::new(0, &format!("ncclGroupStart: {e}")))?; - for (r, buf) in buffers.iter().enumerate() { - let dev = &self.devices[r]; + for (rank, buf) in buffers.iter().enumerate() { + let dev = &self.devices[rank]; dev.bind_thread()?; let stream = dev.active_stream.as_ref().ok_or_else(|| { HipError::new( 0, &format!( - "all_reduce_sum_f32: device {r} has no active_stream — \ - set `gpus.devices[r].active_stream = Some(stream)` before calling.", + "all_reduce_sum_f32: device {rank} has no active_stream — \ + set `gpus.devices[{rank}].active_stream = Some(stream)` before calling.", ), ) })?; // SAFETY: `buf` is a live device buffer of `count` f32 on device - // `r`, and `stream` is that device's active stream. + // `rank`, and `stream` is that device's active stream. unsafe { rccl.all_reduce_sum_f32( - r, + rank, buf.as_ptr() as *const f32, buf.as_ptr() as *mut f32, count, stream.raw_ptr(), ) } - .map_err(|e| HipError::new(0, &format!("ncclAllReduce rank={r}: {e}")))?; + .map_err(|e| HipError::new(0, &format!("ncclAllReduce rank={rank}: {e}")))?; } rccl.group_end() .map_err(|e| HipError::new(0, &format!("ncclGroupEnd: {e}")))?; @@ -1435,23 +1481,15 @@ impl Gpus { Ok(()) } - /// All-reduce-sum of f32 buffers across all ranks via **direct peer copy + - /// local add** — bypassing RCCL. On consumer/prosumer RDNA P2P (no xGMI, - /// e.g. hiptrx 4× gfx1201), `ncclAllReduce` costs ~40 ms/call for these - /// small/medium messages regardless of NCCL_PROTO/CHANNELS/BUFFSIZE/ - /// SOCKET_IFNAME, while this path is ~1 ms. Used by EP prefill and TP; EP - /// decode's tiny per-token reduce stays on RCCL (already fast). PP never - /// all-reduces (it uses `boundary_copy` point-to-point). + /// All-reduce-sum of f32 buffers across a participating device group via + /// direct peer copies and local adds, bypassing RCCL. /// - /// Algorithm (N-rank, race-free): **phase 1** copies every OTHER rank's - /// ORIGINAL buffer into a local temp (all reads, no writes); a barrier - /// (`wait_boundary`); **phase 2** adds the peer temps into the local buffer. - /// All-reads-before-writes ⇒ no cross-device read/write race. `n==1` is the - /// identity (no-op). Requires peer access (caller's `enable_peer_all`) for - /// the fast P2P path; without it `boundary_copy` host-stages (slower but - /// correct). In-place: `buffers[r]` is both input and output. + /// `group` lists global device IDs and `buffers[k]` belongs to + /// `self.devices[group[k]]`. Unlike the RCCL path this is genuinely + /// subgroup-capable, which is required for composed TP×EP meshes. pub fn all_reduce_sum_f32_peer( &mut self, + group: &[usize], buffers: &[&DeviceBuffer], count: usize, ) -> HipResult<()> { @@ -1464,56 +1502,75 @@ impl Gpus { "all_reduce_sum_f32_peer: peer scratch is leased — use leased API or release lease", )); } + let g = group.len(); let n = self.devices.len(); - if buffers.len() != n { + if buffers.len() != g { return Err(HipError::new( 0, &format!( - "all_reduce_sum_f32_peer: buffers.len()={} != n_devices={n}", + "all_reduce_sum_f32_peer: buffers.len()={} != group.len()={g}", buffers.len() ), )); } - if n == 1 { + for (index, &device) in group.iter().enumerate() { + if device >= n || group[..index].contains(&device) { + return Err(HipError::new( + 0, + &format!( + "all_reduce_sum_f32_peer: group contains invalid or duplicate device {device}" + ), + )); + } + } + if g <= 1 { return Ok(()); } - let bytes = count * 4; + let bytes = count + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| HipError::new(0, "all_reduce_sum_f32_peer: count overflow"))?; self.ensure_peer_ar_tmp(bytes)?; + // Phase 1: read every peer's ORIGINAL buffer into a local temp. - let mut evts = Vec::with_capacity(n * (n - 1)); - for r in 0..n { - let mut slot = 0usize; - for j in 0..n { - if j == r { + let mut events = Vec::with_capacity(g * (g - 1)); + for (k, &device_k) in group.iter().enumerate() { + let mut slot = 0; + for (m, &device_m) in group.iter().enumerate() { + if m == k { continue; } - let evt = - self.boundary_copy(j, r, buffers[j], &self.peer_ar_tmp[r][slot], bytes)?; - evts.push(evt); + let event = self.boundary_copy( + device_m, + device_k, + buffers[m], + &self.peer_ar_tmp[device_k][slot], + bytes, + )?; + events.push(event); slot += 1; } } - for evt in evts { - self.wait_boundary(evt)?; + for event in events { + self.wait_boundary(event)?; } - // Phase 2: add the peer temps into each rank's buffer. - for r in 0..n { + // Phase 2: add peer temps into each rank's buffer. + for (k, &device_k) in group.iter().enumerate() { let dst = GpuTensor { - buf: unsafe { buffers[r].alias() }, + buf: unsafe { buffers[k].alias() }, shape: vec![count], dtype: DType::F32, }; - let srcs: Vec = (0..n - 1) + let srcs: Vec = (0..g - 1) .map(|slot| GpuTensor { - buf: unsafe { self.peer_ar_tmp[r][slot].alias() }, + buf: unsafe { self.peer_ar_tmp[device_k][slot].alias() }, shape: vec![count], dtype: DType::F32, }) .collect(); - self.devices[r].bind_thread()?; + self.devices[device_k].bind_thread()?; for src in &srcs { - self.devices[r].add_inplace_f32(&dst, src)?; + self.devices[device_k].add_inplace_f32(&dst, src)?; } } Ok(()) @@ -1848,27 +1905,50 @@ fn uniform_split_counts(n_devices: usize, n_layers: usize) -> Vec { /// Resolve logical device IDs after the physical `hardware.devices` list has /// been installed as `ROCR_VISIBLE_DEVICES` and HIP has received the matching /// post-filter logical IDs. When unset, take the first `n_devices` visible IDs. +/// Map each requested logical device id into the physical range +/// `[0, real_count)` by euclidean remainder. Used by +/// `HIPFIRE_EMULATE_GPUS` to alias N logical devices onto fewer physical +/// devices. A non-positive `real_count` is left untouched so the caller can +/// surface the underlying HIP error. +fn alias_ids(ids: &[i32], real_count: i32) -> Vec { + if real_count <= 0 { + return ids.to_vec(); + } + ids.iter().map(|&id| id.rem_euclid(real_count)).collect() +} + +/// Resolve logical IDs from `HIPFIRE_DEVICES`, or use the first N visible IDs. +/// `HIPFIRE_EMULATE_GPUS` optionally aliases those IDs into the physical +/// runtime device range for debug-only multi-rank emulation. fn resolve_device_ids(n_devices: usize) -> HipResult> { - if let Some(ref s) = crate::config::get().devices { - let ids: Vec = s + let opts = DeviceResolveOpts::from_env(); + let ids: Vec = if let Some(value) = &opts.devices { + let parsed = value .split(',') - .map(|p| p.trim()) - .filter(|p| !p.is_empty()) - .map(|p| p.parse::()) - .collect::>() - .map_err(|e| HipError::new(0, &format!("hardware.devices parse: {e}")))?; - if ids.len() < n_devices { + .map(str::trim) + .filter(|part| !part.is_empty()) + .map(str::parse::) + .collect::, _>>() + .map_err(|error| HipError::new(0, &format!("HIPFIRE_DEVICES parse: {error}")))?; + if parsed.len() < n_devices { return Err(HipError::new( 0, &format!( - "hardware.devices exposes {} ids but n_devices = {n_devices}", - ids.len(), + "HIPFIRE_DEVICES has {} ids but n_devices = {n_devices}", + parsed.len() ), )); } - return Ok(ids[..n_devices].to_vec()); + parsed[..n_devices].to_vec() + } else { + (0..n_devices as i32).collect() + }; + + if opts.emulate_gpus.is_some() { + let real_count = HipRuntime::load()?.device_count()?; + return Ok(alias_ids(&ids, real_count)); } - Ok((0..n_devices as i32).collect()) + Ok(ids) } fn construct_devices(ids: &[i32]) -> HipResult> { @@ -1884,18 +1964,18 @@ fn preflight_vram_with_opts(devices: &[Gpu], check_vram_delta: bool) -> HipResul return Ok(()); } let arch0 = devices[0].arch.clone(); - let allow_mixed = crate::config::get().allow_mixed_arch; + let opts = DeviceResolveOpts::from_env(); let mut frees = Vec::with_capacity(devices.len()); - for d in devices { - if d.arch != arch0 { - if allow_mixed { + for device in devices { + if device.arch != arch0 { + if opts.allow_mixed_arch { eprintln!( "preflight_vram: mixed-arch detected — dev 0 is {arch0}, dev {} is {}. \ Proceeding because HIPFIRE_ALLOW_MIXED_ARCH=1. \ Per-arch JIT cache will be populated on first run; boundary_copy uses \ hipMemcpyPeer / hipMemcpyPeerAsync which fall through to host-staging \ if peer access is unsupported by the pair (correctness holds either way).", - d.device_id, d.arch, + device.device_id, device.arch, ); } else { return Err(HipError::new( @@ -1903,13 +1983,13 @@ fn preflight_vram_with_opts(devices: &[Gpu], check_vram_delta: bool) -> HipResul &format!( "preflight_vram: arch mismatch — dev 0 is {arch0}, dev {} is {}. \ Mixed-arch is not supported by default; set HIPFIRE_ALLOW_MIXED_ARCH=1 to override.", - d.device_id, d.arch, + device.device_id, device.arch, ), )); } } - d.bind_thread()?; - let (free, _total) = d.hip.get_vram_info()?; + device.bind_thread()?; + let (free, _total) = device.hip.get_vram_info()?; frees.push(free); } if !check_vram_delta { @@ -1918,17 +1998,17 @@ fn preflight_vram_with_opts(devices: &[Gpu], check_vram_delta: bool) -> HipResul let max_free = *frees.iter().max().unwrap(); let min_free = *frees.iter().min().unwrap(); let delta_gb = (max_free - min_free) as f64 / 1e9; - let tol_gb = crate::config::get() + let tolerance_gb = opts .uniform_vram_tolerance_gb - .map(|t| t as f64) + .map(|value| value as f64) .unwrap_or(DEFAULT_VRAM_TOLERANCE_GB); - if delta_gb > tol_gb { + if delta_gb > tolerance_gb { return Err(HipError::new( 0, &format!( "preflight_vram: VRAM delta {:.1} GiB exceeds tolerance {:.1} GiB. \ Override via HIPFIRE_UNIFORM_VRAM_TOLERANCE_GB or use init_layers().", - delta_gb, tol_gb, + delta_gb, tolerance_gb, ), )); } diff --git a/crates/hipfire-hardware/src/mesh.rs b/crates/hipfire-hardware/src/mesh.rs new file mode 100644 index 0000000000..76680f29d7 --- /dev/null +++ b/crates/hipfire-hardware/src/mesh.rs @@ -0,0 +1,368 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 alpineq +// hipfire — see LICENSE and NOTICE in the project root. + +//! Pure rectangular device topology for pipeline, tensor, and expert +//! parallelism. +//! +//! A mesh is an ordered set of named axes. Device IDs are the row-major +//! flattening of coordinates (the final axis varies fastest). This module is +//! deliberately independent of GPU handles, carrier policy, loading, and +//! allocation: it only answers placement and collective-group questions. + +use std::sync::atomic::{AtomicU64, Ordering}; + +static NEXT_MESH_EPOCH: AtomicU64 = AtomicU64::new(1); + +/// Identity of one admitted mesh generation. +/// +/// Cloning or squeezing a mesh preserves this identity. A fresh call to +/// [`DeviceMesh::single`] or [`DeviceMesh::rect`] receives a new epoch, even +/// when its shape happens to match another mesh. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)] +pub struct MeshEpoch(u64); + +impl MeshEpoch { + /// Return the process-local epoch number for diagnostics and cache keys. + pub fn as_u64(self) -> u64 { + self.0 + } +} + +/// The named parallelism dimensions of a device coordinate. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)] +pub enum DimKind { + /// Pipeline stages. Layers are banded across this axis and residuals cross + /// stage boundaries point-to-point; PP is never an all-reduce axis. + Pp, + /// Tensor-parallel ranks participating in dense row-sharded collectives. + Tp, + /// Expert-parallel ranks participating in routed-expert collectives. + Ep, +} + +/// One axis in a rectangular mesh. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)] +pub struct Axis { + pub kind: DimKind, + pub size: usize, +} + +/// A collective or point-to-point operation implied by mesh placement. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)] +pub enum CollectiveHint { + /// Reduce values across the named axis group. + AllReduce { kind: DimKind }, + /// Transfer the residual from one pipeline stage to the next. + /// + /// `src` and `dst` are stage coordinates (not physical device IDs). Use + /// [`DeviceMesh::stage_devices`] to expand a stage coordinate into its + /// global device IDs when TP or EP is composed with PP. + BandXfer { src: usize, dst: usize }, +} + +/// A rectangular named-axis mesh. +/// +/// The empty axis list is the single-device topology. Size-one axes remain in +/// the shape supplied to [`DeviceMesh::rect`]; [`DeviceMesh::squeezed`] drops +/// those axes while retaining the same [`MeshEpoch`]. Equality is +/// identity-sensitive: independently constructed meshes compare unequal even +/// when their shapes match. +#[derive(Clone, Debug)] +pub struct DeviceMesh { + axes: Vec, + epoch: MeshEpoch, +} + +impl PartialEq for DeviceMesh { + fn eq(&self, other: &Self) -> bool { + // MeshEpoch is intentionally identity-sensitive: independently + // constructed meshes are distinct even when their shapes match. + self.epoch == other.epoch + } +} + +impl Eq for DeviceMesh {} + +impl DeviceMesh { + /// Build a rectangular mesh from `(kind, size)` pairs. + /// + /// Axis sizes are normalized to at least one so every mesh has a valid + /// coordinate space and at least one logical device. Named axes are kept + /// in caller order; the final axis varies fastest in the flattened ID. + pub fn rect(axes: &[(DimKind, usize)]) -> Self { + Self { + axes: axes + .iter() + .map(|&(kind, size)| Axis { + kind, + size: size.max(1), + }) + .collect(), + epoch: fresh_epoch(), + } + } + + /// The single-device topology: one logical device and no named axes. + pub fn single() -> Self { + Self::rect(&[]) + } + + /// Ordered named axes of this mesh. + pub fn axes(&self) -> &[Axis] { + &self.axes + } + + /// Identity of this mesh generation. + pub fn epoch(&self) -> MeshEpoch { + self.epoch + } + + /// Total number of logical devices (one for an empty mesh). + pub fn n_devices(&self) -> usize { + self.axes + .iter() + .map(|axis| axis.size) + .product::() + .max(1) + } + + /// Size of the first axis with `kind`, or one when it is absent. + pub fn size_of(&self, kind: DimKind) -> usize { + self.axes + .iter() + .find(|axis| axis.kind == kind) + .map_or(1, |axis| axis.size) + } + + /// Whether this mesh has a non-degenerate axis of `kind`. + pub fn has_axis(&self, kind: DimKind) -> bool { + self.axes + .iter() + .any(|axis| axis.kind == kind && axis.size > 1) + } + + /// Convert a flattened row-major device ID to an axis coordinate. + pub fn coord_of(&self, dev: usize) -> Vec { + let mut rem = dev % self.n_devices(); + let mut coord = vec![0; self.axes.len()]; + for (index, axis) in self.axes.iter().enumerate().rev() { + coord[index] = rem % axis.size; + rem /= axis.size; + } + coord + } + + /// Convert an axis coordinate to a flattened row-major device ID. + /// + /// Coordinates are debug-asserted to have the mesh rank. In non-debug + /// builds, missing entries default to zero and out-of-range entries clamp + /// to the final valid index, preserving a total function for diagnostics. + pub fn device_of(&self, coord: &[usize]) -> usize { + debug_assert_eq!(coord.len(), self.axes.len()); + self.axes.iter().enumerate().fold(0, |id, (index, axis)| { + let value = coord.get(index).copied().unwrap_or(0); + id * axis.size + value.min(axis.size - 1) + }) + } + + /// Return the devices sharing all coordinates except `kind`, ordered by + /// their index along that axis. An absent axis yields this device alone. + pub fn group_along(&self, kind: DimKind, coord: &[usize]) -> Vec { + let Some(axis_index) = self.axes.iter().position(|axis| axis.kind == kind) else { + return vec![self.device_of(coord)]; + }; + let size = self.axes[axis_index].size; + let mut base = self.normalized_coord(coord); + (0..size) + .map(|index| { + base[axis_index] = index; + self.device_of(&base) + }) + .collect() + } + + /// Return the pipeline stage containing `layer` under a uniform banding. + /// + /// The first `n_layers % pp_size` stages receive one extra layer. A mesh + /// without a PP axis has one stage. Valid layer indexes always map to a + /// non-empty stage; an out-of-range index is clamped to the final stage. + pub fn stage_for_layer(&self, layer: usize, n_layers: usize) -> usize { + let stages = self.size_of(DimKind::Pp); + if stages <= 1 || n_layers == 0 { + return 0; + } + let base = n_layers / stages; + let remainder = n_layers % stages; + let mut start = 0; + for stage in 0..stages { + let count = base + usize::from(stage < remainder); + if layer < start + count { + return stage; + } + start += count; + } + stages - 1 + } + + /// Return a point-to-point hint when the next layer crosses a PP band. + pub fn band_xfer_after(&self, layer: usize, n_layers: usize) -> Option { + if n_layers == 0 || layer >= n_layers.saturating_sub(1) { + return None; + } + let src = self.stage_for_layer(layer, n_layers); + let dst = self.stage_for_layer(layer + 1, n_layers); + (src != dst).then_some(CollectiveHint::BandXfer { src, dst }) + } + + /// Expand the PP stage in `coord` to its global device IDs. + /// + /// For a mesh without PP, all devices belong to the sole stage. For a + /// composed mesh, every TP/EP coordinate in the selected PP stage is + /// returned in row-major order. + pub fn stage_devices(&self, coord: &[usize]) -> Vec { + let Some(pp_index) = self.axes.iter().position(|axis| axis.kind == DimKind::Pp) else { + return (0..self.n_devices()).collect(); + }; + let stage = coord + .get(pp_index) + .copied() + .unwrap_or(0) + .min(self.axes[pp_index].size - 1); + (0..self.n_devices()) + .filter(|&device| self.coord_of(device)[pp_index] == stage) + .collect() + } + + /// Drop size-one axes while preserving the mesh epoch identity. + pub fn squeezed(&self) -> Self { + Self { + axes: self + .axes + .iter() + .copied() + .filter(|axis| axis.size > 1) + .collect(), + epoch: self.epoch, + } + } + + fn normalized_coord(&self, coord: &[usize]) -> Vec { + debug_assert_eq!(coord.len(), self.axes.len()); + self.axes + .iter() + .enumerate() + .map(|(index, axis)| coord.get(index).copied().unwrap_or(0).min(axis.size - 1)) + .collect() + } +} + +impl Default for DeviceMesh { + fn default() -> Self { + Self::single() + } +} + +fn fresh_epoch() -> MeshEpoch { + let mut current = NEXT_MESH_EPOCH.load(Ordering::Relaxed); + loop { + if current == u64::MAX { + panic!("MeshEpoch exhausted: no remaining issuable epochs"); + } + match NEXT_MESH_EPOCH.compare_exchange_weak( + current, + current + 1, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return MeshEpoch(current), + Err(actual) => current = actual, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn single_is_one_device_and_identity_collectives_are_noops() { + let mesh = DeviceMesh::single(); + assert_eq!(mesh.n_devices(), 1); + assert_eq!(mesh.axes(), &[]); + assert_eq!(mesh.coord_of(0), Vec::::new()); + assert_eq!(mesh.device_of(&[]), 0); + assert_eq!(mesh.group_along(DimKind::Tp, &[]), vec![0]); + assert_eq!(mesh.group_along(DimKind::Ep, &[]), vec![0]); + assert_eq!(mesh.stage_for_layer(0, 32), 0); + assert_eq!(mesh.band_xfer_after(0, 32), None); + assert_eq!(mesh.stage_devices(&[]), vec![0]); + } + + #[test] + fn single_and_empty_rect_have_same_shape_but_fresh_identity() { + let single = DeviceMesh::single(); + let empty_rect = DeviceMesh::rect(&[]); + assert_ne!(single, empty_rect); + assert_eq!(single.axes(), empty_rect.axes()); + assert_eq!(single.n_devices(), empty_rect.n_devices()); + assert_eq!(single.coord_of(0), empty_rect.coord_of(0)); + assert_ne!(single.epoch(), empty_rect.epoch()); + assert_eq!(single.epoch(), single.clone().epoch()); + } + + #[test] + fn pp_bands_and_boundary_hints_are_uniform() { + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 3)]); + assert_eq!((0..6).map(|l| mesh.stage_for_layer(l, 6)).collect::>(), + vec![0, 0, 1, 1, 2, 2]); + assert_eq!(mesh.band_xfer_after(1, 6), Some(CollectiveHint::BandXfer { src: 0, dst: 1 })); + assert_eq!(mesh.band_xfer_after(3, 6), Some(CollectiveHint::BandXfer { src: 1, dst: 2 })); + assert_eq!(mesh.band_xfer_after(5, 6), None); + assert_eq!(mesh.stage_devices(&[1]), vec![1]); + } + + #[test] + fn composed_coordinates_groups_stages_and_squeeze() { + let mesh = DeviceMesh::rect(&[ + (DimKind::Pp, 2), + (DimKind::Tp, 2), + (DimKind::Ep, 2), + ]); + assert_eq!(mesh.n_devices(), 8); + assert_eq!(mesh.coord_of(0), vec![0, 0, 0]); + assert_eq!(mesh.coord_of(7), vec![1, 1, 1]); + assert_eq!(mesh.device_of(&[1, 0, 1]), 6); + assert_eq!(mesh.group_along(DimKind::Tp, &[1, 0, 1]), vec![6, 7]); + assert_eq!(mesh.group_along(DimKind::Ep, &[1, 0, 1]), vec![4, 6]); + assert_eq!(mesh.stage_devices(&[1, 0, 0]), vec![4, 5, 6, 7]); + + let degenerate = DeviceMesh::rect(&[ + (DimKind::Pp, 2), + (DimKind::Tp, 1), + (DimKind::Ep, 2), + ]); + assert_eq!( + degenerate.squeezed().axes(), + &[ + Axis { + kind: DimKind::Pp, + size: 2 + }, + Axis { + kind: DimKind::Ep, + size: 2 + }, + ] + ); + assert_eq!(degenerate.epoch(), degenerate.squeezed().epoch()); + } + + #[test] + fn coordinate_round_trip_holds_for_every_device() { + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 3), (DimKind::Tp, 2), (DimKind::Ep, 2)]); + for device in 0..mesh.n_devices() { + assert_eq!(mesh.device_of(&mesh.coord_of(device)), device); + } + } +} diff --git a/crates/hipfire-hardware/tests/ownership.rs b/crates/hipfire-hardware/tests/ownership.rs new file mode 100644 index 0000000000..2d0f42ff52 --- /dev/null +++ b/crates/hipfire-hardware/tests/ownership.rs @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2026 alpineq +// hipfire — see LICENSE and NOTICE in the project root. + +use hipfire_hardware::{CollectiveHint, DeviceMesh, DimKind, Gpus}; +use std::path::Path; + +#[test] +fn hardware_leaf_exposes_owner_and_named_topology() { + let _owner = std::any::TypeId::of::(); + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + assert_eq!(mesh.n_devices(), 4); + assert_eq!(mesh.group_along(DimKind::Tp, &[1, 0]), vec![2, 3]); + assert_eq!( + mesh.band_xfer_after(0, 2), + Some(CollectiveHint::BandXfer { src: 0, dst: 1 }) + ); +} + +#[test] +fn runtime_has_no_legacy_owner_or_compatibility_reexport() { + let crate_root = Path::new(env!("CARGO_MANIFEST_DIR")); + assert!(!crate_root.join("../hipfire-runtime/src/multi_gpu.rs").exists()); + let runtime_lib = std::fs::read_to_string(crate_root.join("../hipfire-runtime/src/lib.rs")) + .expect("runtime lib source"); + assert!(!runtime_lib.contains("pub mod multi_gpu")); + assert!(!runtime_lib.contains("pub use hipfire_hardware::*")); +} diff --git a/crates/hipfire-loader/Cargo.toml b/crates/hipfire-loader/Cargo.toml index 6e600544b7..4f67d31a77 100644 --- a/crates/hipfire-loader/Cargo.toml +++ b/crates/hipfire-loader/Cargo.toml @@ -25,6 +25,7 @@ ep-fault-inject = [] rdna-compute = { path = "../rdna-compute" } hipfire-config = { path = "../hipfire-config" } hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hip-bridge = { path = "../hip-bridge" } saddle-core = { path = "../saddle-core" } # Optional arch crate deps — feature-gated. diff --git a/crates/hipfire-loader/src/carriers.rs b/crates/hipfire-loader/src/carriers.rs index c7dba2dbaa..6e57ad8a05 100644 --- a/crates/hipfire-loader/src/carriers.rs +++ b/crates/hipfire-loader/src/carriers.rs @@ -267,9 +267,9 @@ fn load_qwen35_pp( sum, config.n_layers )); } - hipfire_runtime::multi_gpu::Gpus::init_layers(&counts).map_err(|e| format!("{e}"))? + hipfire_hardware::Gpus::init_layers(&counts).map_err(|e| format!("{e}"))? } - None => hipfire_runtime::multi_gpu::Gpus::init_uniform(pp, config.n_layers) + None => hipfire_hardware::Gpus::init_uniform(pp, config.n_layers) .map_err(|e| format!("{e}"))?, }; // Discrete GPUs: keep model pages in the page cache across reads — diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index c099ef2818..77e80e2b22 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -32,7 +32,7 @@ use hipfire_runtime::kv_mode; use hipfire_runtime::llama; use hipfire_runtime::llama::{KvCacheExt, KvDims, KvLayers, KvTarget}; use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hipfire_runtime::spec::{SpecEmit, SpecEmitCtx, SpecTargetGuard, Speculator}; use hipfire_runtime::triattn::{EvictionCtx, TriAttnCenters}; use rdna_compute::Gpu; diff --git a/crates/hipfire-runtime/Cargo.toml b/crates/hipfire-runtime/Cargo.toml index a684d70f6e..fcc1773d75 100644 --- a/crates/hipfire-runtime/Cargo.toml +++ b/crates/hipfire-runtime/Cargo.toml @@ -47,6 +47,7 @@ ep-fault-inject = [] serve-fault-inject = [] [dependencies] hip-bridge = { path = "../hip-bridge" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-config = { path = "../hipfire-config" } rdna-compute = { path = "../rdna-compute" } saddle-core = { path = "../saddle-core" } diff --git a/crates/hipfire-runtime/examples/bench_tp_graph_host_overhead_gfx1201.rs b/crates/hipfire-runtime/examples/bench_tp_graph_host_overhead_gfx1201.rs index e00ae4e98c..5b31e33077 100644 --- a/crates/hipfire-runtime/examples/bench_tp_graph_host_overhead_gfx1201.rs +++ b/crates/hipfire-runtime/examples/bench_tp_graph_host_overhead_gfx1201.rs @@ -4,7 +4,7 @@ use std::time::Instant; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; fn percentile(mut values: Vec, quantile: f64) -> f64 { values.sort_by(f64::total_cmp); diff --git a/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs b/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs index 6c4f740421..07b4466070 100644 --- a/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs +++ b/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs @@ -11,7 +11,7 @@ //! a model-throughput claim. use hip_bridge::DeviceBuffer; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use std::time::Instant; const RANKS: usize = 4; diff --git a/crates/hipfire-runtime/examples/ep_decode_parity.rs b/crates/hipfire-runtime/examples/ep_decode_parity.rs index ee3a620ec7..a15e1cdb2f 100644 --- a/crates/hipfire-runtime/examples/ep_decode_parity.rs +++ b/crates/hipfire-runtime/examples/ep_decode_parity.rs @@ -50,7 +50,7 @@ fn main() { use hipfire_arch_qwen35::qwen35::{self, DeltaNetState, Qwen35Scratch}; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::{self, KvCache}; - use hipfire_runtime::multi_gpu::Gpus; + use hipfire_hardware::Gpus; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; use std::path::Path; diff --git a/crates/hipfire-runtime/examples/pp_parity_chatml.rs b/crates/hipfire-runtime/examples/pp_parity_chatml.rs index 76716f208a..4402df52fa 100644 --- a/crates/hipfire-runtime/examples/pp_parity_chatml.rs +++ b/crates/hipfire-runtime/examples/pp_parity_chatml.rs @@ -21,7 +21,7 @@ use hipfire_arch_qwen35::qwen35::{ use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use rdna_compute::Gpu; use std::path::Path; diff --git a/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs b/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs index 6cc609280d..ce6b8af528 100644 --- a/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs +++ b/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs @@ -13,7 +13,7 @@ //! replay rather than passing on a stale record from the prior replay. use hip_bridge::{DeviceBuffer, Graph, GraphExec}; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use rdna_compute::{DType, GpuTensor}; const ELEMS: usize = 4_096; diff --git a/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs b/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs index 2dfac4b4ce..1fab187aff 100644 --- a/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs +++ b/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs @@ -15,7 +15,7 @@ // HIP_VISIBLE_DEVICES=0,1 HIPFIRE_TP_BENCH_N=2 cargo run ... (TP=2) use hip_bridge::DeviceBuffer; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use std::time::Instant; const SIZES_BYTES: &[usize] = &[4 * 1024, 32 * 1024, 128 * 1024, 512 * 1024]; @@ -82,7 +82,8 @@ fn main() { } let refs: Vec<&DeviceBuffer> = buffers.iter().collect(); - gpus.all_reduce_sum_f32(&refs, count) + let group: Vec = (0..refs.len()).collect(); + gpus.all_reduce_sum_f32(&group, &refs, count) .expect("all_reduce_sum_f32"); // Sync all rank streams before readback. @@ -132,9 +133,10 @@ fn main() { for &bytes in SIZES_BYTES { let count = bytes / std::mem::size_of::(); let refs: Vec<&DeviceBuffer> = buffers.iter().collect(); + let group: Vec = (0..refs.len()).collect(); for _ in 0..warmup { - gpus.all_reduce_sum_f32(&refs, count) + gpus.all_reduce_sum_f32(&group, &refs, count) .expect("all_reduce warm"); for dev in &gpus.devices { dev.bind_thread().expect("bind"); @@ -147,7 +149,9 @@ fn main() { let mut samples = Vec::with_capacity(iters); for _ in 0..iters { let t = Instant::now(); - gpus.all_reduce_sum_f32(&refs, count).expect("all_reduce"); + gpus + .all_reduce_sum_f32(&group, &refs, count) + .expect("all_reduce"); for dev in &gpus.devices { dev.bind_thread().expect("bind"); dev.hip diff --git a/crates/hipfire-runtime/src/config.rs b/crates/hipfire-runtime/src/config.rs index b0760bc5a4..8239704165 100644 --- a/crates/hipfire-runtime/src/config.rs +++ b/crates/hipfire-runtime/src/config.rs @@ -73,7 +73,7 @@ pub struct RuntimeConfig { pub lm_head_f16: String, /// Tensor-parallel RCCL all-reduce toggle. `None` (unset) → RCCL is used /// (default). `Some(false)` (HIPFIRE_TP_USE_RCCL=0) → opt out of the RCCL - /// path. `Some(true)` → force on. Read by `multi_gpu::Gpus::ensure_rccl`. + /// path. `Some(true)` → force on. Read by `hipfire_hardware::Gpus::ensure_rccl`. pub tp_use_rccl: Option, pub ngram_loop_threshold: usize, pub ngram_window: usize, diff --git a/crates/hipfire-runtime/src/ep.rs b/crates/hipfire-runtime/src/ep.rs index 15f5ae8e36..eeddad2229 100644 --- a/crates/hipfire-runtime/src/ep.rs +++ b/crates/hipfire-runtime/src/ep.rs @@ -31,7 +31,7 @@ //! driver loops layers (advancing each rank's per-layer binding state) the same //! way the single-GPU lowered driver loops `run_layer_program`. -use crate::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hip_bridge::{DeviceBuffer, HipError}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::superop::{ @@ -67,10 +67,13 @@ fn all_reduce_sum_f32_decode( let use_peer = *PEER_DECODE.get_or_init(|| { hipfire_config::developer_var("HIPFIRE_EP_PEER_ALLREDUCE_DECODE").as_deref() == Ok("1") }); + let group: Vec = (0..refs.len()).collect(); if use_peer { - gpus.all_reduce_sum_f32_peer(refs, count).map_err(hip_err) + gpus + .all_reduce_sum_f32_peer(&group, refs, count) + .map_err(hip_err) } else { - gpus.all_reduce_sum_f32(refs, count).map_err(hip_err) + gpus.all_reduce_sum_f32(&group, refs, count).map_err(hip_err) } } diff --git a/crates/hipfire-runtime/src/lib.rs b/crates/hipfire-runtime/src/lib.rs index 1f98ad8d98..ebf81aff3f 100644 --- a/crates/hipfire-runtime/src/lib.rs +++ b/crates/hipfire-runtime/src/lib.rs @@ -47,7 +47,6 @@ pub mod loader_api; pub mod loop_guard; pub mod model_load; pub mod model_source; -pub mod multi_gpu; pub mod paro; pub mod prefix; pub mod reset_core; diff --git a/crates/hipfire-runtime/src/llama.rs b/crates/hipfire-runtime/src/llama.rs index e9e9091075..2e1e5dc167 100644 --- a/crates/hipfire-runtime/src/llama.rs +++ b/crates/hipfire-runtime/src/llama.rs @@ -11,7 +11,7 @@ use crate::kv_backend::{ KvBackend, KvChunkPlan, DEFAULT_KV_CHUNK_TOKENS, DEFAULT_VMM_PHYSICAL_CHUNK_BYTES, }; use crate::kv_mode::KvMode; -use crate::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hip_bridge::HipResult; use rdna_compute::{DType, Gpu, GpuTensor}; @@ -5563,8 +5563,8 @@ pub use saddle_core::kv::{KvCache, KvDims, KvLayers, VMode}; // `KvMode` and `KvBackend` are re-exported from their canonical modules // (`crate::kv_mode`, `crate::kv_backend`) which themselves re-export from // `saddle-core`; `llama.rs` does not need to re-export them directly. -// `KvTarget` stays here because it depends on `crate::multi_gpu::Gpus`, -// which is runtime-specific and must not leak into `saddle-core`. +// `KvTarget` stays here because it depends on `hipfire_hardware::Gpus`, +// which is hardware-specific and must not leak into `saddle-core`. pub enum KvTarget<'a> { /// pp == 1: one GPU. Sites 1–5. diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index 2f1873a01a..e8349e9e75 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -6,7 +6,7 @@ //! per-tensor dequant), which `WeightSource::read_layer` calls internally. use crate::llama::{EmbeddingFormat, WeightTensor}; -use crate::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use hip_bridge::HipResult; use rdna_compute::{Gpu, GpuTensor}; diff --git a/crates/saddle-lab/Cargo.toml b/crates/saddle-lab/Cargo.toml index d18fb866d4..463f8658fa 100644 --- a/crates/saddle-lab/Cargo.toml +++ b/crates/saddle-lab/Cargo.toml @@ -29,6 +29,7 @@ rdna-compute = { path = "../rdna-compute" } saddle-core = { path = "../saddle-core" } hipfire-dispatch = { path = "../hipfire-dispatch" } hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-arch-qwen35 = { path = "../hipfire-arch-qwen35" } hipfire-arch-qwen35-vl = { path = "../hipfire-arch-qwen35-vl" } hipfire-arch-llama = { path = "../hipfire-arch-llama" } diff --git a/crates/saddle-lab/examples/gpus_smoke.rs b/crates/saddle-lab/examples/gpus_smoke.rs index 7d6ed3b657..06da0913ba 100644 --- a/crates/saddle-lab/examples/gpus_smoke.rs +++ b/crates/saddle-lab/examples/gpus_smoke.rs @@ -7,7 +7,7 @@ //! //! Run: HIP_VISIBLE_DEVICES=0,1 cargo run -p hipfire-runtime --example gpus_smoke -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use rdna_compute::{DType, Gpu}; fn main() { diff --git a/crates/saddle-lab/examples/pp2_vram_probe.rs b/crates/saddle-lab/examples/pp2_vram_probe.rs index cdb2e09383..bf35a946f5 100644 --- a/crates/saddle-lab/examples/pp2_vram_probe.rs +++ b/crates/saddle-lab/examples/pp2_vram_probe.rs @@ -16,7 +16,7 @@ use hipfire_arch_qwen35::qwen35::{self, DeltaNetState, Qwen35ScratchSet, StateQu use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use std::path::Path; fn used_gb(gpus: &Gpus, baseline_free: &[(usize, usize)]) -> Vec { diff --git a/crates/saddle-lab/examples/pp_parity.rs b/crates/saddle-lab/examples/pp_parity.rs index 3ae3336ca9..f8ce79bd95 100644 --- a/crates/saddle-lab/examples/pp_parity.rs +++ b/crates/saddle-lab/examples/pp_parity.rs @@ -18,7 +18,7 @@ use hipfire_arch_qwen35::qwen35::{ use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use rdna_compute::Gpu; use std::path::Path; From d2cc8737c2ad851447440a1200419a98207a36b9 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 18:47:31 +0200 Subject: [PATCH 02/25] fix(hardware): restore resolved device topology authority --- .../examples/ep_deepseek4.rs | 3 +- .../examples/ep_dspark_topology_probe.rs | 3 +- crates/hipfire-arch-deepseek4/src/ep.rs | 4 + crates/hipfire-arch-deepseek4/src/mtp.rs | 3 +- .../examples/ep_minimax.rs | 3 +- crates/hipfire-arch-minimax/src/forward.rs | 3 +- .../examples/qwen_dense_tp2_parity.rs | 5 +- .../examples/test_qwen35_load_multi.rs | 4 +- .../examples/test_qwen35_state_multi.rs | 4 +- .../src/qwen35/ep_batch.rs | 4 +- crates/hipfire-arch-qwen35/tests/pp_parity.rs | 4 +- crates/hipfire-hardware/src/lib.rs | 182 ++++++++---------- crates/hipfire-hardware/src/mesh.rs | 120 +++++++++--- crates/hipfire-hardware/tests/ownership.rs | 2 +- crates/hipfire-loader/src/carriers.rs | 6 +- crates/hipfire-loader/src/lib.rs | 18 +- .../bench_tp_graph_host_overhead_gfx1201.rs | 3 +- .../ds4_gfx1201_owner_worker_transport.rs | 4 +- .../examples/ep_decode_parity.rs | 3 +- .../examples/pp_parity_chatml.rs | 4 +- .../tp4_cross_device_graph_barrier.rs | 4 +- .../examples/tp_allreduce_smoke.rs | 4 +- crates/hipfire-runtime/src/config.rs | 23 +++ crates/hipfire-runtime/src/ep.rs | 18 +- crates/saddle-lab/examples/gpus_smoke.rs | 6 +- crates/saddle-lab/examples/pp2_vram_probe.rs | 4 +- crates/saddle-lab/examples/pp_parity.rs | 4 +- 27 files changed, 274 insertions(+), 171 deletions(-) diff --git a/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs b/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs index 8e57147221..582276bcf2 100644 --- a/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs +++ b/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs @@ -110,7 +110,8 @@ fn main() { drop(hfq0); // ── bring up N ranks ──────────────────────────────────────────────────── - let mut gpus = Gpus::init_tp(tp, cfg.num_hidden_layers).expect("init_tp"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = Gpus::init_tp(&device_opts, tp, cfg.num_hidden_layers).expect("init_tp"); let n = gpus.devices.len(); assert_eq!( n, tp, diff --git a/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs b/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs index 56a66c1f68..faee247042 100644 --- a/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs +++ b/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs @@ -182,7 +182,8 @@ fn main() -> Result<(), String> { "topology: target=TP3 devices 0,1,2 drafter=device 3 hidden={} layers={} verify_B={} position={}", cfg.hidden_size, cfg.num_hidden_layers, args.verify_batch, args.position ); - let mut gpus = Gpus::init_tp(TARGET_RANKS, cfg.num_hidden_layers) + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = Gpus::init_tp(&device_opts, TARGET_RANKS, cfg.num_hidden_layers) .map_err(|e| format!("initialize TP3 target: {e:?}"))?; if gpus.devices.len() != TARGET_RANKS || gpus.devices.iter().any(|gpu| !gpu.arch_caps.is_gfx1201()) diff --git a/crates/hipfire-arch-deepseek4/src/ep.rs b/crates/hipfire-arch-deepseek4/src/ep.rs index 6f2a97a882..dc32a58b30 100644 --- a/crates/hipfire-arch-deepseek4/src/ep.rs +++ b/crates/hipfire-arch-deepseek4/src/ep.rs @@ -125,6 +125,7 @@ fn forward_ep_tp_graph_body( position: u32, ) -> Result<(), String> { let n = gpus.devices.len(); + let group: Vec = (0..n).collect(); let program = ds4_lower_program(); let skip_ffn = config_cache::skip_ffn(); for layer_idx in 0..cfg.num_hidden_layers { @@ -144,6 +145,7 @@ fn forward_ep_tp_graph_body( gpus, bindings.as_mut_slice(), partials, + &group, &program, cfg.hidden_size, ) @@ -381,6 +383,7 @@ fn forward_ep_direct( .unwrap_or(false); let t_layers = std::time::Instant::now(); let program = ds4_lower_program(); + let group: Vec = (0..n).collect(); for l in 0..cfg.num_hidden_layers { { let mut binds: Vec = Vec::with_capacity(n); @@ -399,6 +402,7 @@ fn forward_ep_direct( gpus, binds.as_mut_slice(), partials, + &group, &program, hidden, ) diff --git a/crates/hipfire-arch-deepseek4/src/mtp.rs b/crates/hipfire-arch-deepseek4/src/mtp.rs index fc3cb3fb3a..272cad46d8 100644 --- a/crates/hipfire-arch-deepseek4/src/mtp.rs +++ b/crates/hipfire-arch-deepseek4/src/mtp.rs @@ -550,7 +550,7 @@ pub fn mtp_forward_ep( assert_eq!(h_n_per_rank.len(), n, "mtp_forward_ep: h_n_per_rank len"); let hidden = cfg.hidden_size; let mtp_layer_idx = cfg.num_hidden_layers; - + let group: Vec = (0..n).collect(); // 1. Per-rank pre-FFN (embed/norm/HC + attention), replicated. attn_stub // reads state.n_tokens for the MTP-layer SWA ring slot → set it to // `position` per rank (matches spec_decode's bookkeeping). @@ -592,6 +592,7 @@ pub fn mtp_forward_ep( gpus, binds.as_mut_slice(), partials, + &group, &program, hidden, ) diff --git a/crates/hipfire-arch-minimax/examples/ep_minimax.rs b/crates/hipfire-arch-minimax/examples/ep_minimax.rs index 24c88d6745..ddbfbe3d43 100644 --- a/crates/hipfire-arch-minimax/examples/ep_minimax.rs +++ b/crates/hipfire-arch-minimax/examples/ep_minimax.rs @@ -71,7 +71,8 @@ fn main() { drop(hfq0); // ── bring up N ranks ──────────────────────────────────────────────────── - let mut gpus = Gpus::init_tp(tp, cfg.num_hidden_layers).expect("init_tp"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = Gpus::init_tp(&device_opts, tp, cfg.num_hidden_layers).expect("init_tp"); let n = gpus.devices.len(); assert_eq!(n, tp, "init_tp gave {n} devices (check HIP_VISIBLE_DEVICES)"); for (r, d) in gpus.devices.iter().enumerate() { diff --git a/crates/hipfire-arch-minimax/src/forward.rs b/crates/hipfire-arch-minimax/src/forward.rs index aa42b84f97..a900763272 100644 --- a/crates/hipfire-arch-minimax/src/forward.rs +++ b/crates/hipfire-arch-minimax/src/forward.rs @@ -1653,7 +1653,7 @@ pub fn forward_ep( assert_eq!(partials.len(), n, "forward_ep: partials len"); let hidden = cfg.hidden_size; let eps = cfg.rms_norm_eps; - + let group: Vec = (0..n).collect(); // 1. Embed + stage pos per rank (replicated, deterministic). for r in 0..n { gpus.devices[r] @@ -1695,6 +1695,7 @@ pub fn forward_ep( gpus, binds.as_mut_slice(), partials, + &group, &program, hidden, ) diff --git a/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs b/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs index 93dea57513..ce652feaed 100644 --- a/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs +++ b/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs @@ -147,8 +147,9 @@ fn run_tp(path: &str, seed: &[u32], forced: &[u32]) -> (Vec, Vec>) .iter() .map(|l| qwen35::local_dense_tp_config(&global, l)) .collect(); - let mut gpus = - Gpus::init_tp(tp, global.n_layers).unwrap_or_else(|e| panic!("init tp{tp}: {e:?}")); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = Gpus::init_tp(&device_opts, tp, global.n_layers) + .unwrap_or_else(|e| panic!("init tp{tp}: {e:?}")); for gpu in &mut gpus.devices { gpu.bind_thread().unwrap(); gpu.active_stream = Some(gpu.hip.stream_create().unwrap()); diff --git a/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs b/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs index a5343e0fe4..41c75378e8 100644 --- a/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs +++ b/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs @@ -27,7 +27,9 @@ fn main() { config.n_layers, config.vocab_size, config.dim, config.hidden_dim, ); - let mut gpus = Gpus::init_uniform(2, config.n_layers).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let n = gpus.devices.len(); let out_dev = gpus.output_device; println!( diff --git a/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs b/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs index 1842e45478..76835b0d7e 100644 --- a/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs +++ b/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs @@ -31,7 +31,9 @@ fn main() { config.n_layers, config.head_dim, config.n_kv_heads, config.vocab_size, ); - let mut gpus = Gpus::init_uniform(2, config.n_layers).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let n_dev = gpus.devices.len(); let out_dev = gpus.output_device; println!( diff --git a/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs b/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs index 561d84d74f..16e61a7462 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs @@ -2170,6 +2170,7 @@ pub fn forward_ep( let n_v_heads = config.linear_num_value_heads; let hd = config.linear_key_head_dim; let pos_i32 = pos as i32; + let group: Vec = (0..n).collect(); // 1. Embed token + write pos on each rank (replicated; deterministic, since // weights are byte-identical replicas → s.x is bit-identical per rank). @@ -2236,6 +2237,7 @@ pub fn forward_ep( gpus, binds.as_mut_slice(), partials, + &group, &program, dim, ) @@ -2354,6 +2356,7 @@ pub fn forward_prefill_batch_ep( return Ok(()); } let dim = config.dim; + let group: Vec = (0..n_rank).collect(); // Per-call contract: one window must fit max_batch. Long prompts are driven // by calling this repeatedly with advancing start_pos + persistent kv/dn // (KV + DeltaNet state accumulate in place across calls, identical to the @@ -2458,7 +2461,6 @@ pub fn forward_prefill_batch_ep( let t_a = std::time::Instant::now(); let refs: Vec<&hip_bridge::DeviceBuffer> = partials.iter().map(|partial| &partial.buf).collect(); - let group: Vec = (0..refs.len()).collect(); if ep_peer_ar { gpus .all_reduce_sum_f32_peer(&group, &refs, n * dim) diff --git a/crates/hipfire-arch-qwen35/tests/pp_parity.rs b/crates/hipfire-arch-qwen35/tests/pp_parity.rs index 46b4d59a04..a9ffbbb882 100644 --- a/crates/hipfire-arch-qwen35/tests/pp_parity.rs +++ b/crates/hipfire-arch-qwen35/tests/pp_parity.rs @@ -117,7 +117,9 @@ fn run_pp1(path: &str, prompt: &[u32]) -> Vec { fn run_pp2(path: &str, prompt: &[u32]) -> Vec { let mut hfq = HfqFile::open(Path::new(path)).expect("open hfq"); let config = qwen35::config_from_hfq(&hfq).expect("config"); - let mut gpus = Gpus::init_uniform(2, config.n_layers).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let layout = qwen35::Layout::from_gpus(&gpus, config.n_layers); let mut hfq_source = qwen35::HfqSource::new(&mut hfq, &config); let weights = qwen35::load_weights(&mut hfq_source, &mut gpus.devices, &layout) diff --git a/crates/hipfire-hardware/src/lib.rs b/crates/hipfire-hardware/src/lib.rs index 8363fd7ceb..197196ea57 100644 --- a/crates/hipfire-hardware/src/lib.rs +++ b/crates/hipfire-hardware/src/lib.rs @@ -21,55 +21,29 @@ //! 3. Pass the multi-GPU coherence gate. use hip_bridge::{ - DeviceBuffer, Event, HipError, HipResult, HipRuntime, RcclComms, + DeviceBuffer, Event, HipError, HipResult, RcclComms, HIP_ERROR_PEER_ACCESS_ALREADY_ENABLED, HIP_ERROR_PEER_ACCESS_UNSUPPORTED, HIP_EVENT_DISABLE_TIMING, HIP_EVENT_RELEASE_TO_SYSTEM, }; use rdna_compute::{DType, Gpu, GpuTensor}; mod mesh; -pub use mesh::{Axis, CollectiveHint, DeviceMesh, DimKind, MeshEpoch}; +pub use mesh::{Axis, CollectiveHint, DeviceMesh, DimKind, MeshEpoch, MeshError}; -/// Device-resolution knobs used when constructing a [`Gpus`] owner. +/// Device-resolution knobs supplied by the resolved process/runtime config +/// when constructing a [`Gpus`] owner. /// -/// The hardware leaf reads the same legacy `HIPFIRE_*` environment variables -/// as the runtime did, without depending on runtime configuration. Higher -/// layers may eventually supply an explicit option set; the constructors -/// below intentionally keep the existing process-environment behavior. -#[derive(Clone, Debug, Default)] +/// `devices` contains the post-visibility logical HIP IDs produced by +/// `hipfire-config`, not the original physical ROCr selectors. Hardware +/// consumes this value as-is and never rereads process environment. +#[derive(Clone, Debug, Default, PartialEq)] pub struct DeviceResolveOpts { pub tp_use_rccl: Option, pub devices: Option, - pub emulate_gpus: Option, pub allow_mixed_arch: bool, pub uniform_vram_tolerance_gb: Option, } -impl DeviceResolveOpts { - pub fn from_env() -> Self { - Self { - tp_use_rccl: std::env::var("HIPFIRE_TP_USE_RCCL") - .ok() - .as_deref() - .map(|value| value != "0" && !value.eq_ignore_ascii_case("false")), - devices: std::env::var("HIPFIRE_DEVICES") - .ok() - .filter(|value| !value.is_empty()), - emulate_gpus: std::env::var("HIPFIRE_EMULATE_GPUS") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|&count| count >= 2), - allow_mixed_arch: std::env::var("HIPFIRE_ALLOW_MIXED_ARCH") - .ok() - .as_deref() - == Some("1"), - uniform_vram_tolerance_gb: std::env::var("HIPFIRE_UNIFORM_VRAM_TOLERANCE_GB") - .ok() - .and_then(|value| value.parse().ok()), - } - } -} - /// Stream-event handoff returned by `Gpus::boundary_copy`. When the src /// device has an active stream, `completion` holds a HIP event recorded /// after the async peer copy; `Gpus::wait_boundary` makes the dst stream @@ -147,20 +121,16 @@ pub fn peer_reduce_scratch_total_bytes(rank_count: usize, requested_bytes: usize } /// Internal active lease record stored inside `Gpus` while a lease is live. -#[derive(Debug)] -struct ActivePeerLease { - id: u64, - bytes: usize, - rank_count: usize, -} - pub struct Gpus { /// RCCL communicators (one per rank), lazily initialized on the first /// `all_reduce_sum_*` call. Declared BEFORE `devices` so `Drop` tears /// down comms (via `ncclCommDestroy`) before the underlying HIP - /// devices, which RCCL relies on. `None` means RCCL hasn't been used - /// or `HIPFIRE_TP_USE_RCCL=0` forced the opt-out. + /// devices, which RCCL relies on. `None` means RCCL hasn't been used. rccl_comms: Option, + /// Resolved at construction from the process/runtime config. Keeping this + /// decision on the owner prevents a later environment mutation from + /// changing the collective route. + use_rccl: bool, pub devices: Vec, /// Per-layer device id, length = n_layers. pub layer_to_device: Vec, @@ -214,7 +184,11 @@ impl Gpus { /// uniformly: max-min ≤ 1 layer per band. Pre-flight VRAM check enforces /// arch match and bounded VRAM delta (override /// `HIPFIRE_UNIFORM_VRAM_TOLERANCE_GB`, default 2 GiB). - pub fn init_uniform(n_devices: usize, n_layers: usize) -> HipResult { + pub fn init_uniform( + opts: &DeviceResolveOpts, + n_devices: usize, + n_layers: usize, + ) -> HipResult { if n_devices == 0 { return Err(HipError::new(0, "init_uniform: n_devices must be >= 1")); } @@ -227,18 +201,15 @@ impl Gpus { ), )); } - let device_ids = resolve_device_ids(n_devices)?; + let device_ids = resolve_device_ids(n_devices, opts)?; let devices = construct_devices(&device_ids)?; - preflight_vram_with_opts(&devices, /*check_vram_delta=*/ true)?; + preflight_vram_with_opts(&devices, /*check_vram_delta=*/ true, opts)?; let per_device = uniform_split_counts(n_devices, n_layers); - Self::from_parts(devices, per_device, n_layers) + Self::from_parts(devices, per_device, n_layers, opts) } - - /// Explicit escape hatch for asymmetric VRAM / hand-tuned splits. - /// Keeps arch-mismatch and per-device bind/free pre-flight checks, but /// skips the uniform VRAM-delta gate. `per_device` length determines /// `n_devices`; sum determines `n_layers`. - pub fn init_layers(per_device: &[usize]) -> HipResult { + pub fn init_layers(opts: &DeviceResolveOpts, per_device: &[usize]) -> HipResult { let n_devices = per_device.len(); if n_devices == 0 { return Err(HipError::new( @@ -253,20 +224,22 @@ impl Gpus { )); } let n_layers: usize = per_device.iter().sum(); - let device_ids = resolve_device_ids(n_devices)?; + let device_ids = resolve_device_ids(n_devices, opts)?; let devices = construct_devices(&device_ids)?; // init_layers is the documented escape hatch for asymmetric VRAM // splits — the caller has declared the per-device counts, so skip // the VRAM-delta check (which would otherwise reject 32 GB MI50 + // 12 GB 6700 XT pairs out of the box). Arch-mismatch + per-device // bind+free probe still run. - preflight_vram_with_opts(&devices, /*check_vram_delta=*/ false)?; - Self::from_parts(devices, per_device.to_vec(), n_layers) + preflight_vram_with_opts(&devices, /*check_vram_delta=*/ false, opts)?; + Self::from_parts(devices, per_device.to_vec(), n_layers, opts) } - /// Reserved for v1.1 — automatic VRAM-weighted band assignment. For v1 - /// use `init_layers(...)` with hand-computed counts. - pub fn init_vram_weighted(_n_devices: usize, _n_layers: usize) -> HipResult { + pub fn init_vram_weighted( + _opts: &DeviceResolveOpts, + _n_devices: usize, + _n_layers: usize, + ) -> HipResult { Err(HipError::new( 0, "init_vram_weighted: scheduled for v1.1; use init_layers(per_device) instead", @@ -275,9 +248,10 @@ impl Gpus { /// PP=1 back-compat path: wrap an existing single `Gpu` into a `Gpus` /// with all layers on dev 0. `output_device = 0`. - pub fn single(gpu: Gpu, n_layers: usize) -> Self { + pub fn single(opts: &DeviceResolveOpts, gpu: Gpu, n_layers: usize) -> Self { Self { rccl_comms: None, + use_rccl: opts.tp_use_rccl.unwrap_or(true), devices: vec![gpu], layer_to_device: vec![0; n_layers], band_starts: vec![0], @@ -316,16 +290,20 @@ impl Gpus { /// only validates the device count. Pre-flight runs the arch-match + /// VRAM-delta gate (TP ranks are identical cards, so the uniform delta /// check applies). - pub fn init_tp(tp_size: usize, n_layers: usize) -> HipResult { + pub fn init_tp( + opts: &DeviceResolveOpts, + tp_size: usize, + n_layers: usize, + ) -> HipResult { if tp_size == 0 { return Err(HipError::new(0, "init_tp: tp_size must be >= 1")); } if n_layers == 0 { return Err(HipError::new(0, "init_tp: n_layers must be >= 1")); } - let device_ids = resolve_device_ids(tp_size)?; + let device_ids = resolve_device_ids(tp_size, opts)?; let devices = construct_devices(&device_ids)?; - preflight_vram_with_opts(&devices, /*check_vram_delta=*/ true)?; + preflight_vram_with_opts(&devices, /*check_vram_delta=*/ true, opts)?; let band_starts = tp_band_starts(tp_size, n_layers); // PP=1 TP topology: every device runs every layer. Encode the layer @@ -333,6 +311,7 @@ impl Gpus { // owning empty bands. Ok(Self { rccl_comms: None, + use_rccl: opts.tp_use_rccl.unwrap_or(true), devices, layer_to_device: vec![0u8; n_layers], band_starts, @@ -863,7 +842,12 @@ impl Gpus { Ok(()) } - fn from_parts(devices: Vec, per_device: Vec, n_layers: usize) -> HipResult { + fn from_parts( + devices: Vec, + per_device: Vec, + n_layers: usize, + opts: &DeviceResolveOpts, + ) -> HipResult { debug_assert_eq!(per_device.iter().sum::(), n_layers); debug_assert_eq!(per_device.len(), devices.len()); let n_devices = devices.len(); @@ -879,6 +863,7 @@ impl Gpus { } Ok(Self { rccl_comms: None, + use_rccl: opts.tp_use_rccl.unwrap_or(true), devices, layer_to_device, band_starts, @@ -900,22 +885,16 @@ impl Gpus { }) } - // ────────────────────────────────────────────────────────────────── - // Tensor-parallel collectives (RCCL-backed). See - // docs/plans/multi-gpu-tp-a3b.md §3.3 and the comm baseline at - // docs/investigations/2026-05-28-tp-comm-baseline-hiptrx.md. - // ────────────────────────────────────────────────────────────────── - /// Lazily initialize RCCL communicators across all devices owned by /// this `Gpus`. Cached for process lifetime; subsequent calls are - /// no-ops. `HIPFIRE_TP_USE_RCCL=0` short-circuits with a clear + /// no-ops. A resolved `tp_use_rccl=false` short-circuits with a clear /// error so callers can fall through to a host-driven path (not /// yet implemented — Stage 2 follow-up). pub fn ensure_rccl(&mut self) -> HipResult<()> { if self.rccl_comms.is_some() { return Ok(()); } - if matches!(DeviceResolveOpts::from_env().tp_use_rccl, Some(false)) { + if !self.use_rccl { return Err(HipError::new( 0, "ensure_rccl: HIPFIRE_TP_USE_RCCL=0 — RCCL path opted out. \ @@ -1902,53 +1881,36 @@ fn uniform_split_counts(n_devices: usize, n_layers: usize) -> Vec { .collect() } -/// Resolve logical device IDs after the physical `hardware.devices` list has -/// been installed as `ROCR_VISIBLE_DEVICES` and HIP has received the matching -/// post-filter logical IDs. When unset, take the first `n_devices` visible IDs. -/// Map each requested logical device id into the physical range -/// `[0, real_count)` by euclidean remainder. Used by -/// `HIPFIRE_EMULATE_GPUS` to alias N logical devices onto fewer physical -/// devices. A non-positive `real_count` is left untouched so the caller can -/// surface the underlying HIP error. -fn alias_ids(ids: &[i32], real_count: i32) -> Vec { - if real_count <= 0 { - return ids.to_vec(); - } - ids.iter().map(|&id| id.rem_euclid(real_count)).collect() -} - -/// Resolve logical IDs from `HIPFIRE_DEVICES`, or use the first N visible IDs. -/// `HIPFIRE_EMULATE_GPUS` optionally aliases those IDs into the physical -/// runtime device range for debug-only multi-rank emulation. -fn resolve_device_ids(n_devices: usize) -> HipResult> { - let opts = DeviceResolveOpts::from_env(); - let ids: Vec = if let Some(value) = &opts.devices { +/// Resolve logical device IDs from the already-resolved `hardware.devices` +/// value, or use the first `n_devices` visible IDs. The supplied list is +/// post-visibility logical HIP IDs and is consumed without physical-ID +/// remapping. +fn resolve_device_ids(n_devices: usize, opts: &DeviceResolveOpts) -> HipResult> { + if let Some(value) = &opts.devices { let parsed = value .split(',') .map(str::trim) .filter(|part| !part.is_empty()) .map(str::parse::) .collect::, _>>() - .map_err(|error| HipError::new(0, &format!("HIPFIRE_DEVICES parse: {error}")))?; + .map_err(|error| HipError::new(0, &format!("hardware.devices parse: {error}")))?; if parsed.len() < n_devices { return Err(HipError::new( 0, &format!( - "HIPFIRE_DEVICES has {} ids but n_devices = {n_devices}", + "hardware.devices exposes {} ids but n_devices = {n_devices}", parsed.len() ), )); } - parsed[..n_devices].to_vec() - } else { - (0..n_devices as i32).collect() - }; - - if opts.emulate_gpus.is_some() { - let real_count = HipRuntime::load()?.device_count()?; - return Ok(alias_ids(&ids, real_count)); + return Ok(parsed[..n_devices].to_vec()); } - Ok(ids) + (0..n_devices) + .map(|id| { + i32::try_from(id) + .map_err(|_| HipError::new(0, "hardware.devices logical ID exceeds HIP range")) + }) + .collect() } fn construct_devices(ids: &[i32]) -> HipResult> { @@ -1959,12 +1921,15 @@ fn construct_devices(ids: &[i32]) -> HipResult> { Ok(devices) } -fn preflight_vram_with_opts(devices: &[Gpu], check_vram_delta: bool) -> HipResult<()> { +fn preflight_vram_with_opts( + devices: &[Gpu], + check_vram_delta: bool, + opts: &DeviceResolveOpts, +) -> HipResult<()> { if devices.is_empty() { return Ok(()); } let arch0 = devices[0].arch.clone(); - let opts = DeviceResolveOpts::from_env(); let mut frees = Vec::with_capacity(devices.len()); for device in devices { if device.arch != arch0 { @@ -2165,4 +2130,13 @@ mod tests { Gpus::host_reduce_rows(&mut rows3, 0); assert_eq!(rows3[0], vec![5.0, 6.0]); } + #[test] + fn resolver_consumes_post_visibility_logical_ids_without_aliasing() { + let opts = DeviceResolveOpts { + devices: Some("0,1".into()), + ..DeviceResolveOpts::default() + }; + assert_eq!(resolve_device_ids(2, &opts).unwrap(), vec![0, 1]); + } } + diff --git a/crates/hipfire-hardware/src/mesh.rs b/crates/hipfire-hardware/src/mesh.rs index 76680f29d7..d0fc7f83f1 100644 --- a/crates/hipfire-hardware/src/mesh.rs +++ b/crates/hipfire-hardware/src/mesh.rs @@ -61,6 +61,25 @@ pub enum CollectiveHint { BandXfer { src: usize, dst: usize }, } +/// A rectangular topology cannot represent a cardinality larger than the +/// platform's `usize` range. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)] +pub enum MeshError { + CardinalityOverflow, +} + +impl std::fmt::Display for MeshError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::CardinalityOverflow => { + f.write_str("rectangular device mesh cardinality overflow") + } + } + } +} + +impl std::error::Error for MeshError {} + /// A rectangular named-axis mesh. /// /// The empty axis list is the single-device topology. Size-one axes remain in @@ -71,6 +90,7 @@ pub enum CollectiveHint { #[derive(Clone, Debug)] pub struct DeviceMesh { axes: Vec, + n_devices: usize, epoch: MeshEpoch, } @@ -90,21 +110,29 @@ impl DeviceMesh { /// Axis sizes are normalized to at least one so every mesh has a valid /// coordinate space and at least one logical device. Named axes are kept /// in caller order; the final axis varies fastest in the flattened ID. - pub fn rect(axes: &[(DimKind, usize)]) -> Self { - Self { - axes: axes - .iter() - .map(|&(kind, size)| Axis { - kind, - size: size.max(1), - }) - .collect(), - epoch: fresh_epoch(), + /// + /// The cardinality is checked while constructing the mesh. Callers must + /// handle [`MeshError::CardinalityOverflow`] instead of observing a + /// wrapped device count. + pub fn rect(axes: &[(DimKind, usize)]) -> Result { + let mut normalized = Vec::with_capacity(axes.len()); + let mut n_devices = 1usize; + for &(kind, size) in axes { + let size = size.max(1); + n_devices = n_devices + .checked_mul(size) + .ok_or(MeshError::CardinalityOverflow)?; + normalized.push(Axis { kind, size }); } + Ok(Self { + axes: normalized, + n_devices, + epoch: fresh_epoch(), + }) } /// The single-device topology: one logical device and no named axes. - pub fn single() -> Self { + pub fn single() -> Result { Self::rect(&[]) } @@ -120,11 +148,7 @@ impl DeviceMesh { /// Total number of logical devices (one for an empty mesh). pub fn n_devices(&self) -> usize { - self.axes - .iter() - .map(|axis| axis.size) - .product::() - .max(1) + self.n_devices } /// Size of the first axis with `kind`, or one when it is absent. @@ -243,6 +267,7 @@ impl DeviceMesh { .copied() .filter(|axis| axis.size > 1) .collect(), + n_devices: self.n_devices, epoch: self.epoch, } } @@ -259,7 +284,7 @@ impl DeviceMesh { impl Default for DeviceMesh { fn default() -> Self { - Self::single() + Self::single().expect("single-device mesh cannot overflow") } } @@ -287,7 +312,7 @@ mod tests { #[test] fn single_is_one_device_and_identity_collectives_are_noops() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().unwrap(); assert_eq!(mesh.n_devices(), 1); assert_eq!(mesh.axes(), &[]); assert_eq!(mesh.coord_of(0), Vec::::new()); @@ -301,8 +326,8 @@ mod tests { #[test] fn single_and_empty_rect_have_same_shape_but_fresh_identity() { - let single = DeviceMesh::single(); - let empty_rect = DeviceMesh::rect(&[]); + let single = DeviceMesh::single().unwrap(); + let empty_rect = DeviceMesh::rect(&[]).unwrap(); assert_ne!(single, empty_rect); assert_eq!(single.axes(), empty_rect.axes()); assert_eq!(single.n_devices(), empty_rect.n_devices()); @@ -313,11 +338,21 @@ mod tests { #[test] fn pp_bands_and_boundary_hints_are_uniform() { - let mesh = DeviceMesh::rect(&[(DimKind::Pp, 3)]); - assert_eq!((0..6).map(|l| mesh.stage_for_layer(l, 6)).collect::>(), - vec![0, 0, 1, 1, 2, 2]); - assert_eq!(mesh.band_xfer_after(1, 6), Some(CollectiveHint::BandXfer { src: 0, dst: 1 })); - assert_eq!(mesh.band_xfer_after(3, 6), Some(CollectiveHint::BandXfer { src: 1, dst: 2 })); + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 3)]).unwrap(); + assert_eq!( + (0..6) + .map(|l| mesh.stage_for_layer(l, 6)) + .collect::>(), + vec![0, 0, 1, 1, 2, 2] + ); + assert_eq!( + mesh.band_xfer_after(1, 6), + Some(CollectiveHint::BandXfer { src: 0, dst: 1 }) + ); + assert_eq!( + mesh.band_xfer_after(3, 6), + Some(CollectiveHint::BandXfer { src: 1, dst: 2 }) + ); assert_eq!(mesh.band_xfer_after(5, 6), None); assert_eq!(mesh.stage_devices(&[1]), vec![1]); } @@ -328,20 +363,35 @@ mod tests { (DimKind::Pp, 2), (DimKind::Tp, 2), (DimKind::Ep, 2), - ]); + ]) + .unwrap(); assert_eq!(mesh.n_devices(), 8); assert_eq!(mesh.coord_of(0), vec![0, 0, 0]); assert_eq!(mesh.coord_of(7), vec![1, 1, 1]); - assert_eq!(mesh.device_of(&[1, 0, 1]), 6); - assert_eq!(mesh.group_along(DimKind::Tp, &[1, 0, 1]), vec![6, 7]); - assert_eq!(mesh.group_along(DimKind::Ep, &[1, 0, 1]), vec![4, 6]); + + // The final (Ep) axis varies fastest in the row-major flattening. + let coord = [1, 0, 1]; + assert_eq!(mesh.device_of(&coord), 5); + assert_eq!(mesh.coord_of(5), coord); + assert_eq!(mesh.group_along(DimKind::Tp, &coord), vec![5, 7]); + assert_eq!(mesh.group_along(DimKind::Ep, &coord), vec![4, 5]); assert_eq!(mesh.stage_devices(&[1, 0, 0]), vec![4, 5, 6, 7]); + for pp in 0..2 { + for tp in 0..2 { + for ep in 0..2 { + let coord = [pp, tp, ep]; + assert_eq!(mesh.coord_of(mesh.device_of(&coord)), coord); + } + } + } + let degenerate = DeviceMesh::rect(&[ (DimKind::Pp, 2), (DimKind::Tp, 1), (DimKind::Ep, 2), - ]); + ]) + .unwrap(); assert_eq!( degenerate.squeezed().axes(), &[ @@ -360,9 +410,17 @@ mod tests { #[test] fn coordinate_round_trip_holds_for_every_device() { - let mesh = DeviceMesh::rect(&[(DimKind::Pp, 3), (DimKind::Tp, 2), (DimKind::Ep, 2)]); + let mesh = + DeviceMesh::rect(&[(DimKind::Pp, 3), (DimKind::Tp, 2), (DimKind::Ep, 2)]).unwrap(); for device in 0..mesh.n_devices() { assert_eq!(mesh.device_of(&mesh.coord_of(device)), device); } } + + #[test] + fn rectangular_cardinality_overflow_is_rejected() { + let error = DeviceMesh::rect(&[(DimKind::Pp, usize::MAX), (DimKind::Tp, 2)]) + .expect_err("rectangular cardinality must fail closed"); + assert_eq!(error, MeshError::CardinalityOverflow); + } } diff --git a/crates/hipfire-hardware/tests/ownership.rs b/crates/hipfire-hardware/tests/ownership.rs index 2d0f42ff52..5b41b5439d 100644 --- a/crates/hipfire-hardware/tests/ownership.rs +++ b/crates/hipfire-hardware/tests/ownership.rs @@ -8,7 +8,7 @@ use std::path::Path; #[test] fn hardware_leaf_exposes_owner_and_named_topology() { let _owner = std::any::TypeId::of::(); - let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]).unwrap(); assert_eq!(mesh.n_devices(), 4); assert_eq!(mesh.group_along(DimKind::Tp, &[1, 0]), vec![2, 3]); assert_eq!( diff --git a/crates/hipfire-loader/src/carriers.rs b/crates/hipfire-loader/src/carriers.rs index 6e57ad8a05..4179197650 100644 --- a/crates/hipfire-loader/src/carriers.rs +++ b/crates/hipfire-loader/src/carriers.rs @@ -245,6 +245,7 @@ fn load_qwen35_pp( let pp = ctx.pp; let config = hipfire_arch_qwen35::qwen35::config_from_hfq(&hfq_file) .map_err(|e| format!("failed to read Qwen3.5 config: {e}"))?; + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); let mut gpus = match hipfire_config::developer_var("HIPFIRE_PP_LAYERS") .ok() .filter(|s| !s.is_empty()) @@ -267,9 +268,10 @@ fn load_qwen35_pp( sum, config.n_layers )); } - hipfire_hardware::Gpus::init_layers(&counts).map_err(|e| format!("{e}"))? + hipfire_hardware::Gpus::init_layers(&device_opts, &counts) + .map_err(|e| format!("{e}"))? } - None => hipfire_hardware::Gpus::init_uniform(pp, config.n_layers) + None => hipfire_hardware::Gpus::init_uniform(&device_opts, pp, config.n_layers) .map_err(|e| format!("{e}"))?, }; // Discrete GPUs: keep model pages in the page cache across reads — diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index 77e80e2b22..cadae3c795 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -2990,8 +2990,9 @@ fn load_model_ep_ds4( let chat_template = resolve_chat_template(&hfq, path); let rec = hfq.recommended_sampling(); - let gpus = - Gpus::init_tp(tp, config.num_hidden_layers).map_err(|e| format!("init_tp: {e:?}"))?; + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let gpus = Gpus::init_tp(&device_opts, tp, config.num_hidden_layers) + .map_err(|e| format!("init_tp: {e:?}"))?; let n = gpus.devices.len(); if n != tp { return Err(format!( @@ -3218,8 +3219,9 @@ fn load_model_ep_minimax(path: &str, max_seq: usize, tp: usize) -> Result 0, "samples must be nonzero"); - let mut gpus = Gpus::init_uniform(RANKS, RANKS).expect("init four GPUs"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, RANKS, RANKS).expect("init four GPUs"); assert_eq!(gpus.devices.len(), RANKS, "requires exactly four GPUs"); for (rank, gpu) in gpus.devices.iter().enumerate() { assert_eq!( diff --git a/crates/hipfire-runtime/examples/ep_decode_parity.rs b/crates/hipfire-runtime/examples/ep_decode_parity.rs index a15e1cdb2f..a00e488ac3 100644 --- a/crates/hipfire-runtime/examples/ep_decode_parity.rs +++ b/crates/hipfire-runtime/examples/ep_decode_parity.rs @@ -141,7 +141,8 @@ fn main() { ); // ── bring up N ranks ──────────────────────────────────────────────────── - let mut gpus = Gpus::init_tp(tp, config.n_layers).expect("init_tp"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = Gpus::init_tp(&device_opts, tp, config.n_layers).expect("init_tp"); let n = gpus.devices.len(); assert_eq!( n, tp, diff --git a/crates/hipfire-runtime/examples/pp_parity_chatml.rs b/crates/hipfire-runtime/examples/pp_parity_chatml.rs index 4402df52fa..8a13fa49a0 100644 --- a/crates/hipfire-runtime/examples/pp_parity_chatml.rs +++ b/crates/hipfire-runtime/examples/pp_parity_chatml.rs @@ -135,7 +135,9 @@ fn run_single_gpu(path: &str, prompt_tokens: &[u32]) -> (Vec, Vec> fn run_multi_gpu(path: &str, prompt_tokens: &[u32]) -> (Vec, Vec>) { let mut hfq = HfqFile::open(Path::new(path)).expect("open hfq"); let config = qwen35::config_from_hfq(&hfq).expect("config"); - let mut gpus = Gpus::init_uniform(2, config.n_layers).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let layout = qwen35::Layout::from_gpus(&gpus, config.n_layers); let mut hfq_source = qwen35::HfqSource::new(&mut hfq, &config); let weights = diff --git a/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs b/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs index ce6b8af528..b9f541f0d9 100644 --- a/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs +++ b/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs @@ -40,7 +40,9 @@ fn main() { "stale-event screen requires at least two replays" ); - let mut gpus = Gpus::init_uniform(ranks, ranks).expect("init TP GPUs"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, ranks, ranks).expect("init TP GPUs"); assert_eq!(gpus.devices.len(), ranks, "wrong GPU count"); for (rank, gpu) in gpus.devices.iter().enumerate() { assert_eq!( diff --git a/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs b/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs index 1fab187aff..9961c46986 100644 --- a/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs +++ b/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs @@ -37,7 +37,9 @@ fn main() { // n_layers placeholder — init_uniform requires n_layers >= n_devices. // The TP path doesn't care about layer-to-device; we just need devices. - let mut gpus = Gpus::init_uniform(n_ranks, n_ranks).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, n_ranks, n_ranks).expect("init_uniform"); let peer_ok = gpus.enable_peer_all().expect("enable_peer_all"); if !peer_ok { eprintln!("WARN: peer access incomplete (host-staging fallback applies)"); diff --git a/crates/hipfire-runtime/src/config.rs b/crates/hipfire-runtime/src/config.rs index 8239704165..21a1bd3c31 100644 --- a/crates/hipfire-runtime/src/config.rs +++ b/crates/hipfire-runtime/src/config.rs @@ -202,6 +202,19 @@ impl RuntimeConfig { .unwrap_or(3), } } + + /// Lower the already-resolved hardware settings into the hardware leaf's + /// dependency-free construction value. In particular, `devices` is the + /// post-visibility logical HIP-ID list produced by `hipfire-config`. + pub fn device_resolve_opts(&self) -> hipfire_hardware::DeviceResolveOpts { + hipfire_hardware::DeviceResolveOpts { + tp_use_rccl: self.tp_use_rccl, + devices: self.devices.clone(), + allow_mixed_arch: self.allow_mixed_arch, + uniform_vram_tolerance_gb: self.uniform_vram_tolerance_gb, + } + } + } #[cfg(test)] @@ -242,6 +255,11 @@ mod tests { .set_cli("generation.loop_guard_threshold", "12") .unwrap(); layer.set_cli("hardware.devices", "2,3").unwrap(); + layer.set_cli("hardware.tp_use_rccl", "false").unwrap(); + layer.set_cli("hardware.allow_mixed_arch", "true").unwrap(); + layer + .set_cli("hardware.uniform_vram_tolerance_gb", "1.5") + .unwrap(); let resolved = resolve([NamedLayer { source: ConfigSource::GlobalUser { path: "config.toml".into(), @@ -255,6 +273,11 @@ mod tests { assert!(!config.normalize_prompt); assert_eq!(config.ngram_loop_threshold, 12); assert_eq!(config.devices.as_deref(), Some("0,1")); + let opts = config.device_resolve_opts(); + assert_eq!(opts.devices.as_deref(), Some("0,1")); + assert_eq!(opts.tp_use_rccl, Some(false)); + assert!(opts.allow_mixed_arch); + assert_eq!(opts.uniform_vram_tolerance_gb, Some(1.5)); assert!(config.prefill_batched, "sparse arch defaults remain intact"); } diff --git a/crates/hipfire-runtime/src/ep.rs b/crates/hipfire-runtime/src/ep.rs index eeddad2229..96bddfaf4d 100644 --- a/crates/hipfire-runtime/src/ep.rs +++ b/crates/hipfire-runtime/src/ep.rs @@ -58,6 +58,7 @@ pub fn ensure_rank_streams(gpus: &mut Gpus) -> Result<(), DispatchError> { fn all_reduce_sum_f32_decode( gpus: &mut Gpus, + group: &[usize], refs: &[&DeviceBuffer], count: usize, ) -> Result<(), DispatchError> { @@ -67,13 +68,11 @@ fn all_reduce_sum_f32_decode( let use_peer = *PEER_DECODE.get_or_init(|| { hipfire_config::developer_var("HIPFIRE_EP_PEER_ALLREDUCE_DECODE").as_deref() == Ok("1") }); - let group: Vec = (0..refs.len()).collect(); if use_peer { - gpus - .all_reduce_sum_f32_peer(&group, refs, count) + gpus.all_reduce_sum_f32_peer(group, refs, count) .map_err(hip_err) } else { - gpus.all_reduce_sum_f32(&group, refs, count).map_err(hip_err) + gpus.all_reduce_sum_f32(group, refs, count).map_err(hip_err) } } @@ -106,6 +105,8 @@ fn tp_peer_hc3_admitted(gpus: &Gpus, bindings: &[B]) -> bool /// buffer of length `residual_dim` on `gpus.devices[r]`. The executor owns the /// zero/all-reduce/add lifecycle; the binding only writes its owned-expert /// contribution into it during `run_moe_ep`. +/// - `group` is the retained ordered global-device group for this owner. It is +/// borrowed for every collective and must be `[0, 1, ..., n-1]`. /// - `residual_dim` is the residual width (= hidden size) used for the partial /// memset byte size and the all-reduce element count. /// @@ -114,6 +115,7 @@ pub fn run_layer_program_ep( gpus: &mut Gpus, bindings: &mut [B], partials: &[GpuTensor], + group: &[usize], program: &LayerProgram, residual_dim: usize, ) -> Result<(), DispatchError> { @@ -128,6 +130,10 @@ pub fn run_layer_program_ep( n, "run_layer_program_ep: partials.len() != n_ranks" ); + assert!( + group.len() == n && group.iter().copied().eq(0..n), + "run_layer_program_ep: group must be the ordered full device group" + ); for op in program { if matches!(op.kind, SuperOpKind::Attend) @@ -216,7 +222,7 @@ pub fn run_layer_program_ep( }) }) .collect::>()?; - all_reduce_sum_f32_decode(gpus, &refs, residual_dim)?; + all_reduce_sum_f32_decode(gpus, group, &refs, residual_dim)?; for r in 0..n { gpus.devices[r].bind_thread().map_err(hip_err)?; bindings[r].ep_finish_attend(&mut gpus.devices[r])?; @@ -269,7 +275,7 @@ pub fn run_layer_program_ep( } else { // 3. All-reduce-sum the partials across ranks (in-place, RCCL). let refs: Vec<&DeviceBuffer> = partials.iter().map(|p| &p.buf).collect(); - all_reduce_sum_f32_decode(gpus, &refs, residual_dim)?; + all_reduce_sum_f32_decode(gpus, group, &refs, residual_dim)?; // 4. Fold the reduced partial into each residual stream. for r in 0..n { diff --git a/crates/saddle-lab/examples/gpus_smoke.rs b/crates/saddle-lab/examples/gpus_smoke.rs index 06da0913ba..920ae7d3b2 100644 --- a/crates/saddle-lab/examples/gpus_smoke.rs +++ b/crates/saddle-lab/examples/gpus_smoke.rs @@ -12,7 +12,9 @@ use rdna_compute::{DType, Gpu}; fn main() { println!("── Gpus::init_uniform(2, 24) ─────────────────────────────"); - let mut gpus = Gpus::init_uniform(2, 24).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, 2, 24).expect("init_uniform"); assert_eq!(gpus.devices.len(), 2); assert_eq!(gpus.layer_to_device.len(), 24); assert_eq!(gpus.band_starts, vec![0, 12]); @@ -97,7 +99,7 @@ fn main() { println!("\n── Gpus::single back-compat ──────────────────────────────"); drop(gpus); // release dev 0/1 before re-init for single let solo = Gpu::init_with_device(0).expect("init solo"); - let single = Gpus::single(solo, 24); + let single = Gpus::single(&device_opts, solo, 24); assert_eq!(single.devices.len(), 1); assert_eq!(single.layer_to_device, vec![0u8; 24]); assert_eq!(single.output_device, 0); diff --git a/crates/saddle-lab/examples/pp2_vram_probe.rs b/crates/saddle-lab/examples/pp2_vram_probe.rs index bf35a946f5..e8abc0cbe5 100644 --- a/crates/saddle-lab/examples/pp2_vram_probe.rs +++ b/crates/saddle-lab/examples/pp2_vram_probe.rs @@ -60,7 +60,9 @@ fn main() { config.head_dim, ); - let mut gpus = Gpus::init_uniform(2, config.n_layers).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let baseline_free: Vec<(usize, usize)> = (0..gpus.devices.len()) .map(|i| gpus.devices[i].hip.get_vram_info().unwrap_or((0, 0))) .collect(); diff --git a/crates/saddle-lab/examples/pp_parity.rs b/crates/saddle-lab/examples/pp_parity.rs index f8ce79bd95..874e08d2ec 100644 --- a/crates/saddle-lab/examples/pp_parity.rs +++ b/crates/saddle-lab/examples/pp_parity.rs @@ -84,7 +84,9 @@ fn run_single_gpu(path: &str) -> Vec { fn run_multi_gpu(path: &str) -> Vec { let mut hfq = HfqFile::open(Path::new(path)).expect("open hfq"); let config = qwen35::config_from_hfq(&hfq).expect("config"); - let mut gpus = Gpus::init_uniform(2, config.n_layers).expect("init_uniform"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let mut gpus = + Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let layout = qwen35::Layout::from_gpus(&gpus, config.n_layers); let mut hfq_source = qwen35::HfqSource::new(&mut hfq, &config); let weights = qwen35::load_weights(&mut hfq_source, &mut gpus.devices, &layout) From 61be6397fb0f2604ca6cfced4be51950f40674e1 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 19:16:22 +0200 Subject: [PATCH 03/25] fix(hardware): restore active peer lease record --- crates/hipfire-hardware/src/lib.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/crates/hipfire-hardware/src/lib.rs b/crates/hipfire-hardware/src/lib.rs index 197196ea57..037a582cbc 100644 --- a/crates/hipfire-hardware/src/lib.rs +++ b/crates/hipfire-hardware/src/lib.rs @@ -121,6 +121,13 @@ pub fn peer_reduce_scratch_total_bytes(rank_count: usize, requested_bytes: usize } /// Internal active lease record stored inside `Gpus` while a lease is live. +#[derive(Debug)] +struct ActivePeerLease { + id: u64, + bytes: usize, + rank_count: usize, +} + pub struct Gpus { /// RCCL communicators (one per rank), lazily initialized on the first /// `all_reduce_sum_*` call. Declared BEFORE `devices` so `Drop` tears From 2a831a753f071fefdc1d8e0d69286814d0e96b48 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 00:30:05 +0200 Subject: [PATCH 04/25] style(device-mesh): format hardware cutover --- .../examples/ep_deepseek4.rs | 2 +- .../examples/ep_dspark_topology_probe.rs | 2 +- crates/hipfire-arch-deepseek4/src/ep.rs | 10 +- crates/hipfire-arch-deepseek4/src/forward.rs | 4 +- crates/hipfire-arch-deepseek4/src/mtp.rs | 12 +- .../examples/ep_minimax.rs | 125 ++++++++++++++---- crates/hipfire-arch-minimax/src/forward.rs | 5 +- .../examples/qwen_dense_tp2_parity.rs | 2 +- .../examples/test_qwen35_load_multi.rs | 7 +- .../examples/test_qwen35_state_multi.rs | 5 +- .../src/qwen35/ep_batch.rs | 8 +- .../hipfire-arch-qwen35/src/qwen35/forward.rs | 2 +- .../hipfire-arch-qwen35/src/qwen35/weights.rs | 2 +- crates/hipfire-arch-qwen35/tests/pp_parity.rs | 5 +- crates/hipfire-hardware/src/mesh.rs | 16 +-- crates/hipfire-hardware/tests/ownership.rs | 4 +- crates/hipfire-loader/src/lib.rs | 10 +- .../ds4_gfx1201_owner_worker_transport.rs | 3 +- .../examples/ep_decode_parity.rs | 2 +- .../examples/pp_parity_chatml.rs | 5 +- .../tp4_cross_device_graph_barrier.rs | 3 +- .../examples/tp_allreduce_smoke.rs | 6 +- crates/hipfire-runtime/src/config.rs | 1 - crates/hipfire-runtime/src/ep.rs | 2 +- crates/hipfire-runtime/src/llama.rs | 2 +- crates/hipfire-runtime/src/model_load.rs | 2 +- crates/saddle-lab/examples/gpus_smoke.rs | 3 +- crates/saddle-lab/examples/pp2_vram_probe.rs | 5 +- crates/saddle-lab/examples/pp_parity.rs | 5 +- 29 files changed, 159 insertions(+), 101 deletions(-) diff --git a/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs b/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs index 582276bcf2..fa8b35157d 100644 --- a/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs +++ b/crates/hipfire-arch-deepseek4/examples/ep_deepseek4.rs @@ -30,9 +30,9 @@ fn fnv1a(ids: &[u32]) -> u64 { fn main() { use hipfire_arch_deepseek4::forward; use hipfire_arch_deepseek4::{DeepseekV4, DeepseekV4State}; + use hipfire_hardware::Gpus; use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::HfqFile; - use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; diff --git a/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs b/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs index faee247042..e52461ddf0 100644 --- a/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs +++ b/crates/hipfire-arch-deepseek4/examples/ep_dspark_topology_probe.rs @@ -15,9 +15,9 @@ use hipfire_arch_deepseek4::forward::{ PrefillBatchScratch, }; use hipfire_arch_deepseek4::{DeepseekV4, DeepseekV4State}; +use hipfire_hardware::Gpus; use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::HfqFile; -use hipfire_hardware::Gpus; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::path::{Path, PathBuf}; diff --git a/crates/hipfire-arch-deepseek4/src/ep.rs b/crates/hipfire-arch-deepseek4/src/ep.rs index dc32a58b30..71e9ae76b7 100644 --- a/crates/hipfire-arch-deepseek4/src/ep.rs +++ b/crates/hipfire-arch-deepseek4/src/ep.rs @@ -5,10 +5,13 @@ use crate::config_cache; use crate::deepseek4::{DeepseekV4Config, DeepseekV4State, DeepseekV4Weights}; use crate::forward::{ - Deepseek4Bindings, compressor_cache_uses_vmm, ds4_lower_program, final_norm_and_head, - init_residual_streams, refresh_compressor_cache_shard_tables, + compressor_cache_uses_vmm, ds4_lower_program, final_norm_and_head, init_residual_streams, + refresh_compressor_cache_shard_tables, Deepseek4Bindings, +}; +use crate::forward::{ + ensure_compressor_capacity, precompute_positions, precompute_token_id, update_attn_state_host, + update_pos_array_host, update_token_id_host, }; -use crate::forward::{precompute_positions, precompute_token_id, update_attn_state_host, update_pos_array_host, update_token_id_host, ensure_compressor_capacity}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::superop::{self, SuperOpKind}; use hipfire_hardware::Gpus; @@ -532,4 +535,3 @@ fn forward_ep_direct( } Ok(()) } - diff --git a/crates/hipfire-arch-deepseek4/src/forward.rs b/crates/hipfire-arch-deepseek4/src/forward.rs index 5a0c29c2f2..fe54503b17 100644 --- a/crates/hipfire-arch-deepseek4/src/forward.rs +++ b/crates/hipfire-arch-deepseek4/src/forward.rs @@ -1481,7 +1481,9 @@ pub fn ensure_request_capacity( Ok(scratch_grew || cache_grew) } -pub(crate) fn refresh_compressor_cache_shard_tables(states: &mut [DeepseekV4State]) -> Result<(), String> { +pub(crate) fn refresh_compressor_cache_shard_tables( + states: &mut [DeepseekV4State], +) -> Result<(), String> { let world = states.len(); if !matches!(world, 3 | 4) { return Err(format!( diff --git a/crates/hipfire-arch-deepseek4/src/mtp.rs b/crates/hipfire-arch-deepseek4/src/mtp.rs index 272cad46d8..51a780c35a 100644 --- a/crates/hipfire-arch-deepseek4/src/mtp.rs +++ b/crates/hipfire-arch-deepseek4/src/mtp.rs @@ -5,14 +5,13 @@ use crate::config_cache; use crate::deepseek4::{DeepseekV4Config, DeepseekV4State, DeepseekV4Weights}; use crate::forward::{ - Deepseek4Bindings, OloraSchedule, apply_tail_rope, apply_tail_rope_batched, attn_stub, ds4_superop, ffn_routed, ffn_stub, - gemv_auto, hc_attn_mix, hc_ffn_mix, kv_joint, mhc_pre, q_lora, - weight_needs_fwht, precompute_attn_state_batched, precompute_positions_batched, + apply_tail_rope, apply_tail_rope_batched, attn_stub, ds4_superop, ffn_routed, ffn_stub, + gemv_auto, hc_attn_mix, hc_ffn_mix, kv_joint, mhc_pre, precompute_attn_state_batched, + precompute_positions_batched, q_lora, weight_needs_fwht, Deepseek4Bindings, OloraSchedule, }; use crate::forward::{ - attention_block_batched_swa_only, ffn_batched, gemv_auto_batched_wmma, - hc_attn_mix_batched, hc_ffn_mix_batched, kv_joint_batched, mhc_pre_batched, - q_lora_batched, + attention_block_batched_swa_only, ffn_batched, gemv_auto_batched_wmma, hc_attn_mix_batched, + hc_ffn_mix_batched, kv_joint_batched, mhc_pre_batched, q_lora_batched, }; use hipfire_dispatch::pipeline::superop::SuperOpKind; use rdna_compute::{DType, Gpu, GpuTensor}; @@ -915,4 +914,3 @@ pub fn mtp_forward_batched( Ok(()) } - diff --git a/crates/hipfire-arch-minimax/examples/ep_minimax.rs b/crates/hipfire-arch-minimax/examples/ep_minimax.rs index ddbfbe3d43..59394f5c35 100644 --- a/crates/hipfire-arch-minimax/examples/ep_minimax.rs +++ b/crates/hipfire-arch-minimax/examples/ep_minimax.rs @@ -35,8 +35,8 @@ fn fnv1a(ids: &[u32]) -> u64 { fn main() { use hipfire_arch_minimax::forward; use hipfire_arch_minimax::minimax::{MiniMaxConfig, MiniMaxState, MiniMaxWeights}; - use hipfire_runtime::hfq::HfqFile; use hipfire_hardware::Gpus; + use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; @@ -50,11 +50,26 @@ fn main() { let mut i = 1; while i < argv.len() { match argv[i].as_str() { - "--model" => { model = Some(PathBuf::from(&argv[i + 1])); i += 2; } - "--prompt" => { prompt = argv[i + 1].clone(); i += 2; } - "--max" => { max = argv[i + 1].parse().expect("--max"); i += 2; } - "--tp" => { tp = argv[i + 1].parse().expect("--tp"); i += 2; } - other => { eprintln!("unknown arg {other}"); std::process::exit(1); } + "--model" => { + model = Some(PathBuf::from(&argv[i + 1])); + i += 2; + } + "--prompt" => { + prompt = argv[i + 1].clone(); + i += 2; + } + "--max" => { + max = argv[i + 1].parse().expect("--max"); + i += 2; + } + "--tp" => { + tp = argv[i + 1].parse().expect("--tp"); + i += 2; + } + other => { + eprintln!("unknown arg {other}"); + std::process::exit(1); + } } } let model = model.expect("--model required"); @@ -74,14 +89,22 @@ fn main() { let device_opts = hipfire_runtime::config::get().device_resolve_opts(); let mut gpus = Gpus::init_tp(&device_opts, tp, cfg.num_hidden_layers).expect("init_tp"); let n = gpus.devices.len(); - assert_eq!(n, tp, "init_tp gave {n} devices (check HIP_VISIBLE_DEVICES)"); + assert_eq!( + n, tp, + "init_tp gave {n} devices (check HIP_VISIBLE_DEVICES)" + ); for (r, d) in gpus.devices.iter().enumerate() { eprintln!(" rank {r}: device_id={} arch={}", d.device_id, d.arch); } // ── shard-aware replicated load (each rank uploads only its owned experts) ─ - let shard = ShardConfig::new(tp, /*tp_kv_replicate=*/ true, n_exp, ExpertAssign::Stride) - .expect("ShardConfig"); + let shard = ShardConfig::new( + tp, + /*tp_kv_replicate=*/ true, + n_exp, + ExpertAssign::Stride, + ) + .expect("ShardConfig"); let mut weights_per_rank: Vec = Vec::with_capacity(n); for r in 0..n { gpus.devices[r].bind_thread().expect("bind"); @@ -89,7 +112,10 @@ fn main() { let t = std::time::Instant::now(); let w = MiniMaxWeights::load(&mut hfq, &cfg, &mut gpus.devices[r], Some((&shard, r))) .expect("shard-aware load"); - eprintln!(" [rank {r}] loaded owned shard in {:.1}s", t.elapsed().as_secs_f64()); + eprintln!( + " [rank {r}] loaded owned shard in {:.1}s", + t.elapsed().as_secs_f64() + ); weights_per_rank.push(w); } eprintln!(" all ranks loaded (stride: rank r owns experts e%{tp}==r)"); @@ -104,15 +130,25 @@ fn main() { state_per_rank.push( MiniMaxState::new_with_max_seq(&mut gpus.devices[r], &cfg, max_seq).expect("state"), ); - partials.push(gpus.devices[r].zeros(&[cfg.hidden_size], DType::F32).expect("partial")); + partials.push( + gpus.devices[r] + .zeros(&[cfg.hidden_size], DType::F32) + .expect("partial"), + ); } let peer = gpus.enable_peer_all().expect("enable_peer_all"); eprintln!(" peer_access_enabled={peer}"); hipfire_runtime::ep::ensure_rank_streams(&mut gpus).expect("ensure_rank_streams"); let argmax = |v: &[f32]| -> u32 { - let mut bi = 0u32; let mut bv = f32::NEG_INFINITY; - for (i, &x) in v.iter().enumerate() { if x > bv { bv = x; bi = i as u32; } } + let mut bi = 0u32; + let mut bv = f32::NEG_INFINITY; + for (i, &x) in v.iter().enumerate() { + if x > bv { + bv = x; + bi = i as u32; + } + } bi }; @@ -120,12 +156,26 @@ fn main() { eprintln!("\nprompt {:?} → {} tokens", prompt, prompt_ids.len()); let t0 = std::time::Instant::now(); for (pos, &t) in prompt_ids.iter().enumerate() { - forward::forward_ep(&mut gpus, &weights_per_rank, &cfg, &mut state_per_rank, &partials, t, pos as u32) - .expect("forward_ep prefill"); + forward::forward_ep( + &mut gpus, + &weights_per_rank, + &cfg, + &mut state_per_rank, + &partials, + t, + pos as u32, + ) + .expect("forward_ep prefill"); } gpus.devices[0].bind_thread().expect("bind0"); - let mut logits = gpus.devices[0].download_f32(&state_per_rank[0].logits).expect("dl"); - eprintln!("prefill {} tok in {:.2}s", prompt_ids.len(), t0.elapsed().as_secs_f64()); + let mut logits = gpus.devices[0] + .download_f32(&state_per_rank[0].logits) + .expect("dl"); + eprintln!( + "prefill {} tok in {:.2}s", + prompt_ids.len(), + t0.elapsed().as_secs_f64() + ); let mut gen = Vec::new(); let mut pos = prompt_ids.len(); @@ -138,21 +188,46 @@ fn main() { if matches!(next, 200020 | 151643 | 151645 | 2) { break; } - if step == 2 { steady_t = std::time::Instant::now(); steady = 0; } - forward::forward_ep(&mut gpus, &weights_per_rank, &cfg, &mut state_per_rank, &partials, next, pos as u32) - .expect("forward_ep decode"); + if step == 2 { + steady_t = std::time::Instant::now(); + steady = 0; + } + forward::forward_ep( + &mut gpus, + &weights_per_rank, + &cfg, + &mut state_per_rank, + &partials, + next, + pos as u32, + ) + .expect("forward_ep decode"); gpus.devices[0].bind_thread().expect("bind0"); - logits = gpus.devices[0].download_f32(&state_per_rank[0].logits).expect("dl"); - if step >= 2 { steady += 1; } + logits = gpus.devices[0] + .download_f32(&state_per_rank[0].logits) + .expect("dl"); + if step >= 2 { + steady += 1; + } pos += 1; } let dt = t1.elapsed().as_secs_f64(); - let steady_tps = if steady > 0 { steady as f64 / steady_t.elapsed().as_secs_f64() } else { f64::NAN }; + let steady_tps = if steady > 0 { + steady as f64 / steady_t.elapsed().as_secs_f64() + } else { + f64::NAN + }; eprintln!( "decoded {} tok in {:.2}s ({:.1} tok/s overall, {:.1} tok/s steady)", - gen.len(), dt, gen.len() as f64 / dt, steady_tps, + gen.len(), + dt, + gen.len() as f64 / dt, + steady_tps, + ); + println!( + "=== PROMPT ===\n{prompt}\n=== GENERATION (tp={tp} EP) ===\n{}", + tok.decode(&gen) ); - println!("=== PROMPT ===\n{prompt}\n=== GENERATION (tp={tp} EP) ===\n{}", tok.decode(&gen)); eprintln!("gen ids: {:?}", &gen[..gen.len().min(40)]); eprintln!("gen FNV: 0x{:016x}", fnv1a(&gen)); } diff --git a/crates/hipfire-arch-minimax/src/forward.rs b/crates/hipfire-arch-minimax/src/forward.rs index a900763272..2fd8b38908 100644 --- a/crates/hipfire-arch-minimax/src/forward.rs +++ b/crates/hipfire-arch-minimax/src/forward.rs @@ -34,9 +34,10 @@ use hipfire_dispatch::pipeline::superop::{ }; use hipfire_dispatch::pipeline::{execute_steps, GemvInput, Step}; use hipfire_dispatch::types::{dtype_rotation_plan, DispatchError}; -use hipfire_runtime::llama::{ - fused_silu_mul_rotate_mq_batched_for, rotate_x_mq_batched_for, rotate_x_mq_for, weight_gemv}; use hipfire_runtime::llama::KvCacheExt; +use hipfire_runtime::llama::{ + fused_silu_mul_rotate_mq_batched_for, rotate_x_mq_batched_for, rotate_x_mq_for, weight_gemv, +}; use rdna_compute::{DType, Gpu, GpuTensor}; /// Decode one token (eager); returns the full logits vector. Used for prefill, diff --git a/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs b/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs index ce652feaed..5336b8c466 100644 --- a/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs +++ b/crates/hipfire-arch-qwen35/examples/qwen_dense_tp2_parity.rs @@ -3,10 +3,10 @@ use hipfire_arch_qwen35::qwen35::{ self, DeltaNetState, HfqSource, Layout, Qwen35Scratch, StateQuant, }; +use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::kv_mode::KvMode; use hipfire_runtime::llama::{KvCache, KvCacheExt, KvDims, KvLayers, KvTarget}; -use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::Gpu; diff --git a/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs b/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs index 41c75378e8..d3f6a6920b 100644 --- a/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs +++ b/crates/hipfire-arch-qwen35/examples/test_qwen35_load_multi.rs @@ -13,9 +13,9 @@ //! ~/.hipfire/models/qwen3.5-0.8b.mq4 use hipfire_arch_qwen35::qwen35; -use hipfire_runtime::llama::KvCacheExt; -use hipfire_runtime::hfq::HfqFile; use hipfire_hardware::Gpus; +use hipfire_runtime::hfq::HfqFile; +use hipfire_runtime::llama::KvCacheExt; use std::path::Path; fn main() { @@ -28,8 +28,7 @@ fn main() { ); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let n = gpus.devices.len(); let out_dev = gpus.output_device; println!( diff --git a/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs b/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs index 76835b0d7e..c553093893 100644 --- a/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs +++ b/crates/hipfire-arch-qwen35/examples/test_qwen35_state_multi.rs @@ -16,10 +16,10 @@ //! ~/.hipfire/models/qwen3.5-0.8b.mq4 use hipfire_arch_qwen35::qwen35::{self, DeltaNetState, LayerType, Qwen35ScratchSet, StateQuant}; +use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_hardware::Gpus; use std::path::Path; fn main() { @@ -32,8 +32,7 @@ fn main() { ); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let n_dev = gpus.devices.len(); let out_dev = gpus.output_device; println!( diff --git a/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs b/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs index 16e61a7462..3f46cea213 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/ep_batch.rs @@ -50,13 +50,13 @@ use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::execute_steps; use hipfire_dispatch::pipeline::GemvInput; use hipfire_dispatch::pipeline::Step; +use hipfire_hardware::Gpus; use hipfire_runtime::llama; use hipfire_runtime::llama::fused_rmsnorm_rotate_for_mq; use hipfire_runtime::llama::weight_gemv_prerotated; use hipfire_runtime::llama::weight_gemv_swiglu_residual; use hipfire_runtime::llama::EmbeddingFormat; use hipfire_runtime::llama::WeightTensor; -use hipfire_hardware::Gpus; use rdna_compute::DType; use rdna_compute::GpuTensor; @@ -2462,12 +2462,10 @@ pub fn forward_prefill_batch_ep( let refs: Vec<&hip_bridge::DeviceBuffer> = partials.iter().map(|partial| &partial.buf).collect(); if ep_peer_ar { - gpus - .all_reduce_sum_f32_peer(&group, &refs, n * dim) + gpus.all_reduce_sum_f32_peer(&group, &refs, n * dim) .map_err(|e| HipError::new(0, &e.to_string()))?; } else { - gpus - .all_reduce_sum_f32(&group, &refs, n * dim) + gpus.all_reduce_sum_f32(&group, &refs, n * dim) .map_err(|e| HipError::new(0, &e.to_string()))?; } if ep_timing { diff --git a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs index f45076f570..0997be3659 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/forward.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/forward.rs @@ -45,13 +45,13 @@ use hipfire_dispatch::pipeline::Step; use hipfire_dispatch::types::dtype_rotation_plan; use hipfire_dispatch::types::DispatchError; use hipfire_dispatch::types::RotationPlan; +use hipfire_hardware::Gpus; use hipfire_runtime::llama; use hipfire_runtime::llama::fused_rmsnorm_rotate_for_mq; use hipfire_runtime::llama::EmbeddingFormat; use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::llama::ParoRotation; use hipfire_runtime::llama::WeightTensor; -use hipfire_hardware::Gpus; use hipfire_runtime::tp_shard::ShardConfig; use rdna_compute::DType; use rdna_compute::Gpu; diff --git a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs index fc74135817..5e9716d50e 100644 --- a/crates/hipfire-arch-qwen35/src/qwen35/weights.rs +++ b/crates/hipfire-arch-qwen35/src/qwen35/weights.rs @@ -9,10 +9,10 @@ use super::config::LayerType; use super::config::Qwen35Config; use hip_bridge::HipError; use hip_bridge::HipResult; +use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::EmbeddingFormat; use hipfire_runtime::llama::WeightTensor; -use hipfire_hardware::Gpus; use hipfire_runtime::screen_weight_tensor; use hipfire_runtime::MmqScreenable; use rdna_compute::DType; diff --git a/crates/hipfire-arch-qwen35/tests/pp_parity.rs b/crates/hipfire-arch-qwen35/tests/pp_parity.rs index a9ffbbb882..1a4d47874b 100644 --- a/crates/hipfire-arch-qwen35/tests/pp_parity.rs +++ b/crates/hipfire-arch-qwen35/tests/pp_parity.rs @@ -21,10 +21,10 @@ use hipfire_arch_qwen35::qwen35::{ self, DeltaNetState, Qwen35Scratch, Qwen35ScratchSet, StateQuant, }; +use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use rdna_compute::Gpu; use std::path::Path; @@ -118,8 +118,7 @@ fn run_pp2(path: &str, prompt: &[u32]) -> Vec { let mut hfq = HfqFile::open(Path::new(path)).expect("open hfq"); let config = qwen35::config_from_hfq(&hfq).expect("config"); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let layout = qwen35::Layout::from_gpus(&gpus, config.n_layers); let mut hfq_source = qwen35::HfqSource::new(&mut hfq, &config); let weights = qwen35::load_weights(&mut hfq_source, &mut gpus.devices, &layout) diff --git a/crates/hipfire-hardware/src/mesh.rs b/crates/hipfire-hardware/src/mesh.rs index d0fc7f83f1..eb6fa79172 100644 --- a/crates/hipfire-hardware/src/mesh.rs +++ b/crates/hipfire-hardware/src/mesh.rs @@ -359,12 +359,8 @@ mod tests { #[test] fn composed_coordinates_groups_stages_and_squeeze() { - let mesh = DeviceMesh::rect(&[ - (DimKind::Pp, 2), - (DimKind::Tp, 2), - (DimKind::Ep, 2), - ]) - .unwrap(); + let mesh = + DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2), (DimKind::Ep, 2)]).unwrap(); assert_eq!(mesh.n_devices(), 8); assert_eq!(mesh.coord_of(0), vec![0, 0, 0]); assert_eq!(mesh.coord_of(7), vec![1, 1, 1]); @@ -386,12 +382,8 @@ mod tests { } } - let degenerate = DeviceMesh::rect(&[ - (DimKind::Pp, 2), - (DimKind::Tp, 1), - (DimKind::Ep, 2), - ]) - .unwrap(); + let degenerate = + DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 1), (DimKind::Ep, 2)]).unwrap(); assert_eq!( degenerate.squeezed().axes(), &[ diff --git a/crates/hipfire-hardware/tests/ownership.rs b/crates/hipfire-hardware/tests/ownership.rs index 5b41b5439d..653d945770 100644 --- a/crates/hipfire-hardware/tests/ownership.rs +++ b/crates/hipfire-hardware/tests/ownership.rs @@ -20,7 +20,9 @@ fn hardware_leaf_exposes_owner_and_named_topology() { #[test] fn runtime_has_no_legacy_owner_or_compatibility_reexport() { let crate_root = Path::new(env!("CARGO_MANIFEST_DIR")); - assert!(!crate_root.join("../hipfire-runtime/src/multi_gpu.rs").exists()); + assert!(!crate_root + .join("../hipfire-runtime/src/multi_gpu.rs") + .exists()); let runtime_lib = std::fs::read_to_string(crate_root.join("../hipfire-runtime/src/lib.rs")) .expect("runtime lib source"); assert!(!runtime_lib.contains("pub mod multi_gpu")); diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index cadae3c795..c97fd9d3a1 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -24,6 +24,7 @@ use hipfire_arch_qwen35::qwen35::{self}; use hipfire_arch_qwen35::speculative::DeltaNetSnapshot; use hipfire_arch_qwen35::Qwen35Bundle; use hipfire_arch_qwen35_vl::qwen35_vl; +use hipfire_hardware::Gpus; use hipfire_runtime::arch_model::ArchModel; use hipfire_runtime::cask::CaskCtx; use hipfire_runtime::hfq::HfqFile; @@ -32,7 +33,6 @@ use hipfire_runtime::kv_mode; use hipfire_runtime::llama; use hipfire_runtime::llama::{KvCacheExt, KvDims, KvLayers, KvTarget}; use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; -use hipfire_hardware::Gpus; use hipfire_runtime::spec::{SpecEmit, SpecEmitCtx, SpecTargetGuard, Speculator}; use hipfire_runtime::triattn::{EvictionCtx, TriAttnCenters}; use rdna_compute::Gpu; @@ -3361,8 +3361,8 @@ fn load_model_ep_qwen35( let chat_template = resolve_chat_template(&hfq_probe, path); let rec = hfq_probe.recommended_sampling(); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let gpus = Gpus::init_tp(&device_opts, tp, config.n_layers) - .map_err(|e| format!("init_tp: {e:?}"))?; + let gpus = + Gpus::init_tp(&device_opts, tp, config.n_layers).map_err(|e| format!("init_tp: {e:?}"))?; let n = gpus.devices.len(); if n != tp { return Err(format!( @@ -3494,8 +3494,8 @@ fn load_model_tp_qwen35_dense( drop(hfq); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let gpus = Gpus::init_tp(&device_opts, tp, config.n_layers) - .map_err(|e| format!("init_tp: {e:?}"))?; + let gpus = + Gpus::init_tp(&device_opts, tp, config.n_layers).map_err(|e| format!("init_tp: {e:?}"))?; if gpus.devices.len() != tp { return Err(format!( "init_tp gave {} devices, expected tp={tp}", diff --git a/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs b/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs index c4955e8cc1..0e212631bc 100644 --- a/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs +++ b/crates/hipfire-runtime/examples/ds4_gfx1201_owner_worker_transport.rs @@ -121,8 +121,7 @@ fn main() { assert!(samples > 0, "samples must be nonzero"); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, RANKS, RANKS).expect("init four GPUs"); + let mut gpus = Gpus::init_uniform(&device_opts, RANKS, RANKS).expect("init four GPUs"); assert_eq!(gpus.devices.len(), RANKS, "requires exactly four GPUs"); for (rank, gpu) in gpus.devices.iter().enumerate() { assert_eq!( diff --git a/crates/hipfire-runtime/examples/ep_decode_parity.rs b/crates/hipfire-runtime/examples/ep_decode_parity.rs index a00e488ac3..4fe2b2a1e2 100644 --- a/crates/hipfire-runtime/examples/ep_decode_parity.rs +++ b/crates/hipfire-runtime/examples/ep_decode_parity.rs @@ -48,9 +48,9 @@ fn fnv1a(ids: &[u32]) -> u64 { #[cfg(feature = "deltanet")] fn main() { use hipfire_arch_qwen35::qwen35::{self, DeltaNetState, Qwen35Scratch}; + use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::{self, KvCache}; - use hipfire_hardware::Gpus; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; use std::path::Path; diff --git a/crates/hipfire-runtime/examples/pp_parity_chatml.rs b/crates/hipfire-runtime/examples/pp_parity_chatml.rs index 8a13fa49a0..5b39cfb76b 100644 --- a/crates/hipfire-runtime/examples/pp_parity_chatml.rs +++ b/crates/hipfire-runtime/examples/pp_parity_chatml.rs @@ -18,10 +18,10 @@ use hipfire_arch_qwen35::qwen35::{ self, DeltaNetState, Qwen35Scratch, Qwen35ScratchSet, StateQuant, }; +use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_hardware::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use rdna_compute::Gpu; use std::path::Path; @@ -136,8 +136,7 @@ fn run_multi_gpu(path: &str, prompt_tokens: &[u32]) -> (Vec, Vec>) let mut hfq = HfqFile::open(Path::new(path)).expect("open hfq"); let config = qwen35::config_from_hfq(&hfq).expect("config"); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let layout = qwen35::Layout::from_gpus(&gpus, config.n_layers); let mut hfq_source = qwen35::HfqSource::new(&mut hfq, &config); let weights = diff --git a/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs b/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs index b9f541f0d9..7726af5f4c 100644 --- a/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs +++ b/crates/hipfire-runtime/examples/tp4_cross_device_graph_barrier.rs @@ -41,8 +41,7 @@ fn main() { ); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, ranks, ranks).expect("init TP GPUs"); + let mut gpus = Gpus::init_uniform(&device_opts, ranks, ranks).expect("init TP GPUs"); assert_eq!(gpus.devices.len(), ranks, "wrong GPU count"); for (rank, gpu) in gpus.devices.iter().enumerate() { assert_eq!( diff --git a/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs b/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs index 9961c46986..dc4d975f3a 100644 --- a/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs +++ b/crates/hipfire-runtime/examples/tp_allreduce_smoke.rs @@ -38,8 +38,7 @@ fn main() { // n_layers placeholder — init_uniform requires n_layers >= n_devices. // The TP path doesn't care about layer-to-device; we just need devices. let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, n_ranks, n_ranks).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, n_ranks, n_ranks).expect("init_uniform"); let peer_ok = gpus.enable_peer_all().expect("enable_peer_all"); if !peer_ok { eprintln!("WARN: peer access incomplete (host-staging fallback applies)"); @@ -151,8 +150,7 @@ fn main() { let mut samples = Vec::with_capacity(iters); for _ in 0..iters { let t = Instant::now(); - gpus - .all_reduce_sum_f32(&group, &refs, count) + gpus.all_reduce_sum_f32(&group, &refs, count) .expect("all_reduce"); for dev in &gpus.devices { dev.bind_thread().expect("bind"); diff --git a/crates/hipfire-runtime/src/config.rs b/crates/hipfire-runtime/src/config.rs index 21a1bd3c31..bb5e9e5566 100644 --- a/crates/hipfire-runtime/src/config.rs +++ b/crates/hipfire-runtime/src/config.rs @@ -214,7 +214,6 @@ impl RuntimeConfig { uniform_vram_tolerance_gb: self.uniform_vram_tolerance_gb, } } - } #[cfg(test)] diff --git a/crates/hipfire-runtime/src/ep.rs b/crates/hipfire-runtime/src/ep.rs index 96bddfaf4d..775a63a550 100644 --- a/crates/hipfire-runtime/src/ep.rs +++ b/crates/hipfire-runtime/src/ep.rs @@ -31,13 +31,13 @@ //! driver loops layers (advancing each rank's per-layer binding state) the same //! way the single-GPU lowered driver loops `run_layer_program`. -use hipfire_hardware::Gpus; use hip_bridge::{DeviceBuffer, HipError}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::superop::{ dispatch_super_op, ForwardBindings, LayerProgram, SuperOpKind, }; use hipfire_dispatch::types::DispatchError; +use hipfire_hardware::Gpus; use rdna_compute::GpuTensor; fn hip_err(e: HipError) -> DispatchError { diff --git a/crates/hipfire-runtime/src/llama.rs b/crates/hipfire-runtime/src/llama.rs index 2e1e5dc167..92db944488 100644 --- a/crates/hipfire-runtime/src/llama.rs +++ b/crates/hipfire-runtime/src/llama.rs @@ -11,8 +11,8 @@ use crate::kv_backend::{ KvBackend, KvChunkPlan, DEFAULT_KV_CHUNK_TOKENS, DEFAULT_VMM_PHYSICAL_CHUNK_BYTES, }; use crate::kv_mode::KvMode; -use hipfire_hardware::Gpus; use hip_bridge::HipResult; +use hipfire_hardware::Gpus; use rdna_compute::{DType, Gpu, GpuTensor}; /// Model architecture type. diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index e8349e9e75..dadc63f8b6 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -6,8 +6,8 @@ //! per-tensor dequant), which `WeightSource::read_layer` calls internally. use crate::llama::{EmbeddingFormat, WeightTensor}; -use hipfire_hardware::Gpus; use hip_bridge::HipResult; +use hipfire_hardware::Gpus; use rdna_compute::{Gpu, GpuTensor}; /// Where each piece of the model lands across a device slice. `single` = the diff --git a/crates/saddle-lab/examples/gpus_smoke.rs b/crates/saddle-lab/examples/gpus_smoke.rs index 920ae7d3b2..629a1a3558 100644 --- a/crates/saddle-lab/examples/gpus_smoke.rs +++ b/crates/saddle-lab/examples/gpus_smoke.rs @@ -13,8 +13,7 @@ use rdna_compute::{DType, Gpu}; fn main() { println!("── Gpus::init_uniform(2, 24) ─────────────────────────────"); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, 2, 24).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, 2, 24).expect("init_uniform"); assert_eq!(gpus.devices.len(), 2); assert_eq!(gpus.layer_to_device.len(), 24); assert_eq!(gpus.band_starts, vec![0, 12]); diff --git a/crates/saddle-lab/examples/pp2_vram_probe.rs b/crates/saddle-lab/examples/pp2_vram_probe.rs index e8abc0cbe5..ad228faec2 100644 --- a/crates/saddle-lab/examples/pp2_vram_probe.rs +++ b/crates/saddle-lab/examples/pp2_vram_probe.rs @@ -13,10 +13,10 @@ //! ~/.hipfire/models/qwen3.5-0.8b.mq4 [max_seq=4096] use hipfire_arch_qwen35::qwen35::{self, DeltaNetState, Qwen35ScratchSet, StateQuant}; +use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_hardware::Gpus; use std::path::Path; fn used_gb(gpus: &Gpus, baseline_free: &[(usize, usize)]) -> Vec { @@ -61,8 +61,7 @@ fn main() { ); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let baseline_free: Vec<(usize, usize)> = (0..gpus.devices.len()) .map(|i| gpus.devices[i].hip.get_vram_info().unwrap_or((0, 0))) .collect(); diff --git a/crates/saddle-lab/examples/pp_parity.rs b/crates/saddle-lab/examples/pp_parity.rs index 874e08d2ec..22ba578869 100644 --- a/crates/saddle-lab/examples/pp_parity.rs +++ b/crates/saddle-lab/examples/pp_parity.rs @@ -15,10 +15,10 @@ use hipfire_arch_qwen35::qwen35::{ self, DeltaNetState, Qwen35Scratch, Qwen35ScratchSet, StateQuant, }; +use hipfire_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCache; use hipfire_runtime::llama::KvCacheExt; -use hipfire_hardware::Gpus; use rdna_compute::Gpu; use std::path::Path; @@ -85,8 +85,7 @@ fn run_multi_gpu(path: &str) -> Vec { let mut hfq = HfqFile::open(Path::new(path)).expect("open hfq"); let config = qwen35::config_from_hfq(&hfq).expect("config"); let device_opts = hipfire_runtime::config::get().device_resolve_opts(); - let mut gpus = - Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); + let mut gpus = Gpus::init_uniform(&device_opts, 2, config.n_layers).expect("init_uniform"); let layout = qwen35::Layout::from_gpus(&gpus, config.n_layers); let mut hfq_source = qwen35::HfqSource::new(&mut hfq, &config); let weights = qwen35::load_weights(&mut hfq_source, &mut gpus.devices, &layout) From 8e4f1daa92439f74b285657dd1b09f623107651d Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 13:48:31 +0200 Subject: [PATCH 05/25] style(device-mesh): satisfy current hardware lint --- crates/hipfire-hardware/src/lib.rs | 45 ++++++++++++++---------------- 1 file changed, 21 insertions(+), 24 deletions(-) diff --git a/crates/hipfire-hardware/src/lib.rs b/crates/hipfire-hardware/src/lib.rs index 037a582cbc..1217a322e9 100644 --- a/crates/hipfire-hardware/src/lib.rs +++ b/crates/hipfire-hardware/src/lib.rs @@ -21,9 +21,8 @@ //! 3. Pass the multi-GPU coherence gate. use hip_bridge::{ - DeviceBuffer, Event, HipError, HipResult, RcclComms, - HIP_ERROR_PEER_ACCESS_ALREADY_ENABLED, HIP_ERROR_PEER_ACCESS_UNSUPPORTED, - HIP_EVENT_DISABLE_TIMING, HIP_EVENT_RELEASE_TO_SYSTEM, + DeviceBuffer, Event, HipError, HipResult, RcclComms, HIP_ERROR_PEER_ACCESS_ALREADY_ENABLED, + HIP_ERROR_PEER_ACCESS_UNSUPPORTED, HIP_EVENT_DISABLE_TIMING, HIP_EVENT_RELEASE_TO_SYSTEM, }; use rdna_compute::{DType, Gpu, GpuTensor}; @@ -297,11 +296,7 @@ impl Gpus { /// only validates the device count. Pre-flight runs the arch-match + /// VRAM-delta gate (TP ranks are identical cards, so the uniform delta /// check applies). - pub fn init_tp( - opts: &DeviceResolveOpts, - tp_size: usize, - n_layers: usize, - ) -> HipResult { + pub fn init_tp(opts: &DeviceResolveOpts, tp_size: usize, n_layers: usize) -> HipResult { if tp_size == 0 { return Err(HipError::new(0, "init_tp: tp_size must be >= 1")); } @@ -833,16 +828,15 @@ impl Gpus { for rank in 0..n { devices[rank].tp_graph_signal_store_gfx1201(&signals[rank], epoch)?; } - for destination in 0..n { + for (destination, device) in devices.iter_mut().enumerate() { let peers: Vec<&DeviceBuffer> = (0..n) .filter(|&source| source != destination) .map(|source| &signals[source]) .collect(); if n == 3 { - devices[destination].tp_graph_signal_wait2_gfx1201([peers[0], peers[1]], epoch)?; + device.tp_graph_signal_wait2_gfx1201([peers[0], peers[1]], epoch)?; } else { - devices[destination] - .tp_graph_signal_wait3_gfx1201([peers[0], peers[1], peers[2]], epoch)?; + device.tp_graph_signal_wait3_gfx1201([peers[0], peers[1], peers[2]], epoch)?; } } self.tp_graph_capture_epoch += 1; @@ -1058,8 +1052,13 @@ impl Gpus { } // D2H: synchronize producer stream, then download into row's prefix. // Preserve first error immediately — no partial-success claim. - for r in 0..n { - let dev = &self.devices[r]; + for (r, ((dev, buffer), row)) in self + .devices + .iter() + .zip(buffers.iter()) + .zip(self.host_ar_tmp.iter_mut()) + .enumerate() + { dev.bind_thread()?; let stream = dev.active_stream.as_ref().ok_or_else(|| { HipError::new( @@ -1075,10 +1074,10 @@ impl Gpus { // (bytes) lie within the Vec's initialized length. The derived // `[u8]` slice is exactly `bytes` and is bounded to that prefix. let dst: &mut [u8] = unsafe { - let ptr = self.host_ar_tmp[r].as_mut_ptr() as *mut u8; + let ptr = row.as_mut_ptr() as *mut u8; std::slice::from_raw_parts_mut(ptr, bytes) }; - dev.hip.memcpy_dtoh(dst, buffers[r])?; + dev.hip.memcpy_dtoh(dst, buffer)?; } // Exact left-associated reduction into row 0: rank0+rank1+... Self::host_reduce_rows(&mut self.host_ar_tmp, count); @@ -1088,10 +1087,9 @@ impl Gpus { let ptr = self.host_ar_tmp[0].as_ptr() as *const u8; std::slice::from_raw_parts(ptr, bytes) }; - for r in 0..n { - let dev = &self.devices[r]; + for (dev, buffer) in self.devices.iter().zip(buffers.iter()) { dev.bind_thread()?; - dev.hip.memcpy_htod(buffers[r], src)?; + dev.hip.memcpy_htod(buffer, src)?; } Ok(()) } @@ -1685,11 +1683,11 @@ impl Gpus { self.ensure_peer_ar_tmp(bytes)?; let mut gather_events = Vec::with_capacity(n - 1); - for rank in 1..n { + for (rank, partial) in partials.iter().enumerate().take(n).skip(1) { gather_events.push(self.boundary_copy( rank, 0, - partials[rank], + partial, &self.peer_ar_tmp[0][rank - 1], bytes, )?); @@ -1835,11 +1833,11 @@ impl Gpus { } // Gather N-1 peers into rank-0 lease scratch, never allocating. let mut gather_events = Vec::with_capacity(n - 1); - for rank in 1..n { + for (rank, buffer) in buffers.iter().enumerate().take(n).skip(1) { gather_events.push(self.boundary_copy( rank, 0, - buffers[rank], + buffer, &self.peer_lease_buffers[0][rank - 1], bytes, )?); @@ -2146,4 +2144,3 @@ mod tests { assert_eq!(resolve_device_ids(2, &opts).unwrap(), vec![0, 1]); } } - From b9786d883512ebb45ed3753a6a97d8d5954eff45 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 17:18:17 +0200 Subject: [PATCH 06/25] feat(device-mesh): add pure manifest and llama store pilot --- crates/hipfire-arch-llama/Cargo.toml | 1 + crates/hipfire-arch-llama/src/arch.rs | 154 ++- crates/hipfire-arch-llama/src/arch_model.rs | 19 +- crates/hipfire-arch-llama/src/carrier.rs | 134 ++- crates/hipfire-runtime/src/lib.rs | 2 + crates/hipfire-runtime/src/model_load.rs | 83 ++ crates/hipfire-runtime/src/weight_manifest.rs | 1004 +++++++++++++++++ crates/hipfire-runtime/src/weight_store.rs | 724 ++++++++++++ 8 files changed, 2068 insertions(+), 53 deletions(-) create mode 100644 crates/hipfire-runtime/src/weight_manifest.rs create mode 100644 crates/hipfire-runtime/src/weight_store.rs diff --git a/crates/hipfire-arch-llama/Cargo.toml b/crates/hipfire-arch-llama/Cargo.toml index f063bdf259..265db1f2aa 100644 --- a/crates/hipfire-arch-llama/Cargo.toml +++ b/crates/hipfire-arch-llama/Cargo.toml @@ -11,6 +11,7 @@ lab = [] [dependencies] hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-dispatch = { path = "../hipfire-dispatch", features = ["from-hip-error"] } hip-bridge = { path = "../hip-bridge" } rdna-compute = { path = "../rdna-compute" } diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index ec7322dbf2..920601367c 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -19,7 +19,10 @@ use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::{self, HfqFile}; use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; -use rdna_compute::Gpu; +use hipfire_runtime::weight_manifest::{ + FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, +}; +use rdna_compute::{DType, Gpu}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::{execute_steps, GemvInput, Step}; @@ -43,12 +46,8 @@ impl Architecture for Llama { fn arch_id() -> u32 { // `arch_id = 0` is the canonical LLaMA-family marker. The - // actual arch_id loaded at runtime is on `HfqFile::arch_id` - // and is either 0 (LLaMA / Mistral) or 1 (plain Qwen3 / - // Qwen2); both share this trait impl. The qwen3-norm flag - // is read off the HFQ metadata inside `config_from_hfq`, - // so the bring-up triple does not need a separate marker - // type per arch_id. + // actual id loaded at runtime is on `HfqFile::arch_id` and may + // differ for plain Qwen3/Qwen2; config parsing resolves that. 0 } @@ -57,13 +56,6 @@ impl Architecture for Llama { } fn config_from_hfq(hfq: &HfqFile) -> Result { - // `hfq::config_from_hfq` is the LLaMA-family HFQ metadata - // parser — emits a `LlamaConfig` with the appropriate - // `ModelArch` (Llama vs Qwen3) tag. It lives in the runtime - // crate because the qwen35 hybrid path's pflash drafter also - // calls it via `hfq::config_from_hfq` for its "Plain" - // variant. See arch-llama/src/lib.rs for the colocation - // rationale. hfq::config_from_hfq(hfq) } @@ -72,27 +64,141 @@ impl Architecture for Llama { cfg: &Self::Config, gpu: &mut Gpu, ) -> Result { - // `hfq::load_weights_hfq` is the LLaMA-family HFQ tensor - // loader. Same colocation reasoning as `config_from_hfq`. hfq::load_weights_hfq(hfq, cfg, gpu) .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}")) } fn new_state(gpu: &mut Gpu, cfg: &Self::Config) -> Result { - // The LLaMA-arch "state" is the `ForwardScratch` — persistent - // GPU scratch buffers reused across decode steps. There is no - // separate recurrent state (LLaMA is full-attention only). ForwardScratch::new(gpu, cfg) .map_err(|e| format!("llama: ForwardScratch::new failed: {e:?}")) } // Optional overrides: defaults from `hipfire_runtime::arch` already // assume Qwen3.5 family conventions. LLaMA / Mistral / Qwen3 don't - // emit `` blocks, but PR 11 keeps the override surface - // empty here on purpose — the daemon's existing per-`arch_id` - // policy choices stay unchanged. Future PRs that consolidate - // policy through the trait can populate these (LLaMA: no - // strip_think, no Qwen-specific blocked tokens). + // emit `` blocks, but the existing policy choices stay unchanged. +} + +impl Llama { + /// Pure dense LLaMA-family weight declaration. Source names remain + /// logical; carriers translate them to HFQ/safetensors namespaces. + pub fn weight_manifest(cfg: &LlamaConfig) -> Vec { + use ShardPolicy::*; + let (dim, hidden, head_dim) = (cfg.dim, cfg.hidden_dim, cfg.head_dim); + let (heads, kv_heads) = (cfg.n_heads, cfg.n_kv_heads); + let mut manifest = Vec::with_capacity(cfg.n_layers * 11 + 3); + manifest.push(WeightEntry::model( + "token_embd", + vec![cfg.vocab_size, dim], + DType::F16, + Pin(PinTarget::Embed), + )); + for layer in 0..cfg.n_layers { + manifest.push(WeightEntry::layer( + "wq", + layer, + vec![heads * head_dim, dim], + DType::F16, + FusedQkv { + q_heads: heads, + kv_heads, + head_dim, + layout: FusedQkvLayout::Qkv, + }, + )); + manifest.push(WeightEntry::layer( + "wk", + layer, + vec![kv_heads * head_dim, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "wv", + layer, + vec![kv_heads * head_dim, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "wo", + layer, + vec![dim, heads * head_dim], + DType::F16, + RowShard { axis: 1 }, + )); + manifest.push(WeightEntry::layer( + "ffn_gate", + layer, + vec![hidden, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "ffn_up", + layer, + vec![hidden, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "ffn_down", + layer, + vec![dim, hidden], + DType::F16, + RowShard { axis: 1 }, + )); + manifest.push(WeightEntry::layer( + "attn_norm", + layer, + vec![dim], + DType::F32, + Replicate, + )); + manifest.push(WeightEntry::layer( + "ffn_norm", + layer, + vec![dim], + DType::F32, + Replicate, + )); + if cfg.has_qk_norm { + manifest.push(WeightEntry::layer( + "q_norm", + layer, + vec![head_dim], + DType::F32, + Replicate, + )); + manifest.push(WeightEntry::layer( + "k_norm", + layer, + vec![head_dim], + DType::F32, + Replicate, + )); + } + } + manifest.push(WeightEntry::model( + "output_norm", + vec![dim], + DType::F32, + Replicate, + )); + manifest.push(WeightEntry::model( + "lm_head", + vec![cfg.vocab_size, dim], + DType::F16, + Pin(PinTarget::Output), + )); + manifest + } + + /// Pure state declaration for the full-attention LLaMA family. + pub fn state_manifest(cfg: &LlamaConfig) -> Vec { + (0..cfg.n_layers) + .map(|layer| StateEntry::new(StateKind::Kv { quant: String::new() }, layer)) + .collect() + } } // ── Dispatch integration ───────────────────────────────────────── diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 008b88eaf3..7e7d3d7ce0 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -40,19 +40,22 @@ impl ArchModel for LlamaBundle { weights, scratch, kv, + manifest_plan: _, + weight_store, + mesh, dflash_extract_layers: _, dspark_weights: _, dspark_assets: _, } = *self; - // Mirror unload_model ModelState::Llama arm exactly (lib.rs:3041): - // b.scratch.free_gpu(gpu); - // b.weights.free_gpu(gpu); - // note(b.kv.free_gpu(gpu)…) - // Ordering matters: scratch → weights → kv. dspark sidecars (when - // present) are reclaimed via the speculator/spec scratch paths, not - // here — matching the current unload_model which also does not handle - // them in this arm. + // Mirror the existing unload ordering: scratch → store/weights → kv. + // A committed store is only released here, through the ArchModel owner; + // no store destructor or independent carrier free path exists. scratch.free_gpu(gpu); + if let Some(store) = weight_store { + if let Err((_, error)) = store.release_on_owner(&mesh, gpu) { + eprintln!("llama: refusing weight-store release: {error}"); + } + } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); } diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 736eeb8674..c41b802983 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -4,17 +4,30 @@ use crate::dspark_body::Qwen3DrafterAssets; use crate::Llama; +use hipfire_hardware::DeviceMesh; use hipfire_runtime::arch::Architecture; use hipfire_runtime::dspark_core::DsparkWeights; -use hipfire_runtime::llama::{ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; +use hipfire_runtime::llama::{ + ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LlamaConfig, LlamaWeights, +}; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan}; +use hipfire_runtime::weight_store::WeightStore; pub struct LlamaBundle { pub config: LlamaConfig, pub weights: LlamaWeights, pub scratch: ForwardScratch, pub kv: KvCache, + /// Pure declaration/placement plan captured at load time. The plan has no + /// GPU handles and is immutable after publication. + pub manifest_plan: ManifestPlan, + /// A pilot store is attached only after its handles are assembled under + /// this bundle. It is crate-visible so callers cannot create an independent + /// unload owner; `ArchModel::free_gpu` is the sole release path. + pub(crate) weight_store: Option, + pub(crate) mesh: DeviceMesh, /// Decoder-layer indices whose residual hidden states a hidden-conditioned /// drafter (DFlash / EAGLE) wants captured, ascending order. Empty = no /// capture (the `SpecTarget::dflash_extract_layers` default of `None`). The @@ -25,26 +38,34 @@ pub struct LlamaBundle { /// was found or speculation was disabled. Task-10 wires the speculator build. pub dspark_weights: Option, /// Loaded DSpark drafter body assets (5-layer dense-GQA transformer + - /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. + /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } /// Build the LLaMA GPU bundle from an HFQ or safetensors-directory source. /// -/// Verbatim relocation of the carrier's `(config, weights, kv, scratch)` -/// seam: HFQ via `Architecture` trait, Dir via ParoQuant loaders. Error -/// strings are byte-identical to the prior inline carrier block. +/// The source/config path remains architecture-owned. Once it resolves, the +/// carrier publishes a pure Single manifest plan. Every fallible GPU stage +/// explicitly releases earlier allocations before returning an error; no +/// implicit GPU-buffer destructor is introduced. pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { let (config, weights, kv, scratch) = match src { ModelSource::Hfq(mut hfq) => { - let config = ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; + let config = + ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; let weights = ::load_weights(&mut hfq, &config, ctx.gpu)?; hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - // Size scratch (flash-attention partials) for the runtime KV cap so the - // asym/flash attends, which index partials by ceil(physical_cap/128), don't - // overflow it (the trait `new_state` only knows the model's declared max). - let scratch = ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) - .map_err(|e| format!("llama: ForwardScratch::new_with_max_seq failed: {e:?}"))?; + // The plain LLaMA path has no independent cap resolver. PR #661's + // physical-cap behavior is owned by the existing upstream KV plan. + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + )); + } + }; let dims = KvDims { layers: KvLayers::Flat(config.n_layers), n_kv_heads: config.n_kv_heads, @@ -52,7 +73,7 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result::from_mode( + let kv = match ::from_mode( hipfire_runtime::kv_mode::resolve( ctx.kv_mode_override.unwrap_or(""), &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, @@ -61,8 +82,16 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result::from_mode failed: {e}"))?; + ) { + Ok(kv) => kv, + Err(error) => { + scratch.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ::from_mode failed: {error}" + )); + } + }; (config, weights, kv, scratch) } ModelSource::Dir(source) => { @@ -74,7 +103,6 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Result Result::from_mode( + let kv = match ::from_mode( rr.mode, KvTarget::Single(ctx.gpu), &dims, - ) - .map_err(|e| format!("KvCache: {e}"))?; - let scratch = ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) - .map_err(|e| format!("ForwardScratch::new_with_max_seq: {e:?}"))?; + ) { + Ok(kv) => kv, + Err(error) => { + weights.free_gpu(ctx.gpu); + return Err(format!("KvCache: {error}")); + } + }; + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let _ = kv.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "ForwardScratch::new_with_max_seq: {error:?}" + )); + } + }; (config, weights, kv, scratch) } }; + + // Pure plan publication happens after source/config resolution and before + // the bundle becomes visible to the loader. It performs no GPU or file IO. + let mesh = DeviceMesh::single(); + let manifest = Llama::weight_manifest(&config); + let state = Llama::state_manifest(&config); + let manifest_plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) + .map_err(|e| format!("llama: manifest planning failed: {e}"))?; + Ok(LlamaBundle { config, weights, scratch, kv, + manifest_plan, + weight_store: None, + mesh, dflash_extract_layers: Vec::new(), dspark_weights: None, dspark_assets: None, @@ -121,6 +177,42 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Result<(), (WeightStore, String)> { + if self.weight_store.is_some() { + return Err((store, "llama: weight store already attached".into())); + } + let Some(origin) = store.origin() else { + return Err((store, "llama: weight store has no origin".into())); + }; + if origin.mesh_epoch() != self.mesh.epoch() { + return Err(( + store, + format!( + "llama: weight store origin epoch {:?} does not match bundle epoch {:?}", + origin.mesh_epoch(), + self.mesh.epoch() + ), + )); + } + self.weight_store = Some(store); + Ok(()) + } + + /// The immutable mesh identity used by this bundle's manifest plan. + /// Callers that run the Single pilot must pass this exact mesh to + /// `fulfill_manifest`; constructing a fresh `DeviceMesh::single()` would + /// intentionally fail the origin check. + pub fn manifest_mesh(&self) -> &DeviceMesh { + &self.mesh + } + /// Set the decoder-layer indices whose residual hidden states the /// hidden-conditioned drafter wants captured (ascending order). The /// speculator calls this with `dflash::DflashConfig::target_layer_ids`. diff --git a/crates/hipfire-runtime/src/lib.rs b/crates/hipfire-runtime/src/lib.rs index ebf81aff3f..b8c833190d 100644 --- a/crates/hipfire-runtime/src/lib.rs +++ b/crates/hipfire-runtime/src/lib.rs @@ -74,5 +74,7 @@ pub mod tokenizer; pub mod calibration; pub mod tool_call; pub mod weight_backend; +pub mod weight_manifest; +pub mod weight_store; pub use crate::arch::{maybe_screen_mmq, screen_weight_tensor, MmqScreenable}; diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index dadc63f8b6..46fcad6643 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -6,6 +6,7 @@ //! per-tensor dequant), which `WeightSource::read_layer` calls internally. use crate::llama::{EmbeddingFormat, WeightTensor}; +use hipfire_hardware::{DeviceMesh, DimKind}; use hip_bridge::HipResult; use hipfire_hardware::Gpus; use rdna_compute::{Gpu, GpuTensor}; @@ -30,6 +31,58 @@ impl Layout { layer_to_device: (0..n_layers).map(|i| g.device_for_layer(i)).collect(), } } + + /// Build the canonical stage/rank-0 view from an admitted mesh. The + /// manifest planner owns the full stage grid; this legacy loader view + /// selects rank zero for each layer so existing orchestrators continue to + /// have one deterministic device index until their typed mesh path lands. + pub fn from_mesh(mesh: &DeviceMesh, n_layers: usize) -> Self { + let mut output_coord = mesh.coord_of(0); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + output_coord[index] = mesh.size_of(DimKind::Pp).saturating_sub(1); + } + let layer_to_device = (0..n_layers) + .map(|layer| { + let mut coord = mesh.coord_of(0); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + coord[index] = mesh.stage_for_layer(layer, n_layers); + } + mesh.device_of(&coord) + }) + .collect(); + Self { + output_device: mesh.device_of(&output_coord), + layer_to_device, + } + } + + /// Validate the pure layout before any source preparation or GPU upload. + pub fn validate(&self, n_devices: usize, n_layers: usize) -> Result<(), String> { + if self.output_device >= n_devices { + return Err(format!( + "layout output device {} outside device count {}", + self.output_device, n_devices + )); + } + if self.layer_to_device.len() != n_layers { + return Err(format!( + "layout has {} layer assignments, expected {n_layers}", + self.layer_to_device.len() + )); + } + if let Some((layer, &device)) = self + .layer_to_device + .iter() + .enumerate() + .find(|(_, &device)| device >= n_devices) + { + return Err(format!( + "layout layer {layer} device {device} outside device count {n_devices}" + )); + } + Ok(()) + } + pub fn device_for_layer(&self, i: usize) -> usize { self.layer_to_device[i] } @@ -80,6 +133,15 @@ pub fn load_weights( devices: &mut [Gpu], layout: &Layout, ) -> HipResult> { + if devices.is_empty() { + return Err(hip_bridge::HipError::new( + 0, + "load_weights: at least one device is required", + )); + } + layout + .validate(devices.len(), source.n_layers()) + .map_err(|reason| hip_bridge::HipError::new(0, &reason))?; source.prepare(devices.len())?; let out_dev = layout.output_device(); let can_alias = devices.len() == 1; @@ -114,4 +176,25 @@ mod tests { assert_eq!(l.device_for_layer(i), 0); } } + + #[test] + fn mesh_layout_selects_stage_rank_zero_without_io() { + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let layout = Layout::from_mesh(&mesh, 4); + assert_eq!(layout.output_device(), 2); + assert_eq!( + (0..4) + .map(|layer| layout.device_for_layer(layer)) + .collect::>(), + vec![0, 0, 2, 2] + ); + assert!(layout.validate(mesh.n_devices(), 4).is_ok()); + } + + #[test] + fn invalid_layout_is_rejected_before_source_work() { + let layout = Layout::single(2); + assert!(layout.validate(0, 2).is_err()); + assert!(layout.validate(1, 3).is_err()); + } } diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs new file mode 100644 index 0000000000..530421db3e --- /dev/null +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -0,0 +1,1004 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Pure logical model declarations and device-mesh planning. +//! +//! A manifest describes *what* an architecture needs. [`plan_manifest`] resolves +//! those declarations against one already-admitted rectangular +//! [`hipfire_hardware::DeviceMesh`] and describes *where* each declaration and +//! synchronization point belongs. This module deliberately has no GPU, file, +//! carrier, quantizer, or allocation dependency; fulfillment is separate. +//! +//! The manifest is the single source of truth for collectives. A row-sharded +//! projection contributes one ordered tensor collective over `Tp`, an +//! expert-sharded projection contributes one over `Ep`, and pipeline boundaries +//! come from the mesh. Executors consume this schedule rather than add +//! family-local reductions. + +use crate::tp_shard::ExpertAssign; +use hipfire_hardware::{CollectiveHint, DeviceMesh, DimKind}; +use rdna_compute::DType; +use std::collections::HashSet; + +/// Derive the collective required by one weight policy. +/// +/// The returned hint is per declared operation. Two different row-sharded +/// operations in one layer are two distinct schedule entries and both execute +/// once. +#[inline] +pub fn collective_for_policy(policy: &ShardPolicy) -> Option { + match policy { + ShardPolicy::RowShard { .. } => Some(CollectiveHint::AllReduce { kind: DimKind::Tp }), + ShardPolicy::ExpertSharded { .. } => { + Some(CollectiveHint::AllReduce { kind: DimKind::Ep }) + } + ShardPolicy::ExpertTensorSharded { inner, .. } => collective_for_policy(inner), + _ => None, + } +} + +/// Non-layer placement targets resolved from mesh stage coordinates. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum PinTarget { + /// Token embedding, pinned to pipeline stage zero. + Embed, + /// Final norm/language head, pinned to the final pipeline stage. + Output, +} + +/// Optional placement override. It is separate from [`ShardPolicy`] so a tied +/// logical identity can be materialized at an output stage without changing +/// the source declaration. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +pub enum PlacementHint { + /// Resolve placement from the policy and layer scope. + #[default] + Policy, + /// Resolve placement from a mesh-derived pin target. + Pin(PinTarget), +} + +/// Source dtype acceptance. The logical manifest dtype remains an architecture +/// expectation; fulfillment preserves the source dtype and never silently +/// converts representation. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum SourceDType { + Any, + Exact(DType), + OneOf(Vec), +} + +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct DTypeConstraint { + pub source: SourceDType, +} + +impl DTypeConstraint { + pub fn any_source() -> Self { + Self { + source: SourceDType::Any, + } + } + + pub fn source_exact(dtype: DType) -> Self { + Self { + source: SourceDType::Exact(dtype), + } + } + + pub fn source_from_sources(sources: Vec) -> Self { + Self { + source: SourceDType::OneOf(sources), + } + } + + pub fn accepts(&self, dtype: DType) -> bool { + match &self.source { + SourceDType::Any => true, + SourceDType::Exact(expected) => *expected == dtype, + SourceDType::OneOf(allowed) => allowed.contains(&dtype), + } + } +} + +/// The block ordering of a fused projection. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum FusedQkvLayout { + /// `[Q | K | V]`. + Qkv, + /// `[Q | gate]`. + QGate, + /// `[Q | K | V | Z]`. + QkvZ, +} + +/// How one logical tensor is projected onto mesh devices. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum ShardPolicy { + /// A complete tensor on every device in the owning compute grid. + Replicate, + /// Split the output dimension `axis` across `Tp`. + ColumnShard { axis: usize }, + /// Split the input dimension `axis` across `Tp`; the consumer reduces. + RowShard { axis: usize }, + /// Assign complete expert tensors across `Ep` ranks. + ExpertSharded { + n_experts: usize, + assign: ExpertAssign, + }, + /// Fused QKV projection with head-aware block boundaries. + FusedQkv { + q_heads: usize, + kv_heads: usize, + head_dim: usize, + layout: FusedQkvLayout, + }, + /// Per-head projection (DeltaNet state/projections). + HeadSharded { n_heads: usize, head_dim: usize }, + /// Alias another logical source in the same manifest scope. + Tied { source: String }, + /// Pin to a mesh-derived non-layer stage. + Pin(PinTarget), + /// Split vocabulary rows across `Tp`. + VocabShard { axis: usize }, + /// Split each expert tensor across `Tp`. The inner policy is normally + /// `ColumnShard { axis: 1 }` for gate/up or `RowShard { axis: 2 }` for down. + ExpertTensorSharded { + n_experts: usize, + inner: Box, + }, +} + +/// A logical weight declaration. No source filename or GPU handle belongs +/// here; architecture carriers resolve those at fulfillment time. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct WeightEntry { + pub name: String, + pub layer: Option, + pub logical_shape: Vec, + pub dtype: DType, + pub dtype_constraint: DTypeConstraint, + pub placement: PlacementHint, + pub policy: ShardPolicy, +} + +impl WeightEntry { + pub fn model( + name: impl Into, + logical_shape: Vec, + dtype: DType, + policy: ShardPolicy, + ) -> Self { + Self::model_with_dtype_constraint( + name, + logical_shape, + dtype, + DTypeConstraint::any_source(), + policy, + ) + } + + pub fn model_with_dtype_constraint( + name: impl Into, + logical_shape: Vec, + dtype: DType, + dtype_constraint: DTypeConstraint, + policy: ShardPolicy, + ) -> Self { + Self { + name: name.into(), + layer: None, + logical_shape, + dtype, + dtype_constraint, + placement: PlacementHint::Policy, + policy, + } + } + + pub fn layer( + name: impl Into, + layer: usize, + logical_shape: Vec, + dtype: DType, + policy: ShardPolicy, + ) -> Self { + Self::layer_with_dtype_constraint( + name, + layer, + logical_shape, + dtype, + DTypeConstraint::any_source(), + policy, + ) + } + + pub fn layer_with_dtype_constraint( + name: impl Into, + layer: usize, + logical_shape: Vec, + dtype: DType, + dtype_constraint: DTypeConstraint, + policy: ShardPolicy, + ) -> Self { + Self { + name: name.into(), + layer: Some(layer), + logical_shape, + dtype, + dtype_constraint, + placement: PlacementHint::Policy, + policy, + } + } + + pub fn with_placement(mut self, placement: PlacementHint) -> Self { + self.placement = placement; + self + } + + /// Stable identity used by source resolvers and store keys. + pub fn identity(&self) -> (&str, Option) { + (&self.name, self.layer) + } +} + +/// Per-layer state declaration. Actual cache representation remains in the +/// architecture/model owner; this records logical placement scope only. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +pub enum StateKind { + Kv { quant: String }, + Recurrent, + Conv, +} + +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct StateEntry { + pub kind: StateKind, + pub layer: usize, +} + +impl StateEntry { + pub fn new(kind: StateKind, layer: usize) -> Self { + Self { kind, layer } + } +} + +/// One fully resolved weight placement. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct WeightPlacement { + pub name: String, + pub layer: Option, + pub devices: Vec, +} + +/// One ordered collective implied by one manifest operation. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct CollectiveScheduleEntry { + pub name: String, + pub layer: usize, + pub hint: CollectiveHint, +} + +/// Complete pure compilation of declarations against a mesh. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct ManifestPlan { + pub weights: Vec, + /// State and the global devices on which that state is resident. + pub state: Vec<(StateEntry, Vec)>, + /// Ordered `(layer, hint)` schedule retained for executor integration. + pub layer_collectives: Vec<(usize, CollectiveHint)>, + /// Named schedule entries, allowing an executor to prove no operation was + /// silently omitted or scheduled twice. + pub collective_schedule: Vec, + /// PP boundary hints in ascending after-layer order. + pub band_xfers: Vec<(usize, CollectiveHint)>, +} + +fn base_coord_for(entry: &WeightEntry, mesh: &DeviceMesh, n_layers: usize) -> Vec { + let stage = match (entry.placement, &entry.policy, entry.layer) { + (PlacementHint::Pin(PinTarget::Embed), _, _) + | (PlacementHint::Policy, ShardPolicy::Pin(PinTarget::Embed), _) => 0, + (PlacementHint::Pin(PinTarget::Output), _, _) + | (PlacementHint::Policy, ShardPolicy::Pin(PinTarget::Output), _) => { + mesh.size_of(DimKind::Pp).saturating_sub(1) + } + (PlacementHint::Policy, _, Some(layer)) => mesh.stage_for_layer(layer, n_layers), + (PlacementHint::Policy, _, None) => 0, + }; + let mut coord = mesh.coord_of(0); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + coord[index] = stage; + } + coord +} + +/// Compute global placement without touching a source, GPU, or allocator. +pub fn placement_devices(entry: &WeightEntry, mesh: &DeviceMesh, n_layers: usize) -> Vec { + let coord = base_coord_for(entry, mesh, n_layers); + match &entry.policy { + ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } => vec![mesh.device_of(&coord)], + ShardPolicy::ExpertSharded { .. } => mesh.group_along(DimKind::Ep, &coord), + ShardPolicy::ExpertTensorSharded { .. } => mesh.group_along(DimKind::Tp, &coord), + _ => mesh.stage_devices(&coord), + } +} + +/// Ordered per-operation collective schedule. This deliberately does not +/// deduplicate by `(layer, kind)`: two distinct row-sharded projections in one +/// layer represent two distinct output points and each must reduce once. +pub fn collective_schedule(manifest: &[WeightEntry]) -> Vec { + manifest + .iter() + .filter_map(|entry| { + let layer = entry.layer?; + let hint = collective_for_policy(&entry.policy)?; + Some(CollectiveScheduleEntry { + name: entry.name.clone(), + layer, + hint, + }) + }) + .collect() +} + +/// Compact schedule view consumed by executor adapters. +pub fn layer_collectives(manifest: &[WeightEntry]) -> Vec<(usize, CollectiveHint)> { + collective_schedule(manifest) + .into_iter() + .map(|entry| (entry.layer, entry.hint)) + .collect() +} + +fn validate_shape(entry: &WeightEntry) -> Result<(), String> { + if entry.name.is_empty() { + return Err("manifest entry has an empty name".to_string()); + } + if entry.logical_shape.is_empty() || entry.logical_shape.iter().any(|&d| d == 0) { + return Err(format!( + "{}[layer {:?}]: logical_shape {:?} must be non-empty", + entry.name, entry.layer, entry.logical_shape + )); + } + Ok(()) +} + +/// Validate logical shard math and tied source identity before fulfillment. +pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result<(), String> { + let mut identities = HashSet::new(); + for entry in manifest { + validate_shape(entry)?; + if !identities.insert(entry.identity()) { + return Err(format!( + "duplicate manifest identity ('{}', {:?})", + entry.name, entry.layer + )); + } + } + + let tp = mesh.size_of(DimKind::Tp); + for entry in manifest { + let context = format!("{}[layer {:?}]", entry.name, entry.layer); + match &entry.policy { + ShardPolicy::ColumnShard { axis } + | ShardPolicy::RowShard { axis } + | ShardPolicy::VocabShard { axis } => { + let dim = entry.logical_shape.get(*axis).ok_or_else(|| { + format!("{context}: shard axis {axis} outside logical shape") + })?; + if tp > 1 && dim % tp != 0 { + return Err(format!( + "{context}: shard dim {dim} (axis {axis}) not divisible by Tp={tp}" + )); + } + } + ShardPolicy::FusedQkv { + q_heads, + kv_heads, + head_dim, + .. + } => { + if *q_heads == 0 || *kv_heads == 0 || *head_dim == 0 { + return Err(format!("{context}: fused QKV geometry must be non-zero")); + } + if tp > 1 && (q_heads % tp != 0 || kv_heads % tp != 0) { + return Err(format!( + "{context}: q_heads={q_heads}/kv_heads={kv_heads} not divisible by Tp={tp}" + )); + } + } + ShardPolicy::HeadSharded { n_heads, head_dim } => { + if *n_heads == 0 || *head_dim == 0 { + return Err(format!("{context}: head geometry must be non-zero")); + } + if tp > 1 && n_heads % tp != 0 { + return Err(format!( + "{context}: n_heads={n_heads} not divisible by Tp={tp}" + )); + } + } + ShardPolicy::Tied { source } => { + if source.is_empty() { + return Err(format!("{context}: tied source is empty")); + } + let source_entry = manifest + .iter() + .find(|candidate| candidate.name == *source && candidate.layer == entry.layer) + .ok_or_else(|| { + format!("{context}: Tied source '{source}' has no manifest entry in scope") + })?; + if source_entry.identity() == entry.identity() { + return Err(format!("{context}: an entry cannot tie to itself")); + } + } + ShardPolicy::ExpertSharded { n_experts, .. } => { + if *n_experts == 0 || entry.logical_shape.first() != Some(n_experts) { + return Err(format!( + "{context}: logical_shape {:?} first dimension must equal n_experts={n_experts}", + entry.logical_shape + )); + } + } + ShardPolicy::ExpertTensorSharded { n_experts, inner } => { + if *n_experts == 0 || entry.logical_shape.first() != Some(n_experts) { + return Err(format!( + "{context}: ExpertTensorSharded shape {:?} must start with n_experts={n_experts}", + entry.logical_shape + )); + } + let axis = match inner.as_ref() { + ShardPolicy::ColumnShard { axis: 1 } + | ShardPolicy::RowShard { axis: 2 } => match inner.as_ref() { + ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => *axis, + _ => unreachable!(), + }, + ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => { + return Err(format!( + "{context}: ExpertTensorSharded inner axis {axis} is incompatible with [expert, projection, hidden]" + )); + } + other => { + return Err(format!( + "{context}: ExpertTensorSharded inner policy {other:?} is unsupported" + )); + } + }; + let dim = entry.logical_shape.get(axis).copied().ok_or_else(|| { + format!("{context}: ExpertTensorSharded axis {axis} outside shape") + })?; + if tp > 1 && dim % tp != 0 { + return Err(format!( + "{context}: ExpertTensorSharded dim {dim} (axis {axis}) not divisible by Tp={tp}" + )); + } + } + ShardPolicy::Replicate | ShardPolicy::Pin(_) => {} + } + } + Ok(()) +} + +/// Compile declarations against a mesh. +pub fn plan_manifest( + weights: &[WeightEntry], + state: &[StateEntry], + mesh: &DeviceMesh, + n_layers: usize, +) -> Result { + validate_manifest(weights, mesh)?; + let mut state_ids = HashSet::new(); + for entry in state { + if entry.layer >= n_layers { + return Err(format!( + "state {:?} layer {} outside n_layers={n_layers}", + entry.kind, entry.layer + )); + } + if !state_ids.insert((&entry.kind, entry.layer)) { + return Err(format!( + "duplicate state declaration {:?}[layer {}]", + entry.kind, entry.layer + )); + } + } + let schedule = collective_schedule(weights); + let layer_collectives = schedule + .iter() + .map(|entry| (entry.layer, entry.hint)) + .collect(); + let weight_placements = weights + .iter() + .map(|entry| WeightPlacement { + name: entry.name.clone(), + layer: entry.layer, + devices: placement_devices(entry, mesh, n_layers), + }) + .collect(); + let state_placements = state + .iter() + .map(|entry| { + let mut coord = mesh.coord_of(0); + let stage = mesh.stage_for_layer(entry.layer, n_layers); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + coord[index] = stage; + } + (entry.clone(), mesh.stage_devices(&coord)) + }) + .collect(); + let band_xfers = (0..n_layers) + .filter_map(|layer| mesh.band_xfer_after(layer, n_layers).map(|hint| (layer, hint))) + .collect(); + Ok(ManifestPlan { + weights: weight_placements, + state: state_placements, + layer_collectives, + collective_schedule: schedule, + band_xfers, + }) +} + +// ── Logical expert source identity ───────────────────────────────────────── + +/// How one logical expert group is distributed. This declaration is consumed +/// by the G5 executor-owned sealed plan; no rank assignment is resolved here. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum ExpertParallelism { + Single, + TensorParallel, + ExpertParallel, +} + +/// Stable source identities for expert projections. These names are manifest +/// references, not on-disk paths; the carrier/source resolver owns translation +/// to an artifact namespace. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum ExpertSourceLayout { + PackedFused { + gate_up: String, + down: String, + sidecars: Vec, + }, + PackedSeparate { + gate: String, + up: String, + down: String, + sidecars: Vec, + }, + PerExpertFused { + gate_up: Vec, + down: Vec, + sidecars: Vec, + }, + PerExpertSeparate { + gate: Vec, + up: Vec, + down: Vec, + sidecars: Vec, + }, +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub struct ExpertResourceRequirements { + pub bytes_per_expert: usize, + pub alignment: usize, +} + +/// Architecture-declared identity and source description of one expert group. +/// G5 derives rank ownership and seals the executor plan from this value. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct ExpertGroupSpec { + pub group: String, + pub layer: Option, + pub n_experts: usize, + pub parallelism: ExpertParallelism, + pub assignment: ExpertAssign, + pub source_layout: ExpertSourceLayout, + pub resources: ExpertResourceRequirements, + pub router: String, + pub execution: String, +} + +fn expert_context(spec: &ExpertGroupSpec) -> String { + format!("expert group '{}' layer {:?}", spec.group, spec.layer) +} + +fn source_names(layout: &ExpertSourceLayout) -> Vec<(&'static str, Vec)> { + match layout { + ExpertSourceLayout::PackedFused { + gate_up, + down, + sidecars, + } => vec![ + ("gate_up", vec![gate_up.clone()]), + ("down", vec![down.clone()]), + ("sidecar", sidecars.clone()), + ], + ExpertSourceLayout::PackedSeparate { + gate, + up, + down, + sidecars, + } => vec![ + ("gate", vec![gate.clone()]), + ("up", vec![up.clone()]), + ("down", vec![down.clone()]), + ("sidecar", sidecars.clone()), + ], + ExpertSourceLayout::PerExpertFused { + gate_up, + down, + sidecars, + } => vec![ + ("gate_up", gate_up.clone()), + ("down", down.clone()), + ("sidecar", sidecars.clone()), + ], + ExpertSourceLayout::PerExpertSeparate { + gate, + up, + down, + sidecars, + } => vec![ + ("gate", gate.clone()), + ("up", up.clone()), + ("down", down.clone()), + ("sidecar", sidecars.clone()), + ], + } +} + +fn manifest_entry<'a>( + spec: &ExpertGroupSpec, + manifest: &'a [WeightEntry], + label: &str, + name: &str, +) -> Result<&'a WeightEntry, String> { + let context = expert_context(spec); + if name.is_empty() { + return Err(format!("{context}: {label} reference is empty")); + } + manifest + .iter() + .find(|entry| entry.name == name && entry.layer == spec.layer) + .ok_or_else(|| format!("{context}: {label} reference '{name}' not found")) +} + +fn source_policy_matches( + spec: &ExpertGroupSpec, + label: &str, + policy: &ShardPolicy, +) -> bool { + match spec.parallelism { + ExpertParallelism::Single => matches!( + policy, + ShardPolicy::Replicate + | ShardPolicy::Pin(_) + | ShardPolicy::Tied { .. } + ), + ExpertParallelism::TensorParallel => match (label, policy) { + ( + "gate_up" | "gate" | "up", + ShardPolicy::ExpertTensorSharded { n_experts, inner }, + ) => { + *n_experts == spec.n_experts + && matches!(inner.as_ref(), ShardPolicy::ColumnShard { axis: 1 }) + } + ("down", ShardPolicy::ExpertTensorSharded { n_experts, inner }) => { + *n_experts == spec.n_experts + && matches!(inner.as_ref(), ShardPolicy::RowShard { axis: 2 }) + } + ("sidecar", ShardPolicy::Replicate | ShardPolicy::Tied { .. }) => true, + _ => false, + }, + ExpertParallelism::ExpertParallel => match (label, policy) { + ( + "gate_up" | "gate" | "up" | "down", + ShardPolicy::ExpertSharded { n_experts, assign }, + ) => *n_experts == spec.n_experts && *assign == spec.assignment, + ("sidecar", ShardPolicy::Replicate | ShardPolicy::Tied { .. }) => true, + _ => false, + }, + } +} + +fn source_shape_matches( + spec: &ExpertGroupSpec, + label: &str, + per_expert: bool, + entry: &WeightEntry, +) -> Result<(), String> { + let context = expert_context(spec); + if !source_policy_matches(spec, label, &entry.policy) { + return Err(format!( + "{context}: {label} source '{}' has incompatible policy {:?}", + entry.name, entry.policy + )); + } + if entry.logical_shape.len() < 2 { + return Err(format!( + "{context}: {label} source '{}' shape {:?} is too short", + entry.name, entry.logical_shape + )); + } + if !per_expert && entry.logical_shape.first() != Some(&spec.n_experts) { + return Err(format!( + "{context}: {label} source '{}' shape {:?} must start in n_experts={}", + entry.name, entry.logical_shape, spec.n_experts + )); + } + Ok(()) +} + +fn validate_expert_sources( + spec: &ExpertGroupSpec, + manifest: &[WeightEntry], +) -> Result<(), String> { + let context = expert_context(spec); + let router = manifest_entry(spec, manifest, "router", &spec.router)?; + if !matches!(router.logical_shape.len(), 1 | 2) + || router.logical_shape.last() != Some(&spec.n_experts) + { + return Err(format!( + "{context}: router '{}' shape {:?} must end in n_experts={}", + router.name, router.logical_shape, spec.n_experts + )); + } + if !matches!( + router.policy, + ShardPolicy::Replicate | ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } + ) { + return Err(format!( + "{context}: router '{}' has incompatible policy {:?}", + router.name, router.policy + )); + } + let per_expert = matches!( + spec.source_layout, + ExpertSourceLayout::PerExpertFused { .. } | ExpertSourceLayout::PerExpertSeparate { .. } + ); + if per_expert && spec.parallelism != ExpertParallelism::Single { + return Err(format!( + "{context}: per-expert source layout is only admitted for Single" + )); + } + + for (label, names) in source_names(&spec.source_layout) { + if names.is_empty() { + continue; + } + if per_expert && label != "sidecar" && names.len() != spec.n_experts { + return Err(format!( + "{context}: {label} source count={} != n_experts={}", + names.len(), + spec.n_experts + )); + } + let mut seen = HashSet::new(); + let mut shape: Option> = None; + for (index, name) in names.iter().enumerate() { + if !seen.insert(name.as_str()) { + return Err(format!( + "{context}: duplicate {label} source '{name}' at index {index}" + )); + } + source_shape_matches(spec, label, per_expert, entry)?; + if per_expert { + if let Some(previous) = &shape { + if previous != &entry.logical_shape { + return Err(format!( + "{context}: {label}[{index}] shape {:?} differs from {:?}", + entry.logical_shape, previous + )); + } + } else { + shape = Some(entry.logical_shape.clone()); + } + } + } + } + Ok(()) +} + +/// Validate logical expert source identities. Rank assignment remains owned by +/// G5; this function only proves source names, shapes, and scope are coherent. +pub fn validate_expert_group_specs( + specs: &[ExpertGroupSpec], + manifest: &[WeightEntry], +) -> Result<(), String> { + let mut identities = HashSet::new(); + for entry in manifest { + if !identities.insert(entry.identity()) { + return Err(format!( + "duplicate manifest identity ('{}', {:?})", + entry.name, entry.layer + )); + } + } + let mut groups = HashSet::new(); + for spec in specs { + let context = expert_context(spec); + if spec.group.is_empty() || spec.router.is_empty() || spec.execution.is_empty() { + return Err(format!("{context}: group/router/execution identities must be non-empty")); + } + if spec.n_experts == 0 || spec.resources.bytes_per_expert == 0 { + return Err(format!("{context}: n_experts and bytes_per_expert must be non-zero")); + } + if spec.resources.alignment == 0 || !spec.resources.alignment.is_power_of_two() { + return Err(format!("{context}: alignment must be a non-zero power of two")); + } + if !groups.insert((&spec.group, spec.layer)) { + return Err(format!("{context}: duplicate group/layer identity")); + } + validate_expert_sources(spec, manifest)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn layer_entry(name: &str, layer: usize, policy: ShardPolicy) -> WeightEntry { + WeightEntry::layer(name, layer, vec![8, 8], DType::F16, policy) + } + + #[test] + fn placement_and_boundaries_use_named_mesh() { + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let embed = WeightEntry::model( + "token_embd", + vec![32, 8], + DType::F16, + ShardPolicy::Pin(PinTarget::Embed), + ); + let row = layer_entry("wo", 1, ShardPolicy::RowShard { axis: 1 }); + assert_eq!(placement_devices(&embed, &mesh, 4), vec![0]); + assert_eq!(placement_devices(&row, &mesh, 4), vec![2, 3]); + let plan = plan_manifest( + &[layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 }), row], + &[], + &mesh, + 4, + ) + .unwrap(); + assert_eq!(plan.layer_collectives.len(), 2); + assert_eq!( + plan.band_xfers, + vec![(1, CollectiveHint::BandXfer { src: 0, dst: 1 })] + ); + } + + #[test] + fn schedule_is_ordered_per_operation_not_deduped() { + let manifest = vec![ + layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 }), + layer_entry("down", 0, ShardPolicy::RowShard { axis: 1 }), + ]; + assert_eq!( + layer_collectives(&manifest), + vec![ + (0, CollectiveHint::AllReduce { kind: DimKind::Tp }), + (0, CollectiveHint::AllReduce { kind: DimKind::Tp }), + ] + ); + assert_eq!(collective_schedule(&manifest)[0].name, "wo"); + assert_eq!(collective_schedule(&manifest)[1].name, "down"); + } + + #[test] + fn validation_covers_divisibility_ties_and_expert_shape() { + let tp3 = DeviceMesh::rect(&[(DimKind::Tp, 3)]); + assert!(validate_manifest( + &[layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 })], + &tp3 + ) + .is_err()); + let tied = vec![ + WeightEntry::model( + "embed", + vec![8, 8], + DType::F16, + ShardPolicy::Pin(PinTarget::Embed), + ), + WeightEntry::model( + "lm_head", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "embed".into(), + }, + ), + ]; + assert!(validate_manifest(&tied, &DeviceMesh::single()).is_ok()); + let bad_expert = WeightEntry::layer( + "experts", + 0, + vec![3, 8, 8], + DType::F16, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ); + assert!(validate_manifest(&[bad_expert], &DeviceMesh::single()).is_err()); + } + + #[test] + fn expert_source_identity_and_shape_are_checked() { + let manifest = vec![ + WeightEntry::layer( + "router", + 0, + vec![8, 4], + DType::F16, + ShardPolicy::Replicate, + ), + WeightEntry::layer( + "gate_up", + 0, + vec![4, 8, 8], + DType::F16, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ), + WeightEntry::layer( + "down", + 0, + vec![4, 8, 8], + DType::F16, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ), + ]; + let spec = ExpertGroupSpec { + group: "ffn".into(), + layer: Some(0), + n_experts: 4, + parallelism: ExpertParallelism::ExpertParallel, + assignment: ExpertAssign::Stride, + source_layout: ExpertSourceLayout::PackedFused { + gate_up: "gate_up".into(), + down: "down".into(), + sidecars: Vec::new(), + }, + resources: ExpertResourceRequirements { + bytes_per_expert: 1024, + alignment: 256, + }, + router: "router".into(), + execution: "moe.ffn".into(), + }; + assert!(validate_expert_group_specs(&[spec], &manifest).is_ok()); + let bad = ExpertGroupSpec { + source_layout: ExpertSourceLayout::PackedFused { + gate_up: "missing".into(), + down: "down".into(), + sidecars: Vec::new(), + }, + ..ExpertGroupSpec { + group: "ffn2".into(), + layer: Some(0), + n_experts: 4, + parallelism: ExpertParallelism::ExpertParallel, + assignment: ExpertAssign::Stride, + source_layout: ExpertSourceLayout::PackedFused { + gate_up: "gate_up".into(), + down: "down".into(), + sidecars: Vec::new(), + }, + resources: ExpertResourceRequirements { + bytes_per_expert: 1024, + alignment: 256, + }, + router: "router".into(), + execution: "moe.ffn".into(), + } + }; + assert!(validate_expert_group_specs(&[bad], &manifest).is_err()); + } +} diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs new file mode 100644 index 0000000000..a663369857 --- /dev/null +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -0,0 +1,724 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Transactional fulfillment for the pure weight manifest. +//! +//! [`crate::weight_manifest::plan_manifest`] owns the CPU-only "where". This +//! module owns the narrow "how" pilot for a plain LLaMA Single target: a +//! source callback supplies already-resolved bytes and dtype, the store uploads +//! them, and the first failure explicitly rolls back every resident buffer. +//! +//! The store is not a model owner. It has no `Drop` implementation and never +//! frees GPU buffers implicitly. A carrier may move a committed store into its +//! existing `ArchModel` owner; that owner must call [`WeightStore::release_on_owner`] +//! during its existing teardown path. `take` transfers a resident handle to the +//! owner that is assembling typed weights, and therefore removes the cell from +//! the store's cleanup set. + +use crate::weight_manifest::{placement_devices, ShardPolicy, WeightEntry}; +use hipfire_hardware::{DeviceMesh, MeshEpoch}; +use rdna_compute::{DType, Gpu, GpuTensor}; +use std::collections::HashMap; + +/// Stable logical placement identity. Layer is part of the key because a +/// per-layer name such as `wq` appears once for every decoder block. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +pub struct WeightPlacementKey { + pub name: String, + pub layer: Option, + pub device: usize, +} + +impl WeightPlacementKey { + pub fn new(name: impl Into, layer: Option, device: usize) -> Self { + Self { + name: name.into(), + layer, + device, + } + } +} + +/// The immutable projection applied to one logical source before upload. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum WeightProjectionKind { + Static, + ColumnShard, + RowShard, + FusedQkv, + HeadSharded, + VocabShard, + ExpertCompact, + ExpertTensor, +} + +/// Value-owned placement metadata. It contains no GPU or source-file +/// representation and remains stable after a handle is taken from the store. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct WeightProjection { + pub kind: WeightProjectionKind, + pub axis: Option, + pub rank: usize, + pub world_size: usize, + pub logical_shape: Vec, + pub dtype: DType, +} + +fn projection_for(entry: &WeightEntry, rank: usize, world_size: usize, dtype: DType) -> WeightProjection { + let (kind, axis) = match &entry.policy { + ShardPolicy::ColumnShard { axis } => (WeightProjectionKind::ColumnShard, Some(*axis)), + ShardPolicy::RowShard { axis } => (WeightProjectionKind::RowShard, Some(*axis)), + ShardPolicy::FusedQkv { .. } => (WeightProjectionKind::FusedQkv, None), + ShardPolicy::HeadSharded { .. } => (WeightProjectionKind::HeadSharded, None), + ShardPolicy::VocabShard { axis } => (WeightProjectionKind::VocabShard, Some(*axis)), + ShardPolicy::ExpertSharded { .. } => (WeightProjectionKind::ExpertCompact, None), + ShardPolicy::ExpertTensorSharded { .. } => (WeightProjectionKind::ExpertTensor, None), + ShardPolicy::Replicate | ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } => { + (WeightProjectionKind::Static, None) + } + }; + WeightProjection { + kind, + axis, + rank, + world_size, + logical_shape: entry.logical_shape.clone(), + dtype, + } +} + +/// A resident GPU tensor or a symbolic alias to another logical source. +/// +/// Aliases own no buffer. Resident buffers have no implicit destructor; the +/// current model owner explicitly consumes them through its teardown method. +pub enum WeightHandle { + Resident(GpuTensor), + Alias(String), +} + +/// Identity captured at the start of a load. It is deliberately immutable and +/// contains only mesh generation, logical rank, and physical device identity. +/// No policy or source representation is smuggled into the origin. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +pub struct WeightOrigin { + mesh_epoch: MeshEpoch, + logical_rank: usize, + physical_device: i32, +} + +impl WeightOrigin { + pub fn from_parts(mesh_epoch: MeshEpoch, logical_rank: usize, physical_device: i32) -> Self { + Self { + mesh_epoch, + logical_rank, + physical_device, + } + } + + pub fn for_single(mesh: &DeviceMesh, gpu: &Gpu) -> Self { + Self::from_parts(mesh.epoch(), 0, gpu.device_id) + } + + pub fn mesh_epoch(self) -> MeshEpoch { + self.mesh_epoch + } + + pub fn logical_rank(self) -> usize { + self.logical_rank + } + + pub fn physical_device(self) -> i32 { + self.physical_device + } +} + +/// Errors that are detected before a store is allowed to release a resident +/// buffer. Origin mismatch always returns the store to the caller unchanged. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum WeightStoreError { + OriginMismatch { + expected: WeightOrigin, + actual: WeightOrigin, + }, + UnboundOrigin, + DuplicatePlacement(WeightPlacementKey), + MissingPlacement(WeightPlacementKey), + InvalidTarget(String), +} + +impl std::fmt::Display for WeightStoreError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::OriginMismatch { expected, actual } => write!( + f, + "weight store origin mismatch: expected {:?}, got {:?}", + expected, actual + ), + Self::UnboundOrigin => write!(f, "weight store has no target origin"), + Self::DuplicatePlacement(key) => write!( + f, + "duplicate weight placement {}[layer {:?}] on device {}", + key.name, key.layer, key.device + ), + Self::MissingPlacement(key) => write!( + f, + "missing weight placement {}[layer {:?}] on device {}", + key.name, key.layer, key.device + ), + Self::InvalidTarget(message) => write!(f, "invalid weight store target: {message}"), + } + } +} + +impl std::error::Error for WeightStoreError {} + +/// Error identifying the first failed manifest cell. The store has already +/// been rolled back before this value is returned by [`fulfill_manifest`]. +#[derive(Debug)] +pub struct FulfillError { + pub name: String, + pub layer: Option, + pub device: usize, + pub reason: String, +} + +impl std::fmt::Display for FulfillError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "fulfill_manifest: {}[layer {:?}] on device {}: {}", + self.name, self.layer, self.device, self.reason + ) + } +} + +impl std::error::Error for FulfillError {} + +/// Load-side placement container. It records one immutable projection per +/// `(name, layer, device)` and captures the target origin once. +/// +/// There is intentionally no `Drop` implementation. A `WeightStore` that is +/// abandoned without explicit rollback/release leaks rather than guessing a +/// GPU owner; production callers keep it beneath `ArchModel`. +#[derive(Default)] +pub struct WeightStore { + placements: HashMap, + projections: HashMap, + origin: Option, +} + +impl WeightStore { + pub fn new() -> Self { + Self::default() + } + + pub fn with_origin(origin: WeightOrigin) -> Self { + Self { + placements: HashMap::new(), + projections: HashMap::new(), + origin: Some(origin), + } + } + + pub fn origin(&self) -> Option { + self.origin + } + + pub fn len(&self) -> usize { + self.placements.len() + } + + pub fn is_empty(&self) -> bool { + self.placements.is_empty() + } + + pub fn contains(&self, name: &str, layer: Option, device: usize) -> bool { + self.placements + .contains_key(&WeightPlacementKey::new(name, layer, device)) + } + + pub fn get(&self, name: &str, layer: Option, device: usize) -> Option<&WeightHandle> { + self.placements + .get(&WeightPlacementKey::new(name, layer, device)) + } + + pub fn projection( + &self, + name: &str, + layer: Option, + device: usize, + ) -> Option<&WeightProjection> { + self.projections + .get(&WeightPlacementKey::new(name, layer, device)) + } + + pub fn devices_for(&self, name: &str, layer: Option) -> Vec { + let mut devices: Vec<_> = self + .placements + .keys() + .filter(|key| key.name == name && key.layer == layer) + .map(|key| key.device) + .collect(); + devices.sort_unstable(); + devices + } + + fn insert( + &mut self, + key: WeightPlacementKey, + handle: WeightHandle, + projection: WeightProjection, + ) -> Result<(), WeightStoreError> { + if self.placements.contains_key(&key) { + return Err(WeightStoreError::DuplicatePlacement(key)); + } + self.placements.insert(key.clone(), handle); + self.projections.insert(key, projection); + Ok(()) + } + + /// Stage a symbolic alias without GPU work. Used for tied declarations and + /// CPU ownership tests; aliases never participate in release. + pub fn stage_alias( + &mut self, + name: impl Into, + layer: Option, + device: usize, + source: impl Into, + projection: WeightProjection, + ) -> Result<(), WeightStoreError> { + self.insert( + WeightPlacementKey::new(name, layer, device), + WeightHandle::Alias(source.into()), + projection, + ) + } + + /// Move a handle out of the store. The projection is removed with it so no + /// stale metadata can describe a cell that the store no longer owns. + pub fn take( + &mut self, + name: &str, + layer: Option, + device: usize, + ) -> Option { + let key = WeightPlacementKey::new(name, layer, device); + self.projections.remove(&key); + self.placements.remove(&key) + } + + pub fn take_with_projection( + &mut self, + name: &str, + layer: Option, + device: usize, + ) -> Option<(WeightHandle, WeightProjection)> { + let key = WeightPlacementKey::new(name, layer, device); + let handle = self.placements.remove(&key)?; + let projection = self.projections.remove(&key)?; + Some((handle, projection)) + } + + /// Start a typed assembly transaction. Handles moved through the + /// transaction are restored to this store if the transaction is dropped + /// before `finalize`; no GPU free or second owner is introduced. + pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { + WeightStoreAssembly { + store: self, + taken: Vec::new(), + committed: false, + } + } + + /// Compare a store's captured origin with an already-resolved target + /// identity. This pure seam is used by fault-path tests and by owner + /// teardown after target resolution. + pub fn validate_origin_value( + &self, + expected: WeightOrigin, + ) -> Result<(), WeightStoreError> { + let actual = self.origin.ok_or(WeightStoreError::UnboundOrigin)?; + if actual != expected { + return Err(WeightStoreError::OriginMismatch { expected, actual }); + } + Ok(()) + } + + /// Verify that this store is still being handled by the same mesh/device + /// target. No GPU calls occur on mismatch. + pub fn validate_origin( + &self, + mesh: &DeviceMesh, + gpu: &Gpu, + ) -> Result<(), WeightStoreError> { + self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) + } + + /// Explicit owner teardown. This method is intentionally consuming and has + /// no implicit/drop fallback; `ArchModel::free_gpu` is the production call + /// site. Callers must validate the target first with [`validate_origin`]. + /// On mismatch the original store is returned unchanged for retry by the + /// owner; no GPU call occurs. + pub fn release_on_owner( + self, + mesh: &DeviceMesh, + gpu: &mut Gpu, + ) -> Result<(), (Self, WeightStoreError)> { + if let Err(error) = self.validate_origin(mesh, gpu) { + return Err((self, error)); + } + self.release_unchecked(gpu); + Ok(()) + } + + /// Explicit rollback for a failed transaction. It consumes the partial + /// store and frees every resident buffer on the single owning GPU. + fn rollback(self, gpu: &Gpu) { + self.release_unchecked(gpu); + } + + fn release_unchecked(self, gpu: &Gpu) { + for handle in self.placements.into_values() { + if let WeightHandle::Resident(tensor) = handle { + // Rollback is deliberately direct and best-effort, matching + // the existing loader's explicit owner teardown. The store + // never relies on a destructor to release GPU memory. + let _ = gpu.hip.free(tensor.buf); + } + } + } +} + +/// One resident/alias handle temporarily moved during typed assembly. +pub struct TakenWeight { + pub key: WeightPlacementKey, + pub handle: WeightHandle, + pub projection: WeightProjection, +} + +/// Rollback-owning assembly transaction. Dropping it restores every taken cell +/// to the parent store; it never frees a GPU buffer implicitly. +pub struct WeightStoreAssembly<'a> { + store: &'a mut WeightStore, + taken: Vec, + committed: bool, +} + +impl<'a> WeightStoreAssembly<'a> { + pub fn take( + &mut self, + name: &str, + layer: Option, + device: usize, + ) -> Option { + let key = WeightPlacementKey::new(name, layer, device); + let (handle, projection) = self.store.take_with_projection(name, layer, device)?; + let slot = self.taken.len(); + self.taken.push(TakenWeight { + key, + handle, + projection, + }); + Some(slot) + } + + pub fn commit(self) -> WeightStoreAssemblyGuard<'a> { + WeightStoreAssemblyGuard { inner: self } + } +} + +impl Drop for WeightStoreAssembly<'_> { + fn drop(&mut self) { + if self.committed { + return; + } + for taken in self.taken.drain(..) { + let _ = self + .store + .insert(taken.key, taken.handle, taken.projection); + } + } +} + +/// Guard retained while the typed architecture object is being built. If it +/// is dropped before `finalize`, all handles return to the parent store. +pub struct WeightStoreAssemblyGuard<'a> { + inner: WeightStoreAssembly<'a>, +} + +impl WeightStoreAssemblyGuard<'_> { + pub fn get(&self, slot: usize) -> Option<&WeightHandle> { + self.inner.taken.get(slot).map(|taken| &taken.handle) + } + + pub fn projection(&self, slot: usize) -> Option<&WeightProjection> { + self.inner.taken.get(slot).map(|taken| &taken.projection) + } + + /// Transfer the taken handles to the existing ArchModel-owned typed + /// weights. This is the sole operation that removes them from rollback + /// ownership. + pub fn finalize(mut self) -> Vec { + self.inner.committed = true; + std::mem::take(&mut self.inner.taken) + } +} + +fn target_error(mesh: &DeviceMesh) -> Option { + (mesh.n_devices() != 1).then(|| FulfillError { + name: "".to_string(), + layer: None, + device: 0, + reason: format!( + "plain LLaMA Single fulfillment requires one logical device, got {}", + mesh.n_devices() + ), + }) +} + +/// Fulfill a manifest for a plain LLaMA Single target. +/// +/// The source callback is the architecture-owned namespace seam and returns +/// raw bytes plus the actual source dtype. No file/GGUF/HFQ type crosses this +/// API. On the first source, dtype, or upload failure every earlier resident is +/// explicitly released before the error is returned. +pub fn fulfill_manifest_single( + weights: &[WeightEntry], + mesh: &DeviceMesh, + n_layers: usize, + gpu: &Gpu, + source: F, +) -> Result +where + F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, +{ + if let Some(error) = target_error(mesh) { + return Err(error); + } + if let Err(reason) = crate::weight_manifest::validate_manifest(weights, mesh) { + return Err(FulfillError { + name: "".to_string(), + layer: None, + device: 0, + reason, + }); + } + + let origin = WeightOrigin::for_single(mesh, gpu); + let mut store = WeightStore::with_origin(origin); + for entry in weights { + let devices = placement_devices(entry, mesh, n_layers); + if devices.as_slice() != [0] { + let error = FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: devices.first().copied().unwrap_or(0), + reason: format!( + "Single placement resolved to {:?}, expected [0]", + devices + ), + }; + store.rollback(gpu); + return Err(error); + } + let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); + if let ShardPolicy::Tied { source: source_name } = &entry.policy { + let projection = projection_for(entry, 0, 1, entry.dtype); + if let Err(reason) = store.insert( + key, + WeightHandle::Alias(source_name.clone()), + projection, + ) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: reason.to_string(), + }); + } + continue; + } + + let (bytes, dtype) = match source(entry) { + Ok(value) => value, + Err(reason) => { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!("source read failed: {reason}"), + }); + } + }; + if !entry.dtype_constraint.accepts(dtype) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "source dtype {dtype:?} violates constraint {:?}", + entry.dtype_constraint + ), + }); + } + let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { + Ok(tensor) => tensor, + Err(error) => { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!("upload_raw failed: {error}"), + }); + } + }; + tensor.dtype = dtype; + let projection = projection_for(entry, 0, 1, dtype); + if let Err(reason) = store.insert(key, WeightHandle::Resident(tensor), projection) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: reason.to_string(), + }); + } + } + Ok(store) +} + +/// Canonical name used by the manifest fulfillment seam. The target is +/// deliberately Single-only in this pilot; multi-device fulfillment belongs to +/// the admitted mesh/G5 integration and must not grow a second owner here. +pub fn fulfill_manifest( + weights: &[WeightEntry], + mesh: &DeviceMesh, + n_layers: usize, + gpu: &Gpu, + source: F, +) -> Result +where + F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, +{ + fulfill_manifest_single(weights, mesh, n_layers, gpu, source) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::weight_manifest::{PinTarget, ShardPolicy}; + use hipfire_hardware::DimKind; + + fn projection(dtype: DType) -> WeightProjection { + WeightProjection { + kind: WeightProjectionKind::Static, + axis: None, + rank: 0, + world_size: 1, + logical_shape: vec![1], + dtype, + } + } + + #[test] + fn origin_mismatch_is_detected_before_gpu_release() { + let first = DeviceMesh::single(); + let second = DeviceMesh::single(); + let actual = WeightOrigin::from_parts(first.epoch(), 0, 0); + let expected = WeightOrigin::from_parts(second.epoch(), 0, 0); + let store = WeightStore::with_origin(actual); + let error = store.validate_origin_value(expected).unwrap_err(); + assert!(matches!( + error, + WeightStoreError::OriginMismatch { + expected: got_expected, + actual: got_actual + } if got_expected == expected && got_actual == actual + )); + } + + #[test] + fn staged_rollback_removes_handles_and_projection_together() { + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); + let mut store = WeightStore::with_origin(origin); + store + .stage_alias("first", None, 0, "source", projection(DType::F16)) + .unwrap(); + store + .stage_alias("second", Some(2), 0, "source", projection(DType::F16)) + .unwrap(); + assert_eq!(store.len(), 2); + let first = store.take_with_projection("first", None, 0).unwrap(); + assert!(matches!(first.0, WeightHandle::Alias(_))); + assert!(store.projection("first", None, 0).is_none()); + assert_eq!(store.len(), 1); + let second = store.take("second", Some(2), 0).unwrap(); + assert!(matches!(second, WeightHandle::Alias(_))); + assert!(store.is_empty()); + } + + #[test] + fn assembly_drop_restores_staged_handles() { + let mesh = DeviceMesh::single(); + let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); + store + .stage_alias("x", None, 0, "source", projection(DType::F16)) + .unwrap(); + { + let mut assembly = store.begin_assembly(); + assert_eq!(assembly.take("x", None, 0), Some(0)); + assert!(assembly.get(0).is_some()); + } + assert!(store.contains("x", None, 0)); + assert!(store.projection("x", None, 0).is_some()); + } + + #[test] + fn repeated_unload_lookup_cannot_reclaim_a_transferred_cell() { + let mesh = DeviceMesh::single(); + let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); + store + .stage_alias("x", None, 0, "source", projection(DType::F16)) + .unwrap(); + let _owned = store.take("x", None, 0).unwrap(); + assert!(store.take("x", None, 0).is_none()); + assert!(store.projection("x", None, 0).is_none()); + assert!(store.is_empty()); + } + + #[test] + fn duplicate_projection_is_rejected_without_replacing_identity() { + let mesh = DeviceMesh::single(); + let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); + store + .stage_alias("x", None, 0, "source-a", projection(DType::F16)) + .unwrap(); + let error = store + .stage_alias("x", None, 0, "source-b", projection(DType::F32)) + .unwrap_err(); + assert!(matches!(error, WeightStoreError::DuplicatePlacement(_))); + assert!(matches!(store.get("x", None, 0), Some(WeightHandle::Alias(source)) if source == "source-a")); + assert_eq!(store.projection("x", None, 0).unwrap().dtype, DType::F16); + } + + #[test] + fn single_target_refuses_multi_device_before_source_or_gpu_work() { + let mesh = DeviceMesh::rect(&[(DimKind::Tp, 2)]); + let entry = WeightEntry::model( + "embed", + vec![2, 2], + DType::F16, + ShardPolicy::Pin(PinTarget::Embed), + ); + // The target guard is pure and can be checked without constructing a + // Gpu; the closure would be unreachable on this path. + assert!(target_error(&mesh).is_some()); + assert_eq!(placement_devices(&entry, &mesh, 1), vec![0]); + } +} From f907ef7501ccaa401a50d00e8e9d6f42b67676df Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 17:49:21 +0200 Subject: [PATCH 07/25] fix(device-mesh): bind expert manifest source entries --- crates/hipfire-runtime/src/weight_manifest.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 530421db3e..26c64ad2ee 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -782,6 +782,7 @@ fn validate_expert_sources( "{context}: duplicate {label} source '{name}' at index {index}" )); } + let entry = manifest_entry(spec, manifest, &format!("{label}[{index}]"), name)?; source_shape_matches(spec, label, per_expert, entry)?; if per_expert { if let Some(previous) = &shape { From ccd50edc8995e220dfedfa122d7d08f7c7d1625d Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 18:55:36 +0200 Subject: [PATCH 08/25] fix(device-mesh): complete llama weight store pilot --- crates/hipfire-arch-llama/src/arch.rs | 99 ++- crates/hipfire-arch-llama/src/arch_model.rs | 61 +- crates/hipfire-arch-llama/src/carrier.rs | 758 +++++++++++++++--- crates/hipfire-runtime/src/weight_backend.rs | 9 + crates/hipfire-runtime/src/weight_manifest.rs | 129 +++ crates/hipfire-runtime/src/weight_store.rs | 254 +++++- 6 files changed, 1146 insertions(+), 164 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index 920601367c..f154f0c1b2 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -20,7 +20,7 @@ use hipfire_runtime::hfq::{self, HfqFile}; use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::weight_manifest::{ - FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, + DTypeConstraint, FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, }; use rdna_compute::{DType, Gpu}; @@ -39,6 +39,58 @@ use hipfire_runtime::llama::{attention_family, AttnParams, KvTierInputs, KvTierP /// see [`hipfire_arch_qwen35::Qwen35`] for those. pub struct Llama; +fn linear_source_constraint() -> DTypeConstraint { + DTypeConstraint::source_from_sources(vec![ + DType::F32, + DType::Q4F16G64, + DType::Q8_0, + DType::Q4K, + DType::Q8HFQ, + DType::HFQ4G256, + DType::HFQ4G128, + DType::HFQ6G256, + DType::HFQ2G256, + DType::HFQ2G128, + DType::HFQ3G256, + DType::HFQ3G128, + DType::MQ4G256, + DType::MQ8G256, + DType::MQ6G256, + DType::MQ3G256, + DType::MQ2G256, + DType::MQ2G256Lloyd, + DType::MQ2G256LloydU, + DType::MQ3G256Lloyd, + DType::HFP4G32, + DType::MFP4G32, + DType::MQ4G256Lloyd, + DType::MQ2G256GL, + DType::MQ3G256GL, + DType::TQ2G128, + DType::BQ1G128, + DType::MQ4G256V2, + DType::MQ4CG256, + DType::MQ6G256V2, + DType::MQ5G256V2, + DType::MQ3G256V2, + DType::MQ2G256V2, + ]) +} + +fn embedding_source_constraint() -> DTypeConstraint { + DTypeConstraint::source_from_sources(vec![ + DType::F32, + DType::Q8_0, + DType::Q4K, + DType::HFQ4G256, + DType::HFQ4G128, + ]) +} + +fn norm_source_constraint() -> DTypeConstraint { + DTypeConstraint::source_exact(DType::F32) +} + impl Architecture for Llama { type Weights = LlamaWeights; type State = ForwardScratch; @@ -85,19 +137,24 @@ impl Llama { use ShardPolicy::*; let (dim, hidden, head_dim) = (cfg.dim, cfg.hidden_dim, cfg.head_dim); let (heads, kv_heads) = (cfg.n_heads, cfg.n_kv_heads); + let linear = linear_source_constraint(); + let embedding = embedding_source_constraint(); + let norm = norm_source_constraint(); let mut manifest = Vec::with_capacity(cfg.n_layers * 11 + 3); - manifest.push(WeightEntry::model( + manifest.push(WeightEntry::model_with_dtype_constraint( "token_embd", vec![cfg.vocab_size, dim], DType::F16, + embedding, Pin(PinTarget::Embed), )); for layer in 0..cfg.n_layers { - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wq", layer, vec![heads * head_dim, dim], DType::F16, + linear.clone(), FusedQkv { q_heads: heads, kv_heads, @@ -105,89 +162,101 @@ impl Llama { layout: FusedQkvLayout::Qkv, }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wk", layer, vec![kv_heads * head_dim, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wv", layer, vec![kv_heads * head_dim, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wo", layer, vec![dim, heads * head_dim], DType::F16, + linear.clone(), RowShard { axis: 1 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_gate", layer, vec![hidden, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_up", layer, vec![hidden, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_down", layer, vec![dim, hidden], DType::F16, + linear.clone(), RowShard { axis: 1 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "attn_norm", layer, vec![dim], DType::F32, + norm.clone(), Replicate, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_norm", layer, vec![dim], DType::F32, + norm.clone(), Replicate, )); if cfg.has_qk_norm { - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "q_norm", layer, vec![head_dim], DType::F32, + norm.clone(), Replicate, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "k_norm", layer, vec![head_dim], DType::F32, + norm.clone(), Replicate, )); } } - manifest.push(WeightEntry::model( + manifest.push(WeightEntry::model_with_dtype_constraint( "output_norm", vec![dim], DType::F32, + norm, Replicate, )); - manifest.push(WeightEntry::model( + manifest.push(WeightEntry::model_with_dtype_constraint( "lm_head", vec![cfg.vocab_size, dim], DType::F16, + linear, Pin(PinTarget::Output), )); manifest diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 7e7d3d7ce0..9d1297df95 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -4,10 +4,26 @@ use hipfire_runtime::arch_model::ArchModel; use hipfire_runtime::llama::KvCache; +use hipfire_runtime::weight_store::{ + WeightHandle, WeightOrigin, WeightStore, WeightStoreError, +}; use rdna_compute::Gpu; use crate::carrier::LlamaBundle; +fn drain_weight_store( + store: WeightStore, + origin: WeightOrigin, + gpu: &mut Gpu, +) -> Result<(), (WeightStore, WeightStoreError)> { + for handle in store.take_all(origin)? { + if let WeightHandle::Resident(tensor) = handle { + let _ = gpu.free_tensor(tensor); + } + } + Ok(()) +} + impl ArchModel for LlamaBundle { fn dim(&self) -> usize { self.config.dim @@ -35,6 +51,17 @@ impl ArchModel for LlamaBundle { } fn free_gpu(self: Box, gpu: &mut Gpu) { + // Validate before destructuring the consuming owner. A mismatch must + // leave the resident store attached to an owner that can be retried; + // leaking the boxed owner is safer than dropping the only cleanup + // authority. + if let Some(store) = self.weight_store.as_ref() { + if let Err(error) = store.validate_origin_value(self.weight_origin) { + eprintln!("llama: refusing weight-store release: {error}"); + let _ = Box::into_raw(self); + return; + } + } let LlamaBundle { config: _, weights, @@ -42,21 +69,47 @@ impl ArchModel for LlamaBundle { kv, manifest_plan: _, weight_store, - mesh, + weight_origin, + mesh: _, dflash_extract_layers: _, dspark_weights: _, dspark_assets: _, } = *self; // Mirror the existing unload ordering: scratch → store/weights → kv. - // A committed store is only released here, through the ArchModel owner; - // no store destructor or independent carrier free path exists. scratch.free_gpu(gpu); if let Some(store) = weight_store { - if let Err((_, error)) = store.release_on_owner(&mesh, gpu) { + if let Err((store, error)) = drain_weight_store(store, weight_origin, gpu) { + // This is defensive because the pre-check above used the + // same immutable origin. Never discard a rejected store. eprintln!("llama: refusing weight-store release: {error}"); + std::mem::forget(store); } } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); } } + +#[cfg(test)] +mod tests { + use super::*; + use hipfire_hardware::DeviceMesh; + use hipfire_runtime::weight_manifest::{ShardPolicy, WeightEntry}; + use hipfire_runtime::weight_store::fulfill_manifest_single; + + #[test] + fn arch_owner_unload_drains_residents_and_repeated_empty_unload_is_safe() { + let Ok(mut gpu) = Gpu::init() else { + return; + }; + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::for_single(&mesh, &gpu); + let entry = WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); + let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], rdna_compute::DType::F32)) + }) + .unwrap(); + assert!(drain_weight_store(store, origin, &mut gpu).is_ok()); + assert!(drain_weight_store(WeightStore::with_origin(origin), origin, &mut gpu).is_ok()); + } +} diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index c41b802983..86b4dcfe51 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -7,13 +7,21 @@ use crate::Llama; use hipfire_hardware::DeviceMesh; use hipfire_runtime::arch::Architecture; use hipfire_runtime::dspark_core::DsparkWeights; +use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::llama::{ - ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LlamaConfig, LlamaWeights, + EmbeddingFormat, ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LayerWeights, + LlamaConfig, LlamaWeights, WeightTensor, }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; -use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan}; -use hipfire_runtime::weight_store::WeightStore; +use hipfire_runtime::weight_backend::hfq_weight_dtype; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; +use hipfire_runtime::weight_store::{ + TakenWeight, WeightHandle, WeightStore, WeightStoreAssembly, WeightStoreAssemblyGuard, + WeightOrigin, +}; +use rdna_compute::{DType, GpuTensor}; +use std::collections::HashMap; pub struct LlamaBundle { pub config: LlamaConfig, @@ -27,6 +35,9 @@ pub struct LlamaBundle { /// this bundle. It is crate-visible so callers cannot create an independent /// unload owner; `ArchModel::free_gpu` is the sole release path. pub(crate) weight_store: Option, + /// Exact target identity captured when this bundle was admitted. The + /// owner uses it to validate every store-origin component before teardown. + pub(crate) weight_origin: WeightOrigin, pub(crate) mesh: DeviceMesh, /// Decoder-layer indices whose residual hidden states a hidden-conditioned /// drafter (DFlash / EAGLE) wants captured, ascending order. Empty = no @@ -38,139 +49,539 @@ pub struct LlamaBundle { /// was found or speculation was disabled. Task-10 wires the speculator build. pub dspark_weights: Option, /// Loaded DSpark drafter body assets (5-layer dense-GQA transformer + + /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } +fn plan_single(config: &LlamaConfig) -> Result<(DeviceMesh, ManifestPlan), String> { + let mesh = DeviceMesh::single(); + let manifest = Llama::weight_manifest(config); + let state = Llama::state_manifest(config); + let plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) + .map_err(|e| format!("llama: manifest planning failed: {e}"))?; + Ok((mesh, plan)) +} -/// Build the LLaMA GPU bundle from an HFQ or safetensors-directory source. -/// -/// The source/config path remains architecture-owned. Once it resolves, the -/// carrier publishes a pure Single manifest plan. Every fallible GPU stage -/// explicitly releases earlier allocations before returning an error; no -/// implicit GPU-buffer destructor is introduced. -pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { - let (config, weights, kv, scratch) = match src { - ModelSource::Hfq(mut hfq) => { - let config = - ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; - let weights = ::load_weights(&mut hfq, &config, ctx.gpu)?; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - // The plain LLaMA path has no independent cap resolver. PR #661's - // physical-cap behavior is owned by the existing upstream KV plan. - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" - )); - } - }; - let dims = KvDims { - layers: KvLayers::Flat(config.n_layers), - n_kv_heads: config.n_kv_heads, - head_dim: config.head_dim, - max_seq: ctx.max_seq, - physical_cap: None, - }; - let kv = match ::from_mode( - hipfire_runtime::kv_mode::resolve( - ctx.kv_mode_override.unwrap_or(""), - &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, - config.head_dim, - ) - .mode, - KvTarget::Single(ctx.gpu), - &dims, +fn llama_kv_dims(config: &LlamaConfig, max_seq: usize, physical_cap: Option) -> KvDims { + KvDims { + layers: KvLayers::Flat(config.n_layers), + n_kv_heads: config.n_kv_heads, + head_dim: config.head_dim, + max_seq, + physical_cap, + } +} + +fn hfq_layer_names(layer: usize, relative: &str) -> Vec { + vec![ + format!("model.layers.{layer}.{relative}.weight"), + format!("layers.{layer}.{relative}.weight"), + ] +} + +fn hfq_entry_names(entry: &WeightEntry) -> Result, String> { + let names = match (entry.name.as_str(), entry.layer) { + ("token_embd", None) => vec!["model.embed_tokens.weight".to_string()], + ("output_norm", None) => vec!["model.norm.weight".to_string()], + ("lm_head", None) => vec![ + "lm_head.weight".to_string(), + "model.lm_head.weight".to_string(), + "model.language_model.lm_head.weight".to_string(), + ], + ("wq", Some(layer)) => hfq_layer_names(layer, "self_attn.q_proj"), + ("wk", Some(layer)) => hfq_layer_names(layer, "self_attn.k_proj"), + ("wv", Some(layer)) => hfq_layer_names(layer, "self_attn.v_proj"), + ("wo", Some(layer)) => hfq_layer_names(layer, "self_attn.o_proj"), + ("ffn_gate", Some(layer)) => hfq_layer_names(layer, "mlp.gate_proj"), + ("ffn_up", Some(layer)) => hfq_layer_names(layer, "mlp.up_proj"), + ("ffn_down", Some(layer)) => hfq_layer_names(layer, "mlp.down_proj"), + ("attn_norm", Some(layer)) => hfq_layer_names(layer, "input_layernorm"), + ("ffn_norm", Some(layer)) => hfq_layer_names(layer, "post_attention_layernorm"), + ("q_norm", Some(layer)) => hfq_layer_names(layer, "self_attn.q_norm"), + ("k_norm", Some(layer)) => hfq_layer_names(layer, "self_attn.k_norm"), + (name, layer) => { + return Err(format!( + "llama: manifest entry {name}[layer {layer:?}] has no HFQ source mapping" + )); + } + }; + Ok(names) +} + +fn hfq_entry_data(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, u8), String> { + for name in hfq_entry_names(entry)? { + if let Some((info, data)) = hfq.tensor_data_vec(&name) { + if !matches!( + entry.name.as_str(), + "token_embd" | "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm" ) { - Ok(kv) => kv, - Err(error) => { - scratch.free_gpu(ctx.gpu); - weights.free_gpu(ctx.gpu); + let sidecar = match name.strip_suffix(".weight") { + Some(stem) => format!("{stem}.awq_scale.weight"), + None => format!("{name}.awq_scale.weight"), + }; + if hfq.find_tensor_info(&sidecar).is_some() { return Err(format!( - "llama: ::from_mode failed: {error}" + "llama: AWQ sidecar {sidecar} is not represented by the manifest pilot" )); } - }; - (config, weights, kv, scratch) + } + return Ok((data, info.quant_type)); + } + } + if entry.name == "lm_head" && entry.layer.is_none() { + if let Some((info, data)) = hfq.tensor_data_vec("model.embed_tokens.weight") { + return Ok((data, info.quant_type)); + } + } + Err(format!( + "llama: source tensor for {}[layer {:?}] is missing", + entry.name, entry.layer + )) +} + +fn f32_bytes_from_hfq(quant_type: u8, data: &[u8], name: &str) -> Result, String> { + let mut bytes = Vec::with_capacity(match quant_type { + 1 | 16 => data.len() * 2, + 2 => data.len(), + _ => 0, + }); + match quant_type { + 1 => { + let chunks = data.chunks_exact(2); + if !chunks.remainder().is_empty() { + return Err(format!("{name}: truncated F16 payload")); + } + for chunk in chunks { + bytes.extend_from_slice( + &hipfire_runtime::llama::f16_to_f32(u16::from_le_bytes([chunk[0], chunk[1]])) + .to_le_bytes(), + ); + } + } + 2 => { + if !data.len().is_multiple_of(4) { + return Err(format!("{name}: truncated F32 payload")); + } + bytes.extend_from_slice(data); } - ModelSource::Dir(source) => { - let config = - hipfire_runtime::hfq::config_from_safetensors_llama(&source).map_err(|e| { - format!("failed to parse LLaMA/Qwen3 config from config.json: {e}") - })?; - let weights = - hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) - .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - let kv_mode_str = ctx - .kv_mode_override - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); - let rr = hipfire_runtime::kv_mode::resolve( - &kv_mode_str, - &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, - config.head_dim, - ); - if let Some(w) = rr.warning { - eprintln!( - " KV cache: {w} (site {})", - hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site + 16 => { + let chunks = data.chunks_exact(2); + if !chunks.remainder().is_empty() { + return Err(format!("{name}: truncated BF16 payload")); + } + for chunk in chunks { + bytes.extend_from_slice( + &f32::from_bits( + u16::from_le_bytes([chunk[0], chunk[1]]) as u32 * (1 << 16), + ) + .to_le_bytes(), ); } - let dims = KvDims { - layers: KvLayers::Flat(config.n_layers), - n_kv_heads: config.n_kv_heads, - head_dim: config.head_dim, - max_seq: ctx.max_seq, - physical_cap: Some(ctx.max_seq), - }; - let kv = match ::from_mode( - rr.mode, - KvTarget::Single(ctx.gpu), - &dims, - ) { - Ok(kv) => kv, - Err(error) => { - weights.free_gpu(ctx.gpu); - return Err(format!("KvCache: {error}")); - } - }; - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - let _ = kv.free_gpu(ctx.gpu); - weights.free_gpu(ctx.gpu); - return Err(format!( - "ForwardScratch::new_with_max_seq: {error:?}" - )); - } - }; - (config, weights, kv, scratch) } - }; + other => { + return Err(format!( + "{name}: quant_type={other} is not a host float payload" + )); + } + } + Ok(bytes) +} - // Pure plan publication happens after source/config resolution and before - // the bundle becomes visible to the loader. It performs no GPU or file IO. - let mesh = DeviceMesh::single(); - let manifest = Llama::weight_manifest(&config); - let state = Llama::state_manifest(&config); - let manifest_plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) - .map_err(|e| format!("llama: manifest planning failed: {e}"))?; +fn hfq_source(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, DType), String> { + let (data, quant_type) = hfq_entry_data(hfq, entry)?; + let name = format!("{}[layer {:?}]", entry.name, entry.layer); + if entry.name == "token_embd" { + return match quant_type { + 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), + 3 => Ok((data, DType::Q8_0)), + 4 => Ok((data, DType::Q4K)), + 6 => Ok((data, DType::HFQ4G256)), + 7 => Ok((data, DType::HFQ4G128)), + other => Err(format!( + "{name}: quant_type={other} is unsupported for a LLaMA embedding" + )), + }; + } + if matches!(entry.name.as_str(), "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm") + { + return Ok(( + f32_bytes_from_hfq(quant_type, &data, &name)?, + DType::F32, + )); + } + match quant_type { + 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), + other => hfq_weight_dtype(other) + .map(|dtype| (data, dtype)) + .ok_or_else(|| format!("{name}: unsupported HFQ quant_type={other}")), + } +} - Ok(LlamaBundle { +fn take_slot( + assembly: &mut WeightStoreAssembly<'_>, + slots: &mut HashMap<(String, Option), usize>, + name: &str, + layer: Option, +) -> Result<(), String> { + let slot = assembly + .take(name, layer, 0) + .ok_or_else(|| format!("llama: fulfilled store is missing {name}[layer {layer:?}]"))?; + slots.insert((name.to_string(), layer), slot); + Ok(()) +} + +fn require_resident( + assembly: &WeightStoreAssemblyGuard<'_>, + name: &str, + layer: Option, + slot: usize, +) -> Result<(), String> { + if matches!(assembly.get(slot), Some(WeightHandle::Resident(_))) { + Ok(()) + } else { + Err(format!( + "llama: {name}[layer {layer:?}] is an alias; typed LLaMA assembly requires a resident handle" + )) + } +} + +fn resident_cell( + cells: &mut HashMap<(String, Option), TakenWeight>, + name: &str, + layer: Option, +) -> Result { + let taken = cells + .remove(&(name.to_string(), layer)) + .ok_or_else(|| format!("llama: assembled store is missing {name}[layer {layer:?}]"))?; + match taken.handle { + WeightHandle::Resident(tensor) => Ok(tensor), + WeightHandle::Alias(source) => Err(format!( + "llama: {name}[layer {layer:?}] aliases {source}, expected resident handle" + )), + } +} + +fn resident_weight( + cells: &mut HashMap<(String, Option), TakenWeight>, + name: &str, + layer: Option, + m: usize, + k: usize, +) -> Result { + let tensor = resident_cell(cells, name, layer)?; + let dtype = tensor.dtype; + Ok(WeightTensor { + buf: tensor, + gpu_dtype: dtype, + m, + k, + row_stride: dtype.row_stride(k), + paro: None, + awq_scale: None, + }) +} + +fn embedding_format(dtype: DType) -> Result { + match dtype { + DType::F32 => Ok(EmbeddingFormat::F32), + DType::Q4K => Ok(EmbeddingFormat::Q4K), + DType::HFQ4G256 => Ok(EmbeddingFormat::HFQ4G256), + DType::HFQ4G128 => Ok(EmbeddingFormat::HFQ4G128), + DType::Q8_0 => Ok(EmbeddingFormat::Q8_0), + other => Err(format!( + "llama: unsupported assembled embedding dtype {other:?}" + )), + } +} + +fn assemble_llama_weights( + config: &LlamaConfig, + store: &mut WeightStore, +) -> Result { + let mut assembly = store.begin_assembly(); + let mut slots = HashMap::new(); + let mut take = |name: &str, layer: Option| take_slot(&mut assembly, &mut slots, name, layer); + + take("token_embd", None)?; + take("output_norm", None)?; + take("lm_head", None)?; + for layer in 0..config.n_layers { + for name in ["wq", "wk", "wv", "wo", "ffn_gate", "ffn_up", "ffn_down", "attn_norm", "ffn_norm"] { + take(name, Some(layer))?; + } + if config.has_qk_norm { + take("q_norm", Some(layer))?; + take("k_norm", Some(layer))?; + } + } + + drop(take); + let guard = assembly.commit(); + for ((name, layer), slot) in &slots { + require_resident(&guard, name, *layer, *slot)?; + } + let cells: HashMap<_, _> = guard + .finalize() + .into_iter() + .map(|taken| ((taken.key.name.clone(), taken.key.layer), taken)) + .collect(); + let mut cells = cells; + let token_embd = resident_cell(&mut cells, "token_embd", None)?; + let embd_format = embedding_format(token_embd.dtype)?; + let output_norm = resident_cell(&mut cells, "output_norm", None)?; + let output = resident_weight( + &mut cells, + "lm_head", + None, + config.vocab_size, + config.dim, + )?; + let mut layers = Vec::with_capacity(config.n_layers); + for layer in 0..config.n_layers { + let q_norm = if config.has_qk_norm { + Some(resident_cell(&mut cells, "q_norm", Some(layer))?) + } else { + None + }; + let k_norm = if config.has_qk_norm { + Some(resident_cell(&mut cells, "k_norm", Some(layer))?) + } else { + None + }; + layers.push(LayerWeights { + attn_norm: resident_cell(&mut cells, "attn_norm", Some(layer))?, + wq: resident_weight( + &mut cells, + "wq", + Some(layer), + config.n_heads * config.head_dim, + config.dim, + )?, + wk: resident_weight( + &mut cells, + "wk", + Some(layer), + config.n_kv_heads * config.head_dim, + config.dim, + )?, + wv: resident_weight( + &mut cells, + "wv", + Some(layer), + config.n_kv_heads * config.head_dim, + config.dim, + )?, + wo: resident_weight( + &mut cells, + "wo", + Some(layer), + config.dim, + config.n_heads * config.head_dim, + )?, + q_norm, + k_norm, + ffn_norm: resident_cell(&mut cells, "ffn_norm", Some(layer))?, + w_gate: resident_weight( + &mut cells, + "ffn_gate", + Some(layer), + config.hidden_dim, + config.dim, + )?, + w_up: resident_weight( + &mut cells, + "ffn_up", + Some(layer), + config.hidden_dim, + config.dim, + )?, + w_down: resident_weight( + &mut cells, + "ffn_down", + Some(layer), + config.dim, + config.hidden_dim, + )?, + }); + } + Ok(LlamaWeights { + token_embd, + embd_format, + output_norm, + output, + layers, + lm_head_aliases_embd: false, + }) +} + +/// Build the LLaMA GPU bundle from an HFQ or safetensors-directory source. +/// +/// The HFQ plain-LLaMA Single path is the production manifest pilot: planning +/// and source admission happen first, fulfillment uploads transactionally, and +/// typed handles are moved into `LlamaWeights` before the committed remainder +/// is published beneath this bundle's owner. The directory path remains on its +/// existing ParoQuant loader until that source has an equivalent representation +/// resolver. +pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { + let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = + match src { + ModelSource::Hfq(hfq) => { + let config = + ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; + let (mesh, manifest_plan) = plan_single(&config)?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let mut store = hipfire_runtime::weight_store::fulfill_manifest( + &Llama::weight_manifest(&config), + &mesh, + config.n_layers, + ctx.gpu, + |entry| hfq_source(&hfq, entry), + ) + .map_err(|e| format!("llama: {e}"))?; + let weights = match assemble_llama_weights(&config, &mut store) { + Ok(weights) => weights, + Err(error) => { + store.rollback_unpublished(ctx.gpu); + return Err(error); + } + }; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + // The plain LLaMA path has no independent cap resolver. PR + // #661's physical-cap behavior is owned by the existing + // upstream KV plan. + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + store.rollback_unpublished(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + )); + } + }; + let dims = llama_kv_dims(&config, ctx.max_seq, None); + let kv = match ::from_mode( + hipfire_runtime::kv_mode::resolve( + ctx.kv_mode_override.unwrap_or(""), + &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, + config.head_dim, + ) + .mode, + KvTarget::Single(ctx.gpu), + &dims, + ) { + Ok(kv) => kv, + Err(error) => { + scratch.free_gpu(ctx.gpu); + store.rollback_unpublished(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ::from_mode failed: {error}" + )); + } + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + Some(store), + mesh, + weight_origin, + ) + } + ModelSource::Dir(source) => { + let config = + hipfire_runtime::hfq::config_from_safetensors_llama(&source).map_err(|e| { + format!("failed to parse LLaMA/Qwen3 config from config.json: {e}") + })?; + let (mesh, manifest_plan) = plan_single(&config)?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let weights = + hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) + .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + let kv_mode_str = ctx + .kv_mode_override + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); + let rr = hipfire_runtime::kv_mode::resolve( + &kv_mode_str, + &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, + config.head_dim, + ); + if let Some(w) = rr.warning { + eprintln!( + " KV cache: {w} (site {})", + hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site + ); + } + let dims = llama_kv_dims(&config, ctx.max_seq, Some(ctx.max_seq)); + let kv = match ::from_mode( + rr.mode, + KvTarget::Single(ctx.gpu), + &dims, + ) { + Ok(kv) => kv, + Err(error) => { + weights.free_gpu(ctx.gpu); + return Err(format!("KvCache: {error}")); + } + }; + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let _ = kv.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "ForwardScratch::new_with_max_seq: {error:?}" + )); + } + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + None, + mesh, + weight_origin, + ) + } + }; + + let mut bundle = LlamaBundle { config, weights, scratch, kv, manifest_plan, weight_store: None, + weight_origin, mesh, dflash_extract_layers: Vec::new(), dspark_weights: None, dspark_assets: None, - }) + }; + if let Some(store) = weight_store { + if let Err((store, error)) = bundle.attach_weight_store(store) { + let LlamaBundle { + weights, + scratch, + kv, + .. + } = bundle; + store.rollback_unpublished(ctx.gpu); + scratch.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + let _ = kv.free_gpu(ctx.gpu); + return Err(error); + } + } + Ok(bundle) } /// Alias matching the `load__bundle` naming convention in the task. @@ -178,8 +589,9 @@ pub use load_bundle as load_llama_bundle; impl LlamaBundle { /// Attach a store whose resident handles have been assembled for this - /// bundle. The target mesh epoch is checked before publication; teardown - /// remains exclusively in `ArchModel::free_gpu`. On rejection the store is + /// bundle. The complete target origin (mesh epoch, logical rank, and + /// physical device) is checked before publication; teardown remains + /// exclusively in `ArchModel::free_gpu`. On rejection the store is /// returned unchanged so the caller can retry against the right owner. pub fn attach_weight_store( &mut self, @@ -188,18 +600,8 @@ impl LlamaBundle { if self.weight_store.is_some() { return Err((store, "llama: weight store already attached".into())); } - let Some(origin) = store.origin() else { - return Err((store, "llama: weight store has no origin".into())); - }; - if origin.mesh_epoch() != self.mesh.epoch() { - return Err(( - store, - format!( - "llama: weight store origin epoch {:?} does not match bundle epoch {:?}", - origin.mesh_epoch(), - self.mesh.epoch() - ), - )); + if let Err(error) = store.validate_origin_value(self.weight_origin) { + return Err((store, format!("llama: weight store origin rejected: {error}"))); } self.weight_store = Some(store); Ok(()) @@ -224,3 +626,107 @@ impl LlamaBundle { self.dflash_extract_layers = layers; } } + +#[cfg(test)] +mod tests { + use super::*; + use hipfire_runtime::llama::ModelArch; + use hipfire_runtime::weight_store::{WeightProjection, WeightProjectionKind}; + + fn config() -> LlamaConfig { + LlamaConfig { + arch: ModelArch::Llama, + dim: 4, + hidden_dim: 8, + n_layers: 1, + n_heads: 1, + n_kv_heads: 1, + vocab_size: 8, + head_dim: 4, + norm_eps: 1e-5, + max_seq_len: 32, + rope_freq_base: 10_000.0, + bos_token: 1, + eos_token: 2, + has_qk_norm: false, + } + } + + fn alias_projection() -> WeightProjection { + WeightProjection { + kind: WeightProjectionKind::Static, + axis: None, + rank: 0, + world_size: 1, + logical_shape: vec![1], + dtype: DType::F32, + } + } + + #[test] + fn single_plan_covers_every_typed_llama_handle() { + let (mesh, plan) = plan_single(&config()).unwrap(); + let manifest = Llama::weight_manifest(&config()); + assert_eq!(mesh.n_devices(), 1); + assert_eq!(plan.weights.len(), 12); + assert_eq!(plan.state.len(), 1); + assert!(plan.collective_schedule.iter().any(|entry| entry.name == "wo")); + assert!(manifest[0].dtype_constraint.accepts(DType::HFQ4G256)); + assert!(manifest[1].dtype_constraint.accepts(DType::MQ4G256)); + assert!(manifest[9].dtype_constraint.accepts(DType::F32)); + assert!(!manifest[9].dtype_constraint.accepts(DType::F16)); + } + + #[test] + fn typed_assembly_rolls_back_when_a_cell_is_not_resident() { + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); + let mut store = WeightStore::with_origin(origin); + for name in ["token_embd", "output_norm", "lm_head"] { + store + .stage_alias(name, None, 0, "source", alias_projection()) + .unwrap(); + } + let error = match assemble_llama_weights( + &LlamaConfig { + n_layers: 0, + ..config() + }, + &mut store, + ) { + Ok(_) => panic!("alias unexpectedly assembled as typed weights"), + Err(error) => error, + }; + assert!(error.contains("alias")); + assert_eq!(store.len(), 3); + assert!(store.contains("token_embd", None, 0)); + assert!(store.projection("lm_head", None, 0).is_some()); + } + + + #[test] + fn hfq_float_widening_matches_legacy_f32_representation() { + let f16_one = [0x00, 0x3c, 0x00, 0xc0]; + let actual = f32_bytes_from_hfq(1, &f16_one, "test").unwrap(); + let expected = [1.0f32, -2.0f32] + .into_iter() + .flat_map(f32::to_le_bytes) + .collect::>(); + assert_eq!(actual, expected); + } + + #[test] + fn manifest_constraints_admit_every_pilot_representation() { + let manifest = Llama::weight_manifest(&config()); + assert!(manifest[0].dtype_constraint.accepts(DType::HFQ4G256)); + assert!(manifest[1].dtype_constraint.accepts(DType::MQ4G256)); + assert!(manifest[9].dtype_constraint.accepts(DType::F32)); + assert!(!manifest[9].dtype_constraint.accepts(DType::F16)); + } + #[test] + fn physical_cap_remains_separate_from_configured_max_seq() { + let dims = llama_kv_dims(&config(), 32_768, Some(4_096)); + assert_eq!(dims.max_seq, 32_768); + assert_eq!(dims.physical_cap, Some(4_096)); + } +} diff --git a/crates/hipfire-runtime/src/weight_backend.rs b/crates/hipfire-runtime/src/weight_backend.rs index 43a7d09dc9..b06a6076cc 100644 --- a/crates/hipfire-runtime/src/weight_backend.rs +++ b/crates/hipfire-runtime/src/weight_backend.rs @@ -470,6 +470,15 @@ pub(crate) fn raw_codec(quant_type: u8) -> Option<&'static RawCodec> { RAW_CODECS.iter().find(|c| c.quant_type == quant_type) } +/// Return the compute representation used for a raw HFQ weight payload. +/// +/// Host-decoded source types (F16/F32/BF16) intentionally return `None`; +/// callers must widen those payloads before upload so the result matches the +/// established LLaMA loader semantics. +pub fn hfq_weight_dtype(quant_type: u8) -> Option { + raw_codec(quant_type).map(|codec| codec.dtype) +} + /// Decode a passthrough quant format: enforce the K%256 guard (via DType), /// upload bytes verbatim, build the `WeightTensor` with the dtype + its /// DType-derived row_stride. `name` is the caller context for the guard panic. diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 26c64ad2ee..4e9f91023a 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -364,6 +364,23 @@ fn validate_shape(entry: &WeightEntry) -> Result<(), String> { Ok(()) } +pub(crate) fn validate_weight_layers( + manifest: &[WeightEntry], + n_layers: usize, +) -> Result<(), String> { + for entry in manifest { + if let Some(layer) = entry.layer { + if layer >= n_layers { + return Err(format!( + "{} layer {} outside n_layers={n_layers}", + entry.name, layer + )); + } + } + } + Ok(()) +} + /// Validate logical shard math and tied source identity before fulfillment. pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result<(), String> { let mut identities = HashSet::new(); @@ -431,6 +448,30 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< if source_entry.identity() == entry.identity() { return Err(format!("{context}: an entry cannot tie to itself")); } + if source_entry.logical_shape != entry.logical_shape { + return Err(format!( + "{context}: tied source '{source}' shape {:?} does not match {:?}", + source_entry.logical_shape, entry.logical_shape + )); + } + if source_entry.dtype != entry.dtype { + return Err(format!( + "{context}: tied source '{source}' dtype {:?} does not match {:?}", + source_entry.dtype, entry.dtype + )); + } + if !entry.dtype_constraint.accepts(source_entry.dtype) + || !source_entry.dtype_constraint.accepts(entry.dtype) + { + return Err(format!( + "{context}: tied source '{source}' violates the source dtype contract" + )); + } + if matches!(&source_entry.policy, ShardPolicy::Tied { .. }) { + return Err(format!( + "{context}: tied source '{source}' is itself tied; chains and cycles are unsupported" + )); + } } ShardPolicy::ExpertSharded { n_experts, .. } => { if *n_experts == 0 || entry.logical_shape.first() != Some(n_experts) { @@ -486,6 +527,7 @@ pub fn plan_manifest( mesh: &DeviceMesh, n_layers: usize, ) -> Result { + validate_weight_layers(weights, n_layers)?; validate_manifest(weights, mesh)?; let mut state_ids = HashSet::new(); for entry in state { @@ -1002,4 +1044,91 @@ mod tests { }; assert!(validate_expert_group_specs(&[bad], &manifest).is_err()); } + + #[test] + fn planning_rejects_weight_layer_at_n_layers_and_accepts_last_layer() { + let mesh = DeviceMesh::single(); + let valid = layer_entry("w", 2, ShardPolicy::Replicate); + assert!(plan_manifest(&[valid], &[], &mesh, 3).is_ok()); + let out_of_range = layer_entry("w", 3, ShardPolicy::Replicate); + let error = plan_manifest(&[out_of_range], &[], &mesh, 3).unwrap_err(); + assert!(error.contains("outside n_layers=3")); + } + + #[test] + fn tied_entries_require_matching_representation_and_no_tied_chain() { + let source = WeightEntry::model( + "source", + vec![8, 8], + DType::F16, + ShardPolicy::Replicate, + ); + let shape_mismatch = WeightEntry::model( + "shape_mismatch", + vec![8, 4], + DType::F16, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + assert!(validate_manifest( + &[source.clone(), shape_mismatch], + &DeviceMesh::single() + ) + .is_err()); + + let dtype_mismatch = WeightEntry::model( + "dtype_mismatch", + vec![8, 8], + DType::F32, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + assert!(validate_manifest( + &[source.clone(), dtype_mismatch], + &DeviceMesh::single() + ) + .is_err()); + + let chained_source = WeightEntry::model( + "chained_source", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + let chain = WeightEntry::model( + "chain", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "chained_source".into(), + }, + ); + assert!(validate_manifest( + &[source, chained_source, chain], + &DeviceMesh::single() + ) + .is_err()); + + let cycle_a = WeightEntry::model( + "cycle_a", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "cycle_b".into(), + }, + ); + let cycle_b = WeightEntry::model( + "cycle_b", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "cycle_a".into(), + }, + ); + assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single()).is_err()); + } } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index a663369857..80b8b9bd89 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -11,16 +11,21 @@ //! //! The store is not a model owner. It has no `Drop` implementation and never //! frees GPU buffers implicitly. A carrier may move a committed store into its -//! existing `ArchModel` owner; that owner must call [`WeightStore::release_on_owner`] -//! during its existing teardown path. `take` transfers a resident handle to the -//! owner that is assembling typed weights, and therefore removes the cell from -//! the store's cleanup set. +//! existing `ArchModel` owner; that owner must transfer its resident handles +//! through [`WeightStore::take_all`] during the existing teardown path. +//! `take` transfers a resident handle to the owner that is assembling typed +//! weights, and therefore removes the cell from the store's cleanup set. use crate::weight_manifest::{placement_devices, ShardPolicy, WeightEntry}; use hipfire_hardware::{DeviceMesh, MeshEpoch}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::collections::HashMap; +#[cfg(test)] +thread_local! { + static RESIDENT_RELEASES: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + /// Stable logical placement identity. Layer is part of the key because a /// per-layer name such as `wq` appears once for every decoder block. #[derive(Clone, PartialEq, Eq, Hash, Debug)] @@ -199,8 +204,8 @@ impl std::error::Error for FulfillError {} /// `(name, layer, device)` and captures the target origin once. /// /// There is intentionally no `Drop` implementation. A `WeightStore` that is -/// abandoned without explicit rollback/release leaks rather than guessing a -/// GPU owner; production callers keep it beneath `ArchModel`. +/// abandoned without explicit rollback or owner transfer leaks rather than +/// guessing a GPU owner; production callers keep it beneath `ArchModel`. #[derive(Default)] pub struct WeightStore { placements: HashMap, @@ -355,21 +360,25 @@ impl WeightStore { self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) } - /// Explicit owner teardown. This method is intentionally consuming and has - /// no implicit/drop fallback; `ArchModel::free_gpu` is the production call - /// site. Callers must validate the target first with [`validate_origin`]. - /// On mismatch the original store is returned unchanged for retry by the - /// owner; no GPU call occurs. - pub fn release_on_owner( + /// Roll back a fulfilled store before it is published beneath a model + /// owner. This is the only public consuming GPU-free operation: callers + /// may use it while a load transaction is still unpublished, but an + /// attached store can only be drained by the model owner via `take_all`. + pub fn rollback_unpublished(self, gpu: &Gpu) { + self.release_unchecked(gpu); + } + + /// Transfer every resident/alias handle to the model owner after checking + /// the complete captured origin. On mismatch, the original store is + /// returned unchanged so the owner can retry against the correct target. + pub fn take_all( self, - mesh: &DeviceMesh, - gpu: &mut Gpu, - ) -> Result<(), (Self, WeightStoreError)> { - if let Err(error) = self.validate_origin(mesh, gpu) { + expected: WeightOrigin, + ) -> Result, (Self, WeightStoreError)> { + if let Err(error) = self.validate_origin_value(expected) { return Err((self, error)); } - self.release_unchecked(gpu); - Ok(()) + Ok(self.placements.into_values().collect()) } /// Explicit rollback for a failed transaction. It consumes the partial @@ -385,6 +394,8 @@ impl WeightStore { // the existing loader's explicit owner teardown. The store // never relies on a destructor to release GPU memory. let _ = gpu.hip.free(tensor.buf); + #[cfg(test)] + RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); } } } @@ -496,7 +507,9 @@ where if let Some(error) = target_error(mesh) { return Err(error); } - if let Err(reason) = crate::weight_manifest::validate_manifest(weights, mesh) { + if let Err(reason) = crate::weight_manifest::validate_weight_layers(weights, n_layers) + .and_then(|_| crate::weight_manifest::validate_manifest(weights, mesh)) + { return Err(FulfillError { name: "".to_string(), layer: None, @@ -565,6 +578,27 @@ where ), }); } + if matches!(dtype, DType::F32 | DType::F16 | DType::BF16) { + let expected_bytes = entry + .logical_shape + .iter() + .try_fold(1usize, |count, &dim| count.checked_mul(dim)) + .and_then(|elements| elements.checked_mul(dtype.size())); + if expected_bytes != Some(bytes.len()) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "source payload has {} bytes, expected {:?} for {dtype:?} {:?}", + bytes.len(), + expected_bytes, + entry.logical_shape + ), + }); + } + } let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { Ok(tensor) => tensor, Err(error) => { @@ -721,4 +755,186 @@ mod tests { assert!(target_error(&mesh).is_some()); assert_eq!(placement_devices(&entry, &mesh, 1), vec![0]); } + + #[test] + fn successful_single_fulfillment_commits_resident_projection() { + let Ok(gpu) = Gpu::init() else { + return; + }; + let mesh = DeviceMesh::single(); + let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); + let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], DType::F32)) + }) + .unwrap(); + assert_eq!(store.len(), 1); + assert!(matches!( + store.get("resident", None, 0), + Some(WeightHandle::Resident(tensor)) if tensor.dtype == DType::F32 + )); + assert_eq!( + store.projection("resident", None, 0).unwrap().dtype, + DType::F32 + ); + store.rollback_unpublished(&gpu); + } + + #[test] + fn full_origin_mismatch_returns_resident_store_unchanged() { + let first = DeviceMesh::single(); + let second = DeviceMesh::single(); + let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); + let expected = WeightOrigin::from_parts(second.epoch(), 4, 12); + let mut store = WeightStore::with_origin(actual); + store + .stage_alias("resident", None, 0, "source", projection(DType::F16)) + .unwrap(); + let (store, error) = match store.take_all(expected) { + Ok(_) => panic!("origin mismatch unexpectedly succeeded"), + Err(value) => value, + }; + assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); + assert_eq!(store.origin(), Some(actual)); + assert!(store.contains("resident", None, 0)); + assert!(store.projection("resident", None, 0).is_some()); + } + + #[test] + fn full_origin_mismatch_does_not_free_a_resident_store() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); + let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], DType::F32)) + }) + .unwrap(); + let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); + let (store, error) = match store.take_all(expected) { + Ok(_) => panic!("origin mismatch unexpectedly succeeded"), + Err(value) => value, + }; + assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); + assert_eq!(store.len(), 1); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + 0, + "origin rejection must not free resident buffers" + ); + store.rollback_unpublished(&gpu); + assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); + } + + #[test] + fn owner_transfer_is_consuming_and_empty_transfer_is_idempotent() { + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); + let mut store = WeightStore::with_origin(origin); + store + .stage_alias("owned", None, 0, "source", projection(DType::F16)) + .unwrap(); + let handles = store.take_all(origin).unwrap(); + assert_eq!(handles.len(), 1); + assert!(matches!( + handles.into_iter().next(), + Some(WeightHandle::Alias(_)) + )); + let second = WeightStore::with_origin(origin).take_all(origin).unwrap(); + assert!(second.is_empty()); + } + + #[test] + fn source_failure_after_resident_upload_rolls_back_everything() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let entries = vec![ + WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), + WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), + ]; + let error = fulfill_manifest_single(&entries, &mesh, 1, &gpu, |entry| { + if entry.name == "first" { + Ok((vec![0; 4], DType::F32)) + } else { + Err("injected source failure".into()) + } + }) + .unwrap_err(); + assert_eq!(error.name, "second"); + assert!(error.reason.contains("source read failed")); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + 1, + "the first resident allocation must be explicitly freed" + ); + } + + #[test] + fn dtype_failure_after_resident_upload_rolls_back_everything() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let constraint = DTypeConstraint::source_exact(DType::F32); + let entries = vec![ + WeightEntry::model_with_dtype_constraint( + "first", + vec![1], + DType::F32, + constraint.clone(), + ShardPolicy::Replicate, + ), + WeightEntry::model_with_dtype_constraint( + "second", + vec![1], + DType::F32, + constraint, + ShardPolicy::Replicate, + ), + ]; + let error = fulfill_manifest_single(&entries, &mesh, 1, &gpu, |entry| { + if entry.name == "first" { + Ok((vec![0; 4], DType::F32)) + } else { + Ok((vec![0; 2], DType::F16)) + } + }) + .unwrap_err(); + assert_eq!(error.name, "second"); + assert!(error.reason.contains("violates constraint")); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + 1, + "the first resident allocation must be explicitly freed" + ); + } + + #[test] + fn malformed_upload_payload_after_resident_allocation_rolls_back() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let entries = vec![ + WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), + WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), + ]; + let error = fulfill_manifest_single(&entries, &mesh, 1, &gpu, |entry| { + if entry.name == "first" { + Ok((vec![0; 4], DType::F32)) + } else { + Ok((vec![0; 1], DType::F32)) + } + }) + .unwrap_err(); + assert_eq!(error.name, "second"); + assert!(error.reason.contains("payload")); + assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); + } } From 6cee389b882a0ff3536a75fd8ca8c6fe86df9032 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 19:36:49 +0200 Subject: [PATCH 09/25] fix(device-mesh): close llama manifest pilot review gaps --- crates/hipfire-arch-llama/src/arch.rs | 28 + crates/hipfire-arch-llama/src/arch_model.rs | 59 +- crates/hipfire-arch-llama/src/carrier.rs | 582 ++++++++++++++---- crates/hipfire-runtime/src/hfq.rs | 53 +- crates/hipfire-runtime/src/model_load.rs | 14 +- crates/hipfire-runtime/src/weight_manifest.rs | 46 +- crates/hipfire-runtime/src/weight_store.rs | 254 +++++--- 7 files changed, 790 insertions(+), 246 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index f154f0c1b2..27c60aa9e6 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -262,6 +262,34 @@ impl Llama { manifest } + /// Build the manifest for an HFQ source after source classification. + /// + /// A separate `lm_head.weight` is a resident output projection. When the + /// source omits it, the declaration is a true tie to `token_embd`; the + /// output placement remains pinned to the final stage while the source + /// representation contract is copied from the embedding entry. + pub fn weight_manifest_for_hfq( + cfg: &LlamaConfig, + has_separate_lm_head: bool, + ) -> Vec { + let mut manifest = Self::weight_manifest(cfg); + if !has_separate_lm_head { + let embedding_constraint = manifest + .first() + .expect("LLaMA manifest always contains token_embd") + .dtype_constraint + .clone(); + let output = manifest + .last_mut() + .expect("LLaMA manifest always contains lm_head"); + output.dtype_constraint = embedding_constraint; + output.policy = ShardPolicy::Tied { + source: "token_embd".into(), + }; + } + manifest + } + /// Pure state declaration for the full-attention LLaMA family. pub fn state_manifest(cfg: &LlamaConfig) -> Vec { (0..cfg.n_layers) diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 9d1297df95..763dbbb293 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -4,26 +4,10 @@ use hipfire_runtime::arch_model::ArchModel; use hipfire_runtime::llama::KvCache; -use hipfire_runtime::weight_store::{ - WeightHandle, WeightOrigin, WeightStore, WeightStoreError, -}; use rdna_compute::Gpu; use crate::carrier::LlamaBundle; -fn drain_weight_store( - store: WeightStore, - origin: WeightOrigin, - gpu: &mut Gpu, -) -> Result<(), (WeightStore, WeightStoreError)> { - for handle in store.take_all(origin)? { - if let WeightHandle::Resident(tensor) = handle { - let _ = gpu.free_tensor(tensor); - } - } - Ok(()) -} - impl ArchModel for LlamaBundle { fn dim(&self) -> usize { self.config.dim @@ -51,17 +35,6 @@ impl ArchModel for LlamaBundle { } fn free_gpu(self: Box, gpu: &mut Gpu) { - // Validate before destructuring the consuming owner. A mismatch must - // leave the resident store attached to an owner that can be retried; - // leaking the boxed owner is safer than dropping the only cleanup - // authority. - if let Some(store) = self.weight_store.as_ref() { - if let Err(error) = store.validate_origin_value(self.weight_origin) { - eprintln!("llama: refusing weight-store release: {error}"); - let _ = Box::into_raw(self); - return; - } - } let LlamaBundle { config: _, weights, @@ -69,7 +42,7 @@ impl ArchModel for LlamaBundle { kv, manifest_plan: _, weight_store, - weight_origin, + weight_origin: _, mesh: _, dflash_extract_layers: _, dspark_weights: _, @@ -78,12 +51,10 @@ impl ArchModel for LlamaBundle { // Mirror the existing unload ordering: scratch → store/weights → kv. scratch.free_gpu(gpu); if let Some(store) = weight_store { - if let Err((store, error)) = drain_weight_store(store, weight_origin, gpu) { - // This is defensive because the pre-check above used the - // same immutable origin. Never discard a rejected store. - eprintln!("llama: refusing weight-store release: {error}"); - std::mem::forget(store); - } + // Attachment already checked the complete origin and created this + // owner capability. There is no mismatch branch to leak the model: + // an attached store can only be drained by this consuming owner. + store.drain(gpu); } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); @@ -95,21 +66,27 @@ mod tests { use super::*; use hipfire_hardware::DeviceMesh; use hipfire_runtime::weight_manifest::{ShardPolicy, WeightEntry}; - use hipfire_runtime::weight_store::fulfill_manifest_single; + use hipfire_runtime::weight_store::{ + fulfill_manifest_single, WeightLoadTransaction, WeightOrigin, WeightStore, + }; #[test] - fn arch_owner_unload_drains_residents_and_repeated_empty_unload_is_safe() { - let Ok(mut gpu) = Gpu::init() else { + fn attached_owner_drain_is_consuming_and_empty_drain_is_safe() { + let Ok(gpu) = Gpu::init() else { return; }; let mesh = DeviceMesh::single(); let origin = WeightOrigin::for_single(&mesh, &gpu); - let entry = WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); - let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + let entry = + WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); + let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], rdna_compute::DType::F32)) }) .unwrap(); - assert!(drain_weight_store(store, origin, &mut gpu).is_ok()); - assert!(drain_weight_store(WeightStore::with_origin(origin), origin, &mut gpu).is_ok()); + transaction.publish(origin).unwrap().drain(&gpu); + WeightLoadTransaction::new(WeightStore::with_origin(origin)) + .publish(origin) + .unwrap() + .drain(&gpu); } } diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 86b4dcfe51..5a8b5ebe1e 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -15,10 +15,9 @@ use hipfire_runtime::llama::{ }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; use hipfire_runtime::weight_backend::hfq_weight_dtype; -use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_store::{ - TakenWeight, WeightHandle, WeightStore, WeightStoreAssembly, WeightStoreAssemblyGuard, - WeightOrigin, + AttachedWeightStore, TakenWeight, WeightHandle, WeightLoadTransaction, + WeightStoreAssembly, WeightStoreAssemblyGuard, WeightOrigin, }; use rdna_compute::{DType, GpuTensor}; use std::collections::HashMap; @@ -34,11 +33,11 @@ pub struct LlamaBundle { /// A pilot store is attached only after its handles are assembled under /// this bundle. It is crate-visible so callers cannot create an independent /// unload owner; `ArchModel::free_gpu` is the sole release path. - pub(crate) weight_store: Option, - /// Exact target identity captured when this bundle was admitted. The - /// owner uses it to validate every store-origin component before teardown. + pub(crate) weight_store: Option, + /// Exact target identity captured before publication. The attached store + /// binds this identity into its private drain capability, so teardown + /// cannot encounter an origin mismatch. pub(crate) weight_origin: WeightOrigin, - pub(crate) mesh: DeviceMesh, /// Decoder-layer indices whose residual hidden states a hidden-conditioned /// drafter (DFlash / EAGLE) wants captured, ascending order. Empty = no /// capture (the `SpecTarget::dflash_extract_layers` default of `None`). The @@ -53,9 +52,12 @@ pub struct LlamaBundle { /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } -fn plan_single(config: &LlamaConfig) -> Result<(DeviceMesh, ManifestPlan), String> { +fn plan_single( + config: &LlamaConfig, + has_separate_lm_head: bool, +) -> Result<(DeviceMesh, ManifestPlan), String> { let mesh = DeviceMesh::single(); - let manifest = Llama::weight_manifest(config); + let manifest = Llama::weight_manifest_for_hfq(config, has_separate_lm_head); let state = Llama::state_manifest(config); let plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) .map_err(|e| format!("llama: manifest planning failed: {e}"))?; @@ -230,18 +232,25 @@ fn take_slot( Ok(()) } -fn require_resident( +fn require_materialized( assembly: &WeightStoreAssemblyGuard<'_>, name: &str, layer: Option, slot: usize, ) -> Result<(), String> { - if matches!(assembly.get(slot), Some(WeightHandle::Resident(_))) { - Ok(()) - } else { - Err(format!( - "llama: {name}[layer {layer:?}] is an alias; typed LLaMA assembly requires a resident handle" - )) + match assembly.get(slot) { + Some(WeightHandle::Resident(_)) => Ok(()), + Some(WeightHandle::Alias(source)) + if name == "lm_head" && layer.is_none() && source == "token_embd" => + { + Ok(()) + } + Some(WeightHandle::Alias(source)) => Err(format!( + "llama: {name}[layer {layer:?}] aliases {source}; only lm_head may tie token_embd" + )), + None => Err(format!( + "llama: {name}[layer {layer:?}] assembly slot {slot} is missing" + )), } } @@ -249,15 +258,13 @@ fn resident_cell( cells: &mut HashMap<(String, Option), TakenWeight>, name: &str, layer: Option, -) -> Result { - let taken = cells - .remove(&(name.to_string(), layer)) - .ok_or_else(|| format!("llama: assembled store is missing {name}[layer {layer:?}]"))?; - match taken.handle { - WeightHandle::Resident(tensor) => Ok(tensor), - WeightHandle::Alias(source) => Err(format!( - "llama: {name}[layer {layer:?}] aliases {source}, expected resident handle" - )), +) -> GpuTensor { + match cells.remove(&(name.to_string(), layer)) { + Some(TakenWeight { + handle: WeightHandle::Resident(tensor), + .. + }) => tensor, + _ => unreachable!("validated LLaMA assembly lost resident {name}[layer {layer:?}]"), } } @@ -267,10 +274,10 @@ fn resident_weight( layer: Option, m: usize, k: usize, -) -> Result { - let tensor = resident_cell(cells, name, layer)?; +) -> WeightTensor { + let tensor = resident_cell(cells, name, layer); let dtype = tensor.dtype; - Ok(WeightTensor { + WeightTensor { buf: tensor, gpu_dtype: dtype, m, @@ -278,7 +285,27 @@ fn resident_weight( row_stride: dtype.row_stride(k), paro: None, awq_scale: None, - }) + } +} + +fn tied_weight( + cells: &mut HashMap<(String, Option), TakenWeight>, + token_embd: &GpuTensor, + embd_format: EmbeddingFormat, + name: &str, + layer: Option, + m: usize, + k: usize, +) -> WeightTensor { + match cells.remove(&(name.to_string(), layer)) { + Some(TakenWeight { + handle: WeightHandle::Alias(source), + .. + }) if source == "token_embd" => { + hipfire_runtime::weight_backend::tied_lm_head_alias(token_embd, embd_format, m, k) + } + _ => unreachable!("validated LLaMA assembly lost tied {name}[layer {layer:?}]"), + } } fn embedding_format(dtype: DType) -> Result { @@ -296,17 +323,29 @@ fn embedding_format(dtype: DType) -> Result { fn assemble_llama_weights( config: &LlamaConfig, - store: &mut WeightStore, + transaction: &mut WeightLoadTransaction, ) -> Result { - let mut assembly = store.begin_assembly(); + let mut assembly = transaction.begin_assembly(); let mut slots = HashMap::new(); - let mut take = |name: &str, layer: Option| take_slot(&mut assembly, &mut slots, name, layer); + let mut take = |name: &str, layer: Option| { + take_slot(&mut assembly, &mut slots, name, layer) + }; take("token_embd", None)?; take("output_norm", None)?; take("lm_head", None)?; for layer in 0..config.n_layers { - for name in ["wq", "wk", "wv", "wo", "ffn_gate", "ffn_up", "ffn_down", "attn_norm", "ffn_norm"] { + for name in [ + "wq", + "wk", + "wv", + "wo", + "ffn_gate", + "ffn_up", + "ffn_down", + "attn_norm", + "ffn_norm", + ] { take(name, Some(layer))?; } if config.has_qk_norm { @@ -318,99 +357,124 @@ fn assemble_llama_weights( drop(take); let guard = assembly.commit(); for ((name, layer), slot) in &slots { - require_resident(&guard, name, *layer, *slot)?; + require_materialized(&guard, name, *layer, *slot)?; } + let token_slot = slots[&("token_embd".to_string(), None)]; + let token_dtype = match guard.get(token_slot) { + Some(WeightHandle::Resident(tensor)) => tensor.dtype, + _ => unreachable!("validated token_embd is not resident"), + }; + let embd_format = embedding_format(token_dtype)?; let cells: HashMap<_, _> = guard .finalize() .into_iter() .map(|taken| ((taken.key.name.clone(), taken.key.layer), taken)) .collect(); let mut cells = cells; - let token_embd = resident_cell(&mut cells, "token_embd", None)?; - let embd_format = embedding_format(token_embd.dtype)?; - let output_norm = resident_cell(&mut cells, "output_norm", None)?; - let output = resident_weight( - &mut cells, - "lm_head", - None, - config.vocab_size, - config.dim, - )?; + let token_embd = resident_cell(&mut cells, "token_embd", None); + let output_norm = resident_cell(&mut cells, "output_norm", None); + let lm_head_aliases_embd = matches!( + cells.get(&("lm_head".to_string(), None)), + Some(TakenWeight { + handle: WeightHandle::Alias(_), + .. + }) + ); + let output = if lm_head_aliases_embd { + tied_weight( + &mut cells, + &token_embd, + embd_format, + "lm_head", + None, + config.vocab_size, + config.dim, + ) + } else { + resident_weight( + &mut cells, + "lm_head", + None, + config.vocab_size, + config.dim, + ) + }; let mut layers = Vec::with_capacity(config.n_layers); for layer in 0..config.n_layers { let q_norm = if config.has_qk_norm { - Some(resident_cell(&mut cells, "q_norm", Some(layer))?) + Some(resident_cell(&mut cells, "q_norm", Some(layer))) } else { None }; let k_norm = if config.has_qk_norm { - Some(resident_cell(&mut cells, "k_norm", Some(layer))?) + Some(resident_cell(&mut cells, "k_norm", Some(layer))) } else { None }; layers.push(LayerWeights { - attn_norm: resident_cell(&mut cells, "attn_norm", Some(layer))?, + attn_norm: resident_cell(&mut cells, "attn_norm", Some(layer)), wq: resident_weight( &mut cells, "wq", Some(layer), config.n_heads * config.head_dim, config.dim, - )?, + ), wk: resident_weight( &mut cells, "wk", Some(layer), config.n_kv_heads * config.head_dim, config.dim, - )?, + ), wv: resident_weight( &mut cells, "wv", Some(layer), config.n_kv_heads * config.head_dim, config.dim, - )?, + ), wo: resident_weight( &mut cells, "wo", Some(layer), config.dim, config.n_heads * config.head_dim, - )?, + ), q_norm, k_norm, - ffn_norm: resident_cell(&mut cells, "ffn_norm", Some(layer))?, + ffn_norm: resident_cell(&mut cells, "ffn_norm", Some(layer)), w_gate: resident_weight( &mut cells, "ffn_gate", Some(layer), config.hidden_dim, config.dim, - )?, + ), w_up: resident_weight( &mut cells, "ffn_up", Some(layer), config.hidden_dim, config.dim, - )?, + ), w_down: resident_weight( &mut cells, "ffn_down", Some(layer), config.dim, config.hidden_dim, - )?, + ), }); } + debug_assert!(cells.is_empty(), "validated LLaMA assembly left cells"); Ok(LlamaWeights { token_embd, embd_format, output_norm, output, layers, - lm_head_aliases_embd: false, + lm_head_aliases_embd, }) } @@ -422,43 +486,84 @@ fn assemble_llama_weights( /// is published beneath this bundle's owner. The directory path remains on its /// existing ParoQuant loader until that source has an equivalent representation /// resolver. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum HfqLoadRoute { + /// Plain, non-AWQ files admitted to the manifest/typed-assembly pilot. + ManifestPlainLlama, + /// Files carrying AWQ scale sidecars retain the established loader until + /// sidecar ownership is represented by the manifest transaction. + LegacyAwq, +} + +fn classify_hfq_route(hfq: &HfqFile) -> HfqLoadRoute { + if hfq.has_awq_sidecars() { + HfqLoadRoute::LegacyAwq + } else { + HfqLoadRoute::ManifestPlainLlama + } +} + pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = match src { ModelSource::Hfq(hfq) => { let config = ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; - let (mesh, manifest_plan) = plan_single(&config)?; + // Admission and route classification are pure source checks. + // They must run before any manifest fulfillment or GPU upload. + hipfire_runtime::hfq::validate_llama_hfq_admission(&hfq) + .map_err(|e| e.to_string())?; + let has_separate_lm_head = + hfq.find_tensor_info("lm_head.weight").is_some(); + let route = classify_hfq_route(&hfq); + eprintln!("llama: HFQ source route = {route:?}"); + let (mesh, manifest_plan) = plan_single(&config, has_separate_lm_head)?; let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); - let mut store = hipfire_runtime::weight_store::fulfill_manifest( - &Llama::weight_manifest(&config), - &mesh, - config.n_layers, - ctx.gpu, - |entry| hfq_source(&hfq, entry), - ) - .map_err(|e| format!("llama: {e}"))?; - let weights = match assemble_llama_weights(&config, &mut store) { - Ok(weights) => weights, - Err(error) => { - store.rollback_unpublished(ctx.gpu); - return Err(error); + let (weights, mut weight_store) = match route { + HfqLoadRoute::LegacyAwq => { + let weights = + hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, ctx.gpu) + .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}"))?; + (weights, None) + } + HfqLoadRoute::ManifestPlainLlama => { + let manifest = + Llama::weight_manifest_for_hfq(&config, has_separate_lm_head); + let mut transaction = hipfire_runtime::weight_store::fulfill_manifest( + &manifest, + &mesh, + config.n_layers, + ctx.gpu, + |entry| hfq_source(&hfq, entry), + ) + .map_err(|e| format!("llama: {e}"))?; + let weights = match assemble_llama_weights(&config, &mut transaction) { + Ok(weights) => weights, + Err(error) => { + transaction.rollback(ctx.gpu); + return Err(error); + } + }; + (weights, Some(transaction)) } }; hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); // The plain LLaMA path has no independent cap resolver. PR // #661's physical-cap behavior is owned by the existing // upstream KV plan. - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - store.rollback_unpublished(ctx.gpu); - weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" - )); - } - }; + let scratch = + match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu); + } + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + )); + } + }; let dims = llama_kv_dims(&config, ctx.max_seq, None); let kv = match ::from_mode( hipfire_runtime::kv_mode::resolve( @@ -473,7 +578,9 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result kv, Err(error) => { scratch.free_gpu(ctx.gpu); - store.rollback_unpublished(ctx.gpu); + if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu); + } weights.free_gpu(ctx.gpu); return Err(format!( "llama: ::from_mode failed: {error}" @@ -486,7 +593,7 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Result Result Result Result<(), (WeightStore, String)> { + transaction: WeightLoadTransaction, + ) -> Result<(), (WeightLoadTransaction, String)> { if self.weight_store.is_some() { - return Err((store, "llama: weight store already attached".into())); - } - if let Err(error) = store.validate_origin_value(self.weight_origin) { - return Err((store, format!("llama: weight store origin rejected: {error}"))); + return Err(( + transaction, + "llama: weight store already attached".into(), + )); } - self.weight_store = Some(store); + let attached = match transaction.publish(self.weight_origin) { + Ok(attached) => attached, + Err((transaction, error)) => { + return Err((transaction, format!("llama: weight store origin rejected: {error}"))); + } + }; + self.weight_store = Some(attached); Ok(()) } @@ -631,7 +743,135 @@ impl LlamaBundle { mod tests { use super::*; use hipfire_runtime::llama::ModelArch; - use hipfire_runtime::weight_store::{WeightProjection, WeightProjectionKind}; + use hipfire_runtime::arch_model::ArchModel; + use hipfire_runtime::hfq::{ + write_hfqm_package_mem, HfqFile, HfqMemTensor, + }; + use hipfire_runtime::kv_backend::KvBackend; + use hipfire_runtime::kv_mode::KvMode; + use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; + use hipfire_runtime::llama::{KvCache, KvCacheExt, KvDims, KvLayers, KvTarget}; + use hipfire_runtime::weight_store::{ + WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, + }; + use std::path::{Path, PathBuf}; + use std::time::{SystemTime, UNIX_EPOCH}; + + fn hfq_tensor(name: &str, shape: &[u32], quant_type: u8, bytes: usize) -> HfqMemTensor { + HfqMemTensor { + name: name.into(), + quant_type, + shape: shape.to_vec(), + group_size: 0, + data: vec![0; bytes], + } + } + + fn f32_hfq_tensor(name: &str, shape: &[u32], malformed: bool) -> HfqMemTensor { + let elements = shape.iter().map(|&dim| dim as usize).product::(); + hfq_tensor(name, shape, 2, if malformed { 4 } else { elements * 4 }) + } + + fn fixture_hfq( + with_awq_sidecar: bool, + with_q_proj_bias: bool, + malformed_output_norm: bool, + separate_lm_head: bool, + ) -> (PathBuf, HfqFile) { + let mut tensors = vec![ + f32_hfq_tensor("model.embed_tokens.weight", &[2, 32], false), + f32_hfq_tensor("model.norm.weight", &[32], false), + f32_hfq_tensor("model.layers.0.self_attn.q_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.self_attn.k_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.self_attn.v_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.self_attn.o_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32], false), + f32_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32], false), + f32_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64], false), + f32_hfq_tensor( + "model.layers.0.input_layernorm.weight", + &[32], + false, + ), + f32_hfq_tensor( + "model.layers.0.post_attention_layernorm.weight", + &[32], + false, + ), + ]; + if malformed_output_norm { + tensors[1] = f32_hfq_tensor("model.norm.weight", &[32], true); + } + if with_awq_sidecar { + tensors.push(hfq_tensor( + "model.layers.0.self_attn.q_proj.awq_scale.weight", + &[32], + 1, + 32 * 2, + )); + } + if with_q_proj_bias { + tensors.push(hfq_tensor( + "model.layers.0.self_attn.q_proj.bias", + &[32], + 1, + 32 * 2, + )); + } + if separate_lm_head { + tensors.push(f32_hfq_tensor("lm_head.weight", &[2, 32], false)); + } + let metadata = r#"{ + "config": { + "model_type": "llama", + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "intermediate_size": 64, + "vocab_size": 2, + "head_dim": 32, + "rms_norm_eps": 0.00001, + "max_position_embeddings": 8, + "rope_theta": 10000.0 + } + }"#; + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before epoch") + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "hipfire-g3-{}-{nonce}.hfq", + std::process::id() + )); + write_hfqm_package_mem(&path, 0, metadata, &tensors).expect("write HFQ fixture"); + let hfq = HfqFile::open(&path).expect("open HFQ fixture"); + (path, hfq) + } + + fn load_ctx<'a>( + path: &'a Path, + gpu: &'a mut rdna_compute::Gpu, + cask: &'a CaskConfig, + ) -> LoadCtx<'a> { + LoadCtx { + path: path.to_str().expect("fixture path is UTF-8"), + max_seq: 8, + deepseek4_compute_placement: Default::default(), + deepseek4_experts_per_token: None, + draft_path: None, + kv_mode_override: Some("q8"), + kv_backend: KvBackend::Contiguous, + kv_adaptive_override: None, + state_quant_override: None, + cask, + pp: 1, + spec: SpecLoadCfg::default(), + gpu, + gemma4_drafter_path: None, + gemma4_draft_len: 3, + } + } fn config() -> LlamaConfig { LlamaConfig { @@ -662,10 +902,9 @@ mod tests { dtype: DType::F32, } } - #[test] fn single_plan_covers_every_typed_llama_handle() { - let (mesh, plan) = plan_single(&config()).unwrap(); + let (mesh, plan) = plan_single(&config(), true).unwrap(); let manifest = Llama::weight_manifest(&config()); assert_eq!(mesh.n_devices(), 1); assert_eq!(plan.weights.len(), 12); @@ -687,20 +926,21 @@ mod tests { .stage_alias(name, None, 0, "source", alias_projection()) .unwrap(); } + let mut transaction = WeightLoadTransaction::new(store); let error = match assemble_llama_weights( &LlamaConfig { n_layers: 0, ..config() }, - &mut store, + &mut transaction, ) { Ok(_) => panic!("alias unexpectedly assembled as typed weights"), Err(error) => error, }; assert!(error.contains("alias")); - assert_eq!(store.len(), 3); - assert!(store.contains("token_embd", None, 0)); - assert!(store.projection("lm_head", None, 0).is_some()); + assert_eq!(transaction.len(), 3); + assert!(transaction.contains("token_embd", None, 0)); + assert!(transaction.projection("lm_head", None, 0).is_some()); } @@ -729,4 +969,136 @@ mod tests { assert_eq!(dims.max_seq, 32_768); assert_eq!(dims.physical_cap, Some(4_096)); } + + #[test] + fn missing_lm_head_manifest_declares_a_tied_embedding_alias() { + let manifest = Llama::weight_manifest_for_hfq(&config(), false); + let token = &manifest[0]; + let output = manifest.last().expect("manifest has lm_head"); + assert!(matches!( + output.policy, + ShardPolicy::Tied { ref source } if source == "token_embd" + )); + assert!(token + .dtype_constraint + .same_source_set(&output.dtype_constraint)); + } + + #[test] + fn production_hfq_single_route_aliases_missing_lm_head_without_second_allocation() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, false, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); + drop(ctx); + assert!(bundle.weights.lm_head_aliases_embd); + assert_eq!( + bundle.weights.output.buf.buf.as_ptr(), + bundle.weights.token_embd.buf.as_ptr() + ); + assert!(bundle.weight_store.is_some()); + Box::new(bundle).free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_awq_route_preserves_legacy_loader() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(true, false, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) + .expect("AWQ fixture must use the legacy HFQ loader"); + drop(ctx); + assert!(bundle.weight_store.is_none()); + assert!(bundle.weights.lm_head_aliases_embd); + Box::new(bundle).free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_biased_hfq_is_rejected_before_manifest_upload() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, true, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let error = match load_bundle(ModelSource::Hfq(hfq), &mut ctx) { + Ok(_) => panic!("biased HFQ unexpectedly loaded"), + Err(error) => error, + }; + drop(ctx); + assert!(error.contains("q_proj.bias")); + assert!(error.contains("refusing to load Qwen2")); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_post_resident_failure_returns_clean_load_error() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, false, true, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let error = match load_bundle(ModelSource::Hfq(hfq), &mut ctx) { + Ok(_) => panic!("malformed output norm unexpectedly loaded"), + Err(error) => error, + }; + drop(ctx); + assert!(error.contains("source payload") || error.contains("output_norm")); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_manifest_matches_legacy_alias_contract() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, false, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); + let manifest_alias = bundle.weights.lm_head_aliases_embd; + drop(ctx); + Box::new(bundle).free_gpu(&mut gpu); + + let hfq = HfqFile::open(&path).expect("reopen HFQ fixture"); + let config = ::config_from_hfq(&hfq).expect("fixture config"); + let legacy = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, &mut gpu) + .expect("legacy loader fixture"); + assert_eq!(manifest_alias, legacy.lm_head_aliases_embd); + assert_eq!(legacy.embd_format, EmbeddingFormat::F32); + legacy.free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn physical_cap_is_honored_by_upstream_kv_constructor() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let dims = KvDims { + layers: KvLayers::Flat(1), + n_kv_heads: 1, + head_dim: 32, + max_seq: 8, + physical_cap: Some(4), + }; + let cache = ::from_mode( + KvMode::Q8, + KvTarget::Single(&mut gpu), + &dims, + ) + .expect("upstream Q8 constructor"); + assert_eq!(cache.max_seq, 8); + assert_eq!(cache.physical_cap, 4); + let _ = cache.free_gpu(&mut gpu); + } } diff --git a/crates/hipfire-runtime/src/hfq.rs b/crates/hipfire-runtime/src/hfq.rs index c15678e420..54a68a0b67 100644 --- a/crates/hipfire-runtime/src/hfq.rs +++ b/crates/hipfire-runtime/src/hfq.rs @@ -432,6 +432,19 @@ impl HfqFile { self.overlay.is_some() } + /// Whether this source carries any AWQ scale sidecar. The carrier uses + /// this classification before allocation so supported sidecars stay on the + /// legacy loader until the manifest resolver can represent them. + pub fn has_awq_sidecars(&self) -> bool { + self.tensors + .iter() + .any(|tensor| tensor.name.ends_with(".awq_scale.weight")) + || self + .overlay + .as_ref() + .is_some_and(|overlay| overlay.has_awq_sidecars()) + } + /// Open an HFQM container that lives inside a larger file, starting at /// `base_offset`. Used by the bundled `.mq4-mtp` loader to parse the /// MTP section embedded after the trunk's tensor data. @@ -1638,6 +1651,7 @@ fn load_embedding_llama( .expect("embed_tokens not found"); // Q4K embeddings are llama-family-only (GGUF-derived). qwen2/qwen35 have no // Q4K embedding-lookup kernel — that is why the shared `load_embedding` / + // `embed_classify` deliberately rejects qt 4 (rejecting at load gives a clean // error instead of an "unsupported embedding format" panic deep in the qwen // forward pass). So Q4K stays an explicit llama-only branch here; everything @@ -1651,28 +1665,13 @@ fn load_embedding_llama( load_embedding(gpu, info.quant_type, data, config.vocab_size, config.dim) } -/// Load LLaMA weights from an HFQ file onto GPU. -pub fn load_weights_hfq( - hfq: &HfqFile, - config: &LlamaConfig, - gpu: &mut Gpu, -) -> HipResult { - // R2 guard: the LLaMA-family loader does NOT read Q/K/V proj bias — - // `LayerWeights` has no `wq_bias` / `wk_bias` / `wv_bias` fields and - // the per-layer load below only names `*.q_proj.weight`. Qwen2 - // requires those biases (`attention_bias=true` is the modeling - // default). The quantiser used to auto-tag every Qwen2 model as - // `arch_id=1`, which the daemon dispatches to this loader; the - // result was silently-wrong outputs with no warning. As of the - // `--arch-id` flag (see `hipfire-quantize`), Qwen2 models should be - // tagged `arch_id=7` and dispatched to `hipfire-arch-qwen2`. - // - // If we see `q_proj.bias` while loading as the LLaMA family, the - // input is a mis-tagged Qwen2 HFQ. Refuse hard with a pointer at - // the correct path. (Detection by manifest is robust to either the - // model_type tag or the model family — both LLaMA and Qwen3 lack - // these bias tensors, so any HFQ with `model.layers.0.self_attn.q_proj.bias` - // is by definition a Qwen2-family input.) +/// Reject a mis-tagged Qwen2 HFQ before any model allocation. +/// +/// The LLaMA-family `LayerWeights` type has no attention-bias tensors. A +/// Qwen2 file carrying `q_proj.bias` would therefore load and produce +/// silently-wrong output unless this admission check runs before every loader +/// route, including the manifest pilot. +pub fn validate_llama_hfq_admission(hfq: &HfqFile) -> HipResult<()> { if hfq .find_tensor_info("model.layers.0.self_attn.q_proj.bias") .is_some() @@ -1696,6 +1695,16 @@ pub fn load_weights_hfq( ), )); } + Ok(()) +} + +/// Load LLaMA weights from an HFQ file onto GPU. +pub fn load_weights_hfq( + hfq: &HfqFile, + config: &LlamaConfig, + gpu: &mut Gpu, +) -> HipResult { + validate_llama_hfq_admission(hfq)?; let mut source = LlamaHfqSource { hfq, cfg: config }; let layout = crate::model_load::Layout::single(config.n_layers); diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index 46fcad6643..d233630a62 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -92,15 +92,15 @@ impl Layout { } /// Neutral result of the orchestrator. Each arch assembles its own weights -/// struct from this (qwen35 adds `pager`; llama drops `lm_head_aliases_embd`). +/// struct from this (qwen35 adds `pager`). pub struct LoadedWeights { pub token_embd: GpuTensor, pub embd_format: EmbeddingFormat, pub output_norm: GpuTensor, pub output: WeightTensor, pub layers: Vec, - /// True iff the tied lm_head aliases the embedding buffer (qwen35 single-GPU); - /// llama always returns `false` (it reuploads). + /// True iff the tied lm_head aliases the embedding buffer on this + /// single-device route; false means a separate output allocation exists. pub lm_head_aliases_embd: bool, } @@ -112,12 +112,10 @@ pub trait WeightSource { fn n_layers(&self) -> usize; /// Pre-load hook. HFQ drops the mmap when n==1; PaRo rejects n>1; llama no-op. fn prepare(&mut self, n_devices: usize) -> HipResult<()>; - fn read_embed(&mut self, gpu: &mut Gpu) -> HipResult<(GpuTensor, EmbeddingFormat)>; - fn read_final_norm(&mut self, gpu: &mut Gpu) -> HipResult; - /// `can_alias` is true iff embed and output share a device (n==1); the impl - /// decides whether to use it (qwen35 aliases; llama ignores it and reuploads). + /// `can_alias` is true iff embed and output share a device (n==1); the + /// implementation decides whether to use it (single-device LLaMA and + /// qwen35 alias tied embeddings; multi-device routes re-materialize). fn read_output( - &mut self, gpu: &mut Gpu, embd: &GpuTensor, embd_fmt: EmbeddingFormat, diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 4e9f91023a..044bc62c6b 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -100,6 +100,27 @@ impl DTypeConstraint { SourceDType::OneOf(allowed) => allowed.contains(&dtype), } } + + /// Whether two source constraints admit exactly the same representation + /// set. Variant spelling is not part of the contract: `Exact(F16)` and + /// `OneOf([F16])` are equivalent, while `Any` is never equivalent to a + /// finite list. + pub fn same_source_set(&self, other: &Self) -> bool { + fn finite_equal(left: &[DType], right: &[DType]) -> bool { + left.iter().all(|dtype| right.contains(dtype)) + && right.iter().all(|dtype| left.contains(dtype)) + } + match (&self.source, &other.source) { + (SourceDType::Any, SourceDType::Any) => true, + (SourceDType::Any, _) | (_, SourceDType::Any) => false, + (SourceDType::Exact(left), SourceDType::Exact(right)) => left == right, + (SourceDType::Exact(dtype), SourceDType::OneOf(values)) + | (SourceDType::OneOf(values), SourceDType::Exact(dtype)) => { + values.iter().all(|value| value == dtype) + } + (SourceDType::OneOf(left), SourceDType::OneOf(right)) => finite_equal(left, right), + } + } } /// The block ordering of a fused projection. @@ -460,7 +481,10 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< source_entry.dtype, entry.dtype )); } - if !entry.dtype_constraint.accepts(source_entry.dtype) + if !source_entry + .dtype_constraint + .same_source_set(&entry.dtype_constraint) + || !entry.dtype_constraint.accepts(source_entry.dtype) || !source_entry.dtype_constraint.accepts(entry.dtype) { return Err(format!( @@ -1131,4 +1155,24 @@ mod tests { ); assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single()).is_err()); } + #[test] + fn tied_entries_reject_different_source_sets_with_equal_logical_dtype() { + let source = WeightEntry::model( + "source", + vec![8, 8], + DType::F16, + ShardPolicy::Replicate, + ); + let tied = WeightEntry::model_with_dtype_constraint( + "tied", + vec![8, 8], + DType::F16, + DTypeConstraint::source_exact(DType::F16), + ShardPolicy::Tied { + source: "source".into(), + }, + ); + let error = validate_manifest(&[source, tied], &DeviceMesh::single()).unwrap_err(); + assert!(error.contains("source dtype contract")); + } } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index 80b8b9bd89..99ac6cc886 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -10,12 +10,12 @@ //! them, and the first failure explicitly rolls back every resident buffer. //! //! The store is not a model owner. It has no `Drop` implementation and never -//! frees GPU buffers implicitly. A carrier may move a committed store into its -//! existing `ArchModel` owner; that owner must transfer its resident handles -//! through [`WeightStore::take_all`] during the existing teardown path. -//! `take` transfers a resident handle to the owner that is assembling typed -//! weights, and therefore removes the cell from the store's cleanup set. - +//! frees GPU buffers implicitly. A carrier moves a committed transaction into +//! its existing `ArchModel` owner; that owner consumes the private drain +//! capability during the existing teardown path. +//! `WeightStoreAssembly::take` transfers a resident handle to the owner that is +//! assembling typed weights, and therefore removes the cell from the store's +//! cleanup set. use crate::weight_manifest::{placement_devices, ShardPolicy, WeightEntry}; use hipfire_hardware::{DeviceMesh, MeshEpoch}; use rdna_compute::{DType, Gpu, GpuTensor}; @@ -201,11 +201,9 @@ impl std::fmt::Display for FulfillError { impl std::error::Error for FulfillError {} /// Load-side placement container. It records one immutable projection per -/// `(name, layer, device)` and captures the target origin once. -/// -/// There is intentionally no `Drop` implementation. A `WeightStore` that is -/// abandoned without explicit rollback or owner transfer leaks rather than -/// guessing a GPU owner; production callers keep it beneath `ArchModel`. +/// `(name, layer, device)` and captures the target origin once. The container +/// itself has no consuming teardown API: lifecycle transitions are represented +/// by [`WeightLoadTransaction`] and [`AttachedWeightStore`]. #[derive(Default)] pub struct WeightStore { placements: HashMap, @@ -213,6 +211,145 @@ pub struct WeightStore { origin: Option, } +/// The only owner that may roll back resident allocations before publication. +/// +/// A transaction owns the store until [`Self::publish`] transfers it into the +/// attached owner held by the architecture bundle. It deliberately has no +/// implicit `Drop` cleanup because the GPU is not available to a destructor. +pub struct WeightLoadTransaction { + store: Option, +} + +impl std::fmt::Debug for WeightLoadTransaction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WeightLoadTransaction") + .field("origin", &self.origin()) + .field("len", &self.len()) + .finish() + } +} + +/// The resident-store capability returned by a committed load transaction. +/// +/// The backing store and its drain capability are private. Architecture +/// owners receive this value during attachment and consume it exactly once +/// during unload; no public `WeightStore` method can drain an attached store. +pub struct AttachedWeightStore { + store: WeightStore, + capability: WeightStoreDrainCapability, +} + +struct WeightStoreDrainCapability { + origin: WeightOrigin, +} + +impl WeightLoadTransaction { + pub fn new(store: WeightStore) -> Self { + Self { store: Some(store) } + } + + pub fn origin(&self) -> Option { + self.store.as_ref().and_then(WeightStore::origin) + } + + pub fn len(&self) -> usize { + self.store.as_ref().map_or(0, WeightStore::len) + } + + pub fn is_empty(&self) -> bool { + self.store.as_ref().map_or(true, WeightStore::is_empty) + } + + pub fn contains(&self, name: &str, layer: Option, device: usize) -> bool { + self.store + .as_ref() + .is_some_and(|store| store.contains(name, layer, device)) + } + + pub fn get( + &self, + name: &str, + layer: Option, + device: usize, + ) -> Option<&WeightHandle> { + self.store + .as_ref() + .and_then(|store| store.get(name, layer, device)) + } + + pub fn projection( + &self, + name: &str, + layer: Option, + device: usize, + ) -> Option<&WeightProjection> { + self.store + .as_ref() + .and_then(|store| store.projection(name, layer, device)) + } + + pub fn devices_for(&self, name: &str, layer: Option) -> Vec { + self.store + .as_ref() + .map_or_else(Vec::new, |store| store.devices_for(name, layer)) + } + + /// Start typed assembly while this load is still unpublished. + pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { + self.store + .as_mut() + .expect("weight load transaction was already consumed") + .begin_assembly() + } + + /// Consume this transaction and release every resident handle it owns. + /// This is intentionally the only rollback operation exposed by the + /// lifecycle API. + pub fn rollback(mut self, gpu: &Gpu) { + if let Some(store) = self.store.take() { + store.rollback(gpu); + } + } + + /// Publish the store beneath an architecture owner after checking the + /// complete immutable target identity. A mismatch returns this + /// transaction unchanged so the caller can retry or roll it back. + pub fn publish( + mut self, + expected: WeightOrigin, + ) -> Result { + let store = self + .store + .take() + .expect("weight load transaction was already consumed"); + if let Err(error) = store.validate_origin_value(expected) { + self.store = Some(store); + return Err((self, error)); + } + Ok(AttachedWeightStore { + store, + capability: WeightStoreDrainCapability { origin: expected }, + }) + } +} + +impl AttachedWeightStore { + /// Drain resident handles through the private owner capability. The + /// capability is established only by `WeightLoadTransaction::publish`, so + /// origin mismatch is impossible after attachment. + pub fn drain(self, gpu: &Gpu) { + let Self { store, capability } = self; + capability.drain(store, gpu); + } +} + +impl WeightStoreDrainCapability { + fn drain(self, store: WeightStore, gpu: &Gpu) { + debug_assert_eq!(store.origin, Some(self.origin)); + store.release_unchecked(gpu); + } +} + impl WeightStore { pub fn new() -> Self { Self::default() @@ -300,9 +437,10 @@ impl WeightStore { ) } - /// Move a handle out of the store. The projection is removed with it so no - /// stale metadata can describe a cell that the store no longer owns. - pub fn take( + /// Move a handle out of the store. This is private to the assembly + /// capability so arbitrary store holders cannot independently tear down a + /// resident allocation. + fn take( &mut self, name: &str, layer: Option, @@ -313,7 +451,7 @@ impl WeightStore { self.placements.remove(&key) } - pub fn take_with_projection( + fn take_with_projection( &mut self, name: &str, layer: Option, @@ -325,10 +463,7 @@ impl WeightStore { Some((handle, projection)) } - /// Start a typed assembly transaction. Handles moved through the - /// transaction are restored to this store if the transaction is dropped - /// before `finalize`; no GPU free or second owner is introduced. - pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { + fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { WeightStoreAssembly { store: self, taken: Vec::new(), @@ -337,8 +472,7 @@ impl WeightStore { } /// Compare a store's captured origin with an already-resolved target - /// identity. This pure seam is used by fault-path tests and by owner - /// teardown after target resolution. + /// identity. This read-only seam cannot release or extract any handle. pub fn validate_origin_value( &self, expected: WeightOrigin, @@ -360,27 +494,6 @@ impl WeightStore { self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) } - /// Roll back a fulfilled store before it is published beneath a model - /// owner. This is the only public consuming GPU-free operation: callers - /// may use it while a load transaction is still unpublished, but an - /// attached store can only be drained by the model owner via `take_all`. - pub fn rollback_unpublished(self, gpu: &Gpu) { - self.release_unchecked(gpu); - } - - /// Transfer every resident/alias handle to the model owner after checking - /// the complete captured origin. On mismatch, the original store is - /// returned unchanged so the owner can retry against the correct target. - pub fn take_all( - self, - expected: WeightOrigin, - ) -> Result, (Self, WeightStoreError)> { - if let Err(error) = self.validate_origin_value(expected) { - return Err((self, error)); - } - Ok(self.placements.into_values().collect()) - } - /// Explicit rollback for a failed transaction. It consumes the partial /// store and frees every resident buffer on the single owning GPU. fn rollback(self, gpu: &Gpu) { @@ -500,7 +613,7 @@ pub fn fulfill_manifest_single( n_layers: usize, gpu: &Gpu, source: F, -) -> Result +) -> Result where F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, { @@ -623,7 +736,7 @@ where }); } } - Ok(store) + Ok(WeightLoadTransaction::new(store)) } /// Canonical name used by the manifest fulfillment seam. The target is @@ -635,7 +748,7 @@ pub fn fulfill_manifest( n_layers: usize, gpu: &Gpu, source: F, -) -> Result +) -> Result where F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, { @@ -763,24 +876,24 @@ mod tests { }; let mesh = DeviceMesh::single(); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) }) .unwrap(); - assert_eq!(store.len(), 1); + assert_eq!(transaction.len(), 1); assert!(matches!( - store.get("resident", None, 0), + transaction.get("resident", None, 0), Some(WeightHandle::Resident(tensor)) if tensor.dtype == DType::F32 )); assert_eq!( - store.projection("resident", None, 0).unwrap().dtype, + transaction.projection("resident", None, 0).unwrap().dtype, DType::F32 ); - store.rollback_unpublished(&gpu); + transaction.rollback(&gpu); } #[test] - fn full_origin_mismatch_returns_resident_store_unchanged() { + fn full_origin_mismatch_returns_unpublished_transaction_unchanged() { let first = DeviceMesh::single(); let second = DeviceMesh::single(); let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); @@ -789,60 +902,63 @@ mod tests { store .stage_alias("resident", None, 0, "source", projection(DType::F16)) .unwrap(); - let (store, error) = match store.take_all(expected) { + let transaction = WeightLoadTransaction::new(store); + let (transaction, error) = match transaction.publish(expected) { Ok(_) => panic!("origin mismatch unexpectedly succeeded"), Err(value) => value, }; assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); - assert_eq!(store.origin(), Some(actual)); - assert!(store.contains("resident", None, 0)); - assert!(store.projection("resident", None, 0).is_some()); + assert_eq!(transaction.origin(), Some(actual)); + assert!(transaction.contains("resident", None, 0)); + assert!(transaction.projection("resident", None, 0).is_some()); } #[test] - fn full_origin_mismatch_does_not_free_a_resident_store() { + fn full_origin_mismatch_does_not_free_a_resident_transaction() { let Ok(gpu) = Gpu::init() else { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); let mesh = DeviceMesh::single(); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) }) .unwrap(); let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); - let (store, error) = match store.take_all(expected) { + let (transaction, error) = match transaction.publish(expected) { Ok(_) => panic!("origin mismatch unexpectedly succeeded"), Err(value) => value, }; assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); - assert_eq!(store.len(), 1); + assert_eq!(transaction.len(), 1); assert_eq!( RESIDENT_RELEASES.with(std::cell::Cell::get), 0, "origin rejection must not free resident buffers" ); - store.rollback_unpublished(&gpu); + transaction.rollback(&gpu); assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } #[test] - fn owner_transfer_is_consuming_and_empty_transfer_is_idempotent() { + fn attached_owner_transfer_is_consuming_and_empty_transfer_is_safe() { let mesh = DeviceMesh::single(); let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); let mut store = WeightStore::with_origin(origin); store .stage_alias("owned", None, 0, "source", projection(DType::F16)) .unwrap(); - let handles = store.take_all(origin).unwrap(); - assert_eq!(handles.len(), 1); - assert!(matches!( - handles.into_iter().next(), - Some(WeightHandle::Alias(_)) - )); - let second = WeightStore::with_origin(origin).take_all(origin).unwrap(); - assert!(second.is_empty()); + let transaction = WeightLoadTransaction::new(store); + let attached = transaction.publish(origin).unwrap(); + let Ok(gpu) = Gpu::init() else { + return; + }; + attached.drain(&gpu); + let empty = WeightLoadTransaction::new(WeightStore::with_origin(origin)) + .publish(origin) + .unwrap(); + empty.drain(&gpu); } #[test] From 5173c5e37cce4ffd89434e365b0d749109659fe1 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 20:03:16 +0200 Subject: [PATCH 10/25] fix(device-mesh): close final G3 manifest review gaps --- crates/hipfire-arch-llama/src/arch_model.rs | 29 --- crates/hipfire-arch-llama/src/carrier.rs | 127 ++++++++-- crates/hipfire-runtime/src/weight_store.rs | 250 +++++++++++++------- 3 files changed, 268 insertions(+), 138 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 763dbbb293..f21e75b725 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -61,32 +61,3 @@ impl ArchModel for LlamaBundle { } } -#[cfg(test)] -mod tests { - use super::*; - use hipfire_hardware::DeviceMesh; - use hipfire_runtime::weight_manifest::{ShardPolicy, WeightEntry}; - use hipfire_runtime::weight_store::{ - fulfill_manifest_single, WeightLoadTransaction, WeightOrigin, WeightStore, - }; - - #[test] - fn attached_owner_drain_is_consuming_and_empty_drain_is_safe() { - let Ok(gpu) = Gpu::init() else { - return; - }; - let mesh = DeviceMesh::single(); - let origin = WeightOrigin::for_single(&mesh, &gpu); - let entry = - WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); - let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { - Ok((vec![0; 4], rdna_compute::DType::F32)) - }) - .unwrap(); - transaction.publish(origin).unwrap().drain(&gpu); - WeightLoadTransaction::new(WeightStore::with_origin(origin)) - .publish(origin) - .unwrap() - .drain(&gpu); - } -} diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 5a8b5ebe1e..f80caf8b58 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -16,8 +16,8 @@ use hipfire_runtime::llama::{ use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; use hipfire_runtime::weight_backend::hfq_weight_dtype; use hipfire_runtime::weight_store::{ - AttachedWeightStore, TakenWeight, WeightHandle, WeightLoadTransaction, - WeightStoreAssembly, WeightStoreAssemblyGuard, WeightOrigin, + TakenWeight, WeightHandle, WeightLoadTransaction, WeightOrigin, WeightStoreAssembly, + WeightStoreAssemblyGuard, WeightStoreError, }; use rdna_compute::{DType, GpuTensor}; use std::collections::HashMap; @@ -52,6 +52,32 @@ pub struct LlamaBundle { /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } + +/// Crate-private attached owner for the manifest transaction. +/// +/// The runtime transaction stays public only long enough for the load carrier +/// to assemble or roll it back. Once wrapped here, the only consuming path is +/// the crate's [`hipfire_runtime::arch_model::ArchModel::free_gpu`] implementation. +pub(crate) struct AttachedWeightStore { + transaction: WeightLoadTransaction, +} + +impl AttachedWeightStore { + fn from_transaction( + transaction: WeightLoadTransaction, + expected: WeightOrigin, + ) -> Result { + if let Err(error) = transaction.validate_origin_value(expected) { + return Err((transaction, error)); + } + Ok(Self { transaction }) + } + + pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) { + self.transaction.rollback(gpu); + } +} + fn plan_single( config: &LlamaConfig, has_separate_lm_head: bool, @@ -697,9 +723,9 @@ pub use load_bundle as load_llama_bundle; impl LlamaBundle { /// Attach an unpublished load transaction after validating the complete - /// target identity. Publication creates the sole resident-store drain - /// capability; a rejected transaction is returned unchanged. - pub fn attach_weight_store( + /// target identity. The resulting owner is crate-private and can only be + /// consumed by `ArchModel::free_gpu`. + fn attach_weight_store( &mut self, transaction: WeightLoadTransaction, ) -> Result<(), (WeightLoadTransaction, String)> { @@ -709,10 +735,16 @@ impl LlamaBundle { "llama: weight store already attached".into(), )); } - let attached = match transaction.publish(self.weight_origin) { + let attached = match AttachedWeightStore::from_transaction( + transaction, + self.weight_origin, + ) { Ok(attached) => attached, Err((transaction, error)) => { - return Err((transaction, format!("llama: weight store origin rejected: {error}"))); + return Err(( + transaction, + format!("llama: weight store origin rejected: {error}"), + )); } }; self.weight_store = Some(attached); @@ -750,7 +782,10 @@ mod tests { use hipfire_runtime::kv_backend::KvBackend; use hipfire_runtime::kv_mode::KvMode; use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; - use hipfire_runtime::llama::{KvCache, KvCacheExt, KvDims, KvLayers, KvTarget}; + use hipfire_runtime::llama::{ + weight_gemv, KvCache, KvCacheExt, KvDims, KvLayers, KvTarget, + }; + use hipfire_runtime::weight_store::test_support; use hipfire_runtime::weight_store::{ WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, }; @@ -769,7 +804,20 @@ mod tests { fn f32_hfq_tensor(name: &str, shape: &[u32], malformed: bool) -> HfqMemTensor { let elements = shape.iter().map(|&dim| dim as usize).product::(); - hfq_tensor(name, shape, 2, if malformed { 4 } else { elements * 4 }) + let data = if malformed { + vec![0; 4] + } else { + (0..elements) + .flat_map(|value| ((value as f32) + 1.0).to_le_bytes()) + .collect() + }; + HfqMemTensor { + name: name.into(), + quant_type: 2, + shape: shape.to_vec(), + group_size: 0, + data, + } } fn fixture_hfq( @@ -902,6 +950,24 @@ mod tests { dtype: DType::F32, } } + + fn output_for( + gpu: &mut rdna_compute::Gpu, + weights: &LlamaWeights, + hidden: &[f32], + ) -> Vec { + let input = gpu + .upload_f32(hidden, &[hidden.len()]) + .expect("upload deterministic output input"); + let output = gpu + .alloc_tensor(&[weights.output.m], DType::F32) + .expect("allocate deterministic output"); + weight_gemv(gpu, &weights.output, &input, &output).expect("run output projection"); + let values = gpu.download_f32(&output).expect("download output projection"); + let _ = gpu.free_tensor(output); + let _ = gpu.free_tensor(input); + values + } #[test] fn single_plan_covers_every_typed_llama_handle() { let (mesh, plan) = plan_single(&config(), true).unwrap(); @@ -1040,42 +1106,67 @@ mod tests { } #[test] - fn production_post_resident_failure_returns_clean_load_error() { + fn production_post_resident_failure_reclaims_every_uploaded_allocation() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { return; }; - let (path, hfq) = fixture_hfq(false, false, true, false); + let (path, hfq) = fixture_hfq(false, false, false, false); + test_support::reset(); + test_support::arm_fail_after_upload(1); let cask = CaskConfig::default(); let mut ctx = load_ctx(&path, &mut gpu, &cask); let error = match load_bundle(ModelSource::Hfq(hfq), &mut ctx) { - Ok(_) => panic!("malformed output norm unexpectedly loaded"), + Ok(_) => panic!("post-upload fault unexpectedly succeeded"), Err(error) => error, }; drop(ctx); - assert!(error.contains("source payload") || error.contains("output_norm")); + test_support::clear_faults(); + assert!(error.contains("test fault injected after resident upload")); + let allocations = test_support::resident_allocations(); + assert!(allocations > 0, "fault must follow a resident upload"); + assert_eq!( + allocations, + test_support::resident_releases(), + "every resident allocation must be reclaimed on load failure" + ); std::fs::remove_file(path).expect("remove HFQ fixture"); } #[test] - fn production_manifest_matches_legacy_alias_contract() { + fn production_manifest_matches_legacy_numerical_output() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { return; }; let (path, hfq) = fixture_hfq(false, false, false, false); let cask = CaskConfig::default(); let mut ctx = load_ctx(&path, &mut gpu, &cask); - let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); - let manifest_alias = bundle.weights.lm_head_aliases_embd; + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) + .expect("load plain HFQ fixture through manifest path"); drop(ctx); + assert!(bundle.weights.lm_head_aliases_embd); + let hidden: Vec = (0..bundle.config.dim) + .map(|index| (index as f32 + 1.0) / 17.0) + .collect(); + let manifest_output = output_for(&mut gpu, &bundle.weights, &hidden); Box::new(bundle).free_gpu(&mut gpu); let hfq = HfqFile::open(&path).expect("reopen HFQ fixture"); let config = ::config_from_hfq(&hfq).expect("fixture config"); let legacy = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, &mut gpu) - .expect("legacy loader fixture"); - assert_eq!(manifest_alias, legacy.lm_head_aliases_embd); + .expect("load legacy HFQ fixture"); + assert!(legacy.lm_head_aliases_embd); assert_eq!(legacy.embd_format, EmbeddingFormat::F32); + let legacy_output = output_for(&mut gpu, &legacy, &hidden); legacy.free_gpu(&mut gpu); + + assert_eq!( + manifest_output, legacy_output, + "manifest and legacy output projections must agree for the same input" + ); + assert!( + manifest_output.iter().any(|value| *value != 0.0), + "parity assertion must observe a non-zero numerical output" + ); std::fs::remove_file(path).expect("remove HFQ fixture"); } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index 99ac6cc886..76b919b08c 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -11,8 +11,8 @@ //! //! The store is not a model owner. It has no `Drop` implementation and never //! frees GPU buffers implicitly. A carrier moves a committed transaction into -//! its existing `ArchModel` owner; that owner consumes the private drain -//! capability during the existing teardown path. +//! its existing `ArchModel` owner; that owner consumes the architecture-private +//! attached owner during the existing teardown path. //! `WeightStoreAssembly::take` transfers a resident handle to the owner that is //! assembling typed weights, and therefore removes the cell from the store's //! cleanup set. @@ -21,9 +21,64 @@ use hipfire_hardware::{DeviceMesh, MeshEpoch}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::collections::HashMap; -#[cfg(test)] thread_local! { + static RESIDENT_ALLOCATIONS: std::cell::Cell = + const { std::cell::Cell::new(0) }; static RESIDENT_RELEASES: std::cell::Cell = const { std::cell::Cell::new(0) }; + static FAIL_AFTER_UPLOAD: std::cell::Cell> = + const { std::cell::Cell::new(None) }; +} + +/// Test-only allocation accounting and deterministic post-upload fault seam. +/// +/// The production loader calls the same release path regardless of whether +/// this seam is armed. Callers should use [`reset`] before a scenario and +/// [`clear_faults`] after it so a failed test cannot poison a later one. +#[doc(hidden)] +pub mod test_support { + use super::{ + FAIL_AFTER_UPLOAD, RESIDENT_ALLOCATIONS, RESIDENT_RELEASES, + }; + + pub fn reset() { + RESIDENT_ALLOCATIONS.with(|count| count.set(0)); + RESIDENT_RELEASES.with(|count| count.set(0)); + clear_faults(); + } + + pub fn arm_fail_after_upload(upload_number: usize) { + assert!(upload_number > 0, "upload fault threshold must be non-zero"); + FAIL_AFTER_UPLOAD.with(|fault| fault.set(Some(upload_number))); + } + + pub fn clear_faults() { + FAIL_AFTER_UPLOAD.with(|fault| fault.set(None)); + } + + pub fn resident_allocations() -> usize { + RESIDENT_ALLOCATIONS.with(std::cell::Cell::get) + } + + pub fn resident_releases() -> usize { + RESIDENT_RELEASES.with(std::cell::Cell::get) + } + + pub(super) fn record_resident_upload() -> bool { + let allocation = RESIDENT_ALLOCATIONS.with(|count| { + let next = count.get() + 1; + count.set(next); + next + }); + FAIL_AFTER_UPLOAD.with(|fault| { + let should_fail = fault + .get() + .is_some_and(|upload_number| allocation >= upload_number); + if should_fail { + fault.set(None); + } + should_fail + }) + } } /// Stable logical placement identity. Layer is part of the key because a @@ -202,8 +257,8 @@ impl std::error::Error for FulfillError {} /// Load-side placement container. It records one immutable projection per /// `(name, layer, device)` and captures the target origin once. The container -/// itself has no consuming teardown API: lifecycle transitions are represented -/// by [`WeightLoadTransaction`] and [`AttachedWeightStore`]. +/// itself has no consuming teardown API; lifecycle transitions are represented +/// by [`WeightLoadTransaction`] and the architecture-private attached owner. #[derive(Default)] pub struct WeightStore { placements: HashMap, @@ -213,9 +268,9 @@ pub struct WeightStore { /// The only owner that may roll back resident allocations before publication. /// -/// A transaction owns the store until [`Self::publish`] transfers it into the -/// attached owner held by the architecture bundle. It deliberately has no -/// implicit `Drop` cleanup because the GPU is not available to a destructor. +/// A transaction owns the store until the architecture carrier consumes it +/// into its crate-private attached owner. It deliberately has no implicit +/// `Drop` cleanup because the GPU is not available to a destructor. pub struct WeightLoadTransaction { store: Option, } @@ -229,20 +284,6 @@ impl std::fmt::Debug for WeightLoadTransaction { } } -/// The resident-store capability returned by a committed load transaction. -/// -/// The backing store and its drain capability are private. Architecture -/// owners receive this value during attachment and consume it exactly once -/// during unload; no public `WeightStore` method can drain an attached store. -pub struct AttachedWeightStore { - store: WeightStore, - capability: WeightStoreDrainCapability, -} - -struct WeightStoreDrainCapability { - origin: WeightOrigin, -} - impl WeightLoadTransaction { pub fn new(store: WeightStore) -> Self { Self { store: Some(store) } @@ -294,6 +335,19 @@ impl WeightLoadTransaction { .map_or_else(Vec::new, |store| store.devices_for(name, layer)) } + /// Compare the unpublished transaction's captured target with an admitted + /// owner identity. This read-only check is used before the carrier wraps + /// the transaction in its private attached owner. + pub fn validate_origin_value( + &self, + expected: WeightOrigin, + ) -> Result<(), WeightStoreError> { + self.store.as_ref().map_or( + Err(WeightStoreError::UnboundOrigin), + |store| store.validate_origin_value(expected), + ) + } + /// Start typed assembly while this load is still unpublished. pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { self.store @@ -310,44 +364,6 @@ impl WeightLoadTransaction { store.rollback(gpu); } } - - /// Publish the store beneath an architecture owner after checking the - /// complete immutable target identity. A mismatch returns this - /// transaction unchanged so the caller can retry or roll it back. - pub fn publish( - mut self, - expected: WeightOrigin, - ) -> Result { - let store = self - .store - .take() - .expect("weight load transaction was already consumed"); - if let Err(error) = store.validate_origin_value(expected) { - self.store = Some(store); - return Err((self, error)); - } - Ok(AttachedWeightStore { - store, - capability: WeightStoreDrainCapability { origin: expected }, - }) - } -} - -impl AttachedWeightStore { - /// Drain resident handles through the private owner capability. The - /// capability is established only by `WeightLoadTransaction::publish`, so - /// origin mismatch is impossible after attachment. - pub fn drain(self, gpu: &Gpu) { - let Self { store, capability } = self; - capability.drain(store, gpu); - } -} - -impl WeightStoreDrainCapability { - fn drain(self, store: WeightStore, gpu: &Gpu) { - debug_assert_eq!(store.origin, Some(self.origin)); - store.release_unchecked(gpu); - } } impl WeightStore { @@ -507,7 +523,6 @@ impl WeightStore { // the existing loader's explicit owner teardown. The store // never relies on a destructor to release GPU memory. let _ = gpu.hip.free(tensor.buf); - #[cfg(test)] RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); } } @@ -650,7 +665,34 @@ where } let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); if let ShardPolicy::Tied { source: source_name } = &entry.policy { - let projection = projection_for(entry, 0, 1, entry.dtype); + let source_dtype = match store.get(source_name, entry.layer, 0) { + Some(WeightHandle::Resident(tensor)) => Some(tensor.dtype), + Some(WeightHandle::Alias(_)) | None => None, + }; + let Some(actual_dtype) = source_dtype else { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "tied source '{source_name}' is unresolved or has no actual resident dtype" + ), + }); + }; + if !entry.dtype_constraint.accepts(actual_dtype) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "tied source '{source_name}' actual dtype {actual_dtype:?} is excluded by constraint {:?}", + entry.dtype_constraint + ), + }); + } + let projection = projection_for(entry, 0, 1, actual_dtype); if let Err(reason) = store.insert( key, WeightHandle::Alias(source_name.clone()), @@ -735,6 +777,15 @@ where reason: reason.to_string(), }); } + if test_support::record_resident_upload() { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: "test fault injected after resident upload".into(), + }); + } } Ok(WeightLoadTransaction::new(store)) } @@ -758,7 +809,7 @@ where #[cfg(test)] mod tests { use super::*; - use crate::weight_manifest::{PinTarget, ShardPolicy}; + use crate::weight_manifest::{DTypeConstraint, PinTarget, ShardPolicy}; use hipfire_hardware::DimKind; fn projection(dtype: DType) -> WeightProjection { @@ -869,6 +920,48 @@ mod tests { assert_eq!(placement_devices(&entry, &mesh, 1), vec![0]); } + #[test] + fn tied_projection_preserves_fulfilled_source_dtype() { + let Ok(gpu) = Gpu::init() else { + return; + }; + let mesh = DeviceMesh::single(); + let constraint = DTypeConstraint::source_from_sources(vec![DType::F16, DType::F32]); + let source = WeightEntry::model_with_dtype_constraint( + "source", + vec![1], + DType::F16, + constraint.clone(), + ShardPolicy::Replicate, + ); + let alias = WeightEntry::model_with_dtype_constraint( + "alias", + vec![1], + DType::F16, + constraint, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + let transaction = fulfill_manifest_single( + &[source, alias], + &mesh, + 1, + &gpu, + |_| Ok((vec![0; 4], DType::F32)), + ) + .unwrap(); + assert_eq!( + transaction.projection("alias", None, 0).unwrap().dtype, + DType::F32 + ); + assert!(matches!( + transaction.get("alias", None, 0), + Some(WeightHandle::Alias(source)) if source == "source" + )); + transaction.rollback(&gpu); + } + #[test] fn successful_single_fulfillment_commits_resident_projection() { let Ok(gpu) = Gpu::init() else { @@ -893,7 +986,7 @@ mod tests { } #[test] - fn full_origin_mismatch_returns_unpublished_transaction_unchanged() { + fn full_origin_mismatch_leaves_unpublished_transaction_unchanged() { let first = DeviceMesh::single(); let second = DeviceMesh::single(); let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); @@ -903,10 +996,7 @@ mod tests { .stage_alias("resident", None, 0, "source", projection(DType::F16)) .unwrap(); let transaction = WeightLoadTransaction::new(store); - let (transaction, error) = match transaction.publish(expected) { - Ok(_) => panic!("origin mismatch unexpectedly succeeded"), - Err(value) => value, - }; + let error = transaction.validate_origin_value(expected).unwrap_err(); assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); assert_eq!(transaction.origin(), Some(actual)); assert!(transaction.contains("resident", None, 0)); @@ -926,10 +1016,7 @@ mod tests { }) .unwrap(); let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); - let (transaction, error) = match transaction.publish(expected) { - Ok(_) => panic!("origin mismatch unexpectedly succeeded"), - Err(value) => value, - }; + let error = transaction.validate_origin_value(expected).unwrap_err(); assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); assert_eq!(transaction.len(), 1); assert_eq!( @@ -941,25 +1028,6 @@ mod tests { assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } - #[test] - fn attached_owner_transfer_is_consuming_and_empty_transfer_is_safe() { - let mesh = DeviceMesh::single(); - let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); - let mut store = WeightStore::with_origin(origin); - store - .stage_alias("owned", None, 0, "source", projection(DType::F16)) - .unwrap(); - let transaction = WeightLoadTransaction::new(store); - let attached = transaction.publish(origin).unwrap(); - let Ok(gpu) = Gpu::init() else { - return; - }; - attached.drain(&gpu); - let empty = WeightLoadTransaction::new(WeightStore::with_origin(origin)) - .publish(origin) - .unwrap(); - empty.drain(&gpu); - } #[test] fn source_failure_after_resident_upload_rolls_back_everything() { From f9cb8ec1de06089d5c57ba788a6e3dcfe406312f Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 20:30:48 +0200 Subject: [PATCH 11/25] fix(device-mesh): repair llama manifest APIs --- crates/hipfire-arch-llama/src/arch_model.rs | 4 +- crates/hipfire-arch-llama/src/carrier.rs | 151 ++++++++------ crates/hipfire-runtime/src/model_load.rs | 6 +- crates/hipfire-runtime/src/weight_manifest.rs | 34 +++- crates/hipfire-runtime/src/weight_store.rs | 186 ++++++++++++------ 5 files changed, 255 insertions(+), 126 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index f21e75b725..1de6953da5 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -54,7 +54,9 @@ impl ArchModel for LlamaBundle { // Attachment already checked the complete origin and created this // owner capability. There is no mismatch branch to leak the model: // an attached store can only be drained by this consuming owner. - store.drain(gpu); + if let Err(error) = store.drain(gpu) { + eprintln!("llama: failed to release attached weight store: {error}"); + } } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index f80caf8b58..a47bd13b99 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -73,8 +73,18 @@ impl AttachedWeightStore { Ok(Self { transaction }) } - pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) { - self.transaction.rollback(gpu); + pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) -> hip_bridge::HipResult<()> { + self.transaction.rollback(gpu) + } +} + +fn with_weight_rollback_error( + reason: String, + rollback: hip_bridge::HipResult<()>, +) -> String { + match rollback { + Ok(()) => reason, + Err(error) => format!("{reason}; resident rollback failed: {error}"), } } @@ -82,7 +92,7 @@ fn plan_single( config: &LlamaConfig, has_separate_lm_head: bool, ) -> Result<(DeviceMesh, ManifestPlan), String> { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().map_err(|error| format!("llama: device mesh: {error}"))?; let manifest = Llama::weight_manifest_for_hfq(config, has_separate_lm_head); let state = Llama::state_manifest(config); let plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) @@ -566,8 +576,10 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result weights, Err(error) => { - transaction.rollback(ctx.gpu); - return Err(error); + return Err(with_weight_rollback_error( + error, + transaction.rollback(ctx.gpu), + )); } }; (weights, Some(transaction)) @@ -581,12 +593,17 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result scratch, Err(error) => { - if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu); - } + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + return Err(with_weight_rollback_error( + format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + ), + rollback, )); } }; @@ -604,12 +621,17 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result kv, Err(error) => { scratch.free_gpu(ctx.gpu); - if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu); - } + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ::from_mode failed: {error}" + return Err(with_weight_rollback_error( + format!( + "llama: ::from_mode failed: {error}" + ), + rollback, )); } }; @@ -708,11 +730,11 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Vec { - let input = gpu - .upload_f32(hidden, &[hidden.len()]) - .expect("upload deterministic output input"); - let output = gpu - .alloc_tensor(&[weights.output.m], DType::F32) - .expect("allocate deterministic output"); - weight_gemv(gpu, &weights.output, &input, &output).expect("run output projection"); - let values = gpu.download_f32(&output).expect("download output projection"); - let _ = gpu.free_tensor(output); - let _ = gpu.free_tensor(input); - values - } #[test] fn single_plan_covers_every_typed_llama_handle() { let (mesh, plan) = plan_single(&config(), true).unwrap(); @@ -984,7 +990,7 @@ mod tests { #[test] fn typed_assembly_rolls_back_when_a_cell_is_not_resident() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); let mut store = WeightStore::with_origin(origin); for name in ["token_embd", "output_norm", "lm_head"] { @@ -1133,43 +1139,78 @@ mod tests { } #[test] - fn production_manifest_matches_legacy_numerical_output() { + fn production_manifest_matches_legacy_forward_logits() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { return; }; let (path, hfq) = fixture_hfq(false, false, false, false); let cask = CaskConfig::default(); let mut ctx = load_ctx(&path, &mut gpu, &cask); - let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) - .expect("load plain HFQ fixture through manifest path"); + let mut bundle = + load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); drop(ctx); - assert!(bundle.weights.lm_head_aliases_embd); - let hidden: Vec = (0..bundle.config.dim) - .map(|index| (index as f32 + 1.0) / 17.0) - .collect(); - let manifest_output = output_for(&mut gpu, &bundle.weights, &hidden); + + let manifest_logits = { + forward_scratch_embed( + &mut gpu, + &bundle.weights, + &bundle.config, + 1, + 0, + &bundle.scratch, + ) + .expect("manifest embedding forward"); + forward_scratch_compute( + &mut gpu, + &bundle.weights, + &bundle.config, + 0, + &mut bundle.kv, + &bundle.scratch, + ) + .expect("manifest model forward"); + gpu.download_f32(&bundle.scratch.logits) + .expect("download manifest logits") + }; Box::new(bundle).free_gpu(&mut gpu); let hfq = HfqFile::open(&path).expect("reopen HFQ fixture"); let config = ::config_from_hfq(&hfq).expect("fixture config"); let legacy = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, &mut gpu) .expect("load legacy HFQ fixture"); - assert!(legacy.lm_head_aliases_embd); - assert_eq!(legacy.embd_format, EmbeddingFormat::F32); - let legacy_output = output_for(&mut gpu, &legacy, &hidden); + let scratch = ForwardScratch::new_with_max_seq(&mut gpu, &config, 8) + .expect("allocate legacy forward scratch"); + let dims = llama_kv_dims(&config, 8, None); + let mut kv = ::from_mode( + KvMode::Q8, + KvTarget::Single(&mut gpu), + &dims, + ) + .expect("allocate legacy KV cache"); + forward_scratch_embed(&mut gpu, &legacy, &config, 1, 0, &scratch) + .expect("legacy embedding forward"); + forward_scratch_compute(&mut gpu, &legacy, &config, 0, &mut kv, &scratch) + .expect("legacy model forward"); + let legacy_logits = gpu + .download_f32(&scratch.logits) + .expect("download legacy logits"); + scratch.free_gpu(&mut gpu); + let _ = kv.free_gpu(&mut gpu); legacy.free_gpu(&mut gpu); - assert_eq!( - manifest_output, legacy_output, - "manifest and legacy output projections must agree for the same input" - ); - assert!( - manifest_output.iter().any(|value| *value != 0.0), - "parity assertion must observe a non-zero numerical output" - ); + assert_eq!(manifest_logits.len(), legacy_logits.len()); + for (index, (manifest, legacy)) in + manifest_logits.iter().zip(&legacy_logits).enumerate() + { + assert!( + (manifest - legacy).abs() <= 1e-5, + "logit mismatch at index {index}: manifest={manifest} legacy={legacy}" + ); + } std::fs::remove_file(path).expect("remove HFQ fixture"); } + #[test] fn physical_cap_is_honored_by_upstream_kv_constructor() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index d233630a62..26178a1229 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -112,10 +112,13 @@ pub trait WeightSource { fn n_layers(&self) -> usize; /// Pre-load hook. HFQ drops the mmap when n==1; PaRo rejects n>1; llama no-op. fn prepare(&mut self, n_devices: usize) -> HipResult<()>; + fn read_embed(&mut self, gpu: &mut Gpu) -> HipResult<(GpuTensor, EmbeddingFormat)>; + fn read_final_norm(&mut self, gpu: &mut Gpu) -> HipResult; /// `can_alias` is true iff embed and output share a device (n==1); the /// implementation decides whether to use it (single-device LLaMA and /// qwen35 alias tied embeddings; multi-device routes re-materialize). fn read_output( + &mut self, gpu: &mut Gpu, embd: &GpuTensor, embd_fmt: EmbeddingFormat, @@ -177,7 +180,8 @@ mod tests { #[test] fn mesh_layout_selects_stage_rank_zero_without_io() { - let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]) + .expect("small test mesh construction cannot overflow"); let layout = Layout::from_mesh(&mesh, 4); assert_eq!(layout.output_device(), 2); assert_eq!( diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 044bc62c6b..1a4f62a7d6 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -912,7 +912,8 @@ mod tests { #[test] fn placement_and_boundaries_use_named_mesh() { - let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]) + .expect("small test mesh construction cannot overflow"); let embed = WeightEntry::model( "token_embd", vec![32, 8], @@ -955,7 +956,8 @@ mod tests { #[test] fn validation_covers_divisibility_ties_and_expert_shape() { - let tp3 = DeviceMesh::rect(&[(DimKind::Tp, 3)]); + let tp3 = DeviceMesh::rect(&[(DimKind::Tp, 3)]) + .expect("small test mesh construction cannot overflow"); assert!(validate_manifest( &[layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 })], &tp3 @@ -977,7 +979,11 @@ mod tests { }, ), ]; - assert!(validate_manifest(&tied, &DeviceMesh::single()).is_ok()); + assert!(validate_manifest( + &tied, + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .is_ok()); let bad_expert = WeightEntry::layer( "experts", 0, @@ -988,7 +994,11 @@ mod tests { assign: ExpertAssign::Stride, }, ); - assert!(validate_manifest(&[bad_expert], &DeviceMesh::single()).is_err()); + assert!(validate_manifest( + &[bad_expert], + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .is_err()); } #[test] @@ -1071,7 +1081,7 @@ mod tests { #[test] fn planning_rejects_weight_layer_at_n_layers_and_accepts_last_layer() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let valid = layer_entry("w", 2, ShardPolicy::Replicate); assert!(plan_manifest(&[valid], &[], &mesh, 3).is_ok()); let out_of_range = layer_entry("w", 3, ShardPolicy::Replicate); @@ -1097,7 +1107,7 @@ mod tests { ); assert!(validate_manifest( &[source.clone(), shape_mismatch], - &DeviceMesh::single() + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") ) .is_err()); @@ -1111,7 +1121,7 @@ mod tests { ); assert!(validate_manifest( &[source.clone(), dtype_mismatch], - &DeviceMesh::single() + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") ) .is_err()); @@ -1133,7 +1143,7 @@ mod tests { ); assert!(validate_manifest( &[source, chained_source, chain], - &DeviceMesh::single() + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") ) .is_err()); @@ -1153,7 +1163,7 @@ mod tests { source: "cycle_a".into(), }, ); - assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single()).is_err()); + assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single().expect("single-device mesh construction cannot overflow")).is_err()); } #[test] fn tied_entries_reject_different_source_sets_with_equal_logical_dtype() { @@ -1172,7 +1182,11 @@ mod tests { source: "source".into(), }, ); - let error = validate_manifest(&[source, tied], &DeviceMesh::single()).unwrap_err(); + let error = validate_manifest( + &[source, tied], + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .unwrap_err(); assert!(error.contains("source dtype contract")); } } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index 76b919b08c..ce0019441c 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -358,10 +358,13 @@ impl WeightLoadTransaction { /// Consume this transaction and release every resident handle it owns. /// This is intentionally the only rollback operation exposed by the - /// lifecycle API. - pub fn rollback(mut self, gpu: &Gpu) { + /// lifecycle API. Successful frees are reflected in the resident-release + /// accounting; any failed HIP free is returned to the caller. + pub fn rollback(mut self, gpu: &Gpu) -> hip_bridge::HipResult<()> { if let Some(store) = self.store.take() { - store.rollback(gpu); + store.rollback(gpu) + } else { + Ok(()) } } } @@ -512,20 +515,30 @@ impl WeightStore { /// Explicit rollback for a failed transaction. It consumes the partial /// store and frees every resident buffer on the single owning GPU. - fn rollback(self, gpu: &Gpu) { - self.release_unchecked(gpu); + fn rollback(self, gpu: &Gpu) -> hip_bridge::HipResult<()> { + self.release_unchecked(gpu) } - fn release_unchecked(self, gpu: &Gpu) { + fn release_unchecked(self, gpu: &Gpu) -> hip_bridge::HipResult<()> { + let mut first_error = None; for handle in self.placements.into_values() { if let WeightHandle::Resident(tensor) = handle { - // Rollback is deliberately direct and best-effort, matching - // the existing loader's explicit owner teardown. The store - // never relies on a destructor to release GPU memory. - let _ = gpu.hip.free(tensor.buf); - RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); + match gpu.hip.free(tensor.buf) { + Ok(()) => { + RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); + } + Err(error) => { + if first_error.is_none() { + first_error = Some(error); + } + } + } } } + match first_error { + Some(error) => Err(error), + None => Ok(()), + } } } @@ -615,6 +628,18 @@ fn target_error(mesh: &DeviceMesh) -> Option { ), }) } +fn rollback_fulfill_error( + store: WeightStore, + gpu: &Gpu, + mut error: FulfillError, +) -> FulfillError { + if let Err(release_error) = store.rollback(gpu) { + error + .reason + .push_str(&format!("; resident rollback failed: {release_error}")); + } + error +} /// Fulfill a manifest for a plain LLaMA Single target. /// @@ -660,8 +685,7 @@ where devices ), }; - store.rollback(gpu); - return Err(error); + return Err(rollback_fulfill_error(store, gpu, error)); } let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); if let ShardPolicy::Tied { source: source_name } = &entry.policy { @@ -670,19 +694,18 @@ where Some(WeightHandle::Alias(_)) | None => None, }; let Some(actual_dtype) = source_dtype else { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: format!( "tied source '{source_name}' is unresolved or has no actual resident dtype" ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); }; if !entry.dtype_constraint.accepts(actual_dtype) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, @@ -690,7 +713,8 @@ where "tied source '{source_name}' actual dtype {actual_dtype:?} is excluded by constraint {:?}", entry.dtype_constraint ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } let projection = projection_for(entry, 0, 1, actual_dtype); if let Err(reason) = store.insert( @@ -698,13 +722,13 @@ where WeightHandle::Alias(source_name.clone()), projection, ) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: reason.to_string(), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } continue; } @@ -712,18 +736,17 @@ where let (bytes, dtype) = match source(entry) { Ok(value) => value, Err(reason) => { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: format!("source read failed: {reason}"), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } }; if !entry.dtype_constraint.accepts(dtype) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, @@ -731,7 +754,8 @@ where "source dtype {dtype:?} violates constraint {:?}", entry.dtype_constraint ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } if matches!(dtype, DType::F32 | DType::F16 | DType::BF16) { let expected_bytes = entry @@ -740,8 +764,7 @@ where .try_fold(1usize, |count, &dim| count.checked_mul(dim)) .and_then(|elements| elements.checked_mul(dtype.size())); if expected_bytes != Some(bytes.len()) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, @@ -751,40 +774,41 @@ where expected_bytes, entry.logical_shape ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } } let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { Ok(tensor) => tensor, Err(error) => { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: format!("upload_raw failed: {error}"), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } }; tensor.dtype = dtype; let projection = projection_for(entry, 0, 1, dtype); if let Err(reason) = store.insert(key, WeightHandle::Resident(tensor), projection) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: reason.to_string(), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } if test_support::record_resident_upload() { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: "test fault injected after resident upload".into(), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } } Ok(WeightLoadTransaction::new(store)) @@ -825,8 +849,8 @@ mod tests { #[test] fn origin_mismatch_is_detected_before_gpu_release() { - let first = DeviceMesh::single(); - let second = DeviceMesh::single(); + let first = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let second = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let actual = WeightOrigin::from_parts(first.epoch(), 0, 0); let expected = WeightOrigin::from_parts(second.epoch(), 0, 0); let store = WeightStore::with_origin(actual); @@ -842,7 +866,7 @@ mod tests { #[test] fn staged_rollback_removes_handles_and_projection_together() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); let mut store = WeightStore::with_origin(origin); store @@ -863,7 +887,7 @@ mod tests { #[test] fn assembly_drop_restores_staged_handles() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); store .stage_alias("x", None, 0, "source", projection(DType::F16)) @@ -871,7 +895,8 @@ mod tests { { let mut assembly = store.begin_assembly(); assert_eq!(assembly.take("x", None, 0), Some(0)); - assert!(assembly.get(0).is_some()); + let guard = assembly.commit(); + assert!(guard.get(0).is_some()); } assert!(store.contains("x", None, 0)); assert!(store.projection("x", None, 0).is_some()); @@ -879,7 +904,7 @@ mod tests { #[test] fn repeated_unload_lookup_cannot_reclaim_a_transferred_cell() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); store .stage_alias("x", None, 0, "source", projection(DType::F16)) @@ -892,7 +917,7 @@ mod tests { #[test] fn duplicate_projection_is_rejected_without_replacing_identity() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); store .stage_alias("x", None, 0, "source-a", projection(DType::F16)) @@ -907,7 +932,8 @@ mod tests { #[test] fn single_target_refuses_multi_device_before_source_or_gpu_work() { - let mesh = DeviceMesh::rect(&[(DimKind::Tp, 2)]); + let mesh = DeviceMesh::rect(&[(DimKind::Tp, 2)]) + .expect("small test mesh construction cannot overflow"); let entry = WeightEntry::model( "embed", vec![2, 2], @@ -925,7 +951,7 @@ mod tests { let Ok(gpu) = Gpu::init() else { return; }; - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let constraint = DTypeConstraint::source_from_sources(vec![DType::F16, DType::F32]); let source = WeightEntry::model_with_dtype_constraint( "source", @@ -959,7 +985,9 @@ mod tests { transaction.get("alias", None, 0), Some(WeightHandle::Alias(source)) if source == "source" )); - transaction.rollback(&gpu); + transaction + .rollback(&gpu) + .expect("resident transaction rollback must free its HIP allocation"); } #[test] @@ -967,7 +995,7 @@ mod tests { let Ok(gpu) = Gpu::init() else { return; }; - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) @@ -982,13 +1010,15 @@ mod tests { transaction.projection("resident", None, 0).unwrap().dtype, DType::F32 ); - transaction.rollback(&gpu); + transaction + .rollback(&gpu) + .expect("resident transaction rollback must free its HIP allocation"); } #[test] fn full_origin_mismatch_leaves_unpublished_transaction_unchanged() { - let first = DeviceMesh::single(); - let second = DeviceMesh::single(); + let first = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let second = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); let expected = WeightOrigin::from_parts(second.epoch(), 4, 12); let mut store = WeightStore::with_origin(actual); @@ -1009,7 +1039,7 @@ mod tests { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) @@ -1024,18 +1054,56 @@ mod tests { 0, "origin rejection must not free resident buffers" ); - transaction.rollback(&gpu); + transaction + .rollback(&gpu) + .expect("resident transaction rollback must free its HIP allocation"); assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } + #[test] + fn rollback_reports_free_failure_without_counting_release() { + let Ok(gpu) = Gpu::init() else { + return; + }; + test_support::reset(); + RESIDENT_ALLOCATIONS.with(|count| count.set(1)); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, gpu.device_id); + let mut store = WeightStore::with_origin(origin); + let borrowed = GpuTensor { + buf: unsafe { + hip_bridge::DeviceBuffer::from_raw( + std::ptr::null_mut::(), + 0, + ) + }, + shape: vec![0], + dtype: DType::F32, + }; + store + .insert( + WeightPlacementKey::new("borrowed", None, 0), + WeightHandle::Resident(borrowed), + projection(DType::F32), + ) + .expect("insert borrowed resident test handle"); + let error = WeightLoadTransaction::new(store) + .rollback(&gpu) + .expect_err("rollback must surface a failed HIP free"); + assert!(error.message.contains("borrowed")); + assert_eq!(test_support::resident_allocations(), 1); + assert_eq!(test_support::resident_releases(), 0); + test_support::reset(); + } + #[test] fn source_failure_after_resident_upload_rolls_back_everything() { let Ok(gpu) = Gpu::init() else { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entries = vec![ WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), @@ -1063,7 +1131,7 @@ mod tests { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let constraint = DTypeConstraint::source_exact(DType::F32); let entries = vec![ WeightEntry::model_with_dtype_constraint( @@ -1104,7 +1172,7 @@ mod tests { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entries = vec![ WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), From df5b8f89953ff97e1d66f647b13c4ba927907c1d Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 00:11:27 +0200 Subject: [PATCH 12/25] fix(device-mesh): restore llama manifest integration --- Cargo.lock | 1 + crates/hipfire-arch-llama/src/carrier.rs | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 0f90443e0c..9692607f07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1196,6 +1196,7 @@ version = "0.3.0" dependencies = [ "hip-bridge", "hipfire-dispatch", + "hipfire-hardware", "hipfire-runtime", "rdna-compute", ] diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index a47bd13b99..7725ab82fe 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -14,6 +14,8 @@ use hipfire_runtime::llama::{ LlamaConfig, LlamaWeights, WeightTensor, }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; +use hipfire_runtime::model_source::ModelSource as ModelSourceTrait; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_backend::hfq_weight_dtype; use hipfire_runtime::weight_store::{ TakenWeight, WeightHandle, WeightLoadTransaction, WeightOrigin, WeightStoreAssembly, @@ -27,6 +29,8 @@ pub struct LlamaBundle { pub weights: LlamaWeights, pub scratch: ForwardScratch, pub kv: KvCache, + /// The admitted mesh that owns this plan and the attached store origin. + pub(crate) mesh: DeviceMesh, /// Pure declaration/placement plan captured at load time. The plan has no /// GPU handles and is immutable after publication. pub manifest_plan: ManifestPlan, From d7d18cac0c0ab4d8e5cd6a7177ac665b2272cc79 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 00:29:14 +0200 Subject: [PATCH 13/25] fix(device-mesh): correct manifest pilot contracts --- crates/hipfire-arch-llama/src/carrier.rs | 45 +++++++++++-------- crates/hipfire-runtime/src/weight_manifest.rs | 4 +- 2 files changed, 27 insertions(+), 22 deletions(-) diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 7725ab82fe..b85a71a3d9 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -813,6 +813,7 @@ mod tests { KvTarget, }; use hipfire_runtime::weight_store::test_support; + use hipfire_runtime::weight_manifest::ShardPolicy; use hipfire_runtime::weight_store::{ WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, }; @@ -846,6 +847,22 @@ mod tests { data, } } + fn f16_hfq_tensor(name: &str, shape: &[u32]) -> HfqMemTensor { + let elements = shape.iter().map(|&dim| dim as usize).product::(); + HfqMemTensor { + name: name.into(), + quant_type: 1, + shape: shape.to_vec(), + group_size: 0, + data: (0..elements) + .flat_map(|index| { + let bits = if index % 2 == 0 { 0x3c00u16 } else { 0x3800u16 }; + bits.to_le_bytes() + }) + .collect(), + } + } + fn fixture_hfq( with_awq_sidecar: bool, @@ -856,13 +873,13 @@ mod tests { let mut tensors = vec![ f32_hfq_tensor("model.embed_tokens.weight", &[2, 32], false), f32_hfq_tensor("model.norm.weight", &[32], false), - f32_hfq_tensor("model.layers.0.self_attn.q_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.self_attn.k_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.self_attn.v_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.self_attn.o_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32], false), - f32_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32], false), - f32_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64], false), + f16_hfq_tensor("model.layers.0.self_attn.q_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.k_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.v_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.o_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32]), + f16_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32]), + f16_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64]), f32_hfq_tensor( "model.layers.0.input_layernorm.weight", &[32], @@ -1081,19 +1098,9 @@ mod tests { } #[test] - fn production_awq_route_preserves_legacy_loader() { - let Ok(mut gpu) = rdna_compute::Gpu::init() else { - return; - }; + fn production_awq_sidecar_selects_legacy_loader() { let (path, hfq) = fixture_hfq(true, false, false, false); - let cask = CaskConfig::default(); - let mut ctx = load_ctx(&path, &mut gpu, &cask); - let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) - .expect("AWQ fixture must use the legacy HFQ loader"); - drop(ctx); - assert!(bundle.weight_store.is_none()); - assert!(bundle.weights.lm_head_aliases_embd); - Box::new(bundle).free_gpu(&mut gpu); + assert_eq!(classify_hfq_route(&hfq), HfqLoadRoute::LegacyAwq); std::fs::remove_file(path).expect("remove HFQ fixture"); } diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 1a4f62a7d6..65a9172f48 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -484,8 +484,6 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< if !source_entry .dtype_constraint .same_source_set(&entry.dtype_constraint) - || !entry.dtype_constraint.accepts(source_entry.dtype) - || !source_entry.dtype_constraint.accepts(entry.dtype) { return Err(format!( "{context}: tied source '{source}' violates the source dtype contract" @@ -920,7 +918,7 @@ mod tests { DType::F16, ShardPolicy::Pin(PinTarget::Embed), ); - let row = layer_entry("wo", 1, ShardPolicy::RowShard { axis: 1 }); + let row = layer_entry("wo", 2, ShardPolicy::RowShard { axis: 1 }); assert_eq!(placement_devices(&embed, &mesh, 4), vec![0]); assert_eq!(placement_devices(&row, &mesh, 4), vec![2, 3]); let plan = plan_manifest( From 3c29ecf75cb44715b6263b9ad908862614c637f7 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 00:32:58 +0200 Subject: [PATCH 14/25] style(device-mesh): format manifest pilot --- crates/hipfire-arch-llama/src/arch.rs | 11 +- crates/hipfire-arch-llama/src/arch_model.rs | 1 - crates/hipfire-arch-llama/src/carrier.rs | 395 ++++++++---------- crates/hipfire-runtime/src/model_load.rs | 3 +- crates/hipfire-runtime/src/weight_manifest.rs | 91 ++-- crates/hipfire-runtime/src/weight_store.rs | 119 ++---- 6 files changed, 263 insertions(+), 357 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index 27c60aa9e6..d02289ef0e 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -17,8 +17,8 @@ use hip_bridge::HipResult; use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::{self, HfqFile}; -use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; +use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::weight_manifest::{ DTypeConstraint, FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, }; @@ -293,7 +293,14 @@ impl Llama { /// Pure state declaration for the full-attention LLaMA family. pub fn state_manifest(cfg: &LlamaConfig) -> Vec { (0..cfg.n_layers) - .map(|layer| StateEntry::new(StateKind::Kv { quant: String::new() }, layer)) + .map(|layer| { + StateEntry::new( + StateKind::Kv { + quant: String::new(), + }, + layer, + ) + }) .collect() } } diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 1de6953da5..7f374f20c7 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -62,4 +62,3 @@ impl ArchModel for LlamaBundle { let _ = kv.free_gpu(gpu); } } - diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index b85a71a3d9..a6bd292c02 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -15,8 +15,8 @@ use hipfire_runtime::llama::{ }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; use hipfire_runtime::model_source::ModelSource as ModelSourceTrait; -use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_backend::hfq_weight_dtype; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_store::{ TakenWeight, WeightHandle, WeightLoadTransaction, WeightOrigin, WeightStoreAssembly, WeightStoreAssemblyGuard, WeightStoreError, @@ -82,10 +82,7 @@ impl AttachedWeightStore { } } -fn with_weight_rollback_error( - reason: String, - rollback: hip_bridge::HipResult<()>, -) -> String { +fn with_weight_rollback_error(reason: String, rollback: hip_bridge::HipResult<()>) -> String { match rollback { Ok(()) => reason, Err(error) => format!("{reason}; resident rollback failed: {error}"), @@ -213,10 +210,8 @@ fn f32_bytes_from_hfq(quant_type: u8, data: &[u8], name: &str) -> Result } for chunk in chunks { bytes.extend_from_slice( - &f32::from_bits( - u16::from_le_bytes([chunk[0], chunk[1]]) as u32 * (1 << 16), - ) - .to_le_bytes(), + &f32::from_bits(u16::from_le_bytes([chunk[0], chunk[1]]) as u32 * (1 << 16)) + .to_le_bytes(), ); } } @@ -244,12 +239,11 @@ fn hfq_source(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, DType), St )), }; } - if matches!(entry.name.as_str(), "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm") - { - return Ok(( - f32_bytes_from_hfq(quant_type, &data, &name)?, - DType::F32, - )); + if matches!( + entry.name.as_str(), + "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm" + ) { + return Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)); } match quant_type { 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), @@ -367,9 +361,8 @@ fn assemble_llama_weights( ) -> Result { let mut assembly = transaction.begin_assembly(); let mut slots = HashMap::new(); - let mut take = |name: &str, layer: Option| { - take_slot(&mut assembly, &mut slots, name, layer) - }; + let mut take = + |name: &str, layer: Option| take_slot(&mut assembly, &mut slots, name, layer); take("token_embd", None)?; take("output_norm", None)?; @@ -431,13 +424,7 @@ fn assemble_llama_weights( config.dim, ) } else { - resident_weight( - &mut cells, - "lm_head", - None, - config.vocab_size, - config.dim, - ) + resident_weight(&mut cells, "lm_head", None, config.vocab_size, config.dim) }; let mut layers = Vec::with_capacity(config.n_layers); for layer in 0..config.n_layers { @@ -544,174 +531,159 @@ fn classify_hfq_route(hfq: &HfqFile) -> HfqLoadRoute { } pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { - let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = - match src { - ModelSource::Hfq(hfq) => { - let config = - ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; - // Admission and route classification are pure source checks. - // They must run before any manifest fulfillment or GPU upload. - hipfire_runtime::hfq::validate_llama_hfq_admission(&hfq) - .map_err(|e| e.to_string())?; - let has_separate_lm_head = - hfq.find_tensor_info("lm_head.weight").is_some(); - let route = classify_hfq_route(&hfq); - eprintln!("llama: HFQ source route = {route:?}"); - let (mesh, manifest_plan) = plan_single(&config, has_separate_lm_head)?; - let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); - let (weights, mut weight_store) = match route { - HfqLoadRoute::LegacyAwq => { - let weights = - hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, ctx.gpu) - .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}"))?; - (weights, None) - } - HfqLoadRoute::ManifestPlainLlama => { - let manifest = - Llama::weight_manifest_for_hfq(&config, has_separate_lm_head); - let mut transaction = hipfire_runtime::weight_store::fulfill_manifest( - &manifest, - &mesh, - config.n_layers, - ctx.gpu, - |entry| hfq_source(&hfq, entry), - ) - .map_err(|e| format!("llama: {e}"))?; - let weights = match assemble_llama_weights(&config, &mut transaction) { - Ok(weights) => weights, - Err(error) => { - return Err(with_weight_rollback_error( - error, - transaction.rollback(ctx.gpu), - )); - } - }; - (weights, Some(transaction)) - } - }; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - // The plain LLaMA path has no independent cap resolver. PR - // #661's physical-cap behavior is owned by the existing - // upstream KV plan. - let scratch = - match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, + let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = match src + { + ModelSource::Hfq(hfq) => { + let config = + ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; + // Admission and route classification are pure source checks. + // They must run before any manifest fulfillment or GPU upload. + hipfire_runtime::hfq::validate_llama_hfq_admission(&hfq).map_err(|e| e.to_string())?; + let has_separate_lm_head = hfq.find_tensor_info("lm_head.weight").is_some(); + let route = classify_hfq_route(&hfq); + eprintln!("llama: HFQ source route = {route:?}"); + let (mesh, manifest_plan) = plan_single(&config, has_separate_lm_head)?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let (weights, mut weight_store) = match route { + HfqLoadRoute::LegacyAwq => { + let weights = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, ctx.gpu) + .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}"))?; + (weights, None) + } + HfqLoadRoute::ManifestPlainLlama => { + let manifest = Llama::weight_manifest_for_hfq(&config, has_separate_lm_head); + let mut transaction = hipfire_runtime::weight_store::fulfill_manifest( + &manifest, + &mesh, + config.n_layers, + ctx.gpu, + |entry| hfq_source(&hfq, entry), + ) + .map_err(|e| format!("llama: {e}"))?; + let weights = match assemble_llama_weights(&config, &mut transaction) { + Ok(weights) => weights, Err(error) => { - let rollback = if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu) - } else { - Ok(()) - }; - weights.free_gpu(ctx.gpu); return Err(with_weight_rollback_error( - format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" - ), - rollback, + error, + transaction.rollback(ctx.gpu), )); } }; - let dims = llama_kv_dims(&config, ctx.max_seq, None); - let kv = match ::from_mode( - hipfire_runtime::kv_mode::resolve( - ctx.kv_mode_override.unwrap_or(""), - &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, - config.head_dim, - ) - .mode, - KvTarget::Single(ctx.gpu), - &dims, - ) { - Ok(kv) => kv, - Err(error) => { - scratch.free_gpu(ctx.gpu); - let rollback = if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu) - } else { - Ok(()) - }; - weights.free_gpu(ctx.gpu); - return Err(with_weight_rollback_error( - format!( - "llama: ::from_mode failed: {error}" - ), - rollback, - )); - } - }; - ( - config, - weights, - kv, - scratch, - manifest_plan, - weight_store, - mesh, - weight_origin, - ) - } - ModelSource::Dir(source) => { - let config = - hipfire_runtime::hfq::config_from_safetensors_llama(&source).map_err(|e| { - format!("failed to parse LLaMA/Qwen3 config from config.json: {e}") - })?; - let (mesh, manifest_plan) = - plan_single(&config, source.tensor_info("lm_head.weight").is_some())?; - let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); - let weights = - hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) - .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - let kv_mode_str = ctx - .kv_mode_override - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); - let rr = hipfire_runtime::kv_mode::resolve( - &kv_mode_str, - &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, + (weights, Some(transaction)) + } + }; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + // The plain LLaMA path has no independent cap resolver. PR + // #661's physical-cap behavior is owned by the existing + // upstream KV plan. + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; + weights.free_gpu(ctx.gpu); + return Err(with_weight_rollback_error( + format!("llama: ForwardScratch::new_with_max_seq failed: {error:?}"), + rollback, + )); + } + }; + let dims = llama_kv_dims(&config, ctx.max_seq, None); + let kv = match ::from_mode( + hipfire_runtime::kv_mode::resolve( + ctx.kv_mode_override.unwrap_or(""), + &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, config.head_dim, - ); - if let Some(w) = rr.warning { - eprintln!( - " KV cache: {w} (site {})", - hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site - ); + ) + .mode, + KvTarget::Single(ctx.gpu), + &dims, + ) { + Ok(kv) => kv, + Err(error) => { + scratch.free_gpu(ctx.gpu); + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; + weights.free_gpu(ctx.gpu); + return Err(with_weight_rollback_error( + format!("llama: ::from_mode failed: {error}"), + rollback, + )); } - let dims = llama_kv_dims(&config, ctx.max_seq, Some(ctx.max_seq)); - let kv = match ::from_mode( - rr.mode, - KvTarget::Single(ctx.gpu), - &dims, - ) { + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + weight_store, + mesh, + weight_origin, + ) + } + ModelSource::Dir(source) => { + let config = hipfire_runtime::hfq::config_from_safetensors_llama(&source) + .map_err(|e| format!("failed to parse LLaMA/Qwen3 config from config.json: {e}"))?; + let (mesh, manifest_plan) = + plan_single(&config, source.tensor_info("lm_head.weight").is_some())?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let weights = + hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) + .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + let kv_mode_str = ctx + .kv_mode_override + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); + let rr = hipfire_runtime::kv_mode::resolve( + &kv_mode_str, + &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, + config.head_dim, + ); + if let Some(w) = rr.warning { + eprintln!( + " KV cache: {w} (site {})", + hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site + ); + } + let dims = llama_kv_dims(&config, ctx.max_seq, Some(ctx.max_seq)); + let kv = + match ::from_mode(rr.mode, KvTarget::Single(ctx.gpu), &dims) + { Ok(kv) => kv, Err(error) => { weights.free_gpu(ctx.gpu); return Err(format!("KvCache: {error}")); } }; - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - let _ = kv.free_gpu(ctx.gpu); - weights.free_gpu(ctx.gpu); - return Err(format!( - "ForwardScratch::new_with_max_seq: {error:?}" - )); - } - }; - ( - config, - weights, - kv, - scratch, - manifest_plan, - None, - mesh, - weight_origin, - ) - } - }; + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let _ = kv.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!("ForwardScratch::new_with_max_seq: {error:?}")); + } + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + None, + mesh, + weight_origin, + ) + } + }; let mut bundle = LlamaBundle { config, @@ -756,15 +728,10 @@ impl LlamaBundle { transaction: WeightLoadTransaction, ) -> Result<(), (WeightLoadTransaction, String)> { if self.weight_store.is_some() { - return Err(( - transaction, - "llama: weight store already attached".into(), - )); + return Err((transaction, "llama: weight store already attached".into())); } - let attached = match AttachedWeightStore::from_transaction( - transaction, - self.weight_origin, - ) { + let attached = match AttachedWeightStore::from_transaction(transaction, self.weight_origin) + { Ok(attached) => attached, Err((transaction, error)) => { return Err(( @@ -800,20 +767,18 @@ impl LlamaBundle { #[cfg(test)] mod tests { use super::*; - use hipfire_runtime::llama::ModelArch; use hipfire_runtime::arch_model::ArchModel; - use hipfire_runtime::hfq::{ - write_hfqm_package_mem, HfqFile, HfqMemTensor, - }; + use hipfire_runtime::hfq::{write_hfqm_package_mem, HfqFile, HfqMemTensor}; use hipfire_runtime::kv_backend::KvBackend; use hipfire_runtime::kv_mode::KvMode; - use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; + use hipfire_runtime::llama::ModelArch; use hipfire_runtime::llama::{ forward_scratch_compute, forward_scratch_embed, KvCache, KvCacheExt, KvDims, KvLayers, KvTarget, }; - use hipfire_runtime::weight_store::test_support; + use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; use hipfire_runtime::weight_manifest::ShardPolicy; + use hipfire_runtime::weight_store::test_support; use hipfire_runtime::weight_store::{ WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, }; @@ -863,7 +828,6 @@ mod tests { } } - fn fixture_hfq( with_awq_sidecar: bool, with_q_proj_bias: bool, @@ -880,11 +844,7 @@ mod tests { f16_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32]), f16_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32]), f16_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64]), - f32_hfq_tensor( - "model.layers.0.input_layernorm.weight", - &[32], - false, - ), + f32_hfq_tensor("model.layers.0.input_layernorm.weight", &[32], false), f32_hfq_tensor( "model.layers.0.post_attention_layernorm.weight", &[32], @@ -932,10 +892,8 @@ mod tests { .duration_since(UNIX_EPOCH) .expect("system clock before epoch") .as_nanos(); - let path = std::env::temp_dir().join(format!( - "hipfire-g3-{}-{nonce}.hfq", - std::process::id() - )); + let path = + std::env::temp_dir().join(format!("hipfire-g3-{}-{nonce}.hfq", std::process::id())); write_hfqm_package_mem(&path, 0, metadata, &tensors).expect("write HFQ fixture"); let hfq = HfqFile::open(&path).expect("open HFQ fixture"); (path, hfq) @@ -1002,7 +960,10 @@ mod tests { assert_eq!(mesh.n_devices(), 1); assert_eq!(plan.weights.len(), 12); assert_eq!(plan.state.len(), 1); - assert!(plan.collective_schedule.iter().any(|entry| entry.name == "wo")); + assert!(plan + .collective_schedule + .iter() + .any(|entry| entry.name == "wo")); assert!(manifest[0].dtype_constraint.accepts(DType::HFQ4G256)); assert!(manifest[1].dtype_constraint.accepts(DType::MQ4G256)); assert!(manifest[9].dtype_constraint.accepts(DType::F32)); @@ -1036,7 +997,6 @@ mod tests { assert!(transaction.projection("lm_head", None, 0).is_some()); } - #[test] fn hfq_float_widening_matches_legacy_f32_representation() { let f16_one = [0x00, 0x3c, 0x00, 0xc0]; @@ -1192,12 +1152,9 @@ mod tests { let scratch = ForwardScratch::new_with_max_seq(&mut gpu, &config, 8) .expect("allocate legacy forward scratch"); let dims = llama_kv_dims(&config, 8, None); - let mut kv = ::from_mode( - KvMode::Q8, - KvTarget::Single(&mut gpu), - &dims, - ) - .expect("allocate legacy KV cache"); + let mut kv = + ::from_mode(KvMode::Q8, KvTarget::Single(&mut gpu), &dims) + .expect("allocate legacy KV cache"); forward_scratch_embed(&mut gpu, &legacy, &config, 1, 0, &scratch) .expect("legacy embedding forward"); forward_scratch_compute(&mut gpu, &legacy, &config, 0, &mut kv, &scratch) @@ -1210,9 +1167,7 @@ mod tests { legacy.free_gpu(&mut gpu); assert_eq!(manifest_logits.len(), legacy_logits.len()); - for (index, (manifest, legacy)) in - manifest_logits.iter().zip(&legacy_logits).enumerate() - { + for (index, (manifest, legacy)) in manifest_logits.iter().zip(&legacy_logits).enumerate() { assert!( (manifest - legacy).abs() <= 1e-5, "logit mismatch at index {index}: manifest={manifest} legacy={legacy}" @@ -1221,7 +1176,6 @@ mod tests { std::fs::remove_file(path).expect("remove HFQ fixture"); } - #[test] fn physical_cap_is_honored_by_upstream_kv_constructor() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { @@ -1234,12 +1188,9 @@ mod tests { max_seq: 8, physical_cap: Some(4), }; - let cache = ::from_mode( - KvMode::Q8, - KvTarget::Single(&mut gpu), - &dims, - ) - .expect("upstream Q8 constructor"); + let cache = + ::from_mode(KvMode::Q8, KvTarget::Single(&mut gpu), &dims) + .expect("upstream Q8 constructor"); assert_eq!(cache.max_seq, 8); assert_eq!(cache.physical_cap, 4); let _ = cache.free_gpu(&mut gpu); diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index 26178a1229..90b38444eb 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -6,9 +6,8 @@ //! per-tensor dequant), which `WeightSource::read_layer` calls internally. use crate::llama::{EmbeddingFormat, WeightTensor}; -use hipfire_hardware::{DeviceMesh, DimKind}; use hip_bridge::HipResult; -use hipfire_hardware::Gpus; +use hipfire_hardware::{DeviceMesh, DimKind, Gpus}; use rdna_compute::{Gpu, GpuTensor}; /// Where each piece of the model lands across a device slice. `single` = the diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 65a9172f48..2f29231941 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -30,9 +30,7 @@ use std::collections::HashSet; pub fn collective_for_policy(policy: &ShardPolicy) -> Option { match policy { ShardPolicy::RowShard { .. } => Some(CollectiveHint::AllReduce { kind: DimKind::Tp }), - ShardPolicy::ExpertSharded { .. } => { - Some(CollectiveHint::AllReduce { kind: DimKind::Ep }) - } + ShardPolicy::ExpertSharded { .. } => Some(CollectiveHint::AllReduce { kind: DimKind::Ep }), ShardPolicy::ExpertTensorSharded { inner, .. } => collective_for_policy(inner), _ => None, } @@ -422,9 +420,10 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } | ShardPolicy::VocabShard { axis } => { - let dim = entry.logical_shape.get(*axis).ok_or_else(|| { - format!("{context}: shard axis {axis} outside logical shape") - })?; + let dim = entry + .logical_shape + .get(*axis) + .ok_or_else(|| format!("{context}: shard axis {axis} outside logical shape"))?; if tp > 1 && dim % tp != 0 { return Err(format!( "{context}: shard dim {dim} (axis {axis}) not divisible by Tp={tp}" @@ -511,11 +510,14 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< )); } let axis = match inner.as_ref() { - ShardPolicy::ColumnShard { axis: 1 } - | ShardPolicy::RowShard { axis: 2 } => match inner.as_ref() { - ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => *axis, - _ => unreachable!(), - }, + ShardPolicy::ColumnShard { axis: 1 } | ShardPolicy::RowShard { axis: 2 } => { + match inner.as_ref() { + ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => { + *axis + } + _ => unreachable!(), + } + } ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => { return Err(format!( "{context}: ExpertTensorSharded inner axis {axis} is incompatible with [expert, projection, hidden]" @@ -591,7 +593,10 @@ pub fn plan_manifest( }) .collect(); let band_xfers = (0..n_layers) - .filter_map(|layer| mesh.band_xfer_after(layer, n_layers).map(|hint| (layer, hint))) + .filter_map(|layer| { + mesh.band_xfer_after(layer, n_layers) + .map(|hint| (layer, hint)) + }) .collect(); Ok(ManifestPlan { weights: weight_placements, @@ -728,23 +733,14 @@ fn manifest_entry<'a>( .ok_or_else(|| format!("{context}: {label} reference '{name}' not found")) } -fn source_policy_matches( - spec: &ExpertGroupSpec, - label: &str, - policy: &ShardPolicy, -) -> bool { +fn source_policy_matches(spec: &ExpertGroupSpec, label: &str, policy: &ShardPolicy) -> bool { match spec.parallelism { ExpertParallelism::Single => matches!( policy, - ShardPolicy::Replicate - | ShardPolicy::Pin(_) - | ShardPolicy::Tied { .. } + ShardPolicy::Replicate | ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } ), ExpertParallelism::TensorParallel => match (label, policy) { - ( - "gate_up" | "gate" | "up", - ShardPolicy::ExpertTensorSharded { n_experts, inner }, - ) => { + ("gate_up" | "gate" | "up", ShardPolicy::ExpertTensorSharded { n_experts, inner }) => { *n_experts == spec.n_experts && matches!(inner.as_ref(), ShardPolicy::ColumnShard { axis: 1 }) } @@ -794,10 +790,7 @@ fn source_shape_matches( Ok(()) } -fn validate_expert_sources( - spec: &ExpertGroupSpec, - manifest: &[WeightEntry], -) -> Result<(), String> { +fn validate_expert_sources(spec: &ExpertGroupSpec, manifest: &[WeightEntry]) -> Result<(), String> { let context = expert_context(spec); let router = manifest_entry(spec, manifest, "router", &spec.router)?; if !matches!(router.logical_shape.len(), 1 | 2) @@ -884,13 +877,19 @@ pub fn validate_expert_group_specs( for spec in specs { let context = expert_context(spec); if spec.group.is_empty() || spec.router.is_empty() || spec.execution.is_empty() { - return Err(format!("{context}: group/router/execution identities must be non-empty")); + return Err(format!( + "{context}: group/router/execution identities must be non-empty" + )); } if spec.n_experts == 0 || spec.resources.bytes_per_expert == 0 { - return Err(format!("{context}: n_experts and bytes_per_expert must be non-zero")); + return Err(format!( + "{context}: n_experts and bytes_per_expert must be non-zero" + )); } if spec.resources.alignment == 0 || !spec.resources.alignment.is_power_of_two() { - return Err(format!("{context}: alignment must be a non-zero power of two")); + return Err(format!( + "{context}: alignment must be a non-zero power of two" + )); } if !groups.insert((&spec.group, spec.layer)) { return Err(format!("{context}: duplicate group/layer identity")); @@ -1002,13 +1001,7 @@ mod tests { #[test] fn expert_source_identity_and_shape_are_checked() { let manifest = vec![ - WeightEntry::layer( - "router", - 0, - vec![8, 4], - DType::F16, - ShardPolicy::Replicate, - ), + WeightEntry::layer("router", 0, vec![8, 4], DType::F16, ShardPolicy::Replicate), WeightEntry::layer( "gate_up", 0, @@ -1089,12 +1082,7 @@ mod tests { #[test] fn tied_entries_require_matching_representation_and_no_tied_chain() { - let source = WeightEntry::model( - "source", - vec![8, 8], - DType::F16, - ShardPolicy::Replicate, - ); + let source = WeightEntry::model("source", vec![8, 8], DType::F16, ShardPolicy::Replicate); let shape_mismatch = WeightEntry::model( "shape_mismatch", vec![8, 4], @@ -1161,16 +1149,15 @@ mod tests { source: "cycle_a".into(), }, ); - assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single().expect("single-device mesh construction cannot overflow")).is_err()); + assert!(validate_manifest( + &[cycle_a, cycle_b], + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .is_err()); } #[test] fn tied_entries_reject_different_source_sets_with_equal_logical_dtype() { - let source = WeightEntry::model( - "source", - vec![8, 8], - DType::F16, - ShardPolicy::Replicate, - ); + let source = WeightEntry::model("source", vec![8, 8], DType::F16, ShardPolicy::Replicate); let tied = WeightEntry::model_with_dtype_constraint( "tied", vec![8, 8], @@ -1182,7 +1169,7 @@ mod tests { ); let error = validate_manifest( &[source, tied], - &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + &DeviceMesh::single().expect("single-device mesh construction cannot overflow"), ) .unwrap_err(); assert!(error.contains("source dtype contract")); diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index ce0019441c..d67bf06c03 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -36,9 +36,7 @@ thread_local! { /// [`clear_faults`] after it so a failed test cannot poison a later one. #[doc(hidden)] pub mod test_support { - use super::{ - FAIL_AFTER_UPLOAD, RESIDENT_ALLOCATIONS, RESIDENT_RELEASES, - }; + use super::{FAIL_AFTER_UPLOAD, RESIDENT_ALLOCATIONS, RESIDENT_RELEASES}; pub fn reset() { RESIDENT_ALLOCATIONS.with(|count| count.set(0)); @@ -125,7 +123,12 @@ pub struct WeightProjection { pub dtype: DType, } -fn projection_for(entry: &WeightEntry, rank: usize, world_size: usize, dtype: DType) -> WeightProjection { +fn projection_for( + entry: &WeightEntry, + rank: usize, + world_size: usize, + dtype: DType, +) -> WeightProjection { let (kind, axis) = match &entry.policy { ShardPolicy::ColumnShard { axis } => (WeightProjectionKind::ColumnShard, Some(*axis)), ShardPolicy::RowShard { axis } => (WeightProjectionKind::RowShard, Some(*axis)), @@ -307,12 +310,7 @@ impl WeightLoadTransaction { .is_some_and(|store| store.contains(name, layer, device)) } - pub fn get( - &self, - name: &str, - layer: Option, - device: usize, - ) -> Option<&WeightHandle> { + pub fn get(&self, name: &str, layer: Option, device: usize) -> Option<&WeightHandle> { self.store .as_ref() .and_then(|store| store.get(name, layer, device)) @@ -338,14 +336,12 @@ impl WeightLoadTransaction { /// Compare the unpublished transaction's captured target with an admitted /// owner identity. This read-only check is used before the carrier wraps /// the transaction in its private attached owner. - pub fn validate_origin_value( - &self, - expected: WeightOrigin, - ) -> Result<(), WeightStoreError> { - self.store.as_ref().map_or( - Err(WeightStoreError::UnboundOrigin), - |store| store.validate_origin_value(expected), - ) + pub fn validate_origin_value(&self, expected: WeightOrigin) -> Result<(), WeightStoreError> { + self.store + .as_ref() + .map_or(Err(WeightStoreError::UnboundOrigin), |store| { + store.validate_origin_value(expected) + }) } /// Start typed assembly while this load is still unpublished. @@ -459,12 +455,7 @@ impl WeightStore { /// Move a handle out of the store. This is private to the assembly /// capability so arbitrary store holders cannot independently tear down a /// resident allocation. - fn take( - &mut self, - name: &str, - layer: Option, - device: usize, - ) -> Option { + fn take(&mut self, name: &str, layer: Option, device: usize) -> Option { let key = WeightPlacementKey::new(name, layer, device); self.projections.remove(&key); self.placements.remove(&key) @@ -492,10 +483,7 @@ impl WeightStore { /// Compare a store's captured origin with an already-resolved target /// identity. This read-only seam cannot release or extract any handle. - pub fn validate_origin_value( - &self, - expected: WeightOrigin, - ) -> Result<(), WeightStoreError> { + pub fn validate_origin_value(&self, expected: WeightOrigin) -> Result<(), WeightStoreError> { let actual = self.origin.ok_or(WeightStoreError::UnboundOrigin)?; if actual != expected { return Err(WeightStoreError::OriginMismatch { expected, actual }); @@ -505,11 +493,7 @@ impl WeightStore { /// Verify that this store is still being handled by the same mesh/device /// target. No GPU calls occur on mismatch. - pub fn validate_origin( - &self, - mesh: &DeviceMesh, - gpu: &Gpu, - ) -> Result<(), WeightStoreError> { + pub fn validate_origin(&self, mesh: &DeviceMesh, gpu: &Gpu) -> Result<(), WeightStoreError> { self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) } @@ -558,12 +542,7 @@ pub struct WeightStoreAssembly<'a> { } impl<'a> WeightStoreAssembly<'a> { - pub fn take( - &mut self, - name: &str, - layer: Option, - device: usize, - ) -> Option { + pub fn take(&mut self, name: &str, layer: Option, device: usize) -> Option { let key = WeightPlacementKey::new(name, layer, device); let (handle, projection) = self.store.take_with_projection(name, layer, device)?; let slot = self.taken.len(); @@ -586,9 +565,7 @@ impl Drop for WeightStoreAssembly<'_> { return; } for taken in self.taken.drain(..) { - let _ = self - .store - .insert(taken.key, taken.handle, taken.projection); + let _ = self.store.insert(taken.key, taken.handle, taken.projection); } } } @@ -628,11 +605,7 @@ fn target_error(mesh: &DeviceMesh) -> Option { ), }) } -fn rollback_fulfill_error( - store: WeightStore, - gpu: &Gpu, - mut error: FulfillError, -) -> FulfillError { +fn rollback_fulfill_error(store: WeightStore, gpu: &Gpu, mut error: FulfillError) -> FulfillError { if let Err(release_error) = store.rollback(gpu) { error .reason @@ -680,15 +653,15 @@ where name: entry.name.clone(), layer: entry.layer, device: devices.first().copied().unwrap_or(0), - reason: format!( - "Single placement resolved to {:?}, expected [0]", - devices - ), + reason: format!("Single placement resolved to {:?}, expected [0]", devices), }; return Err(rollback_fulfill_error(store, gpu, error)); } let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); - if let ShardPolicy::Tied { source: source_name } = &entry.policy { + if let ShardPolicy::Tied { + source: source_name, + } = &entry.policy + { let source_dtype = match store.get(source_name, entry.layer, 0) { Some(WeightHandle::Resident(tensor)) => Some(tensor.dtype), Some(WeightHandle::Alias(_)) | None => None, @@ -717,11 +690,9 @@ where return Err(rollback_fulfill_error(store, gpu, error)); } let projection = projection_for(entry, 0, 1, actual_dtype); - if let Err(reason) = store.insert( - key, - WeightHandle::Alias(source_name.clone()), - projection, - ) { + if let Err(reason) = + store.insert(key, WeightHandle::Alias(source_name.clone()), projection) + { let error = FulfillError { name: entry.name.clone(), layer: entry.layer, @@ -926,7 +897,9 @@ mod tests { .stage_alias("x", None, 0, "source-b", projection(DType::F32)) .unwrap_err(); assert!(matches!(error, WeightStoreError::DuplicatePlacement(_))); - assert!(matches!(store.get("x", None, 0), Some(WeightHandle::Alias(source)) if source == "source-a")); + assert!( + matches!(store.get("x", None, 0), Some(WeightHandle::Alias(source)) if source == "source-a") + ); assert_eq!(store.projection("x", None, 0).unwrap().dtype, DType::F16); } @@ -969,13 +942,9 @@ mod tests { source: "source".into(), }, ); - let transaction = fulfill_manifest_single( - &[source, alias], - &mesh, - 1, - &gpu, - |_| Ok((vec![0; 4], DType::F32)), - ) + let transaction = fulfill_manifest_single(&[source, alias], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], DType::F32)) + }) .unwrap(); assert_eq!( transaction.projection("alias", None, 0).unwrap().dtype, @@ -997,10 +966,9 @@ mod tests { }; let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { - Ok((vec![0; 4], DType::F32)) - }) - .unwrap(); + let transaction = + fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| Ok((vec![0; 4], DType::F32))) + .unwrap(); assert_eq!(transaction.len(), 1); assert!(matches!( transaction.get("resident", None, 0), @@ -1041,10 +1009,9 @@ mod tests { RESIDENT_RELEASES.with(|count| count.set(0)); let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { - Ok((vec![0; 4], DType::F32)) - }) - .unwrap(); + let transaction = + fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| Ok((vec![0; 4], DType::F32))) + .unwrap(); let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); let error = transaction.validate_origin_value(expected).unwrap_err(); assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); @@ -1060,7 +1027,6 @@ mod tests { assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } - #[test] fn rollback_reports_free_failure_without_counting_release() { let Ok(gpu) = Gpu::init() else { @@ -1073,10 +1039,7 @@ mod tests { let mut store = WeightStore::with_origin(origin); let borrowed = GpuTensor { buf: unsafe { - hip_bridge::DeviceBuffer::from_raw( - std::ptr::null_mut::(), - 0, - ) + hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut::(), 0) }, shape: vec![0], dtype: DType::F32, From 04de36cf8249797cb34d2aeb2b303e669342fd8c Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 08:56:53 +0200 Subject: [PATCH 15/25] fix(device-mesh): recognize alternate lm head names --- crates/hipfire-arch-llama/src/carrier.rs | 59 ++++++++++++++++++++---- 1 file changed, 51 insertions(+), 8 deletions(-) diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index a6bd292c02..c40e55ad59 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -117,16 +117,26 @@ fn hfq_layer_names(layer: usize, relative: &str) -> Vec { format!("layers.{layer}.{relative}.weight"), ] } +const HFQ_LM_HEAD_NAMES: &[&str] = &[ + "lm_head.weight", + "model.lm_head.weight", + "model.language_model.lm_head.weight", +]; + +fn hfq_has_separate_lm_head(hfq: &HfqFile) -> bool { + HFQ_LM_HEAD_NAMES + .iter() + .any(|name| hfq.find_tensor_info(name).is_some()) +} fn hfq_entry_names(entry: &WeightEntry) -> Result, String> { let names = match (entry.name.as_str(), entry.layer) { ("token_embd", None) => vec!["model.embed_tokens.weight".to_string()], ("output_norm", None) => vec!["model.norm.weight".to_string()], - ("lm_head", None) => vec![ - "lm_head.weight".to_string(), - "model.lm_head.weight".to_string(), - "model.language_model.lm_head.weight".to_string(), - ], + ("lm_head", None) => HFQ_LM_HEAD_NAMES + .iter() + .map(|name| (*name).to_string()) + .collect(), ("wq", Some(layer)) => hfq_layer_names(layer, "self_attn.q_proj"), ("wk", Some(layer)) => hfq_layer_names(layer, "self_attn.k_proj"), ("wv", Some(layer)) => hfq_layer_names(layer, "self_attn.v_proj"), @@ -539,7 +549,7 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result (PathBuf, HfqFile) { + fixture_hfq_with_lm_head( + with_awq_sidecar, + with_q_proj_bias, + malformed_output_norm, + separate_lm_head.then_some("lm_head.weight"), + ) + } + + fn fixture_hfq_with_lm_head( + with_awq_sidecar: bool, + with_q_proj_bias: bool, + malformed_output_norm: bool, + lm_head_name: Option<&str>, ) -> (PathBuf, HfqFile) { let mut tensors = vec![ f32_hfq_tensor("model.embed_tokens.weight", &[2, 32], false), @@ -870,8 +894,8 @@ mod tests { 32 * 2, )); } - if separate_lm_head { - tensors.push(f32_hfq_tensor("lm_head.weight", &[2, 32], false)); + if let Some(lm_head_name) = lm_head_name { + tensors.push(f32_hfq_tensor(lm_head_name, &[2, 32], false)); } let metadata = r#"{ "config": { @@ -1064,6 +1088,25 @@ mod tests { std::fs::remove_file(path).expect("remove HFQ fixture"); } + #[test] + fn alternate_explicit_lm_head_names_are_not_tied() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + for name in &HFQ_LM_HEAD_NAMES[1..] { + let (path, hfq) = fixture_hfq_with_lm_head(false, false, false, Some(name)); + assert!(hfq_has_separate_lm_head(&hfq)); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = + load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load explicit lm_head"); + drop(ctx); + assert!(!bundle.weights.lm_head_aliases_embd); + Box::new(bundle).free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + } + #[test] fn production_biased_hfq_is_rejected_before_manifest_upload() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { From c241a5dcdd3b4a580081e2af23dd8e0f386ffee3 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 09:33:07 +0200 Subject: [PATCH 16/25] feat: seal generic MoE step ownership --- crates/hipfire-dispatch/Cargo.toml | 1 + crates/hipfire-dispatch/src/families/moe.rs | 719 +++++++++- crates/hipfire-dispatch/src/pipeline/mod.rs | 46 +- crates/hipfire-dispatch/src/pipeline/steps.rs | 1166 ++++++++++++++++- crates/hipfire-dispatch/src/types.rs | 12 + crates/hipfire-runtime/src/lib.rs | 1 + crates/hipfire-runtime/src/moe_plan.rs | 1076 +++++++++++++++ 7 files changed, 2987 insertions(+), 34 deletions(-) create mode 100644 crates/hipfire-runtime/src/moe_plan.rs diff --git a/crates/hipfire-dispatch/Cargo.toml b/crates/hipfire-dispatch/Cargo.toml index cc27bdc504..017f3c4148 100644 --- a/crates/hipfire-dispatch/Cargo.toml +++ b/crates/hipfire-dispatch/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true hip-bridge = { path = "../hip-bridge" } hipfire-config = { path = "../hipfire-config" } rdna-compute = { path = "../rdna-compute" } +hipfire-hardware = { path = "../hipfire-hardware" } [features] default = [] diff --git a/crates/hipfire-dispatch/src/families/moe.rs b/crates/hipfire-dispatch/src/families/moe.rs index b2882225e9..6e2462aff3 100644 --- a/crates/hipfire-dispatch/src/families/moe.rs +++ b/crates/hipfire-dispatch/src/families/moe.rs @@ -26,7 +26,331 @@ use crate::tables::moe_table; use crate::tables::KernelRegistry; use crate::traits::KernelFamily; use crate::types::*; +/// The routing operation selected for one admitted MoE group. +/// +/// The generic executor currently has two concrete forms: a normal softmax +/// top-k launch and a precomputed route supplied by an upstream owner. Family +/// policy (bias, hash, or sigmoid variants) stays outside this substrate. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RouterSelection { + SoftmaxTopK, + Precomputed, +} + +/// Typed routing operands. Each variant carries only the operands required by +/// its routing semantics, so a caller cannot accidentally drop normalization +/// while lowering. +pub enum RouterPlan<'a> { + SoftmaxTopK { + scores: &'a GpuTensor, + topk_indices: &'a GpuTensor, + topk_weights: &'a GpuTensor, + k_top: usize, + normalize: bool, + }, + /// The route was selected by an owner outside this executor. This is a + /// real operation boundary, not a second route implementation. + Precomputed { + topk_indices: &'a GpuTensor, + topk_weights: &'a GpuTensor, + k_top: usize, + }, +} + +impl<'a> RouterPlan<'a> { + pub fn selection(&self) -> RouterSelection { + match self { + Self::SoftmaxTopK { .. } => RouterSelection::SoftmaxTopK, + Self::Precomputed { .. } => RouterSelection::Precomputed, + } + } + + pub fn k_top(&self) -> usize { + match self { + Self::SoftmaxTopK { k_top, .. } | Self::Precomputed { k_top, .. } => *k_top, + } + } + + pub fn normalizes(&self) -> bool { + match self { + Self::SoftmaxTopK { normalize, .. } => *normalize, + Self::Precomputed { .. } => true, + } + } + + + pub fn route_buffers(&self) -> (&'a GpuTensor, &'a GpuTensor) { + match self { + Self::SoftmaxTopK { + topk_indices, + topk_weights, + .. + } + | Self::Precomputed { + topk_indices, + topk_weights, + .. + } => (topk_indices, topk_weights), + } + } + + pub fn batch_size(&self) -> usize { + let indices = self.route_buffers().0; + match indices.shape.as_slice() { + [_, k] if *k == self.k_top() => indices.shape[0], + [_] => 1, + _ => 0, + } + } + + /// Validate route metadata before any expert kernel can launch. + pub fn validate_against( + &self, + n_experts: usize, + batch_size: usize, + ) -> Result<(), DispatchError> { + if n_experts == 0 || batch_size == 0 || self.k_top() == 0 || self.k_top() > n_experts { + return Err(DispatchError::Hip(format!( + "MoE route has invalid n_experts={n_experts}, batch_size={batch_size}, k_top={}", + self.k_top() + ))); + } + let (indices, weights) = self.route_buffers(); + if indices.dtype != DType::F32 || weights.dtype != DType::F32 { + return Err(DispatchError::Hip( + "MoE route indices and weights must use F32 storage".into(), + )); + } + let expected_slots = batch_size + .checked_mul(self.k_top()) + .ok_or_else(|| DispatchError::Hip("MoE route slot count overflow".into()))?; + if indices.numel() < expected_slots || weights.numel() < expected_slots { + return Err(DispatchError::Hip(format!( + "MoE route buffers have insufficient capacity for {expected_slots} slots" + ))); + } + if !matches!(indices.shape.as_slice(), [_] | [_, _]) { + return Err(DispatchError::Hip( + "MoE route index shape must be [K] or [B,K]".into(), + )); + } + if indices.shape.last() != Some(&self.k_top()) + || weights.shape.last() != Some(&self.k_top()) + { + return Err(DispatchError::Hip( + "MoE route index/weight shape must end in k_top".into(), + )); + } + if let Self::SoftmaxTopK { scores, .. } = self { + if scores.dtype != DType::F32 + || scores.shape.last() != Some(&n_experts) + || !matches!(scores.shape.len(), 1 | 2) + || scores.numel() < batch_size.saturating_mul(n_experts) + { + return Err(DispatchError::Hip(format!( + "MoE router score shape/dtype does not match batch={batch_size}, experts={n_experts}" + ))); + } + } + Ok(()) + } +} + +/// Executor shape choice for one admitted routed expert group. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExpertExecutionPlan { + IndexedQuantized, + GroupedQuantized, + PerExpertFallback, +} +/// Executable grammar selected by the sealed plan. There is deliberately no +/// generic fallback grammar: a fallback would bypass the owner and collective +/// checks that make a Step program safe. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MoeProtocolKind { + Indexed, + Grouped, +} + +impl ExpertExecutionPlan { + pub fn protocol(self) -> Result { + match self { + Self::IndexedQuantized => Ok(MoeProtocolKind::Indexed), + Self::GroupedQuantized => Ok(MoeProtocolKind::Grouped), + Self::PerExpertFallback => Err(DispatchError::Hip( + "PerExpertFallback is not an executable MoE Step protocol".into(), + )), + } + } +} + +/// Borrowed view over one resolver-owned rank-local expert table. +/// +/// The view contains no allocation handle, source path, or storage owner. It +/// is created by the manifest/runtime owner and borrowed by a Step until the +/// owner is dropped. Keeping the fields private prevents a family from +/// replacing the canonical placement metadata after binding. +pub struct MoeExpertRef<'a> { + gate_up_ptrs: &'a GpuTensor, + down_ptrs: &'a GpuTensor, + dummy_gate_up: Option<&'a GpuTensor>, + dtype: DType, + n_experts: usize, + expert_m: usize, + expert_k: usize, + owned: &'a [usize], + collective_kind: Option, +} +impl<'a> MoeExpertRef<'a> { + /// Bind a pointer-table view to metadata already resolved by an owner. + /// + /// This constructor deliberately accepts borrowed tables only. Runtime + /// owners should expose a narrower `bind_*` method that supplies the + /// canonical values from its sealed plan; no family receives an allocator + /// or a `WeightStore` representation. + #[doc(hidden)] + pub fn from_resolved( + gate_up_ptrs: &'a GpuTensor, + down_ptrs: &'a GpuTensor, + dummy_gate_up: Option<&'a GpuTensor>, + dtype: DType, + n_experts: usize, + expert_m: usize, + expert_k: usize, + owned: &'a [usize], + collective_kind: Option, + ) -> Self { + Self { + gate_up_ptrs, + down_ptrs, + dummy_gate_up, + dtype, + n_experts, + expert_m, + expert_k, + owned, + collective_kind, + } + } + + pub fn gate_up_ptrs(&self) -> &'a GpuTensor { + self.gate_up_ptrs + } + + pub fn down_ptrs(&self) -> &'a GpuTensor { + self.down_ptrs + } + + pub fn dummy_gate_up(&self) -> Option<&'a GpuTensor> { + self.dummy_gate_up + } + + pub fn dtype(&self) -> DType { + self.dtype + } + + pub fn n_experts(&self) -> usize { + self.n_experts + } + + pub fn expert_m(&self) -> usize { + self.expert_m + } + + pub fn expert_k(&self) -> usize { + self.expert_k + } + + pub fn owned(&self) -> &'a [usize] { + self.owned + } + + pub fn collective_kind(&self) -> Option { + self.collective_kind + } + + /// Validate the dimensions shared by the fused gate/up and down kernels. + /// Gate/up is `[2*expert_m, expert_k]`; down is `[expert_k, expert_m]`. + pub fn validate(&self) -> Result<(), DispatchError> { + if self.n_experts == 0 { + return Err(DispatchError::Hip( + "MoeExpertRef: n_experts must be nonzero".into(), + )); + } + if self.expert_m == 0 || self.expert_k == 0 { + return Err(DispatchError::Hip( + "MoeExpertRef: expert dimensions must be nonzero".into(), + )); + } + if self.gate_up_ptrs.numel() < self.n_experts + || self.down_ptrs.numel() < self.n_experts + { + return Err(DispatchError::Hip(format!( + "MoeExpertRef: pointer-table capacity is too small for {} experts", + self.n_experts + ))); + } + if let Some(dummy) = self.dummy_gate_up { + if dummy.numel() == 0 { + return Err(DispatchError::Hip( + "MoeExpertRef: dummy gate/up table is empty".into(), + )); + } + } + let mut previous = None; + for &expert in self.owned { + if expert >= self.n_experts { + return Err(DispatchError::Hip(format!( + "MoeExpertRef: owned expert {expert} >= n_experts {}", + self.n_experts + ))); + } + if previous == Some(expert) { + return Err(DispatchError::Hip(format!( + "MoeExpertRef: duplicate owned expert {expert}" + ))); + } + previous = Some(expert); + } + Ok(()) + } + + /// Refuse a projection pair whose logical shapes cannot share this + /// executor view. This check is pure and runs before a kernel launch. + pub fn validate_projection_shapes( + &self, + gate_up_shape: &[usize], + down_shape: &[usize], + ) -> Result<(), DispatchError> { + self.validate()?; + let expected_gate_up = [2 * self.expert_m, self.expert_k]; + let expected_down = [self.expert_k, self.expert_m]; + if gate_up_shape != expected_gate_up || down_shape != expected_down { + return Err(DispatchError::Hip(format!( + "MoeExpertRef: projection shape mismatch: gate_up={gate_up_shape:?} \ + expected={expected_gate_up:?}, down={down_shape:?} expected={expected_down:?}" + ))); + } + Ok(()) + } +} + +/// Launch-time activation form for routed experts. The concrete kernel family +/// remains an architecture concern; these are the only semantic forms the +/// generic Step substrate needs to name. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MoeActivationVariant { + SiluMul, + SiluMulRotate, +} +/// Typed routed projection shape. Every down projection is expanded and is +/// followed by the executor-owned `MoeCombine` step. Family kernels therefore +/// cannot hide a second reduction inside this vocabulary. +pub enum MoeProj<'a> { + GateUp { up_out: &'a GpuTensor }, + DownExpanded, +} // ── MoE eligibility lattice ──────────────────────────── /// Routed-expert tiers the mixed-tier graded decode path can execute: the @@ -66,27 +390,13 @@ pub struct MoeDtypes { pub routed_gate_up: DType, // ffn.experts[0].gate_up pub routed_down: DType, // ffn.experts[0].down /// Per-expert mixed routed dtype: experts in one layer carry DIFFERENT - /// gate_up and/or down dtypes (N-tier graded: MQ6 hot / MQ4 mid / MQ2L - /// or MQ3L or E8-family cold), so `routed_gate_up` / `routed_down` - /// (= experts[0]) are NOT representative. Built by the model as - /// `ffn.expert_dtype_tags.is_some()` — the tag table is built iff any - /// expert's gate_up or down dtype differs from experts[0]. Tags: - /// 0 = MQ6G256 (200 B/grp affine) - /// 1 = MQ2G256Lloyd ( 72 B/grp codebook) - /// 2 = MQ4G256 (136 B/grp affine) - /// 3 = MQ3G256Lloyd (112 B/grp codebook) - /// 4 = MFP4G32E8 (16 B hdr + (K/32)*17 B; 4-bit E8 lattice, 4.25 bpw) - /// 5 = MFP3G32E8 (16 B hdr + (K/32)*13 B; 3-bit E8 lattice, 3.25 bpw) - /// 6 = MFP2G32E8 (16 B hdr + (K/32)*9 B; 2-bit E8 lattice, 2.25 bpw) - /// Drives the merged dtype-tag-branched gate_up AND down decode kernels. + /// gate_up and/or down dtypes; the tag table is built iff the layer is + /// heterogeneous and drives the merged decode kernels. pub routed_has_mixed_experts: bool, pub has_paro_shared: bool, // ffn.paro_shared.is_some() - /// Per-expert gate_up tiers for intra-layer mixed-tier dispatch. `None` - /// (default) ⇒ today's uniform path (representative `routed_gate_up` drives - /// resolution). `Some(table)` with >1 distinct DType marks the layer - /// `mixed`; a `Some` table that is all-equal collapses to the uniform path. + /// Per-expert gate_up tiers for intra-layer mixed-tier dispatch. pub per_expert_gate_up: Option>, - /// Per-expert down tiers (parallel to `per_expert_gate_up`). Same semantics. + /// Per-expert down tiers (parallel to `per_expert_gate_up`). pub per_expert_down: Option>, } @@ -100,12 +410,11 @@ impl MoeDtypes { self.routed_down, ] .iter() - // V1 (qt14) and V2 (qt47) are both 6-bit FWHT projections that trip - // the gfx1151 MQ4-i8 grouped fence via `force_mq4_grouped_fp16`. .any(|dt| matches!(*dt, DType::MQ6G256 | DType::MQ6G256V2)) } } + /// Resolved fused-vs-fallback eligibility for one MoE decode layer. This IS the /// routing-config logic, relocated from `moe_ffn_decode_impl` into one typed, /// testable place (review finding #1). Pure function of `MoeDtypes` + k. @@ -929,6 +1238,313 @@ impl MoeFamily { pub fn registry(&self) -> &KernelRegistry { &self.registry } + /// Execute the owner-bound route operation. Precomputed routes are an + /// explicit identity boundary; computed routes use the existing batched + /// k=8 router kernel and never infer a different policy. + pub(crate) fn run_route( + &self, + gpu: &mut Gpu, + plan: &RouterPlan<'_>, + ) -> Result<(), DispatchError> { + let batch_size = plan.batch_size(); + match plan { + RouterPlan::SoftmaxTopK { + scores, + topk_indices, + topk_weights, + k_top, + normalize, + } => { + if *k_top != 8 { + return Err(DispatchError::Hip(format!( + "generic MoE softmax route requires k_top=8, got {k_top}" + ))); + } + let n_experts = *scores.shape.last().ok_or_else(|| { + DispatchError::Hip("generic MoE route scores have no expert axis".into()) + })?; + plan.validate_against(n_experts, batch_size)?; + gpu.moe_softmax_topk_renorm_k8_batched( + scores, + topk_indices, + topk_weights, + n_experts, + *normalize, + batch_size, + ) + .map_err(|error| DispatchError::Hip(error.to_string())) + } + RouterPlan::Precomputed { + topk_indices, + topk_weights, + k_top, + } => { + if *k_top == 0 + || topk_indices.dtype != DType::F32 + || topk_weights.dtype != DType::F32 + { + return Err(DispatchError::Hip( + "generic MoE precomputed route metadata is invalid".into(), + )); + } + Ok(()) + } + } + } + + pub(crate) fn run_indexed( + &self, + gpu: &mut Gpu, + experts: &MoeExpertRef<'_>, + which: &MoeProj<'_>, + topk_indices: &GpuTensor, + input: &crate::pipeline::steps::GemvInput<'_>, + out: &GpuTensor, + k_top: usize, + batch_size: usize, + ) -> Result<(), DispatchError> { + experts.validate()?; + if batch_size != 1 { + return Err(DispatchError::Hip( + "indexed MoE Steps are decode-only; grouped Steps serve batches".into(), + )); + } + let x = match input { + crate::pipeline::steps::GemvInput::Prerotated(x) => *x, + crate::pipeline::steps::GemvInput::Raw(_) => { + return Err(DispatchError::Hip( + "indexed MoE gate/down kernels require a pre-rotated activation".into(), + )) + } + }; + match which { + MoeProj::GateUp { up_out } => crate::pipeline::run_uniform_moe_gate_up( + gpu, + experts.dtype, + experts.gate_up_ptrs, + topk_indices, + x, + out, + up_out, + experts.expert_m, + experts.expert_k, + k_top, + ), + MoeProj::DownExpanded => crate::pipeline::run_uniform_moe_down_expanded( + gpu, + experts.dtype, + experts.down_ptrs, + topk_indices, + x, + out, + experts.expert_k, + experts.expert_m, + k_top, + batch_size, + ), + } + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn run_scatter( + &self, + gpu: &mut Gpu, + topk_indices: &GpuTensor, + expert_token_counts: &GpuTensor, + expert_offsets: &GpuTensor, + sorted_slot_index: &GpuTensor, + expert_tile_ids: &GpuTensor, + inverse_perm: &GpuTensor, + total_slots: usize, + n_experts: usize, + m_total_max: usize, + block_m: usize, + ) -> Result<(), DispatchError> { + gpu.moe_scatter_fused_k8( + topk_indices, + expert_token_counts, + expert_offsets, + sorted_slot_index, + expert_tile_ids, + inverse_perm, + total_slots, + n_experts, + m_total_max, + block_m, + ) + .map_err(|error| DispatchError::Hip(error.to_string())) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn run_grouped( + &self, + gpu: &mut Gpu, + experts: &MoeExpertRef<'_>, + which: &MoeProj<'_>, + sorted_slot_index: &GpuTensor, + expert_tile_ids: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m_total: usize, + batch_size: usize, + k_top: usize, + ) -> Result<(), DispatchError> { + experts.validate()?; + if batch_size == 0 || k_top == 0 || m_total == 0 { + return Err(DispatchError::Hip( + "grouped MoE GEMM dimensions must be nonzero".into(), + )); + } + let (ptrs, m, k, x_row_div, rows) = match which { + MoeProj::GateUp { .. } => ( + experts.gate_up_ptrs, + 2 * experts.expert_m, + experts.expert_k, + k_top, + batch_size, + ), + MoeProj::DownExpanded => ( + experts.down_ptrs, + experts.expert_k, + experts.expert_m, + 1, + batch_size + .checked_mul(k_top) + .ok_or_else(|| DispatchError::Hip("grouped MoE row count overflow".into()))?, + ), + }; + crate::pipeline::run_grouped_moe_gemm( + gpu, + experts.dtype, + ptrs, + expert_tile_ids, + sorted_slot_index, + x, + y, + m, + k, + x_row_div, + m_total, + rows, + ) + } + + pub(crate) fn run_unscatter( + &self, + gpu: &mut Gpu, + y_grouped: &GpuTensor, + sorted_slot_index: &GpuTensor, + gate_batch: &GpuTensor, + up_batch: &GpuTensor, + inter: usize, + k_top: usize, + m_total: usize, + ) -> Result<(), DispatchError> { + gpu.moe_gate_up_unscatter_k8( + y_grouped, + sorted_slot_index, + gate_batch, + up_batch, + inter, + k_top, + m_total, + ) + .map_err(|error| DispatchError::Hip(error.to_string())) + } + + pub(crate) fn run_activation( + &self, + gpu: &mut Gpu, + variant: MoeActivationVariant, + gate: &GpuTensor, + up: &GpuTensor, + rot_out: &GpuTensor, + inter: usize, + rows: usize, + ) -> Result<(), DispatchError> { + if inter == 0 || rows == 0 { + return Err(DispatchError::Hip( + "MoE activation dimensions must be nonzero".into(), + )); + } + match variant { + MoeActivationVariant::SiluMul => gpu + .silu_mul_f32(gate, up, rot_out) + .map_err(|error| DispatchError::Hip(error.to_string())), + MoeActivationVariant::SiluMulRotate => gpu + .fused_silu_mul_rotate_mq_batched(gate, up, rot_out, inter, rows) + .map_err(|error| DispatchError::Hip(error.to_string())), + } + } + + pub(crate) fn run_combine( + &self, + gpu: &mut Gpu, + down_out: &GpuTensor, + topk_weights: &GpuTensor, + out: &GpuTensor, + hidden: usize, + k_top: usize, + batch_size: usize, + inverse_perm: Option<&GpuTensor>, + ) -> Result<(), DispatchError> { + if hidden == 0 || k_top == 0 || batch_size == 0 { + return Err(DispatchError::Hip( + "MoE combine dimensions must be nonzero".into(), + )); + } + let result = if let Some(inverse_perm) = inverse_perm { + gpu.moe_down_combine_grouped_k8( + down_out, + inverse_perm, + topk_weights, + out, + hidden, + k_top, + batch_size, + ) + } else { + gpu.moe_down_combine_k8_batched( + down_out, + topk_weights, + out, + hidden, + k_top, + batch_size, + ) + }; + result.map_err(|error| DispatchError::Hip(error.to_string())) + } + + /// Seal a generic typed program before any launch. + pub fn seal_steps<'a>( + &self, + execution: ExpertExecutionPlan, + steps: Vec>, + collectives: Vec, + ) -> Result, DispatchError> { + crate::pipeline::steps::SealedMoeSchedule::new(execution, steps, collectives) + } + + pub fn execute_sealed<'a>( + &self, + gpu: &mut Gpu, + ctx: &DispatchCtx, + schedule: &crate::pipeline::steps::SealedMoeSchedule<'a>, + ) -> Result<(), DispatchError> { + crate::pipeline::steps::execute_sealed_steps(gpu, ctx, schedule) + } + + pub fn execute_sealed_mesh<'a>( + &self, + gpus: &mut hipfire_hardware::Gpus, + mesh: &hipfire_hardware::DeviceMesh, + ctx: &DispatchCtx, + rank_steps: &[&[crate::pipeline::steps::Step<'a>]], + collectives: &[crate::pipeline::steps::StepCollective], + ) -> Result<(), DispatchError> { + crate::pipeline::steps::execute_steps_mesh(gpus, mesh, ctx, rank_steps, collectives) + } /// Resolve the best kernel key for the given MoE variant. /// @@ -1207,4 +1823,67 @@ mod tests { assert!(r.routed_indexable_mq6v2); assert!(!r.use_gpu_topk); } + + #[test] + fn typed_router_preserves_selection_and_normalization_contract() { + let mut scores = GpuTensor::null_for_test(); + scores.shape = vec![8]; + let mut indices = GpuTensor::null_for_test(); + indices.shape = vec![8]; + let mut weights = GpuTensor::null_for_test(); + weights.shape = vec![8]; + let plan = RouterPlan::SoftmaxTopK { + scores: &scores, + topk_indices: &indices, + topk_weights: &weights, + k_top: 8, + normalize: true, + }; + assert_eq!(plan.selection(), RouterSelection::SoftmaxTopK); + assert_eq!(plan.k_top(), 8); + assert!(plan.normalizes()); + plan.validate_against(8, 1).unwrap(); + } + + #[test] + fn expert_ref_rejects_shape_and_owner_mismatch() { + let mut gate_ptrs = GpuTensor::null_for_test(); + gate_ptrs.shape = vec![4]; + let mut down_ptrs = GpuTensor::null_for_test(); + down_ptrs.shape = vec![4]; + let experts = MoeExpertRef::from_resolved( + &gate_ptrs, + &down_ptrs, + None, + DType::MQ4G256, + 4, + 64, + 128, + &[0, 2], + Some(hipfire_hardware::DimKind::Ep), + ); + experts + .validate_projection_shapes(&[128, 128], &[128, 64]) + .unwrap(); + assert!(experts + .validate_projection_shapes(&[64, 128], &[128, 64]) + .is_err()); + let unknown = MoeExpertRef::from_resolved( + &gate_ptrs, + &down_ptrs, + None, + DType::MQ4G256, + 4, + 64, + 128, + &[4], + Some(hipfire_hardware::DimKind::Ep), + ); + assert!(unknown.validate().is_err()); + } + + #[test] + fn fallback_execution_has_no_protocol() { + assert!(ExpertExecutionPlan::PerExpertFallback.protocol().is_err()); + } } diff --git a/crates/hipfire-dispatch/src/pipeline/mod.rs b/crates/hipfire-dispatch/src/pipeline/mod.rs index b4e80d7f48..aff5b236ee 100644 --- a/crates/hipfire-dispatch/src/pipeline/mod.rs +++ b/crates/hipfire-dispatch/src/pipeline/mod.rs @@ -10,8 +10,10 @@ use hip_bridge; use rdna_compute::{DType, Gpu, GpuTensor}; use std::sync::{LazyLock, OnceLock}; -pub(crate) mod steps; -pub use steps::{execute_steps, FusedPattern, GemvInput, Step}; +pub use steps::{ + execute_steps, execute_steps_mesh, FusedPattern, GemvInput, SealedMoeSchedule, Step, + StepCollective, +}; // #397 Ship 6 — forward-as-pipeline C-design lowered super-op substrate (types // only at this step; not on any live path until wired behind HIPFIRE_FORWARD_LOWERED). @@ -2860,7 +2862,7 @@ pub fn run_moe_prefill_bias_aware( /// MoE grouped-GEMM block size (WMMA tile row count). Must match the /// constant in qwen35.rs and the scatter kernel. -const MOE_GROUPED_BLOCK_M: usize = 16; +pub(crate) const MOE_GROUPED_BLOCK_M: usize = 16; /// Dispatch one grouped-GEMM for the given routed expert dtype. /// @@ -3088,6 +3090,44 @@ fn dispatch_grouped_gemm( } } +/// Execute one generic grouped expert projection through the existing MoE +/// kernel table. The caller supplies the already-resolved typed projection +/// shape; no family or storage representation is inferred here. +#[allow(clippy::too_many_arguments)] +pub(crate) fn run_grouped_moe_gemm( + gpu: &mut Gpu, + dtype: DType, + ptrs: &GpuTensor, + tile_ids: &GpuTensor, + sorted_slot_index: &GpuTensor, + x: &GpuTensor, + y: &GpuTensor, + m: usize, + k: usize, + x_row_div: usize, + m_total: usize, + rows: usize, +) -> Result<(), DispatchError> { + dispatch_grouped_gemm( + gpu, + dtype, + None, + ptrs, + tile_ids, + sorted_slot_index, + x, + y, + m, + k, + x_row_div, + m_total, + rows, + false, + false, + false, + ) +} + /// Qwen3.5 batched MoE prefill routed-expert executor. Verbatim transcription /// of the routed block from `prefill_moe_ffn_body_batched` (qwen35.rs:7281). /// diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index dae7c93d95..a524d810a3 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -4,15 +4,19 @@ //! Op-list interpreter. Phase 2a: GEMV + a fused rmsnorm-rotate producer; empty //! fusion table (all per-op fallback). +use hipfire_hardware::{DeviceMesh, DimKind, Gpus, MeshEpoch}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::sync::OnceLock; use crate::context::DispatchCtx; use crate::families::fused_qkv::{FusedQkvBiasParams, FusedQkvFamily, FusedQkvParams}; use crate::families::gemv::{GemvFamily, GemvParams, RotateInputs, WeightRef}; -use crate::families::rotation::{RotationFamily, RotationParams}; -use crate::types::GemvVariant; -use crate::types::{DispatchError, KernelKey, PipelineOp, RotationPlan, RotationVariant}; +use crate::families::moe::{ + ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeProj, MoeProtocolKind, MoeFamily, + RouterPlan, +}; +/// Routing policy is carried by `RouterPlan`; no standalone score activation +/// can be inserted between route and expert phases. /// Rotation disposition of a Gemv's input. Borrows (never owns a RotatedActivation). pub enum GemvInput<'a> { @@ -81,6 +85,78 @@ pub enum Step<'a> { bias: &'a GpuTensor, dim: usize, }, + /// Typed MoE route. Routing semantics (bias/hash/normalization) are + /// carried by the plan rather than reconstructed by the family. + MoeRoute { + plan: RouterPlan<'a>, + }, + /// Indexed routed expert projection. Every down projection is expanded; + /// the executor-owned `MoeCombine` is the only weighted reduction. + IndexedMoeGemv { + experts: &'a MoeExpertRef<'a>, + which: MoeProj<'a>, + topk_indices: &'a GpuTensor, + input: GemvInput<'a>, + out: &'a GpuTensor, + k_top: usize, + batch_size: usize, + }, + /// Weighted combine for an expanded routed-down result. The executor + /// accepts exactly one combine for a routed chain. + MoeCombine { + down_out: &'a GpuTensor, + topk_weights: &'a GpuTensor, + out: &'a GpuTensor, + hidden: usize, + k_top: usize, + batch_size: usize, + inverse_perm: Option<&'a GpuTensor>, + }, + /// Build the deterministic grouped-GEMM permutation for a prefill batch. + MoeScatter { + topk_indices: &'a GpuTensor, + expert_token_counts: &'a GpuTensor, + expert_offsets: &'a GpuTensor, + sorted_slot_index: &'a GpuTensor, + expert_tile_ids: &'a GpuTensor, + inverse_perm: &'a GpuTensor, + total_slots: usize, + n_experts: usize, + m_total_max: usize, + block_m: usize, + }, + /// Grouped routed expert GEMM. Grouped down is always expanded; a + /// residual-fused grouped down is refused before any GPU work. + GroupedMoeGemm { + experts: &'a MoeExpertRef<'a>, + which: MoeProj<'a>, + sorted_slot_index: &'a GpuTensor, + expert_tile_ids: &'a GpuTensor, + x: &'a GpuTensor, + y: &'a GpuTensor, + m_total: usize, + batch_size: usize, + k_top: usize, + }, + /// Deinterleave grouped gate/up output into per-slot gate and up tensors. + MoeGateUpUnscatter { + y_grouped: &'a GpuTensor, + sorted_slot_index: &'a GpuTensor, + gate_batch: &'a GpuTensor, + up_batch: &'a GpuTensor, + inter: usize, + k_top: usize, + m_total: usize, + }, + /// Activation/rotation between routed gate/up and down. + MoeActivation { + variant: MoeActivationVariant, + gate: &'a GpuTensor, + up: &'a GpuTensor, + rot_out: &'a GpuTensor, + inter: usize, + rows: usize, + }, } /// Op-kind for fusion matching. Total over Step variants. @@ -93,9 +169,781 @@ fn op_kind(step: &Step) -> PipelineOp { Step::Rope { .. } => PipelineOp::Rope, Step::QkNorm { .. } => PipelineOp::QkNorm, Step::BiasAdd { .. } => PipelineOp::BiasAdd, + Step::MoeRoute { .. } => PipelineOp::MoeRoute, + Step::IndexedMoeGemv { .. } => PipelineOp::IndexedMoeGemv, + Step::MoeCombine { .. } => PipelineOp::MoeCombine, + Step::MoeScatter { .. } => PipelineOp::MoeScatter, + Step::GroupedMoeGemm { .. } => PipelineOp::GroupedMoeGemm, + Step::MoeGateUpUnscatter { .. } => PipelineOp::MoeGateUpUnscatter, + Step::MoeActivation { .. } => PipelineOp::MoeActivation, + } +} + +/// Collective attached to one lock-step `Step` position. The descriptor is +/// immutable schedule data: membership and mesh identity are supplied by the +/// manifest/topology owner, never inferred by a family. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum StepCollective { + None, + AllReduce { + kind: DimKind, + dim: usize, + group: Vec, + mesh: MeshEpoch, + rank: usize, + }, +} + +impl StepCollective { + pub fn all_reduce( + kind: DimKind, + dim: usize, + group: Vec, + mesh: MeshEpoch, + rank: usize, + ) -> Self { + Self::AllReduce { + kind, + dim, + group, + mesh, + rank, + } + } +} +fn collective_count(collectives: &[StepCollective]) -> usize { + collectives + .iter() + .filter(|collective| matches!(collective, StepCollective::AllReduce { .. })) + .count() +} + +/// Validate the complete grammar of an admitted typed MoE program before any +/// GPU work. The two executable protocols are intentionally exact: +/// +/// ```text +/// indexed: route → gate/up → activation → down(expanded) → combine +/// grouped: route → scatter → gate/up → unscatter → activation +/// → down(expanded) → combine +/// ``` +/// +/// Every routed operand is checked against the one route and expert view that +/// owns it. This makes a hand-built family schedule fail closed instead of +/// silently selecting a second policy or reduction. +pub fn validate_moe_step_schedule( + steps: &[Step], + collectives: &[StepCollective], +) -> Result<(), DispatchError> { + if steps.len() != collectives.len() { + return Err(DispatchError::Hip(format!( + "MoE schedule has {} steps but {} collective descriptors", + steps.len(), + collectives.len() + ))); + } + if steps.is_empty() { + return Err(DispatchError::Hip("MoE schedule is empty".into())); + } + + let (protocol, route, experts, batch_size, combine_index, hidden, route_indices, route_weights) = + match steps { + [ + Step::MoeRoute { plan }, + Step::IndexedMoeGemv { + experts, + which: MoeProj::GateUp { up_out }, + topk_indices, + input: GemvInput::Prerotated(x), + out: gate, + k_top, + batch_size, + }, + Step::MoeActivation { + variant, + gate: act_gate, + up, + rot_out, + inter, + rows, + }, + Step::IndexedMoeGemv { + experts: down_experts, + which: MoeProj::DownExpanded, + topk_indices: down_indices, + input: GemvInput::Prerotated(down_input), + out: down_out, + k_top: down_k, + batch_size: down_batch, + }, + Step::MoeCombine { + down_out: combine_down, + topk_weights, + out, + hidden, + k_top: combine_k, + batch_size: combine_batch, + inverse_perm, + }, + ] => { + if *batch_size != 1 || *down_batch != 1 || inverse_perm.is_some() { + return Err(DispatchError::Hip( + "indexed MoE grammar requires decode batch=1 and no inverse permutation" + .into(), + )); + } + if *k_top != *down_k || *k_top != *combine_k || *batch_size != *combine_batch { + return Err(DispatchError::Hip( + "indexed MoE route width/batch differs across phases".into(), + )); + } + if !std::ptr::eq(*experts, *down_experts) { + return Err(DispatchError::Hip( + "indexed MoE phases use different expert owner views".into(), + )); + } + if !same_tensor(gate, act_gate) + || !same_tensor(up_out, up) + || !same_tensor(rot_out, down_input) + || !same_tensor(down_out, combine_down) + { + return Err(DispatchError::Hip( + "indexed MoE phase operands are not identity-linked".into(), + )); + } + if !same_tensor(topk_indices, down_indices) { + return Err(DispatchError::Hip( + "indexed MoE phases use different route-index buffers".into(), + )); + } + ( + MoeProtocolKind::Indexed, + plan, + *experts, + *batch_size, + 4usize, + *hidden, + topk_indices, + topk_weights, + ) + } + [ + Step::MoeRoute { plan }, + Step::MoeScatter { + topk_indices, + expert_token_counts, + expert_offsets, + sorted_slot_index, + expert_tile_ids, + inverse_perm, + total_slots, + n_experts, + m_total_max, + block_m, + }, + Step::GroupedMoeGemm { + experts, + which: MoeProj::GateUp { .. }, + sorted_slot_index: gate_sorted, + expert_tile_ids: gate_tiles, + x: gate_x, + y: gate_y, + m_total: gate_m_total, + batch_size, + k_top, + }, + Step::MoeGateUpUnscatter { + y_grouped, + sorted_slot_index: unscatter_sorted, + gate_batch, + up_batch, + inter, + k_top: unscatter_k, + m_total: unscatter_m_total, + }, + Step::MoeActivation { + variant: _, + gate: act_gate, + up: act_up, + rot_out, + inter: act_inter, + rows, + }, + Step::GroupedMoeGemm { + experts: down_experts, + which: MoeProj::DownExpanded, + sorted_slot_index: down_sorted, + expert_tile_ids: down_tiles, + x: down_x, + y: down_y, + m_total: down_m_total, + batch_size: down_batch, + k_top: down_k, + }, + Step::MoeCombine { + down_out: combine_down, + topk_weights, + out, + hidden, + k_top: combine_k, + batch_size: combine_batch, + inverse_perm: combine_inverse, + }, + ] => { + if *k_top != *unscatter_k + || *k_top != *down_k + || *k_top != *combine_k + || *batch_size != *down_batch + || *batch_size != *combine_batch + || *total_slots != batch_size.saturating_mul(*k_top) + || !same_tensor(topk_indices, route.route_buffers().0) + || !same_tensor(topk_indices, gate_sorted) + || !same_tensor(topk_indices, down_sorted) + { + return Err(DispatchError::Hip( + "grouped MoE route width, batch, or route identity mismatch".into(), + )); + } + if !std::ptr::eq(*experts, *down_experts) { + return Err(DispatchError::Hip( + "grouped MoE phases use different expert owner views".into(), + )); + } + if *n_experts == 0 + || *m_total_max == 0 + || *block_m == 0 + || !m_total_max.is_multiple_of(*block_m) + || *gate_m_total == 0 + || *down_m_total != *gate_m_total + || *unscatter_m_total != *gate_m_total + || *gate_m_total > *m_total_max + { + return Err(DispatchError::Hip( + "grouped MoE scatter capacity or tile geometry is invalid".into(), + )); + } + if expert_token_counts.numel() == 0 + || expert_offsets.numel() < *n_experts + 1 + { + return Err(DispatchError::Hip( + "grouped MoE scatter metadata has empty or short buffers".into(), + )); + } + let Some(combine_inverse) = *combine_inverse else { + return Err(DispatchError::Hip( + "grouped MoE combine requires the scatter inverse permutation".into(), + )); + }; + if !same_tensor(gate_y, y_grouped) + || !same_tensor(gate_sorted, unscatter_sorted) + || !same_tensor(gate_batch, act_gate) + || !same_tensor(up_batch, act_up) + || *inter != *act_inter + || *rows != batch_size.saturating_mul(*k_top) + || !same_tensor(rot_out, down_x) + || !same_tensor(down_y, combine_down) + || !same_tensor(inverse_perm, combine_inverse) + { + return Err(DispatchError::Hip( + "grouped MoE phase operands are not identity-linked".into(), + )); + } + ( + MoeProtocolKind::Grouped, + plan, + *experts, + *batch_size, + 6usize, + *hidden, + topk_indices, + topk_weights, + ) + } + _ => { + return Err(DispatchError::Hip( + "MoE Step grammar must be exactly indexed or grouped".into(), + )) + } + }; + + experts.validate()?; + route.validate_against(experts.n_experts(), batch_size)?; + if !same_tensor(route_indices, route.route_buffers().0) + || !same_tensor(route_weights, route.route_buffers().1) + { + return Err(DispatchError::Hip( + "MoE route metadata is not bound to the concrete expert Steps".into(), + )); + } + if hidden == 0 { + return Err(DispatchError::Hip("MoE combine hidden size is zero".into())); + } + let shape_gate = [2 * experts.expert_m(), experts.expert_k()]; + let shape_down = [experts.expert_k(), experts.expert_m()]; + experts.validate_projection_shapes(&shape_gate, &shape_down)?; + let dtype_supported = match protocol { + MoeProtocolKind::Indexed => matches!( + experts.dtype(), + DType::MQ4G256 | DType::MQ6G256 | DType::MQ4G256V2 | DType::MQ6G256V2 + ), + MoeProtocolKind::Grouped => matches!( + experts.dtype(), + DType::MQ2G256Lloyd + | DType::MQ2G256LloydU + | DType::MQ3G256Lloyd + | DType::MQ4G256 + | DType::MQ4G256V2 + | DType::MQ6G256 + | DType::MQ6G256V2 + | DType::MFP4G32E8 + | DType::ParoQ4G128 + ), + }; + if !dtype_supported { + return Err(DispatchError::Hip(format!( + "MoE {:?} protocol has no executable kernel for {:?}", + protocol, + experts.dtype() + ))); + } + validate_step_tensors(steps, experts, batch_size, hidden)?; + validate_collectives(collectives, combine_index, hidden, experts.collective_kind())?; + Ok(()) +} + +/// Couple the immutable plan identity to the exact executable grammar. +pub fn validate_moe_protocol_schedule( + steps: &[Step], + collectives: &[StepCollective], + execution: ExpertExecutionPlan, +) -> Result<(), DispatchError> { + let expected = execution.protocol()?; + validate_moe_step_schedule(steps, collectives)?; + let actual = if matches!(steps.get(1), Some(Step::MoeScatter { .. })) { + MoeProtocolKind::Grouped + } else { + MoeProtocolKind::Indexed + }; + if expected != actual { + return Err(DispatchError::Hip(format!( + "MoE execution identity {:?} does not match {:?} Step grammar", + expected, actual + ))); + } + Ok(()) +} + +/// An immutable, pre-validated typed MoE schedule. +pub struct SealedMoeSchedule<'a> { + execution: ExpertExecutionPlan, + steps: Vec>, + collectives: Vec, +} + +impl<'a> SealedMoeSchedule<'a> { + pub fn new( + execution: ExpertExecutionPlan, + steps: Vec>, + collectives: Vec, + ) -> Result { + validate_moe_protocol_schedule(&steps, &collectives, execution)?; + Ok(Self { + execution, + steps, + collectives, + }) + } + + pub fn execution(&self) -> ExpertExecutionPlan { + self.execution + } + + pub fn steps(&self) -> &[Step<'a>] { + &self.steps + } + + pub fn collectives(&self) -> &[StepCollective] { + &self.collectives } } +pub fn execute_sealed_steps( + gpu: &mut Gpu, + ctx: &DispatchCtx, + schedule: &SealedMoeSchedule<'_>, +) -> Result<(), DispatchError> { + validate_moe_protocol_schedule(&schedule.steps, &schedule.collectives, schedule.execution)?; + execute_steps_inner(gpu, ctx, &schedule.steps) +} + +/// Execute one validated MoE schedule per participating device and perform +/// its single manifest-owned routed reduction. The `rank_steps` order is the +/// exact order of `StepCollective::AllReduce::group`; no flat device ordering +/// is inferred. Single-device plans use `execute_steps` and have no reduction. +pub fn execute_steps_mesh<'a>( + gpus: &mut Gpus, + mesh: &DeviceMesh, + ctx: &DispatchCtx, + rank_steps: &[&[Step<'a>]], + collectives: &[StepCollective], +) -> Result<(), DispatchError> { + let reduction = collectives + .iter() + .find_map(|collective| match collective { + StepCollective::AllReduce { + kind, + dim, + group, + mesh: epoch, + rank, + } => Some((*kind, *dim, group, *epoch, *rank)), + StepCollective::None => None, + }) + .ok_or_else(|| DispatchError::Hip("parallel MoE schedule has no collective".into()))?; + let (kind, dim, group, epoch, rank) = reduction; + if mesh.epoch() != epoch { + return Err(DispatchError::Hip( + "MoE collective belongs to a different mesh generation".into(), + )); + } + if mesh.n_devices() != gpus.devices.len() { + return Err(DispatchError::Hip(format!( + "MoE mesh has {} devices but Gpus owns {}", + mesh.n_devices(), + gpus.devices.len() + ))); + } + if group.len() < 2 || rank >= group.len() || rank_steps.len() != group.len() { + return Err(DispatchError::Hip( + "MoE collective rank count does not match mesh group".into(), + )); + } + if group.iter().enumerate().any(|(index, device)| { + *device >= mesh.n_devices() || group[..index].contains(device) + }) { + return Err(DispatchError::Hip( + "MoE collective group contains invalid or duplicate devices".into(), + )); + } + let group_is_mesh_group = group.iter().any(|device| { + let expected = mesh.group_along(kind, &mesh.coord_of(*device)); + expected == *group + }); + if !group_is_mesh_group { + return Err(DispatchError::Hip( + "MoE collective group is not a named DeviceMesh axis group".into(), + )); + } + + for steps in rank_steps { + validate_moe_parallel_schedule(steps, collectives)?; + } + for (index, &device) in group.iter().enumerate() { + if device >= gpus.devices.len() { + return Err(DispatchError::Hip(format!( + "MoE collective device {device} is outside the Gpus owner" + ))); + } + execute_steps_inner(&mut gpus.devices[device], ctx, rank_steps[index])?; + } + let buffers: Vec<&hip_bridge::DeviceBuffer> = rank_steps + .iter() + .map(|steps| { + steps + .iter() + .find_map(|step| match step { + Step::MoeCombine { out, .. } => Some(&out.buf), + _ => None, + }) + .ok_or_else(|| { + DispatchError::Hip("parallel MoE schedule has no combine output".into()) + }) + }) + .collect::>()?; + gpus.all_reduce_sum_f32_peer(group, &buffers, dim) + .map_err(|error| DispatchError::Hip(error.to_string())) +} + +fn same_tensor(a: &GpuTensor, b: &GpuTensor) -> bool { + std::ptr::eq(a, b) + || (!a.buf.as_ptr().is_null() && a.buf.as_ptr() == b.buf.as_ptr()) +} + +fn require_tensor( + tensor: &GpuTensor, + name: &str, + dtype: DType, + capacity: usize, +) -> Result<(), DispatchError> { + if tensor.dtype != dtype || tensor.numel() < capacity { + return Err(DispatchError::Hip(format!( + "MoE {name} has dtype {:?}/capacity {}, expected {:?}/at least {}", + tensor.dtype, + tensor.numel(), + dtype, + capacity + ))); + } + Ok(()) +} + +fn validate_step_tensors( + steps: &[Step], + experts: &MoeExpertRef<'_>, + batch_size: usize, + hidden: usize, +) -> Result<(), DispatchError> { + let slots = batch_size + .checked_mul( + steps + .iter() + .find_map(|step| match step { + Step::IndexedMoeGemv { k_top, .. } + | Step::GroupedMoeGemm { k_top, .. } + | Step::MoeGateUpUnscatter { k_top, .. } + | Step::MoeCombine { k_top, .. } => Some(*k_top), + _ => None, + }) + .unwrap_or(0), + ) + .ok_or_else(|| DispatchError::Hip("MoE slot capacity overflow".into()))?; + require_tensor( + steps + .iter() + .find_map(|step| match step { + Step::MoeRoute { plan } => Some(plan.route_buffers().0), + _ => None, + }) + .expect("validated grammar has a route"), + "route indices", + DType::F32, + slots, + )?; + require_tensor( + steps + .iter() + .find_map(|step| match step { + Step::MoeRoute { plan } => Some(plan.route_buffers().1), + _ => None, + }) + .expect("validated grammar has a route"), + "route weights", + DType::F32, + slots, + )?; + let gate_capacity = slots + .checked_mul(experts.expert_m()) + .ok_or_else(|| DispatchError::Hip("MoE gate capacity overflow".into()))?; + let down_capacity = slots + .checked_mul(experts.expert_k()) + .ok_or_else(|| DispatchError::Hip("MoE down capacity overflow".into()))?; + for step in steps { + match step { + Step::IndexedMoeGemv { + which: MoeProj::GateUp { up_out }, + input, + out, + .. + } => { + let x = match input { + GemvInput::Prerotated(x) => *x, + GemvInput::Raw(_) => { + return Err(DispatchError::Hip( + "generic indexed MoE requires pre-rotated input".into(), + )) + } + }; + require_tensor(x, "indexed gate/up input", DType::F32, batch_size * experts.expert_k())?; + require_tensor(out, "indexed gate output", DType::F32, gate_capacity)?; + require_tensor(up_out, "indexed up output", DType::F32, gate_capacity)?; + if same_tensor(out, up_out) || same_tensor(x, out) || same_tensor(x, up_out) { + return Err(DispatchError::Hip( + "indexed MoE gate/up buffers must not alias".into(), + )); + } + } + Step::IndexedMoeGemv { + which: MoeProj::DownExpanded, + input, + out, + .. + } => { + let x = match input { + GemvInput::Prerotated(x) => *x, + GemvInput::Raw(_) => { + return Err(DispatchError::Hip( + "generic indexed MoE requires pre-rotated input".into(), + )) + } + }; + require_tensor(x, "indexed down input", DType::F32, gate_capacity)?; + require_tensor(out, "indexed down output", DType::F32, down_capacity)?; + if same_tensor(x, out) { + return Err(DispatchError::Hip( + "indexed MoE down input/output must not alias".into(), + )); + } + } + Step::MoeCombine { out, hidden, .. } => { + require_tensor( + out, + "combine output", + DType::F32, + batch_size + .checked_mul(*hidden) + .ok_or_else(|| DispatchError::Hip("MoE output capacity overflow".into()))?, + )?; + } + Step::MoeScatter { + expert_token_counts, + expert_offsets, + sorted_slot_index, + expert_tile_ids, + inverse_perm, + total_slots, + n_experts, + m_total_max, + block_m, + .. + } => { + require_tensor(expert_token_counts, "expert counts", DType::F32, *n_experts)?; + require_tensor( + expert_offsets, + "expert offsets", + DType::F32, + n_experts + 1, + )?; + require_tensor(sorted_slot_index, "sorted slots", DType::F32, *m_total_max)?; + require_tensor(inverse_perm, "inverse permutation", DType::F32, *total_slots)?; + require_tensor( + expert_tile_ids, + "expert tile ids", + DType::F32, + m_total_max / block_m, + )?; + } + Step::GroupedMoeGemm { x, y, m_total, .. } => { + require_tensor(x, "grouped input", DType::F32, x.numel())?; + require_tensor( + y, + "grouped output", + DType::F32, + m_total + .checked_mul(if matches!( + step, + Step::GroupedMoeGemm { + which: MoeProj::GateUp { .. }, + .. + } + ) { + 2 * experts.expert_m() + } else { + experts.expert_k() + }) + .ok_or_else(|| DispatchError::Hip("grouped output capacity overflow".into()))?, + )?; + } + Step::MoeGateUpUnscatter { + gate_batch, + up_batch, + inter, + .. + } => { + require_tensor(gate_batch, "unscatter gate", DType::F32, slots * inter)?; + require_tensor(up_batch, "unscatter up", DType::F32, slots * inter)?; + } + Step::MoeActivation { + gate, + up, + rot_out, + inter, + rows, + .. + } => { + let capacity = rows + .checked_mul(*inter) + .ok_or_else(|| DispatchError::Hip("MoE activation capacity overflow".into()))?; + require_tensor(gate, "activation gate", DType::F32, capacity)?; + require_tensor(up, "activation up", DType::F32, capacity)?; + require_tensor(rot_out, "activation output", DType::F32, capacity)?; + if same_tensor(gate, up) || same_tensor(gate, rot_out) || same_tensor(up, rot_out) { + return Err(DispatchError::Hip( + "MoE activation buffers must not alias".into(), + )); + } + } + _ => {} + } + } + let _ = hidden; + Ok(()) +} + +fn validate_collectives( + collectives: &[StepCollective], + combine_index: usize, + hidden: usize, + expected_kind: Option, +) -> Result<(), DispatchError> { + let mut reduction = None; + for (index, collective) in collectives.iter().enumerate() { + let StepCollective::AllReduce { + kind, + dim, + group, + mesh: _, + rank, + } = collective + else { + continue; + }; + if index != combine_index { + return Err(DispatchError::Hip( + "MoE routed collective must be attached to combine".into(), + )); + } + if expected_kind != Some(*kind) { + return Err(DispatchError::Hip(format!( + "MoE collective axis {kind:?} does not match owner axis {expected_kind:?}" + ))); + } + if *dim != hidden || group.is_empty() || *rank >= group.len() { + return Err(DispatchError::Hip( + "MoE collective rank/group/output dimension is invalid".into(), + )); + } + if group.iter().enumerate().any(|(offset, device)| { + group[..offset].contains(device) + }) { + return Err(DispatchError::Hip( + "MoE collective group contains duplicate devices".into(), + )); + } + if reduction.replace(*kind).is_some() { + return Err(DispatchError::Hip( + "MoE schedule contains duplicate routed collectives".into(), + )); + } + } + Ok(()) +} + +/// Parallel MoE schedule guard. Single-device execution intentionally leaves +/// all descriptors as `None` (the all-reduce is the identity); parallel +/// execution requires the one post-combine collective emitted by the manifest +/// plan. +pub fn validate_moe_parallel_schedule( + steps: &[Step], + collectives: &[StepCollective], +) -> Result<(), DispatchError> { + validate_moe_step_schedule(steps, collectives)?; + if collective_count(collectives) != 1 { + return Err(DispatchError::Hip( + "parallel MoE schedule requires exactly one routed collective".into(), + )); + } + Ok(()) +} + // ── Guard helpers ────────────────────────────────────────────────────────── /// Extract the dtype of the first Gemv step in the window (step index 1, @@ -645,22 +1493,28 @@ const FUSED_TABLE: &[FusedPattern] = &[ static GEMV: OnceLock = OnceLock::new(); static ROTATION: OnceLock = OnceLock::new(); static FUSED_QKV: OnceLock = OnceLock::new(); +static MOE: std::sync::LazyLock = std::sync::LazyLock::new(MoeFamily::new); pub fn execute_steps( gpu: &mut Gpu, ctx: &DispatchCtx, steps: &[Step], +) -> Result<(), DispatchError> { + if steps.iter().any(is_moe_step) { + let collectives = vec![StepCollective::None; steps.len()]; + validate_moe_step_schedule(steps, &collectives)?; + } + execute_steps_inner(gpu, ctx, steps) +} + +fn execute_steps_inner( + gpu: &mut Gpu, + ctx: &DispatchCtx, + steps: &[Step], ) -> Result<(), DispatchError> { let mut i = 0; while i < steps.len() { if let Some((key, len)) = match_prefix(FUSED_TABLE, &steps[i..], ctx) { - // ── QKV bias fold (HIPFIRE_FUSE_QKV_BIAS) ──────────────────────── - // When the flag is on, the matched window is a per-row 3-way QKV - // decode key whose kernel supports the fold, and the 3 steps right - // after the window are `BiasAdd` on the q/k/v outputs in order, fold - // the bias into the kernel's lane-0 store and skip those 3 steps. - // The fold is `acc + bias[row]` (fp32, same operand order as the - // separate `bias_add`) → byte-identical to the unfused path. if ctx.flags.fuse_qkv_bias && len == QKV3.len() && qkv_bias_fold_supported(key, ctx) { if let Some(biases) = match_trailing_qkv_bias(&steps[i..], len) { launch_fused_qkv_with_bias(gpu, ctx, key, &steps[i..i + len], biases)?; @@ -668,7 +1522,6 @@ pub fn execute_steps( continue; } } - // ───────────────────────────────────────────────────────────────── launch_fused(gpu, ctx, key, &steps[i..i + len])?; i += len; } else { @@ -679,6 +1532,19 @@ pub fn execute_steps( Ok(()) } +fn is_moe_step(step: &Step<'_>) -> bool { + matches!( + step, + Step::MoeRoute { .. } + | Step::IndexedMoeGemv { .. } + | Step::MoeCombine { .. } + | Step::MoeScatter { .. } + | Step::GroupedMoeGemm { .. } + | Step::MoeGateUpUnscatter { .. } + | Step::MoeActivation { .. } + ) +} + /// Keys whose 3-way QKV **decode** dispatch arm folds the optional Q/K/V bias /// into the kernel (a `_with_bias` kernel variant exists and is wired in /// `dispatch_fused_qkv`). The fold is additionally guarded off on dp4a archs @@ -999,6 +1865,115 @@ fn launch_op(gpu: &mut Gpu, ctx: &DispatchCtx, step: &Step) -> Result<(), Dispat } => gpu .rmsnorm_batched(x, weight, x, *n_groups, *head_dim, *eps) .map_err(|e| DispatchError::Hip(e.to_string())), + Step::MoeRoute { plan } => MOE.run_route(gpu, plan), + Step::IndexedMoeGemv { + experts, + which, + topk_indices, + input, + out, + k_top, + batch_size, + } => MOE.run_indexed( + gpu, + experts, + which, + topk_indices, + input, + out, + *k_top, + *batch_size, + ), + Step::MoeScatter { + topk_indices, + expert_token_counts, + expert_offsets, + sorted_slot_index, + expert_tile_ids, + inverse_perm, + total_slots, + n_experts, + m_total_max, + block_m, + } => MOE.run_scatter( + gpu, + topk_indices, + expert_token_counts, + expert_offsets, + sorted_slot_index, + expert_tile_ids, + inverse_perm, + *total_slots, + *n_experts, + *m_total_max, + *block_m, + ), + Step::GroupedMoeGemm { + experts, + which, + sorted_slot_index, + expert_tile_ids, + x, + y, + m_total, + batch_size, + k_top, + } => MOE.run_grouped( + gpu, + experts, + which, + sorted_slot_index, + expert_tile_ids, + x, + y, + *m_total, + *batch_size, + *k_top, + ), + Step::MoeGateUpUnscatter { + y_grouped, + sorted_slot_index, + gate_batch, + up_batch, + inter, + k_top, + m_total, + } => MOE.run_unscatter( + gpu, + y_grouped, + sorted_slot_index, + gate_batch, + up_batch, + *inter, + *k_top, + *m_total, + ), + Step::MoeActivation { + variant, + gate, + up, + rot_out, + inter, + rows, + } => MOE.run_activation(gpu, *variant, gate, up, rot_out, *inter, *rows), + Step::MoeCombine { + down_out, + topk_weights, + out, + hidden, + k_top, + batch_size, + inverse_perm, + } => MOE.run_combine( + gpu, + down_out, + topk_weights, + out, + *hidden, + *k_top, + *batch_size, + *inverse_perm, + ), Step::BiasAdd { x, bias, dim } => gpu .bias_add_f32(x, bias, 1, *dim) .map_err(|e| DispatchError::Hip(e.to_string())), @@ -1556,4 +2531,173 @@ mod tests { "FusedGateUpQ8_0 missing from FUSED_TABLE" ); } + fn tensor(shape: Vec) -> GpuTensor { + let mut tensor = GpuTensor::null_for_test(); + tensor.shape = shape; + tensor + } + + fn indexed_steps<'a>( + experts: &'a MoeExpertRef<'a>, + scores: &'a GpuTensor, + indices: &'a GpuTensor, + weights: &'a GpuTensor, + x: &'a GpuTensor, + gate: &'a GpuTensor, + up: &'a GpuTensor, + rot: &'a GpuTensor, + down: &'a GpuTensor, + out: &'a GpuTensor, + ) -> Vec> { + vec![ + Step::MoeRoute { + plan: RouterPlan::SoftmaxTopK { + scores, + topk_indices: indices, + topk_weights: weights, + k_top: 8, + normalize: true, + }, + }, + Step::IndexedMoeGemv { + experts, + which: MoeProj::GateUp { up_out: up }, + topk_indices: indices, + input: GemvInput::Prerotated(x), + out: gate, + k_top: 8, + batch_size: 1, + }, + Step::MoeActivation { + variant: MoeActivationVariant::SiluMul, + gate, + up, + rot_out: rot, + inter: 64, + rows: 8, + }, + Step::IndexedMoeGemv { + experts, + which: MoeProj::DownExpanded, + topk_indices: indices, + input: GemvInput::Prerotated(rot), + out: down, + k_top: 8, + batch_size: 1, + }, + Step::MoeCombine { + down_out: down, + topk_weights: weights, + out, + hidden: 128, + k_top: 8, + batch_size: 1, + inverse_perm: None, + }, + ] + } + + + #[test] + fn typed_indexed_grammar_requires_one_combine() { + let scores = tensor(vec![8]); + let indices = tensor(vec![8]); + let weights = tensor(vec![8]); + let x = tensor(vec![128]); + let gate = tensor(vec![8 * 64]); + let up = tensor(vec![8 * 64]); + let rot = tensor(vec![8 * 64]); + let down = tensor(vec![8 * 128]); + let out = tensor(vec![128]); + let gate_ptrs = tensor(vec![4]); + let down_ptrs = tensor(vec![4]); + let experts = MoeExpertRef::from_resolved( + &gate_ptrs, + &down_ptrs, + None, + DType::MQ4G256, + 8, + 64, + 128, + &[0], + Some(DimKind::Ep), + ); + let steps = indexed_steps( + &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, + ); + validate_moe_step_schedule( + &steps, + &[ + StepCollective::None, + StepCollective::None, + StepCollective::None, + StepCollective::None, + StepCollective::None, + ], + ) + .unwrap(); + let mut duplicate = steps; + duplicate.push(Step::MoeCombine { + down_out: &down, + topk_weights: &weights, + out: &out, + hidden: 128, + k_top: 8, + batch_size: 1, + inverse_perm: None, + }); + let error = validate_moe_step_schedule( + &duplicate, + &vec![StepCollective::None; duplicate.len()], + ) + .expect_err("duplicate combine must be rejected by exact grammar"); + assert!(error.to_string().contains("exactly indexed or grouped")); + } + + #[test] + fn parallel_schedule_requires_named_axis_collective() { + let scores = tensor(vec![8]); + let indices = tensor(vec![8]); + let weights = tensor(vec![8]); + let x = tensor(vec![128]); + let gate = tensor(vec![8 * 64]); + let up = tensor(vec![8 * 64]); + let rot = tensor(vec![8 * 64]); + let down = tensor(vec![8 * 128]); + let out = tensor(vec![128]); + let gate_ptrs = tensor(vec![8]); + let down_ptrs = tensor(vec![8]); + let experts = MoeExpertRef::from_resolved( + &gate_ptrs, + &down_ptrs, + None, + DType::MQ4G256, + 8, + 64, + 128, + &[0, 1], + Some(DimKind::Ep), + ); + let steps = indexed_steps( + &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, + ); + let mesh = DeviceMesh::rect(&[(DimKind::Ep, 2)]).unwrap(); + let mut collectives = vec![StepCollective::None; steps.len()]; + collectives[4] = StepCollective::all_reduce( + DimKind::Ep, + 128, + vec![0, 1], + mesh.epoch(), + 0, + ); + validate_moe_parallel_schedule(&steps, &collectives).unwrap(); + collectives[4] = StepCollective::all_reduce( + DimKind::Tp, + 128, + vec![0, 1], + mesh.epoch(), + 0, + ); + assert!(validate_moe_parallel_schedule(&steps, &collectives).is_err()); + } } diff --git a/crates/hipfire-dispatch/src/types.rs b/crates/hipfire-dispatch/src/types.rs index a191089a0e..589ad552d0 100644 --- a/crates/hipfire-dispatch/src/types.rs +++ b/crates/hipfire-dispatch/src/types.rs @@ -26,6 +26,18 @@ pub enum PipelineOp { IndexedGateUp, IndexedDownExpanded, MoeCombine, + /// Bias-aware or precomputed top-k routing owned by the executor. + MoeRoute, + /// Typed routed expert projection (gate/up or down). + IndexedMoeGemv, + /// Grouped routed expert prefill scatter. + MoeScatter, + /// Grouped routed expert GEMM. + GroupedMoeGemm, + /// Deinterleave grouped gate/up output. + MoeGateUpUnscatter, + /// Routed SwiGLU/rotation activation. + MoeActivation, /// Fused rmsnorm + optional rotation (MQ-weight producer step). /// rotation=FwhtG256 → rmsnorm + FWHT. rotation=None → rmsnorm only. RmsnormAutomatic, diff --git a/crates/hipfire-runtime/src/lib.rs b/crates/hipfire-runtime/src/lib.rs index b8c833190d..333480dc7d 100644 --- a/crates/hipfire-runtime/src/lib.rs +++ b/crates/hipfire-runtime/src/lib.rs @@ -46,6 +46,7 @@ pub mod llama_spec; pub mod loader_api; pub mod loop_guard; pub mod model_load; +pub mod moe_plan; pub mod model_source; pub mod paro; pub mod prefix; diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs new file mode 100644 index 0000000000..e97aa77169 --- /dev/null +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -0,0 +1,1076 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 hipfire contributors + +//! Manifest-derived MoE placement, rank views, and storage ownership. +//! +//! This module is the only owner of the resolved expert placement contract. +//! It is CPU-only: source identities and logical shapes come from G3, topology +//! comes from G1, and architecture carriers bind borrowed pointer tables through +//! [`ExpertPlan::bind_expert_ref`]. No family receives an allocator or a store +//! representation, and no family-side teardown path exists. + +use std::fmt; + +use hipfire_dispatch::families::moe::{ExpertExecutionPlan, MoeExpertRef}; +use hipfire_dispatch::pipeline::StepCollective; +use hipfire_hardware::{CollectiveHint, DeviceMesh, DimKind, MeshEpoch}; +use rdna_compute::{DType, GpuTensor}; + +use crate::tp_shard::ExpertAssign; +use crate::weight_manifest::{ + collective_schedule, validate_expert_group_specs, ExpertGroupSpec, ExpertParallelism, + ExpertResourceRequirements, ExpertSourceLayout, ShardPolicy, WeightEntry, +}; + +/// One logical expert and its deterministic owner-local slot. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ExpertPlacement { + pub global_id: usize, + /// Rank-local owner index within the named DeviceMesh group. + pub owner: usize, + pub local_slot: usize, +} + +/// Logical dimensions shared by the fused gate/up and down projections. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ExpertShape { + pub expert_m: usize, + pub expert_k: usize, + pub fused_gate_up: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExpertPlanError(String); + +impl ExpertPlanError { + fn new(message: impl Into) -> Self { + Self(message.into()) + } +} + +impl fmt::Display for ExpertPlanError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl std::error::Error for ExpertPlanError {} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct ExpertStorageOwner { + slots: Vec, + resident: Vec, +} + +impl ExpertStorageOwner { + fn new(slots: Vec) -> Self { + Self { + resident: vec![false; slots.len()], + slots, + } + } + + fn index_of(&self, placement: ExpertPlacement) -> Option { + self.slots.iter().position(|candidate| *candidate == placement) + } + + fn clear(&mut self) { + self.resident.fill(false); + } + + fn resident_count(&self) -> usize { + self.resident.iter().filter(|resident| **resident).count() + } +} + +/// A sealed, manifest-derived expert plan. +/// +/// Every placement, source identity, shape, resource requirement, rank-local +/// view, and collective row is private and fixed at construction. The only +/// mutable state is the owner transaction's resident bitmap; it is never +/// transferred to a family. +pub struct ExpertPlan { + group: String, + layer: Option, + n_experts: usize, + parallelism: ExpertParallelism, + assignment: ExpertAssign, + shape: ExpertShape, + source_dtype: DType, + source_layout: ExpertSourceLayout, + resources: ExpertResourceRequirements, + router: String, + execution: String, + execution_plan: ExpertExecutionPlan, + mesh_epoch: MeshEpoch, + group_devices: Vec, + collective: Option, + collective_row: Option, + owner_views: Vec>, + owner: ExpertStorageOwner, +} + +impl fmt::Debug for ExpertPlan { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ExpertPlan") + .field("group", &self.group) + .field("layer", &self.layer) + .field("n_experts", &self.n_experts) + .field("parallelism", &self.parallelism) + .field("assignment", &self.assignment) + .field("shape", &self.shape) + .field("source_dtype", &self.source_dtype) + .field("source_layout", &self.source_layout) + .field("resources", &self.resources) + .field("router", &self.router) + .field("execution", &self.execution) + .field("execution_plan", &self.execution_plan) + .field("mesh_epoch", &self.mesh_epoch) + .field("group_devices", &self.group_devices) + .field("collective", &self.collective) + .field("collective_row", &self.collective_row) + .field("resident_slots", &self.owner.resident_count()) + .finish() + } +} + +impl ExpertPlan { + /// Resolve one expert declaration against the G3 logical manifest and G1 + /// named mesh. No GPU, source file, allocator, or `WeightStore` is touched. + pub fn from_manifest( + spec: &ExpertGroupSpec, + manifest: &[WeightEntry], + mesh: &DeviceMesh, + ) -> Result { + validate_expert_group_specs(std::slice::from_ref(spec), manifest) + .map_err(ExpertPlanError::new)?; + validate_spec(spec, mesh)?; + let (shape, source_dtype) = resolve_shape(spec, manifest)?; + let execution_plan = parse_execution(&spec.execution, spec)?; + validate_execution_dtype(execution_plan, source_dtype, spec)?; + + let group_devices = resolve_group_devices(spec, manifest, mesh); + if group_devices.is_empty() { + return Err(ExpertPlanError::new(format!( + "expert group '{}' resolved to an empty mesh group", + spec.group + ))); + } + let owner = resolve_placements( + spec.n_experts, + group_devices.len(), + spec.parallelism, + spec.assignment, + ); + let owner_views = (0..group_devices.len()) + .map(|rank| { + owner + .iter() + .filter(|placement| placement.owner == rank) + .map(|placement| placement.global_id) + .collect() + }) + .collect(); + let (collective, collective_row) = resolve_collective(spec, manifest)?; + + Ok(Self { + group: spec.group.clone(), + layer: spec.layer, + n_experts: spec.n_experts, + parallelism: spec.parallelism, + assignment: spec.assignment, + shape, + source_dtype, + source_layout: spec.source_layout.clone(), + resources: spec.resources, + router: spec.router.clone(), + execution: spec.execution.clone(), + execution_plan, + mesh_epoch: mesh.epoch(), + group_devices, + collective, + collective_row, + owner_views, + owner: ExpertStorageOwner::new(owner), + }) + } + + pub fn group(&self) -> &str { + &self.group + } + + pub fn layer(&self) -> Option { + self.layer + } + + pub fn n_experts(&self) -> usize { + self.n_experts + } + + pub fn group_size(&self) -> usize { + self.group_devices.len() + } + + pub fn assignment(&self) -> ExpertAssign { + self.assignment + } + + pub fn parallelism(&self) -> ExpertParallelism { + self.parallelism + } + + pub fn shape(&self) -> ExpertShape { + self.shape + } + + pub fn source_dtype(&self) -> DType { + self.source_dtype + } + + pub fn source_layout(&self) -> &ExpertSourceLayout { + &self.source_layout + } + + pub fn resources(&self) -> ExpertResourceRequirements { + self.resources + } + + pub fn router(&self) -> &str { + &self.router + } + + pub fn execution(&self) -> &str { + &self.execution + } + + pub fn execution_plan(&self) -> ExpertExecutionPlan { + self.execution_plan + } + + pub fn mesh_epoch(&self) -> MeshEpoch { + self.mesh_epoch + } + + /// Global device IDs in the exact named-axis order used by collectives. + pub fn group_devices(&self) -> &[usize] { + &self.group_devices + } + + /// The one ordered G3 manifest row that authorizes the routed reduction. + pub fn collective_row(&self) -> Option<&str> { + self.collective_row.as_deref() + } + + /// The single post-combine collective implied by the declared parallelism. + pub fn collective(&self) -> Option { + self.collective + } + + pub fn placements(&self) -> &[ExpertPlacement] { + &self.owner.slots + } + + pub fn placement(&self, global_id: usize, owner: usize) -> Option { + self.owner + .slots + .iter() + .copied() + .find(|placement| placement.global_id == global_id && placement.owner == owner) + } + + pub fn owned_experts(&self, rank: usize) -> Result<&[usize], ExpertPlanError> { + self.owner_views.get(rank).map(Vec::as_slice).ok_or_else(|| { + ExpertPlanError::new(format!( + "expert group '{}' rank {rank} is outside group size {}", + self.group, + self.group_size() + )) + }) + } + + /// Bind canonical shape/dtype/owner metadata to a borrowed pointer-table + /// view. The family can retain only the returned immutable view; all + /// placement and storage state remains owned by this plan. + pub fn bind_expert_ref<'a>( + &'a self, + rank: usize, + gate_up_ptrs: &'a GpuTensor, + down_ptrs: &'a GpuTensor, + dummy_gate_up: Option<&'a GpuTensor>, + ) -> Result, ExpertPlanError> { + let owned = self.owned_experts(rank)?; + let collective_kind = match self.collective { + Some(CollectiveHint::AllReduce { kind }) => Some(kind), + _ => None, + }; + let view = MoeExpertRef::from_resolved( + gate_up_ptrs, + down_ptrs, + dummy_gate_up, + self.source_dtype, + self.n_experts, + self.shape.expert_m, + self.shape.expert_k, + owned, + collective_kind, + ); + view.validate() + .map_err(|error| ExpertPlanError::new(error.to_string()))?; + Ok(view) + } + + /// Build the descriptor for the one collective attached to the combine + /// position. Single is an explicit identity and emits `None`. + pub fn step_collective( + &self, + rank: usize, + dim: usize, + ) -> Result { + if rank >= self.group_size() { + return Err(ExpertPlanError::new(format!( + "expert group '{}' rank {rank} is outside group size {}", + self.group, + self.group_size() + ))); + } + if dim == 0 { + return Err(ExpertPlanError::new( + "expert collective output dimension must be nonzero", + )); + } + match self.collective { + None => Ok(StepCollective::None), + Some(CollectiveHint::AllReduce { kind }) => Ok(StepCollective::all_reduce( + kind, + dim, + self.group_devices.clone(), + self.mesh_epoch, + rank, + )), + Some(CollectiveHint::BandXfer { .. }) => Err(ExpertPlanError::new( + "pipeline band transfer is not an expert reduction", + )), + } + } + + /// Start one owner-controlled load transaction. A dropped or explicitly + /// rolled-back transaction removes only its staged slots. + pub fn begin_load(&mut self) -> ExpertLoadTxn<'_> { + ExpertLoadTxn { + owner: &mut self.owner, + staged: Vec::new(), + finished: false, + } + } + + /// Idempotent teardown. No family-side free path exists; all resident + /// slots return to the owner baseline and repeated unloads are harmless. + pub fn unload(&mut self) { + self.owner.clear(); + } + + pub fn resident_slots(&self) -> usize { + self.owner.resident_count() + } + + pub fn allocated_slots(&self) -> usize { + self.owner.slots.len() + } +} + +/// Owner-scoped transactional expert load state. +pub struct ExpertLoadTxn<'a> { + owner: &'a mut ExpertStorageOwner, + staged: Vec, + finished: bool, +} + +impl ExpertLoadTxn<'_> { + /// Reserve one manifest placement. Duplicate reservations are refused. + pub fn reserve(&mut self, placement: ExpertPlacement) -> Result<(), ExpertPlanError> { + let index = self.owner.index_of(placement).ok_or_else(|| { + ExpertPlanError::new(format!("unknown expert placement {placement:?}")) + })?; + if self.owner.resident[index] { + return Err(ExpertPlanError::new(format!( + "expert placement {placement:?} is already resident" + ))); + } + self.owner.resident[index] = true; + self.staged.push(index); + Ok(()) + } + + /// Commit the staged reservations. The owner remains responsible for + /// teardown; committing never transfers ownership to a family. + pub fn commit(mut self) { + self.finished = true; + } + + /// Roll back staged reservations and close the transaction. + pub fn rollback(&mut self) { + for index in self.staged.drain(..) { + self.owner.resident[index] = false; + } + self.finished = true; + } +} + +impl Drop for ExpertLoadTxn<'_> { + fn drop(&mut self) { + if !self.finished { + for index in self.staged.drain(..) { + self.owner.resident[index] = false; + } + } + } +} + +fn validate_spec(spec: &ExpertGroupSpec, mesh: &DeviceMesh) -> Result<(), ExpertPlanError> { + if spec.group.is_empty() { + return Err(ExpertPlanError::new("expert group identity is empty")); + } + if spec.n_experts == 0 { + return Err(ExpertPlanError::new(format!( + "expert group '{}' has no experts", + spec.group + ))); + } + if spec.resources.bytes_per_expert == 0 { + return Err(ExpertPlanError::new(format!( + "expert group '{}' bytes_per_expert is zero", + spec.group + ))); + } + if spec.resources.alignment == 0 || !spec.resources.alignment.is_power_of_two() { + return Err(ExpertPlanError::new(format!( + "expert group '{}' alignment={} is invalid", + spec.group, spec.resources.alignment + ))); + } + spec.n_experts + .checked_mul(spec.resources.bytes_per_expert) + .ok_or_else(|| ExpertPlanError::new("expert resource capacity overflows usize"))?; + if spec.router.is_empty() || spec.execution.is_empty() { + return Err(ExpertPlanError::new(format!( + "expert group '{}' router/execution identity is empty", + spec.group + ))); + } + let group_size = match spec.parallelism { + ExpertParallelism::Single => 1, + ExpertParallelism::TensorParallel => mesh.size_of(DimKind::Tp), + ExpertParallelism::ExpertParallel => mesh.size_of(DimKind::Ep), + }; + if group_size == 0 { + return Err(ExpertPlanError::new("expert group resolved to zero ranks")); + } + if matches!(spec.parallelism, ExpertParallelism::ExpertParallel) + && !spec.n_experts.is_multiple_of(group_size) + { + return Err(ExpertPlanError::new(format!( + "expert group '{}' n_experts={} is not divisible by group_size={group_size}", + spec.group, spec.n_experts, group_size + ))); + } + Ok(()) +} + +fn parse_execution( + execution: &str, + spec: &ExpertGroupSpec, +) -> Result { + match execution { + "indexed_quantized" => Ok(ExpertExecutionPlan::IndexedQuantized), + "grouped_quantized" => Ok(ExpertExecutionPlan::GroupedQuantized), + "per_expert_fallback" => Err(ExpertPlanError::new(format!( + "expert group '{}' uses PerExpertFallback, which is not a Step protocol", + spec.group + ))), + other => Err(ExpertPlanError::new(format!( + "expert group '{}' has unsupported execution identity '{other}'", + spec.group + ))), + } +} + +fn validate_execution_dtype( + execution: ExpertExecutionPlan, + dtype: DType, + spec: &ExpertGroupSpec, +) -> Result<(), ExpertPlanError> { + let supported = match execution { + ExpertExecutionPlan::IndexedQuantized => matches!( + dtype, + DType::MQ4G256 | DType::MQ6G256 | DType::MQ4G256V2 | DType::MQ6G256V2 + ), + ExpertExecutionPlan::GroupedQuantized => matches!( + dtype, + DType::MQ2G256Lloyd + | DType::MQ2G256LloydU + | DType::MQ3G256Lloyd + | DType::MQ4G256 + | DType::MQ4G256V2 + | DType::MQ6G256 + | DType::MQ6G256V2 + | DType::MFP4G32E8 + | DType::ParoQ4G128 + ), + ExpertExecutionPlan::PerExpertFallback => false, + }; + if supported { + Ok(()) + } else { + Err(ExpertPlanError::new(format!( + "expert group '{}' execution {:?} has no generic kernel for source dtype {dtype:?}", + spec.group, execution + ))) + } +} + +fn resolve_group_devices( + spec: &ExpertGroupSpec, + manifest: &[WeightEntry], + mesh: &DeviceMesh, +) -> Vec { + let n_layers = manifest + .iter() + .filter_map(|entry| entry.layer) + .max() + .map_or(1, |layer| layer.saturating_add(1)); + let mut coord = mesh.coord_of(0); + if let Some(layer) = spec.layer { + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + coord[index] = mesh.stage_for_layer(layer, n_layers); + } + } + match spec.parallelism { + ExpertParallelism::Single => vec![mesh.device_of(&coord)], + ExpertParallelism::TensorParallel => mesh.group_along(DimKind::Tp, &coord), + ExpertParallelism::ExpertParallel => mesh.group_along(DimKind::Ep, &coord), + } +} + +fn resolve_placements( + n_experts: usize, + group_size: usize, + parallelism: ExpertParallelism, + assignment: ExpertAssign, +) -> Vec { + let capacity = match parallelism { + ExpertParallelism::TensorParallel => n_experts.saturating_mul(group_size), + ExpertParallelism::Single | ExpertParallelism::ExpertParallel => n_experts, + }; + let mut next_slot = vec![0usize; group_size]; + let mut placements = Vec::with_capacity(capacity); + for global_id in 0..n_experts { + match parallelism { + ExpertParallelism::Single => placements.push(ExpertPlacement { + global_id, + owner: 0, + local_slot: global_id, + }), + ExpertParallelism::TensorParallel => { + for owner in 0..group_size { + placements.push(ExpertPlacement { + global_id, + owner, + local_slot: global_id, + }); + } + } + ExpertParallelism::ExpertParallel => { + let per = n_experts / group_size; + let owner = match assignment { + ExpertAssign::Contiguous => global_id / per, + ExpertAssign::Stride => global_id % group_size, + }; + let local_slot = next_slot[owner]; + next_slot[owner] += 1; + placements.push(ExpertPlacement { + global_id, + owner, + local_slot, + }); + } + } + } + placements +} + +fn source_names<'a>(layout: &'a ExpertSourceLayout) -> Vec<(&'static str, Vec<&'a str>)> { + match layout { + ExpertSourceLayout::PackedFused { + gate_up, + down, + sidecars, + } => vec![ + ("gate_up", vec![gate_up.as_str()]), + ("down", vec![down.as_str()]), + ("sidecar", sidecars.iter().map(String::as_str).collect()), + ], + ExpertSourceLayout::PackedSeparate { + gate, + up, + down, + sidecars, + } => vec![ + ("gate", vec![gate.as_str()]), + ("up", vec![up.as_str()]), + ("down", vec![down.as_str()]), + ("sidecar", sidecars.iter().map(String::as_str).collect()), + ], + ExpertSourceLayout::PerExpertFused { + gate_up, + down, + sidecars, + } => vec![ + ("gate_up", gate_up.iter().map(String::as_str).collect()), + ("down", down.iter().map(String::as_str).collect()), + ("sidecar", sidecars.iter().map(String::as_str).collect()), + ], + ExpertSourceLayout::PerExpertSeparate { + gate, + up, + down, + sidecars, + } => vec![ + ("gate", gate.iter().map(String::as_str).collect()), + ("up", up.iter().map(String::as_str).collect()), + ("down", down.iter().map(String::as_str).collect()), + ("sidecar", sidecars.iter().map(String::as_str).collect()), + ], + } +} + + +fn entry_for<'a>( + spec: &ExpertGroupSpec, + manifest: &'a [WeightEntry], + label: &str, + name: &str, +) -> Result<&'a WeightEntry, ExpertPlanError> { + manifest + .iter() + .find(|entry| entry.name == name && entry.layer == spec.layer) + .ok_or_else(|| { + ExpertPlanError::new(format!( + "expert group '{}' layer {:?} missing {label} source '{name}'", + spec.group, spec.layer + )) + }) +} + +fn check_shape( + spec: &ExpertGroupSpec, + label: &str, + entry: &WeightEntry, + per_expert: bool, +) -> Result, ExpertPlanError> { + if !entry.dtype_constraint.accepts(entry.dtype) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' {label} source '{}' violates dtype constraint", + spec.group, entry.name + ))); + } + let shape = &entry.logical_shape; + let valid_rank = if per_expert { shape.len() == 2 } else { shape.len() == 3 }; + if !valid_rank || shape.iter().any(|dim| *dim == 0) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' layer {:?} {label} source '{}' has invalid logical_shape {:?}", + spec.group, spec.layer, entry.name, shape + ))); + } + if !per_expert && shape[0] != spec.n_experts { + return Err(ExpertPlanError::new(format!( + "expert group '{}' layer {:?} {label} source '{}' has expert axis {}, expected {}", + spec.group, spec.layer, entry.name, shape[0], spec.n_experts + ))); + } + Ok(if per_expert { + shape.clone() + } else { + shape[1..].to_vec() + }) +} + +fn resolve_shape( + spec: &ExpertGroupSpec, + manifest: &[WeightEntry], +) -> Result<(ExpertShape, DType), ExpertPlanError> { + let per_expert = matches!( + spec.source_layout, + ExpertSourceLayout::PerExpertFused { .. } | ExpertSourceLayout::PerExpertSeparate { .. } + ); + for (label, names) in source_names(&spec.source_layout) { + for name in names { + let entry = entry_for(spec, manifest, label, name)?; + if !entry.dtype_constraint.accepts(entry.dtype) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' {label} source '{}' violates dtype constraint", + spec.group, name + ))); + } + } + } + let (gate_shapes, up_shapes, down_shapes, fused) = match &spec.source_layout { + ExpertSourceLayout::PackedFused { gate_up, down, .. } => ( + vec![check_shape( + spec, + "gate_up", + entry_for(spec, manifest, "gate_up", gate_up)?, + false, + )?], + Vec::new(), + vec![check_shape( + spec, + "down", + entry_for(spec, manifest, "down", down)?, + false, + )?], + true, + ), + ExpertSourceLayout::PackedSeparate { + gate, up, down, .. + } => ( + vec![check_shape( + spec, + "gate", + entry_for(spec, manifest, "gate", gate)?, + false, + )?], + vec![check_shape( + spec, + "up", + entry_for(spec, manifest, "up", up)?, + false, + )?], + vec![check_shape( + spec, + "down", + entry_for(spec, manifest, "down", down)?, + false, + )?], + false, + ), + ExpertSourceLayout::PerExpertFused { gate_up, down, .. } => ( + gate_up + .iter() + .map(|name| check_shape(spec, "gate_up", entry_for(spec, manifest, "gate_up", name)?, true)) + .collect::>()?, + Vec::new(), + down + .iter() + .map(|name| check_shape(spec, "down", entry_for(spec, manifest, "down", name)?, true)) + .collect::>()?, + true, + ), + ExpertSourceLayout::PerExpertSeparate { gate, up, down, .. } => ( + gate + .iter() + .map(|name| check_shape(spec, "gate", entry_for(spec, manifest, "gate", name)?, true)) + .collect::>()?, + up.iter() + .map(|name| check_shape(spec, "up", entry_for(spec, manifest, "up", name)?, true)) + .collect::>()?, + down + .iter() + .map(|name| check_shape(spec, "down", entry_for(spec, manifest, "down", name)?, true)) + .collect::>()?, + false, + ), + }; + if gate_shapes.is_empty() || down_shapes.is_empty() || (!up_shapes.is_empty() && up_shapes[0] != gate_shapes[0]) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' layer {:?} projection source count/shape mismatch", + spec.group, spec.layer + ))); + } + if gate_shapes.iter().any(|shape| shape != &gate_shapes[0]) + || down_shapes.iter().any(|shape| shape != &down_shapes[0]) + { + return Err(ExpertPlanError::new(format!( + "expert group '{}' layer {:?} per-expert projection shape mismatch", + spec.group, spec.layer + ))); + } + let gate = &gate_shapes[0]; + let down = &down_shapes[0]; + if gate.len() != 2 || down.len() != 2 || gate[0] % 2 != 0 { + return Err(ExpertPlanError::new(format!( + "expert group '{}' layer {:?} projection shapes are incompatible", + spec.group, spec.layer + ))); + } + let shape = ExpertShape { + expert_m: gate[0] / 2, + expert_k: gate[1], + fused_gate_up: fused, + }; + if down != &[shape.expert_k, shape.expert_m] { + return Err(ExpertPlanError::new(format!( + "expert group '{}' layer {:?} projection shape mismatch: gate_up={gate:?}, down={down:?}", + spec.group, spec.layer + ))); + } + let source_names = source_names(&spec.source_layout); + let source_dtype = source_names + .iter() + .flat_map(|(_, names)| names.iter()) + .find_map(|name| manifest.iter().find(|entry| entry.name == *name && entry.layer == spec.layer).map(|entry| entry.dtype)) + .ok_or_else(|| ExpertPlanError::new("expert projection has no source dtype"))?; + for (_, names) in source_names { + for name in names { + let entry = entry_for(spec, manifest, "projection", name)?; + if entry.dtype != source_dtype && label_is_projection(name, &spec.source_layout) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' projection source '{}' dtype {:?} differs from {:?}", + spec.group, name, entry.dtype, source_dtype + ))); + } + } + } + Ok((shape, source_dtype)) +} + +fn label_is_projection(name: &str, layout: &ExpertSourceLayout) -> bool { + match layout { + ExpertSourceLayout::PackedFused { gate_up, down, .. } => name == gate_up || name == down, + ExpertSourceLayout::PackedSeparate { gate, up, down, .. } => { + name == gate || name == up || name == down + } + ExpertSourceLayout::PerExpertFused { gate_up, down, .. } => { + gate_up.iter().any(|candidate| candidate == name) + || down.iter().any(|candidate| candidate == name) + } + ExpertSourceLayout::PerExpertSeparate { gate, up, down, .. } => { + gate.iter().any(|candidate| candidate == name) + || up.iter().any(|candidate| candidate == name) + || down.iter().any(|candidate| candidate == name) + } + } +} + +fn down_names(layout: &ExpertSourceLayout) -> Vec<&str> { + match layout { + ExpertSourceLayout::PackedFused { down, .. } + | ExpertSourceLayout::PackedSeparate { down, .. } => vec![down.as_str()], + ExpertSourceLayout::PerExpertFused { down, .. } + | ExpertSourceLayout::PerExpertSeparate { down, .. } => { + down.iter().map(String::as_str).collect() + } + } +} + +fn resolve_collective( + spec: &ExpertGroupSpec, + manifest: &[WeightEntry], +) -> Result<(Option, Option), ExpertPlanError> { + let expected = match spec.parallelism { + ExpertParallelism::Single => None, + ExpertParallelism::TensorParallel => Some(DimKind::Tp), + ExpertParallelism::ExpertParallel => Some(DimKind::Ep), + }; + let Some(expected_kind) = expected else { + return Ok((None, None)); + }; + let rows = collective_schedule(manifest); + let mut selected = None; + for name in down_names(&spec.source_layout) { + let row = rows + .iter() + .find(|row| row.layer == spec.layer.unwrap_or(usize::MAX) && row.name == name) + .ok_or_else(|| { + ExpertPlanError::new(format!( + "expert group '{}' has no ordered G3 collective row for down source '{name}'", + spec.group + )) + })?; + if !matches!( + row.hint, + CollectiveHint::AllReduce { kind } if kind == expected_kind + ) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' down source '{name}' collective {:?} does not match {:?}", + spec.group, row.hint, expected_kind + ))); + } + if let Some(previous) = selected { + if previous != row.hint { + return Err(ExpertPlanError::new(format!( + "expert group '{}' down sources disagree on collective axis", + spec.group + ))); + } + } else { + selected = Some(row.hint); + } + } + Ok((selected, down_names(&spec.source_layout).first().map(|name| (*name).to_string()))) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn mesh_ep() -> DeviceMesh { + DeviceMesh::rect(&[(DimKind::Ep, 2)]).expect("test mesh") + } + + fn manifest() -> Vec { + vec![ + WeightEntry::layer( + "router", + 0, + vec![4, 4], + DType::F32, + ShardPolicy::Replicate, + ), + WeightEntry::layer( + "experts.gate_up", + 0, + vec![4, 128, 64], + DType::MQ4G256, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ), + WeightEntry::layer( + "experts.down", + 0, + vec![4, 64, 64], + DType::MQ4G256, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ), + ] + } + + fn spec(execution: &str, parallelism: ExpertParallelism) -> ExpertGroupSpec { + ExpertGroupSpec { + group: "block-0".into(), + layer: Some(0), + n_experts: 4, + parallelism, + assignment: ExpertAssign::Stride, + source_layout: ExpertSourceLayout::PackedFused { + gate_up: "experts.gate_up".into(), + down: "experts.down".into(), + sidecars: vec![], + }, + resources: ExpertResourceRequirements { + bytes_per_expert: 4096, + alignment: 256, + }, + router: "router".into(), + execution: execution.into(), + } + } + + fn table(shape: Vec) -> GpuTensor { + let mut tensor = GpuTensor::null_for_test(); + tensor.shape = shape; + tensor + } + + #[test] + fn stride_assignment_and_named_group_are_deterministic() { + let plan = ExpertPlan::from_manifest( + &spec("indexed_quantized", ExpertParallelism::ExpertParallel), + &manifest(), + &mesh_ep(), + ) + .unwrap(); + let owners: Vec<_> = plan + .placements() + .iter() + .map(|placement| (placement.global_id, placement.owner, placement.local_slot)) + .collect(); + assert_eq!(owners, vec![(0, 0, 0), (1, 1, 0), (2, 0, 1), (3, 1, 1)]); + assert_eq!(plan.group_devices(), &[0, 1]); + assert_eq!(plan.collective(), Some(CollectiveHint::AllReduce { kind: DimKind::Ep })); + } + + #[test] + fn bound_view_uses_canonical_rank_owner() { + let plan = ExpertPlan::from_manifest( + &spec("indexed_quantized", ExpertParallelism::ExpertParallel), + &manifest(), + &mesh_ep(), + ) + .unwrap(); + let gate = table(vec![4]); + let down = table(vec![4]); + let view = plan.bind_expert_ref(1, &gate, &down, None).unwrap(); + assert_eq!(view.owned(), &[1, 3]); + assert_eq!(view.n_experts(), 4); + } + + #[test] + fn failed_load_rolls_back_and_unload_is_idempotent() { + let mut plan = ExpertPlan::from_manifest( + &spec("indexed_quantized", ExpertParallelism::ExpertParallel), + &manifest(), + &mesh_ep(), + ) + .unwrap(); + let first = plan.placements()[0]; + { + let mut load = plan.begin_load(); + load.reserve(first).unwrap(); + } + assert_eq!(plan.resident_slots(), 0); + { + let mut load = plan.begin_load(); + load.reserve(first).unwrap(); + load.commit(); + } + assert_eq!(plan.resident_slots(), 1); + plan.unload(); + plan.unload(); + assert_eq!(plan.resident_slots(), 0); + } + + #[test] + fn per_expert_fallback_is_refused_at_plan_boundary() { + let error = ExpertPlan::from_manifest( + &spec("per_expert_fallback", ExpertParallelism::ExpertParallel), + &manifest(), + &mesh_ep(), + ) + .expect_err("fallback is not a typed Step protocol"); + assert!(error.to_string().contains("PerExpertFallback")); + } + + #[test] + fn source_shape_mismatch_is_refused_before_owner_creation() { + let mut malformed = manifest(); + malformed[1].logical_shape = vec![4, 64, 64]; + let error = ExpertPlan::from_manifest( + &spec("indexed_quantized", ExpertParallelism::ExpertParallel), + &malformed, + &mesh_ep(), + ) + .expect_err("gate/up and down shapes must agree"); + assert!(error.to_string().contains("shape mismatch")); + } + + #[test] + fn single_plan_emits_identity_collective() { + let mut single_manifest = manifest(); + single_manifest[1].policy = ShardPolicy::Replicate; + single_manifest[2].policy = ShardPolicy::Replicate; + let mesh = DeviceMesh::single().unwrap(); + let plan = ExpertPlan::from_manifest( + &spec("indexed_quantized", ExpertParallelism::Single), + &single_manifest, + &mesh, + ) + .unwrap(); + assert_eq!(plan.collective(), None); + assert_eq!(plan.step_collective(0, 64).unwrap(), StepCollective::None); + } +} From 4f94cd1136401f551a0330ade1d2ce68e7abf515 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 09:54:09 +0200 Subject: [PATCH 17/25] fix: close g5 moe ownership and schedule guards --- crates/hipfire-dispatch/src/families/moe.rs | 146 +++++-- crates/hipfire-dispatch/src/pipeline/mod.rs | 2 +- crates/hipfire-dispatch/src/pipeline/steps.rs | 406 ++++++++++++------ crates/hipfire-runtime/src/moe_plan.rs | 30 +- 4 files changed, 426 insertions(+), 158 deletions(-) diff --git a/crates/hipfire-dispatch/src/families/moe.rs b/crates/hipfire-dispatch/src/families/moe.rs index 6e2462aff3..6d0bd9ec80 100644 --- a/crates/hipfire-dispatch/src/families/moe.rs +++ b/crates/hipfire-dispatch/src/families/moe.rs @@ -17,8 +17,8 @@ //! end-to-end from the call site through every inner GEMV. Scratch stays model-owned. //! Grouped-GEMM prefill is a future arm (gated on `ShapeInfo.batch_size`). -use rdna_compute::DType; -use rdna_compute::{Gpu, GpuTensor}; +use hipfire_hardware::MeshEpoch; +use rdna_compute::{DType, Gpu, GpuTensor}; use crate::context::DispatchCtx; use crate::families::gemv::{GivensRef, WeightRef}; @@ -78,7 +78,6 @@ impl<'a> RouterPlan<'a> { } } - pub fn route_buffers(&self) -> (&'a GpuTensor, &'a GpuTensor) { match self { Self::SoftmaxTopK { @@ -115,12 +114,36 @@ impl<'a> RouterPlan<'a> { self.k_top() ))); } + if !self.normalizes() { + return Err(DispatchError::Hip( + "generic MoE route requires normalized top-k weights".into(), + )); + } let (indices, weights) = self.route_buffers(); if indices.dtype != DType::F32 || weights.dtype != DType::F32 { return Err(DispatchError::Hip( "MoE route indices and weights must use F32 storage".into(), )); } + if indices.shape != weights.shape { + return Err(DispatchError::Hip( + "MoE route index/weight shapes must match".into(), + )); + } + let route_shape_ok = match indices.shape.as_slice() { + [k] => batch_size == 1 && *k == self.k_top(), + [batch, k] => *batch == batch_size && *k == self.k_top(), + _ => false, + }; + if !route_shape_ok { + return Err(DispatchError::Hip(format!( + "MoE route index/weight shape must be [k_top] for batch=1 or [batch,k_top], \ + got indices={:?}, weights={:?}, batch={batch_size}, k_top={}", + indices.shape, + weights.shape, + self.k_top() + ))); + } let expected_slots = batch_size .checked_mul(self.k_top()) .ok_or_else(|| DispatchError::Hip("MoE route slot count overflow".into()))?; @@ -129,22 +152,14 @@ impl<'a> RouterPlan<'a> { "MoE route buffers have insufficient capacity for {expected_slots} slots" ))); } - if !matches!(indices.shape.as_slice(), [_] | [_, _]) { - return Err(DispatchError::Hip( - "MoE route index shape must be [K] or [B,K]".into(), - )); - } - if indices.shape.last() != Some(&self.k_top()) - || weights.shape.last() != Some(&self.k_top()) - { - return Err(DispatchError::Hip( - "MoE route index/weight shape must end in k_top".into(), - )); - } if let Self::SoftmaxTopK { scores, .. } = self { + let score_shape_ok = match scores.shape.as_slice() { + [experts] => batch_size == 1 && *experts == n_experts, + [batch, experts] => *batch == batch_size && *experts == n_experts, + _ => false, + }; if scores.dtype != DType::F32 - || scores.shape.last() != Some(&n_experts) - || !matches!(scores.shape.len(), 1 | 2) + || !score_shape_ok || scores.numel() < batch_size.saturating_mul(n_experts) { return Err(DispatchError::Hip(format!( @@ -200,6 +215,9 @@ pub struct MoeExpertRef<'a> { expert_k: usize, owned: &'a [usize], collective_kind: Option, + owner_rank: usize, + group_devices: &'a [usize], + mesh_epoch: MeshEpoch, } impl<'a> MoeExpertRef<'a> { /// Bind a pointer-table view to metadata already resolved by an owner. @@ -209,6 +227,7 @@ impl<'a> MoeExpertRef<'a> { /// canonical values from its sealed plan; no family receives an allocator /// or a `WeightStore` representation. #[doc(hidden)] + #[allow(clippy::too_many_arguments)] pub fn from_resolved( gate_up_ptrs: &'a GpuTensor, down_ptrs: &'a GpuTensor, @@ -219,6 +238,9 @@ impl<'a> MoeExpertRef<'a> { expert_k: usize, owned: &'a [usize], collective_kind: Option, + owner_rank: usize, + group_devices: &'a [usize], + mesh_epoch: MeshEpoch, ) -> Self { Self { gate_up_ptrs, @@ -230,6 +252,9 @@ impl<'a> MoeExpertRef<'a> { expert_k, owned, collective_kind, + owner_rank, + group_devices, + mesh_epoch, } } @@ -269,6 +294,18 @@ impl<'a> MoeExpertRef<'a> { self.collective_kind } + pub fn owner_rank(&self) -> usize { + self.owner_rank + } + + pub fn group_devices(&self) -> &'a [usize] { + self.group_devices + } + + pub fn mesh_epoch(&self) -> MeshEpoch { + self.mesh_epoch + } + /// Validate the dimensions shared by the fused gate/up and down kernels. /// Gate/up is `[2*expert_m, expert_k]`; down is `[expert_k, expert_m]`. pub fn validate(&self) -> Result<(), DispatchError> { @@ -282,21 +319,65 @@ impl<'a> MoeExpertRef<'a> { "MoeExpertRef: expert dimensions must be nonzero".into(), )); } - if self.gate_up_ptrs.numel() < self.n_experts - || self.down_ptrs.numel() < self.n_experts + let pointer_slots = self + .n_experts + .checked_mul(2) + .ok_or_else(|| DispatchError::Hip("MoeExpertRef: pointer-table size overflows".into()))?; + if self.gate_up_ptrs.dtype != DType::F32 + || self.down_ptrs.dtype != DType::F32 + || self.gate_up_ptrs.shape.as_slice() != [pointer_slots] + || self.down_ptrs.shape.as_slice() != [pointer_slots] { return Err(DispatchError::Hip(format!( - "MoeExpertRef: pointer-table capacity is too small for {} experts", + "MoeExpertRef: pointer tables must be F32 [2*{}]", self.n_experts ))); } if let Some(dummy) = self.dummy_gate_up { - if dummy.numel() == 0 { + if dummy.dtype != DType::F32 || dummy.numel() == 0 { return Err(DispatchError::Hip( - "MoeExpertRef: dummy gate/up table is empty".into(), + "MoeExpertRef: dummy gate/up table is invalid".into(), )); } } + if self.group_devices.is_empty() || self.owner_rank >= self.group_devices.len() { + return Err(DispatchError::Hip( + "MoeExpertRef: owner rank is outside its mesh group".into(), + )); + } + if self.group_devices.iter().enumerate().any(|(index, device)| { + self.group_devices[..index].contains(device) + }) { + return Err(DispatchError::Hip( + "MoeExpertRef: mesh group contains duplicate devices".into(), + )); + } + if self.collective_kind.is_none() + && (self.owner_rank != 0 + || self.group_devices.len() != 1 + || self.group_devices.first().copied() != Some(0) + || self.owned.len() != self.n_experts + || !self + .owned + .iter() + .copied() + .enumerate() + .all(|(expert, global_id)| expert == global_id)) + { + return Err(DispatchError::Hip( + "MoeExpertRef: single-device owner view is not canonical".into(), + )); + } + if self.collective_kind.is_some() && self.group_devices.len() < 2 { + return Err(DispatchError::Hip( + "MoeExpertRef: parallel owner view requires at least two ranks".into(), + )); + } + if self.owned.is_empty() { + return Err(DispatchError::Hip( + "MoeExpertRef: owner view has no experts".into(), + )); + } let mut previous = None; for &expert in self.owned { if expert >= self.n_experts { @@ -305,9 +386,9 @@ impl<'a> MoeExpertRef<'a> { self.n_experts ))); } - if previous == Some(expert) { + if previous.is_some_and(|previous| expert <= previous) { return Err(DispatchError::Hip(format!( - "MoeExpertRef: duplicate owned expert {expert}" + "MoeExpertRef: owned experts are not strictly ordered at {expert}" ))); } previous = Some(expert); @@ -1540,10 +1621,9 @@ impl MoeFamily { gpus: &mut hipfire_hardware::Gpus, mesh: &hipfire_hardware::DeviceMesh, ctx: &DispatchCtx, - rank_steps: &[&[crate::pipeline::steps::Step<'a>]], - collectives: &[crate::pipeline::steps::StepCollective], + schedules: &[&crate::pipeline::steps::SealedMoeSchedule<'a>], ) -> Result<(), DispatchError> { - crate::pipeline::steps::execute_steps_mesh(gpus, mesh, ctx, rank_steps, collectives) + crate::pipeline::steps::execute_sealed_steps_mesh(gpus, mesh, ctx, schedules) } /// Resolve the best kernel key for the given MoE variant. @@ -1848,9 +1928,11 @@ mod tests { #[test] fn expert_ref_rejects_shape_and_owner_mismatch() { let mut gate_ptrs = GpuTensor::null_for_test(); - gate_ptrs.shape = vec![4]; + gate_ptrs.shape = vec![8]; let mut down_ptrs = GpuTensor::null_for_test(); - down_ptrs.shape = vec![4]; + down_ptrs.shape = vec![8]; + let mesh = hipfire_hardware::DeviceMesh::rect(&[(hipfire_hardware::DimKind::Ep, 2)]) + .unwrap(); let experts = MoeExpertRef::from_resolved( &gate_ptrs, &down_ptrs, @@ -1861,6 +1943,9 @@ mod tests { 128, &[0, 2], Some(hipfire_hardware::DimKind::Ep), + 0, + &[0, 1], + mesh.epoch(), ); experts .validate_projection_shapes(&[128, 128], &[128, 64]) @@ -1878,6 +1963,9 @@ mod tests { 128, &[4], Some(hipfire_hardware::DimKind::Ep), + 0, + &[0, 1], + mesh.epoch(), ); assert!(unknown.validate().is_err()); } diff --git a/crates/hipfire-dispatch/src/pipeline/mod.rs b/crates/hipfire-dispatch/src/pipeline/mod.rs index aff5b236ee..5ce988b383 100644 --- a/crates/hipfire-dispatch/src/pipeline/mod.rs +++ b/crates/hipfire-dispatch/src/pipeline/mod.rs @@ -11,7 +11,7 @@ use rdna_compute::{DType, Gpu, GpuTensor}; use std::sync::{LazyLock, OnceLock}; pub use steps::{ - execute_steps, execute_steps_mesh, FusedPattern, GemvInput, SealedMoeSchedule, Step, + execute_sealed_steps_mesh, execute_steps, FusedPattern, GemvInput, SealedMoeSchedule, Step, StepCollective, }; diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index a524d810a3..78fba4ee53 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -285,9 +285,15 @@ pub fn validate_moe_step_schedule( inverse_perm, }, ] => { - if *batch_size != 1 || *down_batch != 1 || inverse_perm.is_some() { + if *batch_size != 1 + || *down_batch != 1 + || inverse_perm.is_some() + || *k_top != plan.k_top() + || *inter != (*experts).expert_m() + || *rows != batch_size.saturating_mul(*k_top) + { return Err(DispatchError::Hip( - "indexed MoE grammar requires decode batch=1 and no inverse permutation" + "indexed MoE grammar requires decode batch=1, normalized route, and no inverse permutation" .into(), )); } @@ -389,12 +395,14 @@ pub fn validate_moe_step_schedule( inverse_perm: combine_inverse, }, ] => { - if *k_top != *unscatter_k + if *k_top != plan.k_top() + || *k_top != *unscatter_k || *k_top != *down_k || *k_top != *combine_k || *batch_size != *down_batch || *batch_size != *combine_batch || *total_slots != batch_size.saturating_mul(*k_top) + || *n_experts != experts.n_experts() || !same_tensor(topk_indices, route.route_buffers().0) || !same_tensor(topk_indices, gate_sorted) || !same_tensor(topk_indices, down_sorted) @@ -474,8 +482,10 @@ pub fn validate_moe_step_schedule( "MoE route metadata is not bound to the concrete expert Steps".into(), )); } - if hidden == 0 { - return Err(DispatchError::Hip("MoE combine hidden size is zero".into())); + if hidden == 0 || hidden != experts.expert_k() { + return Err(DispatchError::Hip( + "MoE combine hidden size does not match the expert down projection".into(), + )); } let shape_gate = [2 * experts.expert_m(), experts.expert_k()]; let shape_down = [experts.expert_k(), experts.expert_m()]; @@ -572,37 +582,26 @@ pub fn execute_sealed_steps( schedule: &SealedMoeSchedule<'_>, ) -> Result<(), DispatchError> { validate_moe_protocol_schedule(&schedule.steps, &schedule.collectives, schedule.execution)?; + if collective_count(&schedule.collectives) != 0 { + return Err(DispatchError::Hip( + "parallel MoE schedules require execute_sealed_steps_mesh".into(), + )); + } execute_steps_inner(gpu, ctx, &schedule.steps) } -/// Execute one validated MoE schedule per participating device and perform -/// its single manifest-owned routed reduction. The `rank_steps` order is the -/// exact order of `StepCollective::AllReduce::group`; no flat device ordering -/// is inferred. Single-device plans use `execute_steps` and have no reduction. -pub fn execute_steps_mesh<'a>( +/// Execute one sealed MoE schedule per participating device and perform its +/// single manifest-owned routed reduction. Schedules are ordered by the +/// collective's rank field; every schedule is validated before any GPU work. +pub fn execute_sealed_steps_mesh<'a>( gpus: &mut Gpus, mesh: &DeviceMesh, ctx: &DispatchCtx, - rank_steps: &[&[Step<'a>]], - collectives: &[StepCollective], + schedules: &[&SealedMoeSchedule<'a>], ) -> Result<(), DispatchError> { - let reduction = collectives - .iter() - .find_map(|collective| match collective { - StepCollective::AllReduce { - kind, - dim, - group, - mesh: epoch, - rank, - } => Some((*kind, *dim, group, *epoch, *rank)), - StepCollective::None => None, - }) - .ok_or_else(|| DispatchError::Hip("parallel MoE schedule has no collective".into()))?; - let (kind, dim, group, epoch, rank) = reduction; - if mesh.epoch() != epoch { + if schedules.is_empty() { return Err(DispatchError::Hip( - "MoE collective belongs to a different mesh generation".into(), + "parallel MoE execution has no rank schedules".into(), )); } if mesh.n_devices() != gpus.devices.len() { @@ -612,7 +611,89 @@ pub fn execute_steps_mesh<'a>( gpus.devices.len() ))); } - if group.len() < 2 || rank >= group.len() || rank_steps.len() != group.len() { + + let mut reduction: Option<(DimKind, usize, Vec, MeshEpoch)> = None; + let mut outputs = Vec::with_capacity(schedules.len()); + for (rank, schedule) in schedules.iter().enumerate() { + validate_moe_protocol_schedule( + schedule.steps(), + schedule.collectives(), + schedule.execution(), + )?; + let (kind, dim, group, epoch, descriptor_rank) = schedule + .collectives() + .iter() + .find_map(|collective| match collective { + StepCollective::AllReduce { + kind, + dim, + group, + mesh, + rank, + } => Some((*kind, *dim, group.clone(), *mesh, *rank)), + StepCollective::None => None, + }) + .ok_or_else(|| { + DispatchError::Hip( + "parallel MoE schedule has no manifest-owned collective".into(), + ) + })?; + if descriptor_rank != rank { + return Err(DispatchError::Hip( + "MoE collective rank does not match schedule order".into(), + )); + } + let experts = schedule + .steps() + .iter() + .find_map(|step| match step { + Step::IndexedMoeGemv { experts, .. } + | Step::GroupedMoeGemm { experts, .. } => Some(*experts), + _ => None, + }) + .ok_or_else(|| DispatchError::Hip("MoE schedule has no expert owner view".into()))?; + if experts.owner_rank() != rank + || experts.group_devices() != group.as_slice() + || experts.mesh_epoch() != epoch + || experts.collective_kind() != Some(kind) + { + return Err(DispatchError::Hip( + "MoE schedule expert owner does not match its collective identity".into(), + )); + } + if let Some((expected_kind, expected_dim, expected_group, expected_epoch)) = &reduction { + if *expected_kind != kind + || *expected_dim != dim + || expected_group != &group + || *expected_epoch != epoch + { + return Err(DispatchError::Hip( + "MoE rank schedules disagree on collective mesh identity".into(), + )); + } + } else { + reduction = Some((kind, dim, group, epoch)); + } + let output = schedule + .steps() + .iter() + .find_map(|step| match step { + Step::MoeCombine { out, .. } => Some(&out.buf), + _ => None, + }) + .ok_or_else(|| { + DispatchError::Hip("parallel MoE schedule has no combine output".into()) + })?; + outputs.push(output); + } + + let (kind, dim, group, epoch) = reduction.expect("validated schedules have a collective"); + if mesh.epoch() != epoch { + return Err(DispatchError::Hip( + "MoE collective belongs to a different mesh generation".into(), + )); + } + if group.len() != schedules.len() || group.len() < 2 { return Err(DispatchError::Hip( "MoE collective rank count does not match mesh group".into(), )); @@ -624,42 +705,23 @@ pub fn execute_steps_mesh<'a>( "MoE collective group contains invalid or duplicate devices".into(), )); } - let group_is_mesh_group = group.iter().any(|device| { - let expected = mesh.group_along(kind, &mesh.coord_of(*device)); - expected == *group - }); - if !group_is_mesh_group { + if mesh.group_along(kind, &mesh.coord_of(group[0])) != group { return Err(DispatchError::Hip( "MoE collective group is not a named DeviceMesh axis group".into(), )); } - for steps in rank_steps { - validate_moe_parallel_schedule(steps, collectives)?; - } - for (index, &device) in group.iter().enumerate() { + for &device in &group { if device >= gpus.devices.len() { return Err(DispatchError::Hip(format!( "MoE collective device {device} is outside the Gpus owner" ))); } - execute_steps_inner(&mut gpus.devices[device], ctx, rank_steps[index])?; } - let buffers: Vec<&hip_bridge::DeviceBuffer> = rank_steps - .iter() - .map(|steps| { - steps - .iter() - .find_map(|step| match step { - Step::MoeCombine { out, .. } => Some(&out.buf), - _ => None, - }) - .ok_or_else(|| { - DispatchError::Hip("parallel MoE schedule has no combine output".into()) - }) - }) - .collect::>()?; - gpus.all_reduce_sum_f32_peer(group, &buffers, dim) + for (rank, &device) in group.iter().enumerate() { + execute_steps_inner(&mut gpus.devices[device], ctx, schedules[rank].steps())?; + } + gpus.all_reduce_sum_f32_peer(&group, &outputs, dim) .map_err(|error| DispatchError::Hip(error.to_string())) } @@ -686,46 +748,59 @@ fn require_tensor( Ok(()) } +fn require_raw_i32( + tensor: &GpuTensor, + name: &str, + elements: usize, +) -> Result<(), DispatchError> { + let bytes = elements + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| DispatchError::Hip(format!("MoE {name} capacity overflows")))?; + if tensor.dtype != DType::Raw || tensor.numel() < bytes { + return Err(DispatchError::Hip(format!( + "MoE {name} has dtype {:?}/capacity {}, expected Raw/at least {} bytes", + tensor.dtype, + tensor.numel(), + bytes + ))); + } + Ok(()) +} + fn validate_step_tensors( steps: &[Step], experts: &MoeExpertRef<'_>, batch_size: usize, hidden: usize, ) -> Result<(), DispatchError> { + let k_top = steps + .iter() + .find_map(|step| match step { + Step::IndexedMoeGemv { k_top, .. } + | Step::GroupedMoeGemm { k_top, .. } + | Step::MoeGateUpUnscatter { k_top, .. } + | Step::MoeCombine { k_top, .. } => Some(*k_top), + _ => None, + }) + .ok_or_else(|| DispatchError::Hip("MoE grammar has no route width".into()))?; let slots = batch_size - .checked_mul( - steps - .iter() - .find_map(|step| match step { - Step::IndexedMoeGemv { k_top, .. } - | Step::GroupedMoeGemm { k_top, .. } - | Step::MoeGateUpUnscatter { k_top, .. } - | Step::MoeCombine { k_top, .. } => Some(*k_top), - _ => None, - }) - .unwrap_or(0), - ) + .checked_mul(k_top) .ok_or_else(|| DispatchError::Hip("MoE slot capacity overflow".into()))?; + let route = steps + .iter() + .find_map(|step| match step { + Step::MoeRoute { plan } => Some(plan), + _ => None, + }) + .ok_or_else(|| DispatchError::Hip("MoE grammar has no route".into()))?; require_tensor( - steps - .iter() - .find_map(|step| match step { - Step::MoeRoute { plan } => Some(plan.route_buffers().0), - _ => None, - }) - .expect("validated grammar has a route"), + route.route_buffers().0, "route indices", DType::F32, slots, )?; require_tensor( - steps - .iter() - .find_map(|step| match step { - Step::MoeRoute { plan } => Some(plan.route_buffers().1), - _ => None, - }) - .expect("validated grammar has a route"), + route.route_buffers().1, "route weights", DType::F32, slots, @@ -752,7 +827,16 @@ fn validate_step_tensors( )) } }; - require_tensor(x, "indexed gate/up input", DType::F32, batch_size * experts.expert_k())?; + require_tensor( + x, + "indexed gate/up input", + DType::F32, + batch_size + .checked_mul(experts.expert_k()) + .ok_or_else(|| { + DispatchError::Hip("MoE indexed input capacity overflow".into()) + })?, + )?; require_tensor(out, "indexed gate output", DType::F32, gate_capacity)?; require_tensor(up_out, "indexed up output", DType::F32, gate_capacity)?; if same_tensor(out, up_out) || same_tensor(x, out) || same_tensor(x, up_out) { @@ -783,15 +867,41 @@ fn validate_step_tensors( )); } } - Step::MoeCombine { out, hidden, .. } => { + Step::MoeCombine { + down_out, + topk_weights, + out, + hidden, + .. + } => { + require_tensor( + down_out, + "combine down input", + DType::F32, + batch_size + .checked_mul(k_top) + .and_then(|slots| slots.checked_mul(*hidden)) + .ok_or_else(|| DispatchError::Hip("MoE combine input capacity overflows"))?, + )?; + require_tensor( + topk_weights, + "combine weights", + DType::F32, + slots, + )?; require_tensor( out, "combine output", DType::F32, batch_size .checked_mul(*hidden) - .ok_or_else(|| DispatchError::Hip("MoE output capacity overflow".into()))?, + .ok_or_else(|| DispatchError::Hip("MoE output capacity overflow"))?, )?; + if same_tensor(down_out, out) || same_tensor(topk_weights, out) { + return Err(DispatchError::Hip( + "MoE combine input/output buffers must not alias".into(), + )); + } } Step::MoeScatter { expert_token_counts, @@ -805,51 +915,89 @@ fn validate_step_tensors( block_m, .. } => { - require_tensor(expert_token_counts, "expert counts", DType::F32, *n_experts)?; - require_tensor( - expert_offsets, - "expert offsets", - DType::F32, - n_experts + 1, - )?; - require_tensor(sorted_slot_index, "sorted slots", DType::F32, *m_total_max)?; - require_tensor(inverse_perm, "inverse permutation", DType::F32, *total_slots)?; - require_tensor( + require_raw_i32(expert_token_counts, "expert counts", *n_experts)?; + require_raw_i32(expert_offsets, "expert offsets", *n_experts + 1)?; + require_raw_i32(sorted_slot_index, "sorted slots", *m_total_max)?; + require_raw_i32(inverse_perm, "inverse permutation", *total_slots)?; + require_raw_i32( expert_tile_ids, "expert tile ids", - DType::F32, - m_total_max / block_m, + *m_total_max / *block_m, )?; } - Step::GroupedMoeGemm { x, y, m_total, .. } => { - require_tensor(x, "grouped input", DType::F32, x.numel())?; + Step::GroupedMoeGemm { + which, + x, + y, + m_total, + .. + } => { + let input_capacity = match which { + MoeProj::GateUp { .. } => batch_size + .checked_mul(experts.expert_k()) + .ok_or_else(|| { + DispatchError::Hip("MoE grouped input capacity overflows".into()) + })?, + MoeProj::DownExpanded => slots + .checked_mul(experts.expert_m()) + .ok_or_else(|| { + DispatchError::Hip("MoE grouped input capacity overflows".into()) + })?, + }; + require_tensor(x, "grouped input", DType::F32, input_capacity)?; + let output_width = match which { + MoeProj::GateUp { .. } => 2usize + .checked_mul(experts.expert_m()) + .ok_or_else(|| { + DispatchError::Hip("MoE grouped output width overflows".into()) + })?, + MoeProj::DownExpanded => experts.expert_k(), + }; require_tensor( y, "grouped output", DType::F32, - m_total - .checked_mul(if matches!( - step, - Step::GroupedMoeGemm { - which: MoeProj::GateUp { .. }, - .. - } - ) { - 2 * experts.expert_m() - } else { - experts.expert_k() - }) - .ok_or_else(|| DispatchError::Hip("grouped output capacity overflow".into()))?, + m_total.checked_mul(output_width).ok_or_else(|| { + DispatchError::Hip("grouped output capacity overflow".into()) + })?, )?; } Step::MoeGateUpUnscatter { + y_grouped, gate_batch, up_batch, inter, + m_total, .. } => { - require_tensor(gate_batch, "unscatter gate", DType::F32, slots * inter)?; - require_tensor(up_batch, "unscatter up", DType::F32, slots * inter)?; + require_tensor( + y_grouped, + "unscatter grouped input", + DType::F32, + m_total + .checked_mul( + 2usize.checked_mul(*inter).ok_or_else(|| { + DispatchError::Hip("MoE unscatter width overflows".into()) + })?, + ) + .ok_or_else(|| DispatchError::Hip("MoE unscatter input overflows"))?, + )?; + require_tensor( + gate_batch, + "unscatter gate", + DType::F32, + slots.checked_mul(*inter).ok_or_else(|| { + DispatchError::Hip("MoE unscatter gate capacity overflows".into()) + })?, + )?; + require_tensor( + up_batch, + "unscatter up", + DType::F32, + slots.checked_mul(*inter).ok_or_else(|| { + DispatchError::Hip("MoE unscatter up capacity overflows".into()) + })?, + )?; } Step::MoeActivation { gate, @@ -859,13 +1007,16 @@ fn validate_step_tensors( rows, .. } => { - let capacity = rows - .checked_mul(*inter) - .ok_or_else(|| DispatchError::Hip("MoE activation capacity overflow".into()))?; + let capacity = rows.checked_mul(*inter).ok_or_else(|| { + DispatchError::Hip("MoE activation capacity overflow".into()) + })?; require_tensor(gate, "activation gate", DType::F32, capacity)?; require_tensor(up, "activation up", DType::F32, capacity)?; require_tensor(rot_out, "activation output", DType::F32, capacity)?; - if same_tensor(gate, up) || same_tensor(gate, rot_out) || same_tensor(up, rot_out) { + if same_tensor(gate, up) + || same_tensor(gate, rot_out) + || same_tensor(up, rot_out) + { return Err(DispatchError::Hip( "MoE activation buffers must not alias".into(), )); @@ -906,7 +1057,7 @@ fn validate_collectives( "MoE collective axis {kind:?} does not match owner axis {expected_kind:?}" ))); } - if *dim != hidden || group.is_empty() || *rank >= group.len() { + if *dim != hidden || group.len() < 2 || *rank >= group.len() { return Err(DispatchError::Hip( "MoE collective rank/group/output dimension is invalid".into(), )); @@ -924,6 +1075,11 @@ fn validate_collectives( )); } } + if expected_kind.is_some() != reduction.is_some() { + return Err(DispatchError::Hip( + "MoE collective count does not match the resolved owner parallelism".into(), + )); + } Ok(()) } @@ -2609,8 +2765,8 @@ mod tests { let rot = tensor(vec![8 * 64]); let down = tensor(vec![8 * 128]); let out = tensor(vec![128]); - let gate_ptrs = tensor(vec![4]); - let down_ptrs = tensor(vec![4]); + let gate_ptrs = tensor(vec![16]); + let down_ptrs = tensor(vec![16]); let experts = MoeExpertRef::from_resolved( &gate_ptrs, &down_ptrs, @@ -2619,8 +2775,11 @@ mod tests { 8, 64, 128, + &[0, 1, 2, 3, 4, 5, 6, 7], + None, + 0, &[0], - Some(DimKind::Ep), + DeviceMesh::single().unwrap().epoch(), ); let steps = indexed_steps( &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, @@ -2665,8 +2824,8 @@ mod tests { let rot = tensor(vec![8 * 64]); let down = tensor(vec![8 * 128]); let out = tensor(vec![128]); - let gate_ptrs = tensor(vec![8]); - let down_ptrs = tensor(vec![8]); + let gate_ptrs = tensor(vec![16]); + let down_ptrs = tensor(vec![16]); let experts = MoeExpertRef::from_resolved( &gate_ptrs, &down_ptrs, @@ -2677,6 +2836,9 @@ mod tests { 128, &[0, 1], Some(DimKind::Ep), + 0, + &[0, 1], + DeviceMesh::rect(&[(DimKind::Ep, 2)]).unwrap().epoch(), ); let steps = indexed_steps( &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs index e97aa77169..9de469a355 100644 --- a/crates/hipfire-runtime/src/moe_plan.rs +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -313,6 +313,9 @@ impl ExpertPlan { self.shape.expert_k, owned, collective_kind, + rank, + &self.group_devices, + self.mesh_epoch, ); view.validate() .map_err(|error| ExpertPlanError::new(error.to_string()))?; @@ -457,11 +460,26 @@ fn validate_spec(spec: &ExpertGroupSpec, mesh: &DeviceMesh) -> Result<(), Expert spec.group ))); } - let group_size = match spec.parallelism { - ExpertParallelism::Single => 1, - ExpertParallelism::TensorParallel => mesh.size_of(DimKind::Tp), - ExpertParallelism::ExpertParallel => mesh.size_of(DimKind::Ep), + let required_axis = match spec.parallelism { + ExpertParallelism::Single => None, + ExpertParallelism::TensorParallel => Some(DimKind::Tp), + ExpertParallelism::ExpertParallel => Some(DimKind::Ep), }; + let group_size = required_axis.map_or(1, |kind| mesh.size_of(kind)); + if let Some(kind) = required_axis { + if !mesh.axes().iter().any(|axis| axis.kind == kind) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' requires a named {:?} mesh axis", + spec.group, kind + ))); + } + if group_size < 2 { + return Err(ExpertPlanError::new(format!( + "expert group '{}' parallel mesh group must have at least two ranks", + spec.group + ))); + } + } if group_size == 0 { return Err(ExpertPlanError::new("expert group resolved to zero ranks")); } @@ -1002,8 +1020,8 @@ mod tests { &mesh_ep(), ) .unwrap(); - let gate = table(vec![4]); - let down = table(vec![4]); + let gate = table(vec![8]); + let down = table(vec![8]); let view = plan.bind_expert_ref(1, &gate, &down, None).unwrap(); assert_eq!(view.owned(), &[1, 3]); assert_eq!(view.n_experts(), 4); From 603447e396af6627ef62b93ba77139da3233ffbe Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 10:23:03 +0200 Subject: [PATCH 18/25] fix: close g5 moe review gaps --- Cargo.lock | 1 + crates/hipfire-dispatch/src/families/moe.rs | 216 +++++++++++++++--- crates/hipfire-dispatch/src/pipeline/mod.rs | 2 + crates/hipfire-dispatch/src/pipeline/steps.rs | 205 ++++++++++++++--- crates/hipfire-runtime/src/moe_plan.rs | 214 +++++++++++++++-- 5 files changed, 553 insertions(+), 85 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9692607f07..91a2bf9d26 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1386,6 +1386,7 @@ version = "0.3.0" dependencies = [ "hip-bridge", "hipfire-config", + "hipfire-hardware", "rdna-compute", ] diff --git a/crates/hipfire-dispatch/src/families/moe.rs b/crates/hipfire-dispatch/src/families/moe.rs index 6d0bd9ec80..fd09093210 100644 --- a/crates/hipfire-dispatch/src/families/moe.rs +++ b/crates/hipfire-dispatch/src/families/moe.rs @@ -57,6 +57,14 @@ pub enum RouterPlan<'a> { }, } +fn checked_numel(tensor: &GpuTensor, name: &str) -> Result { + tensor + .shape + .iter() + .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) + .ok_or_else(|| DispatchError::Hip(format!("MoE {name} logical shape overflows"))) +} + impl<'a> RouterPlan<'a> { pub fn selection(&self) -> RouterSelection { match self { @@ -147,9 +155,18 @@ impl<'a> RouterPlan<'a> { let expected_slots = batch_size .checked_mul(self.k_top()) .ok_or_else(|| DispatchError::Hip("MoE route slot count overflow".into()))?; - if indices.numel() < expected_slots || weights.numel() < expected_slots { + let route_bytes = expected_slots + .checked_mul(DType::F32.size()) + .ok_or_else(|| DispatchError::Hip("MoE route byte capacity overflow".into()))?; + let index_elements = checked_numel(indices, "route indices")?; + let weight_elements = checked_numel(weights, "route weights")?; + if index_elements < expected_slots + || weight_elements < expected_slots + || indices.buf.size() < route_bytes + || weights.buf.size() < route_bytes + { return Err(DispatchError::Hip(format!( - "MoE route buffers have insufficient capacity for {expected_slots} slots" + "MoE route buffers have insufficient logical/physical capacity for {expected_slots} slots" ))); } if let Self::SoftmaxTopK { scores, .. } = self { @@ -158,12 +175,20 @@ impl<'a> RouterPlan<'a> { [batch, experts] => *batch == batch_size && *experts == n_experts, _ => false, }; + let score_elements = checked_numel(scores, "router scores")?; + let score_capacity = batch_size + .checked_mul(n_experts) + .ok_or_else(|| DispatchError::Hip("MoE router score capacity overflow".into()))?; + let score_bytes = score_capacity + .checked_mul(DType::F32.size()) + .ok_or_else(|| DispatchError::Hip("MoE router score byte capacity overflow".into()))?; if scores.dtype != DType::F32 || !score_shape_ok - || scores.numel() < batch_size.saturating_mul(n_experts) + || score_elements < score_capacity + || scores.buf.size() < score_bytes { return Err(DispatchError::Hip(format!( - "MoE router score shape/dtype does not match batch={batch_size}, experts={n_experts}" + "MoE router score shape/dtype/capacity does not match batch={batch_size}, experts={n_experts}" ))); } } @@ -199,13 +224,12 @@ impl ExpertExecutionPlan { } } -/// Borrowed view over one resolver-owned rank-local expert table. +/// Opaque plan-produced data used to create an expert view. /// -/// The view contains no allocation handle, source path, or storage owner. It -/// is created by the manifest/runtime owner and borrowed by a Step until the -/// owner is dropped. Keeping the fields private prevents a family from -/// replacing the canonical placement metadata after binding. -pub struct MoeExpertRef<'a> { +/// The binding is intentionally separate from [`MoeExpertRef`]: callers can +/// retain the view for scheduling, but there is no public constructor that +/// accepts arbitrary canonical metadata directly on the executable reference. +pub struct MoeExpertRefBinding<'a> { gate_up_ptrs: &'a GpuTensor, down_ptrs: &'a GpuTensor, dummy_gate_up: Option<&'a GpuTensor>, @@ -214,21 +238,20 @@ pub struct MoeExpertRef<'a> { expert_m: usize, expert_k: usize, owned: &'a [usize], + ownership_partition: &'a [(usize, usize, usize)], + router_identity: &'a str, collective_kind: Option, owner_rank: usize, group_devices: &'a [usize], mesh_epoch: MeshEpoch, } -impl<'a> MoeExpertRef<'a> { - /// Bind a pointer-table view to metadata already resolved by an owner. - /// - /// This constructor deliberately accepts borrowed tables only. Runtime - /// owners should expose a narrower `bind_*` method that supplies the - /// canonical values from its sealed plan; no family receives an allocator - /// or a `WeightStore` representation. + +impl<'a> MoeExpertRefBinding<'a> { + /// Runtime-plan bridge. The returned binding is opaque to callers and can + /// only be consumed by [`MoeExpertRef::from_binding`]. #[doc(hidden)] #[allow(clippy::too_many_arguments)] - pub fn from_resolved( + pub fn from_plan( gate_up_ptrs: &'a GpuTensor, down_ptrs: &'a GpuTensor, dummy_gate_up: Option<&'a GpuTensor>, @@ -237,6 +260,8 @@ impl<'a> MoeExpertRef<'a> { expert_m: usize, expert_k: usize, owned: &'a [usize], + ownership_partition: &'a [(usize, usize, usize)], + router_identity: &'a str, collective_kind: Option, owner_rank: usize, group_devices: &'a [usize], @@ -251,12 +276,58 @@ impl<'a> MoeExpertRef<'a> { expert_m, expert_k, owned, + ownership_partition, + router_identity, collective_kind, owner_rank, group_devices, mesh_epoch, } } +} + +/// Borrowed view over one resolver-owned rank-local expert table. +/// +/// The view contains no allocation handle, source path, or storage owner. It +/// is created from an opaque plan binding and borrowed by a Step until the +/// owner is dropped. +pub struct MoeExpertRef<'a> { + gate_up_ptrs: &'a GpuTensor, + down_ptrs: &'a GpuTensor, + dummy_gate_up: Option<&'a GpuTensor>, + dtype: DType, + n_experts: usize, + expert_m: usize, + expert_k: usize, + owned: &'a [usize], + ownership_partition: &'a [(usize, usize, usize)], + router_identity: &'a str, + collective_kind: Option, + owner_rank: usize, + group_devices: &'a [usize], + mesh_epoch: MeshEpoch, +} + +impl<'a> MoeExpertRef<'a> { + /// Consume an opaque plan binding into an executable expert view. + pub fn from_binding(binding: MoeExpertRefBinding<'a>) -> Self { + Self { + gate_up_ptrs: binding.gate_up_ptrs, + down_ptrs: binding.down_ptrs, + dummy_gate_up: binding.dummy_gate_up, + dtype: binding.dtype, + n_experts: binding.n_experts, + expert_m: binding.expert_m, + expert_k: binding.expert_k, + owned: binding.owned, + ownership_partition: binding.ownership_partition, + router_identity: binding.router_identity, + collective_kind: binding.collective_kind, + owner_rank: binding.owner_rank, + group_devices: binding.group_devices, + mesh_epoch: binding.mesh_epoch, + } + } pub fn gate_up_ptrs(&self) -> &'a GpuTensor { self.gate_up_ptrs @@ -289,6 +360,14 @@ impl<'a> MoeExpertRef<'a> { pub fn owned(&self) -> &'a [usize] { self.owned } + pub fn ownership_partition(&self) -> &'a [(usize, usize, usize)] { + self.ownership_partition + } + + pub fn router_identity(&self) -> &'a str { + self.router_identity + } + pub fn collective_kind(&self) -> Option { self.collective_kind @@ -319,22 +398,40 @@ impl<'a> MoeExpertRef<'a> { "MoeExpertRef: expert dimensions must be nonzero".into(), )); } + if self.router_identity.is_empty() { + return Err(DispatchError::Hip( + "MoeExpertRef: router identity is empty".into(), + )); + } let pointer_slots = self .n_experts .checked_mul(2) .ok_or_else(|| DispatchError::Hip("MoeExpertRef: pointer-table size overflows".into()))?; + let pointer_bytes = pointer_slots + .checked_mul(DType::F32.size()) + .ok_or_else(|| DispatchError::Hip("MoeExpertRef: pointer-table bytes overflow".into()))?; if self.gate_up_ptrs.dtype != DType::F32 || self.down_ptrs.dtype != DType::F32 || self.gate_up_ptrs.shape.as_slice() != [pointer_slots] || self.down_ptrs.shape.as_slice() != [pointer_slots] + || self.gate_up_ptrs.buf.size() < pointer_bytes + || self.down_ptrs.buf.size() < pointer_bytes { return Err(DispatchError::Hip(format!( - "MoeExpertRef: pointer tables must be F32 [2*{}]", + "MoeExpertRef: pointer tables must be F32 [2*{}] with {pointer_bytes} bytes", self.n_experts ))); } if let Some(dummy) = self.dummy_gate_up { - if dummy.dtype != DType::F32 || dummy.numel() == 0 { + let dummy_elements = dummy + .shape + .iter() + .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) + .ok_or_else(|| DispatchError::Hip("MoeExpertRef: dummy table shape overflows".into()))?; + let dummy_bytes = dummy_elements + .checked_mul(DType::F32.size()) + .ok_or_else(|| DispatchError::Hip("MoeExpertRef: dummy table bytes overflow".into()))?; + if dummy.dtype != DType::F32 || dummy_elements == 0 || dummy.buf.size() < dummy_bytes { return Err(DispatchError::Hip( "MoeExpertRef: dummy gate/up table is invalid".into(), )); @@ -373,6 +470,37 @@ impl<'a> MoeExpertRef<'a> { "MoeExpertRef: parallel owner view requires at least two ranks".into(), )); } + if self.ownership_partition.is_empty() { + return Err(DispatchError::Hip( + "MoeExpertRef: ownership partition is empty".into(), + )); + } + let mut previous_placement = None; + for &(global_id, owner, local_slot) in self.ownership_partition { + if global_id >= self.n_experts || owner >= self.group_devices.len() { + return Err(DispatchError::Hip( + "MoeExpertRef: ownership partition contains an invalid placement".into(), + )); + } + let placement = (global_id, owner, local_slot); + if previous_placement.is_some_and(|previous| placement <= previous) { + return Err(DispatchError::Hip( + "MoeExpertRef: ownership partition is not canonical".into(), + )); + } + previous_placement = Some(placement); + } + if self + .ownership_partition + .iter() + .filter(|(_, owner, _)| *owner == self.owner_rank) + .map(|(global_id, _, _)| *global_id) + .ne(self.owned.iter().copied()) + { + return Err(DispatchError::Hip( + "MoeExpertRef: rank view does not match ownership partition".into(), + )); + } if self.owned.is_empty() { return Err(DispatchError::Hip( "MoeExpertRef: owner view has no experts".into(), @@ -404,7 +532,11 @@ impl<'a> MoeExpertRef<'a> { down_shape: &[usize], ) -> Result<(), DispatchError> { self.validate()?; - let expected_gate_up = [2 * self.expert_m, self.expert_k]; + let gate_m = self + .expert_m + .checked_mul(2) + .ok_or_else(|| DispatchError::Hip("MoeExpertRef: gate/up shape overflows".into()))?; + let expected_gate_up = [gate_m, self.expert_k]; let expected_down = [self.expert_k, self.expert_m]; if gate_up_shape != expected_gate_up || down_shape != expected_down { return Err(DispatchError::Hip(format!( @@ -1904,14 +2036,25 @@ mod tests { assert!(!r.use_gpu_topk); } + fn tensor(shape: Vec) -> GpuTensor { + let elements = shape.iter().product::(); + let mut tensor = GpuTensor::null_for_test(); + tensor.buf = unsafe { + hip_bridge::DeviceBuffer::from_raw( + std::ptr::null_mut(), + elements + .checked_mul(DType::F32.size()) + .expect("test tensor bytes"), + ) + }; + tensor.shape = shape; + tensor + } #[test] fn typed_router_preserves_selection_and_normalization_contract() { - let mut scores = GpuTensor::null_for_test(); - scores.shape = vec![8]; - let mut indices = GpuTensor::null_for_test(); - indices.shape = vec![8]; - let mut weights = GpuTensor::null_for_test(); - weights.shape = vec![8]; + let scores = tensor(vec![8]); + let indices = tensor(vec![8]); + let weights = tensor(vec![8]); let plan = RouterPlan::SoftmaxTopK { scores: &scores, topk_indices: &indices, @@ -1927,13 +2070,12 @@ mod tests { #[test] fn expert_ref_rejects_shape_and_owner_mismatch() { - let mut gate_ptrs = GpuTensor::null_for_test(); - gate_ptrs.shape = vec![8]; - let mut down_ptrs = GpuTensor::null_for_test(); - down_ptrs.shape = vec![8]; + let gate_ptrs = tensor(vec![8]); + let down_ptrs = tensor(vec![8]); let mesh = hipfire_hardware::DeviceMesh::rect(&[(hipfire_hardware::DimKind::Ep, 2)]) .unwrap(); - let experts = MoeExpertRef::from_resolved( + let partition = [(0, 0, 0), (1, 1, 0), (2, 0, 1), (3, 1, 1)]; + let experts = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( &gate_ptrs, &down_ptrs, None, @@ -1942,18 +2084,20 @@ mod tests { 64, 128, &[0, 2], + &partition, + "router", Some(hipfire_hardware::DimKind::Ep), 0, &[0, 1], mesh.epoch(), - ); + )); experts .validate_projection_shapes(&[128, 128], &[128, 64]) .unwrap(); assert!(experts .validate_projection_shapes(&[64, 128], &[128, 64]) .is_err()); - let unknown = MoeExpertRef::from_resolved( + let unknown = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( &gate_ptrs, &down_ptrs, None, @@ -1962,11 +2106,13 @@ mod tests { 64, 128, &[4], + &partition, + "router", Some(hipfire_hardware::DimKind::Ep), 0, &[0, 1], mesh.epoch(), - ); + )); assert!(unknown.validate().is_err()); } diff --git a/crates/hipfire-dispatch/src/pipeline/mod.rs b/crates/hipfire-dispatch/src/pipeline/mod.rs index 5ce988b383..afc13c4b37 100644 --- a/crates/hipfire-dispatch/src/pipeline/mod.rs +++ b/crates/hipfire-dispatch/src/pipeline/mod.rs @@ -10,6 +10,8 @@ use hip_bridge; use rdna_compute::{DType, Gpu, GpuTensor}; use std::sync::{LazyLock, OnceLock}; +pub mod steps; + pub use steps::{ execute_sealed_steps_mesh, execute_steps, FusedPattern, GemvInput, SealedMoeSchedule, Step, StepCollective, diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index 78fba4ee53..94cc11dccc 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -12,9 +12,13 @@ use crate::context::DispatchCtx; use crate::families::fused_qkv::{FusedQkvBiasParams, FusedQkvFamily, FusedQkvParams}; use crate::families::gemv::{GemvFamily, GemvParams, RotateInputs, WeightRef}; use crate::families::moe::{ - ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeProj, MoeProtocolKind, MoeFamily, - RouterPlan, + ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeExpertRefBinding, MoeProj, + MoeProtocolKind, MoeFamily, RouterPlan, }; +use crate::families::rotation::{RotationFamily, RotationParams}; +use crate::types::GemvVariant; +use crate::types::{DispatchError, KernelKey, PipelineOp, RotationPlan, RotationVariant}; + /// Routing policy is carried by `RouterPlan`; no standalone score activation /// can be inserted between route and expert phases. @@ -211,6 +215,27 @@ impl StepCollective { } } } +/// Pointer-free identity for one executable rank schedule. Mesh execution +/// compares this value for every rank before launching any device work. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MoeExecutionSignature<'a> { + pub protocol: MoeProtocolKind, + pub execution: ExpertExecutionPlan, + pub router_identity: &'a str, + pub router_selection: crate::families::moe::RouterSelection, + pub k_top: usize, + pub normalize: bool, + pub expert_dtype: DType, + pub n_experts: usize, + pub expert_k: usize, + pub expert_m: usize, + pub batch_size: usize, + pub hidden: usize, + /// Canonical `(global expert, owner rank, local slot)` tuples for every + /// rank. Unlike pointer tables, this is safe to compare across devices. + pub ownership_partition: &'a [(usize, usize, usize)], +} + fn collective_count(collectives: &[StepCollective]) -> usize { collectives .iter() @@ -285,12 +310,15 @@ pub fn validate_moe_step_schedule( inverse_perm, }, ] => { + let expected_rows = batch_size + .checked_mul(*k_top) + .ok_or_else(|| DispatchError::Hip("indexed MoE row count overflows".into()))?; if *batch_size != 1 || *down_batch != 1 || inverse_perm.is_some() || *k_top != plan.k_top() || *inter != (*experts).expert_m() - || *rows != batch_size.saturating_mul(*k_top) + || *rows != expected_rows { return Err(DispatchError::Hip( "indexed MoE grammar requires decode batch=1, normalized route, and no inverse permutation" @@ -395,17 +423,24 @@ pub fn validate_moe_step_schedule( inverse_perm: combine_inverse, }, ] => { + let expected_slots = batch_size + .checked_mul(*k_top) + .ok_or_else(|| DispatchError::Hip("grouped MoE slot count overflows".into()))?; if *k_top != plan.k_top() || *k_top != *unscatter_k || *k_top != *down_k || *k_top != *combine_k || *batch_size != *down_batch || *batch_size != *combine_batch - || *total_slots != batch_size.saturating_mul(*k_top) + || *total_slots != expected_slots || *n_experts != experts.n_experts() - || !same_tensor(topk_indices, route.route_buffers().0) - || !same_tensor(topk_indices, gate_sorted) - || !same_tensor(topk_indices, down_sorted) + || !same_tensor(topk_indices, plan.route_buffers().0) + || same_tensor(topk_indices, sorted_slot_index) + || !same_tensor(sorted_slot_index, gate_sorted) + || !same_tensor(sorted_slot_index, unscatter_sorted) + || !same_tensor(sorted_slot_index, down_sorted) + || !same_tensor(expert_tile_ids, gate_tiles) + || !same_tensor(expert_tile_ids, down_tiles) { return Err(DispatchError::Hip( "grouped MoE route width, batch, or route identity mismatch".into(), @@ -429,9 +464,12 @@ pub fn validate_moe_step_schedule( "grouped MoE scatter capacity or tile geometry is invalid".into(), )); } - if expert_token_counts.numel() == 0 - || expert_offsets.numel() < *n_experts + 1 - { + let expected_offsets = n_experts + .checked_add(1) + .ok_or_else(|| DispatchError::Hip("MoE expert offset count overflows".into()))?; + let counts_elements = checked_numel(expert_token_counts, "expert counts")?; + let offset_elements = checked_numel(expert_offsets, "expert offsets")?; + if counts_elements == 0 || offset_elements < expected_offsets { return Err(DispatchError::Hip( "grouped MoE scatter metadata has empty or short buffers".into(), )); @@ -446,7 +484,7 @@ pub fn validate_moe_step_schedule( || !same_tensor(gate_batch, act_gate) || !same_tensor(up_batch, act_up) || *inter != *act_inter - || *rows != batch_size.saturating_mul(*k_top) + || *rows != expected_slots || !same_tensor(rot_out, down_x) || !same_tensor(down_y, combine_down) || !same_tensor(inverse_perm, combine_inverse) @@ -487,7 +525,11 @@ pub fn validate_moe_step_schedule( "MoE combine hidden size does not match the expert down projection".into(), )); } - let shape_gate = [2 * experts.expert_m(), experts.expert_k()]; + let gate_width = experts + .expert_m() + .checked_mul(2) + .ok_or_else(|| DispatchError::Hip("MoE gate/up shape overflows".into()))?; + let shape_gate = [gate_width, experts.expert_k()]; let shape_down = [experts.expert_k(), experts.expert_m()]; experts.validate_projection_shapes(&shape_gate, &shape_down)?; let dtype_supported = match protocol { @@ -542,6 +584,49 @@ pub fn validate_moe_protocol_schedule( Ok(()) } +fn derive_moe_execution_signature<'a>( + steps: &[Step<'a>], + execution: ExpertExecutionPlan, +) -> Result, DispatchError> { + let route = steps + .iter() + .find_map(|step| match step { + Step::MoeRoute { plan } => Some(plan), + _ => None, + }) + .ok_or_else(|| DispatchError::Hip("MoE schedule has no route".into()))?; + let experts = steps + .iter() + .find_map(|step| match step { + Step::IndexedMoeGemv { experts, .. } + | Step::GroupedMoeGemm { experts, .. } => Some(*experts), + _ => None, + }) + .ok_or_else(|| DispatchError::Hip("MoE schedule has no expert owner view".into()))?; + let hidden = steps + .iter() + .find_map(|step| match step { + Step::MoeCombine { hidden, .. } => Some(*hidden), + _ => None, + }) + .ok_or_else(|| DispatchError::Hip("MoE schedule has no combine geometry".into()))?; + Ok(MoeExecutionSignature { + protocol: execution.protocol()?, + execution, + router_identity: experts.router_identity(), + router_selection: route.selection(), + k_top: route.k_top(), + normalize: route.normalizes(), + expert_dtype: experts.dtype(), + n_experts: experts.n_experts(), + expert_k: experts.expert_k(), + expert_m: experts.expert_m(), + batch_size: route.batch_size(), + hidden, + ownership_partition: experts.ownership_partition(), + }) +} + /// An immutable, pre-validated typed MoE schedule. pub struct SealedMoeSchedule<'a> { execution: ExpertExecutionPlan, @@ -549,6 +634,7 @@ pub struct SealedMoeSchedule<'a> { collectives: Vec, } + impl<'a> SealedMoeSchedule<'a> { pub fn new( execution: ExpertExecutionPlan, @@ -574,6 +660,9 @@ impl<'a> SealedMoeSchedule<'a> { pub fn collectives(&self) -> &[StepCollective] { &self.collectives } + pub fn execution_signature(&self) -> Result, DispatchError> { + derive_moe_execution_signature(&self.steps, self.execution) + } } pub fn execute_sealed_steps( @@ -613,6 +702,7 @@ pub fn execute_sealed_steps_mesh<'a>( } let mut reduction: Option<(DimKind, usize, Vec, MeshEpoch)> = None; + let mut execution_signature: Option> = None; let mut outputs = Vec::with_capacity(schedules.len()); for (rank, schedule) in schedules.iter().enumerate() { validate_moe_protocol_schedule( @@ -620,6 +710,16 @@ pub fn execute_sealed_steps_mesh<'a>( schedule.collectives(), schedule.execution(), )?; + let signature = schedule.execution_signature()?; + if let Some(expected) = &execution_signature { + if expected != &signature { + return Err(DispatchError::Hip( + "MoE rank schedules disagree on executable identity".into(), + )); + } + } else { + execution_signature = Some(signature); + } let (kind, dim, group, epoch, descriptor_rank) = schedule .collectives() .iter() @@ -730,19 +830,34 @@ fn same_tensor(a: &GpuTensor, b: &GpuTensor) -> bool { || (!a.buf.as_ptr().is_null() && a.buf.as_ptr() == b.buf.as_ptr()) } +fn checked_numel(tensor: &GpuTensor, name: &str) -> Result { + tensor + .shape + .iter() + .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) + .ok_or_else(|| DispatchError::Hip(format!("MoE {name} logical shape overflows"))) +} + fn require_tensor( tensor: &GpuTensor, name: &str, dtype: DType, capacity: usize, ) -> Result<(), DispatchError> { - if tensor.dtype != dtype || tensor.numel() < capacity { + let logical_elements = checked_numel(tensor, name)?; + let required_bytes = capacity + .checked_mul(dtype.size()) + .ok_or_else(|| DispatchError::Hip(format!("MoE {name} byte capacity overflows")))?; + if tensor.dtype != dtype + || logical_elements < capacity + || tensor.buf.size() < required_bytes + { return Err(DispatchError::Hip(format!( - "MoE {name} has dtype {:?}/capacity {}, expected {:?}/at least {}", + "MoE {name} has dtype {:?}/logical capacity {logical_elements}/physical bytes {}, \ + expected {:?}/{capacity} elements/{required_bytes} bytes", tensor.dtype, - tensor.numel(), - dtype, - capacity + tensor.buf.size(), + dtype ))); } Ok(()) @@ -756,17 +871,22 @@ fn require_raw_i32( let bytes = elements .checked_mul(std::mem::size_of::()) .ok_or_else(|| DispatchError::Hip(format!("MoE {name} capacity overflows")))?; - if tensor.dtype != DType::Raw || tensor.numel() < bytes { + let logical_bytes = checked_numel(tensor, name)?; + if tensor.dtype != DType::Raw + || logical_bytes < bytes + || tensor.buf.size() < bytes + { return Err(DispatchError::Hip(format!( - "MoE {name} has dtype {:?}/capacity {}, expected Raw/at least {} bytes", + "MoE {name} has dtype {:?}/logical bytes {logical_bytes}/physical bytes {}, \ + expected Raw/at least {bytes} bytes", tensor.dtype, - tensor.numel(), - bytes + tensor.buf.size() ))); } Ok(()) } + fn validate_step_tensors( steps: &[Step], experts: &MoeExpertRef<'_>, @@ -881,7 +1001,7 @@ fn validate_step_tensors( batch_size .checked_mul(k_top) .and_then(|slots| slots.checked_mul(*hidden)) - .ok_or_else(|| DispatchError::Hip("MoE combine input capacity overflows"))?, + .ok_or_else(|| DispatchError::Hip("MoE combine input capacity overflows".into()))?, )?; require_tensor( topk_weights, @@ -895,7 +1015,7 @@ fn validate_step_tensors( DType::F32, batch_size .checked_mul(*hidden) - .ok_or_else(|| DispatchError::Hip("MoE output capacity overflow"))?, + .ok_or_else(|| DispatchError::Hip("MoE output capacity overflow".into()))?, )?; if same_tensor(down_out, out) || same_tensor(topk_weights, out) { return Err(DispatchError::Hip( @@ -915,8 +1035,11 @@ fn validate_step_tensors( block_m, .. } => { + let offsets = n_experts + .checked_add(1) + .ok_or_else(|| DispatchError::Hip("MoE expert offsets overflow".into()))?; require_raw_i32(expert_token_counts, "expert counts", *n_experts)?; - require_raw_i32(expert_offsets, "expert offsets", *n_experts + 1)?; + require_raw_i32(expert_offsets, "expert offsets", offsets)?; require_raw_i32(sorted_slot_index, "sorted slots", *m_total_max)?; require_raw_i32(inverse_perm, "inverse permutation", *total_slots)?; require_raw_i32( @@ -980,7 +1103,7 @@ fn validate_step_tensors( DispatchError::Hip("MoE unscatter width overflows".into()) })?, ) - .ok_or_else(|| DispatchError::Hip("MoE unscatter input overflows"))?, + .ok_or_else(|| DispatchError::Hip("MoE unscatter input overflows".into()))?, )?; require_tensor( gate_batch, @@ -2688,7 +2811,16 @@ mod tests { ); } fn tensor(shape: Vec) -> GpuTensor { + let elements = shape.iter().product::(); let mut tensor = GpuTensor::null_for_test(); + tensor.buf = unsafe { + hip_bridge::DeviceBuffer::from_raw( + std::ptr::null_mut(), + elements + .checked_mul(DType::F32.size()) + .expect("test tensor bytes"), + ) + }; tensor.shape = shape; tensor } @@ -2767,7 +2899,8 @@ mod tests { let out = tensor(vec![128]); let gate_ptrs = tensor(vec![16]); let down_ptrs = tensor(vec![16]); - let experts = MoeExpertRef::from_resolved( + let ownership = (0..8).map(|expert| (expert, 0, expert)).collect::>(); + let experts = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( &gate_ptrs, &down_ptrs, None, @@ -2776,11 +2909,13 @@ mod tests { 64, 128, &[0, 1, 2, 3, 4, 5, 6, 7], + &ownership, + "router", None, 0, &[0], DeviceMesh::single().unwrap().epoch(), - ); + )); let steps = indexed_steps( &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, ); @@ -2826,7 +2961,17 @@ mod tests { let out = tensor(vec![128]); let gate_ptrs = tensor(vec![16]); let down_ptrs = tensor(vec![16]); - let experts = MoeExpertRef::from_resolved( + let ownership = vec![ + (0, 0, 0), + (1, 0, 1), + (2, 1, 0), + (3, 1, 1), + (4, 1, 2), + (5, 1, 3), + (6, 1, 4), + (7, 1, 5), + ]; + let experts = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( &gate_ptrs, &down_ptrs, None, @@ -2835,11 +2980,13 @@ mod tests { 64, 128, &[0, 1], + &ownership, + "router", Some(DimKind::Ep), 0, &[0, 1], DeviceMesh::rect(&[(DimKind::Ep, 2)]).unwrap().epoch(), - ); + )); let steps = indexed_steps( &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, ); diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs index 9de469a355..41b0e98299 100644 --- a/crates/hipfire-runtime/src/moe_plan.rs +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -56,17 +56,24 @@ impl fmt::Display for ExpertPlanError { impl std::error::Error for ExpertPlanError {} -#[derive(Clone, Debug, PartialEq, Eq)] +struct ExpertPointerTables { + gate_up: GpuTensor, + down: GpuTensor, + dummy_gate_up: Option, +} + struct ExpertStorageOwner { slots: Vec, resident: Vec, + rank_tables: Vec>, } impl ExpertStorageOwner { - fn new(slots: Vec) -> Self { + fn new(slots: Vec, group_size: usize) -> Self { Self { resident: vec![false; slots.len()], slots, + rank_tables: (0..group_size).map(|_| None).collect(), } } @@ -74,6 +81,14 @@ impl ExpertStorageOwner { self.slots.iter().position(|candidate| *candidate == placement) } + fn rank_is_resident(&self, rank: usize) -> bool { + self.slots + .iter() + .enumerate() + .filter(|(_, placement)| placement.owner == rank) + .all(|(index, _)| self.resident[index]) + } + fn clear(&mut self) { self.resident.fill(false); } @@ -83,6 +98,7 @@ impl ExpertStorageOwner { } } + /// A sealed, manifest-derived expert plan. /// /// Every placement, source identity, shape, resource requirement, rank-local @@ -107,6 +123,7 @@ pub struct ExpertPlan { collective: Option, collective_row: Option, owner_views: Vec>, + owner_partition: Vec<(usize, usize, usize)>, owner: ExpertStorageOwner, } @@ -146,6 +163,12 @@ impl ExpertPlan { .map_err(ExpertPlanError::new)?; validate_spec(spec, mesh)?; let (shape, source_dtype) = resolve_shape(spec, manifest)?; + if !shape.fused_gate_up { + return Err(ExpertPlanError::new(format!( + "expert group '{}' uses separate gate/up sources, unsupported by the generic executor", + spec.group + ))); + } let execution_plan = parse_execution(&spec.execution, spec)?; validate_execution_dtype(execution_plan, source_dtype, spec)?; @@ -156,13 +179,14 @@ impl ExpertPlan { spec.group ))); } + let group_size = group_devices.len(); let owner = resolve_placements( spec.n_experts, - group_devices.len(), + group_size, spec.parallelism, spec.assignment, ); - let owner_views = (0..group_devices.len()) + let owner_views = (0..group_size) .map(|rank| { owner .iter() @@ -171,6 +195,10 @@ impl ExpertPlan { .collect() }) .collect(); + let owner_partition = owner + .iter() + .map(|placement| (placement.global_id, placement.owner, placement.local_slot)) + .collect(); let (collective, collective_row) = resolve_collective(spec, manifest)?; Ok(Self { @@ -191,7 +219,8 @@ impl ExpertPlan { collective, collective_row, owner_views, - owner: ExpertStorageOwner::new(owner), + owner_partition, + owner: ExpertStorageOwner::new(owner, group_size), }) } @@ -288,35 +317,92 @@ impl ExpertPlan { }) } - /// Bind canonical shape/dtype/owner metadata to a borrowed pointer-table - /// view. The family can retain only the returned immutable view; all - /// placement and storage state remains owned by this plan. + /// Commit the rank-local pointer tables to this owner. The plan validates + /// the table ABI before retaining the tensors; binding later borrows only + /// these committed values. + pub fn commit_rank_tables( + &mut self, + rank: usize, + gate_up: GpuTensor, + down: GpuTensor, + dummy_gate_up: Option, + ) -> Result<(), ExpertPlanError> { + if rank >= self.group_size() { + return Err(ExpertPlanError::new(format!( + "expert group '{}' rank {rank} is outside group size {}", + self.group, + self.group_size() + ))); + } + validate_pointer_table(&gate_up, "gate/up", self.n_experts)?; + validate_pointer_table(&down, "down", self.n_experts)?; + if let Some(dummy) = &dummy_gate_up { + validate_dummy_table(dummy)?; + } + let tables = self + .owner + .rank_tables + .get_mut(rank) + .expect("rank checked above"); + if tables.is_some() { + return Err(ExpertPlanError::new(format!( + "expert group '{}' rank {rank} pointer tables are already committed", + self.group + ))); + } + *tables = Some(ExpertPointerTables { + gate_up, + down, + dummy_gate_up, + }); + Ok(()) + } + + /// Bind the committed rank-local pointer tables to an opaque executable + /// view. A rank cannot bind until its owned placements are resident. pub fn bind_expert_ref<'a>( &'a self, rank: usize, - gate_up_ptrs: &'a GpuTensor, - down_ptrs: &'a GpuTensor, - dummy_gate_up: Option<&'a GpuTensor>, ) -> Result, ExpertPlanError> { let owned = self.owned_experts(rank)?; + if !self.owner.rank_is_resident(rank) { + return Err(ExpertPlanError::new(format!( + "expert group '{}' rank {rank} has nonresident placements", + self.group + ))); + } + let tables = self + .owner + .rank_tables + .get(rank) + .and_then(Option::as_ref) + .ok_or_else(|| { + ExpertPlanError::new(format!( + "expert group '{}' rank {rank} has no committed pointer tables", + self.group + )) + })?; let collective_kind = match self.collective { Some(CollectiveHint::AllReduce { kind }) => Some(kind), _ => None, }; - let view = MoeExpertRef::from_resolved( - gate_up_ptrs, - down_ptrs, - dummy_gate_up, + let binding = hipfire_dispatch::families::moe::MoeExpertRefBinding::from_plan( + &tables.gate_up, + &tables.down, + tables.dummy_gate_up.as_ref(), self.source_dtype, self.n_experts, self.shape.expert_m, self.shape.expert_k, owned, + &self.owner_partition, + &self.router, collective_kind, rank, &self.group_devices, self.mesh_epoch, ); + let view = MoeExpertRef::from_binding(binding); view.validate() .map_err(|error| ExpertPlanError::new(error.to_string()))?; Ok(view) @@ -389,8 +475,14 @@ pub struct ExpertLoadTxn<'a> { } impl ExpertLoadTxn<'_> { - /// Reserve one manifest placement. Duplicate reservations are refused. + /// Reserve one manifest placement. Duplicate reservations and reuse after + /// commit/rollback are refused. pub fn reserve(&mut self, placement: ExpertPlacement) -> Result<(), ExpertPlanError> { + if self.finished { + return Err(ExpertPlanError::new( + "expert load transaction is already finished", + )); + } let index = self.owner.index_of(placement).ok_or_else(|| { ExpertPlanError::new(format!("unknown expert placement {placement:?}")) })?; @@ -488,7 +580,7 @@ fn validate_spec(spec: &ExpertGroupSpec, mesh: &DeviceMesh) -> Result<(), Expert { return Err(ExpertPlanError::new(format!( "expert group '{}' n_experts={} is not divisible by group_size={group_size}", - spec.group, spec.n_experts, group_size + spec.group, spec.n_experts ))); } Ok(()) @@ -662,6 +754,51 @@ fn source_names<'a>(layout: &'a ExpertSourceLayout) -> Vec<(&'static str, Vec<&' } +fn checked_numel(tensor: &GpuTensor, label: &str) -> Result { + tensor + .shape + .iter() + .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) + .ok_or_else(|| ExpertPlanError::new(format!("{label} logical shape overflows"))) +} + +fn validate_pointer_table( + tensor: &GpuTensor, + label: &str, + n_experts: usize, +) -> Result<(), ExpertPlanError> { + let pointer_slots = n_experts + .checked_mul(2) + .ok_or_else(|| ExpertPlanError::new("pointer-table slot count overflows"))?; + let required_bytes = pointer_slots + .checked_mul(DType::F32.size()) + .ok_or_else(|| ExpertPlanError::new("pointer-table byte capacity overflows"))?; + let logical_elements = checked_numel(tensor, label)?; + if tensor.dtype != DType::F32 + || tensor.shape.as_slice() != [pointer_slots] + || logical_elements < pointer_slots + || tensor.buf.size() < required_bytes + { + return Err(ExpertPlanError::new(format!( + "{label} pointer table must be F32 [{pointer_slots}] with {required_bytes} physical bytes" + ))); + } + Ok(()) +} + +fn validate_dummy_table(tensor: &GpuTensor) -> Result<(), ExpertPlanError> { + let logical_elements = checked_numel(tensor, "dummy gate/up table")?; + let required_bytes = logical_elements + .checked_mul(DType::F32.size()) + .ok_or_else(|| ExpertPlanError::new("dummy gate/up table byte capacity overflows"))?; + if tensor.dtype != DType::F32 || logical_elements == 0 || tensor.buf.size() < required_bytes { + return Err(ExpertPlanError::new( + "dummy gate/up table must be nonempty F32 storage with physical capacity", + )); + } + Ok(()) +} + fn entry_for<'a>( spec: &ExpertGroupSpec, manifest: &'a [WeightEntry], @@ -805,6 +942,7 @@ fn resolve_shape( ))); } if gate_shapes.iter().any(|shape| shape != &gate_shapes[0]) + || up_shapes.iter().any(|shape| shape != &up_shapes[0]) || down_shapes.iter().any(|shape| shape != &down_shapes[0]) { return Err(ExpertPlanError::new(format!( @@ -814,14 +952,25 @@ fn resolve_shape( } let gate = &gate_shapes[0]; let down = &down_shapes[0]; - if gate.len() != 2 || down.len() != 2 || gate[0] % 2 != 0 { + if gate.len() != 2 || down.len() != 2 { return Err(ExpertPlanError::new(format!( "expert group '{}' layer {:?} projection shapes are incompatible", spec.group, spec.layer ))); } + let expert_m = if fused { + if gate[0] % 2 != 0 { + return Err(ExpertPlanError::new(format!( + "expert group '{}' layer {:?} fused gate/up width is not even", + spec.group, spec.layer + ))); + } + gate[0] / 2 + } else { + gate[0] + }; let shape = ExpertShape { - expert_m: gate[0] / 2, + expert_m, expert_k: gate[1], fused_gate_up: fused, }; @@ -989,7 +1138,16 @@ mod tests { } fn table(shape: Vec) -> GpuTensor { + let elements = shape.iter().product::(); let mut tensor = GpuTensor::null_for_test(); + tensor.buf = unsafe { + hip_bridge::DeviceBuffer::from_raw( + std::ptr::null_mut(), + elements + .checked_mul(DType::F32.size()) + .expect("test tensor bytes"), + ) + }; tensor.shape = shape; tensor } @@ -1014,7 +1172,7 @@ mod tests { #[test] fn bound_view_uses_canonical_rank_owner() { - let plan = ExpertPlan::from_manifest( + let mut plan = ExpertPlan::from_manifest( &spec("indexed_quantized", ExpertParallelism::ExpertParallel), &manifest(), &mesh_ep(), @@ -1022,7 +1180,21 @@ mod tests { .unwrap(); let gate = table(vec![8]); let down = table(vec![8]); - let view = plan.bind_expert_ref(1, &gate, &down, None).unwrap(); + plan.commit_rank_tables(1, gate, down, None).unwrap(); + { + let placements: Vec<_> = plan + .placements() + .iter() + .copied() + .filter(|placement| placement.owner == 1) + .collect(); + let mut load = plan.begin_load(); + for placement in placements { + load.reserve(placement).unwrap(); + } + load.commit(); + } + let view = plan.bind_expert_ref(1).unwrap(); assert_eq!(view.owned(), &[1, 3]); assert_eq!(view.n_experts(), 4); } From 37cd1378291d1cfcc6cfbf8cddb4789f4611d9a1 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 10:42:29 +0200 Subject: [PATCH 19/25] test: close g5 moe sealing evidence --- crates/hipfire-dispatch/src/families/moe.rs | 97 ++-- crates/hipfire-dispatch/src/pipeline/steps.rs | 259 ++------- crates/hipfire-runtime/src/moe_plan.rs | 530 +++++++++++++++++- 3 files changed, 597 insertions(+), 289 deletions(-) diff --git a/crates/hipfire-dispatch/src/families/moe.rs b/crates/hipfire-dispatch/src/families/moe.rs index fd09093210..c0cf03850b 100644 --- a/crates/hipfire-dispatch/src/families/moe.rs +++ b/crates/hipfire-dispatch/src/families/moe.rs @@ -227,8 +227,10 @@ impl ExpertExecutionPlan { /// Opaque plan-produced data used to create an expert view. /// /// The binding is intentionally separate from [`MoeExpertRef`]: callers can -/// retain the view for scheduling, but there is no public constructor that -/// accepts arbitrary canonical metadata directly on the executable reference. +/// retain the view for scheduling, but there is no safe constructor that +/// accepts arbitrary canonical metadata. The only bridge is the hidden, +/// `unsafe` function below; the runtime plan is responsible for proving its +/// inputs before crossing that boundary. pub struct MoeExpertRefBinding<'a> { gate_up_ptrs: &'a GpuTensor, down_ptrs: &'a GpuTensor, @@ -247,11 +249,49 @@ pub struct MoeExpertRefBinding<'a> { } impl<'a> MoeExpertRefBinding<'a> { - /// Runtime-plan bridge. The returned binding is opaque to callers and can - /// only be consumed by [`MoeExpertRef::from_binding`]. + /// Cross-crate bridge for the sealed [`ExpertPlan`](https://docs.rs/hipfire-runtime/latest/hipfire_runtime/moe_plan/struct.ExpertPlan.html). + /// + /// This function is intentionally hidden from the normal API + /// documentation and is not re-exported by the dispatch crate. It has no + /// safe raw-constructor counterpart. + /// + /// A safe external caller cannot invoke this bridge: + /// + /// ```compile_fail + /// use hipfire_dispatch::families::moe::MoeExpertRefBinding; + /// use hipfire_hardware::DeviceMesh; + /// use rdna_compute::{DType, GpuTensor}; + /// + /// fn safe_call(table: &GpuTensor) { + /// let _ = MoeExpertRefBinding::from_validated_plan( + /// table, + /// table, + /// None, + /// DType::F32, + /// 1, + /// 1, + /// 1, + /// &[0], + /// &[(0, 0, 0)], + /// "router", + /// None, + /// 0, + /// &[0], + /// DeviceMesh::single().unwrap().epoch(), + /// ); + /// } + /// ``` + /// + /// # Safety + /// + /// The caller must be the private `ExpertPlan::bind_expert_ref` path and + /// must have already proved that these references are the plan's committed + /// rank-local tables, that the owner placement is resident, and that all + /// metadata is derived from the same sealed plan/mesh epoch. Callers must + /// not retain or mutate a binding after its owner plan is dropped. #[doc(hidden)] #[allow(clippy::too_many_arguments)] - pub fn from_plan( + pub unsafe fn from_validated_plan( gate_up_ptrs: &'a GpuTensor, down_ptrs: &'a GpuTensor, dummy_gate_up: Option<&'a GpuTensor>, @@ -2068,53 +2108,6 @@ mod tests { plan.validate_against(8, 1).unwrap(); } - #[test] - fn expert_ref_rejects_shape_and_owner_mismatch() { - let gate_ptrs = tensor(vec![8]); - let down_ptrs = tensor(vec![8]); - let mesh = hipfire_hardware::DeviceMesh::rect(&[(hipfire_hardware::DimKind::Ep, 2)]) - .unwrap(); - let partition = [(0, 0, 0), (1, 1, 0), (2, 0, 1), (3, 1, 1)]; - let experts = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( - &gate_ptrs, - &down_ptrs, - None, - DType::MQ4G256, - 4, - 64, - 128, - &[0, 2], - &partition, - "router", - Some(hipfire_hardware::DimKind::Ep), - 0, - &[0, 1], - mesh.epoch(), - )); - experts - .validate_projection_shapes(&[128, 128], &[128, 64]) - .unwrap(); - assert!(experts - .validate_projection_shapes(&[64, 128], &[128, 64]) - .is_err()); - let unknown = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( - &gate_ptrs, - &down_ptrs, - None, - DType::MQ4G256, - 4, - 64, - 128, - &[4], - &partition, - "router", - Some(hipfire_hardware::DimKind::Ep), - 0, - &[0, 1], - mesh.epoch(), - )); - assert!(unknown.validate().is_err()); - } #[test] fn fallback_execution_has_no_protocol() { diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index 94cc11dccc..9ea9e83b70 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -12,8 +12,8 @@ use crate::context::DispatchCtx; use crate::families::fused_qkv::{FusedQkvBiasParams, FusedQkvFamily, FusedQkvParams}; use crate::families::gemv::{GemvFamily, GemvParams, RotateInputs, WeightRef}; use crate::families::moe::{ - ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeExpertRefBinding, MoeProj, - MoeProtocolKind, MoeFamily, RouterPlan, + ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeProj, MoeProtocolKind, MoeFamily, + RouterPlan, }; use crate::families::rotation::{RotationFamily, RotationParams}; use crate::types::GemvVariant; @@ -679,25 +679,29 @@ pub fn execute_sealed_steps( execute_steps_inner(gpu, ctx, &schedule.steps) } -/// Execute one sealed MoE schedule per participating device and perform its -/// single manifest-owned routed reduction. Schedules are ordered by the -/// collective's rank field; every schedule is validated before any GPU work. -pub fn execute_sealed_steps_mesh<'a>( - gpus: &mut Gpus, +/// Results of the host-only mesh checks shared by the executor and focused +/// preflight tests. No device method is called while this value is built. +struct MoeMeshPreflight<'a> { + dim: usize, + group: Vec, + outputs: Vec<&'a hip_bridge::DeviceBuffer>, +} + +fn preflight_sealed_steps_mesh<'a>( + gpus_len: usize, mesh: &DeviceMesh, - ctx: &DispatchCtx, schedules: &[&SealedMoeSchedule<'a>], -) -> Result<(), DispatchError> { +) -> Result, DispatchError> { if schedules.is_empty() { return Err(DispatchError::Hip( "parallel MoE execution has no rank schedules".into(), )); } - if mesh.n_devices() != gpus.devices.len() { + if mesh.n_devices() != gpus_len { return Err(DispatchError::Hip(format!( "MoE mesh has {} devices but Gpus owns {}", mesh.n_devices(), - gpus.devices.len() + gpus_len ))); } @@ -812,16 +816,46 @@ pub fn execute_sealed_steps_mesh<'a>( } for &device in &group { - if device >= gpus.devices.len() { + if device >= gpus_len { return Err(DispatchError::Hip(format!( "MoE collective device {device} is outside the Gpus owner" ))); } } - for (rank, &device) in group.iter().enumerate() { + Ok(MoeMeshPreflight { + dim, + group, + outputs, + }) +} +/// Run the same host-only validation used by +/// [`execute_sealed_steps_mesh`] without touching a device. +/// +/// This is a hidden test/integration seam: callers still need fully sealed +/// schedules, and the GPU executor invokes the identical preflight internally. +#[doc(hidden)] +pub fn validate_sealed_steps_mesh_preflight<'a>( + gpus_len: usize, + mesh: &DeviceMesh, + schedules: &[&SealedMoeSchedule<'a>], +) -> Result<(), DispatchError> { + preflight_sealed_steps_mesh(gpus_len, mesh, schedules).map(|_| ()) +} + +/// Execute one sealed MoE schedule per participating device and perform its +/// single manifest-owned routed reduction. Schedules are ordered by the +/// collective's rank field; every schedule is validated before any GPU work. +pub fn execute_sealed_steps_mesh<'a>( + gpus: &mut Gpus, + mesh: &DeviceMesh, + ctx: &DispatchCtx, + schedules: &[&SealedMoeSchedule<'a>], +) -> Result<(), DispatchError> { + let preflight = preflight_sealed_steps_mesh(gpus.devices.len(), mesh, schedules)?; + for (rank, &device) in preflight.group.iter().enumerate() { execute_steps_inner(&mut gpus.devices[device], ctx, schedules[rank].steps())?; } - gpus.all_reduce_sum_f32_peer(&group, &outputs, dim) + gpus.all_reduce_sum_f32_peer(&preflight.group, &preflight.outputs, preflight.dim) .map_err(|error| DispatchError::Hip(error.to_string())) } @@ -2810,203 +2844,6 @@ mod tests { "FusedGateUpQ8_0 missing from FUSED_TABLE" ); } - fn tensor(shape: Vec) -> GpuTensor { - let elements = shape.iter().product::(); - let mut tensor = GpuTensor::null_for_test(); - tensor.buf = unsafe { - hip_bridge::DeviceBuffer::from_raw( - std::ptr::null_mut(), - elements - .checked_mul(DType::F32.size()) - .expect("test tensor bytes"), - ) - }; - tensor.shape = shape; - tensor - } - fn indexed_steps<'a>( - experts: &'a MoeExpertRef<'a>, - scores: &'a GpuTensor, - indices: &'a GpuTensor, - weights: &'a GpuTensor, - x: &'a GpuTensor, - gate: &'a GpuTensor, - up: &'a GpuTensor, - rot: &'a GpuTensor, - down: &'a GpuTensor, - out: &'a GpuTensor, - ) -> Vec> { - vec![ - Step::MoeRoute { - plan: RouterPlan::SoftmaxTopK { - scores, - topk_indices: indices, - topk_weights: weights, - k_top: 8, - normalize: true, - }, - }, - Step::IndexedMoeGemv { - experts, - which: MoeProj::GateUp { up_out: up }, - topk_indices: indices, - input: GemvInput::Prerotated(x), - out: gate, - k_top: 8, - batch_size: 1, - }, - Step::MoeActivation { - variant: MoeActivationVariant::SiluMul, - gate, - up, - rot_out: rot, - inter: 64, - rows: 8, - }, - Step::IndexedMoeGemv { - experts, - which: MoeProj::DownExpanded, - topk_indices: indices, - input: GemvInput::Prerotated(rot), - out: down, - k_top: 8, - batch_size: 1, - }, - Step::MoeCombine { - down_out: down, - topk_weights: weights, - out, - hidden: 128, - k_top: 8, - batch_size: 1, - inverse_perm: None, - }, - ] - } - - - #[test] - fn typed_indexed_grammar_requires_one_combine() { - let scores = tensor(vec![8]); - let indices = tensor(vec![8]); - let weights = tensor(vec![8]); - let x = tensor(vec![128]); - let gate = tensor(vec![8 * 64]); - let up = tensor(vec![8 * 64]); - let rot = tensor(vec![8 * 64]); - let down = tensor(vec![8 * 128]); - let out = tensor(vec![128]); - let gate_ptrs = tensor(vec![16]); - let down_ptrs = tensor(vec![16]); - let ownership = (0..8).map(|expert| (expert, 0, expert)).collect::>(); - let experts = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( - &gate_ptrs, - &down_ptrs, - None, - DType::MQ4G256, - 8, - 64, - 128, - &[0, 1, 2, 3, 4, 5, 6, 7], - &ownership, - "router", - None, - 0, - &[0], - DeviceMesh::single().unwrap().epoch(), - )); - let steps = indexed_steps( - &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, - ); - validate_moe_step_schedule( - &steps, - &[ - StepCollective::None, - StepCollective::None, - StepCollective::None, - StepCollective::None, - StepCollective::None, - ], - ) - .unwrap(); - let mut duplicate = steps; - duplicate.push(Step::MoeCombine { - down_out: &down, - topk_weights: &weights, - out: &out, - hidden: 128, - k_top: 8, - batch_size: 1, - inverse_perm: None, - }); - let error = validate_moe_step_schedule( - &duplicate, - &vec![StepCollective::None; duplicate.len()], - ) - .expect_err("duplicate combine must be rejected by exact grammar"); - assert!(error.to_string().contains("exactly indexed or grouped")); - } - #[test] - fn parallel_schedule_requires_named_axis_collective() { - let scores = tensor(vec![8]); - let indices = tensor(vec![8]); - let weights = tensor(vec![8]); - let x = tensor(vec![128]); - let gate = tensor(vec![8 * 64]); - let up = tensor(vec![8 * 64]); - let rot = tensor(vec![8 * 64]); - let down = tensor(vec![8 * 128]); - let out = tensor(vec![128]); - let gate_ptrs = tensor(vec![16]); - let down_ptrs = tensor(vec![16]); - let ownership = vec![ - (0, 0, 0), - (1, 0, 1), - (2, 1, 0), - (3, 1, 1), - (4, 1, 2), - (5, 1, 3), - (6, 1, 4), - (7, 1, 5), - ]; - let experts = MoeExpertRef::from_binding(MoeExpertRefBinding::from_plan( - &gate_ptrs, - &down_ptrs, - None, - DType::MQ4G256, - 8, - 64, - 128, - &[0, 1], - &ownership, - "router", - Some(DimKind::Ep), - 0, - &[0, 1], - DeviceMesh::rect(&[(DimKind::Ep, 2)]).unwrap().epoch(), - )); - let steps = indexed_steps( - &experts, &scores, &indices, &weights, &x, &gate, &up, &rot, &down, &out, - ); - let mesh = DeviceMesh::rect(&[(DimKind::Ep, 2)]).unwrap(); - let mut collectives = vec![StepCollective::None; steps.len()]; - collectives[4] = StepCollective::all_reduce( - DimKind::Ep, - 128, - vec![0, 1], - mesh.epoch(), - 0, - ); - validate_moe_parallel_schedule(&steps, &collectives).unwrap(); - collectives[4] = StepCollective::all_reduce( - DimKind::Tp, - 128, - vec![0, 1], - mesh.epoch(), - 0, - ); - assert!(validate_moe_parallel_schedule(&steps, &collectives).is_err()); - } } diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs index 41b0e98299..16681e7530 100644 --- a/crates/hipfire-runtime/src/moe_plan.rs +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -386,22 +386,26 @@ impl ExpertPlan { Some(CollectiveHint::AllReduce { kind }) => Some(kind), _ => None, }; - let binding = hipfire_dispatch::families::moe::MoeExpertRefBinding::from_plan( - &tables.gate_up, - &tables.down, - tables.dummy_gate_up.as_ref(), - self.source_dtype, - self.n_experts, - self.shape.expert_m, - self.shape.expert_k, - owned, - &self.owner_partition, - &self.router, - collective_kind, - rank, - &self.group_devices, - self.mesh_epoch, - ); + // SAFETY: this private plan path has checked the committed rank table, + // resident owner placements, and all metadata comes from `self`. + let binding = unsafe { + hipfire_dispatch::families::moe::MoeExpertRefBinding::from_validated_plan( + &tables.gate_up, + &tables.down, + tables.dummy_gate_up.as_ref(), + self.source_dtype, + self.n_experts, + self.shape.expert_m, + self.shape.expert_k, + owned, + &self.owner_partition, + &self.router, + collective_kind, + rank, + &self.group_devices, + self.mesh_epoch, + ) + }; let view = MoeExpertRef::from_binding(binding); view.validate() .map_err(|error| ExpertPlanError::new(error.to_string()))?; @@ -1079,6 +1083,10 @@ fn resolve_collective( #[cfg(test)] mod tests { use super::*; + use hipfire_dispatch::families::moe::{ + MoeActivationVariant, MoeFamily, MoeProj, MoeProtocolKind, RouterPlan, + }; + use hipfire_dispatch::pipeline::{GemvInput, Step, StepCollective}; fn mesh_ep() -> DeviceMesh { DeviceMesh::rect(&[(DimKind::Ep, 2)]).expect("test mesh") @@ -1116,6 +1124,44 @@ mod tests { ] } + fn manifest_for(parallelism: ExpertParallelism) -> Vec { + let mut entries = manifest(); + if matches!(parallelism, ExpertParallelism::Single) { + entries[1].policy = ShardPolicy::Replicate; + entries[2].policy = ShardPolicy::Replicate; + } + entries + } + + fn separate_manifest() -> Vec { + let mut entries = manifest(); + entries[1].name = "experts.gate".into(); + entries[1].logical_shape = vec![4, 64, 64]; + entries[2].name = "experts.down".into(); + entries.push(WeightEntry::layer( + "experts.up", + 0, + vec![4, 64, 64], + DType::MQ4G256, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + )); + entries + } + + fn separate_spec(execution: &str, parallelism: ExpertParallelism) -> ExpertGroupSpec { + let mut value = spec(execution, parallelism); + value.source_layout = ExpertSourceLayout::PackedSeparate { + gate: "experts.gate".into(), + up: "experts.up".into(), + down: "experts.down".into(), + sidecars: vec![], + }; + value + } + fn spec(execution: &str, parallelism: ExpertParallelism) -> ExpertGroupSpec { ExpertGroupSpec { group: "block-0".into(), @@ -1137,21 +1183,275 @@ mod tests { } } - fn table(shape: Vec) -> GpuTensor { - let elements = shape.iter().product::(); + fn tensor_with_bytes(shape: Vec, dtype: DType, bytes: usize) -> GpuTensor { let mut tensor = GpuTensor::null_for_test(); tensor.buf = unsafe { - hip_bridge::DeviceBuffer::from_raw( - std::ptr::null_mut(), - elements - .checked_mul(DType::F32.size()) - .expect("test tensor bytes"), - ) + hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut(), bytes) }; tensor.shape = shape; + tensor.dtype = dtype; tensor } + fn tensor(shape: Vec) -> GpuTensor { + let elements = shape.iter().product::(); + tensor_with_bytes( + shape, + DType::F32, + elements.checked_mul(DType::F32.size()).expect("test tensor bytes"), + ) + } + + fn table(shape: Vec) -> GpuTensor { + tensor(shape) + } + + fn raw_i32(elements: usize) -> GpuTensor { + let bytes = elements + .checked_mul(std::mem::size_of::()) + .expect("test raw bytes"); + tensor_with_bytes(vec![bytes], DType::Raw, bytes) + } + + fn resident_plan( + execution: &str, + parallelism: ExpertParallelism, + mesh: &DeviceMesh, + ) -> ExpertPlan { + let mut plan = ExpertPlan::from_manifest( + &spec(execution, parallelism), + &manifest_for(parallelism), + mesh, + ) + .expect("test expert plan"); + let group_size = plan.group_size(); + for rank in 0..group_size { + plan.commit_rank_tables(rank, table(vec![8]), table(vec![8]), None) + .expect("commit rank tables"); + } + let placements = plan.placements().to_vec(); + let mut load = plan.begin_load(); + for placement in placements { + load.reserve(placement).expect("reserve placement"); + } + load.commit(); + plan + } + struct GroupedTensors { + scores: GpuTensor, + indices: GpuTensor, + weights: GpuTensor, + counts: GpuTensor, + offsets: GpuTensor, + sorted: GpuTensor, + tiles: GpuTensor, + inverse: GpuTensor, + x: GpuTensor, + grouped_gate: GpuTensor, + gate_batch: GpuTensor, + up_batch: GpuTensor, + rot_batch: GpuTensor, + grouped_down: GpuTensor, + down_x: GpuTensor, + out: GpuTensor, + } + + fn grouped_tensors() -> GroupedTensors { + GroupedTensors { + scores: tensor(vec![2, 4]), + indices: tensor(vec![2, 2]), + weights: tensor(vec![2, 2]), + counts: raw_i32(4), + offsets: raw_i32(5), + sorted: raw_i32(8), + tiles: raw_i32(2), + inverse: raw_i32(4), + x: tensor(vec![2, 64]), + grouped_gate: tensor(vec![8, 128]), + gate_batch: tensor(vec![4, 64]), + up_batch: tensor(vec![4, 64]), + rot_batch: tensor(vec![4, 64]), + grouped_down: tensor(vec![8, 64]), + down_x: tensor(vec![4, 64]), + out: tensor(vec![2, 64]), + } + } + + fn grouped_steps<'a>( + experts: &'a MoeExpertRef<'a>, + tensors: &'a GroupedTensors, + ) -> Vec> { + vec![ + Step::MoeRoute { + plan: RouterPlan::SoftmaxTopK { + scores: &tensors.scores, + topk_indices: &tensors.indices, + topk_weights: &tensors.weights, + k_top: 2, + normalize: true, + }, + }, + Step::MoeScatter { + topk_indices: &tensors.indices, + expert_token_counts: &tensors.counts, + expert_offsets: &tensors.offsets, + sorted_slot_index: &tensors.sorted, + expert_tile_ids: &tensors.tiles, + inverse_perm: &tensors.inverse, + total_slots: 4, + n_experts: 4, + m_total_max: 8, + block_m: 4, + }, + Step::GroupedMoeGemm { + experts, + which: MoeProj::GateUp { + up_out: &tensors.up_batch, + }, + sorted_slot_index: &tensors.sorted, + expert_tile_ids: &tensors.tiles, + x: &tensors.x, + y: &tensors.grouped_gate, + m_total: 8, + batch_size: 2, + k_top: 2, + }, + Step::MoeGateUpUnscatter { + y_grouped: &tensors.grouped_gate, + sorted_slot_index: &tensors.sorted, + gate_batch: &tensors.gate_batch, + up_batch: &tensors.up_batch, + inter: 64, + k_top: 2, + m_total: 8, + }, + Step::MoeActivation { + variant: MoeActivationVariant::SiluMul, + gate: &tensors.gate_batch, + up: &tensors.up_batch, + rot_out: &tensors.rot_batch, + inter: 64, + rows: 4, + }, + Step::GroupedMoeGemm { + experts, + which: MoeProj::DownExpanded, + sorted_slot_index: &tensors.sorted, + expert_tile_ids: &tensors.tiles, + x: &tensors.down_x, + y: &tensors.grouped_down, + m_total: 8, + batch_size: 2, + k_top: 2, + }, + Step::MoeCombine { + down_out: &tensors.grouped_down, + topk_weights: &tensors.weights, + out: &tensors.out, + hidden: 64, + k_top: 2, + batch_size: 2, + inverse_perm: Some(&tensors.inverse), + }, + ] + } + + fn indexed_steps<'a>( + experts: &'a MoeExpertRef<'a>, + scores: &'a GpuTensor, + indices: &'a GpuTensor, + weights: &'a GpuTensor, + x: &'a GpuTensor, + gate: &'a GpuTensor, + up: &'a GpuTensor, + rot: &'a GpuTensor, + down: &'a GpuTensor, + out: &'a GpuTensor, + ) -> Vec> { + vec![ + Step::MoeRoute { + plan: RouterPlan::SoftmaxTopK { + scores, + topk_indices: indices, + topk_weights: weights, + k_top: 2, + normalize: true, + }, + }, + Step::IndexedMoeGemv { + experts, + which: MoeProj::GateUp { up_out: up }, + topk_indices: indices, + input: GemvInput::Prerotated(x), + out: gate, + k_top: 2, + batch_size: 1, + }, + Step::MoeActivation { + variant: MoeActivationVariant::SiluMul, + gate, + up, + rot_out: rot, + inter: 64, + rows: 2, + }, + Step::IndexedMoeGemv { + experts, + which: MoeProj::DownExpanded, + topk_indices: indices, + input: GemvInput::Prerotated(rot), + out: down, + k_top: 2, + batch_size: 1, + }, + Step::MoeCombine { + down_out: down, + topk_weights: weights, + out, + hidden: 64, + k_top: 2, + batch_size: 1, + inverse_perm: None, + }, + ] + } + + #[test] + fn grouped_chain_seals_with_distinct_route_and_sorted_slot_buffers() { + let mesh = DeviceMesh::single().unwrap(); + let plan = resident_plan("grouped_quantized", ExpertParallelism::Single, &mesh); + let experts = plan.bind_expert_ref(0).unwrap(); + let tensors = grouped_tensors(); + let steps = grouped_steps(&experts, &tensors); + assert!(!std::ptr::eq(&tensors.indices, &tensors.sorted)); + + let schedule = MoeFamily::new() + .seal_steps( + ExpertExecutionPlan::GroupedQuantized, + steps, + vec![StepCollective::None; 7], + ) + .expect("grouped chain should seal"); + assert_eq!(schedule.execution(), ExpertExecutionPlan::GroupedQuantized); + assert_eq!(schedule.steps().len(), 7); + let signature = schedule.execution_signature().unwrap(); + assert_eq!(signature.protocol, MoeProtocolKind::Grouped); + match (&schedule.steps()[0], &schedule.steps()[1]) { + ( + Step::MoeRoute { plan }, + Step::MoeScatter { + topk_indices, + sorted_slot_index, + .. + }, + ) => { + assert!(std::ptr::eq(plan.route_buffers().0, *topk_indices)); + assert!(!std::ptr::eq(plan.route_buffers().0, *sorted_slot_index)); + } + _ => panic!("sealed schedule changed grouped route/scatter order"), + } + } + #[test] fn stride_assignment_and_named_group_are_deterministic() { let plan = ExpertPlan::from_manifest( @@ -1200,7 +1500,7 @@ mod tests { } #[test] - fn failed_load_rolls_back_and_unload_is_idempotent() { + fn rollback_then_reserve_reuse_and_repeated_teardown_are_safe() { let mut plan = ExpertPlan::from_manifest( &spec("indexed_quantized", ExpertParallelism::ExpertParallel), &manifest(), @@ -1211,11 +1511,18 @@ mod tests { { let mut load = plan.begin_load(); load.reserve(first).unwrap(); + load.rollback(); + assert_eq!( + load.reserve(first) + .expect_err("a rolled-back transaction is closed") + .to_string(), + "expert load transaction is already finished" + ); } assert_eq!(plan.resident_slots(), 0); { let mut load = plan.begin_load(); - load.reserve(first).unwrap(); + load.reserve(first).expect("rollback must release placement"); load.commit(); } assert_eq!(plan.resident_slots(), 1); @@ -1224,6 +1531,165 @@ mod tests { assert_eq!(plan.resident_slots(), 0); } + #[test] + fn nonresident_and_mismatched_rank_owner_fail_closed() { + let mesh = mesh_ep(); + let mut missing = ExpertPlan::from_manifest( + &spec("grouped_quantized", ExpertParallelism::ExpertParallel), + &manifest(), + &mesh, + ) + .unwrap(); + missing + .commit_rank_tables(0, table(vec![8]), table(vec![8]), None) + .unwrap(); + let error = missing + .bind_expert_ref(0) + .err() + .expect("nonresident owner must not bind"); + assert!(error.to_string().contains("nonresident")); + + let plan = resident_plan("grouped_quantized", ExpertParallelism::ExpertParallel, &mesh); + let experts = plan.bind_expert_ref(0).unwrap(); + assert_eq!(experts.owner_rank(), 0); + let tensors = grouped_tensors(); + let steps = grouped_steps(&experts, &tensors); + let mut collectives = vec![StepCollective::None; 7]; + collectives[6] = StepCollective::all_reduce( + DimKind::Ep, + 64, + vec![0, 1], + mesh.epoch(), + 1, + ); + let schedule = MoeFamily::new() + .seal_steps(ExpertExecutionPlan::GroupedQuantized, steps, collectives) + .expect("collective descriptor is locally typed"); + let error = hipfire_dispatch::pipeline::steps::validate_sealed_steps_mesh_preflight( + 2, + &mesh, + &[&schedule], + ) + .expect_err("rank-1 collective must not launch rank-0 owner view"); + assert!(error.to_string().contains("collective rank")); + } + + + #[test] + fn sealed_route_rejects_physical_buffer_shorter_than_logical_shape() { + let mesh = DeviceMesh::single().unwrap(); + let plan = resident_plan("grouped_quantized", ExpertParallelism::Single, &mesh); + let experts = plan.bind_expert_ref(0).unwrap(); + let mut tensors = grouped_tensors(); + tensors.indices = tensor_with_bytes(vec![2, 2], DType::F32, 3 * DType::F32.size()); + let steps = grouped_steps(&experts, &tensors); + let error = MoeFamily::new() + .seal_steps( + ExpertExecutionPlan::GroupedQuantized, + steps, + vec![StepCollective::None; 7], + ) + .expect_err("logical shape must not stand in for physical capacity"); + assert!(error.to_string().contains("insufficient logical/physical capacity")); + } + + #[test] + fn mesh_preflight_rejects_cross_rank_protocol_disagreement() { + let mesh = mesh_ep(); + let indexed_plan = + resident_plan("indexed_quantized", ExpertParallelism::ExpertParallel, &mesh); + let grouped_plan = + resident_plan("grouped_quantized", ExpertParallelism::ExpertParallel, &mesh); + let indexed_experts = indexed_plan.bind_expert_ref(0).unwrap(); + let grouped_experts = grouped_plan.bind_expert_ref(1).unwrap(); + + let indexed_scores = tensor(vec![4]); + let indexed_indices = tensor(vec![2]); + let indexed_weights = tensor(vec![2]); + let indexed_x = tensor(vec![64]); + let indexed_gate = tensor(vec![2, 64]); + let indexed_up = tensor(vec![2, 64]); + let indexed_rot = tensor(vec![2, 64]); + let indexed_down = tensor(vec![2, 64]); + let indexed_out = tensor(vec![64]); + let indexed = MoeFamily::new() + .seal_steps( + ExpertExecutionPlan::IndexedQuantized, + indexed_steps( + &indexed_experts, + &indexed_scores, + &indexed_indices, + &indexed_weights, + &indexed_x, + &indexed_gate, + &indexed_up, + &indexed_rot, + &indexed_down, + &indexed_out, + ), + { + let mut collectives = vec![StepCollective::None; 5]; + collectives[4] = + StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); + collectives + }, + ) + .expect("indexed rank schedule should seal"); + + let grouped_tensors = grouped_tensors(); + let grouped = MoeFamily::new() + .seal_steps( + ExpertExecutionPlan::GroupedQuantized, + grouped_steps(&grouped_experts, &grouped_tensors), + { + let mut collectives = vec![StepCollective::None; 7]; + collectives[6] = + StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 1); + collectives + }, + ) + .expect("grouped rank schedule should seal"); + + let error = hipfire_dispatch::pipeline::steps::validate_sealed_steps_mesh_preflight( + 2, + &mesh, + &[&indexed, &grouped], + ) + .expect_err("mixed indexed/grouped ranks must not launch"); + assert!(error.to_string().contains("executable identity")); + } + + #[test] + fn parallel_sealing_rejects_duplicate_or_mixed_collectives() { + let mesh = mesh_ep(); + let plan = resident_plan("grouped_quantized", ExpertParallelism::ExpertParallel, &mesh); + let experts = plan.bind_expert_ref(0).unwrap(); + let tensors = grouped_tensors(); + let family = MoeFamily::new(); + + let mut duplicate = vec![StepCollective::None; 7]; + duplicate[5] = + StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); + duplicate[6] = + StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); + let error = family + .seal_steps( + ExpertExecutionPlan::GroupedQuantized, + grouped_steps(&experts, &tensors), + duplicate, + ) + .expect_err("a routed reduction cannot appear twice"); + assert!(error.to_string().contains("attached to combine")); + + let error = family + .seal_steps( + ExpertExecutionPlan::GroupedQuantized, + grouped_steps(&experts, &tensors), + vec![StepCollective::None; 7], + ) + .expect_err("parallel owner must not silently use identity reduction"); + assert!(error.to_string().contains("collective count")); + } #[test] fn per_expert_fallback_is_refused_at_plan_boundary() { let error = ExpertPlan::from_manifest( @@ -1247,6 +1713,17 @@ mod tests { .expect_err("gate/up and down shapes must agree"); assert!(error.to_string().contains("shape mismatch")); } + #[test] + fn separate_layout_fails_closed_at_plan_boundary() { + let mesh = DeviceMesh::single().unwrap(); + let error = ExpertPlan::from_manifest( + &separate_spec("grouped_quantized", ExpertParallelism::Single), + &separate_manifest(), + &mesh, + ) + .expect_err("generic executor must refuse separate projection sources"); + assert!(error.to_string().contains("separate gate/up sources")); + } #[test] fn single_plan_emits_identity_collective() { @@ -1256,6 +1733,7 @@ mod tests { let mesh = DeviceMesh::single().unwrap(); let plan = ExpertPlan::from_manifest( &spec("indexed_quantized", ExpertParallelism::Single), + &single_manifest, &mesh, ) From f70e9321716142d3eebec531db699dbc03962969 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 11:01:59 +0200 Subject: [PATCH 20/25] style(device-mesh): format MoE substrate --- crates/hipfire-dispatch/src/families/moe.rs | 54 +-- crates/hipfire-dispatch/src/pipeline/steps.rs | 352 ++++++++---------- crates/hipfire-runtime/src/lib.rs | 2 +- crates/hipfire-runtime/src/moe_plan.rs | 161 ++++---- 4 files changed, 281 insertions(+), 288 deletions(-) diff --git a/crates/hipfire-dispatch/src/families/moe.rs b/crates/hipfire-dispatch/src/families/moe.rs index c0cf03850b..b3b765053d 100644 --- a/crates/hipfire-dispatch/src/families/moe.rs +++ b/crates/hipfire-dispatch/src/families/moe.rs @@ -61,7 +61,9 @@ fn checked_numel(tensor: &GpuTensor, name: &str) -> Result tensor .shape .iter() - .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) + .try_fold(1usize, |elements, dimension| { + elements.checked_mul(*dimension) + }) .ok_or_else(|| DispatchError::Hip(format!("MoE {name} logical shape overflows"))) } @@ -181,7 +183,9 @@ impl<'a> RouterPlan<'a> { .ok_or_else(|| DispatchError::Hip("MoE router score capacity overflow".into()))?; let score_bytes = score_capacity .checked_mul(DType::F32.size()) - .ok_or_else(|| DispatchError::Hip("MoE router score byte capacity overflow".into()))?; + .ok_or_else(|| { + DispatchError::Hip("MoE router score byte capacity overflow".into()) + })?; if scores.dtype != DType::F32 || !score_shape_ok || score_elements < score_capacity @@ -408,7 +412,6 @@ impl<'a> MoeExpertRef<'a> { self.router_identity } - pub fn collective_kind(&self) -> Option { self.collective_kind } @@ -443,13 +446,14 @@ impl<'a> MoeExpertRef<'a> { "MoeExpertRef: router identity is empty".into(), )); } - let pointer_slots = self - .n_experts - .checked_mul(2) - .ok_or_else(|| DispatchError::Hip("MoeExpertRef: pointer-table size overflows".into()))?; + let pointer_slots = self.n_experts.checked_mul(2).ok_or_else(|| { + DispatchError::Hip("MoeExpertRef: pointer-table size overflows".into()) + })?; let pointer_bytes = pointer_slots .checked_mul(DType::F32.size()) - .ok_or_else(|| DispatchError::Hip("MoeExpertRef: pointer-table bytes overflow".into()))?; + .ok_or_else(|| { + DispatchError::Hip("MoeExpertRef: pointer-table bytes overflow".into()) + })?; if self.gate_up_ptrs.dtype != DType::F32 || self.down_ptrs.dtype != DType::F32 || self.gate_up_ptrs.shape.as_slice() != [pointer_slots] @@ -466,11 +470,17 @@ impl<'a> MoeExpertRef<'a> { let dummy_elements = dummy .shape .iter() - .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) - .ok_or_else(|| DispatchError::Hip("MoeExpertRef: dummy table shape overflows".into()))?; + .try_fold(1usize, |elements, dimension| { + elements.checked_mul(*dimension) + }) + .ok_or_else(|| { + DispatchError::Hip("MoeExpertRef: dummy table shape overflows".into()) + })?; let dummy_bytes = dummy_elements .checked_mul(DType::F32.size()) - .ok_or_else(|| DispatchError::Hip("MoeExpertRef: dummy table bytes overflow".into()))?; + .ok_or_else(|| { + DispatchError::Hip("MoeExpertRef: dummy table bytes overflow".into()) + })?; if dummy.dtype != DType::F32 || dummy_elements == 0 || dummy.buf.size() < dummy_bytes { return Err(DispatchError::Hip( "MoeExpertRef: dummy gate/up table is invalid".into(), @@ -482,9 +492,12 @@ impl<'a> MoeExpertRef<'a> { "MoeExpertRef: owner rank is outside its mesh group".into(), )); } - if self.group_devices.iter().enumerate().any(|(index, device)| { - self.group_devices[..index].contains(device) - }) { + if self + .group_devices + .iter() + .enumerate() + .any(|(index, device)| self.group_devices[..index].contains(device)) + { return Err(DispatchError::Hip( "MoeExpertRef: mesh group contains duplicate devices".into(), )); @@ -667,7 +680,6 @@ impl MoeDtypes { } } - /// Resolved fused-vs-fallback eligibility for one MoE decode layer. This IS the /// routing-config logic, relocated from `moe_ffn_decode_impl` into one typed, /// testable place (review finding #1). Pure function of `MoeDtypes` + k. @@ -1757,14 +1769,7 @@ impl MoeFamily { batch_size, ) } else { - gpu.moe_down_combine_k8_batched( - down_out, - topk_weights, - out, - hidden, - k_top, - batch_size, - ) + gpu.moe_down_combine_k8_batched(down_out, topk_weights, out, hidden, k_top, batch_size) }; result.map_err(|error| DispatchError::Hip(error.to_string())) } @@ -2075,7 +2080,7 @@ mod tests { assert!(r.routed_indexable_mq6v2); assert!(!r.use_gpu_topk); } - + fn tensor(shape: Vec) -> GpuTensor { let elements = shape.iter().product::(); let mut tensor = GpuTensor::null_for_test(); @@ -2108,7 +2113,6 @@ mod tests { plan.validate_against(8, 1).unwrap(); } - #[test] fn fallback_execution_has_no_protocol() { assert!(ExpertExecutionPlan::PerExpertFallback.protocol().is_err()); diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index 9ea9e83b70..d424fc4ef0 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -12,7 +12,7 @@ use crate::context::DispatchCtx; use crate::families::fused_qkv::{FusedQkvBiasParams, FusedQkvFamily, FusedQkvParams}; use crate::families::gemv::{GemvFamily, GemvParams, RotateInputs, WeightRef}; use crate::families::moe::{ - ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeProj, MoeProtocolKind, MoeFamily, + ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeFamily, MoeProj, MoeProtocolKind, RouterPlan, }; use crate::families::rotation::{RotationFamily, RotationParams}; @@ -91,9 +91,7 @@ pub enum Step<'a> { }, /// Typed MoE route. Routing semantics (bias/hash/normalization) are /// carried by the plan rather than reconstructed by the family. - MoeRoute { - plan: RouterPlan<'a>, - }, + MoeRoute { plan: RouterPlan<'a> }, /// Indexed routed expert projection. Every down projection is expanded; /// the executor-owned `MoeCombine` is the only weighted reduction. IndexedMoeGemv { @@ -182,7 +180,7 @@ fn op_kind(step: &Step) -> PipelineOp { Step::MoeActivation { .. } => PipelineOp::MoeActivation, } } - + /// Collective attached to one lock-step `Step` position. The descriptor is /// immutable schedule data: membership and mesh identity are supplied by the /// manifest/topology owner, never inferred by a family. @@ -272,44 +270,38 @@ pub fn validate_moe_step_schedule( let (protocol, route, experts, batch_size, combine_index, hidden, route_indices, route_weights) = match steps { - [ - Step::MoeRoute { plan }, - Step::IndexedMoeGemv { - experts, - which: MoeProj::GateUp { up_out }, - topk_indices, - input: GemvInput::Prerotated(x), - out: gate, - k_top, - batch_size, - }, - Step::MoeActivation { - variant, - gate: act_gate, - up, - rot_out, - inter, - rows, - }, - Step::IndexedMoeGemv { - experts: down_experts, - which: MoeProj::DownExpanded, - topk_indices: down_indices, - input: GemvInput::Prerotated(down_input), - out: down_out, - k_top: down_k, - batch_size: down_batch, - }, - Step::MoeCombine { - down_out: combine_down, - topk_weights, - out, - hidden, - k_top: combine_k, - batch_size: combine_batch, - inverse_perm, - }, - ] => { + [Step::MoeRoute { plan }, Step::IndexedMoeGemv { + experts, + which: MoeProj::GateUp { up_out }, + topk_indices, + input: GemvInput::Prerotated(x), + out: gate, + k_top, + batch_size, + }, Step::MoeActivation { + variant, + gate: act_gate, + up, + rot_out, + inter, + rows, + }, Step::IndexedMoeGemv { + experts: down_experts, + which: MoeProj::DownExpanded, + topk_indices: down_indices, + input: GemvInput::Prerotated(down_input), + out: down_out, + k_top: down_k, + batch_size: down_batch, + }, Step::MoeCombine { + down_out: combine_down, + topk_weights, + out, + hidden, + k_top: combine_k, + batch_size: combine_batch, + inverse_perm, + }] => { let expected_rows = batch_size .checked_mul(*k_top) .ok_or_else(|| DispatchError::Hip("indexed MoE row count overflows".into()))?; @@ -360,69 +352,61 @@ pub fn validate_moe_step_schedule( topk_weights, ) } - [ - Step::MoeRoute { plan }, - Step::MoeScatter { - topk_indices, - expert_token_counts, - expert_offsets, - sorted_slot_index, - expert_tile_ids, - inverse_perm, - total_slots, - n_experts, - m_total_max, - block_m, - }, - Step::GroupedMoeGemm { - experts, - which: MoeProj::GateUp { .. }, - sorted_slot_index: gate_sorted, - expert_tile_ids: gate_tiles, - x: gate_x, - y: gate_y, - m_total: gate_m_total, - batch_size, - k_top, - }, - Step::MoeGateUpUnscatter { - y_grouped, - sorted_slot_index: unscatter_sorted, - gate_batch, - up_batch, - inter, - k_top: unscatter_k, - m_total: unscatter_m_total, - }, - Step::MoeActivation { - variant: _, - gate: act_gate, - up: act_up, - rot_out, - inter: act_inter, - rows, - }, - Step::GroupedMoeGemm { - experts: down_experts, - which: MoeProj::DownExpanded, - sorted_slot_index: down_sorted, - expert_tile_ids: down_tiles, - x: down_x, - y: down_y, - m_total: down_m_total, - batch_size: down_batch, - k_top: down_k, - }, - Step::MoeCombine { - down_out: combine_down, - topk_weights, - out, - hidden, - k_top: combine_k, - batch_size: combine_batch, - inverse_perm: combine_inverse, - }, - ] => { + [Step::MoeRoute { plan }, Step::MoeScatter { + topk_indices, + expert_token_counts, + expert_offsets, + sorted_slot_index, + expert_tile_ids, + inverse_perm, + total_slots, + n_experts, + m_total_max, + block_m, + }, Step::GroupedMoeGemm { + experts, + which: MoeProj::GateUp { .. }, + sorted_slot_index: gate_sorted, + expert_tile_ids: gate_tiles, + x: gate_x, + y: gate_y, + m_total: gate_m_total, + batch_size, + k_top, + }, Step::MoeGateUpUnscatter { + y_grouped, + sorted_slot_index: unscatter_sorted, + gate_batch, + up_batch, + inter, + k_top: unscatter_k, + m_total: unscatter_m_total, + }, Step::MoeActivation { + variant: _, + gate: act_gate, + up: act_up, + rot_out, + inter: act_inter, + rows, + }, Step::GroupedMoeGemm { + experts: down_experts, + which: MoeProj::DownExpanded, + sorted_slot_index: down_sorted, + expert_tile_ids: down_tiles, + x: down_x, + y: down_y, + m_total: down_m_total, + batch_size: down_batch, + k_top: down_k, + }, Step::MoeCombine { + down_out: combine_down, + topk_weights, + out, + hidden, + k_top: combine_k, + batch_size: combine_batch, + inverse_perm: combine_inverse, + }] => { let expected_slots = batch_size .checked_mul(*k_top) .ok_or_else(|| DispatchError::Hip("grouped MoE slot count overflows".into()))?; @@ -464,9 +448,9 @@ pub fn validate_moe_step_schedule( "grouped MoE scatter capacity or tile geometry is invalid".into(), )); } - let expected_offsets = n_experts - .checked_add(1) - .ok_or_else(|| DispatchError::Hip("MoE expert offset count overflows".into()))?; + let expected_offsets = n_experts.checked_add(1).ok_or_else(|| { + DispatchError::Hip("MoE expert offset count overflows".into()) + })?; let counts_elements = checked_numel(expert_token_counts, "expert counts")?; let offset_elements = checked_numel(expert_offsets, "expert offsets")?; if counts_elements == 0 || offset_elements < expected_offsets { @@ -558,7 +542,12 @@ pub fn validate_moe_step_schedule( ))); } validate_step_tensors(steps, experts, batch_size, hidden)?; - validate_collectives(collectives, combine_index, hidden, experts.collective_kind())?; + validate_collectives( + collectives, + combine_index, + hidden, + experts.collective_kind(), + )?; Ok(()) } @@ -598,8 +587,9 @@ fn derive_moe_execution_signature<'a>( let experts = steps .iter() .find_map(|step| match step { - Step::IndexedMoeGemv { experts, .. } - | Step::GroupedMoeGemm { experts, .. } => Some(*experts), + Step::IndexedMoeGemv { experts, .. } | Step::GroupedMoeGemm { experts, .. } => { + Some(*experts) + } _ => None, }) .ok_or_else(|| DispatchError::Hip("MoE schedule has no expert owner view".into()))?; @@ -634,7 +624,6 @@ pub struct SealedMoeSchedule<'a> { collectives: Vec, } - impl<'a> SealedMoeSchedule<'a> { pub fn new( execution: ExpertExecutionPlan, @@ -738,9 +727,7 @@ fn preflight_sealed_steps_mesh<'a>( StepCollective::None => None, }) .ok_or_else(|| { - DispatchError::Hip( - "parallel MoE schedule has no manifest-owned collective".into(), - ) + DispatchError::Hip("parallel MoE schedule has no manifest-owned collective".into()) })?; if descriptor_rank != rank { return Err(DispatchError::Hip( @@ -751,8 +738,9 @@ fn preflight_sealed_steps_mesh<'a>( .steps() .iter() .find_map(|step| match step { - Step::IndexedMoeGemv { experts, .. } - | Step::GroupedMoeGemm { experts, .. } => Some(*experts), + Step::IndexedMoeGemv { experts, .. } | Step::GroupedMoeGemm { experts, .. } => { + Some(*experts) + } _ => None, }) .ok_or_else(|| DispatchError::Hip("MoE schedule has no expert owner view".into()))?; @@ -802,9 +790,11 @@ fn preflight_sealed_steps_mesh<'a>( "MoE collective rank count does not match mesh group".into(), )); } - if group.iter().enumerate().any(|(index, device)| { - *device >= mesh.n_devices() || group[..index].contains(device) - }) { + if group + .iter() + .enumerate() + .any(|(index, device)| *device >= mesh.n_devices() || group[..index].contains(device)) + { return Err(DispatchError::Hip( "MoE collective group contains invalid or duplicate devices".into(), )); @@ -860,15 +850,16 @@ pub fn execute_sealed_steps_mesh<'a>( } fn same_tensor(a: &GpuTensor, b: &GpuTensor) -> bool { - std::ptr::eq(a, b) - || (!a.buf.as_ptr().is_null() && a.buf.as_ptr() == b.buf.as_ptr()) + std::ptr::eq(a, b) || (!a.buf.as_ptr().is_null() && a.buf.as_ptr() == b.buf.as_ptr()) } fn checked_numel(tensor: &GpuTensor, name: &str) -> Result { tensor .shape .iter() - .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) + .try_fold(1usize, |elements, dimension| { + elements.checked_mul(*dimension) + }) .ok_or_else(|| DispatchError::Hip(format!("MoE {name} logical shape overflows"))) } @@ -882,10 +873,7 @@ fn require_tensor( let required_bytes = capacity .checked_mul(dtype.size()) .ok_or_else(|| DispatchError::Hip(format!("MoE {name} byte capacity overflows")))?; - if tensor.dtype != dtype - || logical_elements < capacity - || tensor.buf.size() < required_bytes - { + if tensor.dtype != dtype || logical_elements < capacity || tensor.buf.size() < required_bytes { return Err(DispatchError::Hip(format!( "MoE {name} has dtype {:?}/logical capacity {logical_elements}/physical bytes {}, \ expected {:?}/{capacity} elements/{required_bytes} bytes", @@ -897,19 +885,12 @@ fn require_tensor( Ok(()) } -fn require_raw_i32( - tensor: &GpuTensor, - name: &str, - elements: usize, -) -> Result<(), DispatchError> { +fn require_raw_i32(tensor: &GpuTensor, name: &str, elements: usize) -> Result<(), DispatchError> { let bytes = elements .checked_mul(std::mem::size_of::()) .ok_or_else(|| DispatchError::Hip(format!("MoE {name} capacity overflows")))?; let logical_bytes = checked_numel(tensor, name)?; - if tensor.dtype != DType::Raw - || logical_bytes < bytes - || tensor.buf.size() < bytes - { + if tensor.dtype != DType::Raw || logical_bytes < bytes || tensor.buf.size() < bytes { return Err(DispatchError::Hip(format!( "MoE {name} has dtype {:?}/logical bytes {logical_bytes}/physical bytes {}, \ expected Raw/at least {bytes} bytes", @@ -920,7 +901,6 @@ fn require_raw_i32( Ok(()) } - fn validate_step_tensors( steps: &[Step], experts: &MoeExpertRef<'_>, @@ -947,18 +927,8 @@ fn validate_step_tensors( _ => None, }) .ok_or_else(|| DispatchError::Hip("MoE grammar has no route".into()))?; - require_tensor( - route.route_buffers().0, - "route indices", - DType::F32, - slots, - )?; - require_tensor( - route.route_buffers().1, - "route weights", - DType::F32, - slots, - )?; + require_tensor(route.route_buffers().0, "route indices", DType::F32, slots)?; + require_tensor(route.route_buffers().1, "route weights", DType::F32, slots)?; let gate_capacity = slots .checked_mul(experts.expert_m()) .ok_or_else(|| DispatchError::Hip("MoE gate capacity overflow".into()))?; @@ -985,11 +955,9 @@ fn validate_step_tensors( x, "indexed gate/up input", DType::F32, - batch_size - .checked_mul(experts.expert_k()) - .ok_or_else(|| { - DispatchError::Hip("MoE indexed input capacity overflow".into()) - })?, + batch_size.checked_mul(experts.expert_k()).ok_or_else(|| { + DispatchError::Hip("MoE indexed input capacity overflow".into()) + })?, )?; require_tensor(out, "indexed gate output", DType::F32, gate_capacity)?; require_tensor(up_out, "indexed up output", DType::F32, gate_capacity)?; @@ -1035,14 +1003,11 @@ fn validate_step_tensors( batch_size .checked_mul(k_top) .and_then(|slots| slots.checked_mul(*hidden)) - .ok_or_else(|| DispatchError::Hip("MoE combine input capacity overflows".into()))?, - )?; - require_tensor( - topk_weights, - "combine weights", - DType::F32, - slots, + .ok_or_else(|| { + DispatchError::Hip("MoE combine input capacity overflows".into()) + })?, )?; + require_tensor(topk_weights, "combine weights", DType::F32, slots)?; require_tensor( out, "combine output", @@ -1076,11 +1041,7 @@ fn validate_step_tensors( require_raw_i32(expert_offsets, "expert offsets", offsets)?; require_raw_i32(sorted_slot_index, "sorted slots", *m_total_max)?; require_raw_i32(inverse_perm, "inverse permutation", *total_slots)?; - require_raw_i32( - expert_tile_ids, - "expert tile ids", - *m_total_max / *block_m, - )?; + require_raw_i32(expert_tile_ids, "expert tile ids", *m_total_max / *block_m)?; } Step::GroupedMoeGemm { which, @@ -1090,24 +1051,24 @@ fn validate_step_tensors( .. } => { let input_capacity = match which { - MoeProj::GateUp { .. } => batch_size - .checked_mul(experts.expert_k()) - .ok_or_else(|| { + MoeProj::GateUp { .. } => { + batch_size.checked_mul(experts.expert_k()).ok_or_else(|| { DispatchError::Hip("MoE grouped input capacity overflows".into()) - })?, - MoeProj::DownExpanded => slots - .checked_mul(experts.expert_m()) - .ok_or_else(|| { + })? + } + MoeProj::DownExpanded => { + slots.checked_mul(experts.expert_m()).ok_or_else(|| { DispatchError::Hip("MoE grouped input capacity overflows".into()) - })?, + })? + } }; require_tensor(x, "grouped input", DType::F32, input_capacity)?; let output_width = match which { - MoeProj::GateUp { .. } => 2usize - .checked_mul(experts.expert_m()) - .ok_or_else(|| { + MoeProj::GateUp { .. } => { + 2usize.checked_mul(experts.expert_m()).ok_or_else(|| { DispatchError::Hip("MoE grouped output width overflows".into()) - })?, + })? + } MoeProj::DownExpanded => experts.expert_k(), }; require_tensor( @@ -1132,12 +1093,12 @@ fn validate_step_tensors( "unscatter grouped input", DType::F32, m_total - .checked_mul( - 2usize.checked_mul(*inter).ok_or_else(|| { - DispatchError::Hip("MoE unscatter width overflows".into()) - })?, - ) - .ok_or_else(|| DispatchError::Hip("MoE unscatter input overflows".into()))?, + .checked_mul(2usize.checked_mul(*inter).ok_or_else(|| { + DispatchError::Hip("MoE unscatter width overflows".into()) + })?) + .ok_or_else(|| { + DispatchError::Hip("MoE unscatter input overflows".into()) + })?, )?; require_tensor( gate_batch, @@ -1164,16 +1125,13 @@ fn validate_step_tensors( rows, .. } => { - let capacity = rows.checked_mul(*inter).ok_or_else(|| { - DispatchError::Hip("MoE activation capacity overflow".into()) - })?; + let capacity = rows + .checked_mul(*inter) + .ok_or_else(|| DispatchError::Hip("MoE activation capacity overflow".into()))?; require_tensor(gate, "activation gate", DType::F32, capacity)?; require_tensor(up, "activation up", DType::F32, capacity)?; require_tensor(rot_out, "activation output", DType::F32, capacity)?; - if same_tensor(gate, up) - || same_tensor(gate, rot_out) - || same_tensor(up, rot_out) - { + if same_tensor(gate, up) || same_tensor(gate, rot_out) || same_tensor(up, rot_out) { return Err(DispatchError::Hip( "MoE activation buffers must not alias".into(), )); @@ -1219,9 +1177,11 @@ fn validate_collectives( "MoE collective rank/group/output dimension is invalid".into(), )); } - if group.iter().enumerate().any(|(offset, device)| { - group[..offset].contains(device) - }) { + if group + .iter() + .enumerate() + .any(|(offset, device)| group[..offset].contains(device)) + { return Err(DispatchError::Hip( "MoE collective group contains duplicate devices".into(), )); @@ -2844,6 +2804,4 @@ mod tests { "FusedGateUpQ8_0 missing from FUSED_TABLE" ); } - - } diff --git a/crates/hipfire-runtime/src/lib.rs b/crates/hipfire-runtime/src/lib.rs index 333480dc7d..f5821cca35 100644 --- a/crates/hipfire-runtime/src/lib.rs +++ b/crates/hipfire-runtime/src/lib.rs @@ -46,8 +46,8 @@ pub mod llama_spec; pub mod loader_api; pub mod loop_guard; pub mod model_load; -pub mod moe_plan; pub mod model_source; +pub mod moe_plan; pub mod paro; pub mod prefix; pub mod reset_core; diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs index 16681e7530..a01125d86e 100644 --- a/crates/hipfire-runtime/src/moe_plan.rs +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -78,7 +78,9 @@ impl ExpertStorageOwner { } fn index_of(&self, placement: ExpertPlacement) -> Option { - self.slots.iter().position(|candidate| *candidate == placement) + self.slots + .iter() + .position(|candidate| *candidate == placement) } fn rank_is_resident(&self, rank: usize) -> bool { @@ -98,7 +100,6 @@ impl ExpertStorageOwner { } } - /// A sealed, manifest-derived expert plan. /// /// Every placement, source identity, shape, resource requirement, rank-local @@ -308,13 +309,16 @@ impl ExpertPlan { } pub fn owned_experts(&self, rank: usize) -> Result<&[usize], ExpertPlanError> { - self.owner_views.get(rank).map(Vec::as_slice).ok_or_else(|| { - ExpertPlanError::new(format!( - "expert group '{}' rank {rank} is outside group size {}", - self.group, - self.group_size() - )) - }) + self.owner_views + .get(rank) + .map(Vec::as_slice) + .ok_or_else(|| { + ExpertPlanError::new(format!( + "expert group '{}' rank {rank} is outside group size {}", + self.group, + self.group_size() + )) + }) } /// Commit the rank-local pointer tables to this owner. The plan validates @@ -360,10 +364,7 @@ impl ExpertPlan { /// Bind the committed rank-local pointer tables to an opaque executable /// view. A rank cannot bind until its owned placements are resident. - pub fn bind_expert_ref<'a>( - &'a self, - rank: usize, - ) -> Result, ExpertPlanError> { + pub fn bind_expert_ref<'a>(&'a self, rank: usize) -> Result, ExpertPlanError> { let owned = self.owned_experts(rank)?; if !self.owner.rank_is_resident(rank) { return Err(ExpertPlanError::new(format!( @@ -757,12 +758,13 @@ fn source_names<'a>(layout: &'a ExpertSourceLayout) -> Vec<(&'static str, Vec<&' } } - fn checked_numel(tensor: &GpuTensor, label: &str) -> Result { tensor .shape .iter() - .try_fold(1usize, |elements, dimension| elements.checked_mul(*dimension)) + .try_fold(1usize, |elements, dimension| { + elements.checked_mul(*dimension) + }) .ok_or_else(|| ExpertPlanError::new(format!("{label} logical shape overflows"))) } @@ -833,7 +835,11 @@ fn check_shape( ))); } let shape = &entry.logical_shape; - let valid_rank = if per_expert { shape.len() == 2 } else { shape.len() == 3 }; + let valid_rank = if per_expert { + shape.len() == 2 + } else { + shape.len() == 3 + }; if !valid_rank || shape.iter().any(|dim| *dim == 0) { return Err(ExpertPlanError::new(format!( "expert group '{}' layer {:?} {label} source '{}' has invalid logical_shape {:?}", @@ -889,9 +895,7 @@ fn resolve_shape( )?], true, ), - ExpertSourceLayout::PackedSeparate { - gate, up, down, .. - } => ( + ExpertSourceLayout::PackedSeparate { gate, up, down, .. } => ( vec![check_shape( spec, "gate", @@ -915,31 +919,44 @@ fn resolve_shape( ExpertSourceLayout::PerExpertFused { gate_up, down, .. } => ( gate_up .iter() - .map(|name| check_shape(spec, "gate_up", entry_for(spec, manifest, "gate_up", name)?, true)) + .map(|name| { + check_shape( + spec, + "gate_up", + entry_for(spec, manifest, "gate_up", name)?, + true, + ) + }) .collect::>()?, Vec::new(), - down - .iter() - .map(|name| check_shape(spec, "down", entry_for(spec, manifest, "down", name)?, true)) + down.iter() + .map(|name| { + check_shape(spec, "down", entry_for(spec, manifest, "down", name)?, true) + }) .collect::>()?, true, ), ExpertSourceLayout::PerExpertSeparate { gate, up, down, .. } => ( - gate - .iter() - .map(|name| check_shape(spec, "gate", entry_for(spec, manifest, "gate", name)?, true)) + gate.iter() + .map(|name| { + check_shape(spec, "gate", entry_for(spec, manifest, "gate", name)?, true) + }) .collect::>()?, up.iter() .map(|name| check_shape(spec, "up", entry_for(spec, manifest, "up", name)?, true)) .collect::>()?, - down - .iter() - .map(|name| check_shape(spec, "down", entry_for(spec, manifest, "down", name)?, true)) + down.iter() + .map(|name| { + check_shape(spec, "down", entry_for(spec, manifest, "down", name)?, true) + }) .collect::>()?, false, ), }; - if gate_shapes.is_empty() || down_shapes.is_empty() || (!up_shapes.is_empty() && up_shapes[0] != gate_shapes[0]) { + if gate_shapes.is_empty() + || down_shapes.is_empty() + || (!up_shapes.is_empty() && up_shapes[0] != gate_shapes[0]) + { return Err(ExpertPlanError::new(format!( "expert group '{}' layer {:?} projection source count/shape mismatch", spec.group, spec.layer @@ -988,7 +1005,12 @@ fn resolve_shape( let source_dtype = source_names .iter() .flat_map(|(_, names)| names.iter()) - .find_map(|name| manifest.iter().find(|entry| entry.name == *name && entry.layer == spec.layer).map(|entry| entry.dtype)) + .find_map(|name| { + manifest + .iter() + .find(|entry| entry.name == *name && entry.layer == spec.layer) + .map(|entry| entry.dtype) + }) .ok_or_else(|| ExpertPlanError::new("expert projection has no source dtype"))?; for (_, names) in source_names { for name in names { @@ -1077,7 +1099,12 @@ fn resolve_collective( selected = Some(row.hint); } } - Ok((selected, down_names(&spec.source_layout).first().map(|name| (*name).to_string()))) + Ok(( + selected, + down_names(&spec.source_layout) + .first() + .map(|name| (*name).to_string()), + )) } #[cfg(test)] @@ -1094,13 +1121,7 @@ mod tests { fn manifest() -> Vec { vec![ - WeightEntry::layer( - "router", - 0, - vec![4, 4], - DType::F32, - ShardPolicy::Replicate, - ), + WeightEntry::layer("router", 0, vec![4, 4], DType::F32, ShardPolicy::Replicate), WeightEntry::layer( "experts.gate_up", 0, @@ -1185,9 +1206,7 @@ mod tests { fn tensor_with_bytes(shape: Vec, dtype: DType, bytes: usize) -> GpuTensor { let mut tensor = GpuTensor::null_for_test(); - tensor.buf = unsafe { - hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut(), bytes) - }; + tensor.buf = unsafe { hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut(), bytes) }; tensor.shape = shape; tensor.dtype = dtype; tensor @@ -1198,7 +1217,9 @@ mod tests { tensor_with_bytes( shape, DType::F32, - elements.checked_mul(DType::F32.size()).expect("test tensor bytes"), + elements + .checked_mul(DType::F32.size()) + .expect("test tensor bytes"), ) } @@ -1467,7 +1488,10 @@ mod tests { .collect(); assert_eq!(owners, vec![(0, 0, 0), (1, 1, 0), (2, 0, 1), (3, 1, 1)]); assert_eq!(plan.group_devices(), &[0, 1]); - assert_eq!(plan.collective(), Some(CollectiveHint::AllReduce { kind: DimKind::Ep })); + assert_eq!( + plan.collective(), + Some(CollectiveHint::AllReduce { kind: DimKind::Ep }) + ); } #[test] @@ -1522,7 +1546,8 @@ mod tests { assert_eq!(plan.resident_slots(), 0); { let mut load = plan.begin_load(); - load.reserve(first).expect("rollback must release placement"); + load.reserve(first) + .expect("rollback must release placement"); load.commit(); } assert_eq!(plan.resident_slots(), 1); @@ -1549,19 +1574,17 @@ mod tests { .expect("nonresident owner must not bind"); assert!(error.to_string().contains("nonresident")); - let plan = resident_plan("grouped_quantized", ExpertParallelism::ExpertParallel, &mesh); + let plan = resident_plan( + "grouped_quantized", + ExpertParallelism::ExpertParallel, + &mesh, + ); let experts = plan.bind_expert_ref(0).unwrap(); assert_eq!(experts.owner_rank(), 0); let tensors = grouped_tensors(); let steps = grouped_steps(&experts, &tensors); let mut collectives = vec![StepCollective::None; 7]; - collectives[6] = StepCollective::all_reduce( - DimKind::Ep, - 64, - vec![0, 1], - mesh.epoch(), - 1, - ); + collectives[6] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 1); let schedule = MoeFamily::new() .seal_steps(ExpertExecutionPlan::GroupedQuantized, steps, collectives) .expect("collective descriptor is locally typed"); @@ -1574,7 +1597,6 @@ mod tests { assert!(error.to_string().contains("collective rank")); } - #[test] fn sealed_route_rejects_physical_buffer_shorter_than_logical_shape() { let mesh = DeviceMesh::single().unwrap(); @@ -1590,16 +1612,24 @@ mod tests { vec![StepCollective::None; 7], ) .expect_err("logical shape must not stand in for physical capacity"); - assert!(error.to_string().contains("insufficient logical/physical capacity")); + assert!(error + .to_string() + .contains("insufficient logical/physical capacity")); } #[test] fn mesh_preflight_rejects_cross_rank_protocol_disagreement() { let mesh = mesh_ep(); - let indexed_plan = - resident_plan("indexed_quantized", ExpertParallelism::ExpertParallel, &mesh); - let grouped_plan = - resident_plan("grouped_quantized", ExpertParallelism::ExpertParallel, &mesh); + let indexed_plan = resident_plan( + "indexed_quantized", + ExpertParallelism::ExpertParallel, + &mesh, + ); + let grouped_plan = resident_plan( + "grouped_quantized", + ExpertParallelism::ExpertParallel, + &mesh, + ); let indexed_experts = indexed_plan.bind_expert_ref(0).unwrap(); let grouped_experts = grouped_plan.bind_expert_ref(1).unwrap(); @@ -1662,16 +1692,18 @@ mod tests { #[test] fn parallel_sealing_rejects_duplicate_or_mixed_collectives() { let mesh = mesh_ep(); - let plan = resident_plan("grouped_quantized", ExpertParallelism::ExpertParallel, &mesh); + let plan = resident_plan( + "grouped_quantized", + ExpertParallelism::ExpertParallel, + &mesh, + ); let experts = plan.bind_expert_ref(0).unwrap(); let tensors = grouped_tensors(); let family = MoeFamily::new(); let mut duplicate = vec![StepCollective::None; 7]; - duplicate[5] = - StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); - duplicate[6] = - StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); + duplicate[5] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); + duplicate[6] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); let error = family .seal_steps( ExpertExecutionPlan::GroupedQuantized, @@ -1733,7 +1765,6 @@ mod tests { let mesh = DeviceMesh::single().unwrap(); let plan = ExpertPlan::from_manifest( &spec("indexed_quantized", ExpertParallelism::Single), - &single_manifest, &mesh, ) From c7bb7d3b01435d08091bdc54ae0b163e58ab637b Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 11:10:59 +0200 Subject: [PATCH 21/25] fix(device-mesh): correct MoE sealing fixtures --- crates/hipfire-runtime/src/moe_plan.rs | 54 +++++++++++++------------- 1 file changed, 28 insertions(+), 26 deletions(-) diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs index a01125d86e..0c297b4962 100644 --- a/crates/hipfire-runtime/src/moe_plan.rs +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -1158,16 +1158,15 @@ mod tests { let mut entries = manifest(); entries[1].name = "experts.gate".into(); entries[1].logical_shape = vec![4, 64, 64]; + entries[1].policy = ShardPolicy::Replicate; entries[2].name = "experts.down".into(); + entries[2].policy = ShardPolicy::Replicate; entries.push(WeightEntry::layer( "experts.up", 0, vec![4, 64, 64], DType::MQ4G256, - ShardPolicy::ExpertSharded { - n_experts: 4, - assign: ExpertAssign::Stride, - }, + ShardPolicy::Replicate, )); entries } @@ -1359,7 +1358,7 @@ mod tests { which: MoeProj::DownExpanded, sorted_slot_index: &tensors.sorted, expert_tile_ids: &tensors.tiles, - x: &tensors.down_x, + x: &tensors.rot_batch, y: &tensors.grouped_down, m_total: 8, batch_size: 2, @@ -1605,13 +1604,14 @@ mod tests { let mut tensors = grouped_tensors(); tensors.indices = tensor_with_bytes(vec![2, 2], DType::F32, 3 * DType::F32.size()); let steps = grouped_steps(&experts, &tensors); - let error = MoeFamily::new() - .seal_steps( - ExpertExecutionPlan::GroupedQuantized, - steps, - vec![StepCollective::None; 7], - ) - .expect_err("logical shape must not stand in for physical capacity"); + let error = match MoeFamily::new().seal_steps( + ExpertExecutionPlan::GroupedQuantized, + steps, + vec![StepCollective::None; 7], + ) { + Ok(_) => panic!("logical shape must not stand in for physical capacity"), + Err(error) => error, + }; assert!(error .to_string() .contains("insufficient logical/physical capacity")); @@ -1704,22 +1704,24 @@ mod tests { let mut duplicate = vec![StepCollective::None; 7]; duplicate[5] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); duplicate[6] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); - let error = family - .seal_steps( - ExpertExecutionPlan::GroupedQuantized, - grouped_steps(&experts, &tensors), - duplicate, - ) - .expect_err("a routed reduction cannot appear twice"); + let error = match family.seal_steps( + ExpertExecutionPlan::GroupedQuantized, + grouped_steps(&experts, &tensors), + duplicate, + ) { + Ok(_) => panic!("a routed reduction cannot appear twice"), + Err(error) => error, + }; assert!(error.to_string().contains("attached to combine")); - let error = family - .seal_steps( - ExpertExecutionPlan::GroupedQuantized, - grouped_steps(&experts, &tensors), - vec![StepCollective::None; 7], - ) - .expect_err("parallel owner must not silently use identity reduction"); + let error = match family.seal_steps( + ExpertExecutionPlan::GroupedQuantized, + grouped_steps(&experts, &tensors), + vec![StepCollective::None; 7], + ) { + Ok(_) => panic!("parallel owner must not silently use identity reduction"), + Err(error) => error, + }; assert!(error.to_string().contains("collective count")); } #[test] From 799563889481ee5d26eaa57c010dd3ee21ae8326 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 11:31:40 +0200 Subject: [PATCH 22/25] fix: close MoE schedule review gaps --- crates/hipfire-dispatch/src/families/moe.rs | 1 - crates/hipfire-dispatch/src/pipeline/steps.rs | 43 +++- crates/hipfire-runtime/src/moe_plan.rs | 220 ++++++++++++++---- 3 files changed, 212 insertions(+), 52 deletions(-) diff --git a/crates/hipfire-dispatch/src/families/moe.rs b/crates/hipfire-dispatch/src/families/moe.rs index b3b765053d..a4e5514ee2 100644 --- a/crates/hipfire-dispatch/src/families/moe.rs +++ b/crates/hipfire-dispatch/src/families/moe.rs @@ -505,7 +505,6 @@ impl<'a> MoeExpertRef<'a> { if self.collective_kind.is_none() && (self.owner_rank != 0 || self.group_devices.len() != 1 - || self.group_devices.first().copied() != Some(0) || self.owned.len() != self.n_experts || !self .owned diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index d424fc4ef0..6c10f97f12 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -494,6 +494,14 @@ pub fn validate_moe_step_schedule( )) } }; + if let RouterPlan::SoftmaxTopK { k_top, .. } = route { + if *k_top != 8 { + return Err(DispatchError::Hip(format!( + "generic MoE softmax route requires k_top=8, got {k_top}" + ))); + } + } + experts.validate()?; route.validate_against(experts.n_experts(), batch_size)?; @@ -541,10 +549,10 @@ pub fn validate_moe_step_schedule( experts.dtype() ))); } - validate_step_tensors(steps, experts, batch_size, hidden)?; validate_collectives( collectives, combine_index, + batch_size, hidden, experts.collective_kind(), )?; @@ -1147,9 +1155,13 @@ fn validate_step_tensors( fn validate_collectives( collectives: &[StepCollective], combine_index: usize, + batch_size: usize, hidden: usize, expected_kind: Option, ) -> Result<(), DispatchError> { + let element_count = batch_size + .checked_mul(hidden) + .ok_or_else(|| DispatchError::Hip("MoE collective element count overflows".into()))?; let mut reduction = None; for (index, collective) in collectives.iter().enumerate() { let StepCollective::AllReduce { @@ -1172,7 +1184,7 @@ fn validate_collectives( "MoE collective axis {kind:?} does not match owner axis {expected_kind:?}" ))); } - if *dim != hidden || group.len() < 2 || *rank >= group.len() { + if *dim != element_count || group.len() < 2 || *rank >= group.len() { return Err(DispatchError::Hip( "MoE collective rank/group/output dimension is invalid".into(), )); @@ -1768,15 +1780,21 @@ static ROTATION: OnceLock = OnceLock::new(); static FUSED_QKV: OnceLock = OnceLock::new(); static MOE: std::sync::LazyLock = std::sync::LazyLock::new(MoeFamily::new); +fn reject_unsealed_moe(steps: &[Step]) -> Result<(), DispatchError> { + if steps.iter().any(is_moe_step) { + return Err(DispatchError::Hip( + "unsealed MoE schedules require execute_sealed_steps or execute_sealed_steps_mesh".into(), + )); + } + Ok(()) +} + pub fn execute_steps( gpu: &mut Gpu, ctx: &DispatchCtx, steps: &[Step], ) -> Result<(), DispatchError> { - if steps.iter().any(is_moe_step) { - let collectives = vec![StepCollective::None; steps.len()]; - validate_moe_step_schedule(steps, &collectives)?; - } + reject_unsealed_moe(steps)?; execute_steps_inner(gpu, ctx, steps) } @@ -2804,4 +2822,17 @@ mod tests { "FusedGateUpQ8_0 missing from FUSED_TABLE" ); } + + #[test] + fn unsealed_moe_schedule_is_rejected_before_gpu_execution() { + let route = GpuTensor::null_for_test(); + let plan = RouterPlan::Precomputed { + topk_indices: &route, + topk_weights: &route, + k_top: 1, + }; + let steps = [Step::MoeRoute { plan }]; + let error = reject_unsealed_moe(&steps).expect_err("MoE must use a sealed executor path"); + assert!(error.to_string().contains("unsealed MoE schedules")); + } } diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs index 0c297b4962..405843de76 100644 --- a/crates/hipfire-runtime/src/moe_plan.rs +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -1154,6 +1154,23 @@ mod tests { entries } + fn grouped_manifest_for(parallelism: ExpertParallelism) -> Vec { + let mut entries = manifest_for(parallelism); + entries[0].logical_shape = vec![4, 8]; + entries[1].logical_shape[0] = 8; + entries[2].logical_shape[0] = 8; + if matches!(parallelism, ExpertParallelism::ExpertParallel) { + for entry in &mut entries[1..=2] { + entry.policy = ShardPolicy::ExpertSharded { + n_experts: 8, + assign: ExpertAssign::Stride, + }; + } + } + entries + } + + fn separate_manifest() -> Vec { let mut entries = manifest(); entries[1].name = "experts.gate".into(); @@ -1203,6 +1220,13 @@ mod tests { } } + fn grouped_spec(execution: &str, parallelism: ExpertParallelism) -> ExpertGroupSpec { + let mut value = spec(execution, parallelism); + value.n_experts = 8; + value + } + + fn tensor_with_bytes(shape: Vec, dtype: DType, bytes: usize) -> GpuTensor { let mut tensor = GpuTensor::null_for_test(); tensor.buf = unsafe { hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut(), bytes) }; @@ -1238,16 +1262,26 @@ mod tests { parallelism: ExpertParallelism, mesh: &DeviceMesh, ) -> ExpertPlan { - let mut plan = ExpertPlan::from_manifest( - &spec(execution, parallelism), - &manifest_for(parallelism), - mesh, - ) - .expect("test expert plan"); + let (group_spec, group_manifest) = if execution == "grouped_quantized" { + ( + grouped_spec(execution, parallelism), + grouped_manifest_for(parallelism), + ) + } else { + (spec(execution, parallelism), manifest_for(parallelism)) + }; + let mut plan = + ExpertPlan::from_manifest(&group_spec, &group_manifest, mesh).expect("test expert plan"); let group_size = plan.group_size(); + let pointer_slots = plan.n_experts().checked_mul(2).expect("test pointer slots"); for rank in 0..group_size { - plan.commit_rank_tables(rank, table(vec![8]), table(vec![8]), None) - .expect("commit rank tables"); + plan.commit_rank_tables( + rank, + table(vec![pointer_slots]), + table(vec![pointer_slots]), + None, + ) + .expect("commit rank tables"); } let placements = plan.placements().to_vec(); let mut load = plan.begin_load(); @@ -1278,21 +1312,21 @@ mod tests { fn grouped_tensors() -> GroupedTensors { GroupedTensors { - scores: tensor(vec![2, 4]), - indices: tensor(vec![2, 2]), - weights: tensor(vec![2, 2]), - counts: raw_i32(4), - offsets: raw_i32(5), - sorted: raw_i32(8), - tiles: raw_i32(2), - inverse: raw_i32(4), + scores: tensor(vec![2, 8]), + indices: tensor(vec![2, 8]), + weights: tensor(vec![2, 8]), + counts: raw_i32(8), + offsets: raw_i32(9), + sorted: raw_i32(16), + tiles: raw_i32(4), + inverse: raw_i32(16), x: tensor(vec![2, 64]), - grouped_gate: tensor(vec![8, 128]), - gate_batch: tensor(vec![4, 64]), - up_batch: tensor(vec![4, 64]), - rot_batch: tensor(vec![4, 64]), - grouped_down: tensor(vec![8, 64]), - down_x: tensor(vec![4, 64]), + grouped_gate: tensor(vec![16, 128]), + gate_batch: tensor(vec![16, 64]), + up_batch: tensor(vec![16, 64]), + rot_batch: tensor(vec![16, 64]), + grouped_down: tensor(vec![16, 64]), + down_x: tensor(vec![16, 64]), out: tensor(vec![2, 64]), } } @@ -1307,7 +1341,7 @@ mod tests { scores: &tensors.scores, topk_indices: &tensors.indices, topk_weights: &tensors.weights, - k_top: 2, + k_top: 8, normalize: true, }, }, @@ -1318,9 +1352,9 @@ mod tests { sorted_slot_index: &tensors.sorted, expert_tile_ids: &tensors.tiles, inverse_perm: &tensors.inverse, - total_slots: 4, - n_experts: 4, - m_total_max: 8, + total_slots: 16, + n_experts: 8, + m_total_max: 16, block_m: 4, }, Step::GroupedMoeGemm { @@ -1332,9 +1366,9 @@ mod tests { expert_tile_ids: &tensors.tiles, x: &tensors.x, y: &tensors.grouped_gate, - m_total: 8, + m_total: 16, batch_size: 2, - k_top: 2, + k_top: 8, }, Step::MoeGateUpUnscatter { y_grouped: &tensors.grouped_gate, @@ -1342,8 +1376,8 @@ mod tests { gate_batch: &tensors.gate_batch, up_batch: &tensors.up_batch, inter: 64, - k_top: 2, - m_total: 8, + k_top: 8, + m_total: 16, }, Step::MoeActivation { variant: MoeActivationVariant::SiluMul, @@ -1351,7 +1385,7 @@ mod tests { up: &tensors.up_batch, rot_out: &tensors.rot_batch, inter: 64, - rows: 4, + rows: 16, }, Step::GroupedMoeGemm { experts, @@ -1360,16 +1394,16 @@ mod tests { expert_tile_ids: &tensors.tiles, x: &tensors.rot_batch, y: &tensors.grouped_down, - m_total: 8, + m_total: 16, batch_size: 2, - k_top: 2, + k_top: 8, }, Step::MoeCombine { down_out: &tensors.grouped_down, topk_weights: &tensors.weights, out: &tensors.out, hidden: 64, - k_top: 2, + k_top: 8, batch_size: 2, inverse_perm: Some(&tensors.inverse), }, @@ -1378,7 +1412,6 @@ mod tests { fn indexed_steps<'a>( experts: &'a MoeExpertRef<'a>, - scores: &'a GpuTensor, indices: &'a GpuTensor, weights: &'a GpuTensor, x: &'a GpuTensor, @@ -1390,12 +1423,10 @@ mod tests { ) -> Vec> { vec![ Step::MoeRoute { - plan: RouterPlan::SoftmaxTopK { - scores, + plan: RouterPlan::Precomputed { topk_indices: indices, topk_weights: weights, k_top: 2, - normalize: true, }, }, Step::IndexedMoeGemv { @@ -1472,6 +1503,72 @@ mod tests { } } + #[test] + fn grouped_parallel_collective_covers_every_batched_output_element() { + let mesh = mesh_ep(); + let plan = resident_plan( + "grouped_quantized", + ExpertParallelism::ExpertParallel, + &mesh, + ); + let experts = plan.bind_expert_ref(0).unwrap(); + let tensors = grouped_tensors(); + let family = MoeFamily::new(); + + let mut short = vec![StepCollective::None; 7]; + short[6] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); + let error = match family.seal_steps( + ExpertExecutionPlan::GroupedQuantized, + grouped_steps(&experts, &tensors), + short, + ) { + Ok(_) => panic!("batched grouped output must not reduce only one row"), + Err(error) => error, + }; + assert!(error.to_string().contains("output dimension")); + + let mut full = vec![StepCollective::None; 7]; + full[6] = StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 0); + let schedule = family + .seal_steps( + ExpertExecutionPlan::GroupedQuantized, + grouped_steps(&experts, &tensors), + full, + ) + .expect("collective must cover batch_size * hidden elements"); + assert!(matches!( + &schedule.collectives()[6], + StepCollective::AllReduce { dim: 128, .. } + )); + } + + + #[test] + fn sealing_rejects_non_k8_softmax_routes() { + let mesh = DeviceMesh::single().unwrap(); + let plan = resident_plan("grouped_quantized", ExpertParallelism::Single, &mesh); + let experts = plan.bind_expert_ref(0).unwrap(); + let tensors = grouped_tensors(); + let mut steps = grouped_steps(&experts, &tensors); + steps[0] = Step::MoeRoute { + plan: RouterPlan::SoftmaxTopK { + scores: &tensors.scores, + topk_indices: &tensors.indices, + topk_weights: &tensors.weights, + k_top: 2, + normalize: true, + }, + }; + let error = match MoeFamily::new().seal_steps( + ExpertExecutionPlan::GroupedQuantized, + steps, + vec![StepCollective::None; 7], + ) { + Ok(_) => panic!("generic softmax routing is only executable at k_top=8"), + Err(error) => error, + }; + assert!(error.to_string().contains("k_top=8")); + } #[test] fn stride_assignment_and_named_group_are_deterministic() { let plan = ExpertPlan::from_manifest( @@ -1583,7 +1680,8 @@ mod tests { let tensors = grouped_tensors(); let steps = grouped_steps(&experts, &tensors); let mut collectives = vec![StepCollective::None; 7]; - collectives[6] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 1); + collectives[6] = + StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 1); let schedule = MoeFamily::new() .seal_steps(ExpertExecutionPlan::GroupedQuantized, steps, collectives) .expect("collective descriptor is locally typed"); @@ -1602,7 +1700,8 @@ mod tests { let plan = resident_plan("grouped_quantized", ExpertParallelism::Single, &mesh); let experts = plan.bind_expert_ref(0).unwrap(); let mut tensors = grouped_tensors(); - tensors.indices = tensor_with_bytes(vec![2, 2], DType::F32, 3 * DType::F32.size()); + tensors.indices = + tensor_with_bytes(vec![2, 8], DType::F32, 15 * DType::F32.size()); let steps = grouped_steps(&experts, &tensors); let error = match MoeFamily::new().seal_steps( ExpertExecutionPlan::GroupedQuantized, @@ -1633,7 +1732,6 @@ mod tests { let indexed_experts = indexed_plan.bind_expert_ref(0).unwrap(); let grouped_experts = grouped_plan.bind_expert_ref(1).unwrap(); - let indexed_scores = tensor(vec![4]); let indexed_indices = tensor(vec![2]); let indexed_weights = tensor(vec![2]); let indexed_x = tensor(vec![64]); @@ -1647,7 +1745,6 @@ mod tests { ExpertExecutionPlan::IndexedQuantized, indexed_steps( &indexed_experts, - &indexed_scores, &indexed_indices, &indexed_weights, &indexed_x, @@ -1674,7 +1771,7 @@ mod tests { { let mut collectives = vec![StepCollective::None; 7]; collectives[6] = - StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 1); + StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 1); collectives }, ) @@ -1702,8 +1799,10 @@ mod tests { let family = MoeFamily::new(); let mut duplicate = vec![StepCollective::None; 7]; - duplicate[5] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); - duplicate[6] = StepCollective::all_reduce(DimKind::Ep, 64, vec![0, 1], mesh.epoch(), 0); + duplicate[5] = + StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 0); + duplicate[6] = + StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 0); let error = match family.seal_steps( ExpertExecutionPlan::GroupedQuantized, grouped_steps(&experts, &tensors), @@ -1712,6 +1811,7 @@ mod tests { Ok(_) => panic!("a routed reduction cannot appear twice"), Err(error) => error, }; + assert!(error.to_string().contains("attached to combine")); let error = match family.seal_steps( @@ -1724,6 +1824,36 @@ mod tests { }; assert!(error.to_string().contains("collective count")); } + + #[test] + fn single_plan_accepts_nonzero_pipeline_stage_device() { + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2)]).unwrap(); + let mut staged_manifest = manifest_for(ExpertParallelism::Single); + for entry in &mut staged_manifest { + entry.layer = Some(1); + } + let mut staged_spec = spec("indexed_quantized", ExpertParallelism::Single); + staged_spec.layer = Some(1); + let mut plan = + ExpertPlan::from_manifest(&staged_spec, &staged_manifest, &mesh).unwrap(); + assert_eq!(plan.group_devices(), &[1]); + + for rank in 0..plan.group_size() { + plan.commit_rank_tables(rank, table(vec![8]), table(vec![8]), None) + .unwrap(); + } + let placements = plan.placements().to_vec(); + let mut load = plan.begin_load(); + for placement in placements { + load.reserve(placement).unwrap(); + } + load.commit(); + + let experts = plan.bind_expert_ref(0).unwrap(); + assert_eq!(experts.owner_rank(), 0); + assert_eq!(experts.group_devices(), &[1]); + assert_eq!(experts.owned(), &[0, 1, 2, 3]); + } #[test] fn per_expert_fallback_is_refused_at_plan_boundary() { let error = ExpertPlan::from_manifest( From 3e037380c8e91a1b3d733e27976e1298bb94f4c7 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 11:44:35 +0200 Subject: [PATCH 23/25] style(device-mesh): format MoE closure --- crates/hipfire-dispatch/src/pipeline/steps.rs | 4 ++-- crates/hipfire-runtime/src/moe_plan.rs | 22 ++++++------------- 2 files changed, 9 insertions(+), 17 deletions(-) diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index 6c10f97f12..d2be72bd3b 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -502,7 +502,6 @@ pub fn validate_moe_step_schedule( } } - experts.validate()?; route.validate_against(experts.n_experts(), batch_size)?; if !same_tensor(route_indices, route.route_buffers().0) @@ -1783,7 +1782,8 @@ static MOE: std::sync::LazyLock = std::sync::LazyLock::new(MoeFamily: fn reject_unsealed_moe(steps: &[Step]) -> Result<(), DispatchError> { if steps.iter().any(is_moe_step) { return Err(DispatchError::Hip( - "unsealed MoE schedules require execute_sealed_steps or execute_sealed_steps_mesh".into(), + "unsealed MoE schedules require execute_sealed_steps or execute_sealed_steps_mesh" + .into(), )); } Ok(()) diff --git a/crates/hipfire-runtime/src/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs index 405843de76..b536edf338 100644 --- a/crates/hipfire-runtime/src/moe_plan.rs +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -1170,7 +1170,6 @@ mod tests { entries } - fn separate_manifest() -> Vec { let mut entries = manifest(); entries[1].name = "experts.gate".into(); @@ -1226,7 +1225,6 @@ mod tests { value } - fn tensor_with_bytes(shape: Vec, dtype: DType, bytes: usize) -> GpuTensor { let mut tensor = GpuTensor::null_for_test(); tensor.buf = unsafe { hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut(), bytes) }; @@ -1270,8 +1268,8 @@ mod tests { } else { (spec(execution, parallelism), manifest_for(parallelism)) }; - let mut plan = - ExpertPlan::from_manifest(&group_spec, &group_manifest, mesh).expect("test expert plan"); + let mut plan = ExpertPlan::from_manifest(&group_spec, &group_manifest, mesh) + .expect("test expert plan"); let group_size = plan.group_size(); let pointer_slots = plan.n_experts().checked_mul(2).expect("test pointer slots"); for rank in 0..group_size { @@ -1542,7 +1540,6 @@ mod tests { )); } - #[test] fn sealing_rejects_non_k8_softmax_routes() { let mesh = DeviceMesh::single().unwrap(); @@ -1680,8 +1677,7 @@ mod tests { let tensors = grouped_tensors(); let steps = grouped_steps(&experts, &tensors); let mut collectives = vec![StepCollective::None; 7]; - collectives[6] = - StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 1); + collectives[6] = StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 1); let schedule = MoeFamily::new() .seal_steps(ExpertExecutionPlan::GroupedQuantized, steps, collectives) .expect("collective descriptor is locally typed"); @@ -1700,8 +1696,7 @@ mod tests { let plan = resident_plan("grouped_quantized", ExpertParallelism::Single, &mesh); let experts = plan.bind_expert_ref(0).unwrap(); let mut tensors = grouped_tensors(); - tensors.indices = - tensor_with_bytes(vec![2, 8], DType::F32, 15 * DType::F32.size()); + tensors.indices = tensor_with_bytes(vec![2, 8], DType::F32, 15 * DType::F32.size()); let steps = grouped_steps(&experts, &tensors); let error = match MoeFamily::new().seal_steps( ExpertExecutionPlan::GroupedQuantized, @@ -1799,10 +1794,8 @@ mod tests { let family = MoeFamily::new(); let mut duplicate = vec![StepCollective::None; 7]; - duplicate[5] = - StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 0); - duplicate[6] = - StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 0); + duplicate[5] = StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 0); + duplicate[6] = StepCollective::all_reduce(DimKind::Ep, 128, vec![0, 1], mesh.epoch(), 0); let error = match family.seal_steps( ExpertExecutionPlan::GroupedQuantized, grouped_steps(&experts, &tensors), @@ -1834,8 +1827,7 @@ mod tests { } let mut staged_spec = spec("indexed_quantized", ExpertParallelism::Single); staged_spec.layer = Some(1); - let mut plan = - ExpertPlan::from_manifest(&staged_spec, &staged_manifest, &mesh).unwrap(); + let mut plan = ExpertPlan::from_manifest(&staged_spec, &staged_manifest, &mesh).unwrap(); assert_eq!(plan.group_devices(), &[1]); for rank in 0..plan.group_size() { From 1a069eab757e4d5f85b86a802fc77dd0abd76933 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 12:00:58 +0200 Subject: [PATCH 24/25] fix(device-mesh): validate MoE router before grammar --- crates/hipfire-dispatch/src/pipeline/steps.rs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index d2be72bd3b..2ee431cd38 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -268,6 +268,17 @@ pub fn validate_moe_step_schedule( return Err(DispatchError::Hip("MoE schedule is empty".into())); } + if let Some(Step::MoeRoute { + plan: RouterPlan::SoftmaxTopK { k_top, .. }, + }) = steps.first() + { + if *k_top != 8 { + return Err(DispatchError::Hip(format!( + "generic MoE softmax route requires k_top=8, got {k_top}" + ))); + } + } + let (protocol, route, experts, batch_size, combine_index, hidden, route_indices, route_weights) = match steps { [Step::MoeRoute { plan }, Step::IndexedMoeGemv { @@ -494,13 +505,6 @@ pub fn validate_moe_step_schedule( )) } }; - if let RouterPlan::SoftmaxTopK { k_top, .. } = route { - if *k_top != 8 { - return Err(DispatchError::Hip(format!( - "generic MoE softmax route requires k_top=8, got {k_top}" - ))); - } - } experts.validate()?; route.validate_against(experts.n_experts(), batch_size)?; From 26fe251394499721679ae6a05a221d8f4997e3ab Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 13:07:39 +0200 Subject: [PATCH 25/25] fix(device-mesh): validate MoE step tensors --- crates/hipfire-dispatch/src/pipeline/steps.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/hipfire-dispatch/src/pipeline/steps.rs b/crates/hipfire-dispatch/src/pipeline/steps.rs index 2ee431cd38..5fc548857f 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -552,6 +552,7 @@ pub fn validate_moe_step_schedule( experts.dtype() ))); } + validate_step_tensors(steps, experts, batch_size, hidden)?; validate_collectives( collectives, combine_index,