| //! dllm-kernels: the engine's CUDA kernels (cuda/*.cu, compiled for sm_120a by build.rs and embedded in the binary) and their Rust launch glue. | |
| //! * `Backend` is the device interface (`gpu::CudaBackend`: cudarc driver API, one main stream + one branch stream, CUDA graphs). | |
| //! * `Gpu` holds the typed launchers: small bf16 ops (ops.rs), FP8 linears (linear), NVFP4 MoE (moe), attention (attn), fused norms (norm), | |
| //! greedy lm_head (head), speculative decoding (spec) and the vision tower (vision). | |
| pub mod attn; | |
| pub mod backend; | |
| pub mod gpu; | |
| pub mod head; | |
| pub mod linear; | |
| pub mod moe; | |
| pub mod norm; | |
| pub mod ops; | |
| pub mod spec; | |
| pub mod vision; | |
| pub use backend::{Arg, Backend, DBuf, GraphExec}; | |
| pub use ops::Gpu; | |
| /// Row-invariant launch parameters: while set, the GEMV K-split and the decode-attention KV split are chosen as if the batch had ONE row, so a row's | |
| /// result does not depend on how many other rows share the launch. The speculative verify forward sets it, which makes a verify of k+1 rows bit-identical | |
| /// to one-row greedy decoding (lossless speculation). Costs nothing measurable. | |
| static ROW_INVARIANT: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); | |
| pub fn set_row_invariant(on: bool) { | |
| ROW_INVARIANT.store(on, std::sync::atomic::Ordering::Relaxed); | |
| } | |
| pub fn row_invariant() -> bool { | |
| ROW_INVARIANT.load(std::sync::atomic::Ordering::Relaxed) | |
| } | |