Skip to main content

TERNARY_GEMM_WGSL

Constant TERNARY_GEMM_WGSL 

Source
pub const TERNARY_GEMM_WGSL: &str = "// \u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\n// QualiaDB \u{2014} BitNet b1.58 ternary GEMM (STELLAR \u{a7}A, task #12).\n//\n//   out[m][i] = scale \u{b7} \u{3a3}_j  trit(W[i][j]) \u{b7} act[m][j]      W \u{2208} {-1, 0, +1}\n//\n// The BitNet win: a ternary weight contributes by **add / subtract only** \u{2014} there is\n// no per-weight multiply in the inner loop. The single per-tensor `scale` multiply\n// happens once per output element, at the end. (Contrast `fused_ffn.wgsl`, which does\n// a dequant-multiply-accumulate per weight.)\n//\n// Weights: the row-major trits of an (n_out \u{d7} n_in) matrix, packed 5-per-byte in\n// base-3 (see `ternary.rs::pack_trits`). The 4-byte f32 scale that prefixes a ternary\n// blob is passed in `params.scale` \u{2014} it is NOT in `trit_words`.\n//\n// CPU oracle / parity reference: `ternary::ternary_gemm_cpu` mirrors this kernel\n// exactly (same trit extraction, same add/subtract, same end-scale).\n// \u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\n\nstruct TernaryParams {\n    n_in: u32,\n    n_out: u32,\n    n_batch: u32,\n    in_row_stride: u32,\n    out_row_stride: u32,\n    scale: f32,\n    _pad0: u32,\n    _pad1: u32,\n}\n\n@group(0) @binding(0) var<storage, read> activations: array<f32>;\n@group(0) @binding(1) var<storage, read> trit_words: array<u32>;\n@group(0) @binding(2) var<uniform> params: TernaryParams;\n@group(0) @binding(3) var<storage, read_write> ternary_output: array<f32>;\n\n// Read one packed byte from the u32 word array (little-endian within the word).\nfn read_byte(idx: u32) -> u32 {\n    let w = trit_words[idx >> 2u];\n    return (w >> ((idx & 3u) * 8u)) & 0xFFu;\n}\n\n// Ternary value {-1, 0, +1} at linear weight index `k` (5 trits/byte, base-3).\nfn trit_at(k: u32) -> i32 {\n    let pos = k % 5u;\n    var b = read_byte(k / 5u);\n    for (var p = 0u; p < pos; p = p + 1u) {\n        b = b / 3u;\n    }\n    return i32(b % 3u) - 1;\n}\n\n@compute @workgroup_size(64)\nfn ternary_gemm(@builtin(global_invocation_id) gid: vec3<u32>) {\n    let i = gid.x; // output feature (column of W)\n    let m = gid.y; // batch row\n    let batch = max(params.n_batch, 1u);\n    if (m >= batch || i >= params.n_out) {\n        return;\n    }\n\n    let in_stride = select(params.n_in, params.in_row_stride, params.in_row_stride > 0u);\n    let out_stride = select(params.n_out, params.out_row_stride, params.out_row_stride > 0u);\n    let in_base = m * in_stride;\n    let row0 = i * params.n_in; // linear trit base for weight row i\n\n    var acc = 0.0;\n    for (var j = 0u; j < params.n_in; j = j + 1u) {\n        let t = trit_at(row0 + j);\n        let x = activations[in_base + j];\n        // add/subtract \u{2014} no multiply by the weight (the \u{a7}A win)\n        if (t > 0) {\n            acc = acc + x;\n        } else if (t < 0) {\n            acc = acc - x;\n        }\n    }\n\n    ternary_output[m * out_stride + i] = params.scale * acc;\n}\n";
Expand description

The WGSL ternary-GEMM compute kernel (STELLAR §A). Its CPU parity reference is ternary_gemm_cpu, which mirrors it byte-for-byte (same trit extraction, add/subtract, end-scale). The GPU pipeline binds: 0 activations (f32), 1 packed trits (u32 words), 2 TernaryParams uniform, 3 output (f32).