diff --git a/Cargo.lock b/Cargo.lock index 4ed89d3d8d..91a2bf9d26 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", @@ -1195,6 +1196,7 @@ version = "0.3.0" dependencies = [ "hip-bridge", "hipfire-dispatch", + "hipfire-hardware", "hipfire-runtime", "rdna-compute", ] @@ -1219,6 +1221,7 @@ dependencies = [ "hip-bridge", "hipfire-config", "hipfire-dispatch", + "hipfire-hardware", "hipfire-reap", "hipfire-runtime", "rdna-compute", @@ -1259,6 +1262,7 @@ dependencies = [ "hipfire-arch-qwen35-vl", "hipfire-config", "hipfire-dispatch", + "hipfire-hardware", "hipfire-reap", "hipfire-runtime", "rdna-compute", @@ -1382,6 +1386,7 @@ version = "0.3.0" dependencies = [ "hip-bridge", "hipfire-config", + "hipfire-hardware", "rdna-compute", ] @@ -1446,6 +1451,7 @@ dependencies = [ "hipfire-config", "hipfire-dispatch", "hipfire-engine", + "hipfire-hardware", "hipfire-loader", "hipfire-pflash", "hipfire-runtime", @@ -1456,6 +1462,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 +1489,7 @@ dependencies = [ "hipfire-arch-qwen35", "hipfire-arch-qwen35-vl", "hipfire-config", + "hipfire-hardware", "hipfire-runtime", "rdna-compute", "saddle-core", @@ -1563,6 +1578,7 @@ dependencies = [ "hipfire-detect", "hipfire-dispatch", "hipfire-engine", + "hipfire-hardware", "hipfire-loader", "hipfire-pflash", "libc", @@ -2902,6 +2918,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..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_runtime::multi_gpu::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; @@ -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 e74f694b10..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_runtime::multi_gpu::Gpus; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::path::{Path, PathBuf}; @@ -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 e0104087d2..71e9ae76b7 100644 --- a/crates/hipfire-arch-deepseek4/src/ep.rs +++ b/crates/hipfire-arch-deepseek4/src/ep.rs @@ -5,13 +5,16 @@ 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_runtime::multi_gpu::Gpus; +use hipfire_hardware::Gpus; use rdna_compute::{Gpu, GpuTensor}; // ───────────────────────── Ship 6 substrate-EP (DeepSeek-V4) ───────────────── @@ -41,7 +44,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 +119,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], @@ -125,6 +128,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 +148,7 @@ fn forward_ep_tp_graph_body( gpus, bindings.as_mut_slice(), partials, + &group, &program, cfg.hidden_size, ) @@ -161,7 +166,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 +181,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 +329,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], @@ -381,6 +386,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 +405,7 @@ fn forward_ep_direct( gpus, binds.as_mut_slice(), partials, + &group, &program, hidden, ) @@ -528,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 a646f869c0..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!( @@ -12875,7 +12877,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..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}; @@ -526,7 +525,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], @@ -550,7 +549,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 +591,7 @@ pub fn mtp_forward_ep( gpus, binds.as_mut_slice(), partials, + &group, &program, hidden, ) @@ -914,4 +914,3 @@ pub fn mtp_forward_batched( Ok(()) } - 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..d02289ef0e 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -17,9 +17,12 @@ 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 rdna_compute::Gpu; +use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; +use hipfire_runtime::weight_manifest::{ + DTypeConstraint, FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, +}; +use rdna_compute::{DType, Gpu}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::{execute_steps, GemvInput, Step}; @@ -36,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; @@ -43,12 +98,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 +108,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 +116,193 @@ 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 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_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_with_dtype_constraint( + "wq", + layer, + vec![heads * head_dim, dim], + DType::F16, + linear.clone(), + FusedQkv { + q_heads: heads, + kv_heads, + head_dim, + layout: FusedQkvLayout::Qkv, + }, + )); + 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_with_dtype_constraint( + "wv", + layer, + vec![kv_heads * head_dim, dim], + DType::F16, + linear.clone(), + ColumnShard { axis: 0 }, + )); + 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_with_dtype_constraint( + "ffn_gate", + layer, + vec![hidden, dim], + DType::F16, + linear.clone(), + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer_with_dtype_constraint( + "ffn_up", + layer, + vec![hidden, dim], + DType::F16, + linear.clone(), + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer_with_dtype_constraint( + "ffn_down", + layer, + vec![dim, hidden], + DType::F16, + linear.clone(), + RowShard { axis: 1 }, + )); + manifest.push(WeightEntry::layer_with_dtype_constraint( + "attn_norm", + layer, + vec![dim], + DType::F32, + norm.clone(), + Replicate, + )); + 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_with_dtype_constraint( + "q_norm", + layer, + vec![head_dim], + DType::F32, + norm.clone(), + Replicate, + )); + manifest.push(WeightEntry::layer_with_dtype_constraint( + "k_norm", + layer, + vec![head_dim], + DType::F32, + norm.clone(), + Replicate, + )); + } + } + manifest.push(WeightEntry::model_with_dtype_constraint( + "output_norm", + vec![dim], + DType::F32, + norm, + Replicate, + )); + manifest.push(WeightEntry::model_with_dtype_constraint( + "lm_head", + vec![cfg.vocab_size, dim], + DType::F16, + linear, + Pin(PinTarget::Output), + )); + 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) + .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..7f374f20c7 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -40,19 +40,24 @@ impl ArchModel for LlamaBundle { weights, scratch, kv, + manifest_plan: _, + weight_store, + weight_origin: _, + 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. scratch.free_gpu(gpu); + if let Some(store) = weight_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. + 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 736eeb8674..c40e55ad59 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -4,17 +4,44 @@ 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::hfq::HfqFile; use hipfire_runtime::llama::KvCacheExt; +use hipfire_runtime::llama::{ + EmbeddingFormat, ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LayerWeights, + LlamaConfig, LlamaWeights, WeightTensor, +}; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; +use hipfire_runtime::model_source::ModelSource as ModelSourceTrait; +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, +}; +use rdna_compute::{DType, GpuTensor}; +use std::collections::HashMap; pub struct LlamaBundle { pub config: LlamaConfig, 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, + /// 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 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, /// 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,34 +52,557 @@ 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, } +/// 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) -> 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}"), + } +} + +fn plan_single( + config: &LlamaConfig, + has_separate_lm_head: bool, +) -> Result<(DeviceMesh, ManifestPlan), String> { + 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) + .map_err(|e| format!("llama: manifest planning failed: {e}"))?; + Ok((mesh, plan)) +} + +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"), + ] +} +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) => 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"), + ("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" + ) { + 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: AWQ sidecar {sidecar} is not represented by the manifest pilot" + )); + } + } + 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); + } + 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(), + ); + } + } + other => { + return Err(format!( + "{name}: quant_type={other} is not a host float payload" + )); + } + } + Ok(bytes) +} + +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}")), + } +} + +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_materialized( + assembly: &WeightStoreAssemblyGuard<'_>, + name: &str, + layer: Option, + slot: usize, +) -> Result<(), String> { + 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" + )), + } +} + +fn resident_cell( + cells: &mut HashMap<(String, Option), TakenWeight>, + name: &str, + layer: Option, +) -> GpuTensor { + match cells.remove(&(name.to_string(), layer)) { + Some(TakenWeight { + handle: WeightHandle::Resident(tensor), + .. + }) => tensor, + _ => unreachable!("validated LLaMA assembly lost resident {name}[layer {layer:?}]"), + } +} + +fn resident_weight( + cells: &mut HashMap<(String, Option), TakenWeight>, + name: &str, + layer: Option, + m: usize, + k: usize, +) -> WeightTensor { + let tensor = resident_cell(cells, name, layer); + let dtype = tensor.dtype; + WeightTensor { + buf: tensor, + gpu_dtype: dtype, + m, + k, + 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 { + 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, + transaction: &mut WeightLoadTransaction, +) -> 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); + + 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_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 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))) + } 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, + ), + }); + } + debug_assert!(cells.is_empty(), "validated LLaMA assembly left cells"); + Ok(LlamaWeights { + token_embd, + embd_format, + output_norm, + output, + layers, + lm_head_aliases_embd, + }) +} + /// 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 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. +#[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) = 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)?; + 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_has_separate_lm_head(&hfq); + 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); - // 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:?}"))?; - 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, + // 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 kv = ::from_mode( + 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, @@ -61,20 +611,43 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result 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, ) - .map_err(|e| format!("llama: ::from_mode failed: {e}"))?; - (config, weights, kv, scratch) } 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 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); - // Replicate carriers.rs `resolve_kv_mode` warning path verbatim. let kv_mode_str = ctx .kv_mode_override .filter(|s| !s.is_empty()) @@ -86,41 +659,109 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result::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:?}")); + } }; - let kv = ::from_mode( - rr.mode, - KvTarget::Single(ctx.gpu), - &dims, + ( + config, + weights, + kv, + scratch, + manifest_plan, + None, + mesh, + weight_origin, ) - .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:?}"))?; - (config, weights, kv, scratch) } }; - Ok(LlamaBundle { + + 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(transaction) = weight_store { + if let Err((transaction, error)) = bundle.attach_weight_store(transaction) { + let LlamaBundle { + weights, + scratch, + kv, + .. + } = bundle; + let rollback = transaction.rollback(ctx.gpu); + scratch.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + let _ = kv.free_gpu(ctx.gpu); + return Err(with_weight_rollback_error(error, rollback)); + } + } + Ok(bundle) } /// Alias matching the `load__bundle` naming convention in the task. pub use load_bundle as load_llama_bundle; impl LlamaBundle { + /// Attach an unpublished load transaction after validating the complete + /// 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)> { + if self.weight_store.is_some() { + return Err((transaction, "llama: weight store already attached".into())); + } + 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}"), + )); + } + }; + self.weight_store = Some(attached); + 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`. @@ -132,3 +773,469 @@ impl LlamaBundle { self.dflash_extract_layers = layers; } } + +#[cfg(test)] +mod tests { + use super::*; + 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::llama::ModelArch; + use hipfire_runtime::llama::{ + forward_scratch_compute, forward_scratch_embed, KvCache, KvCacheExt, KvDims, KvLayers, + KvTarget, + }; + 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, + }; + 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::(); + 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 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, + with_q_proj_bias: bool, + malformed_output_norm: bool, + separate_lm_head: bool, + ) -> (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), + f32_hfq_tensor("model.norm.weight", &[32], 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], 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 let Some(lm_head_name) = lm_head_name { + tensors.push(f32_hfq_tensor(lm_head_name, &[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 { + 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(), true).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().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"] { + store + .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 transaction, + ) { + Ok(_) => panic!("alias unexpectedly assembled as typed weights"), + Err(error) => error, + }; + assert!(error.contains("alias")); + assert_eq!(transaction.len(), 3); + assert!(transaction.contains("token_embd", None, 0)); + 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]; + 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)); + } + + #[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_sidecar_selects_legacy_loader() { + let (path, hfq) = fixture_hfq(true, false, false, false); + assert_eq!(classify_hfq_route(&hfq), HfqLoadRoute::LegacyAwq); + 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 { + 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_reclaims_every_uploaded_allocation() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + 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!("post-upload fault unexpectedly succeeded"), + Err(error) => error, + }; + drop(ctx); + 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_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 mut bundle = + load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); + drop(ctx); + + 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"); + 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_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 { + 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-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..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_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; - use hipfire_runtime::multi_gpu::Gpus; 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"); @@ -71,16 +86,25 @@ 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)"); + 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"); @@ -88,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)"); @@ -103,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 }; @@ -119,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(); @@ -137,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 6c9830de63..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, @@ -1635,7 +1636,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], @@ -1653,7 +1654,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 +1696,7 @@ pub fn forward_ep( gpus, binds.as_mut_slice(), partials, + &group, &program, hidden, ) 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..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_runtime::multi_gpu::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::Gpu; @@ -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 fa06e8e313..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_hardware::Gpus; use hipfire_runtime::hfq::HfqFile; -use hipfire_runtime::multi_gpu::Gpus; +use hipfire_runtime::llama::KvCacheExt; use std::path::Path; fn main() { @@ -27,7 +27,8 @@ 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 ca9c2f18e4..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_runtime::multi_gpu::Gpus; use std::path::Path; fn main() { @@ -31,7 +31,8 @@ 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 39a642d05e..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_runtime::multi_gpu::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 @@ -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 @@ -2452,15 +2455,17 @@ 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(); 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..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_runtime::multi_gpu::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..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_runtime::multi_gpu::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..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_runtime::multi_gpu::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use rdna_compute::Gpu; use std::path::Path; @@ -117,7 +117,8 @@ 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-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..a4e5514ee2 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}; @@ -26,7 +26,596 @@ 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, + }, +} + +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 { + 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() + ))); + } + 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()))?; + 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 logical/physical capacity for {expected_slots} slots" + ))); + } + 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, + }; + 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 + || score_elements < score_capacity + || scores.buf.size() < score_bytes + { + return Err(DispatchError::Hip(format!( + "MoE router score shape/dtype/capacity 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(), + )), + } + } +} + +/// 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 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, + 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> MoeExpertRefBinding<'a> { + /// 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 unsafe fn from_validated_plan( + 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, + ) -> Self { + Self { + gate_up_ptrs, + down_ptrs, + dummy_gate_up, + dtype, + n_experts, + 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 + } + + 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 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 + } + + 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> { + 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.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*{}] with {pointer_bytes} bytes", + self.n_experts + ))); + } + if let Some(dummy) = self.dummy_gate_up { + 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(), + )); + } + } + 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.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.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(), + )); + } + 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.is_some_and(|previous| expert <= previous) { + return Err(DispatchError::Hip(format!( + "MoeExpertRef: owned experts are not strictly ordered at {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 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!( + "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 +655,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,8 +675,6 @@ 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)) } } @@ -929,6 +1502,305 @@ 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, + schedules: &[&crate::pipeline::steps::SealedMoeSchedule<'a>], + ) -> Result<(), DispatchError> { + crate::pipeline::steps::execute_sealed_steps_mesh(gpus, mesh, ctx, schedules) + } /// Resolve the best kernel key for the given MoE variant. /// @@ -1207,4 +2079,41 @@ 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(); + 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 scores = tensor(vec![8]); + let indices = tensor(vec![8]); + let weights = tensor(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 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..afc13c4b37 100644 --- a/crates/hipfire-dispatch/src/pipeline/mod.rs +++ b/crates/hipfire-dispatch/src/pipeline/mod.rs @@ -10,8 +10,12 @@ 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 mod steps; + +pub use steps::{ + execute_sealed_steps_mesh, execute_steps, 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 +2864,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 +3092,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..5fc548857f 100644 --- a/crates/hipfire-dispatch/src/pipeline/steps.rs +++ b/crates/hipfire-dispatch/src/pipeline/steps.rs @@ -4,16 +4,24 @@ //! 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::moe::{ + ExpertExecutionPlan, MoeActivationVariant, MoeExpertRef, MoeFamily, MoeProj, MoeProtocolKind, + 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. + /// Rotation disposition of a Gemv's input. Borrows (never owns a RotatedActivation). pub enum GemvInput<'a> { Raw(&'a GpuTensor), // launch_op self-rotates via run_auto (plan-aware) @@ -81,6 +89,76 @@ 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,7 +171,1066 @@ 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, + } + } +} +/// 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() + .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())); + } + + 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 { + 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()))?; + if *batch_size != 1 + || *down_batch != 1 + || inverse_perm.is_some() + || *k_top != plan.k_top() + || *inter != (*experts).expert_m() + || *rows != expected_rows + { + return Err(DispatchError::Hip( + "indexed MoE grammar requires decode batch=1, normalized route, 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, + }] => { + 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 != expected_slots + || *n_experts != experts.n_experts() + || !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(), + )); + } + 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(), + )); + } + 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(), + )); + } + 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 != expected_slots + || !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 || hidden != experts.expert_k() { + return Err(DispatchError::Hip( + "MoE combine hidden size does not match the expert down projection".into(), + )); + } + 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 { + 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, + batch_size, + 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(()) +} + +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, + 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 execution_signature(&self) -> Result, DispatchError> { + derive_moe_execution_signature(&self.steps, self.execution) + } +} + +pub fn execute_sealed_steps( + gpu: &mut Gpu, + ctx: &DispatchCtx, + 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) +} + +/// 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, + schedules: &[&SealedMoeSchedule<'a>], +) -> Result, DispatchError> { + if schedules.is_empty() { + return Err(DispatchError::Hip( + "parallel MoE execution has no rank schedules".into(), + )); + } + if mesh.n_devices() != gpus_len { + return Err(DispatchError::Hip(format!( + "MoE mesh has {} devices but Gpus owns {}", + mesh.n_devices(), + gpus_len + ))); + } + + 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( + schedule.steps(), + 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() + .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(), + )); + } + 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(), + )); + } + 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 &device in &group { + if device >= gpus_len { + return Err(DispatchError::Hip(format!( + "MoE collective device {device} is outside the Gpus owner" + ))); + } + } + 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(&preflight.group, &preflight.outputs, preflight.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 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> { + 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 {:?}/logical capacity {logical_elements}/physical bytes {}, \ + expected {:?}/{capacity} elements/{required_bytes} bytes", + tensor.dtype, + tensor.buf.size(), + dtype + ))); + } + 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")))?; + 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 {:?}/logical bytes {logical_bytes}/physical bytes {}, \ + expected Raw/at least {bytes} bytes", + tensor.dtype, + tensor.buf.size() + ))); + } + 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(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(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()))?; + 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.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) { + 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 { + 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".into()) + })?, + )?; + 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()))?, + )?; + 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, + expert_offsets, + sorted_slot_index, + expert_tile_ids, + inverse_perm, + total_slots, + n_experts, + m_total_max, + 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", 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)?; + } + 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(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( + 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".into()) + })?, + )?; + 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, + 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, + 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 { + 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 != element_count || group.len() < 2 || *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(), + )); + } + } + if expected_kind.is_some() != reduction.is_some() { + return Err(DispatchError::Hip( + "MoE collective count does not match the resolved owner parallelism".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 ────────────────────────────────────────────────────────── @@ -645,22 +1782,35 @@ 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); + +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> { + reject_unsealed_moe(steps)?; + 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 +1818,6 @@ pub fn execute_steps( continue; } } - // ───────────────────────────────────────────────────────────────── launch_fused(gpu, ctx, key, &steps[i..i + len])?; i += len; } else { @@ -679,6 +1828,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 +2161,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 +2827,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-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-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 89% rename from crates/hipfire-runtime/src/multi_gpu.rs rename to crates/hipfire-hardware/src/lib.rs index d88214d294..1217a322e9 100644 --- a/crates/hipfire-runtime/src/multi_gpu.rs +++ b/crates/hipfire-hardware/src/lib.rs @@ -26,6 +26,23 @@ use hip_bridge::{ }; use rdna_compute::{DType, Gpu, GpuTensor}; +mod mesh; +pub use mesh::{Axis, CollectiveHint, DeviceMesh, DimKind, MeshEpoch, MeshError}; + +/// Device-resolution knobs supplied by the resolved process/runtime config +/// when constructing a [`Gpus`] owner. +/// +/// `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 allow_mixed_arch: bool, + pub uniform_vram_tolerance_gb: Option, +} + /// 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 @@ -114,9 +131,12 @@ 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, @@ -170,7 +190,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")); } @@ -183,18 +207,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( @@ -209,20 +230,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", @@ -231,9 +254,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], @@ -272,16 +296,16 @@ 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 @@ -289,6 +313,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, @@ -803,23 +828,27 @@ 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; 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(); @@ -835,6 +864,7 @@ impl Gpus { } Ok(Self { rccl_comms: None, + use_rccl: opts.tp_use_rccl.unwrap_or(true), devices, layer_to_device, band_starts, @@ -856,22 +886,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!(crate::config::get().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. \ @@ -896,71 +920,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}")))?; @@ -1026,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( @@ -1043,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); @@ -1056,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(()) } @@ -1435,23 +1465,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 +1486,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(()) @@ -1642,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, )?); @@ -1792,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, )?); @@ -1845,30 +1886,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. -fn resolve_device_ids(n_devices: usize) -> HipResult> { - if let Some(ref s) = crate::config::get().devices { - let ids: Vec = s +/// 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(|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!("hardware.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(), + parsed.len() ), )); } - return Ok(ids[..n_devices].to_vec()); + return Ok(parsed[..n_devices].to_vec()); } - Ok((0..n_devices as i32).collect()) + (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> { @@ -1879,23 +1926,26 @@ 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 allow_mixed = crate::config::get().allow_mixed_arch; 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 +1953,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 +1968,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, ), )); } @@ -2085,4 +2135,12 @@ 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 new file mode 100644 index 0000000000..eb6fa79172 --- /dev/null +++ b/crates/hipfire-hardware/src/mesh.rs @@ -0,0 +1,418 @@ +// 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 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 +/// 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, + n_devices: usize, + 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. + /// + /// 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() -> Result { + 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.n_devices + } + + /// 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(), + n_devices: self.n_devices, + 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().expect("single-device mesh cannot overflow") + } +} + +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().unwrap(); + 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().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()); + 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)]).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]); + } + + #[test] + fn composed_coordinates_groups_stages_and_squeeze() { + 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]); + + // 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(), + &[ + 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)]).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 new file mode 100644 index 0000000000..653d945770 --- /dev/null +++ b/crates/hipfire-hardware/tests/ownership.rs @@ -0,0 +1,30 @@ +// 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)]).unwrap(); + 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..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_runtime::multi_gpu::Gpus::init_layers(&counts).map_err(|e| format!("{e}"))? + hipfire_hardware::Gpus::init_layers(&device_opts, &counts) + .map_err(|e| format!("{e}"))? } - None => hipfire_runtime::multi_gpu::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 c099ef2818..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_runtime::multi_gpu::Gpus; use hipfire_runtime::spec::{SpecEmit, SpecEmitCtx, SpecTargetGuard, Speculator}; use hipfire_runtime::triattn::{EvictionCtx, TriAttnCenters}; use rdna_compute::Gpu; @@ -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, quantile: f64) -> f64 { values.sort_by(f64::total_cmp); @@ -16,7 +16,8 @@ fn main() { const WARMUPS: usize = 32; const REPEATS: usize = 512; - let gpus = Gpus::init_tp(RANKS, 43).expect("init exact gfx1201 TP4"); + let device_opts = hipfire_runtime::config::get().device_resolve_opts(); + let gpus = Gpus::init_tp(&device_opts, RANKS, 43).expect("init exact gfx1201 TP4"); assert!( gpus.devices .iter() 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..0e212631bc 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; @@ -120,7 +120,8 @@ fn main() { let samples = env_usize("HIPFIRE_DS4_OWNER_WORKER_SAMPLES", 100); assert!(samples > 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 ee3a620ec7..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_runtime::multi_gpu::Gpus; use hipfire_runtime::tp_shard::{ExpertAssign, ShardConfig}; use rdna_compute::{DType, GpuTensor}; use std::path::Path; @@ -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 76716f208a..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_runtime::multi_gpu::Gpus; use hipfire_runtime::tokenizer::Tokenizer; use rdna_compute::Gpu; use std::path::Path; @@ -135,7 +135,8 @@ 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 6cc609280d..7726af5f4c 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; @@ -40,7 +40,8 @@ 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 2dfac4b4ce..dc4d975f3a 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]; @@ -37,7 +37,8 @@ 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)"); @@ -82,7 +83,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 +134,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 +150,8 @@ 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..bb5e9e5566 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, @@ -202,6 +202,18 @@ 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 +254,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 +272,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 15f5ae8e36..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 crate::multi_gpu::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 { @@ -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> { @@ -68,9 +69,10 @@ fn all_reduce_sum_f32_decode( hipfire_config::developer_var("HIPFIRE_EP_PEER_ALLREDUCE_DECODE").as_deref() == Ok("1") }); 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) } } @@ -103,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. /// @@ -111,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> { @@ -125,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) @@ -213,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])?; @@ -266,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/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/lib.rs b/crates/hipfire-runtime/src/lib.rs index 1f98ad8d98..f5821cca35 100644 --- a/crates/hipfire-runtime/src/lib.rs +++ b/crates/hipfire-runtime/src/lib.rs @@ -47,7 +47,7 @@ pub mod loader_api; pub mod loop_guard; pub mod model_load; pub mod model_source; -pub mod multi_gpu; +pub mod moe_plan; pub mod paro; pub mod prefix; pub mod reset_core; @@ -75,5 +75,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/llama.rs b/crates/hipfire-runtime/src/llama.rs index e9e9091075..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 crate::multi_gpu::Gpus; use hip_bridge::HipResult; +use hipfire_hardware::Gpus; use rdna_compute::{DType, Gpu, GpuTensor}; /// Model architecture type. @@ -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..90b38444eb 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 crate::multi_gpu::Gpus; use hip_bridge::HipResult; +use hipfire_hardware::{DeviceMesh, DimKind, Gpus}; use rdna_compute::{Gpu, GpuTensor}; /// Where each piece of the model lands across a device slice. `single` = the @@ -30,6 +30,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] } @@ -39,15 +91,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, } @@ -61,8 +113,9 @@ pub trait WeightSource { 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, @@ -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,26 @@ 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)]) + .expect("small test mesh construction cannot overflow"); + 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/moe_plan.rs b/crates/hipfire-runtime/src/moe_plan.rs new file mode 100644 index 0000000000..b536edf338 --- /dev/null +++ b/crates/hipfire-runtime/src/moe_plan.rs @@ -0,0 +1,1899 @@ +// 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 {} + +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, group_size: usize) -> Self { + Self { + resident: vec![false; slots.len()], + slots, + rank_tables: (0..group_size).map(|_| None).collect(), + } + } + + fn index_of(&self, placement: ExpertPlacement) -> Option { + 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); + } + + 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_partition: Vec<(usize, usize, usize)>, + 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)?; + 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)?; + + 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 group_size = group_devices.len(); + let owner = resolve_placements( + spec.n_experts, + group_size, + spec.parallelism, + spec.assignment, + ); + let owner_views = (0..group_size) + .map(|rank| { + owner + .iter() + .filter(|placement| placement.owner == rank) + .map(|placement| placement.global_id) + .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 { + 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_partition, + owner: ExpertStorageOwner::new(owner, group_size), + }) + } + + 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() + )) + }) + } + + /// 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) -> 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, + }; + // 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()))?; + 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 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:?}")) + })?; + 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 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")); + } + 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 + ))); + } + 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 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], + 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]) + || up_shapes.iter().any(|shape| shape != &up_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 { + 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, + 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::*; + 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") + } + + 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 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 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(); + 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::Replicate, + )); + 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(), + 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 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) }; + 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 (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![pointer_slots]), + table(vec![pointer_slots]), + 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, 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![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]), + } + } + + 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: 8, + 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: 16, + n_experts: 8, + m_total_max: 16, + 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: 16, + batch_size: 2, + k_top: 8, + }, + 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: 8, + m_total: 16, + }, + Step::MoeActivation { + variant: MoeActivationVariant::SiluMul, + gate: &tensors.gate_batch, + up: &tensors.up_batch, + rot_out: &tensors.rot_batch, + inter: 64, + rows: 16, + }, + Step::GroupedMoeGemm { + experts, + which: MoeProj::DownExpanded, + sorted_slot_index: &tensors.sorted, + expert_tile_ids: &tensors.tiles, + x: &tensors.rot_batch, + y: &tensors.grouped_down, + m_total: 16, + batch_size: 2, + k_top: 8, + }, + Step::MoeCombine { + down_out: &tensors.grouped_down, + topk_weights: &tensors.weights, + out: &tensors.out, + hidden: 64, + k_top: 8, + batch_size: 2, + inverse_perm: Some(&tensors.inverse), + }, + ] + } + + fn indexed_steps<'a>( + experts: &'a MoeExpertRef<'a>, + 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::Precomputed { + topk_indices: indices, + topk_weights: weights, + k_top: 2, + }, + }, + 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 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( + &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 mut plan = ExpertPlan::from_manifest( + &spec("indexed_quantized", ExpertParallelism::ExpertParallel), + &manifest(), + &mesh_ep(), + ) + .unwrap(); + let gate = table(vec![8]); + let down = table(vec![8]); + 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); + } + + #[test] + fn rollback_then_reserve_reuse_and_repeated_teardown_are_safe() { + 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(); + 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) + .expect("rollback must release placement"); + load.commit(); + } + assert_eq!(plan.resident_slots(), 1); + plan.unload(); + plan.unload(); + 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, 128, 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, 8], DType::F32, 15 * DType::F32.size()); + let steps = grouped_steps(&experts, &tensors); + 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")); + } + + #[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_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_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, 128, 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, 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), + duplicate, + ) { + 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( + 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] + 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( + &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 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() { + 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); + } +} 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 new file mode 100644 index 0000000000..2f29231941 --- /dev/null +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -0,0 +1,1177 @@ +// 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), + } + } + + /// 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. +#[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(()) +} + +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(); + 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")); + } + 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 !source_entry + .dtype_constraint + .same_source_set(&entry.dtype_constraint) + { + 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) { + 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_weight_layers(weights, n_layers)?; + 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}" + )); + } + 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 { + 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)]) + .expect("small test mesh construction cannot overflow"); + let embed = WeightEntry::model( + "token_embd", + vec![32, 8], + DType::F16, + ShardPolicy::Pin(PinTarget::Embed), + ); + 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( + &[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)]) + .expect("small test mesh construction cannot overflow"); + 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().expect("single-device mesh construction cannot overflow") + ) + .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().expect("single-device mesh construction cannot overflow") + ) + .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()); + } + + #[test] + fn planning_rejects_weight_layer_at_n_layers_and_accepts_last_layer() { + 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); + 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().expect("single-device mesh construction cannot overflow") + ) + .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().expect("single-device mesh construction cannot overflow") + ) + .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().expect("single-device mesh construction cannot overflow") + ) + .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().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 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().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 new file mode 100644 index 0000000000..d67bf06c03 --- /dev/null +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -0,0 +1,1155 @@ +// 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 moves a committed transaction into +//! 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. +use crate::weight_manifest::{placement_devices, ShardPolicy, WeightEntry}; +use hipfire_hardware::{DeviceMesh, MeshEpoch}; +use rdna_compute::{DType, Gpu, GpuTensor}; +use std::collections::HashMap; + +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 +/// 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. The container +/// 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, + projections: HashMap, + origin: Option, +} + +/// The only owner that may roll back resident allocations before publication. +/// +/// 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, +} + +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() + } +} + +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)) + } + + /// 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 + .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. 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) + } else { + Ok(()) + } + } +} + +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. 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 { + let key = WeightPlacementKey::new(name, layer, device); + self.projections.remove(&key); + self.placements.remove(&key) + } + + 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)) + } + + 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 read-only seam cannot release or extract any handle. + 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 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) -> hip_bridge::HipResult<()> { + self.release_unchecked(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 { + 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(()), + } + } +} + +/// 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() + ), + }) +} +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. +/// +/// 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_weight_layers(weights, n_layers) + .and_then(|_| 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), + }; + 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 + { + 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 { + 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) { + let error = 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 + ), + }; + 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) + { + 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; + } + + let (bytes, dtype) = match source(entry) { + Ok(value) => value, + Err(reason) => { + 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) { + let error = FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "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 + .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()) { + let error = 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 + ), + }; + return Err(rollback_fulfill_error(store, gpu, error)); + } + } + let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { + Ok(tensor) => tensor, + Err(error) => { + 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) { + 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() { + 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)) +} + +/// 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::{DTypeConstraint, 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().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); + 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().expect("single-device mesh construction cannot overflow"); + 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().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)) + .unwrap(); + { + let mut assembly = store.begin_assembly(); + assert_eq!(assembly.take("x", None, 0), Some(0)); + let guard = assembly.commit(); + assert!(guard.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().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)) + .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().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)) + .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)]) + .expect("small test mesh construction cannot overflow"); + 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]); + } + + #[test] + fn tied_projection_preserves_fulfilled_source_dtype() { + let Ok(gpu) = Gpu::init() else { + return; + }; + 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", + 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) + .expect("resident transaction rollback must free its HIP allocation"); + } + + #[test] + fn successful_single_fulfillment_commits_resident_projection() { + let Ok(gpu) = Gpu::init() else { + return; + }; + 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(); + assert_eq!(transaction.len(), 1); + assert!(matches!( + transaction.get("resident", None, 0), + Some(WeightHandle::Resident(tensor)) if tensor.dtype == DType::F32 + )); + assert_eq!( + transaction.projection("resident", None, 0).unwrap().dtype, + DType::F32 + ); + 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().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); + store + .stage_alias("resident", None, 0, "source", projection(DType::F16)) + .unwrap(); + let transaction = WeightLoadTransaction::new(store); + 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)); + assert!(transaction.projection("resident", None, 0).is_some()); + } + + #[test] + 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().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 expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); + let error = transaction.validate_origin_value(expected).unwrap_err(); + assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); + assert_eq!(transaction.len(), 1); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + 0, + "origin rejection must not free resident buffers" + ); + 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().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), + ]; + 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().expect("single-device mesh construction cannot overflow"); + 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().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), + ]; + 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); + } +} 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..629a1a3558 100644 --- a/crates/saddle-lab/examples/gpus_smoke.rs +++ b/crates/saddle-lab/examples/gpus_smoke.rs @@ -7,12 +7,13 @@ //! //! 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() { 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 +98,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 cdb2e09383..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_runtime::multi_gpu::Gpus; use std::path::Path; fn used_gb(gpus: &Gpus, baseline_free: &[(usize, usize)]) -> Vec { @@ -60,7 +60,8 @@ 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 3ae3336ca9..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_runtime::multi_gpu::Gpus; use rdna_compute::Gpu; use std::path::Path; @@ -84,7 +84,8 @@ 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)