turbospark-runtime 0.3.0

//! `gpt-oss`'s routed-expert pass (ROADMAP M5).
//!
//! Structurally `families/llama/moe.rs`: top-k experts, no shared expert, so
//! phase 2's residual input is a ZERO buffer and the routed sum is added to
//! the stream afterwards. It is written out rather than shared for one
//! reason, and the reason is the first six lines of the function --
//! **THE ROUTER BIAS IS ADDED HERE**, to the read-back logits, before the
//! top-k sees them.
//!
//! That placement is not incidental. llama.cpp's `SOFTMAX_WEIGHT` gating adds
//! the bias to the logits, selects top-k on the RAW biased logits, then
//! softmaxes over the selected k alone with no renormalization -- which is
//! `router_topk_gemma4`'s existing contract minus Gemma's `per_expert_scale`.
//! Since this port's top-k is already a host-side function of a read-back f32
//! buffer, the add costs no kernel and no dispatch. Adding it after the top-k
//! would select the wrong experts and still produce fluent text; adding it
//! after the softmax would be a different distribution over the right ones.
//!
//! THE CLAMPED SwiGLU AND THE PER-EXPERT BIASES NEED NOTHING FROM THIS FILE.
//! Both ride inside the MXFP4 pair: `encode_moe_phase{1,2}_any` selects
//! `Mxfp4Activation::GPT_OSS` for that layout and derives `has_bias` from
//! `offsets.gate_b != 0`, which is exactly what `moe_offsets_from_layout`
//! writes once the repack walk puts the bias planes in the blob.
//!
//! The slot order is the ROUTER'S RANKING, as in every other family, and that
//! is a correctness constraint rather than a style choice: phase 2 reduces
//! `blob[slot] * routing_w[slot]` in slot-index order and FP addition is not
//! associative, so whatever the slot order depends on, the generated bytes
//! depend on (AGENTS.md Gotcha 27).
//!
//! Takes a [`RoutedSlot`] since the chunked-prefill driver
//! (`prefill.rs`) pipelines this call across a micro-batch's tokens, the
//! same way `families/gemma4/moe.rs` and `families/llama/moe.rs` do. The
//! sequential decode path passes `RoutedSlot::sequential()`, so its bytes
//! cannot move.

use std::time::Instant;

use half::f16;

use crate::families::gptoss::RealGptOssState;
use crate::moe_prefill_pipeline::RoutedSlot;
use crate::real_forward_dispatch::{
    encode_moe_phase1_any, encode_moe_phase2_any, router_topk_gemma4,
};
use crate::real_forward_init::MappedResidency;
use crate::real_forward_layout::RoutedLayerLayout;
use crate::real_forward_types::{DecodeScratch, PhaseCounters, RealForwardError};
use crate::real_forward_utils::f16_slice_to_le_bytes;

#[allow(clippy::too_many_arguments)]
pub(crate) fn encode_gpt_oss_layer_moe(
    context: &mut gpu::MetalContext,
    pass: &gpu::PassEncoder,
    scratch: &DecodeScratch,
    state: &RealGptOssState,
    streamers: &mut [Option<streaming::PreadExpertStreamer>],
    slot_buffers: &[Vec<gpu::MetalBuffer>],
    mapped: &MappedResidency,
    routed_blobs: Option<&gpu::RoutedBlobsBuffer>,
    routed_blobs_banks: &[gpu::RoutedBlobsBuffer],
    moe_offsets: &[gpu::MoeExpertOffsets],
    routed_layouts: &[RoutedLayerLayout],
    router_hist: &mut Option<crate::router_hist::RouterHistogram>,
    phases: &mut PhaseCounters,
    layer: usize,
    hidden: usize,
    moe_inter: u32,
    num_experts: usize,
    top_k: usize,
    slot: &RoutedSlot,
) -> Result<Vec<usize>, RealForwardError> {
    let gpu_err = RealForwardError::Gpu;
    let x_off = (slot.token * hidden * 2) as u64;
    let rw_off = slot.bank * gpu::MAX_STREAMED_EXPERTS * 2;
    let t_router = Instant::now();
    let mut router_logits = gpu::read_f32_buffer_at(
        &state.router_logits_f32,
        slot.token * num_experts,
        num_experts,
    );
    // THE BIAS, BEFORE THE TOP-K. See the module header.
    for (logit, bias) in router_logits.iter_mut().zip(&state.router_bias[layer]) {
        *logit += bias;
    }
    let (selected, route_weights) =
        router_topk_gemma4(&router_logits, top_k, &state.per_expert_ones);
    if let Some(hist) = router_hist.as_mut() {
        hist.record(layer, &selected);
    }
    // MAPPED RESIDENCY SKIPS BOTH THE PLAN AND THE `pread`, exactly as
    // `families/gemma4/moe.rs` does: every expert is already addressable in
    // its layer's mapping, so there is no cache to consult and nothing to
    // copy. Counted as requests that all hit, so `TURBOSPARK_PHASES`'s hit
    // rate stays comparable across residency modes.
    let mapped_active = mapped.buffers.get(layer).is_some_and(Option::is_some);
    let slots: Vec<usize> = if mapped_active {
        phases.expert_requests += selected.len() as u64;
        phases.expert_hits += selected.len() as u64;
        phases.router_nanos += t_router.elapsed().as_nanos() as u64;
        selected.clone()
    } else {
        let streamer = streamers[layer].as_mut().ok_or_else(|| {
            RealForwardError::Unsupported(format!(
                "gpt-oss layer {layer} has no packed-expert streamer"
            ))
        })?;
        // `protect` is empty on the decode path, so this is the same call it
        // has always made. Inside a chunk it names the slots the previous
        // token's in-flight command buffer is reading; the cache ASSERTS
        // rather than degrades when it cannot honour that plus the misses
        // (`ExpertCache::plan_if_possible`), so the caller has already
        // ensured the arithmetic works or retired the buffer first.
        let plan = streamer.plan_experts_cached(&selected, &slot.protect);
        phases.expert_requests += plan.experts.len() as u64;
        phases.expert_hits += plan.hits as u64;
        phases.router_nanos += t_router.elapsed().as_nanos() as u64;

        let streamer = streamers[layer]
            .as_mut()
            .expect("streamer presence checked above");
        let t_io = Instant::now();
        let slots = streamer
            .execute_expert_cache_plan(&plan)
            .map_err(|e| RealForwardError::Unsupported(format!("expert stream: {e}")))?;
        phases.expert_io_nanos += t_io.elapsed().as_nanos() as u64;
        slots
    };

    let t_bind = Instant::now();
    let ordered: Vec<(usize, f32)> = (0..selected.len())
        .map(|i| (slots[i], route_weights[i]))
        .collect();
    let mut routing16 = vec![f16::from_f32(0.0); gpu::MAX_STREAMED_EXPERTS];
    for (dispatch_slot, &(_, weight)) in ordered.iter().enumerate() {
        routing16[dispatch_slot] = f16::from_f32(weight);
    }
    gpu::write_buffer_bytes(
        &scratch.routing_w,
        rw_off,
        &f16_slice_to_le_bytes(&routing16),
    );

    // The two arms differ only in WHERE a blob lives, never in the order the
    // slots are dispatched: that stays the router's own ranking in both,
    // which is what keeps output byte-identical across residency modes
    // (AGENTS.md Gotcha 27), matching `families/gemma4/moe.rs`.
    let blob_refs: Vec<(&gpu::MetalBuffer, u64)> = if mapped_active {
        let buffer = mapped.buffers[layer]
            .as_ref()
            .expect("mapped residency checked above");
        let mapping = mapped.layers[layer]
            .as_ref()
            .expect("mapped residency checked above");
        ordered
            .iter()
            .map(|&(expert, _)| {
                mapping
                    .expert_offset(expert)
                    .map(|offset| (buffer, offset))
                    .map_err(|e| {
                        RealForwardError::Unsupported(format!(
                            "mapped expert layer {layer}, expert {expert}: {e}"
                        ))
                    })
            })
            .collect::<Result<Vec<_>, _>>()?
    } else {
        let layer_slots = &slot_buffers[layer];
        ordered
            .iter()
            .map(|&(cache_slot, _)| (&layer_slots[cache_slot], 0u64))
            .collect()
    };
    let routed = match slot.bank {
        0 => routed_blobs,
        n => routed_blobs_banks.get(n - 1),
    }
    .ok_or_else(|| {
        RealForwardError::Unsupported("install has no routed-blob buffer".to_string())
    })?;
    let offsets = &moe_offsets[layer];
    // `use_silu` is what the two OLDER activations are selected by, and it is
    // dead for this family: the MXFP4 pair takes its activation from
    // `Mxfp4Activation::GPT_OSS` regardless. Passed `true` to match
    // `hidden_activation`, so a layout that somehow resolved to a non-MXFP4
    // kernel gets silu rather than gelu-tanh, which is the nearer wrong
    // answer.
    let (blob_source, blob_function) = routed_layouts[layer].phase1.source_function();
    routed
        .bind_for(context, blob_source, blob_function, true, &blob_refs)
        .map_err(gpu_err)?;
    phases.bind_nanos += t_bind.elapsed().as_nanos() as u64;

    for &(buffer, _) in &blob_refs {
        pass.use_read_buffer(buffer);
    }

    encode_moe_phase1_any(
        routed_layouts[layer].phase1,
        context,
        pass,
        routed,
        offsets,
        (&state.moe_x, x_off),
        (&scratch.moe_acts, 0),
        hidden as u32,
        moe_inter,
        ordered.len() as u32,
        true,
    )
    .map_err(gpu_err)?;
    // `zero_hidden` is the residual input: there is no shared expert to carry
    // one, so phase 2's `y = residual + sum(w * expert)` becomes the routed
    // sum alone and the add below is the only place the stream grows. Passing
    // `x` here would add the residual twice.
    encode_moe_phase2_any(
        routed_layouts[layer].phase2,
        context,
        pass,
        routed,
        offsets,
        (&scratch.moe_acts, 0),
        (&scratch.routing_w, rw_off as u64),
        (&scratch.zero_hidden, 0),
        (&state.h2, 0),
        hidden as u32,
        moe_inter,
        top_k as u32,
        true,
    )
    .map_err(gpu_err)?;
    gpu::encode_residual_add(
        context,
        pass,
        (&scratch.x, x_off),
        (&state.h2, 0),
        hidden as u32,
    )
    .map_err(gpu_err)?;

    // The cache slots this pass BOUND, so the caller can hand them to the
    // next token as `RoutedSlot::protect` while this command buffer is in
    // flight. Returned rather than recomputed, matching
    // `encode_gemma4_layer_routed_moe`'s exact reasoning.
    Ok(ordered.iter().map(|&(cache_slot, _)| cache_slot).collect())
}

impl crate::real_forward::RealForwardRunner {
    /// Pipelined routed MoE loop across micro-batch tokens for one layer.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn encode_gpt_oss_layer_routed_moe_pipelined(
        &mut self,
        pending_routed: &mut Option<gpu::CommittedPass>,
        layer: usize,
        hidden: usize,
        moe_inter: u32,
        num_experts: usize,
        top_k: usize,
        banks: usize,
        m: usize,
    ) -> Result<(), RealForwardError> {
        let mut previous_slots: std::collections::HashSet<usize> = std::collections::HashSet::new();
        for t in 0..m {
            if banks == 1 {
                self.retire_routed(pending_routed);
            }
            let slot = RoutedSlot {
                token: t,
                bank: t % banks,
                // Empty when pipelining is off: retire-before-plan has
                // already run (banks == 1 above), so no slot is in
                // flight for the plan to avoid.
                protect: if self.routed_pipeline && banks > 1 {
                    previous_slots.clone()
                } else {
                    std::collections::HashSet::new()
                },
            };
            let routed_pass = self.context.begin_pass_labeled("gpt-oss routed cb");
            let (
                context,
                scratch,
                state,
                streamers,
                slot_buffers,
                mapped,
                routed_blobs,
                routed_blobs_banks,
                moe_offsets,
                routed_layouts,
                router_hist,
                phases,
            ) = (
                &mut self.context,
                &self.scratch,
                self.real_gpt_oss
                    .as_ref()
                    .expect("real gpt-oss state present"),
                &mut self.streamers,
                &self.slot_buffers,
                &self.mapped,
                &self.routed_blobs,
                &self.routed_blobs_banks,
                &self.moe_offsets,
                &self.routed_layouts,
                &mut self.router_hist,
                &mut self.phases,
            );
            let used = encode_gpt_oss_layer_moe(
                context,
                &routed_pass,
                scratch,
                state,
                streamers,
                slot_buffers,
                mapped,
                routed_blobs.as_ref(),
                routed_blobs_banks,
                moe_offsets,
                routed_layouts,
                router_hist,
                phases,
                layer,
                hidden,
                moe_inter,
                num_experts,
                top_k,
                &slot,
            )?;

            // This token's OWN row, matching Gemma 4's and llama-MoE's
            // chunked drivers.
            let x_off = (t * hidden * 2) as u64;
            crate::steering::encode_steering(
                &mut self.context,
                &routed_pass,
                &self.scratch,
                self.steering.as_ref(),
                layer,
                hidden,
                1,
                x_off,
            )?;
            crate::resid_capture::encode_resid_capture(
                &mut self.context,
                &routed_pass,
                &self.scratch,
                self.resid_capture.as_ref(),
                layer,
                hidden,
                x_off,
            )?;

            if banks > 1 {
                self.retire_routed(pending_routed);
            }
            debug_assert!(pending_routed.is_none(), "routed pipeline depth is 1");
            *pending_routed = Some(routed_pass.commit());
            previous_slots = used.into_iter().collect();
        }
        Ok(())
    }
}