DJL/JiRackJDLTernary_10b.java

/**
# =============================================================================
# COPYRIGHT © 2025-2026 Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
# CMS Manhattan JiRack Technology — PATENT PENDING
#
# This code is proprietary. 
# Personal and non-commercial research use is allowed.
# Any commercial use, derivative works for profit, or distribution 
# requires a paid license and 5% royalty.
#
# It is Java Developers who likes java
#
# Unauthorized commercial use is strictly prohibited.
# Contact: grabko@cmsmanhattan.com
# =============================================================================
*/

package com.cbsinc.cms.llm.ml;

import ai.djl.Device;
import ai.djl.Model;
import ai.djl.ndarray.*;
import ai.djl.ndarray.index.NDIndex;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.nn.*;
import ai.djl.nn.core.Linear;
import ai.djl.training.GradientCollector;
import ai.djl.training.ParameterStore;
import ai.djl.training.initializer.ConstantInitializer;
import ai.djl.training.initializer.NormalInitializer;
import ai.djl.util.Pair;
import ai.djl.util.PairList;

import java.nio.file.Path;
import java.nio.file.Paths;

/**
 * JiRack Ternary Transformer — DJL 0.36 port of JiRackTernaryPyTorch_10b_fixed.py.
 * Same logic, same fixes:
 *
 *  [FIX-1] Both weights and activations are DEQUANTIZED before the matmul,
 *          so there is no post-matmul rescale anywhere.
 *  [FIX-2] Activation quantization is blended through lambda too — the whole
 *          forward pass is continuous in lambda (smooth QAT warmup;
 *          lambda = 0 is EXACTLY a plain linear).
 *  [FIX-3] BitLinear does NO normalization of its own. The block's RMSNorm
 *          owns normalization; BitLinear does canonical BitNet b1.58
 *          absmax activation quant + absmean ternary weight quant.
 *  [FIX-4] lambda is stored as a frozen (1,)-shaped Parameter, so it is
 *          serialized in the .params file and survives checkpoint resume.
 *  [FIX-5] RMSNorm computes statistics in FP32 and casts back (BF16-safe).
 *  [FIX-6] GQA: K/V heads are broadcast to query heads (DJL has no SDPA
 *          with enable_gqa, so the materializing fallback path is used).
 *  [FIX-7] RoPE cos/sin tables are precomputed ONCE (interleaved layout,
 *          already repeat_interleave(2)-expanded), shared by all layers.
 *  [FIX-8] GPT-2-style init: residual output projections (out_proj, ffn_w2)
 *          scaled by 1/sqrt(2 * n_layers). Applied post-initialize.
 *  [FIX-9] Ternary export helper: codes in {-1, 0, +1} as INT8 plus a
 *          per-tensor FP32 gamma — the artifact a real ternary inference
 *          kernel would consume.
 *
 * NOTE on RoPE layout: this port uses INTERLEAVED rotation (GPT-NeoX /
 * llama.cpp "type 0"), matching the PyTorch file. HuggingFace Llama
 * checkpoints use HALF-SPLIT rotation — if you ever init from HF weights,
 * permute q/k projections first or you will get garbage.
 *
 * Engine note: requires the PyTorch engine (stopGradient, stepped slicing,
 * stack). Run main() for the smoke tests before any real training.
 */
public class JiRackJDLTernary_10b {

    // ===================== CONFIG (mutable for smoke mode) =====================
    public static int VOCAB_SIZE = 128256;
    public static int MODEL_DIM = 4096;          // hidden_size
    public static int FFN_HIDDEN_DIM = 18432;    // intermediate_size (JiRack 10B)
    public static int NUM_LAYERS = 32;
    public static int NUM_HEADS = 32;
    public static int NUM_KV_HEADS = 8;
    public static int HEAD_DIM = 128;
    public static int MAX_SEQ_LEN = 8192;

    public static final float RMS_EPS = 1e-5f;
    public static float ROPE_THETA = 500000.0f;
    public static final float ROPE_SCALE_FACTOR = 1.0f;
    public static final float INIT_STD = 0.02f;
    public static final float QUANT_EPS = 1e-5f;

    private static final Device DEVICE = Device.cpu();

    /** Tiny configuration for graph verification / smoke tests. */
    public static void smokeConfig() {
        VOCAB_SIZE = 256;
        MODEL_DIM = 64;
        FFN_HIDDEN_DIM = 128;
        NUM_LAYERS = 2;
        NUM_HEADS = 4;
        NUM_KV_HEADS = 2;
        HEAD_DIM = 16;
        MAX_SEQ_LEN = 64;
        RopeCache.reset();
    }

    /* ============================ Model builder ============================ */

    public static Block buildModel() {
        SequentialBlock model = new SequentialBlock();
        model.add(new CustomEmbedding(VOCAB_SIZE, MODEL_DIM));
        model.add(new FrozenSignatureLayer());
        for (int i = 0; i < NUM_LAYERS; i++) {
            model.add(new TransformerBlock());
        }
        model.add(new RMSNorm(MODEL_DIM));
        model.add(Linear.builder().setUnits(VOCAB_SIZE).optBias(false).build()); // lm_head
        return model;
    }

    /**
     * [FIX-8] GPT-2-style depth-scaled init. Call AFTER block.initialize().
     * All weights are already N(0, INIT_STD) via NormalInitializer; here the
     * residual-stream projections (out_proj, ffn_w2) are rescaled to
     * INIT_STD / sqrt(2 * NUM_LAYERS).
     */
    public static void applyDepthScaledInit(Block root) {
        float factor = (float) (1.0 / Math.sqrt(2.0 * NUM_LAYERS));
        for (Pair<String, Parameter> p : root.getParameters()) {
            String name = p.getKey();
            if (name.contains("out_proj") || name.contains("ffn_w2")) {
                p.getValue().getArray().muli(factor);
            }
        }
    }

    /** Quantization warmup hook: 0.0 = full precision, 1.0 = full fake-quant. */
    public static void setLambda(Block root, float value) {
        forEachBitLinear(root, bl -> bl.lambdaArray().set(new NDIndex(":"), value));
    }

    public static float getLambda(Block root) {
        final float[] out = {0.0f};
        final boolean[] found = {false};
        forEachBitLinear(root, bl -> {
            if (!found[0]) {
                out[0] = bl.lambdaArray().getFloat(0);
                found[0] = true;
            }
        });
        return out[0];
    }

    private static void forEachBitLinear(Block block, java.util.function.Consumer<BitLinear> fn) {
        if (block instanceof BitLinear) {
            fn.accept((BitLinear) block);
        }
        for (Pair<String, Block> child : block.getChildren()) {
            forEachBitLinear(child.getValue(), fn);
        }
    }

    /* ---------------------- RoPE cache (interleaved) ---------------------- */
    /**
     * [FIX-7] cos/sin of shape (MAX_SEQ_LEN, HEAD_DIM), already expanded for
     * interleaved rotation (each frequency duplicated into positions 2i, 2i+1).
     * Built exactly once; all layers slice [0..T).
     */
    static final class RopeCache {
        private static float[] cosTable;
        private static float[] sinTable;
        private static int builtSeqLen = -1;
        private static int builtHeadDim = -1;

        static synchronized void reset() {
            cosTable = null;
            sinTable = null;
            builtSeqLen = -1;
            builtHeadDim = -1;
        }

        private static synchronized void ensureBuilt() {
            if (cosTable != null && builtSeqLen == MAX_SEQ_LEN && builtHeadDim == HEAD_DIM) {
                return;
            }
            int d = HEAD_DIM;
            int half = d / 2;
            double[] invFreq = new double[half];
            for (int i = 0; i < half; i++) {
                double f = 1.0 / Math.pow(ROPE_THETA, 2.0 * i / d);
                if (ROPE_SCALE_FACTOR > 1.0f) {
                    f /= ROPE_SCALE_FACTOR;
                }
                invFreq[i] = f;
            }
            cosTable = new float[MAX_SEQ_LEN * d];
            sinTable = new float[MAX_SEQ_LEN * d];
            for (int t = 0; t < MAX_SEQ_LEN; t++) {
                for (int i = 0; i < half; i++) {
                    double angle = t * invFreq[i];
                    float c = (float) Math.cos(angle);
                    float s = (float) Math.sin(angle);
                    // Interleaved expansion: repeat_interleave(2)
                    cosTable[t * d + 2 * i] = c;
                    cosTable[t * d + 2 * i + 1] = c;
                    sinTable[t * d + 2 * i] = s;
                    sinTable[t * d + 2 * i + 1] = s;
                }
            }
            builtSeqLen = MAX_SEQ_LEN;
            builtHeadDim = d;
        }

        /** Returns {cos, sin}, each shaped (1, 1, T, HEAD_DIM). */
        static NDArray[] get(NDManager m, long T) {
            ensureBuilt();
            if (T > MAX_SEQ_LEN) {
                throw new IllegalArgumentException(
                        "Sequence length " + T + " exceeds MAX_SEQ_LEN " + MAX_SEQ_LEN);
            }
            int d = HEAD_DIM;
            int n = (int) T * d;
            float[] c = new float[n];
            float[] s = new float[n];
            System.arraycopy(cosTable, 0, c, 0, n);
            System.arraycopy(sinTable, 0, s, 0, n);
            return new NDArray[]{
                    m.create(c, new Shape(1, 1, T, d)),
                    m.create(s, new Shape(1, 1, T, d))
            };
        }
    }

    /* -------------------- BitLinear (b1.58 + lambda warmup) -------------------- */
    /**
     * BitNet b1.58-style fake-quant linear with lambda warmup.
     *
     * Forward at lambda = 1:
     *   gamma = mean(|W|)                       (per-tensor weight scale)
     *   W_q   = clamp(round(W / gamma), -1, 1) * gamma   // ternary, dequantized
     *   s     = 127 / absmax(x, last dim)       (per-token activation scale)
     *   x_q   = clamp(round(x * s), -128, 127) / s       // int8, dequantized
     *   out   = x_q @ W_q^T
     *
     * Both operands dequantized -> no output rescale [FIX-1].
     * STE blend applied to BOTH weights and activations [FIX-2]:
     *   eff = full + lambda * stopGradient(quant - full)
     * At lambda = 0 this is exactly a plain linear.
     * No normalization inside [FIX-3] — the block's RMSNorm owns that.
     */
    static class BitLinear extends AbstractBlock {
        private static final byte VERSION = 1;
        private final Parameter weight;   // (out, in)
        private final Parameter lambda;   // (1,) frozen buffer [FIX-4]
        private final int inFeatures;
        private final int outFeatures;

        public BitLinear(int inFeatures, int outFeatures) {
            super(VERSION);
            this.inFeatures = inFeatures;
            this.outFeatures = outFeatures;
            weight = addParameter(Parameter.builder()
                    .setName("weight")
                    .setType(Parameter.Type.WEIGHT)
                    .optShape(new Shape(outFeatures, inFeatures))
                    .build());
            lambda = addParameter(Parameter.builder()
                    .setName("lambda")
                    .setType(Parameter.Type.OTHER)
                    .optShape(new Shape(1))
                    .optInitializer(new ConstantInitializer(0.0f))
                    .optRequiresGrad(false)   // buffer, not trainable
                    .build());
        }

        NDArray lambdaArray() {
            return lambda.getArray();
        }

        NDArray weightArray() {
            return weight.getArray();
        }

        @Override
        public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {}

        @Override
        protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training,
                                         PairList<String, Object> params) {
            NDArray x = inputs.get(0);
            NDArray w = ps.getValue(weight, x.getDevice(), training);
            NDArray lam = ps.getValue(lambda, x.getDevice(), false);

            // Fast path: at lambda = 0 the blend is exactly identity, so
            // skipping the quant math is a pure compute optimization.
            if (!training && lam.getFloat(0) < 1e-6f) {
                return new NDList(x.matMul(w.transpose(1, 0)));
            }

            // === Weights: per-tensor absmean ternary (stats in FP32, [FIX-5]) ===
            NDArray gamma = w.toType(DataType.FLOAT32, true).abs().mean()
                             .maximum(QUANT_EPS).toType(w.getDataType(), false);
            NDArray wQuant = w.div(gamma).round().clip(-1.0, 1.0).mul(gamma);
            NDArray wEff = w.add(wQuant.sub(w).stopGradient().mul(lam));

            // === Activations: per-token absmax int8 [FIX-3: no norm here] ===
            NDArray xScale = x.abs().max(new int[]{-1}, true).maximum(QUANT_EPS);
            NDArray s = xScale.div(127.0f);                    // = 1 / (127/absmax)
            NDArray xQuant = x.div(s).round().clip(-128.0, 127.0).mul(s);
            NDArray xEff = x.add(xQuant.sub(x).stopGradient().mul(lam));

            // [FIX-1] Dequantized operands -> plain matmul, no rescale.
            return new NDList(xEff.matMul(wEff.transpose(1, 0)));
        }

        @Override
        public Shape[] getOutputShapes(Shape[] in) {
            Shape s = in[0];
            long[] dims = s.getShape().clone();
            dims[dims.length - 1] = outFeatures;
            return new Shape[]{new Shape(dims)};
        }
    }

    /* -------------------------- RMSNorm (FP32 stats) -------------------------- */
    static class RMSNorm extends AbstractBlock {
        private static final byte VERSION = 2;
        private final Parameter weight;

        public RMSNorm(int dim) {
            super(VERSION);
            weight = addParameter(Parameter.builder()
                    .setName("weight")
                    .setType(Parameter.Type.GAMMA)
                    .optShape(new Shape(dim))
                    .optInitializer(new ConstantInitializer(1.0f))
                    .build());
        }

        @Override
        protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training,
                                         PairList<String, Object> params) {
            NDArray x = inputs.get(0);
            DataType dtype = x.getDataType();
            NDArray gamma = ps.getValue(weight, x.getDevice(), training);
            // [FIX-5] statistics in FP32, cast back
            NDArray x32 = x.toType(DataType.FLOAT32, false);
            NDArray normed = x32.mul(
                    x32.pow(2).mean(new int[]{-1}, true).add(RMS_EPS).rsqrt());
            return new NDList(normed.mul(gamma.toType(DataType.FLOAT32, false)).toType(dtype, false));
        }

        @Override
        public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {}

        @Override
        public Shape[] getOutputShapes(Shape[] in) { return in; }
    }

    /* --------------------------- TransformerBlock --------------------------- */
    static class TransformerBlock extends AbstractBlock {
        private static final byte VERSION = 3;
        private final RMSNorm norm1, norm2;
        private final BitLinear qProj, kProj, vProj, outProj;
        private final BitLinear ffnW1, ffnW3, ffnW2;
        private final int h = NUM_HEADS;
        private final int kvH = NUM_KV_HEADS;
        private final int d = HEAD_DIM;

        public TransformerBlock() {
            super(VERSION);
            norm1 = addChildBlock("norm1", new RMSNorm(MODEL_DIM));
            norm2 = addChildBlock("norm2", new RMSNorm(MODEL_DIM));
            qProj = addChildBlock("q_proj", new BitLinear(MODEL_DIM, h * d));
            kProj = addChildBlock("k_proj", new BitLinear(MODEL_DIM, kvH * d));
            vProj = addChildBlock("v_proj", new BitLinear(MODEL_DIM, kvH * d));
            outProj = addChildBlock("out_proj", new BitLinear(h * d, MODEL_DIM));
            ffnW1 = addChildBlock("ffn_w1", new BitLinear(MODEL_DIM, FFN_HIDDEN_DIM));
            ffnW3 = addChildBlock("ffn_w3", new BitLinear(MODEL_DIM, FFN_HIDDEN_DIM));
            ffnW2 = addChildBlock("ffn_w2", new BitLinear(FFN_HIDDEN_DIM, MODEL_DIM));
        }

        @Override
        public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {
            Shape hidden = inputShapes[0];
            norm1.initialize(manager, dataType, hidden);
            norm2.initialize(manager, dataType, hidden);
            qProj.initialize(manager, dataType, hidden);
            kProj.initialize(manager, dataType, hidden);
            vProj.initialize(manager, dataType, hidden);
            outProj.initialize(manager, dataType, hidden);
            ffnW1.initialize(manager, dataType, hidden);
            ffnW3.initialize(manager, dataType, hidden);
            ffnW2.initialize(manager, dataType, ffnW1.getOutputShapes(new Shape[]{hidden}));
        }

        /** Interleaved (NeoX-style) rotation: pairs (x0,x1),(x2,x3)... */
        private NDArray applyRotaryInterleaved(NDArray x, NDArray cos, NDArray sin) {
            long dHead = x.getShape().get(3);
            NDArray xEven = x.get(new NDIndex("..., 0:" + dHead + ":2"));
            NDArray xOdd = x.get(new NDIndex("..., 1:" + dHead + ":2"));
            // rotate = interleave(-x_odd, x_even)
            NDArray rotated = NDArrays.stack(new NDList(xOdd.neg(), xEven), -1)
                                      .reshape(x.getShape());
            return x.mul(cos).add(rotated.mul(sin));
        }

        @Override
        protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training,
                                         PairList<String, Object> params) {
            NDArray x = inputs.get(0);
            NDManager m = x.getManager();
            long B = x.getShape().get(0), T = x.getShape().get(1);

            // Pre-norm attention
            NDArray hIn = norm1.forward(ps, new NDList(x), training).singletonOrThrow();

            NDArray Q = qProj.forward(ps, new NDList(hIn), training).singletonOrThrow()
                    .reshape(B, T, h, d).transpose(0, 2, 1, 3);
            NDArray K = kProj.forward(ps, new NDList(hIn), training).singletonOrThrow()
                    .reshape(B, T, kvH, d).transpose(0, 2, 1, 3);
            NDArray V = vProj.forward(ps, new NDList(hIn), training).singletonOrThrow()
                    .reshape(B, T, kvH, d).transpose(0, 2, 1, 3);

            NDArray[] cs = RopeCache.get(m, T);
            Q = applyRotaryInterleaved(Q, cs[0], cs[1]);
            K = applyRotaryInterleaved(K, cs[0], cs[1]);

            // [FIX-6] GQA fallback path: materialize K/V to query-head count
            int groups = h / kvH;
            if (groups > 1) {
                K = K.reshape(B, kvH, 1, T, d)
                        .broadcast(new Shape(B, kvH, groups, T, d)).reshape(B, h, T, d);
                V = V.reshape(B, kvH, 1, T, d)
                        .broadcast(new Shape(B, kvH, groups, T, d)).reshape(B, h, T, d);
            }

            float scale = (float) (1.0 / Math.sqrt(d));
            NDArray scores = Q.matMul(K.transpose(0, 1, 3, 2)).mul(scale);

            NDArray mask = m.arange(T).reshape(-1, 1).lt(m.arange(T).reshape(1, -1))
                    .reshape(1, 1, T, T).broadcast(new Shape(B, h, T, T));
            NDArray infMask = m.zeros(scores.getShape()).sub(1e9f);
            scores = NDArrays.where(mask, infMask, scores);

            NDArray attnOut = scores.softmax(-1).matMul(V)
                    .transpose(0, 2, 1, 3).reshape(B, T, (long) h * d);
            x = x.add(outProj.forward(ps, new NDList(attnOut), training).singletonOrThrow());

            // Pre-norm SwiGLU FFN
            NDArray mIn = norm2.forward(ps, new NDList(x), training).singletonOrThrow();
            NDArray gate = ffnW1.forward(ps, new NDList(mIn), training).singletonOrThrow();
            gate = gate.mul(gate.getNDArrayInternal().sigmoid());   // SiLU
            NDArray up = ffnW3.forward(ps, new NDList(mIn), training).singletonOrThrow();
            NDArray down = ffnW2.forward(ps, new NDList(gate.mul(up)), training).singletonOrThrow();
            x = x.add(down);

            return new NDList(x);
        }

        @Override
        public Shape[] getOutputShapes(Shape[] in) { return new Shape[]{in[0]}; }
    }

    /* -------------------------- Embedding & signature -------------------------- */
    static class CustomEmbedding extends AbstractBlock {
        private static final byte VERSION = 1;
        private final Parameter weight;
        private final int embedDim;

        public CustomEmbedding(int vocabSize, int embedDim) {
            super(VERSION);
            this.embedDim = embedDim;
            weight = addParameter(Parameter.builder()
                    .setName("embedding_weight")
                    .setType(Parameter.Type.WEIGHT)
                    .optShape(new Shape(vocabSize, embedDim))
                    .build());
        }

        @Override
        protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training,
                                         PairList<String, Object> params) {
            NDArray ids = inputs.get(0);  // integer token IDs (B, T)
            return new NDList(ps.getValue(weight, ids.getDevice(), training).get(ids));
        }

        @Override
        public Shape[] getOutputShapes(Shape[] in) {
            return new Shape[]{new Shape(in[0].get(0), in[0].get(1), embedDim)};
        }
    }

    static class FrozenSignatureLayer extends AbstractBlock {
        private static final byte VERSION = 1;
        @SuppressWarnings("unused")
        private final Parameter signatureWeight;

        public FrozenSignatureLayer() {
            super(VERSION);
            signatureWeight = addParameter(Parameter.builder()
                    .setName("H32_L32_D4096_V128256_ML8192_cmsmanhattan_model_signature")
                    .setType(Parameter.Type.OTHER)
                    .optShape(new Shape(1, 1))
                    .optInitializer(new ConstantInitializer(0.0f))
                    .optRequiresGrad(false)
                    .build());
        }

        @Override
        public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {}

        @Override
        protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training,
                                         PairList<String, Object> params) {
            return inputs;
        }

        @Override
        public Shape[] getOutputShapes(Shape[] in) { return in; }
    }

    /* ------------------------ [FIX-9] Ternary export ------------------------ */
    /**
     * Turns trained fake-quant weights into an actual ternary artifact:
     * for every BitLinear, "<name>.codes" (INT8 in {-1,0,+1}) and
     * "<name>.gamma" (scalar FP32). Full-precision parts (embedding,
     * lm_head, norms) are copied as-is. Returned NDList has named arrays
     * and can be persisted with NDList.encode().
     */
    public static NDList exportTernaryStateDict(Block root, NDManager manager) {
        NDList out = new NDList();
        exportWalk(root, "", out, manager);
        for (Pair<String, Parameter> p : root.getParameters()) {
            String name = p.getKey();
            // Everything not owned by a BitLinear stays full precision.
            if (!name.contains("weight") || name.contains("embedding")
                    || name.contains("norm") || name.contains("signature")
                    || isLmHead(name)) {
                NDArray copy = p.getValue().getArray().duplicate();
                copy.setName(name);
                out.add(copy);
            }
        }
        return out;
    }

    private static boolean isLmHead(String paramName) {
        // The final Linear in the SequentialBlock (lm_head) — its weight
        // must stay FP32; BitLinear weights are exported as codes instead.
        return paramName.endsWith("Linear.weight");
    }

    private static void exportWalk(Block block, String prefix, NDList out, NDManager manager) {
        if (block instanceof BitLinear) {
            BitLinear bl = (BitLinear) block;
            NDArray w = bl.weightArray().toType(DataType.FLOAT32, true);
            NDArray gamma = w.abs().mean().maximum(QUANT_EPS);
            NDArray codes = w.div(gamma).round().clip(-1.0, 1.0).toType(DataType.INT8, false);
            codes.setName(prefix + "codes");
            gamma.setName(prefix + "gamma");
            out.add(codes);
            out.add(gamma);
            return;
        }
        for (Pair<String, Block> child : block.getChildren()) {
            exportWalk(child.getValue(), prefix + child.getKey() + ".", out, manager);
        }
    }

    /* ------------------------------- Smoke tests ------------------------------- */
    public static void main(String[] args) throws Exception {
        System.out.println("SMOKE MODE: tiny config, mirrors the PyTorch smoke tests");
        smokeConfig();

        try (Model model = Model.newInstance("jirack", DEVICE);
             NDManager manager = NDManager.newBaseManager()) {

            Block block = buildModel();
            model.setBlock(block);
            block.setInitializer(new NormalInitializer(INIT_STD), Parameter.Type.WEIGHT);
            block.initialize(manager, DataType.FLOAT32, new Shape(2, 32));
            applyDepthScaledInit(block);   // [FIX-8]

            ParameterStore psEval = new ParameterStore(manager, false);
            NDArray ids = manager.randomInteger(0, VOCAB_SIZE, new Shape(2, 32), DataType.INT64);

            // 1) Continuity in lambda: tiny lambda must NOT jump.
            setLambda(block, 0.0f);
            NDArray y0 = block.forward(psEval, new NDList(ids), false).singletonOrThrow();
            setLambda(block, 1e-4f);
            NDArray yEps = block.forward(psEval, new NDList(ids), false).singletonOrThrow();
            setLambda(block, 1.0f);
            NDArray y1 = block.forward(psEval, new NDList(ids), false).singletonOrThrow();

            float relJump = yEps.sub(y0).norm().getFloat() / y0.norm().getFloat();
            System.out.printf("relative change at lambda=1e-4: %.2e  (must be ~1e-4, not ~1)%n", relJump);
            if (relJump >= 1e-2f) throw new AssertionError("lambda warmup is not continuous!");

            // 2) Output scale sanity at full quantization.
            float ratio = std(y1) / std(y0);
            System.out.printf("std ratio lambda=1 vs lambda=0: %.3f  (must be O(1), not ~1e-4)%n", ratio);
            if (!(ratio > 0.1f && ratio < 10.0f)) {
                throw new AssertionError("output scale collapsed or exploded!");
            }

            // 3) Gradients flow through the STE at lambda = 1.
            setLambda(block, 1.0f);
            ParameterStore psTrain = new ParameterStore(manager, false);
            try (GradientCollector gc = manager.getEngine().newGradientCollector()) {
                NDArray logits = block.forward(psTrain, new NDList(ids), true).singletonOrThrow();
                NDArray loss = logits.pow(2).mean();
                gc.backward(loss);
            }
            NDArray g = findFirstBitLinear(block).weightArray().getGradient();
            if (g == null || !g.isFinite().all().getBoolean() || g.abs().sum().getFloat() == 0f) {
                throw new AssertionError("no finite nonzero gradient through STE!");
            }
            System.out.printf("grad norm on q_proj at lambda=1: %.4f  (nonzero, finite)%n",
                    g.norm().getFloat());

            // 4) lambda survives a save/load round-trip [FIX-4].
            Path tmp = Paths.get("build/smoke_models");
            java.nio.file.Files.createDirectories(tmp);
            model.save(tmp, "jirack");
            try (Model model2 = Model.newInstance("jirack", DEVICE)) {
                Block block2 = buildModel();
                model2.setBlock(block2);
                model2.load(tmp, "jirack");
                float lam2 = getLambda(block2);
                if (Math.abs(lam2 - 1.0f) > 1e-9) {
                    throw new AssertionError("lambda not serialized! got " + lam2);
                }
                System.out.println("lambda serialization: OK");
            }

            // 5) Export produces genuinely ternary codes [FIX-9].
            NDList exported = exportTernaryStateDict(block, manager);
            long ternaryCount = exported.stream()
                    .filter(a -> a.getName() != null && a.getName().endsWith("codes"))
                    .count();
            NDArray anyCodes = exported.stream()
                    .filter(a -> a.getName() != null && a.getName().endsWith("codes"))
                    .findFirst().orElseThrow(() -> new AssertionError("no codes exported"));
            float mn = anyCodes.toType(DataType.FLOAT32, false).min().getFloat();
            float mx = anyCodes.toType(DataType.FLOAT32, false).max().getFloat();
            if (mn < -1f || mx > 1f) throw new AssertionError("codes out of {-1,0,1}!");
            System.out.printf("export: %d ternary tensors, values in [%.0f, %.0f], OK%n",
                    ternaryCount, mn, mx);

            System.out.println("\nAll smoke tests passed.");
        }
    }

    private static float std(NDArray a) {
        NDArray f = a.toType(DataType.FLOAT32, false);
        NDArray mean = f.mean();
        return (float) Math.sqrt(f.sub(mean).pow(2).mean().getFloat());
    }

    private static BitLinear findFirstBitLinear(Block root) {
        final BitLinear[] found = {null};
        forEachBitLinear(root, bl -> {
            if (found[0] == null) found[0] = bl;
        });
        return found[0];
    }
}