moe_weighted_sum_blend

Function moe_weighted_sum_blend 

Source
pub fn moe_weighted_sum_blend(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    output: DevicePtr,
    expert_out: DevicePtr,
    expert_weights: DevicePtr,
    shared_out: DevicePtr,
    input: DevicePtr,
    gate_weight: DevicePtr,
    hidden: u32,
    top_k: u32,
    k: u32,
    stream: u64,
) -> Result<()>
Expand description

Fused SiLU+down expert GEMV, wide variant (16 outputs/block for small K).

Same semantics as moe_expert_gemv_silu_down but 4x more outputs per block with sub-warp reduction. Optimal for K<=512 where the narrow kernel has insufficient inner loop iterations for memory latency hiding.

Fused weighted sum + sigmoid blend + gate scalar GEMV.

Computes gate_scalar = dot(input, gate_weight) inline, then: output[j] = sum_e weights[e] * expert_out[e,j] + sigmoid(gate_scalar) * shared_out[j]

Each block independently computes the gate scalar dot product (redundant but only 8KB per block for K=2048 — negligible). Eliminates the separate dense_gemv kernel for the shared expert gate scalar (saves 48 graph nodes).

Grid: (ceil(hidden/256), 1, 1) Block: (256, 1, 1)