engine/crates/kernels/src/lib.rs

//! 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)
}