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.