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#[path = "ops/dispatch_helpers.rs"]
21mod dispatch_helpers;
22#[path = "ops/dispatch_proj.rs"]
23mod dispatch_proj;
24// Row-wise FP8 routing, split out when it took dispatch_proj.rs over the cap.
25#[path = "ops/dispatch_proj_rowwise.rs"]
26mod dispatch_proj_rowwise;
27#[path = "ops/embeddings.rs"]
28mod embeddings;
29#[path = "ops/fp8_gemv_batch.rs"]
30mod fp8_gemv_batch;
31#[path = "ops/fp8_moe.rs"]
32mod fp8_moe;
33#[path = "ops/fp8_moe_batch_a.rs"]
34mod fp8_moe_batch_a;
35#[path = "ops/fp8_moe_batch_b.rs"]
36mod fp8_moe_batch_b;
37#[path = "ops/gdn_flashinfer.rs"]
38// The FlashInfer GDN bridge uses dlopen/dlsym, which do not exist on Windows.
39// The absent variant is mounted at the SAME module path so both call sites
40// (trait_prefill_recur, trait_prefill_gdn) need no cfg — they already gate on
41// `available()`, which is simply always false there.
42#[cfg(unix)]
43pub mod gdn_flashinfer;
44#[cfg(not(unix))]
45#[path = "ops/gdn_flashinfer_absent.rs"]
46pub mod gdn_flashinfer;
47#[path = "ops/gemm_dense.rs"]
48mod gemm_dense;
49#[path = "ops/gemm_dense_int8.rs"]
50mod gemm_dense_int8;
51#[path = "ops/gemm_fp4.rs"]
52mod gemm_fp4;
53#[path = "ops/model_stats.rs"]
54pub mod model_stats;
55pub use model_stats::ModelStats;
56
57#[path = "ops/gemm_fp8_prefill.rs"]
58mod gemm_fp8_prefill;
59#[path = "ops/gemm_quant.rs"]
60mod gemm_quant;
61#[path = "ops/gemv_q2.rs"]
62mod gemv_q2;
63#[path = "ops/gemv_q2_vec.rs"]
64mod gemv_q2_vec;
65#[path = "ops/gemv_sw.rs"]
66mod gemv_sw;
67/// GLM-5.3-Flash mHC dispatch (Slice 10 gate 0) -- kept out of `hyper_connection.rs` so
68/// DeepSeek-V4's proven dispatch stays byte-untouched.
69#[path = "ops/glm5next_mhc.rs"]
70mod glm5next_mhc;
71#[path = "ops/hyper_connection.rs"]
72mod hyper_connection;
73#[path = "ops/hyper_connection_dispatch.rs"]
74mod hyper_connection_dispatch;
75#[path = "ops/hyper_connection_lowrank.rs"]
76mod hyper_connection_lowrank;
77#[cfg(test)]
78#[path = "ops/hyper_connection_lowrank_tests.rs"]
79mod hyper_connection_lowrank_tests;
80#[path = "ops/kv_cache.rs"]
81mod kv_cache;
82#[path = "ops/kv_cache_fp8k.rs"]
83mod kv_cache_fp8k;
84#[path = "ops/kv_cache_turbok.rs"]
85mod kv_cache_turbok;
86#[path = "ops/lora_delta.rs"]
87pub mod lora_delta;
88#[path = "ops/model_levers.rs"]
89mod model_levers;
90#[path = "ops/moe_atomic_c4.rs"]
91mod moe_atomic_c4;
92#[path = "ops/moe_expert.rs"]
93mod moe_expert;
94#[path = "ops/moe_expert_more.rs"]
95mod moe_expert_more;
96#[path = "ops/moe_gate.rs"]
97mod moe_gate;
98#[path = "ops/moe_grouped_a.rs"]
99mod moe_grouped_a;
100#[path = "ops/moe_grouped_a2.rs"]
101mod moe_grouped_a2;
102#[path = "ops/moe_grouped_b.rs"]
103mod moe_grouped_b;
104#[path = "ops/moe_grouped_fp4.rs"]
105mod moe_grouped_fp4;
106#[path = "ops/moe_lora_grouped.rs"]
107pub mod moe_lora_grouped;
108#[path = "ops/moe_prefill.rs"]
109mod moe_prefill;
110#[path = "ops/norm.rs"]
111mod norm;
112mod nvfp4_mmq;
113#[path = "ops/ple.rs"]
114mod ple;
115#[cfg(test)]
116#[path = "ops/ple_tests.rs"]
117mod ple_tests;
118#[path = "ops/prefill_attn_a.rs"]
119mod prefill_attn_a;
120#[path = "ops/prefill_attn_b.rs"]
121mod prefill_attn_b;
122#[path = "ops/prefill_attn_batched.rs"]
123mod prefill_attn_batched;
124#[path = "ops/prefill_attn_fp8k.rs"]
125mod prefill_attn_fp8k;
126#[path = "ops/prefill_attn_main_a.rs"]
127mod prefill_attn_main_a;
128#[path = "ops/prefill_attn_main_b.rs"]
129mod prefill_attn_main_b;
130#[path = "ops/prefill_attn_turbok.rs"]
131mod prefill_attn_turbok;
132mod q2_0_mmq;
133mod q4k_mmq;
134#[path = "ops/qsa.rs"]
135mod qsa;
136#[path = "ops/quant_dispatch.rs"]
137mod quant_dispatch;
138#[path = "ops/sampling.rs"]
139mod sampling;
140#[path = "ops/ssm_gdn_a.rs"]
141mod ssm_gdn_a;
142#[path = "ops/ssm_gdn_a2.rs"]
143mod ssm_gdn_a2;
144#[path = "ops/ssm_gdn_a3.rs"]
145mod ssm_gdn_a3;
146#[path = "ops/ssm_gdn_b.rs"]
147mod ssm_gdn_b;
148#[path = "ops/ssm_gdn_batched.rs"]
149mod ssm_gdn_batched;
150#[path = "ops/ssm_gdn_snap.rs"]
151mod ssm_gdn_snap;
152#[path = "ops/ssm_mamba.rs"]
153mod ssm_mamba;
154#[path = "ops/ssm_preproc.rs"]
155mod ssm_preproc;
156#[path = "ops/ssm_ssd.rs"]
157mod ssm_ssd;
158pub mod token_overlay;
159#[path = "ops/wide_prefill.rs"]
160mod wide_prefill;
161
162pub use activations::*;
163pub use derived_weights::{Derivation, DerivedWeights};
164pub use dispatch_config::GemmDispatch;
165pub use dispatch_helpers::*;
166pub use dispatch_proj::*;
167pub use dispatch_proj_rowwise::*;
168pub use embeddings::*;
169pub use fp8_gemv_batch::*;
170pub use fp8_moe::*;
171pub use fp8_moe_batch_a::*;
172pub use fp8_moe_batch_b::*;
173pub use gemm_dense::*;
174pub use gemm_dense_int8::*;
175pub use gemm_fp4::*;
176pub use gemm_fp8_prefill::*;
177pub use gemm_quant::*;
178pub use gemv_q2::*;
179pub use gemv_q2_vec::*;
180pub use gemv_sw::*;
181pub use glm5next_mhc::*;
182pub use hyper_connection::*;
183pub use hyper_connection_dispatch::*;
184pub use hyper_connection_lowrank::*;
185pub use kv_cache::*;
186pub use kv_cache_fp8k::*;
187pub use kv_cache_turbok::*;
188pub use model_levers::ModelLevers;
189pub use moe_atomic_c4::*;
190pub use moe_expert::*;
191pub use moe_expert_more::*;
192pub use moe_gate::*;
193pub use moe_grouped_a::*;
194pub use moe_grouped_a2::*;
195#[allow(unused_imports)]
196pub(crate) use moe_grouped_b::*;
197pub use moe_grouped_fp4::*;
198pub use moe_lora_grouped::*;
199pub use moe_prefill::*;
200pub use norm::*;
201pub use nvfp4_mmq::*;
202pub use ple::*;
203pub use prefill_attn_a::*;
204pub use prefill_attn_b::*;
205pub use prefill_attn_batched::*;
206pub use prefill_attn_fp8k::*;
207pub use prefill_attn_main_a::*;
208pub use prefill_attn_main_b::*;
209pub use prefill_attn_turbok::*;
210pub use q2_0_mmq::*;
211pub use q4k_mmq::*;
212pub use qsa::*;
213pub use quant_dispatch::*;
214pub use sampling::*;
215pub use ssm_gdn_a::*;
216pub use ssm_gdn_a2::*;
217pub use ssm_gdn_a3::*;
218pub use ssm_gdn_b::*;
219pub use ssm_gdn_batched::*;
220pub use ssm_gdn_snap::*;
221pub use ssm_mamba::*;
222pub use ssm_preproc::*;
223pub use ssm_ssd::*;
224pub use wide_prefill::*;
spark_model/layers/ops.rs

spark_model/layers/
ops.rs
