Module moe_lora_grouped 

Device-side MoE expert down_proj LoRA fold launcher (moe_lora_grouped_down).

Replaces the host-synced per-expert loop (crate::lora::expert_apply, which D2H-copies expert_offsets and drives a host launch count — both illegal under CUDA-graph capture) with a single two-launch kernel that reads expert_offsets DEVICE-side. The grid is a STATIC worst-case bound (worst_case_m_tiles = ceil(te/64), matching the base grouped GEMM), so the launch shape is constant across capture/replay; per-tile early-return on an empty / out-of-range / unadapted expert span keeps it correct without a host value. See kernels/gb10/common/moe_lora_grouped_down.cu.

The fold math is BYTE-IDENTICAL to apply_lora_bgmv / per-row apply_lora_delta(m=1) (shrink→BF16 xa, expand→BF16 delta, then base += scale·fp32(bf16(delta))), so one kernel serves the nvfp4, bf16, and fp8 grouped prefill paths — all write the same sorted BF16 expert_down_out.

Structs

ExpertTables

Host-side packed tables from pack_expert_tables, ready for H2D upload.

MoeExpertRoute

Per-EXPERT routing tables for the grouped down fold — the expert-keyed analogue of the slot-keyed super::lora_delta::LoraRoute. Built once at adapter install from the layer’s Down pairs; load-time-fixed device addresses, so they are stable kernel args across capture/replay (adapter identity for a mixed batch flows through the per-row moe_row_adapter, not these tables).

Constants

MLG_M_TILE

M_TILE the static worst-case grid pairs with — must match the MLG_M_TILE #define in moe_lora_grouped_down.cu AND the base grouped GEMM’s worst_case_m_tiles = ceil(total_expanded/64) sizing.

Functions

gather_bgmv_grids

PURE (GPU-free, unit-tested): the exact (shrink, expand) grid triples for the decode gather-fold, given the route dims and the flat row count. Each is [ceil(out/4), n_slots, 1] with a (256,1,1) block (one 64-lane group per output, N_PER_BLOCK = 4). n_slots is a host constant per captured graph, so the shape is EXACT (no worst-case tiles) and capture-stable.

gather_row_token

PURE: the owning token index of a flat (token, slot) row — mirrors the kernel’s row / top_k so the per-token row_adapter gather is verifiable.

grouped_down_wc

PURE (GPU-free, unit-tested): the shrink/expand grid grid.y (m-tile count) for one chunk window [row_offset, row_end) — ceil((row_end-row_offset)/64), min 1. Matches the launcher’s wc; exposed so the chunk-boundary math is verifiable without a GPU. A full-window call (row_offset=0, row_end=te) returns exactly the pre-chunk ceil(te/64).

grouped_down_windows

Contiguous [start, end) launch windows covering 0..total_rows.

moe_lora_gather_bgmv

SOLID Incr-4: launch the DECODE-path MoE expert down fold. The unsorted, slot-major analogue of moe_lora_grouped_down — instead of an expert_offsets prefix sum over sorted rows, each flat (token, slot) row gathers its expert from indices[row] (the same indices_dev the fused expert GEMV routed on) and its base/adapt decision from row_adapter[row / top_k] (< 0 = base skip, or DevicePtr::NULL to fold every row on the single-active-adapter path).

moe_lora_grouped_down

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.

pack_expert_tables

PURE (GPU-free, unit-tested): pack a set of adapted-expert (expert_id, a_addr, b_addr, scale) entries into dense [n_experts] tables indexed by expert id, with 0 / 0.0 at every unadapted slot. n_experts = max expert_id + 1. Returns None when entries is empty (a router-only adapter installs no expert route). Duplicate expert ids keep the LAST entry (callers pass at most one Down pair per expert).