metrale_model_layers/layers/
ops.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Kernel launch wrappers that the layers compose into forward
4//! passes, grouped into the submodules under `ops/` and re-exported here
5//! (most by glob).
6//!
7//! Owner: model-layers (ops).
8//! Invariants: none beyond the types.
9
10#[path = "ops/activations.rs"]
11mod activations;
12#[path = "ops/derived_weights.rs"]
13mod derived_weights;
14#[path = "ops/dispatch_config.rs"]
15mod dispatch_config;
16#[cfg(test)]
17#[path = "ops/dispatch_config_routing_tests.rs"]
18mod dispatch_config_routing_tests;
19#[path = "ops/dispatch_helpers.rs"]
20mod dispatch_helpers;
21#[path = "ops/dispatch_proj.rs"]
22mod dispatch_proj;
23#[cfg(test)]
24#[path = "ops/kernel_tree_tests_util.rs"]
25mod kernel_tree_tests_util;
26// 2026-09-25: W8A8 routing of the decode projections at `DECODE_W8A8_ROWS` rows.
27#[path = "ops/dispatch_proj_decode.rs"]
28mod dispatch_proj_decode;
29// 2026-09-25: Resolution of the compiled target's serving defaults (`TargetLevers`).
30#[path = "ops/target_defaults.rs"]
31pub mod target_defaults;
32// 2026-09-25: Row-wise FP8 projection routing (`cublas_fp8_rowwise_proj`).
33#[path = "ops/dispatch_proj_rowwise.rs"]
34mod dispatch_proj_rowwise;
35#[path = "ops/embeddings.rs"]
36mod embeddings;
37#[path = "ops/fp8_act_quant.rs"]
38mod fp8_act_quant;
39// 2026-09-25: The CTA-count floor and the lever of the Hopper FP8 activation-quant launch.
40#[path = "ops/fp8_act_quant_floor.rs"]
41mod fp8_act_quant_floor;
42#[path = "ops/fp8_gemv_batch.rs"]
43mod fp8_gemv_batch;
44// 2026-09-25: Launchers including the batch-2 and batch-3 FP8 MoE expert kernels.
45#[path = "ops/fp8_moe.rs"]
46mod fp8_moe;
47#[path = "ops/fp8_moe_batch_a.rs"]
48mod fp8_moe_batch_a;
49#[path = "ops/fp8_moe_batch_b.rs"]
50mod fp8_moe_batch_b;
51mod fp8_moe_grouped;
52// 2026-09-28: The opt-in W8A8 twin of the tensor-core grouped FP8 MoE decode.
53mod fp8_moe_grouped_tc_w8a8;
54// 2026-09-27: The grouped NVFP4 MoE decode kernels (`moe/forward_nvfp4_grouped_decode.rs`).
55mod nvfp4_moe_grouped;
56// 2026-10-02: The BF16 point of the tensor-core grouped MoE decode.
57mod bf16_moe_grouped;
58// 2026-09-25: The 32-row M-tile twin of `w8a16_gemm_pipelined` and its by-M selector.
59mod w8a16_gemm_pipelined_m32;
60#[path = "ops/w8a8_decode.rs"]
61mod w8a8_decode;
62// 2026-09-28: The tensor-core row-tile W8A16 projection (`w8a16_tc_rows.cu`).
63mod w8a16_tc_rows;
64// 2026-10-02: Its NVFP4 W4A16 point (`w4a16_tc_rows.cu`).
65mod w4a16_tc_rows;
66// 2026-09-25: Tensor-core BF16 GEMM with a 16-row M tile (`dense_gemm_m16_bf16`).
67#[path = "ops/dense_gemm_m16_bf16.rs"]
68mod dense_gemm_m16_bf16;
69#[path = "ops/e4m3_range.rs"]
70mod e4m3_range;
71#[path = "ops/gemm_dense.rs"]
72mod gemm_dense;
73#[path = "ops/gemm_dense_int8.rs"]
74mod gemm_dense_int8;
75#[path = "ops/gemm_fp4.rs"]
76mod gemm_fp4;
77#[path = "ops/model_stats.rs"]
78pub mod model_stats;
79#[path = "ops/n_tile.rs"]
80mod n_tile;
81#[path = "ops/w8a16_gemm_m16.rs"]
82mod w8a16_gemm_m16;
83// 2026-09-25: N-column-blocked W8A16 batch-16 GEMVs (`w8a16_gemv_batch16_ncol2`/`ncol4`).
84#[path = "ops/w8a16_gemv_ncol.rs"]
85mod w8a16_gemv_ncol;
86pub use model_stats::ModelStats;
87
88#[path = "ops/gemm_fp8_prefill.rs"]
89mod gemm_fp8_prefill;
90#[path = "ops/gemm_quant.rs"]
91mod gemm_quant;
92#[path = "ops/gemv_q2.rs"]
93mod gemv_q2;
94#[path = "ops/gemv_q2_vec.rs"]
95mod gemv_q2_vec;
96#[path = "ops/gemv_sw.rs"]
97mod gemv_sw;
98#[path = "ops/hyper_connection.rs"]
99mod hyper_connection;
100#[path = "ops/hyper_connection_dispatch.rs"]
101mod hyper_connection_dispatch;
102#[path = "ops/hyper_connection_lowrank.rs"]
103mod hyper_connection_lowrank;
104#[cfg(test)]
105#[path = "ops/hyper_connection_lowrank_tests.rs"]
106mod hyper_connection_lowrank_tests;
107#[path = "ops/kv_cache.rs"]
108mod kv_cache;
109#[path = "ops/kv_cache_fp8k.rs"]
110mod kv_cache_fp8k;
111#[path = "ops/kv_cache_turbok.rs"]
112mod kv_cache_turbok;
113#[path = "ops/lora_delta.rs"]
114pub mod lora_delta;
115#[path = "ops/model_levers.rs"]
116mod model_levers;
117#[path = "ops/moe_atomic_c4.rs"]
118mod moe_atomic_c4;
119#[path = "ops/moe_expert.rs"]
120mod moe_expert;
121#[path = "ops/moe_expert_more.rs"]
122mod moe_expert_more;
123#[path = "ops/moe_gate.rs"]
124mod moe_gate;
125#[path = "ops/moe_grouped_a.rs"]
126mod moe_grouped_a;
127#[path = "ops/moe_grouped_a2.rs"]
128mod moe_grouped_a2;
129#[path = "ops/moe_grouped_b.rs"]
130mod moe_grouped_b;
131#[path = "ops/moe_grouped_fp4.rs"]
132mod moe_grouped_fp4;
133#[path = "ops/moe_lora_grouped.rs"]
134pub mod moe_lora_grouped;
135#[path = "ops/moe_prefill.rs"]
136mod moe_prefill;
137#[path = "ops/norm.rs"]
138mod norm;
139// 2026-09-25: Tensor-core routing of the MTP drafter's BF16 GEMV.
140#[path = "ops/dense_gemv_tc.rs"]
141pub mod dense_gemv_tc;
142#[path = "ops/gemv_tc.rs"]
143pub mod gemv_tc;
144#[cfg(test)]
145#[path = "ops/kquant_fold_tests.rs"]
146mod kquant_fold_tests;
147mod kquant_mmq;
148#[cfg(test)]
149#[path = "ops/kquant_mmq_tests.rs"]
150mod kquant_mmq_tests;
151#[cfg(test)]
152#[path = "ops/norm_gated_rms_strided_tests.rs"]
153mod norm_gated_rms_strided_tests;
154mod nvfp4_mmq;
155#[path = "ops/ple.rs"]
156mod ple;
157#[path = "ops/prefill_attn_a.rs"]
158mod prefill_attn_a;
159#[path = "ops/prefill_attn_b.rs"]
160mod prefill_attn_b;
161#[path = "ops/prefill_attn_batched.rs"]
162mod prefill_attn_batched;
163#[path = "ops/prefill_attn_fa128.rs"]
164mod prefill_attn_fa128;
165#[path = "ops/prefill_attn_fp8k.rs"]
166mod prefill_attn_fp8k;
167#[path = "ops/prefill_attn_main_a.rs"]
168mod prefill_attn_main_a;
169#[path = "ops/prefill_attn_main_b.rs"]
170mod prefill_attn_main_b;
171#[path = "ops/prefill_attn_turbok.rs"]
172mod prefill_attn_turbok;
173mod q2_0_mmq;
174mod q4k_mmq;
175#[path = "ops/qsa.rs"]
176mod qsa;
177#[path = "ops/quant_dispatch.rs"]
178mod quant_dispatch;
179#[path = "ops/sampling.rs"]
180mod sampling;
181#[path = "ops/ssm_ba_gates_hopper.rs"]
182mod ssm_ba_gates_hopper;
183#[path = "ops/ssm_gdn_a.rs"]
184mod ssm_gdn_a;
185#[path = "ops/ssm_gdn_a2.rs"]
186mod ssm_gdn_a2;
187#[path = "ops/ssm_gdn_a3.rs"]
188mod ssm_gdn_a3;
189#[path = "ops/ssm_gdn_b.rs"]
190mod ssm_gdn_b;
191#[path = "ops/ssm_gdn_batched.rs"]
192mod ssm_gdn_batched;
193#[path = "ops/ssm_gdn_carry.rs"]
194mod ssm_gdn_carry;
195#[path = "ops/ssm_gdn_carry_sizes.rs"]
196mod ssm_gdn_carry_sizes;
197#[path = "ops/ssm_gdn_exact_f16.rs"]
198mod ssm_gdn_exact_f16;
199#[path = "ops/ssm_gdn_hopper_prefill.rs"]
200mod ssm_gdn_hopper_prefill;
201#[path = "ops/ssm_gdn_snap.rs"]
202mod ssm_gdn_snap;
203#[path = "ops/ssm_gdn_tc_route.rs"]
204mod ssm_gdn_tc_route;
205#[path = "ops/ssm_gdn_woa.rs"]
206mod ssm_gdn_woa;
207#[path = "ops/ssm_gdn_wyn.rs"]
208mod ssm_gdn_wyn;
209#[path = "ops/ssm_mamba.rs"]
210mod ssm_mamba;
211#[path = "ops/ssm_preproc.rs"]
212mod ssm_preproc;
213#[path = "ops/ssm_ssd.rs"]
214mod ssm_ssd;
215pub mod token_overlay;
216#[path = "ops/w4a4_proj.rs"]
217pub mod w4a4_proj;
218// 2026-09-25: Host-side tests of the Hopper `w8a16_gemv` override's index math.
219#[cfg(test)]
220#[path = "ops/w8a16_gemv_hopper_tests.rs"]
221mod w8a16_gemv_hopper_tests;
222#[path = "ops/wide_prefill.rs"]
223mod wide_prefill;
224
225pub use activations::*;
226pub use bf16_moe_grouped::*;
227pub use dense_gemm_m16_bf16::*;
228pub use derived_weights::{Derivation, DerivedWeights};
229pub use dispatch_config::{CublasScope, GemmDispatch, parse_cublas_scope};
230pub use dispatch_helpers::*;
231pub use dispatch_proj::*;
232pub use dispatch_proj_decode::*;
233pub use dispatch_proj_rowwise::*;
234pub use e4m3_range::{
235    E4M3_MAX, E4m3Saturation, E4m3SaturationScope, MOE_ROUTED_DOWN_KNOWN,
236    NEMOTRON_SHARED_EXPERT_E4M3, allow_e4m3_saturation, check_e4m3_activation_range,
237    check_e4m3_grouped, check_e4m3_range, e4m3_saturations,
238};
239pub use embeddings::*;
240pub use fp8_act_quant::*;
241pub use fp8_act_quant_floor::*;
242pub use fp8_gemv_batch::*;
243pub use fp8_moe::*;
244pub use fp8_moe_batch_a::*;
245pub use fp8_moe_batch_b::*;
246pub use fp8_moe_grouped::*;
247pub use fp8_moe_grouped_tc_w8a8::*;
248pub use gemm_dense::*;
249pub use gemm_dense_int8::*;
250pub use gemm_fp4::*;
251pub use gemm_fp8_prefill::*;
252pub use gemm_quant::*;
253pub use gemv_q2::*;
254pub use gemv_q2_vec::*;
255pub use gemv_sw::*;
256pub use n_tile::n_tile_blocks;
257pub use nvfp4_moe_grouped::*;
258
259pub use hyper_connection::*;
260pub use hyper_connection_dispatch::*;
261pub use hyper_connection_lowrank::*;
262pub use kquant_mmq::*;
263pub use kv_cache::*;
264pub use kv_cache_fp8k::*;
265pub use kv_cache_turbok::*;
266pub use model_levers::ModelLevers;
267pub use moe_atomic_c4::*;
268pub use moe_expert::*;
269pub use moe_expert_more::*;
270pub use moe_gate::*;
271pub use moe_grouped_a::*;
272pub use moe_grouped_a2::*;
273#[allow(unused_imports)]
274pub(crate) use moe_grouped_b::*;
275pub use moe_grouped_fp4::*;
276pub use moe_lora_grouped::*;
277pub use moe_prefill::*;
278pub use norm::*;
279pub use nvfp4_mmq::*;
280pub use ple::*;
281pub use prefill_attn_a::*;
282pub use prefill_attn_b::*;
283pub use prefill_attn_batched::*;
284pub use prefill_attn_fa128::AttnFa128Kernels;
285pub use prefill_attn_fp8k::*;
286pub use prefill_attn_main_a::*;
287pub use prefill_attn_main_b::*;
288pub use prefill_attn_turbok::*;
289pub use q2_0_mmq::*;
290pub use q4k_mmq::*;
291pub use qsa::*;
292pub use quant_dispatch::*;
293pub use sampling::*;
294pub use ssm_ba_gates_hopper::*;
295pub use ssm_gdn_a::*;
296pub use ssm_gdn_a2::*;
297pub use ssm_gdn_a3::*;
298pub use ssm_gdn_b::*;
299pub use ssm_gdn_batched::*;
300pub use ssm_gdn_carry::*;
301pub use ssm_gdn_carry_sizes::*;
302pub use ssm_gdn_exact_f16::*;
303pub(crate) use ssm_gdn_hopper_prefill::*;
304pub use ssm_gdn_snap::*;
305pub use ssm_gdn_tc_route::*;
306pub use ssm_gdn_woa::*;
307pub use ssm_gdn_wyn::*;
308pub use ssm_mamba::*;
309pub use ssm_preproc::*;
310pub use ssm_ssd::*;
311pub use w4a16_tc_rows::*;
312pub use w8a8_decode::*;
313pub use w8a16_gemm_m16::*;
314pub use w8a16_gemm_pipelined_m32::*;
315pub use w8a16_gemv_ncol::*;
316pub use w8a16_tc_rows::*;
317pub use wide_prefill::*;
metrale_model_layers/layers/ops.rs

metrale_model_layers/layers/
ops.rs
