Module moe_lora_grouped

Module moe_lora_grouped 

Source
Expand description

2026-09-25: Launchers for the device-side MoE expert LoRA folds: moe_lora_grouped_down over the expert-sorted prefill rows and moe_lora_gather_bgmv over the slot-major decode rows, plus the host packing of their per-expert tables.

Both read the expert routing (expert_offsets or indices) on the device. The host loop crate::lora::expert_apply::apply_expert_lora_sorted needs a host copy of expert_offsets and one launch per adapted expert, which a CUDA graph cannot capture.

Owner: model-layers ops. Invariants:

  • An expert whose a_table or b_table cell is 0, or whose id is at or above the table length, folds nothing (moe_lora_grouped_down.cu, moe_lora_gather_bgmv.cu).

Structs§

ExpertTables
2026-09-25: Host-side tables from pack_expert_tables, before upload.
MoeExpertRoute
2026-09-25: Device tables for one projection’s expert LoRA pairs, indexed by expert id: the expert-keyed counterpart of the slot-keyed super::lora_delta::LoraRoute. MoeLayer::build_expert_route builds one per projection (down, gate, up) at adapter install, and the router fold builds a one-entry route whose expert 0 is the router pair.

Constants§

MLG_M_TILE
2026-09-25: Rows per grouped-fold tile. It must equal the MLG_M_TILE #define in moe_lora_grouped_down.cu, or the host grid and the kernel’s tiles cover different rows.

Functions§

gather_bgmv_grids
2026-09-25: (shrink, expand) grids of the gather fold: [ceil(out/4), n_slots, 1], where out is max_rank for the shrink and n_out for the expand. Each 256-thread block computes 4 outputs, 64 threads per output (GBGMV_N_PER_BLOCK in moe_lora_gather_bgmv.cu).
gather_row_token
2026-09-25: The token that owns flat (token, slot) row row: the gather kernels’ row / top_k.
grouped_down_wc
2026-09-25: grid.y of both grouped-fold kernels for the window [row_offset, row_end): ceil((row_end - row_offset) / MLG_M_TILE), at least 1. It is the same expression as the launcher’s wc, kept separate so tests can check the window math without a GPU.
grouped_down_windows
2026-09-25: Contiguous [start, end) windows of at most cap rows covering 0..total_rows. Panics when cap is 0.
moe_lora_gather_bgmv
2026-09-25: Fold base_out[row] += scale_e * (x[x_row] @ A_e^T) @ B_e^T for every flat (token, slot) row in [0, n_slots), where e = indices[row]. This is the unsorted counterpart of moe_lora_grouped_down: the expert comes from indices instead of an expert_offsets span.
moe_lora_grouped_down
2026-09-25: Fold base_out[r] += scale_e * (x[row(r)] @ A_e^T) @ B_e^T for the sorted rows r in the window [row_offset, row_end), where e is the expert whose expert_offsets span holds r.
pack_expert_tables
2026-09-25: Pack (expert_id, a_addr, b_addr, scale) entries into dense tables indexed by expert id, with 0 / 0.0 at every unadapted id and n_experts = max expert_id + 1. Returns None when entries is empty. A duplicate expert id keeps its last entry.