vyre_libs/nn/linear/inner/linear_4bit/
mod.rs

1//! Fused `linear_4bit` constructor  -  unpack-on-demand 4-bit quantized linear.
2//!
3//! Instead of materializing an unpacked f32 weight buffer, this kernel loads
4//! the packed u32 weight, extracts the correct nibble inside the inner `k`
5//! loop, and accumulates directly. This eliminates the 8× memory expansion
6//! of a separate unpack dispatch.
7//!
8//! [`unpack_on_demand`] owns the plain INT4 dot product. The grouped affine
9//! path splits across [`affine_grouped`] and its weight-tile strategy in
10//! [`affine_grouped_weight_reuse`], over the geometry in [`grouped_layout`].
11//! [`quantized_spec`] validates first-class quantized metadata and
12//! [`planner_evidence`] reports the cost of the fused path.
13
14mod affine_grouped;
15mod affine_grouped_weight_reuse;
16mod grouped_layout;
17mod op_registration;
18mod planner_evidence;
19mod quantized_spec;
20mod unpack_on_demand;
21
22pub use affine_grouped::{linear_4bit_affine_grouped, linear_4bit_affine_grouped_batched};
23pub use planner_evidence::linear_4bit_affine_grouped_planner_evidence;
24pub use quantized_spec::{
25    linear_4bit_affine_grouped_batched_typed, linear_4bit_affine_grouped_typed,
26};
27pub use unpack_on_demand::linear_4bit;
28
29use vyre_foundation::ir::DataType;
30
31/// Maximum absolute output drift allowed for grouped INT4 planner evidence tests.
32pub const LINEAR_4BIT_AFFINE_GROUPED_OUTPUT_DRIFT_ABS_TOLERANCE: f32 = 1.0e-4;
33
34/// Planner evidence for fused grouped INT4 linear versus dequantized matmul.
35#[derive(Debug, Clone, PartialEq)]
36pub struct QuantizedLinear4BitPlannerEvidence {
37    /// Input feature dimension.
38    pub in_dim: u32,
39    /// Output feature dimension.
40    pub out_dim: u32,
41    /// Quantization group size.
42    pub group_size: u32,
43    /// Number of quantization groups.
44    pub group_count: u32,
45    /// Packed INT4 weight bytes.
46    pub packed_weight_bytes: u64,
47    /// Bytes that a materialized f32 dequantized weight matrix would require.
48    pub dequantized_weight_bytes: u64,
49    /// Scale plus zero-point sidecar bytes.
50    pub sidecar_bytes: u64,
51    /// Bias bytes.
52    pub bias_bytes: u64,
53    /// Output bytes.
54    pub output_bytes: u64,
55    /// Dequantized weight bytes avoided by the fused path.
56    pub dequant_bytes_elided: u64,
57    /// Equivalent matmul planner M dimension.
58    pub matmul_m: u32,
59    /// Equivalent matmul planner K dimension.
60    pub matmul_k: u32,
61    /// Equivalent matmul planner N dimension.
62    pub matmul_n: u32,
63    /// Equivalent matmul planner K tile.
64    pub matmul_tile: u32,
65    /// Selected shared matmul planner path.
66    pub matmul_selected_path: &'static str,
67    /// Candidate tensor-core path from the shared matmul planner, when any.
68    pub matmul_candidate_path: Option<&'static str>,
69    /// Shared matmul planner fallback reason, when the selected path is cooperative.
70    pub matmul_fallback_reason: Option<&'static str>,
71    /// Whether the shared matmul planner selected a tensor-core path.
72    pub tensor_core_eligible: bool,
73    /// Maximum absolute output drift accepted by evidence tests.
74    pub output_drift_abs_tolerance: f32,
75}
76
77/// Typed metadata for fused grouped INT4 linear.
78///
79/// The actual packed weight buffer is still addressed as `u32` words because
80/// the kernel extracts eight nibbles per word. This spec binds that physical
81/// layout to the first-class `DataType::Quantized` contract so call sites do
82/// not pass an untyped integer buffer and lose the scale/zero-point semantics.
83#[derive(Debug, Clone, PartialEq, Eq)]
84pub struct QuantizedLinear4BitSpec {
85    /// Input feature dimension.
86    pub in_dim: u32,
87    /// Output feature dimension.
88    pub out_dim: u32,
89    /// First-class quantized weight metadata.
90    pub weight_type: DataType,
91}
vyre_libs/nn/linear/inner/linear_4bit/mod.rs

vyre_libs/nn/linear/inner/linear_4bit/
mod.rs
