metrale_model_arch/glm5next_mlp/
forward_prefill_gemm.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Routed-MoE prefill for GLM-5.3 through the shared grouped W4A16 GEMM
4//! (`moe_w4a16_grouped_gemm_ptrtable` and its tile variants): one launch per projection for all
5//! experts. `tile.rs` holds the tile table and the env levers, `dispatch.rs` the launch.
6//!
7//! The GEMM stages each dequantised weight as BF16 for `mma.sync`; the GEMV path keeps it in
8//! FP32, so the two paths do not produce the same bits. `forward_moe` takes this path only for
9//! row groups wider than `MOE_ROW_BATCH_MAX_ROWS` and at least `prefill_gemm_min_rows()` wide.
10//! Measured 2026-09-22 with `examples/glm5next_moe_grouped_prefill_microtest.rs` (M=64, top_k=8,
11//! 16 experts, N=256, K=512): against an FP32 reference that rounds each dequantised weight to
12//! BF16 first, the grouped GEMM's max_rel was 0.003888; the GEMV's against the exact FP32
13//! reference was 0.003889.
14//!
15//! Owner: model-arch (GLM-5.3).
16//! Invariants:
17//! - The GEMM reads the same `Glm5NextMoePtrTables` as the GEMV paths, indexed by GLOBAL expert
18//!   id, and writes nothing for an expert whose `packed` pointer is null.
19//! - Every buffer the path uses is in `Glm5NextMlpWorkspace`; the path allocates nothing.
20
21mod dispatch;
22mod tile;
23
24pub(crate) use tile::*;
25
26pub(super) use dispatch::forward_moe_grouped_prefill;
27
28#[cfg(test)]
29mod tests;