Living Weights on Qwen4 dense, step 1: the QSA backward on the host, against binary64 (#1804)

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- Living Weights on Qwen4 dense, step 1: the QSA backward pass on the host, checked against a binary64 reference; nothing served changes (#2021, part of #1804).
+
 - CUDA kernels share their duplicated device primitives; the machine code is identical on sm_80, 86, 89, 120 and 121 (#2018, closes #1409).
 
 - Metal: Qwen3.8-27B's keep gathers and scatters the binding's KV heads a row, so a keep stages what the layout holds (#2013). Thanks to @cshintov.

docs/recipes/qwen4-dense.md

@@ -20,6 +20,9 @@ Qwen4 uses the shared projection and pointwise operators and geometry specializa
 Existing Flash Next factories retain their fixed checkpoint dimensions and their arithmetic.
 The new native adapters use the existing engine API, lane host, prompt cache and verification machinery.
 
+Living Weights currently has a [host QSA backward and CUDA/Metal implementation plan](qwen4-living-weights.md).
+That mathematical check does not enable a Qwen4 learner or qualify a device gradient.
+
 Weights remain as stored: BF16, MLX affine 4/8-bit groups 32/64/128, ModelOpt FP8 or ModelOpt NVFP4.
 Each projection keeps its own recipe, scales and bias metadata.
 Activations stay BF16 and reductions stay FP32.

docs/recipes/qwen4-living-weights.md

@@ -0,0 +1,101 @@
+# Qwen4 Living Weights: host backward and implementation plan
+
+This is the host mathematical step for #1804. It does not enable learning in the native Qwen4 engine.
+`families/qwen4/qsa_backward.zig` computes the continuous selected-attention adjoint for FP32 or FP64 arrays.
+Q and upstream gradients are `[rows, heads, width]`; K/V and their adjoints are token-major
+`[capacity, kv_heads, width]`. Query head `h` reads KV head `h / (heads / kv_heads)`.
+The served QSA scale is 0.0625 at width 256. The host routine accepts an explicit scale for synthetic checks.
+
+For each causal row, dense attention reads keys `0..end`. Sparse attention reads exactly the supplied unique IDs,
+including any incomplete causal tail. IDs beyond `end`, duplicate IDs and empty selections are refused.
+The selection is a fixed tape. No surrogate derivative is supplied for pooling, score ranking or top-k decisions.
+Unselected and future K/V rows are never read and their adjoints stay zero. Shared K/V adjoints accumulate every
+query-head and row use. Shape, selection, alias and input-finiteness admission precedes gradient writes.
+Numeric overflow clears the output adjoints and returns an error. Inputs and the selection tape remain unchanged.
+
+The proved attention formula is `p = softmax(scale * Q K^T)`, `O = p V` over the fixed selection:
+
+```text
+dP[j] = dot(dO, V[j])
+dS[j] = p[j] * (dP[j] - sum_i p[i] dP[i])
+dQ += scale * dS[j] * K[j]
+dK[j] += scale * dS[j] * Q
+dV[j] += p[j] * dO
+```
+
+The separate output-gate pullback takes saved forward context and sigmoid values. It uses the saved sigmoid for
+the context adjoint and the ordinary sigmoid derivative for the logit adjoint. Forward cuts are straight-through
+in this local mathematical model. Independent finite differences freeze each cast's offset at the base point;
+they do not differentiate a discontinuous rounding function. Both cut and uncut output-gate cases are checked.
+
+```sh
+zig build test-qwen4-qsa-backward-host -j1 --summary all
+```
+
+The independent reference evaluates a dense binary64 masked loss rather than calling the backward's helpers.
+Tiny cases check every Q/K/V and gated-output coordinate at three central step sizes. Directional cases use the
+served 24 query heads, 4 KV heads and width 256, with 129 dense keys and 2,051 sparse keys including a three-token
+tail. Tests also cover poisoned unused cache rows, shared-head accumulation, malformed tapes, aliases and repeats.
+FP64 checks use `1e-8 + 1e-6 * abs(slope)`; FP32 checks use `2e-7 + 3e-5 * abs(slope)`.
+These are host derivative-check limits, not new device fidelity or whole-model acceptance thresholds.
+
+The CUDA tensor-core attention path stores BF16 probability operands; Metal's split softmax retains FP32
+probabilities. This step does not model those backend-specific probability cuts, split reduction orders, Q/K
+RMS/RoPE or their stored intermediates. It is not a gradient receipt for either device profile. The native forward,
+its probability stores and all shader sources remain unchanged. Future backward kernels need their own contracts.
+
+## CUDA pocket and trainer
+
+1. Adapt the dense pocket storage and shard writer used by Qwen3.5, with a Qwen4 family descriptor keyed by its
+   declared config. Preserve original tensor bytes and separate base columns from appended neurons when loading.
+   Use the same serialized family/tensor/quantization description on both backends. Reject unsupported learning
+   layouts before allocation; a group-32 Flash Next transpose must never be applied to group-64/128 or other formats.
+2. Add hooks at the dense MLP branch before the HC writeback in the shared Flash Next layer graph. Every prompt,
+   verify and decode pass must run the same hook. Keep the empty-pocket branch free of allocations, readbacks,
+   synchronization and added launches. Supply per-stream untaught masks, including mixed taught/untaught batches.
+3. Reuse Flash Next's HC branch/readout backward and GDN backward with explicit Qwen4 geometry and gate mode.
+   Qwen4's default SiLU must not inherit Flash Next's sigmoid derivative. Make the backward mode explicit and prove
+   both modes with a planted wrong-gate control before reuse, preserving Flash Next's existing mode and arithmetic.
+   Add dense SwiGLU/base-weight and pocket-input gradients, the LM head/loss
+   path, and optional PLE backward. Until PLE backward is present, learning on PLE configs must refuse explicitly.
+4. Record Q/K norm inputs, stored rotary frequencies, V, output gates and the exact selected IDs/counts/row ends
+   from the same forward. The backward must use this tape, not reselect. Implement selected softmax backward in
+   FP32, accumulating shared K/V uses without lost writes. Integrate the Q/K RMS, rotary, projections and HC chain.
+   Keep each sequence's cache ownership explicit; packed training sequences must not read another sequence's keys.
+5. Define and test the backward convention for every probability, context and gate cut in the actual CUDA profile.
+   Finite differences of the ideal formula above cannot qualify BF16 probability stores. Check local cut offsets
+   and the complete computed graph, with controls that drop a cut or transpose a GQA group.
+
+## Metal pocket and trainer
+
+Reuse the shared family descriptor, learned-neuron layout, fitting policy and shard writer. Follow the existing
+Nemotron Metal trainer's ownership and residency structure, with Qwen4 operators and its own forward tape. Device
+buffers and captured metadata must survive command completion. Keep training scratch separate from served streams;
+drain outstanding work before changing pocket buffers and invalidate prompt states after any successful edit.
+
+Port the qualified CUDA backward stages as Metal operators without replacing inference with CPU work. Reuse HC,
+GDN and dense projection derivatives only where layout, precision and gate mode match. Register any new operators
+and their served shapes, exceptional classes, repeats/aliasing, accumulator stress and launch limits. Pending
+compiler/resource records stay unmeasured until actual images are checked. Include pre-M5 and M5 implementations.
+
+## Gates before advertising learning
+
+Use a tiny nonzero Qwen4-shaped model to check whole trainer/serving loss and directional gradients at early rows,
+layer boundaries, GDN/QSA and PLE layers. A whole-model finite-difference probe must record whether its selected
+IDs change under perturbation. Selection crossings are nondifferentiable witnesses, not permission to widen the
+gradient bound or silently freeze a real serving pass. The local fixed-selection check remains a separate receipt.
+
+Test fitting, held-out recall, unrelated-fact retention, exact take-out/delete/rollback, cancellation and save failure
+through the existing shared learner/server acceptance path. Round outputs exactly as they will be stored before
+judging a fit. Compare the immutable base tensor bytes before/after, and reopen each backend's saved folder on the
+other. Assert that untaught streams read the original model even beside taught streams.
+
+Each affected checkpoint and profile is then run on a CUDA GPU and on pre-M5 and M5 Macs. Match binary/runtime,
+model/tokenizer revisions, prompts, seeds, quantization and cache settings. Require empty-folder tokens and speed
+equal to the base, drafted/plain and concurrent/solo equality, every prompt cut, resume/restart and rejection cleanup.
+Measure learning time and peak host/device memory separately. No device, real-checkpoint or performance gate is
+claimed by this host step; the portable layout and complete CUDA/Metal trainers remain follow-up work.
+
+Existing owners are the MLP HC seam in `families/flashnext/forward_graph.zig:34`, the CUDA HC/GDN trainer chain in
+`families/flashnext/cuda_train.zig:369`, shared GDN backward in `cuda/gdn_train.zig:299`, and Qwen4's indexed-attention
+and output-gate adapters in `families/qwen4/gpu_ops.zig:169` and `families/qwen4/metal_gpu.zig:196` under `zig/src/`.

zig/build/qwen4.zig

@@ -37,6 +37,8 @@ pub fn steps(b: *std.Build, parent: *std.Build.Step) void {
     ple_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{"Qwen4 PLE native contract"} })).step);
     const qsa_step = b.step("test-qwen4-qsa-host", "Generated native QSA bodies against independent scalar references");
     qsa_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{ "Qwen4 QSA native contract", "Qwen4 QSA scores", "Qwen4 QSA selection" } })).step);
+    const backward = b.step("test-qwen4-qsa-backward-host", "Fixed-selection QSA adjoints against independent binary64 finite differences");
+    backward.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{"Qwen4 QSA backward:"} })).step);
     const device_host = b.step("test-qwen4-device-host", "The full native device gate under numeric host instruments and refusal controls");
     device_host.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m, .filters = &.{"Qwen4 device gate:"} })).step);
     const reference_step = b.step("test-qwen4-reference-host", "Qwen4 reference precision, rotary inputs and queued capture fault controls");

zig/src/families/qwen4/qsa_backward.zig

@@ -0,0 +1,143 @@
+//! Host QSA adjoints with fixed selection and token-major GQA caches; no production inference fallback.
+const std = @import("std");
+
+pub const Shape = struct { rows: usize, heads: usize, kv: usize, width: usize, capacity: usize, scale: f64 = 0.0625 };
+/// Dense rows read 0..end; sparse rows read exactly these unique causal token IDs, including the incomplete tail.
+pub const Row = struct { end: usize, ids: ?[]const u32 = null };
+
+fn times(a: usize, b: usize, c: usize) !usize {
+    return std.math.mul(usize, try std.math.mul(usize, a, b), c);
+}
+fn keyAt(row: Row, j: usize) usize {
+    return if (row.ids) |ids| ids[j] else j;
+}
+fn keyCount(row: Row) usize {
+    return if (row.ids) |ids| ids.len else row.end;
+}
+fn scalar(comptime T: type) void {
+    if (T != f32 and T != f64) @compileError("QSA host backward requires f32 or f64");
+}
+fn disjoint(comptime T: type, output: []T, other: []const T) !void {
+    const a = @intFromPtr(output.ptr);
+    const b = @intFromPtr(other.ptr);
+    const n = try std.math.mul(usize, output.len, @sizeOf(T));
+    const m = try std.math.mul(usize, other.len, @sizeOf(T));
+    if (n != 0 and m != 0 and (if (a <= b) b - a < n else a - b < m)) return error.QsaGradientAlias;
+}
+fn finite(comptime T: type, values: []const T) !void {
+    for (values) |v| if (!std.math.isFinite(v)) return error.QsaGradientNonfinite;
+}
+
+/// q/dout/dq [rows,heads,width], k/v/dk/dv [capacity,kv,width]; overwrites adjoints, accumulates shared keys.
+pub fn backward(comptime T: type, a: std.mem.Allocator, g: Shape, rows: []const Row, q: []const T, k: []const T, v: []const T, dout: []const T, dq: []T, dk: []T, dv: []T) !void {
+    scalar(T);
+    if (g.rows == 0 or g.heads == 0 or g.kv == 0 or g.width == 0 or g.capacity == 0 or
+        g.heads % g.kv != 0 or !std.math.isFinite(g.scale) or g.scale <= 0 or rows.len != g.rows)
+        return error.QsaGradientShape;
+    const scale: T = @floatCast(g.scale);
+    if (!std.math.isFinite(scale) or scale <= 0) return error.QsaGradientShape;
+    const nq = try times(g.rows, g.heads, g.width);
+    const nk = try times(g.capacity, g.kv, g.width);
+    if (q.len != nq or dout.len != nq or dq.len != nq or k.len != nk or v.len != nk or dk.len != nk or dv.len != nk)
+        return error.QsaGradientShape;
+    const outputs = [_][]T{ dq, dk, dv };
+    for (outputs, 0..) |out, i| {
+        for ([_][]const T{ q, k, v, dout }) |input| try disjoint(T, out, input);
+        for (outputs[0..i]) |prior| try disjoint(T, out, prior);
+    }
+    var largest: usize = 0;
+    for (rows) |row| {
+        if (row.end == 0 or row.end > g.capacity or keyCount(row) == 0 or keyCount(row) > row.end) return error.QsaGradientSelection;
+        if (row.ids) |ids| for (ids, 0..) |id, i| {
+            if (id >= row.end or std.mem.indexOfScalar(u32, ids[0..i], id) != null) return error.QsaGradientSelection;
+        };
+        largest = @max(largest, keyCount(row));
+    }
+    try finite(T, q);
+    try finite(T, dout);
+    for (rows) |row| for (0..keyCount(row)) |j| {
+        const base = keyAt(row, j) * g.kv * g.width;
+        try finite(T, k[base..][0 .. g.kv * g.width]);
+        try finite(T, v[base..][0 .. g.kv * g.width]);
+    };
+    const probabilities = try a.alloc(T, largest);
+    defer a.free(probabilities);
+    const dp = try a.alloc(T, largest);
+    defer a.free(dp);
+    for (outputs) |out| @memset(out, 0);
+    errdefer for (outputs) |out| @memset(out, 0);
+    for (rows, 0..) |row, r| for (0..g.heads) |h| {
+        const head = h / (g.heads / g.kv);
+        const at = (r * g.heads + h) * g.width;
+        const count = keyCount(row);
+        var top = -std.math.inf(T);
+        for (probabilities[0..count], 0..) |*score, j| {
+            const base = (keyAt(row, j) * g.kv + head) * g.width;
+            score.* = 0;
+            for (0..g.width) |d| score.* += q[at + d] * k[base + d];
+            score.* *= scale;
+            if (!std.math.isFinite(score.*)) return error.QsaGradientNonfinite;
+            top = @max(top, score.*);
+        }
+        var total: T = 0;
+        for (probabilities[0..count]) |*p| {
+            p.* = @exp(p.* - top);
+            total += p.*;
+        }
+        var center: T = 0;
+        for (probabilities[0..count], dp[0..count], 0..) |*p, *derivative, j| {
+            p.* /= total;
+            derivative.* = 0;
+            const base = (keyAt(row, j) * g.kv + head) * g.width;
+            for (0..g.width) |d| derivative.* += dout[at + d] * v[base + d];
+            center += p.* * derivative.*;
+        }
+        for (probabilities[0..count], dp[0..count], 0..) |p, derivative, j| {
+            const ds = p * (derivative - center) * scale;
+            const base = (keyAt(row, j) * g.kv + head) * g.width;
+            for (0..g.width) |d| {
+                dq[at + d] += ds * k[base + d];
+                dk[base + d] += ds * q[at + d];
+                dv[base + d] += p * dout[at + d];
+            }
+        }
+    };
+    for (outputs) |out| try finite(T, out);
+}
+
+fn sigmoid(comptime T: type, x: T) T {
+    const e = @exp(-@abs(x));
+    return if (x >= 0) 1 / (1 + e) else e / (1 + e);
+}
+
+/// Stored context/sigmoid retain forward cuts; cast adjoints are straight-through, not rounding derivatives.
+pub fn gateBackward(comptime T: type, context: []const T, logits: []const T, stored_sigmoid: []const T, dout: []const T, dx: []T, dg: []T) !void {
+    scalar(T);
+    const n = context.len;
+    if (n == 0 or logits.len != n or stored_sigmoid.len != n or dout.len != n or dx.len != n or dg.len != n)
+        return error.QsaGradientShape;
+    for ([_][]T{ dx, dg }) |out| for ([_][]const T{ context, logits, stored_sigmoid, dout }) |input| try disjoint(T, out, input);
+    try disjoint(T, dx, dg);
+    for ([_][]const T{ context, logits, stored_sigmoid, dout }) |input| try finite(T, input);
+    for (stored_sigmoid) |s| if (s < 0 or s > 1) return error.QsaGradientShape;
+    for (context, logits, stored_sigmoid, dout, dx, dg) |x, gate, saved, dy, *to_x, *to_gate| {
+        const s = sigmoid(T, gate);
+        to_x.* = dy * saved;
+        to_gate.* = dy * x * s * (1 - s);
+    }
+    if (finite(T, dx)) |_| {
+        finite(T, dg) catch |err| {
+            @memset(dx, 0);
+            @memset(dg, 0);
+            return err;
+        };
+    } else |err| {
+        @memset(dx, 0);
+        @memset(dg, 0);
+        return err;
+    }
+}
+
+test {
+    _ = @import("qsa_backward_test.zig");
+}