Module layers

Module layers 

Source
Expand description

2026-09-25: The layer types (FfnComponent and the modules below) and the kernel-handle resolvers several layers share.

Owner: model-layers. Invariants: none beyond the types.

Re-exports§

pub use dense_ffn::DenseFfnLayer;
pub use dense_ffn::DenseFfnWeights;
pub use dense_ffn::FfnActivation;
pub use glm_vit::GlmVit;
pub use glm_vit::GlmVitBlock;
pub use glm_vit::GlmVitMerger;
pub use moe::MoeLayer;
pub use moe::moe_expert_fp8_act;
pub use moe::set_moe_expert_fp8_act;
pub use mtp_head::MtpHead;
pub use mtp_head::MtpQuantization;
pub use mtp_head::mtp_drafter_prefill_enabled;
pub use mtp_head::set_mtp_experts_nvfp4_from_cli;
pub use qwen3_attention::Qwen3AttentionLayer;
pub use qwen3_ssm::Qwen3SsmLayer;
pub use vision_encoder::MergerLayer;
pub use vision_encoder::ViTBlock;
pub use vision_encoder::VisionEncoder;
pub use vision_tower::VisionTower;

Modules§

dense_ffn
2026-09-25: Dense gated FFN layer, down(act(gate(x)) * up(x)) with SiLU or GELU, for single-token decode, small verify batches and prefill.
ep_dispatch
2026-09-25: Partition of top-k (token, expert) assignments into this rank’s experts and the other ranks’ experts, for EP token dispatch.
fp8_calibration
2026-09-25: Online per-tensor FP8 KV scale calibration for one attention layer.
glm_vit
2026-09-25: The GLM-5.3-Flash vision tower: ViT blocks with per-head QK-RMSNorm, 2D axial RoPE and clamped SwiGLU, then a 2×2 strided-conv downsample and a SwiGLU merger.
moe
2026-09-25: The MoE feed-forward block MoeLayer: router, routed experts, optional shared expert, and the kernel handles and pointer tables its forward paths use.
mtp_head
2026-09-25: The MTP (multi-token prediction) draft head, a DraftProposer. One decoder layer drafts from the token embedding and the target’s hidden: pre-fc norms, concat, fc, gated attention over its own KV cache, MoE or dense FFN, final norm, the NVFP4 LM head and argmax. MtpQuantization selects the projection weights’ precision.
mtp_meta
2026-09-25: The one definition of an MTP proposer’s attention-metadata slab, packed by MtpHead::forward_one, forward_batch_position and DeepseekV4MtpHead::propose.
mtp_multi
2026-09-25: Multi-module MTP proposer: draft slot i runs its own MtpHead (own weights and KV cache), as modules[i].forward_one(previous draft token, previous module's hidden).
ngram_embed
2026-09-25: N-gram embedding (arXiv 2601.21204): the dimensions, the host-side polynomial-rolling-hash ids that select rows in the K * (N-1) lookup tables (ids), the tables (table) and the GPU embedding (embed). ids_match_python_reference checks the ids against bench/ngram_ref/ngram_id_fixtures.json, written by bench/ngram_ref/make_fixtures.py.
ops
2026-09-25: Kernel launch wrappers that the layers compose into forward passes, grouped into the submodules under ops/ and re-exported here (most by glob).
ple
2026-09-25: PLE, hashed n-gram injection into the hyper-connection highway: the layer (layer, with its per-sequence carry in layer/aux_state), its ids (ids) and debug taps (dump). The reference forward, Qwen4ExpTextPLELayer.forward in bench/qwen4_exp/ref/modeling_qwen4_exp.py:
qsa
2026-09-25: The Qwen3.8-Flash-Next QSA indexer: sparse-attention token selection.
qwen3_attention
2026-09-25: The full-attention transformer layer, Qwen3AttentionLayer: its construction, kernel requirements, decode and prefill attention, and the call into the layer’s FFN.
qwen3_ssm
2026-09-25: Qwen3SsmLayer, the Gated DeltaNet (linear-attention) layer, and the submodules its constructor, forward paths and tests live in.
vision_encoder
2026-09-25: The Qwen3-VL-shaped vision encoder: a ViT with optional DeepStack mergers.
vision_tower
2026-09-25: The vision tower a built model holds, whichever family bound it.
w4a16_gemv_tiers
2026-09-25: SSOT for the narrow w4a16_gemv_batch{M} tier family (M = 4..8).

Structs§

W8a8Ctx
2026-09-28: The kernels and the quantized-activation scratch one model’s W8A8 projections share.
W8a8Ffn
2026-09-28: A SiLU dense FFN’s gate, up and down projections.
W8a8Mixer
2026-09-28: The input and output projections of an attention layer (Q|K|V stacked, O) or a GDN layer (QKV|Z stacked, out_proj).

Enums§

ExpertQuantization
2026-09-27: The expert precision tiers. Measured on GB10 (Qwen3.6-35B-A3B-FP8, canonical tiers, --mtp-gate force):
FfnComponent
2026-09-25: A layer’s FFN: MoE, dense, or none.
MoeExpertTables
2026-10-02: The decision.
RowTiers
2026-09-27: The projection order a process runs.

Functions§

act_route
2026-09-30: The format family runs at rows rows (Adaptive: the site’s own routing).
activation_quantization
2026-09-30: The value in force: adaptive unless the serve published another.
any_fixed
2026-09-30: Whether any family runs a fixed format at any row count. The serve then publishes the canonical row tiers (row_tiers.rs), whose single-order kernels are the fixed bf16 path of the W8A16 projections and of an NVFP4 LM head.
expert_quantization
2026-09-27: The --expert-quantization tier in force: fp8 unless the serve published another.
family_fixed
2026-09-30: Whether family runs a fixed format at some row count.
fixed_act
2026-09-30: Some(format) when family runs one fixed activation format at rows rows, None when that row count is adaptive (the site’s own routing, unchanged).
k64_kernel
2026-09-25: The k64 deep-K tile GEMM. Prefers w4a16_gemm_t_k64_p3; takes w4a16_gemm_t_k64 when METRALE_NO_K64_PIPELINE3 is set to any value (0 included) or the target lacks _p3, and errors when that is missing too.
k64_n64_kernel
2026-09-25: The N_TILE=64 deep-K twin w4a16_gemm_t_k64_n64_p3. KernelHandle(0) when the target lacks it, or, with no lookup, when METRALE_NO_K64_N64 is set to any value (0 included).
k64_n64_wins
2026-09-25: Whether the N_TILE=64 twin serves an m × n launch (K64_N64_MAX_WIDE_CTAS).
kernel_caps
2026-09-28: Which decode families the policy may rely on. On: the dense per-channel W8A8 family (ops::w8a8_decode). 2026-09-29: On as well, validated on Qwen3.6-35B-A3B-FP8 by the full BFCL draw and agentic-webserver (declared against nvfp4, one binary, dgx3): the MoE expert W8A8 family (moe/fp8_grouped_tc_w8a8.rs). The block-scaled attention/GDN W8A8 adoption (qwen35/load_layers/w8a8_adopt.rs) is off again (2026-09-29, see below). Also on: the batched FP8 lm_head (model/lm_head_fp8_rows.rs, one weight pass per 128 rows at every row count), so declared takes a checkpoint’s declared FP8 head. Each W8A8 family still checks at run time that its kernels resolved and the shapes fit.
moe_expert_tables
2026-10-02: The published decision, if any.
moe_grouped_decode_decide
2026-09-25: Whether the arm runs for n rows, given the two switches. It reads no env, so tests do not depend on the latched OnceLocks above.
moe_grouped_decode_enabled
2026-09-25: False when METRALE_NO_MOE_GROUPED_DECODE is set to any value, 0 included. Read once per process.
moe_grouped_decode_for
2026-09-25: Whether the arm runs for n rows under this process’s switches.
moe_grouped_decode_forced
2026-09-25: METRALE_MOE_GROUPED_DECODE=1 runs the arm below moe_grouped_decode_min_rows() too, for measuring narrow widths. Read once per process.
moe_grouped_decode_min_rows
2026-09-25: Minimum rows for the grouped-GEMM MoE decode arm. The SSM stack (qwen3_ssm::trait_decode_multi_seq) and the attention layers’ multi-seq FFN both decide through moe_grouped_decode_for, so they share this width. The arm reads each routed expert once for all rows but first sorts and permutes the rows, a fixed cost that the threshold keeps off narrow batches.
prefill_row_invariant
2026-10-01: Whether a prefill pass picks its kernels without regard to its row count, so a row’s bits are the same in a long pass and in a short pass over the same rows (a warm request replaying a prompt’s tail). On under the canonical tiers: the FP8 MoE then takes the grouped prefill kernels at every row count (MoeLayer::forward_prefill), where it took the per-row decode kernels up to 64 rows. The row-count switches left in a prefill (W8A16 pipe128, the BA-gates twin, the router’s register tile, the gate/up _w1/_w2 entries) choose between kernels that give the same bits.
publish_row_tiers
2026-09-27: Publish the policy of the model about to be built. A model swap publishes again before the next build.
resolve_row_tiers
2026-09-27: A model’s policy: --no-canonical-tiers keeps the scalar order if METRALE_ROW_EXACT_TIERS is present and the by-rows tiers otherwise; else a present lever decides; else FP8 MoE checkpoints run canonical and the others by rows. The dense mixed FP8/NVFP4 checkpoints (Qwen3.8-27B) keep their tiers: canonical tiers cost them 2-4% of decode throughput.
row_invariant
2026-09-27: Whether this process runs a row-invariant policy.
row_tiers
2026-09-27: This process’s RowTiers: the published policy, else the two levers, read once.
row_tiers_from
2026-09-27: RowTiers from the presence of the two levers.
set_activation_quantization_from_cli
2026-09-30: Publish --activation-quantization. Returns the value in force; a caller that gets a different one should warn.
set_expert_quantization_from_cli
2026-09-27: Publish --expert-quantization. Returns the tier in force; a caller that gets a different one should warn.
set_moe_expert_tables_from_plan
2026-10-02: Publish the serve’s decision. Returns the decision in force; a caller that gets a different one should warn.
set_weight_quantization_from_cli
2026-09-28: Publish --weight-quantization and its lever. Returns the tier in force; a caller that gets a different one should warn.
tgemm_kernel
2026-09-25: The N128/M64 tile GEMM. Prefers w4a16_gemm_t_p3; takes w4a16_gemm_t when METRALE_NO_TGEMM_PIPELINE3 is set to any value (0 included) or the target lacks _p3.
tgemm_probe_ok
2026-09-25: False for kimi_k3 and kimi_linear. For those the model builder (model-engine model/impl_a1.rs) skips the tgemm_kernel lookup and keeps a zero handle.
try_kernel
2026-09-25: Look up module::func, returning KernelHandle(0) (and a debug log) when the lookup fails. A zero handle is a slower path or a missing feature, so a caller must check it before launching.
try_target_kernel
2026-09-25: Look up a kernel from a module only some targets compile. Returns KernelHandle(0) without a lookup, and so without an audit row, when the backend does not carry module; otherwise it is try_kernel.
w4a16_v2_kernel
2026-09-26: The w4a16_gemm_t_m128_v2 handle for the layer builders in dense_ffn_init.rs, qwen3_attention/init_prefill_kernels.rs and qwen3_ssm/init.rs.
w4a16_v3_kernel
2026-09-25: The w4a16_gemm_t_m128_v3 handle. KernelHandle(0), with no lookup and so no boot-audit row, unless METRALE_W4A16_VARIANT=v3; then it panics when the target lacks the kernel. Its user (qwen3_attention/prefill_weights.rs) launches it only for v == 3 and a nonzero handle.
weight_quantization
2026-09-28: The tier in force: declared unless the serve published another.