moe_lora_grouped_down

Function moe_lora_grouped_down 

Source
pub fn moe_lora_grouped_down(
    gpu: &dyn GpuBackend,
    kernels: &LoraKernels,
    route: &MoeExpertRoute,
    x: DevicePtr,
    base_out: DevicePtr,
    expert_offsets: DevicePtr,
    sorted_token_ids: DevicePtr,
    moe_row_adapter: DevicePtr,
    xa: DevicePtr,
    row_offset: u32,
    row_end: u32,
    x_gather: u32,
    stream: u64,
) -> Result<()>
Expand description

Launch the device-side grouped fold for ONE chunk window [row_offset, row_end) of the sorted rows. Down (x_gather==0): x = post-SiLU sorted activations ([te, k_in] BF16), base_out = sorted expert_down_out. Gate/up (x_gather==1): x = the TOKEN-MAJOR expert_input ([num_tokens, k_in=hidden] BF16, gathered per sorted row via sorted_token_ids), base_out = sorted expert_gate_out/expert_up_out ([te, n_out=inter]). In both, base_out is [te, n_out] BF16 folded IN PLACE, expert_offsets = the device [num_experts+1] i32 prefix sum, sorted_token_ids = the device [te] i32 sorted-row→token map, moe_row_adapter = [num_tokens] i32 device map (< 0 = base skip) or DevicePtr::NULL for the single-active-adapter path, xa = the fixed-address [cap, max_rank] BF16 shrink scratch indexed at the LOCAL row r-row_offset (so the caller only needs >= (row_end-row_offset) rows, NOT >= te). The hooks loop [0, te) in windows of cap; a single call at row_offset=0, row_end=te (te <= cap) is bit-identical to the pre-chunk kernel.

ARG ORDER is in lockstep with moe_lora_grouped_down.cu (cuLaunchKernel is type-blind; the byte-identity oracle is the only guard — keep both in sync). row_offset/row_end are appended LAST in both kernels, so existing arg offsets are untouched.