feat(compile): export Q8_0 / Q4_0 Linear weights as a block-quant linalg chain

CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ### Added
 
+- **Block-quantized weights through StableHLO export** (`skainet-compile-dag`, `skainet-compile-opt`, `skainet-compile-hlo`): `PackedConstantHandling.KEEP_PACKED` keeps a packed weight as its storage instead of failing or dequantizing; `BlockQuantMatmulFusionPass` folds a Q8_0 / Q4_0 weight used as `x · Wᵀ` into a `block_quant_matmul` node; `BlockQuantMatmulConverter` (registered in every converter set) emits it as `(sitofp(codes) - zp) * scale` over `[nB, 32, N]` `i8` / `i4` codes with `bf16` scales, then `linalg.batch_matmul`, all constants inline. A packed weight left unfolded fails the export. On an SL2610, Torq 2.3.0 with `--iree-input-type=stablehlo-torq` compiles the emitted Q8_0 and Q4_0 projections to its `mixed_quant_matmul` kernel and matches the dequantized weight (cosine 0.999996); `TorqBlockQuantGoldenTest` checks the match when `SKAINET_TORQ_COMPILE` names the compiler. `HloGenerator.generate` takes `packedConstants`.
 - **StableHLO-phase target passes** (`skainet-compile-opt`, `skainet-compile-hlo`): `TargetOptimizer.stableHloPasses()` and `TargetOptimizers.registerStableHloPasses(target) { … }` / `stableHloPassesFor(target)` wire the `STABLE_HLO` phase. A `StableHloPass` rewrites the emitted module's text into valid StableHLO; `toStableHlo(graph, …, passes = …)` and `StableHloModule.applyPasses(passes)` apply them in order. The existing `toStableHlo` signatures are unchanged.
 - **TOSA preparation passes** (`sk.ainet.compile.hlo.tosa`): `DivideByBroadcastPass` (a divide by a broadcast becomes a multiply by the broadcast reciprocal, or by `rsqrt` for a `sqrt`) and `RankTwoWeightMatmulPass` (`[1, M, K] · [K, N]` as a rank-2 dot) shape a module for StableHLO's own `--stablehlo-prepare-for-tosa --stablehlo-legalize-to-tosa`; `TosaFinish` translates what that legalization leaves (`broadcast_in_dim`, `rsqrt`, batched attention `dot_general`s) and throws `TosaFinishException` on anything else. Moonshine v2 tiny streaming encoder (64 frames) on an SL2610 with the Torq NPU's `tosa-torq` input: one dispatch per call, golden transcripts identical to the recipe-side rewrite they replace.
 

docs/modules/ROOT/pages/explanation/compile-flexibility.adoc

@@ -100,6 +100,33 @@ carries what goes around it:
 The legalization runs outside SKaiNET (`stablehlo-opt`, built from the same LLVM as the
 compiler that reads its output).
 
+=== Block-quantized weights for an NPU kernel
+
+A Q8_0 or Q4_0 `Linear` can reach the compiler as integer codes and per-block scales instead of
+a dense FP32 copy. Three pieces carry it:
+
+* the trace keeps the weight packed: `toComputeGraph(…, packedConstants = KEEP_PACKED)` (or
+  `HloGenerator.generate(…, packedConstants = …)`);
+* `BlockQuantMatmulFusionPass` (DAG phase) folds `weight → transpose → matmul` into one
+  `block_quant_matmul` node that holds the packed storage;
+* `BlockQuantMatmulConverter` emits that node as a `linalg` chain: a `linalg.generic` computing
+  `(sitofp(codes) - zp) * scale` over `[nB, 32, N]` (`i8` or `i4` codes, `bf16` scales broadcast
+  over each 32-element block), then `collapse_shape`, `expand_shape` and `linalg.batch_matmul`
+  with a `bf16` activation. All constants are inline.
+
+[source,kotlin]
+----
+TargetOptimizers.registerDagPasses("torq") { listOf(BlockQuantMatmulFusionPass()) }
+----
+
+This is the form the Torq compiler (2.3.0, `--iree-input-type=stablehlo-torq`) matches to its
+`mixed_quant_matmul` kernel when the activation has one row. That kernel takes only positive
+scales, so a block with a negative scale `d` is stored exactly as codes `-1 - c`, zero point
+`-1` and scale `-d` (ggml Q4_0 sets `d = max / -8`, about half its blocks are negative). A
+packed weight the pass did not fold, such as a tied embedding, makes the export fail; trace
+with `DEQUANTIZE` for those. Stock IREE `llvm-cpu` has given wrong results for `i4` weights, so
+keep Q4_0 to targets that were checked.
+
 == Why exact GELU needs care
 
 [cols="1,3"]

skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api

@@ -530,13 +530,19 @@ public final class sk/ainet/lang/trace/PackedConstantException : java/lang/Illeg
 }
 
 public final class sk/ainet/lang/trace/PackedConstantHandling : java/lang/Enum {
+	public static final field Companion Lsk/ainet/lang/trace/PackedConstantHandling$Companion;
 	public static final field DEQUANTIZE Lsk/ainet/lang/trace/PackedConstantHandling;
 	public static final field FAIL Lsk/ainet/lang/trace/PackedConstantHandling;
+	public static final field KEEP_PACKED Lsk/ainet/lang/trace/PackedConstantHandling;
+	public static final field PACKED_WEIGHT_PARAM Ljava/lang/String;
 	public static fun getEntries ()Lkotlin/enums/EnumEntries;
 	public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/trace/PackedConstantHandling;
 	public static fun values ()[Lsk/ainet/lang/trace/PackedConstantHandling;
 }
 
+public final class sk/ainet/lang/trace/PackedConstantHandling$Companion {
+}
+
 public final class sk/ainet/lang/trace/TapeSink : sk/ainet/lang/trace/OpSink {
 	public fun <init> (Lsk/ainet/tape/ExecutionTape;)V
 	public fun onOpExecuted (Lsk/ainet/lang/trace/OpTrace;)V

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt

@@ -25,6 +25,21 @@ public enum class PackedConstantHandling {
      * use only when the export target genuinely needs dense constants.
      */
     DEQUANTIZE,
+
+    /**
+     * Keep the packed data as is: the weight node carries the
+     * [sk.ainet.lang.tensor.storage.PackedBlockStorage] under [PACKED_WEIGHT_PARAM]
+     * and keeps its encoding, with no float copy. Something downstream must
+     * consume it: `BlockQuantMatmulFusionPass` folds a Q8_0 / Q4_0 weight into
+     * the matmul that reads it, and the StableHLO emitter refuses a packed
+     * weight that is still a plain constant.
+     */
+    KEEP_PACKED;
+
+    public companion object {
+        /** Weight-node parameter holding the packed storage under [KEEP_PACKED]. */
+        public const val PACKED_WEIGHT_PARAM: String = "packed_weight"
+    }
 }
 
 /**

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt

@@ -292,26 +292,32 @@ public class TraceToGraphBuilder(
                         // packed encoding forward would misdescribe it.
                         encoding = null
                     }
+                    PackedConstantHandling.KEEP_PACKED -> Unit
                 }
             }
+            val keptPacked = if (constantValues == null && packedConstants == PackedConstantHandling.KEEP_PACKED) packedData else null
 
             val syntheticNode: GraphNode
             val producedSpec: TensorSpec
-            if (constantValues != null) {
-                // Create a constant/weight node with embedded values
+            if (constantValues != null || keptPacked != null) {
+                // Create a constant/weight node with embedded values (or, under
+                // KEEP_PACKED, the packed storage itself)
                 val weightShape = tensor!!.shape.dimensions.toList()
                 val weightDtype = tensor.dtype.simpleName ?: "FP32"
                 val nodeId = "n${nextNodeId++}_weight"
                 val op = TraceBackedOperation(
                     name = "weight",
                     type = "constant",
-                    parameters = mapOf(
+                    parameters = if (keptPacked != null) mapOf(
+                        PackedConstantHandling.PACKED_WEIGHT_PARAM to keptPacked,
+                        "trainable" to false
+                    ) else mapOf(
                         // Store the primitive FloatArray, NOT .toList(): boxing a
                         // real LLM weight (e.g. 262153x640 embedding) into a
                         // List<Float> is ~2.7GB and OOMs the trace. The HLO
                         // converter handles FloatArray for both inline and
                         // external (.irpa) materialization.
-                        "initial_value" to constantValues,
+                        "initial_value" to constantValues!!,
                         "trainable" to false
                     )
                 )

skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api

@@ -398,6 +398,12 @@ public final class sk/ainet/compile/hlo/converters/BasicMathConverter : sk/ainet
 	public fun getSupportedOperations ()Ljava/util/Set;
 }
 
+public final class sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter : sk/ainet/compile/hlo/StableHloOperationConverter {
+	public fun <init> ()V
+	public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult;
+	public fun getSupportedOperations ()Ljava/util/Set;
+}
+
 public final class sk/ainet/compile/hlo/converters/ConstantOperationsConverter : sk/ainet/compile/hlo/StableHloOperationConverter {
 	public fun <init> ()V
 	public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult;
@@ -570,8 +576,8 @@ public final class sk/ainet/compile/hlo/examples/ValidationDemonstrationResult {
 public final class sk/ainet/compile/hlo/generate/HloGenerator {
 	public static final field INSTANCE Lsk/ainet/compile/hlo/generate/HloGenerator;
 	public final fun corePasses (Ljava/lang/String;)Ljava/util/List;
-	public final fun generate (Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
-	public static synthetic fun generate$default (Lsk/ainet/compile/hlo/generate/HloGenerator;Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
+	public final fun generate (Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lsk/ainet/lang/trace/PackedConstantHandling;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static synthetic fun generate$default (Lsk/ainet/compile/hlo/generate/HloGenerator;Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lsk/ainet/lang/trace/PackedConstantHandling;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 }
 
 public final class sk/ainet/compile/hlo/generate/HloGeneratorJvmKt {

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt

@@ -2,6 +2,7 @@ package sk.ainet.compile.hlo
 
 import sk.ainet.compile.hlo.converters.ActivationOperationsConverter
 import sk.ainet.compile.hlo.converters.ArgMaxOperationsConverter
+import sk.ainet.compile.hlo.converters.BlockQuantMatmulConverter
 import sk.ainet.compile.hlo.converters.AttentionOperationsConverter
 import sk.ainet.compile.hlo.converters.ConstantOperationsConverter
 import sk.ainet.compile.hlo.converters.GatherOperationsConverter
@@ -81,6 +82,8 @@ public object StableHloConverterFactory {
 
         // Register constant operations converter
         registry.register(ConstantOperationsConverter())
+        // block_quant_matmul exists only where BlockQuantMatmulFusionPass ran; packed weights stay packed
+        registry.register(BlockQuantMatmulConverter())
 
         // Register attention (scaledDotProductAttention) converter
         registry.register(AttentionOperationsConverter())
@@ -141,6 +144,8 @@ public object StableHloConverterFactory {
 
         // Register constant operations converter
         registry.register(ConstantOperationsConverter())
+        // block_quant_matmul exists only where BlockQuantMatmulFusionPass ran; packed weights stay packed
+        registry.register(BlockQuantMatmulConverter())
 
         // Register attention (scaledDotProductAttention) converter
         registry.register(AttentionOperationsConverter())
@@ -178,6 +183,8 @@ public object StableHloConverterFactory {
         registry.register(UnaryMathConverter())
         registry.register(ScalarOperationsConverter())
         registry.register(ConstantOperationsConverter())
+        // block_quant_matmul exists only where BlockQuantMatmulFusionPass ran; packed weights stay packed
+        registry.register(BlockQuantMatmulConverter())
 
         return StableHloConverter(registry, typeMapper, null, policy, errorPolicy = errorPolicy)
     }