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 atlas_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 ATLAS_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 ATLAS_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 ATLAS_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/norm_gated_rms_strided_tests.rs"]
146mod norm_gated_rms_strided_tests;
147mod nvfp4_mmq;
148#[path = "ops/ple.rs"]
149mod ple;
150#[cfg(test)]
151#[path = "ops/ple_tests.rs"]
152mod ple_tests;
153#[path = "ops/prefill_attn_a.rs"]
154mod prefill_attn_a;
155#[path = "ops/prefill_attn_b.rs"]
156mod prefill_attn_b;
157#[path = "ops/prefill_attn_batched.rs"]
158mod prefill_attn_batched;
159#[path = "ops/prefill_attn_fp8k.rs"]
160mod prefill_attn_fp8k;
161#[path = "ops/prefill_attn_main_a.rs"]
162mod prefill_attn_main_a;
163#[path = "ops/prefill_attn_main_b.rs"]
164mod prefill_attn_main_b;
165#[path = "ops/prefill_attn_turbok.rs"]
166mod prefill_attn_turbok;
167mod q2_0_mmq;
168mod q4k_mmq;
169#[path = "ops/qsa.rs"]
170mod qsa;
171#[path = "ops/quant_dispatch.rs"]
172mod quant_dispatch;
173#[path = "ops/sampling.rs"]
174mod sampling;
175#[path = "ops/ssm_ba_gates_hopper.rs"]
176mod ssm_ba_gates_hopper;
177#[path = "ops/ssm_gdn_a.rs"]
178mod ssm_gdn_a;
179#[path = "ops/ssm_gdn_a2.rs"]
180mod ssm_gdn_a2;
181#[path = "ops/ssm_gdn_a3.rs"]
182mod ssm_gdn_a3;
183#[path = "ops/ssm_gdn_b.rs"]
184mod ssm_gdn_b;
185#[path = "ops/ssm_gdn_batched.rs"]
186mod ssm_gdn_batched;
187#[path = "ops/ssm_gdn_hopper_prefill.rs"]
188mod ssm_gdn_hopper_prefill;
189#[path = "ops/ssm_gdn_snap.rs"]
190mod ssm_gdn_snap;
191#[path = "ops/ssm_gdn_tc_route.rs"]
192mod ssm_gdn_tc_route;
193#[path = "ops/ssm_mamba.rs"]
194mod ssm_mamba;
195#[path = "ops/ssm_preproc.rs"]
196mod ssm_preproc;
197#[path = "ops/ssm_ssd.rs"]
198mod ssm_ssd;
199pub mod token_overlay;
200/// HOST SIMULATION of the Hopper `w8a16_gemv` override's loop order against the
201/// gb10 kernel's, so a GPU-free `cargo test` still judges the one claim the
202/// device microtest cannot make cheaply: that the UNROLL-wide prefetch did not
203/// reorder the FP32 accumulation every batch oracle in the tree compares to.
204#[cfg(test)]
205#[path = "ops/w8a16_gemv_hopper_tests.rs"]
206mod w8a16_gemv_hopper_tests;
207#[path = "ops/wide_prefill.rs"]
208mod wide_prefill;
209
210pub use activations::*;
211pub use dense_gemm_m16_bf16::*;
212pub use derived_weights::{Derivation, DerivedWeights};
213pub use dispatch_config::{CublasScope, GemmDispatch, parse_cublas_scope};
214pub use dispatch_helpers::*;
215pub use dispatch_proj::*;
216pub use dispatch_proj_decode::*;
217pub use dispatch_proj_rowwise::*;
218pub use embeddings::*;
219pub use fp8_act_quant::*;
220pub use fp8_act_quant_floor::*;
221pub use fp8_gemv_batch::*;
222pub use fp8_moe::*;
223pub use fp8_moe_batch_a::*;
224pub use fp8_moe_batch_b::*;
225pub use gemm_dense::*;
226pub use gemm_dense_int8::*;
227pub use gemm_fp4::*;
228pub use gemm_fp8_prefill::*;
229pub use gemm_quant::*;
230pub use gemv_q2::*;
231pub use gemv_q2_vec::*;
232pub use gemv_sw::*;
233pub use glm5next_mhc::*;
234pub use hyper_connection::*;
235pub use hyper_connection_dispatch::*;
236pub use hyper_connection_lowrank::*;
237pub use kv_cache::*;
238pub use kv_cache_fp8k::*;
239pub use kv_cache_turbok::*;
240pub use model_levers::ModelLevers;
241pub use moe_atomic_c4::*;
242pub use moe_expert::*;
243pub use moe_expert_more::*;
244pub use moe_gate::*;
245pub use moe_grouped_a::*;
246pub use moe_grouped_a2::*;
247#[allow(unused_imports)]
248pub(crate) use moe_grouped_b::*;
249pub use moe_grouped_fp4::*;
250pub use moe_lora_grouped::*;
251pub use moe_prefill::*;
252pub use norm::*;
253pub use nvfp4_mmq::*;
254pub use ple::*;
255pub use prefill_attn_a::*;
256pub use prefill_attn_b::*;
257pub use prefill_attn_batched::*;
258pub use prefill_attn_fp8k::*;
259pub use prefill_attn_main_a::*;
260pub use prefill_attn_main_b::*;
261pub use prefill_attn_turbok::*;
262pub use q2_0_mmq::*;
263pub use q4k_mmq::*;
264pub use qsa::*;
265pub use quant_dispatch::*;
266pub use sampling::*;
267pub use ssm_ba_gates_hopper::*;
268pub use ssm_gdn_a::*;
269pub use ssm_gdn_a2::*;
270pub use ssm_gdn_a3::*;
271pub use ssm_gdn_b::*;
272pub use ssm_gdn_batched::*;
273pub(crate) use ssm_gdn_hopper_prefill::*;
274pub use ssm_gdn_snap::*;
275pub use ssm_gdn_tc_route::*;
276pub use ssm_mamba::*;
277pub use ssm_preproc::*;
278pub use ssm_ssd::*;
279pub use w8a16_gemm_m16::*;
280pub use w8a16_gemv_ncol::*;
281pub use wide_prefill::*;
spark_model/layers/ops.rs

spark_model/layers/
ops.rs
