spark_model/layers/
ops.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Shared kernel dispatch operations.
4//!
5//! Freestanding functions wrapping CUDA kernel launches via `KernelLaunch`.
6//! Layer implementations compose these to build forward passes.
7//!
8//! Each function's parameters exactly match the corresponding CUDA kernel
9//! signature. Grid/block dimensions are computed from the problem size.
10//!
11//! Refactor wave 4a (2026-05-03): split into `ops/` sub-modules with thematic
12//! groupings. All public functions remain available at this path via re-export.
13
14#[path = "ops/activations.rs"]
15mod activations;
16#[path = "ops/derived_weights.rs"]
17mod derived_weights;
18#[path = "ops/dispatch_config.rs"]
19mod dispatch_config;
20#[cfg(test)]
21#[path = "ops/dispatch_config_routing_tests.rs"]
22mod dispatch_config_routing_tests;
23#[path = "ops/dispatch_helpers.rs"]
24mod dispatch_helpers;
25#[path = "ops/dispatch_proj.rs"]
26mod dispatch_proj;
27// W8A8 block-scaled cuBLASLt routing for the 5..16-row DECODE projections
28// (#927), a sibling of dispatch_proj.rs so neither file crosses the cap.
29#[path = "ops/dispatch_proj_decode.rs"]
30mod dispatch_proj_decode;
31// The compiled target's serving defaults (`kernels/<hw>/HARDWARE.toml`
32// `[defaults]`, baked into avarok_kernels), resolved BEFORE the environment.
33// SSOT for every lever that differs between one target and another.
34#[path = "ops/target_defaults.rs"]
35pub mod target_defaults;
36// Row-wise FP8 routing, split out when it took dispatch_proj.rs over the cap.
37#[path = "ops/dispatch_proj_rowwise.rs"]
38mod dispatch_proj_rowwise;
39#[path = "ops/embeddings.rs"]
40mod embeddings;
41#[path = "ops/fp8_act_quant.rs"]
42mod fp8_act_quant;
43// WHEN the Hopper FP8 act-quant twin runs: the CTA floor, its lever and the
44// route line (#928, round-16 receipt § 2.1). A sibling so neither file crosses
45// the cap.
46#[path = "ops/fp8_act_quant_floor.rs"]
47mod fp8_act_quant_floor;
48#[path = "ops/fp8_gemv_batch.rs"]
49mod fp8_gemv_batch;
50// Tensor-core W8A16 decode GEMM with a 16-row M tile (#927), the ALU-bound
51// `w8a16_gemv_batch16`'s replacement at 5..=16 rows. Behind AVAROK_FFN_M16_TC.
52#[path = "ops/fp8_moe.rs"]
53mod fp8_moe;
54#[path = "ops/fp8_moe_batch_a.rs"]
55mod fp8_moe_batch_a;
56#[path = "ops/fp8_moe_batch_b.rs"]
57mod fp8_moe_batch_b;
58#[path = "ops/gdn_flashinfer.rs"]
59// The FlashInfer GDN bridge uses dlopen/dlsym, which do not exist on Windows.
60// The absent variant is mounted at the SAME module path so both call sites
61// (trait_prefill_recur, trait_prefill_gdn) need no cfg — they already gate on
62// `available()`, which is simply always false there.
63#[cfg(unix)]
64pub mod gdn_flashinfer;
65#[cfg(not(unix))]
66#[path = "ops/gdn_flashinfer_absent.rs"]
67pub mod gdn_flashinfer;
68// Tensor-core BF16 decode GEMM with a 16-row M tile — the BF16 LM-head arm
69// (#927/#928). Behind AVAROK_LM_HEAD_M16_TC; SSOT for its launch geometry.
70#[path = "ops/dense_gemm_m16_bf16.rs"]
71mod dense_gemm_m16_bf16;
72#[path = "ops/gemm_dense.rs"]
73mod gemm_dense;
74#[path = "ops/gemm_dense_int8.rs"]
75mod gemm_dense_int8;
76#[path = "ops/gemm_fp4.rs"]
77mod gemm_fp4;
78#[path = "ops/model_stats.rs"]
79pub mod model_stats;
80#[path = "ops/w8a16_gemm_m16.rs"]
81mod w8a16_gemm_m16;
82// The bit-exact N-column-blocked sibling of `w8a16_gemv_batch16` (#927),
83// for the attention decode projections. Behind AVAROK_ATTN_NCOL_GEMV.
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/// GLM-5.3-Flash mHC dispatch (Slice 10 gate 0) -- kept out of `hyper_connection.rs` so
99/// DeepSeek-V4's proven dispatch stays byte-untouched.
100#[path = "ops/glm5next_mhc.rs"]
101mod glm5next_mhc;
102#[path = "ops/hyper_connection.rs"]
103mod hyper_connection;
104#[path = "ops/hyper_connection_dispatch.rs"]
105mod hyper_connection_dispatch;
106#[path = "ops/hyper_connection_lowrank.rs"]
107mod hyper_connection_lowrank;
108#[cfg(test)]
109#[path = "ops/hyper_connection_lowrank_tests.rs"]
110mod hyper_connection_lowrank_tests;
111#[path = "ops/kv_cache.rs"]
112mod kv_cache;
113#[path = "ops/kv_cache_fp8k.rs"]
114mod kv_cache_fp8k;
115#[path = "ops/kv_cache_turbok.rs"]
116mod kv_cache_turbok;
117#[path = "ops/lora_delta.rs"]
118pub mod lora_delta;
119#[path = "ops/model_levers.rs"]
120mod model_levers;
121#[path = "ops/moe_atomic_c4.rs"]
122mod moe_atomic_c4;
123#[path = "ops/moe_expert.rs"]
124mod moe_expert;
125#[path = "ops/moe_expert_more.rs"]
126mod moe_expert_more;
127#[path = "ops/moe_gate.rs"]
128mod moe_gate;
129#[path = "ops/moe_grouped_a.rs"]
130mod moe_grouped_a;
131#[path = "ops/moe_grouped_a2.rs"]
132mod moe_grouped_a2;
133#[path = "ops/moe_grouped_b.rs"]
134mod moe_grouped_b;
135#[path = "ops/moe_grouped_fp4.rs"]
136mod moe_grouped_fp4;
137#[path = "ops/moe_lora_grouped.rs"]
138pub mod moe_lora_grouped;
139#[path = "ops/moe_prefill.rs"]
140mod moe_prefill;
141#[path = "ops/norm.rs"]
142mod norm;
143// The gated-RMS-norm launch-count pin (#927): 48 per step, not 768.
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#[cfg(test)]
158#[path = "ops/ple_tests.rs"]
159mod ple_tests;
160#[path = "ops/prefill_attn_a.rs"]
161mod prefill_attn_a;
162#[path = "ops/prefill_attn_b.rs"]
163mod prefill_attn_b;
164#[path = "ops/prefill_attn_batched.rs"]
165mod prefill_attn_batched;
166#[path = "ops/prefill_attn_fp8k.rs"]
167mod prefill_attn_fp8k;
168#[path = "ops/prefill_attn_main_a.rs"]
169mod prefill_attn_main_a;
170#[path = "ops/prefill_attn_main_b.rs"]
171mod prefill_attn_main_b;
172#[path = "ops/prefill_attn_turbok.rs"]
173mod prefill_attn_turbok;
174mod q2_0_mmq;
175mod q4k_mmq;
176#[path = "ops/qsa.rs"]
177mod qsa;
178#[path = "ops/quant_dispatch.rs"]
179mod quant_dispatch;
180#[path = "ops/sampling.rs"]
181mod sampling;
182#[path = "ops/ssm_ba_gates_hopper.rs"]
183mod ssm_ba_gates_hopper;
184#[path = "ops/ssm_gdn_a.rs"]
185mod ssm_gdn_a;
186#[path = "ops/ssm_gdn_a2.rs"]
187mod ssm_gdn_a2;
188#[path = "ops/ssm_gdn_a3.rs"]
189mod ssm_gdn_a3;
190#[path = "ops/ssm_gdn_b.rs"]
191mod ssm_gdn_b;
192#[path = "ops/ssm_gdn_batched.rs"]
193mod ssm_gdn_batched;
194#[path = "ops/ssm_gdn_hopper_prefill.rs"]
195mod ssm_gdn_hopper_prefill;
196#[path = "ops/ssm_gdn_snap.rs"]
197mod ssm_gdn_snap;
198#[path = "ops/ssm_gdn_tc_route.rs"]
199mod ssm_gdn_tc_route;
200#[path = "ops/ssm_gdn_woa.rs"]
201mod ssm_gdn_woa;
202#[path = "ops/ssm_gdn_wyn.rs"]
203mod ssm_gdn_wyn;
204#[path = "ops/ssm_mamba.rs"]
205mod ssm_mamba;
206#[path = "ops/ssm_preproc.rs"]
207mod ssm_preproc;
208#[path = "ops/ssm_ssd.rs"]
209mod ssm_ssd;
210pub mod token_overlay;
211/// HOST SIMULATION of the Hopper `w8a16_gemv` override's loop order against the
212/// gb10 kernel's, so a GPU-free `cargo test` still judges the one claim the
213/// device microtest cannot make cheaply: that the UNROLL-wide prefetch did not
214/// reorder the FP32 accumulation every batch oracle in the tree compares to.
215#[cfg(test)]
216#[path = "ops/w8a16_gemv_hopper_tests.rs"]
217mod w8a16_gemv_hopper_tests;
218#[path = "ops/wide_prefill.rs"]
219mod wide_prefill;
220
221pub use activations::*;
222pub use dense_gemm_m16_bf16::*;
223pub use derived_weights::{Derivation, DerivedWeights};
224pub use dispatch_config::{CublasScope, GemmDispatch, parse_cublas_scope};
225pub use dispatch_helpers::*;
226pub use dispatch_proj::*;
227pub use dispatch_proj_decode::*;
228pub use dispatch_proj_rowwise::*;
229pub use embeddings::*;
230pub use fp8_act_quant::*;
231pub use fp8_act_quant_floor::*;
232pub use fp8_gemv_batch::*;
233pub use fp8_moe::*;
234pub use fp8_moe_batch_a::*;
235pub use fp8_moe_batch_b::*;
236pub use gemm_dense::*;
237pub use gemm_dense_int8::*;
238pub use gemm_fp4::*;
239pub use gemm_fp8_prefill::*;
240pub use gemm_quant::*;
241pub use gemv_q2::*;
242pub use gemv_q2_vec::*;
243pub use gemv_sw::*;
244pub use glm5next_mhc::*;
245pub use hyper_connection::*;
246pub use hyper_connection_dispatch::*;
247pub use hyper_connection_lowrank::*;
248pub use kquant_mmq::*;
249pub use kv_cache::*;
250pub use kv_cache_fp8k::*;
251pub use kv_cache_turbok::*;
252pub use model_levers::ModelLevers;
253pub use moe_atomic_c4::*;
254pub use moe_expert::*;
255pub use moe_expert_more::*;
256pub use moe_gate::*;
257pub use moe_grouped_a::*;
258pub use moe_grouped_a2::*;
259#[allow(unused_imports)]
260pub(crate) use moe_grouped_b::*;
261pub use moe_grouped_fp4::*;
262pub use moe_lora_grouped::*;
263pub use moe_prefill::*;
264pub use norm::*;
265pub use nvfp4_mmq::*;
266pub use ple::*;
267pub use prefill_attn_a::*;
268pub use prefill_attn_b::*;
269pub use prefill_attn_batched::*;
270pub use prefill_attn_fp8k::*;
271pub use prefill_attn_main_a::*;
272pub use prefill_attn_main_b::*;
273pub use prefill_attn_turbok::*;
274pub use q2_0_mmq::*;
275pub use q4k_mmq::*;
276pub use qsa::*;
277pub use quant_dispatch::*;
278pub use sampling::*;
279pub use ssm_ba_gates_hopper::*;
280pub use ssm_gdn_a::*;
281pub use ssm_gdn_a2::*;
282pub use ssm_gdn_a3::*;
283pub use ssm_gdn_b::*;
284pub use ssm_gdn_batched::*;
285pub(crate) use ssm_gdn_hopper_prefill::*;
286pub use ssm_gdn_snap::*;
287pub use ssm_gdn_tc_route::*;
288pub use ssm_gdn_woa::*;
289pub use ssm_gdn_wyn::*;
290pub use ssm_mamba::*;
291pub use ssm_preproc::*;
292pub use ssm_ssd::*;
293pub use w8a16_gemm_m16::*;
294pub use w8a16_gemv_ncol::*;
295pub use wide_prefill::*;
spark_model/layers/ops.rs

spark_model/layers/
ops.rs
