feat(dispatch): seal collective-aware MoE step execution

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",

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" }

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,

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())

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<usize> = (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<usize> = (0..n).collect();
     for l in 0..cfg.num_hidden_layers {
         {
             let mut binds: Vec<Deepseek4Bindings> = 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(())
 }
-

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],

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<usize> = (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(())
 }
-

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" }