diff --git a/CHANGELOG.md b/CHANGELOG.md index 41d25ebdfc..50a2c7adca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Drama 3 on Metal reads the pruned NVFP4 checkpoint's packed E2M1 experts, block-16 E4M3 scales and F32 global scales without requantizing or expanding the experts. Routing masks dropped original IDs before softmax and maps selected IDs to compact slots. The full FP8-expert path retains its existing arithmetic. +- Metal Drama speech supports opt-in semantic classifier-free guidance with shared weights and independent conditional/unconditional backbone state. `--cfg-gamma`, `--cfg-allow-request-gamma`, `--cfg-force-guided` and `--cfg-batched` preserve builtin unguided voices, validate request gamma, and reset guidance on cancellation or reference changes. Paired mode submits both states in one command buffer, not a vectorized GEMM; admission includes the extra state. CUDA guidance is not enabled by this change. + - Drama speech on Metal reuses a catalog voice's reference prefill across requests on the same lane, so a request with the same voice prefills only its own text. Audio is byte-identical to an uncached request. On an M3 Ultra with Drama 3 and a catalog voice, the first audio byte of a streamed reply comes 90% sooner (15.6 s to 1.5 s over four texts) and whole replies take about half as long; the default voice and the first request with a voice are unchanged (#1237). - Flash Next on Macs without tensor units decodes faster with the same tokens: the hyper-connection act pair runs for windows of two rows and more, the HC down projection runs without staging or barriers at four outputs a threadgroup, rounds are gated on the GPU with a GPU-only event between them, drafts are gated at tau 0.4, the decode and draft argmaxes take two launches, and wide windows use grouped experts and exact tree windows. On an M3 Ultra with Qwen3.8 Flash Next MLX 4-bit MTP, the served sweep goes from 120.2 to 147.6 tokens/s on one stream (+22.7%), 129.7 to 155.3 at two (+19.8%), 143.8 to 167.0 at four (+16.2%) and 154.4 to 176.1 at eight (+14.1%); prefill is unchanged and every token is identical. diff --git a/README.md b/README.md index 795ffe2b6c..cf69ddec8c 100644 --- a/README.md +++ b/README.md @@ -179,6 +179,7 @@ Conversions are explicitly marked: their upstream weights need the linked conver | Kolibri-1, block FP8 (research) | `tensorfold pull Aleph-Alpha/Kolibri-1@e52eb4627d11516b0c01de49210ab5a4e4061444` | One NVIDIA GPU; GB10 measured; 74 GB checkpoint. | [Kolibri](docs/recipes/kolibri-1.md). Text only; copy drafts, no EAGLE-3/neural drafter. | | Liquid AI d1-3B, decisions | Not documented yet. | Native CUDA; RAM minimum not published. | [API section below](#serve-flags). `/v1/systemone` and `/v1/decisions`. | | Fish Audio Drama 3, speech | Not documented yet. | Native Metal or CUDA; GB10 measured. | [Release audio instructions](CHANGELOG.md#110). `/v1/audio/speech`; checkpoint and `codec.pth` required. | +| Drama 3 pruned NVFP4, speech | `hf download plotarmordev/drama-3-NVFP4-Pruned --revision 494db08481db952f834b2241bce63fb3dfa72554 --local-dir drama-3-NVFP4-Pruned` | Native Metal; roughly 21 GiB checkpoint plus runtime state. Not a minimum-RAM or cross-backend quality claim. | Packed experts, BF16 dense/tower weights and FP32 accumulation; no conversion. Use `serve drama-3-NVFP4-Pruned --cfg-gamma 1.5 --cfg-batched` with an attested catalog voice. | ### GLM-5.3: one machine or two? @@ -287,6 +288,10 @@ The binary's `capabilities --json` response lists its supported flags and platfo | `--drafter-bits 0\|4` | Drafter weight precision where the model offers a choice: 4 bits (the default), or 0 to keep bf16. Qwen3.8-27B on CUDA takes 4 only. | | `--record DIR` | Write each stream's committed tokens, positions and tapped rows under DIR, for drafter training. CUDA Nemotron only; a single stream runs through the lane rounds, and recording stops below 25 GiB free disk. | | `--keep-warm SECONDS` | Keep the Metal GPU active while idle for this long after a request, default 900; zero disables it. | +| `--cfg-gamma GAMMA` | Metal Drama speech: semantic guidance strength, finite 1–5; 0/default disables it. Only reference-conditioned requests are guided. Enabled replies report `X-Cfg-Gamma`; codec heads stay conditional. | +| `--cfg-allow-request-gamma` | Metal Drama speech: allow a finite 1–5 `X-Cfg-Gamma` request header; invalid values return 400 before engine submission. Without this opt-in, the server's gamma is fixed. | +| `--cfg-force-guided` | With an enabled Metal Drama guidance policy, exercise the unconditional state even at gamma 1. The semantic mix at 1 remains exact conditional-logit identity. | +| `--cfg-batched` | With enabled Metal Drama guidance, submit both independent backbone states in one command buffer rather than separate submissions. Weights are shared; extra state is included in admission. Not vectorized batching. | | `--silence-floor-dbfs DBFS` | Drama speech only: explicit frame-power floor for the silence guard, with no default floor. Calibrate below quiet intended speech using public or built-in speech. | | `--silence-seconds SECONDS` | Drama speech only: stop after this continuous duration below the explicit floor, retaining audio and reporting `silence`. Unset by default: off; calibrate above intentional pauses. | | `--silence-onset-s SECONDS` | Drama speech only: with the silence guard on, stop once this many seconds run below its floor while the synthesis has voiced less than `--silence-onset-voiced-s`, retaining audio and reporting `silence`. Off unless given; must be below `--silence-seconds`. | diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index c696b5254e..5f9393acd0 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -419,6 +419,8 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st .imports = &.{ .{ .name = "engine_api", .module = native_modules.api }, .{ .name = "tokenizer", .module = mods.tokenizer }, .{ .name = "template", .module = template }, .{ .name = "living", .module = mods.living } }, }); const server_json = @import("server.zig").addTo(b, server_tests, host, .debug); + b.step("test-drama-cfg-host", "Drama CFG CLI, HTTP and voice transitions without GPU execution").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"CFG"} })).step); + b.step("test-speech-redraw-deadline-host", "Busy speech redraw expiry without counting an unsubmitted take").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"a busy redraw stops at its latest start"} })).step); @import("proposal_vocab_host.zig").steps(b, mods.core); const eos_policy = b.step("test-eos-policy-host", "Shared EOS, real HfText and HTTP policy without a GPU"); eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"EOS policy:"} })).step); diff --git a/zig/build/drama3_host.zig b/zig/build/drama3_host.zig index 48e6b9385c..9108859473 100644 --- a/zig/build/drama3_host.zig +++ b/zig/build/drama3_host.zig @@ -46,6 +46,11 @@ pub fn metal(b: *std.Build, core: *std.Build.Module, api: *std.Build.Module) voi const step = b.step("check-drama-metal-host", "Compile Drama's Metal host entry points without compiling shaders"); compile(b, step, family, "drama_metal_host"); native(b, step, family, core, api, null, "drama_metal_native_host"); + const quant = b.addRunArtifact(b.addTest(.{ + .root_module = family, + .filters = &.{ "NVFP4", "expert metadata", "independent scalar reference", "CFG" }, + })); + b.step("test-drama-metal-quant-host", "Packed expert formats, pruning and CFG state without GPU execution").dependOn(&quant.step); } /// Drama's prompt plan on the host: the single-row default and the opt-in batched plan. diff --git a/zig/build/drama3_metal.zig b/zig/build/drama3_metal.zig index 7932179e45..a74e027016 100644 --- a/zig/build/drama3_metal.zig +++ b/zig/build/drama3_metal.zig @@ -2,7 +2,7 @@ const std = @import("std"); const dir = "zig/kernels/metal/drama3"; -const sources = [_][]const u8{ "dense.metal", "gdn.metal", "attn.metal", "moe.metal", "sample.metal", "codec.metal" }; +const sources = [_][]const u8{ "dense.metal", "gdn.metal", "attn.metal", "moe.metal", "nvfp4.metal", "guidance.metal", "sample.metal", "codec.metal" }; /// Drama in the Metal engines: its `core` is the engine's safetensors, the server's tokenizer and frame_rounds. pub fn addNative(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, native: *std.Build.Module, metal_std: []const u8) void { @@ -65,6 +65,17 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "core", .module = core }, .{ .name = "drama3_metallib", .module = bytes } }, }); + const quant_check = b.addExecutable(.{ + .name = "drama3-metal-quant-check", + .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/drama_nvfp4_metal.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "drama3", .module = family }, .{ .name = "drama3_metallib", .module = bytes } }, + }), + }); + b.step("drama3-metal-quant-check", "Build the NVFP4 Metal operator checker; does not execute GPU work").dependOn(&b.addInstallArtifact(quant_check, .{}).step); const tool = b.createModule(.{ .root_source_file = b.path("zig/src/families/drama3/generate.zig"), .target = target, diff --git a/zig/build/kernel_contract.zig b/zig/build/kernel_contract.zig index 0a60668baa..d8a7ca1c8b 100644 --- a/zig/build/kernel_contract.zig +++ b/zig/build/kernel_contract.zig @@ -36,6 +36,17 @@ pub fn hostTests(b: *std.Build, family: *std.Build.Step) void { }) })); b.step("test-delta-slots-contract-host", "FP64 DeltaNet and slot corruption controls without shaders").dependOn(&delta.step); step.dependOn(&delta.step); + const nvfp4_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/contract/drama3_nvfp4_ref.zig"), + .target = b.graph.host, + .optimize = .debug, + .imports = &.{.{ .name = "contract", .module = module(b, b.graph.host, .debug) }}, + }); + const nvfp4 = b.addRunArtifact(b.addTest(.{ .root_module = nvfp4_module })); + b.step("test-drama-nvfp4-reference-host", "NVFP4 scalar reference and reduced-accumulator discrimination").dependOn(&nvfp4.step); + step.dependOn(&nvfp4.step); + const shared_gate = b.addRunArtifact(b.addTest(.{ .root_module = nvfp4_module, .filters = &.{"shared gate stress"} })); + b.step("test-drama-combine-stress-host", "Shared-gate stress remains discriminating after nonlinear output").dependOn(&shared_gate.step); } /// macOS: tf-kernel-contract over the engine's Metal kernels; `kernel-contract` runs it on this Mac's GPU. diff --git a/zig/kernels/metal/drama3/guidance.metal b/zig/kernels/metal/drama3/guidance.metal new file mode 100644 index 0000000000..78117694f3 --- /dev/null +++ b/zig/kernels/metal/drama3/guidance.metal @@ -0,0 +1,14 @@ +#include +using namespace metal; + +struct D3CfgArgs { uint count; float gamma; }; + +// Candidate logits include semantic tokens and end-of-speech; residual codec heads remain conditional. +[[max_total_threads_per_threadgroup(256)]] kernel void d3_cfg_mix(device float* conditional [[buffer(0)]], + device const float* unconditional [[buffer(1)]], + constant D3CfgArgs& args [[buffer(2)]], + uint i [[thread_position_in_grid]]) { + if (i >= args.count || args.gamma == 1.0f) return; + float c = conditional[i]; + conditional[i] = c + (args.gamma - 1.0f) * (c - unconditional[i]); +} diff --git a/zig/kernels/metal/drama3/nvfp4.metal b/zig/kernels/metal/drama3/nvfp4.metal new file mode 100644 index 0000000000..451869bf6c --- /dev/null +++ b/zig/kernels/metal/drama3/nvfp4.metal @@ -0,0 +1,135 @@ +// Packed row-major NVFP4: E2M1 nibbles, unsigned E4M3 block16 scales, FP32 projection globals. +#include "d3_common.h" + +struct Fp4Projection { device const uchar* w; device const uchar* s; device const float* global; }; +struct Fp4Expert { Fp4Projection g, u, d; }; +struct Fp4Args { uint K, N, prompt, cap, slots; }; + +inline float fp4_value(uchar code) { + const float values[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + return (code & 8) ? -values[code & 7] : values[code & 7]; +} +inline float fp4_scale(uchar code) { + const uint exp = (code >> 3) & 15, mant = code & 7; + return exp ? ldexp(1.0f + float(mant) / 8.0f, int(exp) - 7) : float(mant) / 512.0f; +} +inline float fp4_row(Fp4Projection w, uint n, uint K, device const bf16* x, uint lane) { + float acc = 0.0f; + for (uint k = lane * 16; k < K; k += 512) { + float block = 0.0f; + for (uint j = 0; j < 16; ++j) { + const uchar b = w.w[(ulong)n * (K / 2) + (k + j) / 2]; + const uchar code = (j & 1) ? (b >> 4) : (b & 15); + block = fma(fp4_value(code), b2f(x[k + j]), block); + } + acc = fma(block, fp4_scale(w.s[(ulong)n * (K / 16) + k / 16]), acc); + } + return simd_sum(acc) * w.global[0]; +} + +[[max_total_threads_per_threadgroup(128)]] kernel void d3_nvfp4_up( + device const bf16* x [[buffer(0)]], device const Fp4Expert* table [[buffer(1)]], + device const int* picks [[buffer(2)]], device const int* counts [[buffer(3)]], + device const int* members [[buffer(4)]], device bf16* out [[buffer(5)]], + constant Fp4Args& a [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + const uint n = tg.x * 4 + sg; + const int e = a.prompt ? int(tg.y) : picks[tg.y]; + const int cnt = a.prompt ? counts[e] : 1; + if (!cnt) return; + const Fp4Expert w = table[e]; + for (int j = 0; j < cnt; ++j) { + const uint pair = a.prompt ? uint(members[(ulong)e * a.cap + j]) : tg.y; + device const bf16* row = x + (ulong)(pair / a.slots) * a.K; + const float g = rb(fp4_row(w.g, n, a.K, row, lane)); + const float u = rb(fp4_row(w.u, n, a.K, row, lane)); + if (lane == 0) out[(ulong)pair * a.N + n] = f2b(g / (1.0f + precise::exp(-g)) * u); + } +} + +[[max_total_threads_per_threadgroup(128)]] kernel void d3_nvfp4_down( + device const bf16* x [[buffer(0)]], device const Fp4Expert* table [[buffer(1)]], + device const int* picks [[buffer(2)]], device const int* counts [[buffer(3)]], + device const int* members [[buffer(4)]], device float* out [[buffer(5)]], + constant Fp4Args& a [[buffer(6)]], uint2 tg [[threadgroup_position_in_grid]], + uint sg [[simdgroup_index_in_threadgroup]], uint lane [[thread_index_in_simdgroup]]) { + const uint n = tg.x * 4 + sg; + const int e = a.prompt ? int(tg.y) : picks[tg.y]; + const int cnt = a.prompt ? counts[e] : 1; + if (!cnt) return; + const Fp4Expert w = table[e]; + for (int j = 0; j < cnt; ++j) { + const uint pair = a.prompt ? uint(members[(ulong)e * a.cap + j]) : tg.y; + const float v = fp4_row(w.d, n, a.K, x + (ulong)pair * a.K, lane); + if (lane == 0) out[(ulong)pair * a.N + n] = v; + } +} + +struct CompactRouterArgs { uint ss, off; }; +// Retained original IDs compete before softmax/top8; only the selected IDs are translated to compact slots. +[[max_total_threads_per_threadgroup(256)]] kernel void d3_router_compact( + device const bf16* sp [[buffer(0)]], device int* picks [[buffer(1)]], + device float* weights [[buffer(2)]], device const int* slots [[buffer(3)]], + constant CompactRouterArgs& a [[buffer(4)]], uint t [[threadgroup_position_in_grid]], + uint e [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float scratch[32], wv[8], best_v[8]; + threadgroup int wi[8], best_i[8]; + const bool kept = slots[e] >= 0; + const float lg = kept ? b2f(sp[(ulong)t * a.ss + a.off + e]) : -INFINITY; + const float m = group_max(lg, scratch, sg, lane, 8); + const float p = precise::exp(lg - m); + const float sum = group_sum(p, scratch + 8, sg, lane, 8); + // Poisoned probabilities must propagate NaNs without indexing a dropped expert or selecting one twice. + const bool poisoned = !isfinite(sum) || sum <= 0.0f; + float mine = kept ? (poisoned ? 0.0f : p / sum) : -1.0f; + for (int k = 0; k < 8; ++k) { + float v = mine; + int idx = int(e); + for (ushort o = 16; o > 0; o >>= 1) { + const float ov = simd_shuffle_xor(v, o); + const int oi = simd_shuffle_xor(idx, o); + if (ov > v || (ov == v && oi < idx)) { v = ov; idx = oi; } + } + if (lane == 0) { wv[sg] = v; wi[sg] = idx; } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (e == 0) { + float bv = wv[0]; + int bi = wi[0]; + for (int j = 1; j < 8; ++j) + if (wv[j] > bv || (wv[j] == bv && wi[j] < bi)) { bv = wv[j]; bi = wi[j]; } + best_v[k] = bv; + best_i[k] = bi; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (int(e) == best_i[k]) mine = -1.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (e < 8) { + float total = 0.0f; + for (int j = 0; j < 8; ++j) total += best_v[j]; + picks[t * 8 + e] = slots[best_i[e]]; + weights[t * 8 + e] = poisoned ? NAN : best_v[e] / total; + } +} + +// CUDA drama_moe_combine_f32: round each weighted pair, sum, then fuse the FP32 shared gate and add residual. +[[max_total_threads_per_threadgroup(256)]] kernel void d3_nvfp4_combine( + device bf16* residual [[buffer(0)]], device const float* down [[buffer(1)]], + device const float* weights [[buffer(2)]], device const bf16* shared [[buffer(3)]], + device const bf16* x [[buffer(4)]], device const bf16* wg [[buffer(5)]], + constant uint& D [[buffer(6)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_index_in_threadgroup]], uint sg [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) { + threadgroup float scratch[32]; + float part = 0.0f; + for (uint i = t; i < D; i += 256) part += b2f(x[(ulong)row * D + i]) * b2f(wg[i]); + const float gate = d3_sigmoid(group_sum(part, scratch, sg, lane, 8)); + for (uint i = t; i < D; i += 256) { + float sum = 0.0f; + for (uint slot = 0; slot < 8; ++slot) + sum += rb(weights[row * 8 + slot] * down[((ulong)row * 8 + slot) * D + i]); + const float moe = rb(gate * b2f(shared[(ulong)row * D + i]) + rb(sum)); + residual[(ulong)row * D + i] = f2b(b2f(residual[(ulong)row * D + i]) + moe); + } +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 89047df093..b71562cb11 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -178,6 +178,8 @@ pub const drama3 = struct { pub const gdn = @embedFile("drama3/gdn.metal"); pub const attn = @embedFile("drama3/attn.metal"); pub const moe = @embedFile("drama3/moe.metal"); + pub const nvfp4 = @embedFile("drama3/nvfp4.metal"); + pub const guidance = @embedFile("drama3/guidance.metal"); pub const sample = @embedFile("drama3/sample.metal"); pub const codec = @embedFile("drama3/codec.metal"); }; diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index 4d427aece2..1d5aff877f 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -6,6 +6,7 @@ pub const gates = @import("contract/gates.zig"); pub const rows = @import("contract/rows.zig"); pub const affine = @import("contract/affine.zig"); pub const dense = @import("contract/dense.zig"); +pub const bf16_round = @import("bf16_round.zig"); pub const limits = @import("contract/limits.zig"); pub const registry = @import("contract/registry.zig"); pub const cubin = @import("contract/cubin.zig"); diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index d5cb7cbe24..c7e8f24faf 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -508,6 +508,11 @@ pub const metal = [_]M{ .{ .name = "drama3 moe d3_plan", .threads = 1024, .tg_memory = 1024, .reserved = 1024 }, .{ .name = "drama3 moe d3_expert_up", .threads = 128, .tg_memory = 0, .reserved = 128 }, .{ .name = "drama3 moe d3_expert_down", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "drama3 nvfp4 d3_nvfp4_up", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "drama3 nvfp4 d3_nvfp4_down", .threads = 128, .tg_memory = 0, .reserved = 128 }, + .{ .name = "drama3 nvfp4 d3_router_compact", .threads = 256, .tg_memory = 256, .reserved = 256 }, + .{ .name = "drama3 nvfp4 d3_nvfp4_combine", .threads = 256, .tg_memory = 128, .reserved = 256 }, + .{ .name = "drama3 guidance d3_cfg_mix", .threads = 256, .tg_memory = 0, .reserved = 256 }, .{ .name = "drama3 sample d3_argmax", .threads = 1024, .tg_memory = 256, .reserved = 1024 }, .{ .name = "drama3 sample d3_sample_main", .threads = 256, .tg_memory = 3632, .reserved = 256 }, .{ .name = "drama3 sample d3_sample_code", .threads = 256, .tg_memory = 3632, .reserved = 256 }, diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index e1b5ece864..21a962d309 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -264,6 +264,10 @@ pub const Open = struct { context: ?i64 = null, max_frames: ?u32 = null, silence: speech_guard.Config = .{}, + /// Allocate an independent backbone state per speech lane; Metal Drama only. + cfg_enabled: bool = false, + /// Submit the two independent CFG sequences in one Metal command buffer, not vectorized GEMM. + cfg_batched: bool = false, lanes: u32 = 8, /// --parallel named a number: an engine whose memory fits fewer streams refuses instead of serving fewer. lanes_fixed: bool = false, diff --git a/zig/src/core/speech_api.zig b/zig/src/core/speech_api.zig index d593239a45..08ac16a710 100644 --- a/zig/src/core/speech_api.zig +++ b/zig/src/core/speech_api.zig @@ -16,6 +16,9 @@ pub const SpeechRequest = struct { stream: bool = false, /// End of the invariant reference prompt, excluding request text; 0 disables reuse; backend saves no lookahead row. reference_rows: u32 = 0, + /// Plain text frame rows for reference-conditioned semantic CFG; empty means no guidance. + unconditioned: []const [11]i32 = &.{}, + cfg_gamma: f32 = 1, }; pub const SpeechInfo = struct { diff --git a/zig/src/families/drama3/drama3_metal.zig b/zig/src/families/drama3/drama3_metal.zig index 41b2e059e4..444f17d7cd 100644 --- a/zig/src/families/drama3/drama3_metal.zig +++ b/zig/src/families/drama3/drama3_metal.zig @@ -16,6 +16,8 @@ pub const kernels = @import("metal/kernels.zig"); pub const weights = @import("metal/weights.zig"); pub const load_limit = @import("metal/load_limit.zig"); pub const engine = @import("metal/engine.zig"); +pub const guidance = @import("metal/guidance.zig"); +pub const quant_check = @import("metal/quant_check.zig"); pub const codec = @import("metal/codec.zig"); pub const tts = @import("metal/tts.zig"); @@ -39,6 +41,8 @@ test { _ = wav; _ = kernels; _ = load_limit; + _ = guidance; + _ = @import("metal/quant_format_test.zig"); } pub const residency = @import("metal/residency.zig"); diff --git a/zig/src/families/drama3/metal/engine.zig b/zig/src/families/drama3/metal/engine.zig index 54f761c37d..5ce06ebced 100644 --- a/zig/src/families/drama3/metal/engine.zig +++ b/zig/src/families/drama3/metal/engine.zig @@ -108,7 +108,7 @@ pub const Engine = struct { } e.state_start = e.g.bufs.items.len; errdefer e.g.releaseFrom(e.state_start); - try e.allocState(); + try e.allocState(false); e.state_end = e.g.bufs.items.len; if (o.warmup) { var rows: [4][W]i32 = @splat(@splat(0)); @@ -131,13 +131,23 @@ pub const Engine = struct { } pub fn fork(base: *Engine) !*Engine { + return base.forkState(false); + } + + /// CFG's independent backbone state, without codec tower caches, output frames or tower scratch. + pub fn forkBackbone(base: *Engine) !*Engine { + return base.forkState(true); + } + + fn forkState(base: *Engine, backbone_only: bool) !*Engine { const e = try base.gpa.create(Engine); errdefer base.gpa.destroy(e); e.* = base.*; e.owns_model = false; + if (backbone_only) e.o.max_rows = 1; e.state_start = e.g.bufs.items.len; errdefer e.g.releaseFrom(e.state_start); - try e.allocState(); + try e.allocState(backbone_only); e.state_end = e.g.bufs.items.len; return e; } @@ -151,12 +161,12 @@ pub const Engine = struct { return e.g.alloc(bytes); } - fn allocState(e: *Engine) !void { + fn allocState(e: *Engine, backbone_only: bool) !void { const R = e.o.max_rows; const ctxn = e.o.max_ctx; const kvb = ctxn * cfg.kv_heads * cfg.head_dim * 2; for (0..cfg.layers) |li| e.kv_bb[li] = if (e.c.full_attention[li]) .{ try e.alloc(kvb), try e.alloc(kvb) } else .{ 0, 0 }; - for (0..cfg.tower_layers) |ti| e.kv_tw[ti] = .{ try e.alloc(kvb), try e.alloc(kvb) }; + for (0..cfg.tower_layers) |ti| e.kv_tw[ti] = if (backbone_only) .{ 0, 0 } else .{ try e.alloc(kvb), try e.alloc(kvb) }; var n_gdn: usize = 0; for (e.c.full_attention) |f| n_gdn += @intFromBool(!f); const conv_b = cfg.conv_channels * 3 * 2; @@ -176,13 +186,15 @@ pub const Engine = struct { } inline for (.{ "pos", "pos_bb", "step", "count" }) |f| @field(e, f) = try e.alloc(8); e.bb_input = try e.alloc(W * 4); - inline for (.{ "t_xn", "t_tmp", "x", "xn", "tmp", "shout", "hfin", "mixed" }) |f| @field(e, f) = try e.alloc(R * H * 2); - e.t_proj = try e.alloc(R * 2 * cfg.tower_ffn * 2); - inline for (.{ "t_y", "t_q", "y", "q" }) |f| @field(e, f) = try e.alloc(R * 4096 * 2); + inline for (.{ "t_xn", "t_tmp", "mixed" }) |f| @field(e, f) = if (backbone_only) 0 else try e.alloc(R * H * 2); + inline for (.{ "x", "xn", "tmp", "shout", "hfin" }) |f| @field(e, f) = try e.alloc(R * H * 2); + e.t_proj = if (backbone_only) 0 else try e.alloc(R * 2 * cfg.tower_ffn * 2); + inline for (.{ "t_y", "t_q" }) |f| @field(e, f) = if (backbone_only) 0 else try e.alloc(R * 4096 * 2); + inline for (.{ "y", "q" }) |f| @field(e, f) = try e.alloc(R * 4096 * 2); e.window = try e.alloc(16 * 4); e.cur_input = try e.alloc(W * 4); - e.frame = try e.alloc(W * 4); - e.frames_out = try e.alloc(e.o.max_frames * W * 4); + e.frame = if (backbone_only) 0 else try e.alloc(W * 4); + e.frames_out = if (backbone_only) 0 else try e.alloc(e.o.max_frames * W * 4); e.pframes = try e.alloc((ctxn + 1) * W * 4); e.proj = try e.alloc(R * cfg.gdn_proj * 2); e.conv_out = try e.alloc(R * cfg.conv_channels * 2); @@ -199,12 +211,12 @@ pub const Engine = struct { e.gsc = try e.alloc(R * 8 * (cfg.expert_width / 128) * 4); e.counts = try e.alloc(cfg.experts * 4); e.members = try e.alloc(cfg.experts * R * 8 * 4); - e.th = .{ try e.alloc(R * H * 2), try e.alloc(R * H * 2) }; - e.mix = try e.alloc(R * 2 * H * 2); - e.act = try e.alloc(R * cfg.tower_ffn * 2); + e.th = if (backbone_only) .{ 0, 0 } else .{ try e.alloc(R * H * 2), try e.alloc(R * H * 2) }; + e.mix = if (backbone_only) 0 else try e.alloc(R * 2 * H * 2); + e.act = if (backbone_only) 0 else try e.alloc(R * cfg.tower_ffn * 2); e.hlast = try e.alloc(H * 2); e.cand_logits = try e.alloc(e.w.n_cand * 4); - e.code_logits = try e.alloc(cfg.codebook_size * 4); + e.code_logits = if (backbone_only) 0 else try e.alloc(cfg.codebook_size * 4); } pub fn flush(e: *Engine) !void { @@ -260,8 +272,9 @@ pub const Engine = struct { try e.put(e.pframes, std.mem.sliceAsBytes(rows)); } - pub fn promptChunk(e: *Engine, n: usize, c0: usize, t: usize, probe: ?Probe) !void { - if (t == 0 or t > e.o.max_rows or c0 >= n or t > n - c0) return error.BadPromptChunk; + /// Backbone-only prompt chunk for the unconditional CFG sequence; no codec tower work. + pub fn backboneChunk(e: *Engine, c0: usize, t: usize, probe: ?Probe) !void { + if (t == 0 or t > e.o.max_rows) return error.BadPromptChunk; const g = e.g; var nb: [64]u8 = undefined; try e.setPos(c0); @@ -273,6 +286,12 @@ pub const Engine = struct { } k.rmsnorm(g, e.x, H, e.w.final_norm, e.hfin, H, t, H, cfg.eps, true); if (probe) |p| try p.f(p.ctx, "backbone_final_norm", e.hfin, t, H, false); + } + + pub fn promptChunk(e: *Engine, n: usize, c0: usize, t: usize, probe: ?Probe) !void { + if (t == 0 or t > e.o.max_rows or c0 >= n or t > n - c0) return error.BadPromptChunk; + try e.backboneChunk(c0, t, probe); + const g = e.g; // the tower over this chunk's rows that have a next prompt row const tt = @min(c0 + t, n - 1) - c0; if (tt > 0) try e.tower(tt, c0, e.hfin, probe); diff --git a/zig/src/families/drama3/metal/forward.zig b/zig/src/families/drama3/metal/forward.zig index 5f79ccd4b0..72f1f4bbd5 100644 --- a/zig/src/families/drama3/metal/forward.zig +++ b/zig/src/families/drama3/metal/forward.zig @@ -54,24 +54,35 @@ pub fn backboneLayer(e: *Engine, ln: Lane, li: usize, t: usize) void { moe(e, L, t); } -/// The MoE block: router, gated shared expert and the routed FP8 experts, added to the residual stream. +/// The MoE block: original-ID router, gated shared expert and packed FP8/NVFP4 experts, added to the residual. fn moe(e: *Engine, L: *const wts.Layer, t: usize) void { const g = e.g; const mp = cfg.moe_proj; k.linear(g, L.moe_in, e.xn, H, H, mp, e.sp, mp, t, false); - k.router(g, e.sp, mp, 2 * cfg.shared_width, e.picks, e.rw, t); + if (L.expert_count == cfg.experts) + k.router(g, e.sp, mp, 2 * cfg.shared_width, e.picks, e.rw, t) + else + k.routerCompact(g, e.sp, mp, 2 * cfg.shared_width, e.picks, e.rw, L.expert_slots, t); k.swiglu(g, e.sp, mp, 0, cfg.shared_width, cfg.shared_width, e.sact, t); k.linear(g, L.shared_down, e.sact, cfg.shared_width, cfg.shared_width, H, e.shout, H, t, false); const pairs = t * cfg.top_k; const prompt = t > 1; if (prompt) k.plan(g, e.picks, pairs, e.counts, e.members, e.o.max_rows * cfg.top_k); - const r: k.Route = .{ .table = L.experts, .picks = e.picks, .counts = e.counts, .members = e.members, .cap = e.o.max_rows * cfg.top_k, .pairs = pairs, .prompt = prompt }; - // activations to FP8 per row and 128 inputs before each expert projection, as the reference serves them - k.actQuant(g, e.xn, H, e.xq, H, e.xsc, t, H); - k.experts(g, true, e.xq, e.xsc, r, H, cfg.expert_width, e.gu); - k.actQuant(g, e.gu, cfg.expert_width, e.gq, cfg.expert_width, e.gsc, pairs, cfg.expert_width); - k.experts(g, false, e.gq, e.gsc, r, cfg.expert_width, H, e.down); - k.moeCombine(g, e.x, e.down, e.rw, e.shout, e.sp, mp, 2 * cfg.shared_width + cfg.experts, H, t); + const r: k.Route = .{ .table = L.experts, .picks = e.picks, .counts = e.counts, .members = e.members, .cap = e.o.max_rows * cfg.top_k, .pairs = pairs, .prompt = prompt, .experts = L.expert_count }; + if (L.nvfp4) { + k.expertsNvfp4(g, true, e.xn, r, H, cfg.expert_width, e.gu); + k.expertsNvfp4(g, false, e.gu, r, cfg.expert_width, H, e.down); + } else { + // The legacy FP8 path retains its activation quantization and reduction arithmetic. + k.actQuant(g, e.xn, H, e.xq, H, e.xsc, t, H); + k.experts(g, true, e.xq, e.xsc, r, H, cfg.expert_width, e.gu); + k.actQuant(g, e.gu, cfg.expert_width, e.gq, cfg.expert_width, e.gsc, pairs, cfg.expert_width); + k.experts(g, false, e.gq, e.gsc, r, cfg.expert_width, H, e.down); + } + if (L.nvfp4) + k.moeCombineNvfp4(g, e.x, e.down, e.rw, e.shout, e.xn, L.moe_in + (2 * cfg.shared_width + cfg.experts) * H * 2, H, t) + else + k.moeCombine(g, e.x, e.down, e.rw, e.shout, e.sp, mp, 2 * cfg.shared_width + cfg.experts, H, t); } /// Tower layer ti over t rows: hidden hin -> hout; label frames lf (rows lfs ints apart), input frames inf. @@ -96,14 +107,21 @@ fn probeName(buf: []u8, comptime f: []const u8, i: usize) []const u8 { return std.fmt.bufPrint(buf, f, .{i}) catch unreachable; } +pub const Guidance = struct { unconditioned: u64, gamma: f32 }; + /// A frame step's dispatches: sample, the tower, record, then the backbone row at the next position. pub fn enqueueFrame(e: *Engine, smp: Sampling, probe: ?Probe, force: ?Engine.Force) !void { + return enqueueFrameGuided(e, smp, probe, force, null); +} + +pub fn enqueueFrameGuided(e: *Engine, smp: Sampling, probe: ?Probe, force: ?Engine.Force, guidance: ?Guidance) !void { const g = e.g; const sem_begin = e.c.ids.sem_begin; const greedy = smp.greedy(); const sp = k.sampling(smp); if (!greedy and (smp.top_k < 1 or smp.top_k > 64)) return error.TopKUnsupportedOnMetal; k.linear(g, e.w.cand_head, e.hlast, H, H, e.w.n_cand, e.cand_logits, e.w.n_cand, 1, true); + if (guidance) |guide| @import("guidance.zig").mix(e, guide.unconditioned, guide.gamma); if (probe) |p| try p.f(p.ctx, "tower_decode0.candidate_logits", e.cand_logits, 1, e.w.n_cand, true); if (force) |fc| { k.argmax(g, e.cand_logits, e.w.n_cand, e.w.cand_ids, fc.picks, sem_begin, cfg.codebook_size); diff --git a/zig/src/families/drama3/metal/gpu.zig b/zig/src/families/drama3/metal/gpu.zig index 294d845b52..fefffdba06 100644 --- a/zig/src/families/drama3/metal/gpu.zig +++ b/zig/src/families/drama3/metal/gpu.zig @@ -22,6 +22,11 @@ pub const P = enum { d3_plan, d3_expert_up, d3_expert_down, + d3_nvfp4_up, + d3_nvfp4_down, + d3_nvfp4_combine, + d3_router_compact, + d3_cfg_mix, d3_argmax, d3_sample_main, d3_sample_code, diff --git a/zig/src/families/drama3/metal/guidance.zig b/zig/src/families/drama3/metal/guidance.zig new file mode 100644 index 0000000000..516ec30f09 --- /dev/null +++ b/zig/src/families/drama3/metal/guidance.zig @@ -0,0 +1,152 @@ +//! Semantic-only CFG with independent recurrent/KV state and one shared checkpoint weight set. +const std = @import("std"); +const cfg = @import("../config.zig"); +const Engine = @import("engine.zig").Engine; +const k = @import("kernels.zig"); +const fwd = @import("forward.zig"); +const H = cfg.hidden; + +pub fn validGamma(gamma: f32) bool { + return std.math.isFinite(gamma) and gamma >= 1 and gamma <= 5; +} + +pub fn mix(e: *Engine, unconditional: u64, gamma: f32) void { + mixLogits(e.g, e.cand_logits, unconditional, e.w.n_cand, gamma); +} + +fn mixLogits(g: *@import("gpu.zig").Gpu, conditional: u64, unconditional: u64, count: usize, gamma: f32) void { + const Args = extern struct { count: u32, gamma: f32 }; + g.go(.d3_cfg_mix, .{ (count + 255) / 256, 1, 1 }, 256, &.{ conditional, unconditional }, Args{ .count = @intCast(count), .gamma = gamma }); +} + +/// Host lifecycle is separate from GPU storage; no old reference may activate a later plain request. +pub const State = struct { + active: bool = false, + gamma: f32 = 1, + + pub fn begin(s: *State, rows: usize, gamma: f32) !void { + s.clear(); + if (!validGamma(gamma)) return error.BadDramaGuidance; + if (rows == 0) return; + s.* = .{ .active = true, .gamma = gamma }; + } + + pub fn clear(s: *State) void { + s.* = .{}; + } +}; + +pub const Guide = struct { + e: *Engine, + batched: bool, + state: State = .{}, + + pub fn create(base: *Engine, batched: bool) !Guide { + return .{ .e = try base.forkBackbone(), .batched = batched }; + } + + pub fn destroy(g: *Guide) void { + g.e.destroy(); + } + + pub fn begin(g: *Guide, rows: []const [cfg.frame_width]i32, gamma: f32) !void { + try g.state.begin(rows.len, gamma); + errdefer g.state.clear(); + // A guide is never restored from the conditional reference cache. + if (g.state.active) try g.e.startPrompt(rows); + } + + /// One prompt row per scheduler round keeps cancellation and other lanes responsive. + pub fn promptStep(g: *Guide, rows: []const [cfg.frame_width]i32) !bool { + if (!g.state.active or g.e.rows_done == rows.len) return false; + try g.e.backboneChunk(g.e.rows_done, 1, null); + g.e.rows_done += 1; + if (g.e.rows_done == rows.len) { + g.e.g.copy(g.e.hfin, g.e.hlast, H * 2); + try g.e.finishPrompt(rows); + } + return true; + } + + pub fn clear(g: *Guide) void { + g.state.clear(); + } + /// Paired means one command-buffer submission, not a two-row GEMM or simultaneous GPU execution. + pub fn frame(g: *Guide, conditional: *Engine, sampling: @import("../common.zig").Sampling) !void { + if (!g.state.active) return conditional.frameStep(sampling, true, null); + const u = g.e; + k.linear(u.g, u.w.cand_head, u.hlast, H, H, u.w.n_cand, u.cand_logits, u.w.n_cand, 1, true); + if (!g.batched) try u.g.commit(); + try fwd.enqueueFrameGuided(conditional, sampling, null, null, .{ .unconditioned = u.cand_logits, .gamma = g.state.gamma }); + if (!g.batched) try u.g.commit(); + // Feed the actual sampled conditional frame to the unconditional backbone, never its codec tower. + k.nextInput(u.g, conditional.frame, u.bb_input, u.pos, u.pos_bb); + k.embedBackbone(u.g, u.bb_input, u.w.embed, u.w.cb_emb, u.x, 1, H, u.c.ids.sem_begin, u.c.ids.sem_end, cfg.codebook_size); + for (0..cfg.layers) |li| fwd.backboneLayer(u, fwd.laneBB(u, u.pos_bb), li, 1); + k.rmsnorm(u.g, u.x, H, u.w.final_norm, u.hlast, H, 1, H, cfg.eps, true); + k.copyI32(u.g, u.pos_bb, u.pos, 1); + try u.g.commit(); + } +}; + +test "CFG rejects nonfinite and out-of-range gamma before state activation" { + for ([_]f32{ 0, 0.99, 5.01, std.math.nan(f32), std.math.inf(f32) }) |g| + try std.testing.expect(!validGamma(g)); +} + +test "CFG transitions clear old voice on plain request cancellation and invalid begin" { + var state: State = .{}; + try state.begin(12, 1.5); + try std.testing.expect(state.active); + try state.begin(0, 1.5); + try std.testing.expect(!state.active and state.gamma == 1); + try state.begin(8, 2); + state.clear(); + try std.testing.expect(!state.active); + try state.begin(9, 3); + try std.testing.expectError(error.BadDramaGuidance, state.begin(9, std.math.nan(f32))); + try std.testing.expect(!state.active and state.gamma == 1); + try state.begin(7, 1); + try std.testing.expect(state.active and state.gamma == 1); +} + +/// Synthetic device qualification called only by the explicit Metal quant-check executable. +pub fn check(g: *@import("gpu.zig").Gpu) !void { + const first = g.bufs.items.len; + defer g.releaseFrom(first); + const n = cfg.codebook_size + 1; + // The trailing sentinels model codec logits, which must never be touched by semantic CFG. + var conditional: [n + 8]f32 = undefined; + var unconditional: [n + 8]f32 = undefined; + var actual: [n + 8]f32 = undefined; + for (&conditional, &unconditional, 0..) |*c, *u, i| { + c.* = @as(f32, @floatFromInt(@as(i32, @intCast(i % 97)) - 48)) / 8; + u.* = @as(f32, @floatFromInt(@as(i32, @intCast(i % 73)) - 36)) / 16; + } + conditional[0] = -0.0; + const cbuf = try g.alloc(@sizeOf(@TypeOf(conditional))); + const ubuf = try g.alloc(@sizeOf(@TypeOf(unconditional))); + for ([_]f32{ 1, 1.5, 5 }) |gamma| { + // Gamma one must preserve bits even if an unused unconditional value is NaN. + unconditional[0] = if (gamma == 1) std.math.nan(f32) else -2; + try g.put(cbuf, std.mem.asBytes(&conditional)); + try g.put(ubuf, std.mem.asBytes(&unconditional)); + mixLogits(g, cbuf, ubuf, n, gamma); + try g.get(cbuf, std.mem.asBytes(&actual)); + for (actual, conditional, unconditional, 0..) |got, c, u, i| { + if (gamma == 1 or i >= n) { + if (@as(u32, @bitCast(got)) != @as(u32, @bitCast(c))) return error.GuidanceIdentityMismatch; + } else { + const expected = @as(f64, c) + (@as(f64, gamma) - 1) * (@as(f64, c) - @as(f64, u)); + if (!std.math.isFinite(got) or @abs(@as(f64, got) - expected) > 1e-6) + return error.GuidanceFormulaMismatch; + } + } + } + for ([_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32), 0, 5.01 }) |bad| { + var state: State = .{}; + if (state.begin(2, bad)) |_| return error.GuidanceAcceptedInvalidGamma else |err| { + if (err != error.BadDramaGuidance or state.active) return error.GuidanceAcceptedInvalidGamma; + } + } +} diff --git a/zig/src/families/drama3/metal/kernels.zig b/zig/src/families/drama3/metal/kernels.zig index cb6311cb3f..1895731147 100644 --- a/zig/src/families/drama3/metal/kernels.zig +++ b/zig/src/families/drama3/metal/kernels.zig @@ -50,6 +50,10 @@ pub fn moeCombine(g: *Gpu, r: u64, down: u64, w: u64, shared: u64, sp: u64, ss: g.go(.d3_moe_combine, .{ t, 1, 1 }, 256, &.{ r, down, w, shared, sp }, A{ .ss = u(ss), .sg = u(sg_off), .D = u(dim) }); } +pub fn moeCombineNvfp4(g: *Gpu, residual: u64, down: u64, weights: u64, shared: u64, x: u64, gate: u64, dim: usize, rows: usize) void { + g.go(.d3_nvfp4_combine, .{ rows, 1, 1 }, 256, &.{ residual, down, weights, shared, x, gate }, u(dim)); +} + pub const Mix = struct { lf: u64, lfs: usize, code_col: usize, inf: u64, with_shared: bool, table: u64, cb_emb: u64, h: u64, enorm: u64, hnorm: u64, out: u64 }; pub fn towerMix(g: *Gpu, m: Mix, dim: usize, sem_begin: i32, sem_end: i32, cb: i32, eps: f32, t: usize) void { @@ -89,6 +93,11 @@ pub fn router(g: *Gpu, sp: u64, ss: usize, r_off: usize, picks: u64, weights: u6 g.go(.d3_router, .{ t, 1, 1 }, 256, &.{ sp, picks, weights }, A{ .ss = u(ss), .off = u(r_off) }); } +pub fn routerCompact(g: *Gpu, sp: u64, ss: usize, r_off: usize, picks: u64, weights: u64, slots: u64, t: usize) void { + const A = extern struct { ss: u32, off: u32 }; + g.go(.d3_router_compact, .{ t, 1, 1 }, 256, &.{ sp, picks, weights, slots }, A{ .ss = u(ss), .off = u(r_off) }); +} + /// x rows to FP8 values (kept in bf16) and their scales [rows, kdim / 128]. pub fn actQuant(g: *Gpu, x: u64, xs: usize, q: u64, qs: usize, sc: u64, rows: usize, kdim: usize) void { const A = extern struct { xs: u32, qs: u32, groups: u32 }; @@ -100,16 +109,23 @@ pub fn plan(g: *Gpu, picks: u64, pairs: usize, counts: u64, members: u64, cap: u g.go(.d3_plan, .{ 1, 1, 1 }, 1024, &.{ picks, counts, members }, A{ .pairs = u(pairs), .cap = u(cap) }); } -pub const Route = struct { table: u64, picks: u64, counts: u64, members: u64, cap: usize, pairs: usize, prompt: bool }; +pub const Route = struct { table: u64, picks: u64, counts: u64, members: u64, cap: usize, pairs: usize, prompt: bool, experts: usize = 256 }; /// Routed experts: `up` maps token rows (x) to SwiGLU pair rows, else pair rows to fp32 down rows. pub fn experts(g: *Gpu, up: bool, x: u64, xsc: u64, r: Route, kdim: usize, n: usize, out: u64) void { const A = extern struct { K: u32, N: u32, prompt: u32, cap: u32, slots: u32 }; const a: A = .{ .K = u(kdim), .N = u(n), .prompt = @intFromBool(r.prompt), .cap = u(r.cap), .slots = 8 }; - const gy: usize = if (r.prompt) 256 else r.pairs; + const gy: usize = if (r.prompt) r.experts else r.pairs; g.go(if (up) .d3_expert_up else .d3_expert_down, .{ n / @as(usize, if (up) 8 else 16), gy, 1 }, 128, &.{ x, xsc, r.table, r.picks, r.counts, r.members, out }, a); } +/// NVFP4 reads BF16 activation rows directly and writes BF16 SwiGLU / FP32 down; no activation quantization. +pub fn expertsNvfp4(g: *Gpu, up: bool, x: u64, r: Route, kdim: usize, n: usize, out: u64) void { + const A = extern struct { K: u32, N: u32, prompt: u32, cap: u32, slots: u32 }; + const a: A = .{ .K = u(kdim), .N = u(n), .prompt = @intFromBool(r.prompt), .cap = u(r.cap), .slots = 8 }; + g.go(if (up) .d3_nvfp4_up else .d3_nvfp4_down, .{ n / 4, if (r.prompt) r.experts else r.pairs, 1 }, 128, &.{ x, r.table, r.picks, r.counts, r.members, out }, a); +} + /// Greedy: cand_ids non-zero picks a main token (frame[0..2]), zero a code (frame[0]). pub fn argmax(g: *Gpu, logits: u64, count: usize, cand_ids: u64, frame: u64, sem_begin: i32, cb: i32) void { const A = extern struct { count: u32, ids: u32, sb: i32, cb: i32 }; diff --git a/zig/src/families/drama3/metal/load_limit.zig b/zig/src/families/drama3/metal/load_limit.zig index 0e024444e9..957dd1096c 100644 --- a/zig/src/families/drama3/metal/load_limit.zig +++ b/zig/src/families/drama3/metal/load_limit.zig @@ -60,6 +60,11 @@ pub fn cacheBytes(c: cfg.Config, context: u32) u64 { /// The bytes a load needs: the checkpoint's tensors (read in place, the bf16 ones copied too) and the caches. pub fn planBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, context: u32) !u64 { + return planBytesStates(gpa, io, dir, context, 1); +} + +pub fn planBytesStates(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, context: u32, states: u32) !u64 { + if (states == 0) return error.BadStateCount; const c = try cfg.Config.load(gpa, io, dir); var sh = try shards.Shards.open(gpa, io, dir); defer sh.close(io); @@ -67,9 +72,9 @@ pub fn planBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, context: u for (sh.files) |f| for (f.names.values()) |e| { const bytes: u64 = e.end - e.begin; resident += bytes; - if (e.dtype == .bf16) resident += bytes; // copied to the device beside the mapped pages + if (e.dtype == .bf16) resident += bytes; }; - return resident + cacheBytes(c, context); + return resident + cacheBytes(c, context) * states; } /// A plan past the limit is refused before a weight is read. @@ -79,10 +84,15 @@ pub fn admitUnder(plan: u64, limit: Limit) error{OverMemoryLimit}!void { /// This Mac's ceiling and the load's plan: a load that does not fit is refused before a weight is read. pub fn admit(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, context: u32) !void { + return admitStates(gpa, io, dir, context, 1); +} + +/// Independent guidance/serving states share weights, but each owns its caches and row scratch. +pub fn admitStates(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, context: u32, states: u32) !void { const device = try mtl.Device.init(); defer device.deinit(); const limit = try loadLimit(device); - const plan = try planBytes(gpa, io, dir, context); + const plan = try planBytesStates(gpa, io, dir, context, states); if (plan > limit.bytes) { std.log.err("drama3: {d:.1} GB of weights and caches pass this Mac's {d:.1} GB load limit ({s}); free memory, or set DRAMA_LOAD_LIMIT_GB (past 70% of RAM a load risks a watchdog panic)", .{ toGb(plan), toGb(limit.bytes), limit.source() }); return error.OverMemoryLimit; diff --git a/zig/src/families/drama3/metal/quant_check.zig b/zig/src/families/drama3/metal/quant_check.zig new file mode 100644 index 0000000000..36d4f2649f --- /dev/null +++ b/zig/src/families/drama3/metal/quant_check.zig @@ -0,0 +1,232 @@ +//! Deterministic device gate: packed projections, SwiGLU, pruning and mixed FP8/NVFP4 vs independent FP64. +const std = @import("std"); +const Gpu = @import("gpu.zig").Gpu; +const k = @import("kernels.zig"); +const q = @import("quant_format.zig"); +const ref = @import("quant_reference.zig"); + +const H = 2048; +const I = 512; +const E = 8; +const max_rows = 3; +const pairs = max_rows * E; + +const Score = struct { + count: usize = 0, + worst: f64 = 0, + worst_column: usize = 0, + squared: f64 = 0, + fn add(s: *Score, got: f64, want: f64, bound: f64, col: usize) !void { + if (!std.math.isFinite(got) or !std.math.isFinite(want)) return error.NonfiniteQuantResult; + const delta = @abs(got - want); + s.count += 1; + s.squared += delta * delta; + if (delta > s.worst) { + s.worst = delta; + s.worst_column = col; + } + if (delta > bound) { + std.debug.print("FAIL quant column={d} got={d} FP64={d} abs={d} bound={d}\n", .{ col, got, want, delta, bound }); + return error.QuantReferenceMismatch; + } + } + fn report(s: Score, label: []const u8) void { + std.debug.print("SUMMARY {s} PASS count={d} max_abs={d} worst_column={d} rms={d}\n", .{ label, s.count, s.worst, s.worst_column, @sqrt(s.squared / @as(f64, @floatFromInt(s.count))) }); + } +}; + +fn put(g: *Gpu, bytes: []const u8, indirect: bool) !u64 { + const a = try g.alloc(bytes.len); + @memcpy(g.host(a, bytes.len), bytes); + if (indirect) try g.indirect.append(g.gpa, g.ref(a).buf); + return a; +} + +fn bfRows(g: *Gpu, a: u64, n: usize) []align(1) u16 { + return std.mem.bytesAsSlice(u16, @as([]align(1) u8, g.host(a, n * 2))); +} + +const Projection = struct { + w: u64, + s: u64, + global: u64, + n: usize, + dim: usize, + fn init(g: *Gpu, n: usize, dim: usize, seed: usize) !Projection { + const w = try g.alloc(n * dim / 2); + const s = try g.alloc(n * dim / 16); + for (g.host(w, n * dim / 2), 0..) |*b, j| b.* = @truncate((j * 37 + seed * 19) % 256); + for (g.host(s, n * dim / 16), 0..) |*b, j| b.* = @intCast((j + seed * 11) % 127); + const global: f32 = @as(f32, @floatFromInt(seed + 1)) / 1024; + const gp = try put(g, std.mem.asBytes(&global), true); + try g.indirect.append(g.gpa, g.ref(w).buf); + try g.indirect.append(g.gpa, g.ref(s).buf); + return .{ .w = w, .s = s, .global = gp, .n = n, .dim = dim }; + } + fn scalar(p: Projection, g: *Gpu, row: usize, x: []align(1) const u16) f64 { + const global: f32 = @bitCast(std.mem.readInt(u32, g.host(p.global, 4)[0..4], .little)); + return ref.projection(g.host(p.w, p.n * p.dim / 2), g.host(p.s, p.n * p.dim / 16), global, row, p.dim, x); + } +}; + +/// Run only in an exclusive GPU window. The caller opens the newly built drama3 metallib, not a model checkpoint. +pub fn run(g: *Gpu) !void { + try g.flush(); + const first = g.bufs.items.len; + const indirect_first = g.indirect.items.len; + defer { + g.flush() catch {}; + g.indirect.shrinkRetainingCapacity(indirect_first); + g.releaseFrom(first); + } + const gate = try Projection.init(g, I, H, 0); + const up = try Projection.init(g, I, H, 1); + const down = try Projection.init(g, H, I, 2); + var pointers: [E * 9]u64 = undefined; + for (0..E) |e| for ([_]Projection{ gate, up, down }, 0..) |p, j| { + pointers[e * 9 + j * 3] = g.gpuAddress(p.w); + pointers[e * 9 + j * 3 + 1] = g.gpuAddress(p.s); + pointers[e * 9 + j * 3 + 2] = g.gpuAddress(p.global); + }; + const table = try put(g, std.mem.sliceAsBytes(&pointers), false); + const x = try g.alloc(max_rows * H * 2); + for (bfRows(g, x, max_rows * H), 0..) |*b, j| b.* = q.bf16(@as(f32, @floatFromInt(@as(i32, @intCast((j * 13) % 29)) - 14)) / 512); + const logits = try g.alloc(max_rows * 256 * 2); + const original = [_]usize{ 1, 7, 19, 32, 65, 127, 200, 255 }; + var slot_map: [256]i32 = @splat(-1); + for (original, 0..) |id, slot| slot_map[id] = @intCast(slot); + const slots = try put(g, std.mem.sliceAsBytes(&slot_map), false); + // Excluded experts deliberately have the largest logits; retained ones tie, proving original-ID tie order. + for (bfRows(g, logits, max_rows * 256), 0..) |*b, j| b.* = q.bf16(if (slot_map[j % 256] >= 0) -2 else 100); + const picks = try g.alloc(pairs * 4); + const rw = try g.alloc(pairs * 4); + const counts = try g.alloc(256 * 4); + const members = try g.alloc(256 * pairs * 4); + const gu = try g.alloc(pairs * I * 2); + const output = try g.alloc(pairs * H * 4); + var up_score: Score = .{}; + var down_score: Score = .{}; + for ([_]usize{ 1, max_rows }) |rows| { + k.routerCompact(g, logits, 256, 0, picks, rw, slots, rows); + if (rows > 1) k.plan(g, picks, rows * E, counts, members, pairs); + const route: k.Route = .{ .table = table, .picks = picks, .counts = counts, .members = members, .cap = pairs, .pairs = rows * E, .prompt = rows > 1, .experts = E }; + k.expertsNvfp4(g, true, x, route, H, I, gu); + k.expertsNvfp4(g, false, gu, route, I, H, output); + try g.flush(); + for (0..rows * E) |pair| { + const picked = std.mem.readInt(i32, g.host(picks + pair * 4, 4)[0..4], .little); + if (picked != pair % E) return error.PrunedRoutingMismatch; + const weight: f32 = @bitCast(std.mem.readInt(u32, g.host(rw + pair * 4, 4)[0..4], .little)); + if (weight != 0.125) return error.PrunedRoutingMismatch; + const xr = bfRows(g, x + (pair / E) * H * 2, H); + const gr = bfRows(g, gu + pair * I * 2, I); + for (0..I) |col| { + const want = ref.swiglu(gate.scalar(g, col, xr), up.scalar(g, col, xr)); + try up_score.add(ref.fromBf16(gr[col]), want, 0.000002 + @abs(want) / 64, col); + } + for (0..H) |col| { + const want = down.scalar(g, col, gr); + const got: f32 = @bitCast(std.mem.readInt(u32, g.host(output + (pair * H + col) * 4, 4)[0..4], .little)); + try down_score.add(got, want, 0.0002 + @abs(want) * 0.00002, col); + } + } + try combine(g, output, rw, x, rows); + try rawProjection(g, route, gate, x, output, rows, "drama3 NVFP4 gate FP32"); + try rawProjection(g, route, up, x, output, rows, "drama3 NVFP4 up FP32"); + try fp8Layer(g, route, x, gu, output, rows); + } + up_score.report("drama3 NVFP4 gate/up SwiGLU decode+prompt"); + down_score.report("drama3 NVFP4 down FP32 decode+prompt"); +} + +fn combine(g: *Gpu, down: u64, weights: u64, x: u64, rows: usize) !void { + const residual = try g.alloc(rows * H * 2); + const shared = try g.alloc(rows * H * 2); + const wg = try g.alloc(H * 2); + for (bfRows(g, residual, rows * H), 0..) |*v, i| v.* = q.bf16(@as(f32, @floatFromInt(i % 13)) / 32); + for (bfRows(g, shared, rows * H), 0..) |*v, i| v.* = q.bf16(@as(f32, @floatFromInt(@as(i32, @intCast(i % 17)) - 8)) / 16); + for (bfRows(g, wg, H), 0..) |*v, i| v.* = q.bf16(@as(f32, @floatFromInt(@as(i32, @intCast(i % 11)) - 5)) / 64); + k.moeCombineNvfp4(g, residual, down, weights, shared, x, wg, H, rows); + try g.flush(); + var score: Score = .{}; + for (0..rows) |row| { + var dot: f64 = 0; + for (0..H) |j| dot += ref.fromBf16(bfRows(g, x, rows * H)[row * H + j]) * ref.fromBf16(bfRows(g, wg, H)[j]); + const gate = 1 / (1 + @exp(-dot)); + for (0..H) |col| { + var sum: f64 = 0; + for (0..E) |slot| { + const weight: f32 = @bitCast(std.mem.readInt(u32, g.host(weights + (row * E + slot) * 4, 4)[0..4], .little)); + const d: f32 = @bitCast(std.mem.readInt(u32, g.host(down + ((row * E + slot) * H + col) * 4, 4)[0..4], .little)); + sum += ref.roundBf16(@as(f32, weight * d)); + } + const sh = ref.fromBf16(bfRows(g, shared, rows * H)[row * H + col]); + const old = ref.roundBf16(@as(f64, @floatFromInt((row * H + col) % 13)) / 32); + const want = ref.roundBf16(old + ref.roundBf16(gate * sh + ref.roundBf16(sum))); + const got = ref.fromBf16(bfRows(g, residual, rows * H)[row * H + col]); + try score.add(got, want, 0.00001 + @abs(want) / 128, col); + } + } + score.report("drama3 NVFP4 pair rounding shared gate residual"); +} + +fn rawProjection(g: *Gpu, original_route: k.Route, p: Projection, x: u64, output: u64, rows: usize, label: []const u8) !void { + var ptrs: [E * 9]u64 = @splat(0); + for (0..E) |e| { + ptrs[e * 9 + 6] = g.gpuAddress(p.w); + ptrs[e * 9 + 7] = g.gpuAddress(p.s); + ptrs[e * 9 + 8] = g.gpuAddress(p.global); + } + var route = original_route; + route.table = try put(g, std.mem.sliceAsBytes(&ptrs), false); + const repeated = try g.alloc(rows * E * H * 2); + for (0..rows * E) |pair| + @memcpy(g.host(repeated + pair * H * 2, H * 2), g.host(x + (pair / E) * H * 2, H * 2)); + k.expertsNvfp4(g, false, repeated, route, H, I, output); + try g.flush(); + var score: Score = .{}; + for (0..rows * E) |pair| for (0..I) |col| { + const want = p.scalar(g, col, bfRows(g, repeated + pair * H * 2, H)); + const got: f32 = @bitCast(std.mem.readInt(u32, g.host(output + (pair * I + col) * 4, 4)[0..4], .little)); + try score.add(got, want, 0.00002 + @abs(want) * 0.00002, col); + }; + score.report(label); +} + +fn fp8Layer(g: *Gpu, original_route: k.Route, x: u64, gu: u64, output: u64, rows: usize) !void { + const w = try g.alloc(H * I); + for (g.host(w, H * I), 0..) |*b, j| b.* = @intCast(0x20 + j % 32); + const scales = try g.alloc((H / 128) * (I / 128) * 4); + const fs = std.mem.bytesAsSlice(f32, @as([]align(1) u8, g.host(scales, (H / 128) * (I / 128) * 4))); + @memset(fs, 0.0078125); + try g.indirect.append(g.gpa, g.ref(w).buf); + try g.indirect.append(g.gpa, g.ref(scales).buf); + var ptrs: [E * 6]u64 = undefined; + for (0..E * 3) |i| { + ptrs[i * 2] = g.gpuAddress(w); + ptrs[i * 2 + 1] = g.gpuAddress(scales); + } + var route = original_route; + route.table = try put(g, std.mem.sliceAsBytes(&ptrs), false); + const xq = try g.alloc(rows * H * 2); + const xs = try g.alloc(rows * (H / 128) * 4); + const gq = try g.alloc(rows * E * I * 2); + const gs = try g.alloc(rows * E * (I / 128) * 4); + k.actQuant(g, x, H, xq, H, xs, rows, H); + k.experts(g, true, xq, xs, route, H, I, gu); + k.actQuant(g, gu, I, gq, I, gs, rows * E, I); + k.experts(g, false, gq, gs, route, I, H, output); + try g.flush(); + var score: Score = .{}; + for (0..rows * E) |pair| for (0..H) |col| { + var want: f64 = 0; + const act = bfRows(g, gq + pair * I * 2, I); + for (0..I) |j| { + const scale: f32 = @bitCast(std.mem.readInt(u32, g.host(gs + (pair * (I / 128) + j / 128) * 4, 4)[0..4], .little)); + want += ref.fp8(g.host(w, H * I)[col * I + j]) * ref.fromBf16(act[j]) * scale * 0.0078125; + } + const got: f32 = @bitCast(std.mem.readInt(u32, g.host(output + (pair * H + col) * 4, 4)[0..4], .little)); + try score.add(got, want, 0.0001 + @abs(want) * 0.00002, col); + }; + score.report("drama3 mixed-layer legacy FP8 down"); +} diff --git a/zig/src/families/drama3/metal/quant_format.zig b/zig/src/families/drama3/metal/quant_format.zig new file mode 100644 index 0000000000..a3e4959b37 --- /dev/null +++ b/zig/src/families/drama3/metal/quant_format.zig @@ -0,0 +1,73 @@ +//! Packed expert metadata: checkpoint tensor headers, not config quantization labels, select the format. +const std = @import("std"); +const core = @import("core"); +const cfg = @import("../config.zig"); +const Tensor = core.safetensors.Tensor; + +pub const Held = struct { + counts: [cfg.layers]u16 = @splat(cfg.experts), + ids: [cfg.layers][cfg.experts]u8 = @splat(identity), + pub const identity = blk: { + var ids: [cfg.experts]u8 = undefined; + for (&ids, 0..) |*id, i| id.* = @intCast(i); + break :blk ids; + }; + pub fn slots(h: *const Held, layer: usize) [cfg.experts]i32 { + var out: [cfg.experts]i32 = @splat(-1); + for (h.ids[layer][0..h.counts[layer]], 0..) |id, slot| out[id] = @intCast(slot); + return out; + } +}; + +pub fn parseHeld(value: std.json.Value) !Held { + if (value != .object or value.object.count() != cfg.layers) return error.BadExpertMap; + var h: Held = .{}; + for (0..cfg.layers) |li| { + var key: [8]u8 = undefined; + const ids = value.object.get(try std.fmt.bufPrint(&key, "{d}", .{li})) orelse return error.BadExpertMap; + if (ids != .array or ids.array.items.len < cfg.top_k or ids.array.items.len > cfg.experts) return error.BadExpertMap; + for (ids.array.items, 0..) |id, s| { + if (id != .integer or id.integer < 0 or id.integer >= cfg.experts) return error.BadExpertMap; + if (s > 0 and id.integer <= h.ids[li][s - 1]) return error.BadExpertMap; + h.ids[li][s] = @intCast(id.integer); + } + h.counts[li] = @intCast(ids.array.items.len); + } + return h; +} + +pub fn heldMetadata(metadata: ?std.json.Value) !Held { + const meta = metadata orelse return .{}; + if (meta != .object) return error.BadExpertMap; + if (meta.object.get("drama_head_rows") != null) return error.UnsupportedCompactHead; + return if (meta.object.get("drama_experts")) |v| try parseHeld(v) else .{}; +} + +pub fn checkNames(h: *const Held, names: anytype) !void { + var it = names; + const pre = "model.language_model.layers."; + while (it.next()) |key| { + const name = key.*; + if (!std.mem.startsWith(u8, name, pre)) continue; + var parts = std.mem.splitScalar(u8, name[pre.len..], '.'); + const li = parts.next() orelse continue; + if (!std.mem.eql(u8, parts.next() orelse continue, "mlp")) continue; + if (!std.mem.eql(u8, parts.next() orelse continue, "experts")) continue; + const l = std.fmt.parseInt(usize, li, 10) catch return error.BadExpertMap; + const e = std.fmt.parseInt(u8, parts.next() orelse return error.BadExpertMap, 10) catch return error.BadExpertMap; + if (l >= cfg.layers or std.mem.indexOfScalar(u8, h.ids[l][0..h.counts[l]], e) == null) return error.UnlistedExpert; + } +} + +pub fn nvfp4(w: Tensor, s: Tensor, g: Tensor, n: usize, k: usize) !void { + if (k % 16 != 0 or !w.is(.u8, &.{ n, k / 2 }) or !s.is(.f8_e4m3, &.{ n, k / 16 }) or !g.is(.f32, &.{})) return error.BadNvfp4Shape; + if (w.bytes.len != n * k / 2 or s.bytes.len != n * k / 16 or g.bytes.len != 4) return error.BadNvfp4Shape; + for (s.bytes) |b| if (b == 0x7f or b > 0x80) return error.BadNvfp4Scale; + const global: f32 = @bitCast(std.mem.readInt(u32, g.bytes[0..4], .little)); + if (!(global > 0 and std.math.isFinite(global))) return error.BadScale; +} + +pub fn bf16(x: f32) u16 { + const bits: u32 = @bitCast(x); + return @truncate((bits +% (0x7fff + ((bits >> 16) & 1))) >> 16); +} diff --git a/zig/src/families/drama3/metal/quant_format_test.zig b/zig/src/families/drama3/metal/quant_format_test.zig new file mode 100644 index 0000000000..747124a519 --- /dev/null +++ b/zig/src/families/drama3/metal/quant_format_test.zig @@ -0,0 +1,90 @@ +const std = @import("std"); +const st = @import("core").safetensors; +const q = @import("quant_format.zig"); +const ref = @import("quant_reference.zig"); +const cfg = @import("../config.zig"); + +fn tensor(dtype: st.DType, shape: []const usize, bytes: []const u8) st.Tensor { + var t: st.Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +test "NVFP4 rejects shape dtype negative nan scales and invalid globals" { + const codes: [32]u8 = @splat(0x98); + var scales: [4]u8 = @splat(0x38); + var global: f32 = 1; + const w = tensor(.u8, &.{ 1, 32 }, &codes); + const s = tensor(.f8_e4m3, &.{ 1, 4 }, &scales); + const g = tensor(.f32, &.{}, std.mem.asBytes(&global)); + try q.nvfp4(w, s, g, 1, 64); + for ([_]u8{ 0x7f, 0x81, 0xff }) |bad| { + scales[0] = bad; + try std.testing.expectError(error.BadNvfp4Scale, q.nvfp4(w, s, g, 1, 64)); + } + for ([_]u8{ 0, 0x80 }) |zero| { + scales[0] = zero; + try q.nvfp4(w, s, g, 1, 64); + } + for ([_]f32{ 0, -1, std.math.inf(f32), std.math.nan(f32) }) |bad| { + global = bad; + try std.testing.expectError(error.BadScale, q.nvfp4(w, s, g, 1, 64)); + } + global = 1; + try std.testing.expectError(error.BadNvfp4Shape, q.nvfp4(tensor(.f8_e4m3, &.{ 1, 32 }, &codes), s, g, 1, 64)); + try std.testing.expectError(error.BadNvfp4Shape, q.nvfp4(w, tensor(.f8_e4m3, &.{ 4, 1 }, &scales), g, 1, 64)); + try std.testing.expectError(error.BadNvfp4Shape, q.nvfp4(w, s, tensor(.f32, &.{1}, std.mem.asBytes(&global)), 1, 64)); + try std.testing.expectError(error.BadNvfp4Shape, q.nvfp4(tensor(.u8, &.{ 1, 32 }, codes[0..31]), s, g, 1, 64)); + try std.testing.expectError(error.BadNvfp4Shape, q.nvfp4(w, s, tensor(.f32, &.{}, &.{}), 1, 64)); +} + +fn mapText(a: std.mem.Allocator, ids: []const u8, omit: bool) ![]u8 { + var text: std.ArrayList(u8) = .empty; + errdefer text.deinit(a); + try text.append(a, '{'); + for (0..cfg.layers - @as(usize, @intFromBool(omit))) |i| try text.print(a, "{s}\"{d}\":{s}", .{ if (i == 0) "" else ",", i, ids }); + try text.append(a, '}'); + return text.toOwnedSlice(a); +} + +test "expert metadata is complete ascending unique and keeps at least eight original IDs" { + const a = std.testing.allocator; + const good = "[1,3,5,7,9,11,13,255]"; + for ([_][]const u8{ good, "[1,3,5,7,9,11,13,13]", "[3,1,5,7,9,11,13,15]", "[1,3,5,7,9,11,13,256]", "[1,3]", "[1,3,5,7,9,11,13,1.5]" }, 0..) |ids, i| { + const text = try mapText(a, ids, false); + defer a.free(text); + const tree = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer tree.deinit(); + if (i == 0) { + const h = try q.parseHeld(tree.value); + const slots = h.slots(0); + try std.testing.expectEqual(@as(i32, -1), slots[0]); + try std.testing.expectEqual(@as(i32, 7), slots[255]); + var names = std.StringHashMapUnmanaged(void).empty; + defer names.deinit(a); + try names.put(a, "model.language_model.layers.0.mlp.experts.0.gate_proj.weight", {}); + try std.testing.expectError(error.UnlistedExpert, q.checkNames(&h, names.keyIterator())); + } else try std.testing.expectError(error.BadExpertMap, q.parseHeld(tree.value)); + } + const text = try mapText(a, good, true); + defer a.free(text); + const tree = try std.json.parseFromSlice(std.json.Value, a, text, .{}); + defer tree.deinit(); + try std.testing.expectError(error.BadExpertMap, q.parseHeld(tree.value)); +} + +test "independent scalar reference covers E2M1 signs and E4M3 boundary scales" { + const values = [_]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + for (values, 0..) |v, i| { + try std.testing.expectEqual(v, ref.fp4(@intCast(i))); + try std.testing.expectEqual(-v, ref.fp4(@intCast(i + 8))); + } + try std.testing.expectEqual(@as(f64, 1.0 / 512.0), ref.fp8(1)); + try std.testing.expectEqual(@as(f64, 7.0 / 512.0), ref.fp8(7)); + try std.testing.expectEqual(@as(f64, 1.0 / 64.0), ref.fp8(8)); + try std.testing.expectEqual(@as(f64, 1), ref.fp8(0x38)); + try std.testing.expectEqual(@as(f64, 448), ref.fp8(0x7e)); + for (0..127) |code| try std.testing.expectEqual(-ref.fp8(@intCast(code)), ref.fp8(@intCast(code + 128))); + try std.testing.expectEqual(@as(f64, 1), ref.roundBf16(1.00390625)); + try std.testing.expectEqual(@as(f64, 1.015625), ref.roundBf16(1.01171875)); +} diff --git a/zig/src/families/drama3/metal/quant_reference.zig b/zig/src/families/drama3/metal/quant_reference.zig new file mode 100644 index 0000000000..51a7d72078 --- /dev/null +++ b/zig/src/families/drama3/metal/quant_reference.zig @@ -0,0 +1,63 @@ +//! Independent FP64 scalar references for the quantized Metal operator gate, never used for inference. +const std = @import("std"); + +pub fn fp8(code: u8) f64 { + const negative = code >= 128; + const bits = code % 128; + const exponent = bits / 8; + const fraction = bits % 8; + if (bits == 127) return std.math.nan(f64); + var value: f64 = @floatFromInt(if (exponent == 0) fraction else 8 + @as(u16, fraction)); + var shift: i32 = if (exponent == 0) -9 else @as(i32, exponent) - 10; + while (shift < 0) : (shift += 1) value /= 2; + while (shift > 0) : (shift -= 1) value *= 2; + return if (negative) -value else value; +} + +pub fn fp4(code: u8) f64 { + const bits = code % 8; + const value: f64 = switch (bits) { + 0 => 0, + 1 => 0.5, + 2 => 1, + 3 => 1.5, + 4 => 2, + 5 => 3, + 6 => 4, + 7 => 6, + else => unreachable, + }; + return if (code >= 8) -value else value; +} + +pub fn fromBf16(bits: u16) f64 { + return @as(f32, @bitCast(@as(u32, bits) << 16)); +} + +/// Nearest BF16 by neighboring representable values, independent of the production integer rounding trick. +pub fn roundBf16(x: f64) f64 { + const f: f32 = @floatCast(x); + const bits: u32 = @bitCast(f); + const low: u16 = @truncate(bits >> 16); + const hi = low +% 1; + const a = fromBf16(low); + const b = fromBf16(hi); + const da = @abs(@as(f64, f) - a); + const db = @abs(@as(f64, f) - b); + return if (da < db or (da == db and low % 2 == 0)) a else b; +} + +pub fn projection(codes: []const u8, scales: []const u8, global: f32, n: usize, k: usize, x: []align(1) const u16) f64 { + var sum: f64 = 0; + for (0..k) |i| { + const byte = codes[n * (k / 2) + i / 2]; + const code = if (i % 2 == 0) byte % 16 else byte / 16; + sum += fp4(code) * fp8(scales[n * (k / 16) + i / 16]) * fromBf16(x[i]); + } + return sum * global; +} + +pub fn swiglu(gate: f64, up: f64) f64 { + const g = roundBf16(gate); + return roundBf16(g / (1 + @exp(-g)) * roundBf16(up)); +} diff --git a/zig/src/families/drama3/metal/weights.zig b/zig/src/families/drama3/metal/weights.zig index d5da9d4d33..36f543a1b3 100644 --- a/zig/src/families/drama3/metal/weights.zig +++ b/zig/src/families/drama3/metal/weights.zig @@ -1,10 +1,11 @@ -//! Drama 3's checkpoint on Metal: bf16 projections copied and stacked, FP8 experts read in place from the shards. +//! Drama 3 Metal weights: BF16 dense projections and packed FP8/NVFP4 experts with original-ID compaction. const std = @import("std"); const mtl = @import("metal"); const core = @import("core"); const cfg = @import("../config.zig"); const shards = @import("../shards.zig"); const gpu = @import("gpu.zig"); +const quant = @import("quant_format.zig"); const Gpu = gpu.Gpu; const Shards = shards.Shards; @@ -20,7 +21,10 @@ pub const Layer = struct { attn: Attn = undefined, moe_in: u64, // [1288, 2048] shared_down: u64, // [2048, 512] - experts: u64, // [256] ExpertPtrs: gate, gate scales, up, up scales, down, down scales (GPU addresses) + experts: u64, // Compact slot table: 6 FP8 pointers or 9 NVFP4 pointers per held expert. + expert_slots: u64, + expert_count: usize, + nvfp4: bool, }; pub const TowerLayer = struct { @@ -118,30 +122,57 @@ fn mapShards(g: *Gpu, sh: *const Shards) ![]u64 { return out; } -/// The 256 experts of layer li: FP8 [512, 2048] gate / up, [2048, 512] down and their fp32 block scales, checked. -fn expertTable(l: *Loader, li: usize, maps: []const u64) !u64 { - const table = try l.g.alloc(cfg.experts * 6 * 8); - const ptrs = std.mem.bytesAsSlice(u64, l.g.host(table, cfg.experts * 6 * 8)); +/// Compact slots preserve ascending original IDs; each layer has one expert format, selected from tensor headers. +fn expertTable(l: *Loader, li: usize, maps: []const u64, held: *const quant.Held) !void { + const layer = &l.w.layers[li]; + const ids = held.ids[li][0..held.counts[li]]; + layer.expert_count = ids.len; + layer.expert_slots = 0; + if (ids.len != cfg.experts) { + const slots = held.slots(li); + layer.expert_slots = try l.g.alloc(@sizeOf(@TypeOf(slots))); + @memcpy(l.g.host(layer.expert_slots, @sizeOf(@TypeOf(slots))), std.mem.asBytes(&slots)); + l.w.bytes += @sizeOf(@TypeOf(slots)); + } const n = cfg.expert_width; const k = cfg.hidden; const parts = [_]struct { []const u8, [2]usize }{ .{ "gate_proj", .{ n, k } }, .{ "up_proj", .{ n, k } }, .{ "down_proj", .{ k, n } } }; - var nb: [160]u8 = undefined; - for (0..cfg.experts) |e| { + var nb: [3][192]u8 = undefined; + const first = try l.sh.get(fmt(&nb[0], "model.language_model.layers.{d}.mlp.experts.{d}.gate_proj.weight", .{ li, ids[0] })); + layer.nvfp4 = first.dtype == .u8; + if (!layer.nvfp4 and first.dtype != .f8_e4m3) return error.UnsupportedExpertDtype; + const stride: usize = if (layer.nvfp4) 9 else 6; + layer.experts = try l.g.alloc(ids.len * stride * 8); + const ptrs = std.mem.bytesAsSlice(u64, l.g.host(layer.experts, ids.len * stride * 8)); + for (ids, 0..) |e, slot| { for (parts, 0..) |p, j| { - const wname = fmt(&nb, "model.language_model.layers.{d}.mlp.experts.{d}.{s}.weight", .{ li, e, p[0] }); - try Loader.expect(try l.sh.get(wname), .f8_e4m3, &p[1], wname); + const wname = fmt(&nb[0], "model.language_model.layers.{d}.mlp.experts.{d}.{s}.weight", .{ li, e, p[0] }); + const sname = fmt(&nb[1], "{s}{s}", .{ wname, if (layer.nvfp4) "_scale" else "_scale_inv" }); + const wt = try l.sh.get(wname); + const sc = try l.sh.get(sname); const wl = try l.sh.locate(wname); - if (wl.offset % 8 != 0) return error.MisalignedExpert; - ptrs[e * 6 + 2 * j] = l.g.gpuAddress(maps[wl.file] + wl.offset); - const sname = fmt(&nb, "model.language_model.layers.{d}.mlp.experts.{d}.{s}.weight_scale_inv", .{ li, e, p[0] }); - try Loader.expect(try l.sh.get(sname), .f32, &.{ p[1][0] / 128, p[1][1] / 128 }, sname); const sl = try l.sh.locate(sname); - if (sl.offset % 4 != 0) return error.MisalignedScale; - ptrs[e * 6 + 2 * j + 1] = l.g.gpuAddress(maps[sl.file] + sl.offset); - l.w.mapped += p[1][0] * p[1][1] + p[1][0] / 128 * p[1][1] / 128 * 4; + const base = slot * stride + j * (stride / 3); + if (layer.nvfp4) { + const gname = fmt(&nb[2], "{s}_scale_2", .{wname}); + const gl = try l.sh.get(gname); + try quant.nvfp4(wt, sc, gl, p[1][0], p[1][1]); + const loc = try l.sh.locate(gname); + if (loc.offset % 4 != 0) return error.MisalignedScale; + ptrs[base + 2] = l.g.gpuAddress(maps[loc.file] + loc.offset); + l.w.mapped += gl.bytes.len; + } else { + try Loader.expect(wt, .f8_e4m3, &p[1], wname); + try Loader.expect(sc, .f32, &.{ p[1][0] / 128, p[1][1] / 128 }, sname); + if (wl.offset % 8 != 0) return error.MisalignedExpert; + if (sl.offset % 4 != 0) return error.MisalignedScale; + } + ptrs[base] = l.g.gpuAddress(maps[wl.file] + wl.offset); + ptrs[base + 1] = l.g.gpuAddress(maps[sl.file] + sl.offset); + l.w.mapped += wt.bytes.len + sc.bytes.len; } } - return table; + l.w.bytes += ids.len * stride * 8; } pub fn load(g: *Gpu, io: std.Io, dir: []const u8, c: cfg.Config) !*Weights { @@ -152,8 +183,14 @@ pub fn load(g: *Gpu, io: std.Io, dir: []const u8, c: cfg.Config) !*Weights { w.mapped = 0; w.sh = try Shards.open(g.gpa, io, dir); errdefer w.sh.close(io); + const held = try quant.heldMetadata(w.sh.metadata); + try quant.checkNames(&held, w.sh.map.keyIterator()); const first = g.bufs.items.len; - errdefer g.releaseFrom(first); + const indirect_first = g.indirect.items.len; + errdefer { + g.indirect.shrinkRetainingCapacity(indirect_first); + g.releaseFrom(first); + } const maps = try mapShards(g, &w.sh); defer g.gpa.free(maps); var l: Loader = .{ .g = g, .sh = &w.sh, .w = w }; @@ -201,7 +238,7 @@ pub fn load(g: *Gpu, io: std.Io, dir: []const u8, c: cfg.Config) !*Weights { } L.moe_in = try l.stack(&.{ fmt(&nb[0], pre ++ "mlp.shared_expert.gate_proj.weight", .{li}), fmt(&nb[1], pre ++ "mlp.shared_expert.up_proj.weight", .{li}), fmt(&nb[2], pre ++ "mlp.gate.weight", .{li}), fmt(&nb[3], pre ++ "mlp.shared_expert_gate.weight", .{li}) }, &.{ cfg.shared_width, cfg.shared_width, cfg.experts, 1 }, H, cfg.moe_proj); L.shared_down = try l.one(fmt(&nb[0], pre ++ "mlp.shared_expert.down_proj.weight", .{li}), &.{ H, cfg.shared_width }); - L.experts = try expertTable(&l, li, maps); + try expertTable(&l, li, maps, &held); } for (0..cfg.tower_layers) |ti| { const shared = ti < cfg.shared_layers; @@ -228,6 +265,6 @@ pub fn load(g: *Gpu, io: std.Io, dir: []const u8, c: cfg.Config) !*Weights { T.final_norm = try l.one(tn(&nb[0], p, idx, "final_layernorm.weight"), &.{H}); } for (0..cfg.codebooks - 1) |hi| w.heads[hi] = try l.one(fmt(&nb[0], "audio_decoder_mtp.heads.{d}.weight", .{hi}), &.{ cfg.codebook_size, H }); - std.log.info("drama3: {d:.2} GB of bf16 weights copied, {d:.2} GB of FP8 experts read in place", .{ @as(f64, @floatFromInt(w.bytes)) / 1e9, @as(f64, @floatFromInt(w.mapped)) / 1e9 }); + std.log.info("drama3: {d:.2} GB of BF16/copied weights, {d:.2} GB of packed FP8/NVFP4 experts read in place", .{ @as(f64, @floatFromInt(w.bytes)) / 1e9, @as(f64, @floatFromInt(w.mapped)) / 1e9 }); return w; } diff --git a/zig/src/families/drama3/shards.zig b/zig/src/families/drama3/shards.zig index 85b63eff99..6e3475f723 100644 --- a/zig/src/families/drama3/shards.zig +++ b/zig/src/families/drama3/shards.zig @@ -10,6 +10,7 @@ pub const Shards = struct { files: []st.File, map: std.StringHashMapUnmanaged(usize), arena: std.heap.ArenaAllocator, + metadata: ?std.json.Value = null, pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Shards { var s: Shards = .{ .files = &.{}, .map = .empty, .arena = std.heap.ArenaAllocator.init(gpa) }; @@ -19,6 +20,7 @@ pub const Shards = struct { const text = try std.Io.Dir.cwd().readFileAlloc(io, ipath, a, .limited(64 << 20)); const tree = try std.json.parseFromSliceLeaky(std.json.Value, a, text, .{}); const wm = try st.index.weightMap(tree); + s.metadata = tree.object.get("metadata"); var names: std.ArrayList([]const u8) = .empty; var it = wm.iterator(); while (it.next()) |e| { diff --git a/zig/src/native/drama_host.zig b/zig/src/native/drama_host.zig index df77baaddb..fde1bfc020 100644 --- a/zig/src/native/drama_host.zig +++ b/zig/src/native/drama_host.zig @@ -19,9 +19,11 @@ const Session = struct { /// Valid only while this lane's prefix KV remains untouched; snapshots own no KV copies. prefix: reference.Prefix = .{}, cache: Snap = .{}, + guide: ?Guide = null, }; /// Metal snapshot buffers belong to the shared GPU owner; CUDA keeps its uncached path. const Snap = if (is_cuda) struct {} else d3.Snapshot; +const Guide = if (is_cuda) struct {} else d3.guidance.Guide; const Job = struct { id: api.Id, request: *const api.SpeechRequest, @@ -42,6 +44,8 @@ pub const Options = struct { silence: api.speech_guard.Config = .{}, max_frames: ?u32 = null, prefill_rows: u32 = @intCast(d3.common.single_row), // rows a prompt chunk: one row until a caller opts in + cfg_enabled: bool = false, + cfg_batched: bool = false, }; const Host = struct { @@ -91,6 +95,16 @@ const Host = struct { if (!std.math.isFinite(r.temperature) or r.temperature < 0) return false; if (!std.math.isFinite(r.top_p) or r.top_p <= 0 or r.top_p > 1) return false; if (r.top_k == 0 or r.top_k > (if (is_cuda) @as(u32, std.math.maxInt(i32)) else 64)) return false; + if (!std.math.isFinite(r.cfg_gamma) or r.cfg_gamma < 1 or r.cfg_gamma > 5) return false; + if (r.unconditioned.len > 0) { + if (is_cuda or h.sessions[0].guide == null or r.unconditioned.len > sp.max_prompt_tokens) return false; + const ids_ = h.sessions[0].e.c.ids; + d3.model.validate(r.unconditioned, ids_, h.info_.context_window, r.max_frames) catch return false; + for (r.unconditioned) |row| { + if (row[0] >= ids_.sem_begin and row[0] <= ids_.sem_end) return false; + for (row[1..]) |code| if (code != 0) return false; + } + } d3.model.validate(r.frames, h.sessions[0].e.c.ids, h.info_.context_window, r.max_frames) catch return false; return true; } @@ -200,6 +214,7 @@ const Host = struct { } else { const ref_rows = reference.prefillRows(r.frames.len, r.reference_rows, s.e.o.max_rows); if (!j.state.started) { + if (lane.guide) |*guide| try guide.begin(r.unconditioned, r.cfg_gamma); // Invalidate before reset, GPU work or any fallible preparation for a miss. if (try lane.prefix.begin(h.gpa, r.frames, ref_rows)) { try s.e.restore(&lane.cache); @@ -213,6 +228,10 @@ const Host = struct { try s.t.c.reset(); j.state.started = true; } + if (lane.guide) |*guide| if (try guide.promptStep(r.unconditioned)) { + try sync(s); + return; + }; var count = s.e.o.promptRows(r.frames.len - j.state.prompt); const want_snapshot = ref_rows > 0 and !lane.prefix.valid; if (want_snapshot and j.state.prompt < ref_rows) count = @min(count, ref_rows - j.state.prompt); @@ -235,7 +254,10 @@ const Host = struct { return; } const smp: d3.Sampling = .{ .temperature = r.temperature, .top_k = @intCast(r.top_k), .top_p = r.top_p, .seed = r.seed }; - try s.e.frameStep(smp, true, null); + if (!is_cuda and lane.guide != null) + try lane.guide.?.frame(s.e, smp) + else + try s.e.frameStep(smp, true, null); var frame: d3.model.Frame = undefined; try s.e.download(s.e.frames_out + j.state.frames * 11 * 4, std.mem.asBytes(&frame)); try h.produced(i, s.t, frame[0] == s.e.c.ids.eos); @@ -269,6 +291,10 @@ const Host = struct { const completed = state.settle(&.{.{ .ctx = &h.sessions[lane], .call = waitSession }}); const final_reason: api.Reason = if (completed) reason else .failed; const final_message = if (completed) message else "speech drain failed: engine closed; storage retained"; + if (!is_cuda) { + if (h.sessions[lane].guide) |*guide| guide.clear(); + if (final_reason == .cancelled or final_reason == .failed) h.sessions[lane].prefix.invalidate(); + } h.lock(); if (!completed) h.failed.closed = true; const j = h.jobs[lane].?; @@ -313,6 +339,7 @@ const Host = struct { while (i > 0) { i -= 1; h.sessions[i].t.destroy(); + if (!is_cuda) if (h.sessions[i].guide) |*guide| guide.destroy(); h.sessions[i].e.destroy(); } h.limit.deinit(); @@ -322,9 +349,11 @@ const Host = struct { pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, _: []const u8, o: Options) !api.Opened { if (o.context < 1025 or o.context > 8192) return error.BadDramaContext; + if (is_cuda and (o.cfg_enabled or o.cfg_batched)) return error.GuidanceRequiresMetal; + if (o.cfg_batched and !o.cfg_enabled) return error.BadDramaGuidance; const budget = try @import("drama_budget.zig").budget(o.context, o.max_frames); try (d3.Options{ .max_rows = @intCast(o.prefill_rows), .attn_chunk = @intCast(o.prefill_rows) }).validatePrefill(); - if (!is_cuda) try d3.load_limit.admit(gpa, io, dir, o.context); // refused before a weight is read: the floor's one-Mac limit + if (!is_cuda) try d3.load_limit.admitStates(gpa, io, dir, o.context, o.lanes * (1 + @as(u32, @intFromBool(o.cfg_enabled)))); try o.silence.validate(); const h = try gpa.create(Host); errdefer gpa.destroy(h); @@ -343,6 +372,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, _: []const u8, made -= 1; if (!is_cuda) h.sessions[made].prefix.deinit(gpa); h.sessions[made].t.destroy(); + if (!is_cuda) if (h.sessions[made].guide) |*guide| guide.destroy(); h.sessions[made].e.destroy(); }; const codec = try std.fs.path.join(gpa, &.{ dir, "codec.pth" }); @@ -359,6 +389,10 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, _: []const u8, errdefer e.destroy(); const t = if (made == 0) try d3.tts.Tts.create(gpa, io, e, codec, 1) else try h.sessions[0].t.fork(e); h.sessions[made] = .{ .e = e, .t = t }; + if (!is_cuda and o.cfg_enabled) h.sessions[made].guide = d3.guidance.Guide.create(e, o.cfg_batched) catch |err| { + t.destroy(); + return err; + }; } const frame_cap = budget.frames; h.info_ = .{ @@ -374,6 +408,9 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, _: []const u8, errdefer if (!is_cuda) h.warm.drop(h.sessions[0].e.g); h.worker = try std.Thread.spawn(.{}, Host.run, .{h}); if (o.silence.enabled()) std.log.info("drama3: silence guard on: {f}", .{o.silence}); + if (o.cfg_enabled) std.log.info("drama3: semantic CFG enabled; shared weights, independent states; {s}", .{ + if (o.cfg_batched) "paired command-buffer submission (not vectorized GEMM)" else "separate sequence submissions", + }); return .{ .engine = .{ .ctx = h, .vtable = &Host.vtable }, .ctx = h, .close = Host.close }; } diff --git a/zig/src/native/drama_open.zig b/zig/src/native/drama_open.zig index 441302d82f..8432a9590c 100644 --- a/zig/src/native/drama_open.zig +++ b/zig/src/native/drama_open.zig @@ -6,6 +6,10 @@ const setting = @import("cuda_device.zig").setting; const getenv = @import("cuda_device.zig").getenv; pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { + if (o.cfg_enabled or o.cfg_batched) { + problem.* = "this engine revision implements Drama CFG on Metal only; use the pinned CUDA quant engine on Linux"; + return null; + } if (o.tp != 1 or o.rank != 0 or (o.lanes_fixed and (o.lanes == 0 or o.lanes > 8))) { problem.* = "Drama 3 serves 1 through 8 frame lanes on one CUDA device; tensor parallelism is unsupported"; return null; @@ -77,6 +81,8 @@ pub fn openMetal(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: .lanes = if (o.lanes_fixed) o.lanes else 2, .silence = o.silence, .max_frames = o.max_frames, + .cfg_enabled = o.cfg_enabled, + .cfg_batched = o.cfg_batched, }) catch |err| { problem.* = if (err == error.OverMemoryLimit) try std.fmt.allocPrint(a, "Drama 3 needs more memory than this Mac's load limit allows (70% of RAM, or DRAMA_LOAD_LIMIT_GB); the log names the sizes", .{}) diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index d7fbc6eff0..b6903bea84 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -156,6 +156,10 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c return null; } if (std.mem.eql(u8, o.model_type, "fish_omni")) return @import("drama_open.zig").openMetal(a, gpa, io, o, problem); + if (o.cfg_enabled or o.cfg_batched) { + problem.* = "--cfg-* guidance requires a Drama 3 speech model"; + return null; + } if (@import("save_admission.zig").refused(io, o.dir, problem)) return null; const dense = std.mem.eql(u8, o.model_type, "qwen3_5") and qwen27Engine(a, io, o.dir); if (try refusal(a, o, dense)) |why| { diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 5b77bc1e2f..94b9ed379f 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -49,6 +49,7 @@ pub const flags = [_]Flag{ .{ .name = "--chat-template", .native = true }, .{ .name = "--voices", .native = true }, .{ .name = "--default-voice", .native = true }, +} ++ @import("drama_cfg.zig").flags ++ [_]Flag{ .{ .name = "--alias", .kind = .append, .native = true }, .{ .name = "--api-key", .kind = .append, .native = true }, .{ .name = "--api-key-file", .native = true }, @@ -173,6 +174,7 @@ pub const Args = struct { context: ?i64 = null, max_frames: ?u32 = null, silence: api.speech_guard.Config = .{}, + cfg: @import("drama_cfg.zig").Flags = .{}, redraw: @import("speech_redraw.zig").Policy = .{}, speed_up: ?[]const u8 = null, // speed-up mode: this Mac's settings for the two-Mac link (Flash Next and the GLMs) prompt_cache_gib: ?f64 = null, // kept prompt states' budget (0: none; null: what 70% of RAM leaves past the loaded server) @@ -282,6 +284,7 @@ pub fn parse(a: Allocator, argv: []const []const u8, u: *Usage) error{ Usage, Ou try apply(a, &out, name, value, u, &repeated); } out.model = model orelse return fail(u, a, "the following arguments are required: model", .{}); + if ((out.cfg.batched or out.cfg.force_guided) and out.cfg.options() == null) return fail(u, a, "--cfg-batched and --cfg-force-guided need --cfg-gamma or --cfg-allow-request-gamma", .{}); try ranks(u, a, out); out.silence.validate() catch |err| return fail(u, a, "invalid silence guard configuration: {s}", .{@errorName(err)}); @@ -348,6 +351,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa } if (try cudaFlag(a, out, name, v, u) or try gpuFlag(a, out, name, v, u)) return; if (try recallFlag(a, out, name, v, u)) return; + if (out.cfg.set(name, v) catch return fail(u, a, "--cfg-gamma must be 0 (off) or 1 (unguided) through 5", .{})) return; if (is(name, "--voices")) { out.voices = v; return; @@ -503,7 +507,7 @@ pub fn parallelFixed(text: []const u8) bool { /// What a server asks of its engine, from its flags. pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args) !api.Open { - return .{ + return @import("drama_cfg.zig").withGuidance(args.cfg, .{ .dir = dir, .model_type = model_type, .context = args.context, @@ -539,7 +543,7 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .master_port = args.master_port, .policy = args.policy, .vision = if (args.vision) .{ .image_tokens = args.vision_image_tokens, .image_tokens_set = args.vision_image_tokens_set, .max_images = args.vision_max_images, .max_videos = args.vision_max_videos, .video_tokens = args.vision_video_tokens } else null, - }; + }); } /// What ``capabilities --json`` reports about the engine side: its version, chip, backends and families. diff --git a/zig/src/server/drama_cfg.zig b/zig/src/server/drama_cfg.zig new file mode 100644 index 0000000000..0c98506fd1 --- /dev/null +++ b/zig/src/server/drama_cfg.zig @@ -0,0 +1,121 @@ +//! Reference-conditioned semantic CFG; built-in voices remain unguided and disabled mode ignores the header. +const std = @import("std"); +const fields = @import("fields.zig"); +const model_text = @import("model_text.zig"); +const drama_prompt = @import("drama_prompt.zig"); + +pub const flags = [_]@import("cli.zig").Flag{ + .{ .name = "--cfg-gamma", .native = @import("builtin").os.tag == .macos }, + .{ .name = "--cfg-allow-request-gamma", .kind = .store_true, .native = @import("builtin").os.tag == .macos }, + .{ .name = "--cfg-force-guided", .kind = .store_true, .native = @import("builtin").os.tag == .macos }, + .{ .name = "--cfg-batched", .kind = .store_true, .native = @import("builtin").os.tag == .macos }, +}; + +pub const Options = struct { gamma: f32 = 1, request_gamma: bool = false, force_guided: bool = false }; +pub const Flags = struct { + gamma: f32 = 0, + request_gamma: bool = false, + force_guided: bool = false, + batched: bool = false, + + pub fn options(f: Flags) ?Options { + if (f.gamma == 0 and !f.request_gamma) return null; + return .{ .gamma = if (f.gamma == 0) 1 else f.gamma, .request_gamma = f.request_gamma, .force_guided = f.force_guided }; + } + + pub fn set(f: *Flags, name: []const u8, text: []const u8) error{BadGamma}!bool { + const is = std.mem.eql; + if (is(u8, name, "--cfg-gamma")) { + const g = fields.pyFloat(text) orelse return error.BadGamma; + if (!(g == 0 or valid(g))) return error.BadGamma; + f.gamma = @floatCast(g); + } else if (is(u8, name, "--cfg-allow-request-gamma")) { + f.request_gamma = true; + } else if (is(u8, name, "--cfg-force-guided")) { + f.force_guided = true; + } else if (is(u8, name, "--cfg-batched")) { + f.batched = true; + } else return false; + return true; + } +}; + +pub fn withGuidance(f: Flags, request: @import("engine_api").Open) !@import("engine_api").Open { + if (f.options() != null and !std.mem.eql(u8, request.model_type, "fish_omni")) return error.GuidanceRequiresDrama; + var out = request; + out.cfg_enabled = f.options() != null; + out.cfg_batched = f.batched; + return out; +} + +pub fn valid(g: f64) bool { + return g >= 1 and g <= 5; +} + +pub fn gamma(options: Options, sent: ?[]const u8) ?f32 { + if (!options.request_gamma) return options.gamma; + const g = fields.pyFloat(sent orelse return options.gamma) orelse return null; + return if (valid(g)) @floatCast(g) else null; +} + +pub fn active(options: ?Options, g: f32, referenced: bool) bool { + const o = options orelse return false; + return referenced and (g != 1 or o.force_guided); +} + +pub const Guide = struct { rows: []const [11]i32 = &.{}, gamma: f32 = 1 }; + +pub fn guide(a: std.mem.Allocator, text: model_text.Text, options: ?Options, g: f32, input: []const u8, referenced: bool) model_text.Error!Guide { + if (!active(options, g, referenced)) return .{}; + const tokens = try drama_prompt.encode(a, text, input); + const rows = drama_prompt.rows(a, tokens) catch |err| return switch (err) { + error.OutOfMemory => error.OutOfMemory, + error.InvalidToken => error.Template, + }; + return .{ .rows = rows, .gamma = g }; +} + +test "CFG header validation and fixed policy match the CUDA recipe" { + const fixed: Options = .{ .gamma = 1.5 }; + try std.testing.expectEqual(@as(?f32, 1.5), gamma(fixed, "nonsense")); + const open: Options = .{ .gamma = 1.5, .request_gamma = true }; + try std.testing.expectEqual(@as(?f32, 1.5), gamma(open, null)); + for ([_][]const u8{ "1", " 1.5 ", "3", "5" }, [_]f32{ 1, 1.5, 3, 5 }) |sent, want| + try std.testing.expectEqual(@as(?f32, want), gamma(open, sent)); + for ([_][]const u8{ "0", "0.99", "5.01", "nan", "inf", "", "2x" }) |bad| + try std.testing.expectEqual(@as(?f32, null), gamma(open, bad)); +} + +test "CFG activation requires reference and configured policy including forced gamma one" { + try std.testing.expect((Flags{}).options() == null); + const options = (Flags{ .gamma = 1.5 }).options(); + try std.testing.expect(active(options, 1.5, true)); + try std.testing.expect(!active(options, 1.5, false)); + try std.testing.expect(!active(options, 1, true)); + try std.testing.expect(!active(null, 1.5, true)); + try std.testing.expect(active(.{ .force_guided = true }, 1, true)); + try std.testing.expect(!active(.{ .force_guided = true }, 1, false)); +} + +test "CFG native flags allocate guidance states only when enabled and refuse unsupported combinations" { + const cli = @import("cli.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var usage: cli.Usage = .{}; + if (@import("builtin").os.tag != .macos) { + try std.testing.expectError(error.Usage, cli.parse(a, &.{ "m", "--cfg-gamma", "1.5" }, &usage)); + return; + } + const enabled = try cli.parse(a, &.{ "m", "--cfg-gamma", "1.5", "--cfg-batched" }, &usage); + const request = try cli.request(a, "checkpoint", "fish_omni", enabled); + try std.testing.expect(request.cfg_enabled and request.cfg_batched); + try std.testing.expectEqual(@as(f32, 1.5), enabled.cfg.options().?.gamma); + try std.testing.expectError(error.GuidanceRequiresDrama, cli.request(a, "checkpoint", "text", enabled)); + const off = try cli.parse(a, &.{ "m", "--cfg-gamma", "0" }, &usage); + try std.testing.expect(!(try cli.request(a, "checkpoint", "fish_omni", off)).cfg_enabled); + try std.testing.expectError(error.Usage, cli.parse(a, &.{ "m", "--cfg-batched" }, &usage)); + try std.testing.expectError(error.Usage, cli.parse(a, &.{ "m", "--cfg-gamma", "nan" }, &usage)); + try std.testing.expectError(error.Usage, cli.parse(a, &.{ "m", "--cfg-gamma", "6" }, &usage)); + try std.testing.expectError(error.Usage, cli.parse(a, &.{ "m", "--cfg-gamma", "1.5", "--cfg-every", "2" }, &usage)); +} diff --git a/zig/src/server/drama_prompt.zig b/zig/src/server/drama_prompt.zig index 3e466c86d2..7acfadc108 100644 --- a/zig/src/server/drama_prompt.zig +++ b/zig/src/server/drama_prompt.zig @@ -25,6 +25,16 @@ pub const top_k: u32 = 30; /// The segments one prompt is encoded from, in order. pub const segment_count = 7; +/// One frame row a prompt token, the token in the text column. +pub fn rows(a: Allocator, prompt: []const u32) error{ OutOfMemory, InvalidToken }![]const [11]i32 { + const frames = try a.alloc([11]i32, prompt.len); + for (frames, prompt) |*row, token| { + row.* = @splat(0); + row[0] = std.math.cast(i32, token) orelse return error.InvalidToken; + } + return frames; +} + /// The Unicode White_Space property (``UNICODE_WHITE_SPACE``); U+001C..U+001F are not in it. pub fn isWhiteSpace(cp: u21) bool { return switch (cp) { diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index c996f9d371..3fdd321ada 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -50,6 +50,8 @@ test { _ = @import("chunk_plan.zig"); _ = @import("status_routes.zig"); _ = @import("speech.zig"); + _ = @import("drama_cfg.zig"); + if (comptime @import("builtin").os.tag != .windows) _ = @import("speech_cfg_test.zig"); if (comptime @import("builtin").os.tag != .windows) _ = @import("speech_stream_test.zig"); // over a Unix socket pair _ = @import("drama_voice_tool.zig"); _ = @import("speech_redraw.zig"); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index 8d2f741176..642cca332e 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -174,6 +174,7 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .seed_salt = salt, .silence_mid = args.silence.mid_seconds != null, .silence_guard = args.silence.enabled(), + .cfg = args.cfg.options(), .prompt_seed = args.prompt_seed, .sse_keepalive_s = keepalive_s, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 4ea16657a7..9c554fe635 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -67,6 +67,7 @@ pub const Config = struct { seed_salt: i64 = 0, silence_mid: bool = false, silence_guard: bool = false, + cfg: ?@import("drama_cfg.zig").Options = null, /// --prompt-seed keeps historical prompt-keyed sampling when a request omits its seed. prompt_seed: bool = false, /// TENSORFOLD_REQUEST_LOG: where chat and completion bodies are appended; null: nowhere. diff --git a/zig/src/server/speech.zig b/zig/src/server/speech.zig index 201e993a05..2897cd0e24 100644 --- a/zig/src/server/speech.zig +++ b/zig/src/server/speech.zig @@ -12,6 +12,7 @@ const reply_text = @import("reply_text.zig"); const drama_prompt = @import("drama_prompt.zig"); const speech_redraw = @import("speech_redraw.zig"); const speech_register = @import("speech_register_log.zig"); +const drama_cfg = @import("drama_cfg.zig"); const ids = @import("ids.zig"); const log = @import("log.zig"); const clock = @import("clock.zig"); @@ -62,6 +63,8 @@ pub fn post(srv: *Server, c: *Conn, a: Allocator) !void { var metadata: Metadata = .{ .enabled = srv.config.redraw.count() > 0, .register_on = srv.config.redraw.register.on() }; const raw = http_body.readJson(c, &cx) catch |e| return refuse(c, &cx, id, metadata, e); var plan = prepare(srv, &cx, sp, raw) catch |e| return refuse(c, &cx, id, metadata, e); + metadata.cfg_gamma = if (srv.config.cfg) |options| drama_cfg.gamma(options, c.header("X-Cfg-Gamma")) orelse + return fail(c, &cx, id, metadata, .request, "X-Cfg-Gamma must be a number from 1 (unguided) to 5") else null; // only chunked framing tells a client that a stream which ends early is incomplete if (plan.stream and !std.mem.eql(u8, c.version, "HTTP/1.1")) return fail(c, &cx, id, metadata, .request, "stream=true requires HTTP/1.1"); if (srv.config.redraw.count() > 0 or srv.config.redraw.register.on()) @@ -139,16 +142,6 @@ fn promptOf(srv: *Server, cx: *Cx, sp: api.SpeechInfo, input: []const u8) errors return prompt; } -/// One frame row a prompt token, the token in the text column. -fn rows(a: Allocator, prompt: []const u32) error{ OutOfMemory, InvalidToken }![]const [11]i32 { - const frames = try a.alloc([11]i32, prompt.len); - for (frames, prompt) |*row, token| { - row.* = @splat(0); - row[0] = std.math.cast(i32, token) orelse return error.InvalidToken; - } - return frames; -} - /// One synthesis: submit, then send the whole audio after the finished event, or stream it as the engine publishes it. fn synthesize(srv: *Server, c: *Conn, cx: *Cx, id: []const u8, sp: api.SpeechInfo, plan: Plan, received: i96, metadata: *Metadata) void { var held = srv.speech_inflight.load(.acquire); @@ -163,10 +156,12 @@ fn synthesize(srv: *Server, c: *Conn, cx: *Cx, id: []const u8, sp: api.SpeechInf const pcm = srv.gpa.alloc(f32, pcm_len) catch return fail(c, cx, id, metadata.*, .server, "out of memory"); defer srv.gpa.free(pcm); var result: api.SpeechResult = .{}; - const frames = plan.frames orelse rows(cx.a, plan.prompt) catch |e| return switch (e) { + const frames = plan.frames orelse drama_prompt.rows(cx.a, plan.prompt) catch |e| return switch (e) { error.OutOfMemory => fail(c, cx, id, metadata.*, .server, "out of memory"), error.InvalidToken => fail(c, cx, id, metadata.*, .request, "invalid speech token"), }; + const guided = drama_cfg.guide(cx.a, srv.text, srv.config.cfg, metadata.cfg_gamma orelse 1, plan.input, plan.refs.len > 0) catch + return fail(c, cx, id, metadata.*, .server, "the tokenizer cannot encode the unconditional prompt"); const request: api.SpeechRequest = .{ .prompt = plan.prompt, .frames = frames, @@ -179,6 +174,8 @@ fn synthesize(srv: *Server, c: *Conn, cx: *Cx, id: []const u8, sp: api.SpeechInf .top_k = drama_prompt.top_k, .stream = plan.stream, .reference_rows = plan.reference_rows, + .unconditioned = guided.rows, + .cfg_gamma = guided.gamma, }; var box: Mailbox = .{ .io = srv.io, .gpa = srv.gpa }; defer box.deinit(); @@ -226,8 +223,12 @@ fn redraw(srv: *Server, c: *Conn, cx: *Cx, id: []const u8, sp: api.SpeechInfo, p next.reference_rows = voice.reference_rows; } else { next.prompt = promptOf(srv, cx, sp, input) catch |e| return unused(id, k, fallback, if (e == error.OutOfMemory) "out of memory" else cx.message); - next.frames = rows(cx.a, next.prompt) catch |e| return unused(id, k, fallback, @errorName(e)); + next.frames = drama_prompt.rows(cx.a, next.prompt) catch |e| return unused(id, k, fallback, @errorName(e)); } + const guided = drama_cfg.guide(cx.a, srv.text, srv.config.cfg, metadata.cfg_gamma orelse 1, input, plan.refs.len > 0) catch + return unused(id, k, fallback, "the tokenizer cannot encode the unconditional prompt"); + next.unconditioned = guided.rows; + next.cfg_gamma = guided.gamma; } }; } else { diff --git a/zig/src/server/speech_cfg_test.zig b/zig/src/server/speech_cfg_test.zig new file mode 100644 index 0000000000..0cf802e566 --- /dev/null +++ b/zig/src/server/speech_cfg_test.zig @@ -0,0 +1,59 @@ +//! Exercise real speech routing, request frames and response headers with a deterministic engine sink. +const std = @import("std"); +const harness = @import("speech_stream_test.zig"); +const cfg = @import("drama_cfg.zig"); +const voice = @import("drama_voice.zig"); +const family = @import("drama_prompt"); +const gpa = std.testing.allocator; + +test "whole and streamed reference CFG carries only the plain prompt; default stays unguided" { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var text: harness.Bytes = .{}; + const guide = try cfg.guide(a, text.text(), .{ .gamma = 1.5 }, 1.5, "hi", true); + const codes = [_][10]i32{.{ 7, 1, 2, 3, 4, 5, 6, 7, 8, 9 }}; + const refs = [_]family.Reference{.{ .transcript = "A small cloud passes.", .codes = &codes }}; + const voices = [_]voice.Voice{.{ .name = "tone", .references = &refs }}; + const Case = struct { options: ?cfg.Options, header: []const u8, referenced: bool, gamma: f32, guided: bool, says: ?[]const u8 }; + const cases = [_]Case{ + .{ .options = null, .header = "X-Cfg-Gamma: nonsense\r\n", .referenced = true, .gamma = 1, .guided = false, .says = null }, + .{ .options = .{ .gamma = 1.5 }, .header = "X-Cfg-Gamma: nonsense\r\n", .referenced = true, .gamma = 1.5, .guided = true, .says = "X-Cfg-Gamma: 1.5\r\n" }, + .{ .options = .{ .gamma = 1.5 }, .header = "", .referenced = false, .gamma = 1, .guided = false, .says = "X-Cfg-Gamma: 1.5\r\n" }, + .{ .options = .{ .gamma = 1.5, .request_gamma = true }, .header = "X-Cfg-Gamma: 1\r\n", .referenced = true, .gamma = 1, .guided = false, .says = "X-Cfg-Gamma: 1\r\n" }, + .{ .options = .{ .gamma = 1.5, .request_gamma = true }, .header = "X-Cfg-Gamma: 3\r\n", .referenced = true, .gamma = 3, .guided = true, .says = "X-Cfg-Gamma: 3\r\n" }, + .{ .options = .{ .force_guided = true }, .header = "", .referenced = true, .gamma = 1, .guided = true, .says = "X-Cfg-Gamma: 1\r\n" }, + }; + for ([_]bool{ false, true }) |stream| for (cases) |case| { + const input = try std.fmt.allocPrint(a, "{{\"model\":\"m\",\"input\":\"hi\",\"voice\":\"{s}\",\"response_format\":\"pcm\",\"stream\":{s}}}", .{ if (case.referenced) "tone" else "default", if (stream) "true" else "false" }); + var script: harness.Script = .{ .steps = &.{harness.frame}, .inline_ = true, .voices = &voices, .expected_unconditioned = if (case.guided) guide.rows else &.{} }; + var served: harness.Served = undefined; + try served.open(&script, .{ .served_name = "m", .model_ids = &.{"m"}, .cfg = case.options }, case.header, "HTTP/1.1", input, 2000); + defer served.deinit(); + const head = try served.readHead(); + const success = std.mem.indexOf(u8, head, " 200 ") != null; + const header_ok = if (case.says) |says| std.mem.indexOf(u8, head, says) != null else std.mem.indexOf(u8, head, "X-Cfg-Gamma") == null; + _ = try served.readRest(); + served.join(&script); + try std.testing.expect(success and header_ok); + try std.testing.expectEqual(@as(usize, 1), script.takes); + try std.testing.expectEqual(case.gamma, script.cfg_gamma); + try std.testing.expect(script.unconditioned_equal); + }; +} + +test "invalid opted-in CFG header is refused before engine submission" { + for ([_][]const u8{ "nan", "inf", "0.5", "6", "wrong" }) |bad| { + var header_buf: [64]u8 = undefined; + const header = try std.fmt.bufPrint(&header_buf, "X-Cfg-Gamma: {s}\r\n", .{bad}); + var script: harness.Script = .{ .steps = &.{}, .inline_ = true }; + var served: harness.Served = undefined; + try served.open(&script, .{ .served_name = "m", .model_ids = &.{"m"}, .cfg = .{ .request_gamma = true } }, header, "HTTP/1.1", "{\"model\":\"m\",\"input\":\"hi\"}", 2000); + defer served.deinit(); + _ = try served.readHead(); + const body = try served.readRest(); + served.join(&script); + try std.testing.expectEqual(@as(usize, 0), script.takes); + try std.testing.expect(std.mem.indexOf(u8, body, "X-Cfg-Gamma must be a number") != null); + } +} diff --git a/zig/src/server/speech_metadata.zig b/zig/src/server/speech_metadata.zig index 3dfc8cfdd5..3e164d5246 100644 --- a/zig/src/server/speech_metadata.zig +++ b/zig/src/server/speech_metadata.zig @@ -7,8 +7,13 @@ pub const Metadata = struct { register_on: bool = false, register_tripped: bool = false, register_redrawn: bool = false, + cfg_gamma: ?f32 = null, pub fn add(m: Metadata, c: *Conn) !void { + if (m.cfg_gamma) |gamma| { + var gamma_buf: [24]u8 = undefined; + try c.addHeader("X-Cfg-Gamma", std.fmt.bufPrint(&gamma_buf, "{d}", .{gamma}) catch unreachable); + } if (!m.enabled) return; var buf: [3]u8 = undefined; try c.addHeader("X-Silence-Redraws", std.fmt.bufPrint(&buf, "{d}", .{m.redraws}) catch unreachable); diff --git a/zig/src/server/speech_redraw_test.zig b/zig/src/server/speech_redraw_test.zig index 8ab85c7aa0..cf45aaca84 100644 --- a/zig/src/server/speech_redraw_test.zig +++ b/zig/src/server/speech_redraw_test.zig @@ -226,5 +226,4 @@ test "a busy redraw stops at its latest start without counting an unsubmitted ta .busy_redraws = std.math.maxInt(usize), }; try expectReply(&script, .{ .max = 1 }, "X-Deadline-S: 0.1\r\n", "hi", "X-Finish-Reason: silence\r\n", 0, 1); - try std.testing.expect(script.busy_replies > 1); } diff --git a/zig/src/server/speech_stream_test.zig b/zig/src/server/speech_stream_test.zig index b245715b96..647167e382 100644 --- a/zig/src/server/speech_stream_test.zig +++ b/zig/src/server/speech_stream_test.zig @@ -38,6 +38,9 @@ pub const Script = struct { get_voices: bool = false, expected_rows: ?[]const [11]i32 = null, rows_equal: bool = false, + expected_unconditioned: []const [11]i32 = &.{}, + unconditioned_equal: bool = false, + cfg_gamma: f32 = 1, fn engine(e: *Script) api.Engine { return .{ .ctx = e, .vtable = &.{ .info = info, .submit = submit, .cancel = cancel, .status = status, .memory = memory, .submit_speech = speech } }; @@ -49,6 +52,8 @@ pub const Script = struct { fn speech(ctx: *anyopaque, id: api.Id, request: *const api.SpeechRequest, sink: api.Sink) api.SpeechSubmitError!void { const e: *Script = @ptrCast(@alignCast(ctx)); if (e.expected_rows) |expected| e.rows_equal = std.mem.eql(u8, std.mem.sliceAsBytes(expected), std.mem.sliceAsBytes(request.frames)); + e.unconditioned_equal = std.mem.eql(u8, std.mem.sliceAsBytes(e.expected_unconditioned), std.mem.sliceAsBytes(request.unconditioned)); + e.cfg_gamma = request.cfg_gamma; if (e.takes >= e.busy_from) return error.Busy; if (e.takes > 0 and e.busy_replies < e.busy_redraws) { e.busy_replies += 1; diff --git a/zig/tests/contract/drama3_nvfp4_ref.zig b/zig/tests/contract/drama3_nvfp4_ref.zig new file mode 100644 index 0000000000..ac766bc669 --- /dev/null +++ b/zig/tests/contract/drama3_nvfp4_ref.zig @@ -0,0 +1,248 @@ +//! Independent packed-format FP64 projections and signed ranges at Drama's BF16 SwiGLU nodes. +const std = @import("std"); +const contract = @import("contract"); +const c = contract.spec; +const rows = contract.rows; +const bf = contract.bf16_round.bf16; +pub const Interval = @import("gemma4_math.zig").Interval; + +pub const Variant = enum { up, down }; +pub const slots = 8; +pub const experts = 9; + +pub fn depth(s: c.Shape) usize { + return 16 + (s.k + 511) / 512 + 5; +} + +pub fn fp4(code: u8) f64 { + const magnitude: f64 = switch (code % 8) { + 0 => 0, + 1 => 0.5, + 2 => 1, + 3 => 1.5, + 4 => 2, + 5 => 3, + 6 => 4, + 7 => 6, + else => unreachable, + }; + return if (code >= 8) -magnitude else magnitude; +} + +pub fn scale(code: u8) f64 { + const bits = code % 128; + if (bits == 127) return std.math.nan(f64); + const exponent = bits / 8; + const mantissa = bits % 8; + const value = if (exponent == 0) + @as(f64, @floatFromInt(mantissa)) / 512 + else + std.math.ldexp(@as(f64, @floatFromInt(8 + @as(u16, mantissa))), @as(i32, exponent) - 10); + return if (code >= 128) -value else value; +} + +pub const Dot = struct { value: f64, magnitude: f64 }; +pub const Weights = struct { + a: std.mem.Allocator, + codes: []u8, + scales: []u8, + global: f32, + k: usize, + + pub fn init(a: std.mem.Allocator, shape: c.Shape, fill: rows.Fill, seed: u64, projection: usize) !Weights { + const codes = try a.alloc(u8, shape.n * shape.k / 2); + errdefer a.free(codes); + const scales = try a.alloc(u8, shape.n * shape.k / 16); + var rng_state = std.Random.DefaultPrng.init(seed); + const rng = rng_state.random(); + for (codes) |*v| v.* = rng.int(u8); + for (scales, 0..) |*v, i| v.* = switch (i % 131) { + 0 => 0, + 1 => 0x80, + else => @intCast((i +% seed) % 127), + }; + if (fill == .stress) { + @memset(scales, 0x38); + for (0..shape.n) |n| for (0..shape.k / 2) |j| { + const constant = projection < 2 and n % 2 != projection; + const lo: u8 = if (constant) (if (j == 0) 2 else 0) else if (2 * j < shape.k / 2) 2 else 10; + const hi: u8 = if (constant) 0 else if (2 * j + 1 < shape.k / 2) 2 else 10; + codes[n * (shape.k / 2) + j] = lo | (hi << 4); + }; + } + return .{ + .a = a, + .codes = codes, + .scales = scales, + .k = shape.k, + .global = if (fill == .stress) 1 else @as(f32, @floatFromInt(1 + seed % 7)) / 4096, + }; + } + + pub fn deinit(w: Weights) void { + w.a.free(w.codes); + w.a.free(w.scales); + } + + fn code(w: Weights, n: usize, k: usize) u8 { + const byte = w.codes[n * (w.k / 2) + k / 2]; + return if (k % 2 == 0) byte % 16 else byte / 16; + } + + pub fn dot(w: Weights, n: usize, x: []const u16) Dot { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..w.k / 16) |block| { + var subtotal: f64 = 0; + var absolute: f64 = 0; + for (0..16) |j| { + const k = block * 16 + j; + const term = fp4(w.code(n, k)) * @as(f64, c.fromBf16(x[k])); + subtotal += term; + absolute += @abs(term); + } + const s = scale(w.scales[n * (w.k / 16) + block]); + sum += subtotal * s; + magnitude += absolute * @abs(s); + } + return .{ .value = sum * w.global, .magnitude = magnitude * @abs(w.global) }; + } + + pub fn simulate(w: Weights, n: usize, x: []const u16, acc: c.Accum, order: rows.Order) f64 { + var blocks: [128]f64 = undefined; + for (blocks[0 .. w.k / 16], 0..) |*out, block| { + var terms: [16]f64 = undefined; + for (&terms, 0..) |*term, j| { + const k = block * 16 + j; + term.* = fp4(w.code(n, k)) * @as(f64, c.fromBf16(x[k])); + } + out.* = rows.sumIn(acc, order, &terms) * scale(w.scales[n * (w.k / 16) + block]); + } + return rows.sumIn(acc, order, blocks[0 .. w.k / 16]) * w.global; + } +}; + +pub fn fillX(x: []u16, k: usize, variant: Variant, fill: rows.Fill, seed: u64) void { + var state = std.Random.DefaultPrng.init(seed); + const rng = state.random(); + for (x) |*v| v.* = c.toBf16((rng.float(f32) * 2 - 1) / 32); + const count = x.len / k; + if (fill == .stress) { + // Integer partials stay below 2^24; a unit residue is lost inside a block in reduced precision. + for (0..count) |r| { + for (0..k / 2) |i| { + const v: f32 = @floatFromInt(256 * (1 + rng.uintLessThan(u32, 16))); + x[r * k + i] = c.toBf16(v); + x[r * k + k / 2 + i] = c.toBf16(v); + } + x[r * k] = c.toBf16(1); + x[r * k + k / 2] = 0; + } + } else if (fill == .nonfinite) { + for (0..count) |r| { + const token = if (variant == .up) r else r / slots; + switch (token) { + 0 => x[r * k + k / 3] = c.toBf16(std.math.nan(f32)), + 1 => x[r * k + k / 2] = c.toBf16(std.math.inf(f32)), + 2 => x[r * k + k - 1] = c.toBf16(-std.math.inf(f32)), + 3 => { + x[r * k] = c.toBf16(std.math.inf(f32)); + x[r * k + k - 1] = c.toBf16(-std.math.inf(f32)); + }, + else => {}, + } + } + } +} + +pub fn swiglu(gate: f64, up: f64) f64 { + const g = bf(gate); + return bf(g / (1 + @exp(-g)) * bf(up)); +} + +fn fp32(x: Interval) Interval { + const error_bound = c.unit(.f32) * @max(@abs(x.lo), @abs(x.hi)) + 0x1p-149; + return .{ .lo = x.lo - error_bound, .hi = x.hi + error_bound }; +} + +pub fn projectionRange(d: Dot, shape: c.Shape, exact: bool) Interval { + const reduction = if (exact) 0 else c.gamma(@floatFromInt(depth(shape)), c.unit(.f32)) * d.magnitude; + // The projection global is a separate FP32 multiply after the SIMD reduction. + return fp32(Interval.around(d.value, reduction)); +} + +pub fn upRange(gate: Interval, up: Interval) Interval { + const g = Interval{ .lo = bf(gate.lo), .hi = bf(gate.hi) }; + const u = Interval{ .lo = bf(up.lo), .hi = bf(up.hi) }; + // precise::exp's three-ulp FP32 envelope; the following add, divide and multiply each round once. + const exp_error = 6 * c.unit(.f32); + const ex = Interval{ .lo = @exp(-g.hi) * (1 - exp_error), .hi = @exp(-g.lo) * (1 + exp_error) }; + const denom = fp32(ex.add(Interval.point(1))); + const inverse = Interval{ .lo = 1 / denom.hi, .hi = 1 / denom.lo }; + return fp32(g.mul(inverse)).mul(u).rounded(); +} + +pub const shared_gate_terms = 65; + +pub fn sharedGateStressTerm(i: usize) f32 { + return if (i < 16) 32768 else if (i < 48) 1 else if (i < 64) -32768 else if (i == 64) -31 else 0; +} + +pub fn sharedGateOutput(dot: f64) f64 { + return bf(2 / (1 + @exp(-dot))); +} + +test "shared gate stress remains discriminating after sigmoid and BF16 output" { + var terms: [2048]f64 = @splat(0); + for (terms[0..shared_gate_terms], 0..) |*v, i| v.* = sharedGateStressTerm(i); + const want = sharedGateOutput(1); + try std.testing.expectEqual(@as(f64, 1.4609375), want); + for ([_]rows.Order{ .sequential, .pairwise, .strided }) |order| { + try std.testing.expectEqual(@as(f64, 1), rows.sumIn(.f32, order, &terms)); + for ([_]c.Accum{ .f16, .bf16 }) |acc| { + const got = sharedGateOutput(rows.sumIn(acc, order, &terms)); + // Exact FP32 sums yield one stable BF16 value, so the stress interval has zero width. + try std.testing.expect(!std.math.isFinite(got) or got != want); + } + } +} + +test "NVFP4 independent decoding includes signs, subnormal scales, signed zero and maximum finite scale" { + try std.testing.expectEqual(@as(f64, 0.5), fp4(1)); + try std.testing.expectEqual(@as(f64, -6), fp4(15)); + try std.testing.expect(std.math.signbit(fp4(8))); + try std.testing.expectEqual(@as(f64, 0x1p-9), scale(1)); + try std.testing.expectEqual(@as(f64, 1), scale(0x38)); + try std.testing.expectEqual(@as(f64, 448), scale(126)); + try std.testing.expectEqual(@as(f64, 0), scale(0x80)); + try std.testing.expect(std.math.signbit(scale(0x80))); +} + +test "block16 integer cancellation discriminates both reduced accumulators in all three orders" { + const a = std.testing.allocator; + for ([_]usize{ 512, 2048 }) |k| { + const shape = c.Shape{ .rows = 1, .n = 2, .k = k }; + const x = try a.alloc(u16, k); + defer a.free(x); + fillX(x, k, .up, .stress, 17); + for (0..3) |projection| { + const w = try Weights.init(a, shape, .stress, 23, projection); + defer w.deinit(); + const n = if (projection < 2) projection else 0; + const dot = w.dot(n, x); + try std.testing.expectEqual(@as(f64, 1), dot.value); + try std.testing.expect(dot.magnitude < 0x1p24); + for ([_]rows.Order{ .sequential, .pairwise, .strided }) |order| { + try std.testing.expectEqual(dot.value, w.simulate(n, x, .f32, order)); + for ([_]c.Accum{ .f16, .bf16 }) |acc| { + const got = w.simulate(n, x, acc, order); + const want_up = if (projection == 0) swiglu(dot.value, 1) else swiglu(1, dot.value); + const got_up = if (projection == 0) swiglu(got, 1) else swiglu(1, got); + const miss = if (projection < 2) @abs(got_up - want_up) else @abs(got - dot.value); + const bound = if (projection < 2) c.unit(.bf16) * @abs(want_up) else c.unit(.f32); + try std.testing.expect(!std.math.isFinite(got) or miss >= 100 * bound); + } + } + } + } +} diff --git a/zig/tests/contract/metal_drama3.zig b/zig/tests/contract/metal_drama3.zig index f17564a1ed..80a52d2793 100644 --- a/zig/tests/contract/metal_drama3.zig +++ b/zig/tests/contract/metal_drama3.zig @@ -12,7 +12,7 @@ const rows = contract.rows; /// A drama3 source with d3_common.h inlined, by the file its gate 5 label names ("drama3 "). pub fn source(a: std.mem.Allocator, label: []const u8) ![]u8 { const d = sources.drama3; - const files = .{ .{ "dense", d.dense }, .{ "gdn", d.gdn }, .{ "attn", d.attn }, .{ "moe", d.moe }, .{ "sample", d.sample }, .{ "codec", d.codec } }; + const files = .{ .{ "dense", d.dense }, .{ "gdn", d.gdn }, .{ "attn", d.attn }, .{ "moe", d.moe }, .{ "nvfp4", d.nvfp4 }, .{ "guidance", d.guidance }, .{ "sample", d.sample }, .{ "codec", d.codec } }; var it = std.mem.tokenizeScalar(u8, label, ' '); _ = it.next(); const file = it.next() orelse return error.UnknownLaunch; diff --git a/zig/tests/contract/metal_drama3_nvfp4.zig b/zig/tests/contract/metal_drama3_nvfp4.zig new file mode 100644 index 0000000000..28a96485b1 --- /dev/null +++ b/zig/tests/contract/metal_drama3_nvfp4.zig @@ -0,0 +1,239 @@ +//! Production NVFP4 gate/up and down: guarded indirect tables, routed prompt rows and independent FP64 ranges. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const c = contract.spec; +const g = contract.gates; +const rig_ = @import("rig.zig"); +const source = @import("metal_drama3.zig").source; +const ref = @import("drama3_nvfp4_ref.zig"); +pub const Variant = ref.Variant; + +pub fn spec(variant: Variant) c.Spec { + return .{ + .name = if (variant == .up) "metal drama3 d3_nvfp4_up (block16, BF16 SwiGLU)" else "metal drama3 d3_nvfp4_down (block16, FP32)", + .accum = .f32, + .depth = ref.depth, + .ops = if (variant == .up) 10 else 1, + .out = if (variant == .up) .bf16 else .f32, + .tiles = .{ .{ .size = 16 }, .{ .size = 4, .step = 4, .min = 4 }, .{ .size = 512, .step = 16, .min = 16 } }, + .served = if (variant == .up) &.{ + .{ .rows = 1, .n = 512, .k = 2048 }, + .{ .rows = 3, .n = 512, .k = 2048 }, + .{ .rows = 17, .n = 512, .k = 2048 }, + } else &.{ + .{ .rows = 1, .n = 2048, .k = 512 }, + .{ .rows = 3, .n = 2048, .k = 512 }, + .{ .rows = 17, .n = 2048, .k = 512 }, + }, + }; +} + +const Args = extern struct { K: u32, N: u32, prompt: u32, cap: u32, slots: u32 }; +const Packed = struct { w: rig_.Ref, s: rig_.Ref, global: rig_.Ref }; +const Sample = struct { want: f64, magnitude: f64, range: ref.Interval }; + +pub const Proj = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + variant: Variant, + shape: c.Shape = undefined, + fill: g.Fill = .random, + host: []u16 = &.{}, + samples: []Sample = &.{}, + weights: std.ArrayList(ref.Weights) = .empty, + packed_buffers: [16]Packed = undefined, + x: rig_.Ref = undefined, + table: rig_.Ref = undefined, + picks: rig_.Ref = undefined, + counts: rig_.Ref = undefined, + members: rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, variant: Variant) !Proj { + const symbol = if (variant == .up) "d3_nvfp4_up" else "d3_nvfp4_down"; + const label = if (variant == .up) "drama3 nvfp4 d3_nvfp4_up" else "drama3 nvfp4 d3_nvfp4_down"; + const src = try source(a, label); + defer a.free(src); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .pipe = try rig_.pipeline(device, src, symbol), .variant = variant }; + } + + fn clear(p: *Proj) void { + p.a.free(p.host); + p.host = &.{}; + p.a.free(p.samples); + p.samples = &.{}; + for (p.weights.items) |w| w.deinit(); + p.weights.clearRetainingCapacity(); + p.rig.reset(); + } + + pub fn deinit(p: *Proj) void { + p.clear(); + p.weights.deinit(p.a); + p.rig.deinit(); + p.pipe.deinit(); + } + + pub fn kernel(p: *Proj) g.Kernel { + return .{ .spec = spec(p.variant), .ctx = p, .vt = &vt }; + } + + const vt = g.hooks(Proj); + + fn perExpert(p: *const Proj) usize { + return if (p.variant == .up) 2 else 1; + } + + fn expert(pair: usize) usize { + return (3 * (pair % ref.slots) + pair / ref.slots) % ref.slots; + } + + fn address(b: rig_.Ref) u64 { + return b.buf.gpuAddress() + b.off; + } + + fn triple(w: Packed) [3]u64 { + return .{ address(w.w), address(w.s), address(w.global) }; + } + + pub fn prepare(p: *Proj, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const largest = spec(p.variant).served[2]; + if (shape.rows == 0 or shape.rows > largest.rows or shape.n == 0 or shape.n > largest.n or + shape.n % 4 != 0 or shape.k == 0 or shape.k > largest.k or shape.k % 16 != 0) return false; + p.clear(); + p.shape = shape; + p.fill = fill; + const pairs = shape.rows * ref.slots; + p.host = try p.a.alloc(u16, (if (p.variant == .up) shape.rows else pairs) * shape.k); + ref.fillX(p.host, shape.k, p.variant, fill, seed); + p.x = try p.rig.input(std.mem.sliceAsBytes(p.host)); + try p.weights.ensureTotalCapacity(p.a, ref.slots * p.perExpert()); + for (0..ref.slots * p.perExpert()) |i| { + const projection = if (p.variant == .up) i % 2 else 2; + const w = try ref.Weights.init(p.a, shape, fill, seed +% (i * 13), projection); + p.weights.appendAssumeCapacity(w); + p.packed_buffers[i] = .{ + .w = try p.rig.input(w.codes), + .s = try p.rig.input(w.scales), + .global = try p.rig.input(std.mem.asBytes(&w.global)), + }; + } + var table: [ref.experts][9]u64 = undefined; + for (&table, 0..) |*entry, e| { + const base = (e % ref.slots) * p.perExpert(); + const first = triple(p.packed_buffers[base]); + const second = triple(p.packed_buffers[base + p.perExpert() - 1]); + @memcpy(entry[0..3], &first); + @memcpy(entry[3..6], &second); + @memcpy(entry[6..9], &first); + } + p.table = try p.rig.input(std.mem.asBytes(&table)); + const picks = try p.a.alloc(i32, pairs); + defer p.a.free(picks); + const members = try p.a.alloc(i32, ref.experts * pairs); + defer p.a.free(members); + @memset(members, -1); + var counts: [ref.experts]i32 = @splat(0); + for (picks, 0..) |*pick, pair| pick.* = @intCast(expert(pair)); + // Reverse order proves the prompt scatter is not accidentally reading contiguous token rows. + for (0..pairs) |i| { + const pair = pairs - 1 - i; + const e = expert(pair); + members[e * pairs + @as(usize, @intCast(counts[e]))] = @intCast(pair); + counts[e] += 1; + } + p.picks = try p.rig.input(std.mem.sliceAsBytes(picks)); + p.counts = try p.rig.input(std.mem.asBytes(&counts)); + p.members = try p.rig.input(std.mem.sliceAsBytes(members)); + p.out = try p.rig.output(p.count() * c.size(spec(p.variant).out)); + p.samples = try p.a.alloc(Sample, p.count()); + p.computeReference(); + return true; + } + + pub fn run(p: *Proj) !void { + const pass = p.rig.begin(); + const e = pass.e; + e.setPipeline(p.pipe); + for (p.packed_buffers[0..p.weights.items.len]) |w| { + e.useResource(w.w.buf, 1); + e.useResource(w.s.buf, 1); + e.useResource(w.global.buf, 1); + } + for ([_]rig_.Ref{ p.x, p.table, p.picks, p.counts, p.members, p.out }, 0..) |b, i| + e.setBuffer(b.buf, b.off, i); + e.setValue(Args{ + .K = @intCast(p.shape.k), + .N = @intCast(p.shape.n), + .prompt = @intFromBool(p.shape.rows > 1), + .cap = @intCast(p.shape.rows * ref.slots), + .slots = ref.slots, + }, 6); + const groups_y: usize = if (p.shape.rows > 1) ref.experts else ref.slots; + e.dispatchGroups(mtl.Size.of(p.shape.n / 4, groups_y, 1), mtl.Size.of(128, 1, 1)); + try p.rig.finish(pass); + } + + pub fn count(p: *Proj) usize { + return p.shape.rows * ref.slots * p.shape.n; + } + + fn inputRow(p: *const Proj, pair: usize) []const u16 { + const r = if (p.variant == .up) pair / ref.slots else pair; + return p.host[r * p.shape.k ..][0..p.shape.k]; + } + + fn computeReference(p: *Proj) void { + for (0..p.shape.rows * ref.slots) |pair| { + const base = expert(pair) * p.perExpert(); + const x = p.inputRow(pair); + for (0..p.shape.n) |n| { + const first = p.weights.items[base].dot(n, x); + const range = ref.projectionRange(first, p.shape, p.fill == .stress); + var sample = Sample{ .want = first.value, .magnitude = first.magnitude, .range = range }; + if (p.variant == .up) { + const second = p.weights.items[base + 1].dot(n, x); + sample.want = ref.swiglu(first.value, second.value); + sample.magnitude += second.magnitude; + sample.range = ref.upRange(range, ref.projectionRange(second, p.shape, p.fill == .stress)); + } + p.samples[pair * p.shape.n + n] = sample; + } + } + } + + pub fn reference(p: *Proj, want: []f64, mag: []f64) !void { + for (p.samples, want, mag) |sample, *w, *m| { + w.* = sample.want; + m.* = sample.magnitude; + } + } + + pub fn bounds(p: *Proj, want: []const f64, _: []const f64, allow: []f64) !void { + for (p.samples, want, allow) |sample, w, *limit| + limit.* = @max(@abs(sample.range.lo - w), @abs(sample.range.hi - w)); + } + + pub fn intervals(p: *Proj, low: []f64, high: []f64) !void { + for (p.samples, low, high) |sample, *lo, *hi| { + lo.* = sample.range.lo; + hi.* = sample.range.hi; + } + } + + pub fn simulate(p: *Proj, acc: c.Accum, order: g.Order, out: []f64) !void { + for (0..p.shape.rows * ref.slots) |pair| { + const base = expert(pair) * p.perExpert(); + const x = p.inputRow(pair); + for (0..p.shape.n) |n| { + const first = p.weights.items[base].simulate(n, x, acc, order); + out[pair * p.shape.n + n] = if (p.variant == .up) + ref.swiglu(first, p.weights.items[base + 1].simulate(n, x, acc, order)) + else + first; + } + } + } +}; diff --git a/zig/tests/contract/metal_drama_cfg.zig b/zig/tests/contract/metal_drama_cfg.zig new file mode 100644 index 0000000000..23cd7eb952 --- /dev/null +++ b/zig/tests/contract/metal_drama_cfg.zig @@ -0,0 +1,154 @@ +//! Production semantic CFG: independent rows, conditional-only aliasing and untouched codec tails. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const source = @import("metal_drama3.zig").source; +const c = contract.spec; +const g = contract.gates; +const tail = 17; +pub const gammas = [_]f32{ 1, 1.5, 2, 3, 4, 5 }; + +fn depth(_: c.Shape) usize { + return 0; +} + +pub fn spec(gamma: f32) c.Spec { + return .{ + .name = if (gamma == 1) "metal drama3 d3_cfg_mix gamma=1 identity" else if (gamma == 1.5) "metal drama3 d3_cfg_mix gamma=1.5" else if (gamma == 2) "metal drama3 d3_cfg_mix gamma=2" else if (gamma == 3) "metal drama3 d3_cfg_mix gamma=3" else if (gamma == 4) "metal drama3 d3_cfg_mix gamma=4" else "metal drama3 d3_cfg_mix gamma=5", + // Pointwise subtraction, multiply and addition; no reduction or accumulated partial sums. + .accum = .none, + .depth = depth, + .ops = 3, + .out = .f32, + .in_place = true, + .nan = .{ .stated = "gamma=1 preserves all conditional bits without reading unconditional; otherwise IEEE formula" }, + .tiles = .{ .{}, .{ .size = 256 }, .{} }, + .served = &.{ .{ .rows = 1, .n = 4097, .k = 1 }, .{ .rows = 5, .n = 4097, .k = 1 } }, + }; +} + +pub const Mix = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + gamma: f32, + shape: c.Shape = undefined, + conditional: []f32 = &.{}, + unconditional: []f32 = &.{}, + u: rig_.Ref = undefined, + out: rig_.Ref = undefined, + msg: [200]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, gamma: f32) !Mix { + if (!std.math.isFinite(gamma) or gamma < 1 or gamma > 5) return error.InvalidGamma; + const src = try source(a, "drama3 guidance d3_cfg_mix"); + defer a.free(src); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .gamma = gamma, .pipe = try rig_.pipeline(device, src, "d3_cfg_mix") }; + } + + pub fn deinit(t: *Mix) void { + t.a.free(t.conditional); + t.a.free(t.unconditional); + t.rig.deinit(); + t.pipe.deinit(); + } + + pub fn kernel(t: *Mix) g.Kernel { + return .{ .spec = spec(t.gamma), .ctx = t, .vt = &vt }; + } + const vt = blk: { + var hooks = g.hooks(Mix); + hooks.overlap = overlap; + break :blk hooks; + }; + fn overlap(_: *anyopaque) bool { + return true; + } + + pub fn prepare(t: *Mix, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.k != 1 or shape.n == 0 or shape.n > 4097) return false; + t.rig.reset(); + t.shape = shape; + t.a.free(t.conditional); + t.conditional = &.{}; + t.a.free(t.unconditional); + t.unconditional = &.{}; + t.conditional = try t.a.alloc(f32, t.count()); + t.unconditional = try t.a.alloc(f32, t.count()); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (t.conditional, t.unconditional) |*cv, *uv| { + cv.* = rng.float(f32) * 16 - 8; + uv.* = rng.float(f32) * 16 - 8; + } + t.conditional[0] = -0.0; + if (t.gamma == 1) t.unconditional[0] = std.math.nan(f32); + if (fill == .nonfinite) { + const values = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (values, 0..) |v, row| { + t.conditional[row * (shape.n + tail) + 1] = v; + t.unconditional[row * (shape.n + tail) + 2] = v; + } + t.conditional[3 * (shape.n + tail) + 3] = std.math.inf(f32); + t.unconditional[3 * (shape.n + tail) + 3] = std.math.inf(f32); + } + t.u = try t.rig.input(std.mem.sliceAsBytes(t.unconditional)); + // Only conditional logits are mutable. Their initial bytes live on the host, not in an immutable rig slot. + t.out = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Mix) !void { + @memcpy(t.rig.bytes(t.out, t.count() * 4), std.mem.sliceAsBytes(t.conditional)); + const p = t.rig.begin(); + p.e.setPipeline(t.pipe); + const Args = extern struct { count: u32, gamma: f32 }; + p.e.setValue(Args{ .count = @intCast(t.shape.n), .gamma = t.gamma }, 2); + for (0..t.shape.rows) |row| { + const off = row * (t.shape.n + tail) * 4; + p.e.setBuffer(t.out.buf, t.out.off + off, 0); + p.e.setBuffer(t.u.buf, t.u.off + off, 1); + p.e.dispatchGroups(mtl.Size.of((t.shape.n + 255) / 256, 1, 1), mtl.Size.of(256, 1, 1)); + } + try t.rig.finish(p); + } + + pub fn count(t: *Mix) usize { + return t.shape.rows * (t.shape.n + tail); + } + + pub fn reference(t: *Mix, want: []f64, mag: []f64) !void { + for (want, mag, t.conditional, t.unconditional, 0..) |*w, *m, cv, uv, i| { + const active = t.gamma != 1 and i % (t.shape.n + tail) < t.shape.n; + w.* = if (active) @as(f64, cv) + (@as(f64, t.gamma) - 1) * (@as(f64, cv) - uv) else cv; + m.* = if (active) @abs(cv) + (@as(f64, t.gamma) - 1) * (@abs(cv) + @abs(uv)) else 0; + } + } + + pub fn check(t: *Mix, at: g.Gate) !?[]const u8 { + const out = t.rig.bytes(t.out, t.count() * 4); + var nonfinite: usize = 0; + for (t.conditional, t.unconditional, 0..) |cv, uv, i| { + const got = c.read(.f32, out, i); + if (t.gamma == 1 or i % (t.shape.n + tail) >= t.shape.n) { + const bits = std.mem.readInt(u32, out[i * 4 ..][0..4], .little); + if (bits != @as(u32, @bitCast(cv))) + return try std.fmt.bufPrint(&t.msg, "conditional identity/codec tail changed at {d}", .{i}); + nonfinite += @intFromBool(!std.math.isFinite(cv)); + continue; + } + const alpha = @as(f64, t.gamma) - 1; + const want = @as(f64, cv) + alpha * (@as(f64, cv) - uv); + nonfinite += @intFromBool(!std.math.isFinite(want)); + if (c.classOf(got) != c.classOf(want)) + return try std.fmt.bufPrint(&t.msg, "CFG class mismatch at {d}", .{i}); + // Three rounded FP32 operations (or their contraction), bounded by the term magnitudes. + const bound = c.gamma(3, c.unit(.f32)) * (@abs(cv) + alpha * (@abs(cv) + @abs(uv))); + if (at != .nonfinite and std.math.isFinite(want) and @abs(got - want) > bound) + return try std.fmt.bufPrint(&t.msg, "CFG FP64 bound exceeded at {d}", .{i}); + } + if (at == .nonfinite and (nonfinite == 0 or nonfinite == t.count())) return "nonfinite fixture proves nothing"; + return null; + } +}; diff --git a/zig/tests/contract/metal_drama_combine.zig b/zig/tests/contract/metal_drama_combine.zig new file mode 100644 index 0000000000..21291a6bd7 --- /dev/null +++ b/zig/tests/contract/metal_drama_combine.zig @@ -0,0 +1,286 @@ +//! CUDA-node-compatible NVFP4 combine: independent FP64 nodes and separate pair/shared-dot sum stress. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const bfr = contract.bf16_round; +const rig_ = @import("rig.zig"); +const source = @import("metal_drama3.zig").source; +const ref = @import("drama3_nvfp4_ref.zig"); +const c = contract.spec; +const g = contract.gates; +pub const Variant = enum { pairs, shared_gate }; +fn depth(s: c.Shape) usize { + return (s.k + 255) / 256 + 13; +} + +pub fn spec(v: Variant) c.Spec { + return .{ + .name = switch (v) { + .pairs => "metal drama3 d3_nvfp4_combine pair sums", + .shared_gate => "metal drama3 d3_nvfp4_combine shared gate sums", + }, + .accum = .f32, + .depth = depth, + .out = .bf16, + .in_place = true, + .tiles = .{ .{}, .{}, .{ .size = 256 } }, + .served = &.{ .{ .rows = 8, .n = 2048, .k = 2048 }, .{ .rows = 1, .n = 2048, .k = 2048 }, .{ .rows = 168, .n = 2048, .k = 2048 } }, + }; +} + +const Range = struct { + lo: f64, + hi: f64, + fn point(v: f64) Range { + return .{ .lo = v, .hi = v }; + } + fn around(v: f64, e: f64) Range { + return .{ .lo = v - e, .hi = v + e }; + } + fn add(a: Range, b: Range) Range { + return .{ .lo = a.lo + b.lo, .hi = a.hi + b.hi }; + } + fn scale(a: Range, b: f64) Range { + return if (b >= 0) .{ .lo = a.lo * b, .hi = a.hi * b } else .{ .lo = a.hi * b, .hi = a.lo * b }; + } + fn fp32(a: Range) Range { + const err = c.unit(.f32) * @max(@abs(a.lo), @abs(a.hi)) + 0x1p-149; + return .{ .lo = a.lo - err, .hi = a.hi + err }; + } + fn bf16(a: Range) Range { + return .{ .lo = rounded(a.lo), .hi = rounded(a.hi) }; + } +}; +fn rounded(v: f64) f64 { + return bfr.bf16(@as(f32, @floatCast(v))); +} +fn sigmoid(v: f64) f64 { + return 1 / (1 + @exp(-v)); +} + +pub const Combine = struct { + a: std.mem.Allocator, + arena: std.heap.ArenaAllocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + variant: Variant, + shape: c.Shape = undefined, + fill: g.Fill = .random, + old: []u16 = &.{}, + shared: []u16 = &.{}, + x: []u16 = &.{}, + wg: []u16 = &.{}, + down: []f32 = &.{}, + weights: []f32 = &.{}, + refs: [5]rig_.Ref = undefined, + out: rig_.Ref = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, variant: Variant) !Combine { + const src = try source(a, "drama3 nvfp4 d3_nvfp4_combine"); + defer a.free(src); + return .{ .a = a, .arena = std.heap.ArenaAllocator.init(a), .rig = .{ .a = a, .device = device, .queue = queue }, .variant = variant, .pipe = try rig_.pipeline(device, src, "d3_nvfp4_combine") }; + } + pub fn deinit(t: *Combine) void { + t.arena.deinit(); + t.rig.deinit(); + t.pipe.deinit(); + } + pub fn kernel(t: *Combine) g.Kernel { + return .{ .spec = spec(t.variant), .ctx = t, .vt = &vt }; + } + const vt = blk: { + var hooks = g.hooks(Combine); + hooks.overlap = overlap; + break :blk hooks; + }; + fn overlap(_: *anyopaque) bool { + return true; + } + + pub fn prepare(t: *Combine, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.k == 0 or shape.k > 2048 or shape.n != shape.k) return false; + t.rig.reset(); + _ = t.arena.reset(.retain_capacity); + t.shape = shape; + t.fill = fill; + const a = t.arena.allocator(); + const n = t.count(); + t.old = try a.alloc(u16, n); + t.shared = try a.alloc(u16, n); + t.x = try a.alloc(u16, n); + t.wg = try a.alloc(u16, shape.k); + t.down = try a.alloc(f32, n * 8); + t.weights = try a.alloc(f32, shape.rows * 8); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (t.old, t.shared, t.x) |*old, *shared, *x| { + old.* = c.toBf16(rng.float(f32) * 2 - 1); + shared.* = c.toBf16(rng.float(f32) * 2 - 1); + x.* = c.toBf16((rng.float(f32) * 2 - 1) / 8); + } + for (t.wg) |*v| v.* = c.toBf16((rng.float(f32) * 2 - 1) / 8); + for (t.down) |*v| v.* = rng.float(f32) * 8 - 4; + for (t.weights) |*v| v.* = (1 + rng.float(f32)) / 16; + if (fill == .stress) { + @memset(t.old, 0); + @memset(t.shared, c.toBf16(if (t.variant == .shared_gate) 2 else 0)); + @memset(t.x, 0); + @memset(t.wg, c.toBf16(1)); + @memset(t.down, 0); + @memset(t.weights, 0.125); + if (t.variant == .pairs) { + const terms = [_]f32{ 65536, 1, 1, 1, 1, 1, 1, -65536 }; + for (0..shape.rows) |row| for (terms, 0..) |v, slot| + @memset(t.down[(row * 8 + slot) * shape.k ..][0..shape.k], v * 8); + } else { + // Exact dot 1 avoids sigmoid saturation hiding a reduced-precision sum's overflow. + for (0..shape.rows) |row| for (0..ref.shared_gate_terms) |i| { + t.x[row * shape.k + i] = c.toBf16(ref.sharedGateStressTerm(i)); + }; + } + } + if (fill == .nonfinite) { + const values = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (values, 0..) |v, row| { + t.down[(row * 8) * shape.k] = v; + t.shared[row * shape.k + 1] = c.toBf16(v); + t.old[row * shape.k + 2] = c.toBf16(v); + t.x[row * shape.k + 3] = c.toBf16(v); + } + t.x[3 * shape.k] = c.toBf16(std.math.inf(f32)); + t.x[3 * shape.k + 1] = c.toBf16(-std.math.inf(f32)); + t.wg[0] = c.toBf16(1); + t.wg[1] = c.toBf16(1); + t.wg[3] = c.toBf16(1); + } + t.out = try t.rig.output(n * 2); + t.refs[0] = try t.rig.input(std.mem.sliceAsBytes(t.down)); + t.refs[1] = try t.rig.input(std.mem.sliceAsBytes(t.weights)); + t.refs[2] = try t.rig.input(std.mem.sliceAsBytes(t.shared)); + t.refs[3] = try t.rig.input(std.mem.sliceAsBytes(t.x)); + t.refs[4] = try t.rig.input(std.mem.sliceAsBytes(t.wg)); + return true; + } + + pub fn run(t: *Combine) !void { + // Residual alone is allowed in place. Restore its original payload before each independent repeat. + @memcpy(t.rig.bytes(t.out, t.count() * 2), std.mem.sliceAsBytes(t.old)); + const p = t.rig.begin(); + p.e.setPipeline(t.pipe); + p.e.setBuffer(t.out.buf, t.out.off, 0); + for (t.refs, 1..) |buffer, slot| p.e.setBuffer(buffer.buf, buffer.off, slot); + p.e.setValue(@as(u32, @intCast(t.shape.k)), 6); + p.e.dispatchGroups(mtl.Size.of(t.shape.rows, 1, 1), mtl.Size.of(256, 1, 1)); + try t.rig.finish(p); + } + pub fn count(t: *Combine) usize { + return t.shape.rows * t.shape.k; + } + + fn dot(t: *Combine, row: usize) struct { value: f64, magnitude: f64 } { + var value: f64 = 0; + var magnitude: f64 = 0; + for (t.x[row * t.shape.k ..][0..t.shape.k], t.wg) |x, w| { + const term = @as(f64, c.fromBf16(x)) * c.fromBf16(w); + value += term; + magnitude += @abs(term); + } + return .{ .value = value, .magnitude = magnitude }; + } + + fn pairs(t: *Combine, row: usize, col: usize, terms: *[8]f64) void { + for (terms, 0..) |*term, slot| { + const product = @as(f64, t.weights[row * 8 + slot]) * t.down[(row * 8 + slot) * t.shape.k + col]; + term.* = rounded(product); + } + } + fn finish(t: *Combine, row: usize, col: usize, gate: f64, sum: f64) f64 { + const i = row * t.shape.k + col; + const moe = rounded(gate * c.fromBf16(t.shared[i]) + rounded(sum)); + return rounded(@as(f64, c.fromBf16(t.old[i])) + moe); + } + + pub fn reference(t: *Combine, want: []f64, mag: []f64) !void { + for (0..t.shape.rows) |row| { + const d = t.dot(row); + for (0..t.shape.k) |col| { + var terms: [8]f64 = undefined; + t.pairs(row, col, &terms); + var sum: f64 = 0; + var magnitude: f64 = 0; + for (terms) |term| { + sum += term; + magnitude += @abs(term); + } + want[row * t.shape.k + col] = t.finish(row, col, sigmoid(d.value), sum); + mag[row * t.shape.k + col] = magnitude + d.magnitude; + } + } + } + + fn range(t: *Combine, row: usize, col: usize, gate: Range) Range { + var sum = Range.point(0); + var magnitude: f64 = 0; + for (0..8) |slot| { + const product = @as(f64, t.weights[row * 8 + slot]) * t.down[(row * 8 + slot) * t.shape.k + col]; + // A single FP32 product has an exact independently rounded value; CUDA rounds it again to BF16. + const pair = rounded(product); + sum = sum.add(Range.point(pair)); + magnitude += @abs(pair); + } + const error_sum = c.gamma(8, c.unit(.f32)) * magnitude; + sum = Range.around(sum.lo, error_sum).bf16(); + const i = row * t.shape.k + col; + const shared = gate.scale(c.fromBf16(t.shared[i])).fp32(); + const moe = shared.add(sum).fp32().bf16(); + return moe.add(Range.point(c.fromBf16(t.old[i]))).fp32().bf16(); + } + + pub fn intervals(t: *Combine, low: []f64, high: []f64) !void { + if (t.fill == .stress) { + // Exact FP32 integer sums stay away from BF16 rounding boundaries after the sigmoid. + const exact: f64 = if (t.variant == .pairs) 6 else ref.sharedGateOutput(1); + @memset(low, exact); + @memset(high, exact); + return; + } + for (0..t.shape.rows) |row| { + const d = t.dot(row); + const error_dot = c.gamma(@floatFromInt(depth(t.shape)), c.unit(.f32)) * d.magnitude; + const lo = sigmoid(d.value - error_dot); + const hi = sigmoid(d.value + error_dot); + // precise exp's two ulps, the denominator addition and reciprocal; propagate through every BF16 node. + const error_gate = c.gamma(6, c.unit(.f32)) * hi; + const gate = Range{ .lo = @max(0, lo - error_gate), .hi = @min(1, hi + error_gate) }; + for (0..t.shape.k) |col| { + const r = t.range(row, col, gate); + low[row * t.shape.k + col] = r.lo; + high[row * t.shape.k + col] = r.hi; + } + } + } + + pub fn bounds(t: *Combine, want: []const f64, _: []const f64, allow: []f64) !void { + const ranges = try t.a.alloc(f64, t.count() * 2); + defer t.a.free(ranges); + const low = ranges[0..t.count()]; + const high = ranges[t.count()..]; + try t.intervals(low, high); + for (allow, want, low, high) |*v, w, lo, hi| v.* = @max(@abs(w - lo), @abs(hi - w)); + } + + pub fn simulate(t: *Combine, acc: c.Accum, order: g.Order, out: []f64) !void { + const terms = try t.a.alloc(f64, t.shape.k); + defer t.a.free(terms); + for (0..t.shape.rows) |row| { + for (terms, t.wg, 0..) |*v, w, i| + v.* = @as(f64, c.fromBf16(t.x[row * t.shape.k + i])) * c.fromBf16(w); + const gate = sigmoid(contract.rows.sumIn(acc, order, terms)); + for (0..t.shape.k) |col| { + var pairs_: [8]f64 = undefined; + t.pairs(row, col, &pairs_); + out[row * t.shape.k + col] = t.finish(row, col, gate, contract.rows.sumIn(acc, order, &pairs_)); + } + } + } +}; diff --git a/zig/tests/contract/metal_drama_router.zig b/zig/tests/contract/metal_drama_router.zig new file mode 100644 index 0000000000..9f96411251 --- /dev/null +++ b/zig/tests/contract/metal_drama_router.zig @@ -0,0 +1,224 @@ +//! Compact routing: original-ID ordering, masking before softmax and compact slot translation. +const std = @import("std"); +const mtl = @import("metal"); +const contract = @import("contract"); +const rig_ = @import("rig.zig"); +const source = @import("metal_drama3.zig").source; +const c = contract.spec; +const g = contract.gates; +pub const Variant = enum { pruned8, pruned73, all256 }; +const offset = 19; +const stride = 291; + +fn depth(_: c.Shape) usize { + return 21; +} +pub fn spec(v: Variant) c.Spec { + return .{ + .name = switch (v) { + .pruned8 => "metal drama3 d3_router_compact retained=8", + .pruned73 => "metal drama3 d3_router_compact retained=73", + .all256 => "metal drama3 d3_router_compact retained=256", + }, + .accum = .f32, + .depth = depth, + .ops = 12, + .out = .f32, + .nan = .{ .stated = "dropped nonfinites ignored; invalid softmax gives NaN weights and first8 retained original IDs" }, + .served = &.{ .{ .rows = 8, .n = 8, .k = 256 }, .{ .rows = 1, .n = 8, .k = 256 }, .{ .rows = 168, .n = 8, .k = 256 } }, + // Seven exp(-10) terms survive FP32 but both reduced-precision denominators round to 1 in every order. + .sums = .{ .exact = false }, + }; +} + +pub const Router = struct { + a: std.mem.Allocator, + rig: rig_.Rig, + pipe: mtl.Pipeline, + variant: Variant, + shape: c.Shape = undefined, + logits: []u16 = &.{}, + slots: [256]i32 = undefined, + previous: []u8 = &.{}, + runs: usize = 0, + x: rig_.Ref = undefined, + map: rig_.Ref = undefined, + picks: rig_.Ref = undefined, + weights: rig_.Ref = undefined, + msg: [200]u8 = undefined, + + pub fn init(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, variant: Variant) !Router { + const src = try source(a, "drama3 nvfp4 d3_router_compact"); + defer a.free(src); + return .{ .a = a, .rig = .{ .a = a, .device = device, .queue = queue }, .variant = variant, .pipe = try rig_.pipeline(device, src, "d3_router_compact") }; + } + pub fn deinit(t: *Router) void { + t.a.free(t.logits); + t.a.free(t.previous); + t.rig.deinit(); + t.pipe.deinit(); + } + pub fn kernel(t: *Router) g.Kernel { + return .{ .spec = spec(t.variant), .ctx = t, .vt = &vt }; + } + const vt = g.hooks(Router); + + pub fn prepare(t: *Router, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.n != 8 or shape.k != 256) return false; + t.rig.reset(); + t.shape = shape; + t.runs = 0; + t.a.free(t.logits); + t.logits = &.{}; + t.a.free(t.previous); + t.previous = &.{}; + t.logits = try t.a.alloc(u16, shape.rows * stride); + t.previous = try t.a.alloc(u8, t.count() * 4); + @memset(&t.slots, -1); + const kept: usize = switch (t.variant) { + .pruned8 => 8, + .pruned73 => 73, + .all256 => 256, + }; + // Permuted compact slots prove that ties are resolved in original ID space, not slot space. + for (0..kept) |slot| t.slots[(slot * 37 + 7) % 256] = @intCast(kept - 1 - slot); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + @memset(t.logits, c.toBf16(91)); + for (0..shape.rows) |row| for (0..256) |id| { + const score: f32 = if (t.slots[id] < 0) 100 else if (row == 0 and fill != .stress) -2 else @as(f32, @floatFromInt(rng.uintLessThan(u32, 33))) / 8 - 2; + t.logits[row * stride + offset + id] = c.toBf16(score); + }; + if (fill == .stress) for (0..shape.rows) |row| { + var rank: usize = 0; + for (t.slots, 0..) |slot, id| if (slot >= 0) { + t.logits[row * stride + offset + id] = c.toBf16(if (rank == 0) 0 else if (rank < 8) -10 else -80); + rank += 1; + }; + }; + if (fill == .nonfinite) { + var ids: [256]usize = undefined; + var n: usize = 0; + for (t.slots, 0..) |slot, id| if (slot >= 0) { + ids[n] = id; + n += 1; + }; + const values = [_]f32{ std.math.nan(f32), std.math.inf(f32), -std.math.inf(f32) }; + for (values, 0..) |plant, row| t.logits[row * stride + offset + ids[1]] = c.toBf16(plant); + t.logits[3 * stride + offset + ids[0]] = c.toBf16(std.math.inf(f32)); + t.logits[3 * stride + offset + ids[1]] = c.toBf16(-std.math.inf(f32)); + if (kept < 256) for (t.slots, 0..) |slot, id| if (slot < 0) { + t.logits[4 * stride + offset + id] = c.toBf16(values[id % 3]); + }; + for (t.slots, 0..) |slot, id| if (slot >= 0) { + t.logits[5 * stride + offset + id] = c.toBf16(-std.math.inf(f32)); + }; + } + t.x = try t.rig.input(std.mem.sliceAsBytes(t.logits)); + t.map = try t.rig.input(std.mem.sliceAsBytes(&t.slots)); + t.picks = try t.rig.scratch(t.count() * 4); + t.weights = try t.rig.output(t.count() * 4); + return true; + } + + pub fn run(t: *Router) !void { + const p = t.rig.begin(); + p.e.setPipeline(t.pipe); + p.e.setBuffer(t.x.buf, t.x.off, 0); + p.e.setBuffer(t.picks.buf, t.picks.off, 1); + p.e.setBuffer(t.weights.buf, t.weights.off, 2); + p.e.setBuffer(t.map.buf, t.map.off, 3); + p.e.setValue([2]u32{ stride, offset }, 4); + p.e.dispatchGroups(mtl.Size.of(t.shape.rows, 1, 1), mtl.Size.of(256, 1, 1)); + try t.rig.finish(p); + const picked = t.rig.bytes(t.picks, t.count() * 4); + if (t.runs > 0 and !std.mem.eql(u8, t.previous, picked)) return error.RouterPicksNotRepeatable; + @memcpy(t.previous, picked); + t.runs += 1; + } + pub fn count(t: *Router) usize { + return t.shape.rows * 8; + } + + fn selected(t: *Router, row: usize, ids: *[8]usize) bool { + var used: [256]bool = @splat(false); + var poisoned = false; + var finite = false; + for (t.slots, 0..) |slot, id| if (slot >= 0) { + const v = c.fromBf16(t.logits[row * stride + offset + id]); + poisoned = poisoned or std.math.isNan(v) or v == std.math.inf(f32); + finite = finite or std.math.isFinite(v); + }; + poisoned = poisoned or !finite; + for (ids) |*out| { + var best: ?usize = null; + for (t.slots, 0..) |slot, id| { + if (slot < 0 or used[id]) continue; + if (best == null or (!poisoned and t.value(row, id) > t.value(row, best.?))) best = id; + } + out.* = best.?; + used[out.*] = true; + } + return poisoned; + } + fn value(t: *Router, row: usize, id: usize) f64 { + return c.fromBf16(t.logits[row * stride + offset + id]); + } + + pub fn reference(t: *Router, want: []f64, mag: []f64) !void { + for (0..t.shape.rows) |row| { + var ids: [8]usize = undefined; + const poisoned = t.selected(row, &ids); + var terms: [8]f64 = undefined; + var sum: f64 = 0; + for (ids, &terms) |id, *term| { + term.* = @exp(t.value(row, id) - t.value(row, ids[0])); + sum += term.*; + } + for (terms, 0..) |v, j| { + want[row * 8 + j] = if (poisoned) std.math.nan(f64) else v / sum; + mag[row * 8 + j] = @abs(want[row * 8 + j]); + } + } + } + + pub fn simulate(t: *Router, acc: c.Accum, order: g.Order, out: []f64) !void { + for (0..t.shape.rows) |row| { + var ids: [8]usize = undefined; + _ = t.selected(row, &ids); + var all: [256]f64 = @splat(0); + for (t.slots, 0..) |slot, id| if (slot >= 0) { + all[id] = @exp(t.value(row, id) - t.value(row, ids[0])); + }; + const total = contract.rows.sumIn(acc, order, &all); + var chosen: [8]f64 = undefined; + for (ids, &chosen) |id, *v| v.* = all[id] / total; + const denom = contract.rows.sumIn(acc, order, &chosen); + for (chosen, 0..) |v, j| out[row * 8 + j] = v / denom; + } + } + + pub fn check(t: *Router, at: g.Gate) !?[]const u8 { + const picked = t.rig.bytes(t.picks, t.count() * 4); + const weights = t.rig.bytes(t.weights, t.count() * 4); + const values = try t.a.alloc(f64, t.count() * 2); + defer t.a.free(values); + const want = values[0..t.count()]; + const mag = values[t.count()..]; + try t.reference(want, mag); + for (0..t.shape.rows) |row| { + var ids: [8]usize = undefined; + _ = t.selected(row, &ids); + for (ids, 0..) |id, j| { + const i = row * 8 + j; + const slot = std.mem.readInt(i32, picked[i * 4 ..][0..4], .little); + if (slot != t.slots[id]) return try std.fmt.bufPrint(&t.msg, "row {d} rank {d}: compact slot {d}, expected {d} (original ID {d})", .{ row, j, slot, t.slots[id], id }); + const got = c.read(.f32, weights, i); + if (c.classOf(got) != c.classOf(want[i])) return "router weight class differs from masked FP64 softmax"; + if (at != .nonfinite and @abs(got - want[i]) > c.bound(spec(t.variant), t.shape, want[i], mag[i])) + return "router weight exceeds FP64 reduction bound"; + } + } + return null; + } +}; diff --git a/zig/tests/contract/metal_main.zig b/zig/tests/contract/metal_main.zig index 32eedeb698..c544e4beaa 100644 --- a/zig/tests/contract/metal_main.zig +++ b/zig/tests/contract/metal_main.zig @@ -11,6 +11,10 @@ const limits_k = @import("metal_limits.zig"); const h3_k = @import("metal_h3.zig"); const h3_topk = @import("metal_h3_topk.zig"); const drama3_k = @import("metal_drama3.zig"); +const drama_nvfp4 = @import("metal_drama3_nvfp4.zig"); +const drama_router = @import("metal_drama_router.zig"); +const drama_combine = @import("metal_drama_combine.zig"); +const drama_cfg = @import("metal_drama_cfg.zig"); const glm_eh_k = @import("metal_glm_eh.zig"); const flash_pool = @import("metal_flashnext_pool.zig"); @@ -187,6 +191,42 @@ pub fn main(init: std.process.Init) !u8 { d3norm = try drama3_k.Norm.init(a, device, queue); try kernels.append(a, d3norm.?.kernel()); } + var d3projections: [std.enums.values(drama_nvfp4.Variant).len]drama_nvfp4.Proj = undefined; + var d3projection_count: usize = 0; + defer for (d3projections[0..d3projection_count]) |*p| p.deinit(); + for (std.enums.values(drama_nvfp4.Variant)) |v| { + if (!wanted(o, drama_nvfp4.spec(v).name)) continue; + d3projections[d3projection_count] = try drama_nvfp4.Proj.init(a, device, queue, v); + d3projection_count += 1; + try kernels.append(a, d3projections[d3projection_count - 1].kernel()); + } + var d3routers: [std.enums.values(drama_router.Variant).len]drama_router.Router = undefined; + var d3router_count: usize = 0; + defer for (d3routers[0..d3router_count]) |*p| p.deinit(); + for (std.enums.values(drama_router.Variant)) |v| { + if (!wanted(o, drama_router.spec(v).name)) continue; + d3routers[d3router_count] = try drama_router.Router.init(a, device, queue, v); + d3router_count += 1; + try kernels.append(a, d3routers[d3router_count - 1].kernel()); + } + var d3combines: [std.enums.values(drama_combine.Variant).len]drama_combine.Combine = undefined; + var d3combine_count: usize = 0; + defer for (d3combines[0..d3combine_count]) |*p| p.deinit(); + for (std.enums.values(drama_combine.Variant)) |v| { + if (!wanted(o, drama_combine.spec(v).name)) continue; + d3combines[d3combine_count] = try drama_combine.Combine.init(a, device, queue, v); + d3combine_count += 1; + try kernels.append(a, d3combines[d3combine_count - 1].kernel()); + } + var d3mixes: [drama_cfg.gammas.len]drama_cfg.Mix = undefined; + var d3mix_count: usize = 0; + defer for (d3mixes[0..d3mix_count]) |*p| p.deinit(); + for (drama_cfg.gammas) |gamma| { + if (!wanted(o, drama_cfg.spec(gamma).name)) continue; + d3mixes[d3mix_count] = try drama_cfg.Mix.init(a, device, queue, gamma); + d3mix_count += 1; + try kernels.append(a, d3mixes[d3mix_count - 1].kernel()); + } var glm_eh: ?glm_eh_k.Nax = null; defer if (glm_eh) |*n| n.deinit(); if (wanted(o, glm_eh_k.spec().name)) { diff --git a/zig/tests/drama_nvfp4_metal.zig b/zig/tests/drama_nvfp4_metal.zig new file mode 100644 index 0000000000..5d5e6399a9 --- /dev/null +++ b/zig/tests/drama_nvfp4_metal.zig @@ -0,0 +1,10 @@ +//! Check the production NVFP4 kernels against independent references, without loading model weights. +const std = @import("std"); +const drama = @import("drama3"); + +pub fn main() !void { + const gpu = try drama.gpu.Gpu.open(std.heap.page_allocator, @import("drama3_metallib").bytes); + defer gpu.close(); + try drama.quant_check.run(gpu); + try drama.guidance.check(gpu); +}