memra_engine/

parallel.rs

//! Model-specific parallel topology contracts.
//!
//! The rank planner is reusable, but model support is never inferred from a loader or a few
//! scalar dimensions. Each family must register the complete geometry that its TP/EP program
//! shards. Step-3.7-Flash is the first registered contract because its query-head count varies by
//! layer (64 full-attention / 96 sliding-attention), while KV heads stay at 8. Step-3.5 and other
//! siblings do not inherit this contract merely because they share the `step35` architecture tag.

use std::fmt;
use std::ops::Range;

use memra_gguf::config::{Arch, ModelConfig};

/// The execution planner's supported rank envelope. Hardware qualification and tuned defaults
/// remain model x rig evidence, but the placement/runtime contract must not stop at the three
/// cards currently available on Pod B.
pub const PRODUCT_MAX_CARDS: usize = 8;
const STEP_FP8_BLOCK: usize = 128;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HardwareTarget {
    Rtx5090,
    RtxPro6000Blackwell,
}

impl HardwareTarget {
    fn max_cards(self) -> usize {
        match self {
            Self::Rtx5090 => 1,
            Self::RtxPro6000Blackwell => PRODUCT_MAX_CARDS,
        }
    }

    fn label(self) -> &'static str {
        match self {
            Self::Rtx5090 => "rtx-5090",
            Self::RtxPro6000Blackwell => "rtx-pro-6000-blackwell",
        }
    }

    fn from_device_name(name: &str) -> Result<Self, TopologyError> {
        if name.contains("RTX PRO 6000") && name.contains("Blackwell") {
            return Ok(Self::RtxPro6000Blackwell);
        }
        if name.contains("RTX 5090") {
            return Ok(Self::Rtx5090);
        }
        Err(TopologyError::new(format!(
            "unqualified CUDA device {name:?}; first-class targets are RTX 5090 and RTX PRO 6000 \
             Blackwell"
        )))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TopologyRequest {
    pub pipeline: usize,
    pub tensor: usize,
    /// Routed experts are partitioned across the TP group. When false, each rank owns every
    /// expert and tensor-shards the expert projections instead.
    pub expert_parallel: bool,
    pub available_devices: usize,
    pub hardware: HardwareTarget,
}

impl TopologyRequest {
    pub fn world_size(self) -> Result<usize, TopologyError> {
        self.pipeline
            .checked_mul(self.tensor)
            .ok_or_else(|| TopologyError::new("PP x TP world size overflow"))
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelParallelContract {
    pub family: &'static str,
    pub variant: String,
    pub trunk_layers: usize,
    pub mtp_layers: usize,
    pub hidden_size: usize,
    pub vocab_size: usize,
    pub dense_ffn_size: usize,
    pub dense_prefix_layers: usize,
    pub head_dim: usize,
    pub query_heads: Vec<usize>,
    pub kv_heads: Vec<usize>,
    pub expert_count: usize,
    pub experts_per_token: usize,
    pub expert_ffn_size: usize,
    pub shared_expert_ffn_size: usize,
    pub hardware_targets: Vec<HardwareTarget>,
}

impl ModelParallelContract {
    /// Build the model-specific contract. Unregistered families refuse rather than inheriting a
    /// generic transformer assumption.
    pub fn from_model(cfg: &ModelConfig) -> Result<Self, TopologyError> {
        match cfg.arch {
            Arch::Step35 => Self::step35(cfg),
            _ => Err(TopologyError::new(format!(
                "no parallel contract registered for model family {:?}; loading/running does not \
                 establish TP/EP support",
                cfg.arch
            ))),
        }
    }

    fn step35(cfg: &ModelConfig) -> Result<Self, TopologyError> {
        let step = cfg.step35.as_ref().ok_or_else(|| {
            TopologyError::new(
                "step35 parallel contract requires the model-specific per-layer Step geometry",
            )
        })?;
        let total_layers = cfg.n_layer as usize;
        let mtp_layers = cfg.nextn_predict_layers as usize;
        let trunk_layers = total_layers.checked_sub(mtp_layers).ok_or_else(|| {
            TopologyError::new(format!(
                "step35 layer geometry invalid: total={total_layers} mtp={mtp_layers}"
            ))
        })?;
        if trunk_layers == 0 {
            return Err(TopologyError::new("step35 contract has no trunk layers"));
        }
        if step.head_count.len() < total_layers || step.head_count_kv.len() < total_layers {
            return Err(TopologyError::new(format!(
                "step35 per-layer head geometry incomplete: q={} kv={} need={total_layers}",
                step.head_count.len(),
                step.head_count_kv.len()
            )));
        }
        let moe = cfg.moe.as_ref().ok_or_else(|| {
            TopologyError::new("step35 parallel contract requires routed-expert geometry")
        })?;
        let query_heads: Vec<usize> = (0..total_layers)
            .map(|il| cfg.n_head_at(il as u32) as usize)
            .collect();
        let kv_heads: Vec<usize> = (0..total_layers)
            .map(|il| cfg.n_head_kv_at(il as u32) as usize)
            .collect();
        let is_step37 = trunk_layers == 45
            && mtp_layers == 3
            && cfg.n_embd == 4096
            && cfg.n_ff == 11_264
            && cfg.n_vocab == 128_896
            && query_heads
                .iter()
                .enumerate()
                .all(|(il, &heads)| heads == if il % 4 == 0 { 64 } else { 96 })
            && kv_heads.iter().all(|&heads| heads == 8)
            && moe.expert_count == 288
            && moe.expert_used_count == 8
            && moe.expert_ff_length == 1280
            && moe.expert_shared_ff_length == 1280
            && step.first_k_dense_replace == 3;
        if !is_step37 {
            return Err(TopologyError::new(format!(
                "no qualified parallel contract for step35 variant {:?}: only the exact \
                 Step-3.7-Flash geometry is registered",
                cfg.name
            )));
        }

        Ok(Self {
            family: "step35",
            variant: "Step-3.7-Flash-FP8".to_string(),
            trunk_layers,
            mtp_layers,
            hidden_size: cfg.n_embd as usize,
            vocab_size: cfg.n_vocab as usize,
            dense_ffn_size: cfg.n_ff as usize,
            dense_prefix_layers: step.first_k_dense_replace as usize,
            head_dim: cfg.head_dim_k as usize,
            query_heads,
            kv_heads,
            expert_count: moe.expert_count as usize,
            experts_per_token: moe.expert_used_count as usize,
            expert_ffn_size: moe.expert_ff_length as usize,
            shared_expert_ffn_size: moe.expert_shared_ff_length as usize,
            hardware_targets: vec![HardwareTarget::RtxPro6000Blackwell],
        })
    }

    pub fn plan(&self, request: TopologyRequest) -> Result<ParallelPlan, TopologyError> {
        let pp = request.pipeline;
        let tp = request.tensor;
        if !(1..=PRODUCT_MAX_CARDS).contains(&pp) {
            return Err(TopologyError::new(format!(
                "PP size {pp} outside product range 1..={PRODUCT_MAX_CARDS}"
            )));
        }
        if !(1..=PRODUCT_MAX_CARDS).contains(&tp) {
            return Err(TopologyError::new(format!(
                "TP size {tp} outside product range 1..={PRODUCT_MAX_CARDS}"
            )));
        }
        let world = request.world_size()?;
        if world > PRODUCT_MAX_CARDS {
            return Err(TopologyError::new(format!(
                "PP={pp} x TP={tp} requires {world} cards; product envelope is \
                 {PRODUCT_MAX_CARDS}"
            )));
        }
        if !self.hardware_targets.contains(&request.hardware) {
            return Err(TopologyError::new(format!(
                "{} has no qualified {} contract",
                self.variant,
                request.hardware.label()
            )));
        }
        if world > request.hardware.max_cards() {
            return Err(TopologyError::new(format!(
                "{} target permits at most {} card(s), requested {world}",
                request.hardware.label(),
                request.hardware.max_cards()
            )));
        }
        if request.available_devices < world {
            return Err(TopologyError::new(format!(
                "PP={pp} x TP={tp} requires {world} cards, only {} available",
                request.available_devices
            )));
        }
        if pp > self.trunk_layers {
            return Err(TopologyError::new(format!(
                "PP={pp} exceeds {} trunk layers",
                self.trunk_layers
            )));
        }
        if request.expert_parallel && tp == 1 {
            return Err(TopologyError::new(
                "expert parallelism requires TP group size greater than one",
            ));
        }

        // Check the family-specific, per-layer attention geometry before generic dimensions so a
        // refused topology names the model program that actually makes it invalid.
        for (il, (&q, &kv)) in self.query_heads.iter().zip(&self.kv_heads).enumerate() {
            require_divisible(&format!("layer {il} query heads"), q, tp)?;
            require_divisible(&format!("layer {il} KV heads"), kv, tp)?;
        }
        require_divisible("hidden size", self.hidden_size, tp)?;
        require_divisible("vocabulary size", self.vocab_size, tp)?;
        require_fp8_block_shard("dense FFN size", self.dense_ffn_size, tp)?;
        if request.expert_parallel {
            require_divisible("routed expert count", self.expert_count, tp)?;
        } else {
            require_fp8_block_shard("routed expert FFN size", self.expert_ffn_size, tp)?;
        }

        let stage_ranges = (0..pp)
            .map(|stage| stage * self.trunk_layers / pp..(stage + 1) * self.trunk_layers / pp)
            .collect();

        Ok(ParallelPlan {
            contract: self.clone(),
            request,
            world_size: world,
            stage_ranges,
            mtp_owner_stage: self.mtp_layers.gt(&0).then_some(pp - 1),
            // Step's 1280-wide shared expert cannot be split four or eight ways without cutting
            // through checkpoint 128-row E4M3 scale blocks. Replication is the exact program for
            // those TP/EP layouts; only routed experts are distributed.
            shared_expert_replicated: tp > 1 && self.shared_expert_ffn_size > 0,
        })
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParallelPlan {
    pub contract: ModelParallelContract,
    pub request: TopologyRequest,
    pub world_size: usize,
    pub stage_ranges: Vec<Range<usize>>,
    /// MTP layers are not pipeline stages of their own; the final PP stage owns them.
    pub mtp_owner_stage: Option<usize>,
    pub shared_expert_replicated: bool,
}

impl ParallelPlan {
    pub fn global_rank(&self, pipeline_rank: usize, tensor_rank: usize) -> Option<usize> {
        if pipeline_rank >= self.request.pipeline || tensor_rank >= self.request.tensor {
            return None;
        }
        Some(pipeline_rank * self.request.tensor + tensor_rank)
    }

    pub fn query_head_range(&self, layer: usize, tensor_rank: usize) -> Option<Range<usize>> {
        split_range(
            *self.contract.query_heads.get(layer)?,
            self.request.tensor,
            tensor_rank,
        )
    }

    pub fn kv_head_range(&self, layer: usize, tensor_rank: usize) -> Option<Range<usize>> {
        split_range(
            *self.contract.kv_heads.get(layer)?,
            self.request.tensor,
            tensor_rank,
        )
    }

    /// Column-parallel Q output range and the matching row-parallel O input range.
    pub fn query_feature_range(&self, layer: usize, tensor_rank: usize) -> Option<Range<usize>> {
        let heads = self.query_head_range(layer, tensor_rank)?;
        Some(heads.start * self.contract.head_dim..heads.end * self.contract.head_dim)
    }

    /// Column-parallel K/V output range. Step-3.7 has eight KV heads, so TP2 and TP4 partition
    /// them exactly; no KV-head replication is part of this registered contract.
    pub fn kv_feature_range(&self, layer: usize, tensor_rank: usize) -> Option<Range<usize>> {
        let heads = self.kv_head_range(layer, tensor_rank)?;
        Some(heads.start * self.contract.head_dim..heads.end * self.contract.head_dim)
    }

    /// Column-parallel dense gate/up output range and matching row-parallel down input range.
    pub fn dense_ffn_range(&self, tensor_rank: usize) -> Option<Range<usize>> {
        split_range(
            self.contract.dense_ffn_size,
            self.request.tensor,
            tensor_rank,
        )
    }

    pub fn routed_expert_range(&self, tensor_rank: usize) -> Option<Range<usize>> {
        self.request
            .expert_parallel
            .then(|| split_range(self.contract.expert_count, self.request.tensor, tensor_rank))?
    }

    pub fn routed_expert_ffn_range(&self, tensor_rank: usize) -> Option<Range<usize>> {
        (!self.request.expert_parallel).then(|| {
            split_range(
                self.contract.expert_ffn_size,
                self.request.tensor,
                tensor_rank,
            )
        })?
    }
}

/// Validate the live Step PP request before the loader allocates CUDA state. Checkpoint tensor
/// census is deliberately a separate loader gate: topology legality must remain testable without
/// opening model files, while serving requires both gates.
pub fn validate_step_pp_request(cfg: &ModelConfig) -> Result<Option<ParallelPlan>, TopologyError> {
    let pp = match std::env::var("MEMRA_PP_STAGES") {
        Err(_) => return Ok(None),
        Ok(value) if value.is_empty() || value == "0" || value == "1" => return Ok(None),
        Ok(value) => value.parse::<usize>().map_err(|_| {
            TopologyError::new(format!("MEMRA_PP_STAGES={value} is not a positive integer"))
        })?,
    };
    let devices = selected_pp_devices(pp)?;
    let hardware = detect_uniform_hardware(&devices)?;
    let contract = ModelParallelContract::from_model(cfg)?;
    let trunk_layers = contract.trunk_layers;
    let plan = contract.plan(TopologyRequest {
        pipeline: pp,
        tensor: 1,
        expert_parallel: false,
        available_devices: devices.len(),
        hardware,
    })?;
    let fence = crate::pp::pp_cuts(trunk_layers).ok_or_else(|| {
        TopologyError::new(format!(
            "Step PP={pp} has no valid runtime stage fence over {trunk_layers} trunk layers"
        ))
    })?;
    let plan = apply_stage_fence(plan, &fence)?;
    Ok(Some(plan))
}

fn apply_stage_fence(
    mut plan: ParallelPlan,
    fence: &[usize],
) -> Result<ParallelPlan, TopologyError> {
    let expected = plan.request.pipeline + 1;
    if fence.len() != expected
        || fence.first() != Some(&0)
        || fence.last() != Some(&plan.contract.trunk_layers)
        || fence.windows(2).any(|window| window[0] >= window[1])
    {
        return Err(TopologyError::new(format!(
            "invalid PP fence {fence:?} for {} stages over {} trunk layers",
            plan.request.pipeline, plan.contract.trunk_layers
        )));
    }
    plan.stage_ranges = fence
        .windows(2)
        .map(|window| window[0]..window[1])
        .collect();
    Ok(plan)
}

fn selected_pp_devices(pp: usize) -> Result<Vec<usize>, TopologyError> {
    let raw = std::env::var("MEMRA_PP_DEVICES").map_err(|_| {
        TopologyError::new(format!(
            "Step PP={pp} requires explicit MEMRA_PP_DEVICES with one distinct CUDA ordinal per \
             stage; same-device diagnostics do not qualify the multi-card product"
        ))
    })?;
    let devices: Result<Vec<usize>, _> = raw
        .split(',')
        .map(|part| part.trim().parse::<usize>())
        .collect();
    let devices = devices.map_err(|_| {
        TopologyError::new(format!(
            "MEMRA_PP_DEVICES={raw:?} is not a comma-separated CUDA ordinal list"
        ))
    })?;
    if devices.len() != pp {
        return Err(TopologyError::new(format!(
            "MEMRA_PP_DEVICES lists {} devices but MEMRA_PP_STAGES={pp}",
            devices.len()
        )));
    }
    let mut unique = devices.clone();
    unique.sort_unstable();
    unique.dedup();
    if unique.len() != devices.len() {
        return Err(TopologyError::new(format!(
            "Step PP={pp} requires {pp} distinct devices; MEMRA_PP_DEVICES={raw:?} repeats an \
             ordinal"
        )));
    }
    Ok(devices)
}

fn detect_uniform_hardware(devices: &[usize]) -> Result<HardwareTarget, TopologyError> {
    cudarc::driver::result::init().map_err(|error| {
        TopologyError::new(format!("CUDA driver initialization failed: {error}"))
    })?;
    let mut target = None;
    for &ordinal in devices {
        let device = cudarc::driver::result::device::get(ordinal as i32).map_err(|error| {
            TopologyError::new(format!("CUDA device {ordinal} lookup failed: {error}"))
        })?;
        let name = cudarc::driver::result::device::get_name(device).map_err(|error| {
            TopologyError::new(format!("CUDA device {ordinal} name lookup failed: {error}"))
        })?;
        let current = HardwareTarget::from_device_name(&name)?;
        if let Some(expected) = target {
            if current != expected {
                return Err(TopologyError::new(format!(
                    "mixed hardware targets in MEMRA_PP_DEVICES: expected {}, device {ordinal} is \
                     {}",
                    expected.label(),
                    current.label()
                )));
            }
        } else {
            target = Some(current);
        }
    }
    target.ok_or_else(|| TopologyError::new("MEMRA_PP_DEVICES is empty"))
}

fn require_divisible(label: &str, value: usize, parts: usize) -> Result<(), TopologyError> {
    if value == 0 {
        return Err(TopologyError::new(format!("{label} is zero")));
    }
    if value % parts != 0 {
        return Err(TopologyError::new(format!(
            "{label} {value} is not divisible by TP={parts}"
        )));
    }
    Ok(())
}

fn require_fp8_block_shard(label: &str, value: usize, parts: usize) -> Result<(), TopologyError> {
    require_divisible(label, value, parts)?;
    let local = value / parts;
    if local % STEP_FP8_BLOCK != 0 {
        return Err(TopologyError::new(format!(
            "{label} shard {local} for TP={parts} cuts through the Step E4M3 block size \
             {STEP_FP8_BLOCK}"
        )));
    }
    Ok(())
}

fn split_range(total: usize, parts: usize, rank: usize) -> Option<Range<usize>> {
    if parts == 0 || rank >= parts || total % parts != 0 {
        return None;
    }
    let width = total / parts;
    Some(rank * width..(rank + 1) * width)
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TopologyError {
    message: String,
}

impl TopologyError {
    fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl fmt::Display for TopologyError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.message.fmt(f)
    }
}

impl std::error::Error for TopologyError {}

#[cfg(test)]
mod tests {
    use super::*;
    use memra_gguf::config::{MoeConfig, Step35Config};

    fn step37_contract() -> ModelParallelContract {
        let total_layers = 48;
        ModelParallelContract {
            family: "step35",
            variant: "Step-3.7-Flash-FP8".to_string(),
            trunk_layers: 45,
            mtp_layers: 3,
            hidden_size: 4096,
            vocab_size: 128_896,
            dense_ffn_size: 11_264,
            dense_prefix_layers: 3,
            head_dim: 128,
            query_heads: (0..total_layers)
                .map(|il| if il % 4 == 0 { 64 } else { 96 })
                .collect(),
            kv_heads: vec![8; total_layers],
            expert_count: 288,
            experts_per_token: 8,
            expert_ffn_size: 1280,
            shared_expert_ffn_size: 1280,
            hardware_targets: vec![HardwareTarget::RtxPro6000Blackwell],
        }
    }

    fn step37_model_config() -> ModelConfig {
        let total_layers = 48;
        let head_count: Vec<u32> = (0..total_layers)
            .map(|il| if il % 4 == 0 { 64 } else { 96 })
            .collect();
        ModelConfig {
            arch: Arch::Step35,
            name: "Step-3.7-Flash-FP8".to_string(),
            n_layer: total_layers,
            n_embd: 4096,
            n_head: 96,
            n_head_kv: 8,
            head_dim_k: 128,
            head_dim_v: 128,
            n_ff: 11_264,
            n_vocab: 128_896,
            context_length: 262_144,
            rms_eps: 1e-6,
            rope_freq_base: 5_000_000.0,
            rope_dim_count: 128,
            rope_sections: Vec::new(),
            full_attention_interval: 0,
            ssm: None,
            moe: Some(MoeConfig {
                expert_count: 288,
                expert_used_count: 8,
                expert_ff_length: 1280,
                expert_shared_ff_length: 1280,
            }),
            m3: None,
            hy3: None,
            gemma4: None,
            mla: None,
            step35: Some(Step35Config {
                head_count,
                head_count_kv: vec![8; total_layers as usize],
                swa_pattern: (0..total_layers).map(|il| il % 4 != 0).collect(),
                sliding_window: 512,
                rope_base_global: 5_000_000.0,
                rope_base_swa: 10_000.0,
                rope_dims_full: 64,
                rope_dims_swa: 128,
                swiglu_clamp_exp: vec![0.0; total_layers as usize],
                swiglu_clamp_shexp: vec![0.0; total_layers as usize],
                sigmoid_routing: true,
                routed_scaling_factor: 3.0,
                route_norm: true,
                first_k_dense_replace: 3,
            }),
            geometry: None,
            nextn_predict_layers: 3,
            n_layer_total: total_layers,
        }
    }

    fn request(pp: usize, tp: usize, expert_parallel: bool) -> TopologyRequest {
        TopologyRequest {
            pipeline: pp,
            tensor: tp,
            expert_parallel,
            available_devices: pp * tp,
            hardware: HardwareTarget::RtxPro6000Blackwell,
        }
    }

    #[test]
    fn step_pp3_maps_fifteen_trunk_layers_per_card() {
        let plan = step37_contract().plan(request(3, 1, false)).unwrap();
        assert_eq!(plan.world_size, 3);
        assert_eq!(plan.stage_ranges, vec![0..15, 15..30, 30..45]);
        assert_eq!(plan.mtp_owner_stage, Some(2));
    }

    #[test]
    fn step_pp_marker_uses_the_runtime_stage_fence() {
        let plan = step37_contract().plan(request(3, 1, false)).unwrap();
        let plan = apply_stage_fence(plan, &[0, 10, 28, 45]).unwrap();
        assert_eq!(plan.stage_ranges, vec![0..10, 10..28, 28..45]);
    }

    #[test]
    fn step_contract_is_extracted_from_model_specific_geometry() {
        let contract = ModelParallelContract::from_model(&step37_model_config()).unwrap();
        assert_eq!(contract.family, "step35");
        assert_eq!(contract.trunk_layers, 45);
        assert_eq!(contract.mtp_layers, 3);
        assert_eq!(contract.query_heads[0], 64);
        assert_eq!(contract.query_heads[1], 96);
        assert_eq!(contract.kv_heads[47], 8);
        assert_eq!(contract.expert_count, 288);
        assert_eq!(contract.experts_per_token, 8);
    }

    #[test]
    fn step_sibling_does_not_inherit_the_step37_contract() {
        let mut sibling = step37_model_config();
        sibling.name = "Step-3.5-Flash".to_string();
        sibling.n_vocab = 128_000;
        let error = ModelParallelContract::from_model(&sibling).unwrap_err();
        assert!(
            error
                .to_string()
                .contains("only the exact Step-3.7-Flash geometry is registered")
        );
    }

    #[test]
    fn step_without_the_official_mtp_geometry_does_not_inherit_the_contract() {
        let mut stripped = step37_model_config();
        stripped.nextn_predict_layers = 0;
        let error = ModelParallelContract::from_model(&stripped).unwrap_err();
        assert!(
            error
                .to_string()
                .contains("only the exact Step-3.7-Flash geometry is registered")
        );
    }

    #[test]
    fn hardware_target_classification_is_exact() {
        assert_eq!(
            HardwareTarget::from_device_name("NVIDIA RTX PRO 6000 Blackwell Server Edition")
                .unwrap(),
            HardwareTarget::RtxPro6000Blackwell
        );
        assert_eq!(
            HardwareTarget::from_device_name("NVIDIA GeForce RTX 5090 Laptop GPU").unwrap(),
            HardwareTarget::Rtx5090
        );
        assert!(HardwareTarget::from_device_name("NVIDIA H100 80GB HBM3").is_err());
    }

    #[test]
    fn step_tp2_tp4_tp8_and_hybrid_plans_are_geometry_valid() {
        let tp2 = step37_contract().plan(request(1, 2, true)).unwrap();
        assert_eq!(tp2.query_head_range(0, 1), Some(32..64));
        assert_eq!(tp2.query_head_range(1, 1), Some(48..96));
        assert_eq!(tp2.kv_head_range(0, 1), Some(4..8));
        assert_eq!(tp2.routed_expert_range(1), Some(144..288));

        let tp4 = step37_contract().plan(request(1, 4, true)).unwrap();
        assert_eq!(tp4.query_head_range(0, 3), Some(48..64));
        assert_eq!(tp4.query_head_range(1, 3), Some(72..96));
        assert_eq!(tp4.kv_head_range(0, 3), Some(6..8));
        assert_eq!(tp4.query_feature_range(0, 3), Some(6144..8192));
        assert_eq!(tp4.query_feature_range(1, 3), Some(9216..12_288));
        assert_eq!(tp4.kv_feature_range(0, 3), Some(768..1024));
        assert_eq!(tp4.dense_ffn_range(3), Some(8448..11_264));
        assert_eq!(tp4.routed_expert_range(3), Some(216..288));
        assert!(tp4.shared_expert_replicated);

        let tp8 = step37_contract().plan(request(1, 8, true)).unwrap();
        assert_eq!(tp8.query_head_range(0, 7), Some(56..64));
        assert_eq!(tp8.query_head_range(1, 7), Some(84..96));
        assert_eq!(tp8.kv_head_range(0, 7), Some(7..8));
        assert_eq!(tp8.dense_ffn_range(7), Some(9856..11_264));
        assert_eq!(tp8.routed_expert_range(7), Some(252..288));
        assert!(tp8.shared_expert_replicated);

        let hybrid = step37_contract().plan(request(2, 4, true)).unwrap();
        assert_eq!(hybrid.world_size, 8);
        assert_eq!(hybrid.stage_ranges, vec![0..22, 22..45]);
        assert_eq!(hybrid.global_rank(1, 3), Some(7));
        assert_eq!(hybrid.global_rank(2, 0), None);
    }

    #[test]
    fn step_tp4_requires_whole_expert_parallelism() {
        let error = step37_contract().plan(request(1, 4, false)).unwrap_err();
        assert!(
            error
                .to_string()
                .contains("routed expert FFN size shard 320")
        );
        let tp2 = step37_contract().plan(request(1, 2, false)).unwrap();
        assert_eq!(tp2.routed_expert_ffn_range(1), Some(640..1280));
        assert!(tp2.shared_expert_replicated);
    }

    #[test]
    fn step_tp3_refuses_the_real_per_layer_head_geometry() {
        let error = step37_contract()
            .plan(TopologyRequest {
                pipeline: 1,
                tensor: 3,
                expert_parallel: true,
                available_devices: 3,
                hardware: HardwareTarget::RtxPro6000Blackwell,
            })
            .unwrap_err();
        assert!(error.to_string().contains("layer 0 query heads 64"));
    }

    #[test]
    fn product_envelope_accepts_eight_and_refuses_more() {
        let pp8 = step37_contract().plan(request(8, 1, false)).unwrap();
        assert_eq!(pp8.world_size, 8);
        assert_eq!(pp8.stage_ranges.len(), 8);
        assert!(pp8.stage_ranges.iter().all(|range| !range.is_empty()));

        let error = step37_contract()
            .plan(TopologyRequest {
                pipeline: 3,
                tensor: 4,
                expert_parallel: true,
                available_devices: 12,
                hardware: HardwareTarget::RtxPro6000Blackwell,
            })
            .unwrap_err();
        assert!(error.to_string().contains("product envelope is 8"));
    }

    #[test]
    fn expert_parallel_requires_a_multi_rank_tp_group() {
        let error = step37_contract().plan(request(3, 1, true)).unwrap_err();
        assert!(
            error
                .to_string()
                .contains("expert parallelism requires TP group size greater than one")
        );
    }

    #[test]
    fn step_does_not_inherit_the_5090_hardware_contract() {
        let error = step37_contract()
            .plan(TopologyRequest {
                pipeline: 1,
                tensor: 1,
                expert_parallel: false,
                available_devices: 1,
                hardware: HardwareTarget::Rtx5090,
            })
            .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("has no qualified rtx-5090 contract")
        );
    }
}