Signed Quantized Weights

crates/backend-uzu/BENCHMARKS.md

@@ -32,6 +32,7 @@ target is `--lib`.
 |-----------------------------------------|-------------------------------------|
 | `Metal/Kernel/Matmul/GEMM`              | `Metal/Kernel/Matmul/GEMM`          |
 | `Metal/Kernel/Matmul/GEMM_MXU`          | `Metal/Kernel/Matmul/GEMM_MXU`      |
+| `Metal/Kernel/A8W/w4`, `.../w8`         | `Metal/Kernel/A8W`                  |
 | `Metal/Kernel/UnifiedQuantizedGemm/...` | `Metal/Kernel/UnifiedQuantizedGemm` |
 | `Metal/Kernel/Gemv/...`                 | `Metal/Kernel/Gemv`                 |
 | `Metal/Kernel/Qwen3Layers/...`          | `Metal/Kernel/Qwen3Layers`          |
@@ -72,31 +73,39 @@ directory.
 Run one benchmark group at a time to avoid the iOS watchdog killing the
 app.
 
+Set `IPHONEOS_DEPLOYMENT_TARGET` (value from `platforms.toml` `[envs]`)
+for all iOS builds; without it the link step fails with undefined
+symbols (e.g. `___chkstk_darwin`) because objects are built for a newer
+SDK than the default deployment target.
+
 Key flags:
 
-- `-e CRITERION_HOME=target/criterion/a19` — on-device env var. Path is
+- `-e CRITERION_HOME=criterion/a19` — on-device env var. Path is
   relative to the app's cwd (`Documents/`), so this becomes
-  `Documents/target/criterion/a19/` on device.
-- `--copy-back "Documents/target=$(pwd)/target"` — after the run,
-  `cargo-dinghy` pulls `Documents/target` from the device into your
-  repo's `target/`. `$(pwd)` is required (absolute DST) because the
+  `Documents/criterion/a19/` on device. Keep it directly under
+  `Documents/` — nested parents (e.g. `Documents/target/`) do not exist
+  on a fresh install and the pre-run sync cannot create them.
+- `--sync-dirs "$(pwd)/target/criterion=Documents/criterion"` —
+  syncs the criterion tree between host and device before and after the
+  run, so results written on device land back in the repo's
+  `target/criterion/`. `$(pwd)` is required (absolute path) because the
   cargo runner is launched with cwd set to the package dir, not the
   workspace root.
 
 ```bash
 DEVICE=<DEVICE_ID>
 
-cargo dinghy \
+IPHONEOS_DEPLOYMENT_TARGET=26.4 cargo dinghy \
   -d "$DEVICE" \
-  -e CRITERION_HOME=target/criterion/a19 \
-  --copy-back "Documents/target=$(pwd)/target" \
+  -e CRITERION_HOME=criterion/a19 \
+  --sync-dirs "$(pwd)/target/criterion=Documents/criterion" \
   bench -p backend-uzu --lib -- \
     "Metal/Kernel/Matmul" \
     --save-baseline matmul_baseline_a19
 ```
 
-After the run completes you'll have
-`target/criterion/a19/Metal/Kernel/Matmul/<GEMM|GEMM_MXU>/…/matmul_baseline_a19/`
+On-device criterion sanitizes group path separators, so results land in
+`target/criterion/a19/Metal_Kernel_Matmul_<GEMM|GEMM_MXU>/…/`
 on the host, next to any `m2_max/` baselines.
 
 ## Viewing reports

crates/backend-uzu/src/backends/common/gpu_types/activation_transform.rs

@@ -0,0 +1,30 @@
+use bitflags::bitflags;
+
+bitflags! {
+    #[repr(transparent)]
+    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    pub struct ActivationTransformOp: u32 {
+        const INPUT_RHT  = 1 << 0;
+        const OUTPUT_RHT = 1 << 1;
+        const QUANTIZE   = 1 << 2;
+        const GROUP_SUMS = 1 << 3;
+    }
+}
+
+impl ActivationTransformOp {
+    pub fn validate(self) -> Self {
+        assert!(
+            self.contains(Self::INPUT_RHT) ^ self.contains(Self::OUTPUT_RHT),
+            "exactly one of INPUT_RHT / OUTPUT_RHT is required, got {self:?}"
+        );
+        assert!(
+            !self.contains(Self::QUANTIZE) || self.contains(Self::INPUT_RHT),
+            "QUANTIZE requires INPUT_RHT, got {self:?}"
+        );
+        assert!(
+            !self.contains(Self::GROUP_SUMS) || self.contains(Self::QUANTIZE),
+            "GROUP_SUMS requires QUANTIZE, got {self:?}"
+        );
+        self
+    }
+}

crates/backend-uzu/src/backends/common/gpu_types/matmul.rs

@@ -13,4 +13,5 @@ pub struct GemmParams {
     pub aligned_inner_iterations: u32,
     pub use_morton: bool,
     pub ab_scale: f32,
+    pub weight_codes_sign_flip_mask: u32,
 }

crates/backend-uzu/src/backends/common/gpu_types/mod.rs

@@ -3,6 +3,7 @@
 //! These `#[repr(C)]` structs are the source of truth. The build system uses
 //! cbindgen to generate C headers for Metal shaders.
 
+pub mod activation_transform;
 pub mod activation_type;
 pub mod argmax;
 pub mod attention;
@@ -16,6 +17,7 @@ pub mod ring;
 pub mod trie;
 pub mod weaver;
 
+pub use activation_transform::*;
 pub use activation_type::*;
 pub use argmax::*;
 pub use attention::*;

crates/backend-uzu/src/backends/common/gpu_types/quantization.rs

@@ -29,6 +29,14 @@ impl QuantizationMode {
             QuantizationMode::U8 => DataType::U8,
         }
     }
+
+    pub fn weight_codes_sign_flip_mask(&self) -> Option<u8> {
+        match self {
+            QuantizationMode::U4 => Some(0x88),
+            QuantizationMode::U8 => Some(0x80),
+            QuantizationMode::I8 => None,
+        }
+    }
 }
 
 impl From<QuantizationMode> for DataType {

crates/backend-uzu/src/backends/common/kernel/activation_transform.rs

@@ -0,0 +1,121 @@
+use crate::backends::common::{
+    Allocation, Backend, Encoder, Kernels,
+    gpu_types::{ActivationTransformOp, HADAMARD_TRANSFORM_BLOCK_SIZE},
+    kernel::ActivationTransformKernel,
+};
+
+pub struct ActivationTransform<B: Backend> {
+    kernel: <B::Kernels as Kernels>::ActivationTransformKernel,
+    ops: ActivationTransformOp,
+    group_size: u32,
+}
+
+impl<B: Backend> ActivationTransform<B> {
+    pub fn new(
+        context: &B::Context,
+        data_type: crate::data_type::DataType,
+        ops: ActivationTransformOp,
+        group_size: u32,
+    ) -> Result<Self, B::Error> {
+        let ops = ops.validate();
+        let kernel = <B::Kernels as Kernels>::ActivationTransformKernel::new(context, data_type, ops)?;
+        Ok(Self {
+            kernel,
+            ops,
+            group_size,
+        })
+    }
+
+    pub fn input_rht(
+        context: &B::Context,
+        data_type: crate::data_type::DataType,
+    ) -> Result<Self, B::Error> {
+        Self::new(context, data_type, ActivationTransformOp::INPUT_RHT, HADAMARD_TRANSFORM_BLOCK_SIZE as u32)
+    }
+
+    pub fn output_rht(
+        context: &B::Context,
+        data_type: crate::data_type::DataType,
+    ) -> Result<Self, B::Error> {
+        Self::new(context, data_type, ActivationTransformOp::OUTPUT_RHT, HADAMARD_TRANSFORM_BLOCK_SIZE as u32)
+    }
+
+    pub fn quantize(
+        context: &B::Context,
+        data_type: crate::data_type::DataType,
+        group_size: u32,
+        emit_group_sums: bool,
+    ) -> Result<Self, B::Error> {
+        let ops = if emit_group_sums {
+            ActivationTransformOp::INPUT_RHT | ActivationTransformOp::QUANTIZE | ActivationTransformOp::GROUP_SUMS
+        } else {
+            ActivationTransformOp::INPUT_RHT | ActivationTransformOp::QUANTIZE
+        };
+        Self::new(context, data_type, ops, group_size)
+    }
+
+    /// FP Hadamard (input- or output-order depending on construction).
+    /// `input` and `output` must be distinct buffers. The two scratch buffers are
+    /// placeholders for the quantized outputs this mode does not write.
+    pub fn encode_fp(
+        &self,
+        input: &Allocation<B>,
+        output: &mut Allocation<B>,
+        q_scratch: &mut Allocation<B>,
+        scales_scratch: &mut Allocation<B>,
+        rht_factors: &Allocation<B>,
+        element_count: u32,
+        batch_size: u32,
+        encoder: &mut Encoder<B>,
+    ) {
+        assert!(!self.ops.contains(ActivationTransformOp::QUANTIZE));
+        self.kernel.encode(
+            input,
+            output,
+            q_scratch,
+            scales_scratch,
+            None::<&mut Allocation<B>>,
+            rht_factors,
+            batch_size,
+            element_count,
+            self.group_size,
+            encoder,
+        );
+    }
+
+    /// Input RHT + symmetric int8 quantization.
+    pub fn encode_quantize(
+        &self,
+        input: &Allocation<B>,
+        fp_scratch: &mut Allocation<B>,
+        q_out: &mut Allocation<B>,
+        scales_out: &mut Allocation<B>,
+        group_sums_out: Option<&mut Allocation<B>>,
+        rht_factors: &Allocation<B>,
+        batch_size: u32,
+        element_count: u32,
+        encoder: &mut Encoder<B>,
+    ) {
+        assert!(self.ops.contains(ActivationTransformOp::QUANTIZE));
+        self.kernel.encode(
+            input,
+            fp_scratch,
+            q_out,
+            scales_out,
+            group_sums_out,
+            rht_factors,
+            batch_size,
+            element_count,
+            self.group_size,
+            encoder,
+        );
+    }
+
+    pub fn ops(&self) -> ActivationTransformOp {
+        self.ops
+    }
+
+    pub fn emit_group_sums(&self) -> bool {
+        self.ops.contains(ActivationTransformOp::GROUP_SUMS)
+    }
+}

crates/backend-uzu/src/backends/common/kernel/matmul/kernel.rs

@@ -1,5 +1,12 @@
 use crate::{
-    backends::common::{Backend, BufferArg, Encoder, Kernels, kernel::matmul::arguments::MatmulArguments},
+    backends::common::{
+        Backend, BufferArg, Encoder, Kernels,
+        gpu_types::gemm::GemmBPrologueKind,
+        kernel::matmul::{
+            arguments::MatmulArguments,
+            routing::{MatmulPath, MatmulShape},
+        },
+    },
     data_type::DataType,
 };
 
@@ -18,4 +25,13 @@ pub trait MatmulKernel: Sized + Send + Sync {
         arguments: MatmulArguments<'a, 'b, 'd, Self::Backend, TB>,
         encoder: &mut Encoder<Self::Backend>,
     ) -> Result<(), <Self::Backend as Backend>::Error>;
+
+    fn select_path(
+        &self,
+        _shape: &MatmulShape,
+        _b_prologue: GemmBPrologueKind,
+        _context: &<Self::Backend as Backend>::Context,
+    ) -> MatmulPath {
+        MatmulPath::Gemm
+    }
 }

crates/backend-uzu/src/backends/common/kernel/matmul/matmul_b.rs

@@ -16,19 +16,22 @@ pub enum MatmulB<'a, B: Backend, TB: BufferArg<'a, B> = &'a Allocation<B>> {
         biases: &'a Allocation<B>,
         mode: QuantizationMode,
         group_size: u32,
+        signed_codes: bool,
     },
     ScaleZeroPointDequant {
         b: &'a Allocation<B>,
         scales: &'a Allocation<B>,
         zero_points: &'a Allocation<B>,
         mode: QuantizationMode,
         group_size: u32,
+        signed_codes: bool,
     },
     ScaleSymmetricDequant {
         b: &'a Allocation<B>,
         scales: &'a Allocation<B>,
         mode: QuantizationMode,
         group_size: u32,
+        signed_codes: bool,
     },
 }
 
@@ -89,4 +92,24 @@ impl<'a, B: Backend, TB: BufferArg<'a, B>> MatmulB<'a, B, TB> {
             } => Some(*group_size),
         }
     }
+
+    pub fn signed_codes(&self) -> bool {
+        match self {
+            Self::FullPrecision {
+                ..
+            } => false,
+            Self::ScaleBiasDequant {
+                signed_codes,
+                ..
+            }
+            | Self::ScaleZeroPointDequant {
+                signed_codes,
+                ..
+            }
+            | Self::ScaleSymmetricDequant {
+                signed_codes,
+                ..
+            } => *signed_codes,
+        }
+    }
 }

crates/backend-uzu/src/backends/common/kernel/matmul/mod.rs

@@ -4,10 +4,12 @@ mod error;
 mod kernel;
 mod matmul_a;
 mod matmul_b;
+pub mod routing;
 
 pub use arguments::MatmulArguments;
 pub use d_ops::MatmulDOps;
 pub use error::MatmulError;
 pub use kernel::MatmulKernel;
 pub use matmul_a::MatmulA;
 pub use matmul_b::MatmulB;
+pub use routing::{MatmulPath, MatmulShape};

crates/backend-uzu/src/backends/common/kernel/matmul/routing.rs

@@ -0,0 +1,21 @@
+use crate::backends::common::gpu_types::gemm::GemmDTransform;
+
+#[derive(Debug, Clone, Copy)]
+pub struct MatmulShape {
+    pub m: u32,
+    pub n: u32,
+    pub k: u32,
+    pub b_transpose: bool,
+    pub b_leading_dimension: Option<u32>,
+    pub is_quant: bool,
+    pub b_bits: Option<u32>,
+    pub b_group_size: Option<u32>,
+    pub gathered: bool,
+    pub d_transform: GemmDTransform,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum MatmulPath {
+    Gemv,
+    Gemm,
+}

crates/backend-uzu/src/backends/common/kernel/mod.rs

@@ -1,11 +1,14 @@
 use crate::backends::common::Backend;
 
+pub mod activation_transform;
 pub mod attention_gemm;
 pub mod delta_net_chunked_prefill;
 pub mod delta_net_tree_verify;
 pub mod matmul;
 pub mod radix_top_k_small;
 
+pub use activation_transform::ActivationTransform;
+
 include!(concat!(env!("OUT_DIR"), "/traits.rs"));
 
 pub trait Kernels: Sized {