From 39de0393a2a3d52844d8d0d7714504d0261aedfb Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Wed, 14 Oct 2026 18:36:00 +0000 Subject: [PATCH 001/115] lane core and native CUDA host for two-rank families with a vision tower: batched prompt passes (prefill_many, addStreams), a backend's refusals for a memory guard, follow() for a tensor-parallel rank 1, shared-pool admission fields, image and video inputs for the native server (vision_inputs, media_fetch, qwen_image, stb_image) --- build.zig | 5 +- zig/build/cuda.zig | 6 +- zig/build/vision.zig | 10 + zig/src/core/engine_api.zig | 21 +- zig/src/core/lane_host.zig | 91 +- zig/src/core/lanes/backend.zig | 6 + zig/src/core/lanes/engine.zig | 68 + zig/src/core/lanes/fake.zig | 8 + zig/src/core/lanes/lanes.zig | 1 + zig/src/core/lanes/sampling.zig | 8 +- zig/src/core/lanes/stream.zig | 7 + zig/src/native/cuda.zig | 86 +- zig/src/native/cuda_bound.zig | 12 + zig/src/native/no_family.zig | 2 + zig/src/native/no_family_glm.zig | 2 + zig/src/server/chat.zig | 35 +- zig/src/server/cli.zig | 31 +- zig/src/server/main.zig | 4 + zig/src/server/media_fetch.zig | 336 ++ zig/src/server/openai.zig | 7 +- zig/src/server/prompt.zig | 14 + zig/src/server/serve.zig | 4 + zig/src/server/server.zig | 5 + zig/src/server/vision_inputs.zig | 342 ++ zig/src/vision/glm_image.zig | 194 + zig/src/vision/qwen_image.zig | 432 ++ zig/src/vision/tf_video.c | 220 + zig/vendor/stb/README.md | 3 + zig/vendor/stb/stb_image.h | 7988 ++++++++++++++++++++++++++++++ zig/vendor/stb/stb_image_impl.c | 10 + 30 files changed, 9922 insertions(+), 36 deletions(-) create mode 100644 zig/build/vision.zig create mode 100644 zig/src/native/no_family.zig create mode 100644 zig/src/native/no_family_glm.zig create mode 100644 zig/src/server/media_fetch.zig create mode 100644 zig/src/server/vision_inputs.zig create mode 100644 zig/src/vision/glm_image.zig create mode 100644 zig/src/vision/qwen_image.zig create mode 100644 zig/src/vision/tf_video.c create mode 100644 zig/vendor/stb/README.md create mode 100644 zig/vendor/stb/stb_image.h create mode 100644 zig/vendor/stb/stb_image_impl.c diff --git a/build.zig b/build.zig index 08765bd1df..481ad34d26 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig"); const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); +const vision_build = @import("zig/build/vision.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -267,7 +268,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b counters_build.targets(b, target, optimize, metal, mods.sources, test_step); // The engines registry: the families this platform's binaries serve; the CLI reads it for models/info/pull. - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = mods.lanes }} }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = mods.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, @@ -572,7 +573,7 @@ pub fn distMetalServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_id m.metal.addSystemFrameworkPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "System/Library/Frameworks" }))); m.metal.addLibraryPath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/lib" }))); m.metal.addSystemIncludePath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/include" }))); - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = m.lanes }} }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = m.lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, .ReleaseFast) } } }); const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = m.metal }, .{ .name = "engine_api", .module = api }, .{ .name = "tensorfold", .module = m.engine } } }); return nativeServer(b, target, .fast, m.metal, m.engine, api, engines, build_options, null); } diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..4b4fdbb2dc 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -3,6 +3,7 @@ const std = @import("std"); const exlr8_build = @import("exlr8.zig"); const golden_build = @import("golden.zig"); +const vision_build = @import("vision.zig"); /// Each .cu in zig/kernels/cuda (`src`, else `name`) with the flags its Python extension passes in `extra_cuda_cflags`. const Kernel = struct { name: []const u8, flags: []const []const u8, src: ?[]const u8 = null, arch_specific: bool = false }; @@ -129,6 +130,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); runner.addImport("cuda", cuda); runner.addImport("core", mods.core); + runner.addImport("qwen_image", vision_build.module(b, target, optimize)); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); @@ -144,7 +146,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { - const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); + const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "lanes", .module = lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, @@ -210,7 +212,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug), vision_build.module(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/build/vision.zig b/zig/build/vision.zig new file mode 100644 index 0000000000..11eb50c143 --- /dev/null +++ b/zig/build/vision.zig @@ -0,0 +1,10 @@ +//! The image-input module (zig/src/vision/qwen_image.zig with stb_image compiled in), shared by every build that +//! carries the engine API (the server reaches it as engine_api.qwen_image) and by the CUDA test runner. +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addCSourceFile(.{ .file = b.path("zig/vendor/stb/stb_image_impl.c"), .flags = &.{ "-std=c11", "-O2" } }); + m.addIncludePath(b.path("zig/vendor/stb")); + return m; +} diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 903dc4d40b..3c633640ad 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -1,5 +1,8 @@ //! The engine a server drives (``Engine``), and ``LaneHost``: the lane core served on one thread, rounds left to it. const std = @import("std"); +/// Image inputs (decode, resize, patches) for families with a vision tower; the server reaches them here. +pub const Image = lanes.Image; +pub const qwen_image = @import("qwen_image"); const lanes = @import("lanes"); const Allocator = std.mem.Allocator; @@ -88,6 +91,8 @@ pub const Request = struct { logprobs: ?u8 = null, /// Rows for each prompt token after the first, with ranks and this many best tokens, when ``Info.prompt_logprobs``. prompt_logprobs: ?u8 = null, + /// An image prompt's inputs in prompt order (Info.vision engines only; the prompt holds their placeholder rows). + images: []const lanes.Image = &.{}, }; pub const Reason = enum { stop, length, cancelled, failed }; @@ -162,6 +167,10 @@ pub const Info = struct { prompt_cache_plan: ?PromptCachePlan = null, /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, + /// The engine encodes image inputs (Request.images): it was opened with ``vision`` and its family has a tower. + vision: bool = false, + /// Its tower takes GLM-5.3-Flash's preparation and placeholders (qwen_image.glm) instead of Qwen's. + vision_glm: bool = false, /// Whether a request's caches fit beside the running streams' now (null: a free lane does); else it waits in order. admit: ?Admit = null, }; @@ -225,10 +234,20 @@ pub const Open = struct { master_port: u16 = 29551, /// --policy: the engine's policy as `key=value,...`; empty: its defaults. policy: []const u8 = "", + /// --vision: load the family's vision tower and accept image inputs. + vision: bool = false, + /// --kv-dtype: bits of the attention cache's codes (8 or 4) for families that keep a quantized cache. + kv_bits: u8 = 8, }; /// An opened engine; ``close`` stops its thread and frees its backend. -pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque }; +pub const Opened = struct { + engine: Engine, + close: *const fn (ctx: *anyopaque) void, + ctx: *anyopaque, + /// Rank 1 of two: replay rank 0's engine calls until it stops (the server then serves no HTTP); null elsewhere. + follow: ?*const fn (ctx: *anyopaque) anyerror!void = null, +}; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0 }; diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4394fdac05..7a557c20c3 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -286,15 +286,16 @@ pub const LaneHost = struct { } /// Prefills the next queued request into a free lane; false when none waits or no lane is free. - fn admitOne(h: *LaneHost) bool { + /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit yet). + fn takeJob(h: *LaneHost) ?*Job { h.lock(); if (h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) { h.unlock(); - return false; + return null; } const at = admit.pick(h) orelse { // a shared cache pool's head waits for room h.unlock(); - return false; + return null; }; const job = h.queued.orderedRemove(at); job.seq = h.admissions; @@ -302,9 +303,14 @@ pub const LaneHost = struct { h.admitted.append(h.gpa, job) catch { h.unlock(); h.finish(job, .failed, "out of memory"); - return true; + return null; }; h.unlock(); + return job; + } + + /// The admitted request's stream (its kept state, drafter and grammar); false when it was dropped instead. + fn setup(h: *LaneHost, job: *Job) bool { const r = job.request; if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran var reuse: lanes.stream.Reuse = .{}; @@ -315,7 +321,10 @@ pub const LaneHost = struct { job.marks = l.marks; reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); + job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch { + _ = h.drop(job, "the drafter could not start"); + return false; + }; job.stream = lanes.Stream.init(h.gpa, .{ .id = "request", .prompt = r.prompt, @@ -339,18 +348,31 @@ pub const LaneHost = struct { .reuse = reuse, .grammar = if (r.structure) |st| h.constraint(st) catch |e| { job.proposer.deinit(); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + { + _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return false; + } } else null, + .images = r.images, .decode_spans = h.spans(r), .decide = r.decide, .untaught = r.untaught, }) catch { job.proposer.deinit(); - return h.drop(job, "out of memory"); + { + _ = h.drop(job, "out of memory"); + return false; + } }; job.started = true; - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return true; + } + + /// An admitted, set-up request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + fn launch(h: *LaneHost, job: *Job) bool { + const r = job.request; + const began = job.began; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (lone.fits(h, job)) return lone.run(h, job, began, false); if (h.core.fills()) { @@ -366,6 +388,55 @@ pub const LaneHost = struct { if (h.deliver(job)) h.remove(job); return true; } + fn admitOne(h: *LaneHost) bool { + const job = h.takeJob() orelse return false; + if (!h.setup(job)) return true; + return h.launch(job); + } + + /// Every queued request a lane is free for, their prompts in one pass where the backend takes several + /// (lanes.Engine.addStreams); the engine's special paths (chunked prompts, warm passes, the lone driver) go + /// through admitOne as before. + fn admitMany(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return h.admitOne(); + h.lock(); + const waiting = h.queued.items.len; + const free = h.info_.lanes -| h.admitted.items.len; + h.unlock(); + if (waiting < 2 or free < 2) return h.admitOne(); + var jobs: [lanes.Engine.max_together]*Job = undefined; + var n: usize = 0; + while (n < jobs.len and n < free) { + const job = h.takeJob() orelse break; + if (!h.setup(job)) continue; + if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { + _ = h.launch(job); + continue; + } + jobs[n] = job; + n += 1; + } + if (n == 0) return true; + if (n == 1) return h.launch(jobs[0]); + var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; + for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; + var errs: [lanes.Engine.max_together]?anyerror = undefined; + h.core.addStreams(streams[0..n], errs[0..n]); + for (jobs[0..n], errs[0..n]) |job, err| { + if (err) |e| { + _ = switch (e) { + error.Cancelled => h.cancel(job), + error.LaneMemory => h.requeue(job), + else => h.drop(job, h.words(e)), + }; + continue; + } + h.prefilled(job, job.began); + if (h.deliver(job)) h.remove(job); + } + return true; + } + /// The next filling job's next prompt chunk: the other streams' rounds run between chunks, not after the prompt. fn fillOne(h: *LaneHost) void { @@ -494,7 +565,7 @@ pub const LaneHost = struct { fn run(h: *LaneHost) void { while (true) { h.takeCancels(); - while (h.admitOne()) {} + while (h.admitMany()) {} h.fillOne(); status.noteLive(h); h.lock(); diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index fe26d3f003..723dc5effc 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -68,6 +68,9 @@ pub const Backend = struct { prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, /// One chunk of the prompt pass, false while chunks remain; null: the pass runs whole. prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// Prefill several streams' prompts in one pass (each then as after `prefill`); false, having done nothing, + /// when they do not go together (each is then prefilled alone). + prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, /// Feed one token and queue the draw of the next at `position` (a one-token round); a handle to it. @@ -102,6 +105,9 @@ pub const Backend = struct { prompt_rows: bool = false, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, + /// After a failed round: the streams it failed for alone (their caches could not grow), which the caller + /// ends while the round's other streams go on; 0 when the failure is the round's. + refused: ?*const fn (ptr: *anyopaque, out: []*Stream) usize = null, }; pub fn prefill(b: Backend, s: *Stream) !void { diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index cf0a97ea4a..d5be72c0bb 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -83,6 +83,68 @@ pub const Engine = struct { pub const fills = open.fills; pub const fillStream = open.fillStream; + /// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. + pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { + _ = e.arena.reset(.retain_capacity); + s.context.shrinkRetainingCapacity(s.prompt_len); + s.pending = null; + s.cache_len = s.prompt_len; + if (e.cfg.family_mtp and s.drafts) { + const d: u32 = @intCast(try e.rule.depth(win.who(s))); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { + s.next = held; + }; + } + if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row + _ = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.finished) return e.release(s); + try e.live.append(e.gpa, s); + } + + /// Streams a backend prefills in one pass at most. + pub const max_together = 16; + + /// Whether the backend prefills several prompts in one pass (addStreams batches them). + pub fn batchesPrompts(e: *const Engine) bool { + return e.backend.vtable.prefill_many != null; + } + + /// open.add for several streams: their prompts in one pass when the backend takes them together (else each + /// alone), then each opened as after its own prefill. ``errs[i]``: stream i's error (a cancelled one is released + /// already; the caller discards a failed one). + pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (errs) |*x| x.* = null; + const many = e.backend.vtable.prefill_many orelse return e.addEach(ss, errs); + if (ss.len < 2 or ss.len > max_together) return e.addEach(ss, errs); + _ = e.arena.reset(.retain_capacity); + for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; + const together = many(e.backend.ptr, ss) catch |err| { + for (errs) |*x| x.* = err; + return; + }; + if (!together) return e.addEach(ss, errs); + for (ss, errs) |s, *x| { + if (s.isCancelled()) { // cancelled in its pass: no first token + e.backend.release(s); + x.* = error.Cancelled; + continue; + } + open.opened(e, s) catch |err| { + x.* = err; + }; + } + } + + fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| open.add(e, s) catch |err| { + x.* = err; + }; + } + /// Take over a stream another driver decoded so far: its cache settled, `pending` and `cache_len` set, no drafts held. pub fn adopt(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); @@ -100,6 +162,12 @@ pub const Engine = struct { return false; } + /// After ``step`` failed: the streams the backend refused alone (the caller ends them; the rest go on). + pub fn refused(e: *Engine, out: []*Stream) usize { + const ask = e.backend.vtable.refused orelse return 0; + return ask(e.backend.ptr, out); + } + /// Release a cancelled stream between rounds: no pending draws or drafts. pub fn discard(e: *Engine, s: *Stream) void { s.finished = true; diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index 9434736c46..fc90ee7483 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -243,6 +243,14 @@ pub const Fake = struct { return self(ptr).drawn.items[handle]; } + fn refusedFn(ptr: *anyopaque, out: []*Stream) usize { + const x = self(ptr); + const n = @min(out.len, x.refused.items.len); + @memcpy(out[0..n], x.refused.items[0..n]); + x.refused.clearRetainingCapacity(); + return n; + } + fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { const x = self(ptr); x.rounds += 1; diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 5d50a16ef3..57d8ae9eb7 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -29,6 +29,7 @@ pub const grammar = @import("grammar.zig"); pub const Engine = engine.Engine; pub const Config = config.Config; +pub const Image = @import("stream.zig").Image; pub const Model = config.Model; pub const Stream = stream.Stream; pub const Sampling = sampling.Sampling; diff --git a/zig/src/core/lanes/sampling.zig b/zig/src/core/lanes/sampling.zig index e9d469756c..43bddfa3a1 100644 --- a/zig/src/core/lanes/sampling.zig +++ b/zig/src/core/lanes/sampling.zig @@ -75,7 +75,7 @@ pub fn pairwiseSum(a: []const f64) f64 { return pairwiseSum(a[0..n2]) + pairwiseSum(a[n2..]); } -const Candidate = struct { value: f64, id: u64 }; +pub const Candidate = struct { value: f64, id: u64 }; fn byValueThenId(_: void, a: Candidate, b: Candidate) bool { return a.value > b.value or (a.value == b.value and a.id < b.id); @@ -88,6 +88,12 @@ pub fn choose(gpa: Allocator, values: []const f64, ids: []const u64, position: u defer gpa.free(order); for (order, values, ids) |*o, v, id| o.* = .{ .value = v, .id = id }; std.mem.sort(Candidate, order, {}, byValueThenId); + return chooseOrdered(gpa, order, position, s); +} + +/// ``choose`` over candidates already in its order (by value, then the lower id at a tie). +pub fn chooseOrdered(gpa: Allocator, order: []const Candidate, position: u64, s: Sampling) !u64 { + const width = order.len; if (s.temperature <= 0) return order[0].id; // temperature 0 is greedy: the first maximum, ties to the lower id const k = @max(1, @min(if (s.top_k != 0) @as(usize, s.top_k) else width, width)); const scaled = try gpa.alloc(f64, k); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index f028d10cc4..aade9ad50c 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -97,6 +97,10 @@ pub const Held = struct { parents: ?[]i32 = null, }; +/// An image input prepared for a vision tower: its patches (fp32 [grid h * w, 1536], merge-block order) and grid +/// [1, h, w]; the prompt holds its placeholder rows. +pub const Image = struct { patches: []const f32, grid: [3]i64 }; + pub const Spec = struct { id: []const u8, prompt: []const u32, @@ -123,6 +127,7 @@ pub const Spec = struct { prompt_logprobs: ?u8 = null, // and for each prompt token after the first, with ranks (the prompt pass starts at 0) grammar: ?*grammar.Constraint = null, // structured output: the reply stays in it (the creator frees it) decode_spans: []const [2]u32 = &.{}, // prompt rows an earlier reply decoded (Request.decode_spans), for backends that keep their bits + images: []const Image = &.{}, // an image prompt's inputs, in prompt order (the backend encodes them at prefill) }; pub const Stream = struct { @@ -140,6 +145,7 @@ pub const Stream = struct { think_close: []const u32, think_end: i64, think_open: bool, + images: []const Image = &.{}, loop_guard: bool, loop_period: ?u32 = null, chunks: []const u32, @@ -202,6 +208,7 @@ pub const Stream = struct { .cancel_check = spec.cancel_check, .yield_check = spec.yield_check, .think_budget = spec.think_budget, + .images = spec.images, .think_close = spec.think_close, .think_end = spec.think_end, .think_open = spec.think_open orelse (spec.think_budget > 0 or (spec.loop_guard and spec.think_end >= 0)), diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..a34186a4a3 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,4 +1,7 @@ //! CUDA family registration, device setup, memory admission and the native lane host. +/// The serve flags this build's engine takes past the common ones (server/cli.zig served()): two ranks. +pub const serves: []const []const u8 = &.{"--tp"}; + const std = @import("std"); const cuda = @import("cuda"); const api = @import("engine_api"); @@ -152,6 +155,14 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + follow_fn: ?*const fn (*anyopaque) anyerror!void = null, // rank 1 of two: the family's replay loop + + /// Rank 1: the family replays rank 0's calls on this thread (the context was made current here at open). + fn followFamily(p: *anyopaque) anyerror!void { + const h: *Host = @ptrCast(@alignCast(p)); + if (h.gpu) |g| try g.ctx.makeCurrent(); + return h.follow_fn.?(h.family); + } fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -241,6 +252,22 @@ const Host = struct { const refusals = @import("cuda_refusal.zig"); +/// A tensor-parallel rank above 0 (server/engines.zig follow): the family opened as that rank, then its replay loop +/// until rank 0 stops; false with ``problem`` set when it cannot. +pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { + const opened = (try open(a, gpa, io, o, problem)) orelse return false; + defer opened.close(opened.ctx); + const f = opened.follow orelse { + problem.* = "this family has no tensor-parallel ranks"; + return false; + }; + f(opened.ctx) catch |e| { + problem.* = try std.fmt.allocPrint(a, "rank {d} stopped: {s}", .{ o.rank, @errorName(e) }); + return false; + }; + return true; +} + /// The engine for `o.dir`, or null with `problem` set when `refusal` names a flag or no CUDA family reads it. pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { if (refusals.backend(o)) |why| { @@ -322,20 +349,42 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }, }; const before = try pool(a, io, &g.ctx, problem) orelse return null; - const weights = weightBytes(io, o.dir); + const two = @hasDecl(F, "two_ranks") and F.two_ranks; + if (two and o.tp != 2) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on two GPUs here: serve it with --tp 2 --rank 0 on one host and --tp 2 --rank 1 --master on the other", .{F.model_type}); + return null; + } + if (!two and o.tp != 1) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on one GPU: drop --tp", .{F.model_type}); + return null; + } + const weights = if (@hasDecl(F, "deviceWeightBytes")) F.deviceWeightBytes(io, o.dir) else weightBytes(io, o.dir); const held0 = cuda.usage(false).device; if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); return null; } - const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }) catch |e| { + var fo: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments }; + if (@hasField(F.Options, "vision")) fo.vision = o.vision; + if (@hasField(F.Options, "kv_bits")) fo.kv_bits = o.kv_bits else if (o.kv_bits != 8) { + problem.* = try std.fmt.allocPrint(a, "{s} keeps an int8 attention cache only: drop --kv-dtype", .{F.model_type}); + return null; + } + if (@hasField(F.Options, "slide")) fo.slide = o.slide; + if (@hasField(F.Options, "rank")) { + fo.tp = @intCast(o.tp); + fo.rank = @intCast(o.rank); + fo.master = if (o.master.len > 0) o.master else null; + fo.master_port = o.master_port; + } + const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, fo) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null; }; defer if (!opened) loaded.deinit(loaded.ctx); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| loaded.learn_bytes; + const room = after.room(model) -| learnBytes(loaded); const free = loaded.free_streams; const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { @@ -344,16 +393,22 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; return null; }; + if (@hasDecl(F, "setStreams")) F.setStreams(loaded.ctx, streams); + if (@hasDecl(F, "setRoom")) F.setRoom(loaded.ctx, room) catch |e| { // a family that sizes each stream's caches out of one pool + problem.* = try std.fmt.allocPrint(a, "the ranks could not agree on the stream pool ({s})", .{@errorName(e)}); + return null; + }; const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); - h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + // the lane core's kept states (Nemotron's snapshots); a family that keeps its own prompt states has no `target` + const cache = if (@hasField(@TypeOf(loaded), "target")) budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)) else 0; + h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in {d}-row chunks{s}{s}; {s}", .{ + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learnBytes(loaded) > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learnBytes(loaded))}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasField(@TypeOf(loaded), "target") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet @@ -364,7 +419,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (comptime @hasField(@TypeOf(loaded), "learns")) { // a family with a Living Weights learner + if (loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); @@ -378,17 +434,29 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O } }.f; h.learn = s.hook(); + } } errdefer if (h.slide) |s| { h.slide_free.?(gpa, s); gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, + .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; + const follow_loop: ?*const fn (*anyopaque) anyerror!void = if (@hasField(@TypeOf(loaded), "follow")) loaded.follow else null; + if (follow_loop) |f| { + h.follow_fn = f; + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h, .follow = Host.followFamily }; + } return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } +/// A family's bytes kept for --slide (0 for families without a learner). +fn learnBytes(loaded: anytype) u64 { + return if (@hasField(@TypeOf(loaded), "learn_bytes")) loaded.learn_bytes else 0; +} + fn toGib(bytes: u64) f64 { return @as(f64, @floatFromInt(bytes)) / gib; } diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 40977dfa50..2144981686 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -18,6 +18,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.prefill_step.?(x.inner.ptr, s); } }.f else null, + .prefill_many = if (v.prefill_many != null) struct { + fn f(p: *anyopaque, ss: []const *lanes.Stream) anyerror!bool { + const x = H.bind(p); + return x.inner.vtable.prefill_many.?(x.inner.ptr, ss); + } + }.f else null, .first = struct { fn f(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const x = H.bind(p); @@ -108,6 +114,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.held.?(x.inner.ptr, s, out); } }.f else null, + .refused = if (v.refused != null) struct { + fn f(p: *anyopaque, out: []*lanes.Stream) usize { + const x = H.bind(p); + return x.inner.vtable.refused.?(x.inner.ptr, out); + } + }.f else null, .prompt_rows = v.prompt_rows, .release = struct { fn f(p: *anyopaque, s: *lanes.Stream) void { diff --git a/zig/src/native/no_family.zig b/zig/src/native/no_family.zig new file mode 100644 index 0000000000..d81d7567ac --- /dev/null +++ b/zig/src/native/no_family.zig @@ -0,0 +1,2 @@ +//! Stands in for a CUDA family the build leaves out (-Dflashnext=false): nothing is registered. +pub const native = struct {}; diff --git a/zig/src/native/no_family_glm.zig b/zig/src/native/no_family_glm.zig new file mode 100644 index 0000000000..2dcb07a292 --- /dev/null +++ b/zig/src/native/no_family_glm.zig @@ -0,0 +1,2 @@ +//! Stands in for the GLM CUDA family when the build leaves it out (-Dglm=false): nothing is registered. +pub const native = struct {}; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 7c702ef432..82c72d52ad 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -35,6 +35,8 @@ pub const Input = struct { logprobs: ?u8 = null, /// ``/v1/completions``' ``prompt_logprobs``: rows for the prompt with this many alternatives; null: not asked. prompt_logprobs: ?u8 = null, + /// Image inputs (vision_inputs.extract), in prompt order; the rendered prompt's image pads expand to their tokens. + images: []const api.Image = &.{}, }; /// A streamed piece: content text (a string) or a delta object (reasoning or tool calls). @@ -90,9 +92,31 @@ pub const Prepared = struct { preparing: bool, }; +/// --tool-system: the server's instruction as the first message of a chat request that offers tools and sends no +/// system message (the client's own instruction wins); other requests pass unchanged. +fn withToolSystem(srv: *Server, a: Allocator, input: Input) Allocator.Error!Input { + const text = srv.config.tool_system orelse return input; + if (input.tools.len == 0 or input.prompt != null or input.messages != .array) return input; + for (input.messages.array) |m| { + if (m != .object) continue; + const role = m.object.get("role") orelse continue; + if (role == .string and std.mem.eql(u8, role.string, "system")) return input; + } + const sys = try json.newObject(a); + try sys.put(a, "role", .{ .string = "system" }); + try sys.put(a, "content", .{ .string = text }); + const list = try a.alloc(Value, input.messages.array.len + 1); + list[0] = .{ .object = sys }; + @memcpy(list[1..], input.messages.array); + var out = input; + out.messages = .{ .array = list }; + return out; +} + /// Render ``input`` and run every check that can refuse it, before anything reaches the client or the engine. -pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepared { +pub fn prepare(srv: *Server, cx: *Cx, input_in: Input, gone: anytype) Failure!Prepared { const a = cx.a; + const input = try withToolSystem(srv, a, input_in); const io = srv.io; const received = nowNs(io); const f = input.fields; @@ -108,7 +132,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); - const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + var rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + if (input.images.len > 0) rendered.history_len = 0; // image prompts keep no prefix for a later turn (as Python) if (gone.check()) return error.Cancelled; if (rendered.ids.len == 0) return cx.refuse("rendered prompt is empty"); const window: i64 = srv.info.context_window; @@ -119,7 +144,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (input.max_tokens != null and limit > room) return cx.fail(.context_length, "{s} {d} tokens, but the rendered prompt has {d} tokens and requests {d} reply tokens, which exceeds the context window. Reduce the prompt to at most {d} prompt tokens or request at most {d} reply tokens, including chat template and thinking tokens.", .{ errors.context_limit, window, n, limit, @max(0, window - limit), room }); limit = @min(limit, room); } - const system_len: usize = if (input.prompt != null) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); + const system_len: usize = if (input.prompt != null or input.images.len > 0) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); var shared: std.ArrayList(u32) = .empty; if (system_len > 0) for ([_]i64{ @as(i64, @intCast(system_len)) - 2048, @as(i64, @intCast(system_len)) - 512, @intCast(system_len) }) |cut| { if (cut >= 512) try shared.append(a, @intCast(cut)); @@ -137,7 +162,9 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa .history_len = @intCast(rendered.history_len), .rewind_len = @intCast(rendered.rewind_len), .shared_prefixes = shared.items, - .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), // a cut before the conversation's own text: fresh sessions resume their whole harness + .images = input.images, + // a cut just before the conversation's own text: fresh sessions resume their whole harness + .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), .decode_spans = try prompt_mod.replySpans(srv, cx, rendered.ids), // from the tokens alone, for every request: cached states are keyed by tokens .tools_json = if (input.tools.len > 0) try json.stringify(a, .{ .array = @constCast(input.tools) }, .{ .ascii = false }) else "", .decide = if (input.prompt == null and !background) prompt_mod.decideOf(srv, rendered.ids) else null, diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..852b3730b7 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -49,15 +49,16 @@ pub const flags = [_]Flag{ .{ .name = "--api-key-file", .native = true }, .{ .name = "--metrics-open", .kind = .store_true, .native = true }, .{ .name = "--dashboard", .kind = .store_true, .native = true }, - .{ .name = "--vision", .kind = .store_true }, - .{ .name = "--vision-urls", .kind = .store_true }, + .{ .name = "--vision", .kind = .store_true, .native = true }, + .{ .name = "--vision-urls", .kind = .store_true, .native = true }, .{ .name = "--vision-offload", .kind = .store_true }, - .{ .name = "--vision-max-images" }, - .{ .name = "--vision-image-tokens" }, + .{ .name = "--vision-max-images", .native = true }, + .{ .name = "--vision-image-tokens", .native = true }, .{ .name = "--context", .native = true }, .{ .name = "--speed-up", .native = !cuda_build }, .{ .name = "--max-tokens", .native = true }, .{ .name = "--temperature", .native = true }, + .{ .name = "--tool-system", .native = true }, .{ .name = "--top-p", .native = true }, .{ .name = "--top-k", .native = true }, .{ .name = "--min-p", .native = true }, @@ -99,7 +100,7 @@ pub const flags = [_]Flag{ .{ .name = "--master-port", .native = served("--tp") }, // the GPU engine's policy switches as key=value pairs, repeatable .{ .name = "--policy", .kind = .append, .native = served("--policy") }, - .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" } }, + .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" }, .native = served("--kv-dtype"), .native_values = &.{ "int8", "int4" } }, .{ .name = "--prefill-fp8", .kind = .store_true }, .{ .name = "--no-prefill-fp8", .kind = .store_true }, .{ .name = "--precision", .choices = &.{ "checkpoint", "full" } }, @@ -140,6 +141,11 @@ pub const Args = struct { learn_gib: f64 = 32, // disk for learned states, every model and build together max_tokens: i64 = 4096, temperature: ?f64 = null, + vision: bool = false, // accept image inputs (an engine whose family has a vision tower) + vision_urls: bool = false, // also public HTTPS image and video URLs + vision_max_images: i64 = 4, + vision_image_tokens: i64 = 4096, // visual tokens a request's images share + tool_system: ?[]const u8 = null, // an instruction added as the system message of tool requests that send none top_p: ?f64 = null, top_k: ?i64 = null, min_p: ?f64 = null, @@ -170,6 +176,7 @@ pub const Args = struct { master_port: u16 = 29551, policy: []const u8 = "", // every --policy joined by commas cortheon: bool = false, // Cortheon's program, research before building, and cortheon_search/cortheon_fetch run by the server + kv_dtype: []const u8 = "int8", // --kv-dtype (CUDA): the attention cache's codes, int8 or int4 }; /// A usage error's message (argparse's ``error:`` line); the caller exits 2. @@ -265,7 +272,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa } else if (is(name, "--learn-gib")) { out.learn = true; out.learn_gib = try gib(u, a, name, v); - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--tool-system")) out.tool_system = v else if (is(name, "--vision")) out.vision = true else if (is(name, "--vision-urls")) out.vision_urls = true else if (is(name, "--vision-max-images")) out.vision_max_images = try int(u, a, name, v) else if (is(name, "--vision-image-tokens")) out.vision_image_tokens = try int(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -282,9 +289,19 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa fn cudaFlag(a: Allocator, out: *Args, name: []const u8, v: []const u8, u: *Usage) error{ Usage, OutOfMemory }!bool { if (std.mem.eql(u8, name, "--device")) { out.device = std.math.cast(u32, try int(u, a, name, v)) orelse return fail(u, a, "argument --device: a GPU ordinal from 0: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--kv-dtype")) { + out.kv_dtype = v; } else if (std.mem.eql(u8, name, "--segments")) { const n = try int(u, a, name, v); out.segments = if (n >= 1) @intCast(@min(n, std.math.maxInt(u32))) else return fail(u, a, "argument --segments: a count from 1: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--tp")) { + out.tp = if (std.mem.eql(u8, v, "2")) 2 else 1; + } else if (std.mem.eql(u8, name, "--rank")) { + out.rank = if (std.mem.eql(u8, v, "1")) 1 else 0; + } else if (std.mem.eql(u8, name, "--master")) { + out.master = v; + } else if (std.mem.eql(u8, name, "--master-port")) { + out.master_port = std.math.cast(u16, try int(u, a, name, v)) orelse return fail(u, a, "argument --master-port: a TCP port: '{s}'", .{v}); } else return false; return true; } @@ -362,6 +379,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .mtp_drafts = args.mtp_drafts, .mtp_confidence = args.mtp_confidence, .keep = args.checkpoint_slots, + .vision = args.vision, + .kv_bits = if (std.mem.eql(u8, args.kv_dtype, "int4")) 4 else 8, .tp = args.tp, .rank = args.rank, .master = args.master, diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 4b85fca8cd..0cd1b03c6f 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -78,6 +78,10 @@ pub fn main(init: std.process.Init) !u8 { } orelse return fail(problem); defer up.text.deinit(); defer up.engine.close(up.engine.ctx); + if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls + follow(up.engine.ctx) catch |e| return fail(@errorName(e)); + return 0; + } return serve.run(gpa, io, args, .{ .engine = up.engine.engine, .text = up.text.text(), diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig new file mode 100644 index 0000000000..e175033bf6 --- /dev/null +++ b/zig/src/server/media_fetch.zig @@ -0,0 +1,336 @@ +//! Public HTTPS media for --vision-urls, as the Python server fetches it (vision/images_http.py): an https URL on +//! port 443 without credentials or a fragment, its host resolved and every address checked public, a TLS connection +//! verified for the host on an address that was checked, GET with identity encoding, up to 3 redirects (each checked +//! the same way), a declared media type the caller accepts, the bytes within the caller's limit, all inside a deadline. + +const std = @import("std"); +const Io = std.Io; +const net = std.Io.net; + +pub const Error = error{ OutOfMemory, Media }; + +/// The failure's words for the client (the request is refused with them). +pub const Failure = struct { text: []const u8 = "" }; + +pub const max_url_chars = 4096; +pub const max_redirects = 3; +pub const timeout_s = 10; // one download +pub const total_s = 30; // a request's downloads together + +fn fail(f: *Failure, text: []const u8) Error { + f.text = text; + return error.Media; +} + +/// Python ipaddress: is_global and not multicast or reserved, nor 192.0.0.0/24 or 168.63.129.16; no IPv6 address +/// that embeds an IPv4 one (mapped, 6to4, Teredo). +pub fn publicIp(a: net.IpAddress) bool { + switch (a) { + .ip4 => |v| { + const b = v.bytes; + const in = struct { + fn net4(x: [4]u8, base: [4]u8, bits: u5) bool { + const xi = std.mem.readInt(u32, &x, .big); + const bi = std.mem.readInt(u32, &base, .big); + const mask: u32 = if (bits == 0) 0 else ~@as(u32, 0) << @intCast(32 - @as(u6, bits)); + return xi & mask == bi & mask; + } + }.net4; + const blocked = [_]struct { [4]u8, u5 }{ + .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, + .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, + .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, + .{ .{ 224, 0, 0, 0 }, 4 }, .{ .{ 240, 0, 0, 0 }, 4 }, + }; + for (blocked) |blk| if (in(b, blk[0], blk[1])) return false; + if (std.mem.eql(u8, &b, &.{ 168, 63, 129, 16 })) return false; + return true; + }, + .ip6 => |v| { + const b = v.bytes; + // global unicast 2000::/3 only, less documentation (2001:db8::/32), the 2001::/23 IETF block (Teredo + // 2001::/32 among it), 6to4 (2002::/16) + if (b[0] & 0xe0 != 0x20) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] == 0x0d and b[3] == 0xb8) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] < 0x02) return false; + if (b[0] == 0x20 and b[1] == 0x02) return false; + return true; + }, + } +} + +const Url = struct { host: []const u8, target: []const u8 }; + +/// grammar of images_http._url: https, a host, port 443, no credentials, fragment, backslash, whitespace or controls. +fn parseUrl(a: std.mem.Allocator, value: []const u8, f: *Failure) Error!Url { + if (value.len > max_url_chars) return fail(f, "image URL is too long or contains whitespace/control characters"); + for (value) |c| if (c <= 32 or c == 127) return fail(f, "image URL is too long or contains whitespace/control characters"); + const bad = "image URL must be HTTPS on port 443, without credentials or a fragment"; + if (!std.ascii.startsWithIgnoreCase(value, "https://")) return fail(f, bad); + if (std.mem.indexOfScalar(u8, value, '\\') != null or std.mem.indexOfScalar(u8, value, '#') != null) return fail(f, bad); + const rest = value["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const authority = rest[0..end]; + if (std.mem.indexOfScalar(u8, authority, '@') != null) return fail(f, bad); + var host = authority; + if (std.mem.lastIndexOfScalar(u8, authority, ':')) |c| { + if (authority.len > 0 and authority[0] == '[') return fail(f, bad); // a literal IPv6 host: not a name + if (!std.mem.eql(u8, authority[c + 1 ..], "443") and authority[c + 1 ..].len > 0) return fail(f, bad); + host = authority[0..c]; + } + if (host.len == 0 or std.mem.indexOfScalar(u8, host, '%') != null) return fail(f, bad); + for (host) |c| if (c >= 128) return fail(f, bad); // IDNA names: their ASCII (xn--) form + const lower = try std.ascii.allocLowerString(a, std.mem.trimEnd(u8, host, ".")); + for ([_][]const u8{ "localhost", "metadata.google.internal", "instance-data" }) |n| if (std.mem.eql(u8, lower, n)) return fail(f, "image URLs must use public internet hosts"); + // the path and query, percent-quoted as urllib.parse.quote(safe=...) leaves them + var t: std.ArrayList(u8) = .empty; + var tail = rest[end..]; + if (tail.len == 0 or tail[0] == '?') try t.append(a, '/'); + const safe_path = "/%:@!$&'()*+,;=-._~"; + var in_query = false; + for (tail) |c| { + if (c == '?' and !in_query) { + in_query = true; + try t.append(a, '?'); + continue; + } + const safe = std.ascii.isAlphanumeric(c) or std.mem.indexOfScalar(u8, safe_path, c) != null or (in_query and c == '?'); + if (safe) try t.append(a, c) else try t.print(a, "%{X:0>2}", .{c}); + } + tail = t.items; + return .{ .host = lower, .target = tail }; +} + +/// The location a redirect names, resolved against ``base`` (urljoin for absolute and absolute-path forms). +fn join(a: std.mem.Allocator, base: []const u8, location: []const u8) ![]const u8 { + if (std.ascii.startsWithIgnoreCase(location, "https://") or std.ascii.startsWithIgnoreCase(location, "http://")) return a.dupe(u8, location); + if (std.mem.startsWith(u8, location, "//")) return std.fmt.allocPrint(a, "https:{s}", .{location}); + const rest = base["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const origin = base[0 .. "https://".len + end]; + if (std.mem.startsWith(u8, location, "/")) return std.fmt.allocPrint(a, "{s}{s}", .{ origin, location }); + const path_end = std.mem.indexOfScalar(u8, rest[end..], '?') orelse rest.len - end; + const path = rest[end..][0..path_end]; + const dir = if (std.mem.lastIndexOfScalar(u8, path, '/')) |s| path[0 .. s + 1] else "/"; + return std.fmt.allocPrint(a, "{s}{s}{s}", .{ origin, dir, location }); +} + +fn remaining(io: Io, deadline: Io.Timestamp, f: *Failure) Error!i64 { + const left = Io.Clock.awake.now(io).durationTo(deadline).toMilliseconds(); + if (left <= 0) return fail(f, "image download timed out"); + return left; +} + +fn setTimeouts(handle: net.Socket.Handle, ms: i64) void { + const tv: std.posix.timeval = .{ .sec = @intCast(@divTrunc(ms, 1000)), .usec = @intCast(@mod(ms, 1000) * 1000) }; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.RCVTIMEO, std.mem.asBytes(&tv)) catch {}; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.SNDTIMEO, std.mem.asBytes(&tv)) catch {}; +} + +/// A TCP connection to the checked address, its connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds a +/// Linux connect too; std's connect has no timeout yet). +fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + const posix = std.posix; + const linux = std.os.linux; + const family: u32 = switch (addr) { + .ip4 => posix.AF.INET, + .ip6 => posix.AF.INET6, + }; + const rc = linux.socket(family, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (linux.errno(rc) != .SUCCESS) return fail(f, "image download failed or timed out"); + const fd: posix.fd_t = @intCast(rc); + errdefer _ = linux.close(fd); + setTimeouts(fd, ms); + var sa4: posix.sockaddr.in = undefined; + var sa6: posix.sockaddr.in6 = undefined; + const r = switch (addr) { + .ip4 => |v| blk: { + sa4 = .{ .port = std.mem.nativeToBig(u16, 443), .addr = @bitCast(v.bytes) }; + break :blk linux.connect(fd, @ptrCast(&sa4), @sizeOf(posix.sockaddr.in)); + }, + .ip6 => |v| blk: { + sa6 = .{ .port = std.mem.nativeToBig(u16, 443), .flowinfo = 0, .addr = v.bytes, .scope_id = 0 }; + break :blk linux.connect(fd, @ptrCast(&sa6), @sizeOf(posix.sockaddr.in6)); + }, + }; + if (linux.errno(r) != .SUCCESS) return fail(f, "image download failed or timed out"); + return .{ .socket = .{ .handle = fd, .address = addr } }; +} + +/// The bytes of ``url`` (and its media type, one of ``media_types``), at most ``max_bytes``, before ``deadline``. +pub fn fetch(a: std.mem.Allocator, io: Io, url_in: []const u8, max_bytes: usize, deadline_in: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!struct { data: []u8, media: []const u8 } { + var url = url_in; + const deadline = blk: { + const one = Io.Clock.awake.now(io).addDuration(.fromSeconds(timeout_s)); + break :blk if (one.nanoseconds < deadline_in.nanoseconds) one else deadline_in; + }; + var hop: usize = 0; + while (true) : (hop += 1) { + const u = try parseUrl(a, url, f); + const got = try request(a, io, u, max_bytes, deadline, media_types, f); + switch (got) { + .body => |b| return .{ .data = b.data, .media = b.media }, + .redirect => |loc| { + if (hop == max_redirects) return fail(f, "image download has too many redirects"); + url = join(a, url, loc) catch return error.OutOfMemory; + }, + } + } +} + +const Got = union(enum) { body: struct { data: []u8, media: []const u8 }, redirect: []const u8 }; + +var bundle_lock: Io.RwLock = .init; +var bundle: std.crypto.Certificate.Bundle = .empty; +var bundle_loaded = false; + +fn request(a: std.mem.Allocator, io: Io, u: Url, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Got { + // every address the name resolves to must be public; the connection goes to the first of them + const name = net.HostName.init(u.host) catch return fail(f, "image host could not be resolved"); + var results: [32]net.HostName.LookupResult = undefined; + var q: Io.Queue(net.HostName.LookupResult) = .init(&results); + name.lookup(io, &q, .{ .port = 443 }) catch return fail(f, "image host could not be resolved"); + var first: ?net.IpAddress = null; + var count: usize = 0; + while (q.getOneUncancelable(io)) |r| switch (r) { + .address => |addr| { + if (!publicIp(addr)) return fail(f, "image URLs must resolve only to public internet addresses"); + if (first == null) first = addr; + count += 1; + }, + .canonical_name => {}, + } else |_| {} + const addr = first orelse return fail(f, "image host could not be resolved"); + _ = try remaining(io, deadline, f); + // the CA bundle, read once + { + bundle_lock.lockUncancelable(io); + defer bundle_lock.unlock(io); + if (!bundle_loaded) { + bundle.rescan(std.heap.page_allocator, io, Io.Clock.real.now(io)) catch return fail(f, "image download failed: no CA certificates on the server"); + bundle_loaded = true; + } + } + const stream = try connect(addr, try remaining(io, deadline, f), f); + defer stream.close(io); + const tls_len = std.crypto.tls.Client.min_buffer_len; + const bufs = try a.alloc(u8, 4 * tls_len); + var sr = stream.reader(io, bufs[0..tls_len]); + var sw = stream.writer(io, bufs[tls_len .. 2 * tls_len]); + var entropy: [std.crypto.tls.Client.Options.entropy_len]u8 = undefined; + io.random(&entropy); + var tls = std.crypto.tls.Client.init(&sr.interface, &sw.interface, .{ + .host = .{ .explicit = u.host }, + .ca = .{ .bundle = .{ .gpa = std.heap.page_allocator, .io = io, .lock = &bundle_lock, .bundle = &bundle } }, + .read_buffer = bufs[2 * tls_len .. 3 * tls_len], + .write_buffer = bufs[3 * tls_len ..], + .entropy = &entropy, + .realtime_now = Io.Clock.real.now(io), + .allow_truncation_attacks = true, // HTTP's own lengths and framing end the body + }) catch return fail(f, "image download failed: TLS could not be verified for the host"); + // the request + const accept = std.mem.join(a, ", ", media_types) catch return error.OutOfMemory; + tls.writer.print("GET {s} HTTP/1.1\r\nHost: {s}\r\nAccept: {s}\r\nAccept-Encoding: identity\r\nUser-Agent: TensorFold-native\r\nConnection: close\r\n\r\n", .{ u.target, u.host, accept }) catch return fail(f, "image download failed or timed out"); + tls.writer.flush() catch return fail(f, "image download failed or timed out"); + sw.interface.flush() catch return fail(f, "image download failed or timed out"); + const r = &tls.reader; + // status and headers + const status_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + if (status_line.len < 12 or !std.mem.startsWith(u8, status_line, "HTTP/1.")) return fail(f, "image download failed"); + const status = std.fmt.parseInt(u16, status_line[9..12], 10) catch return fail(f, "image download failed"); + var location: ?[]const u8 = null; + var media: []const u8 = ""; + var length: ?usize = null; + var chunked = false; + var encoded = false; + var header_bytes: usize = 0; + while (true) { + const line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + header_bytes += line.len; + if (header_bytes > 64 * 1024) return fail(f, "image download failed"); + const l = std.mem.trimEnd(u8, line, "\r\n"); + if (l.len == 0) break; + const colon = std.mem.indexOfScalar(u8, l, ':') orelse continue; + const key = std.mem.trim(u8, l[0..colon], " \t"); + const val = std.mem.trim(u8, l[colon + 1 ..], " \t"); + if (std.ascii.eqlIgnoreCase(key, "location")) location = try a.dupe(u8, val); + if (std.ascii.eqlIgnoreCase(key, "content-type")) { + const semi = std.mem.indexOfScalar(u8, val, ';') orelse val.len; + media = try std.ascii.allocLowerString(a, std.mem.trim(u8, val[0..semi], " \t")); + } + if (std.ascii.eqlIgnoreCase(key, "content-length")) { + for (val) |c| if (!std.ascii.isDigit(c)) return fail(f, "image response exceeds the encoded byte limit"); + length = std.fmt.parseInt(usize, val, 10) catch return fail(f, "image response exceeds the encoded byte limit"); + } + if (std.ascii.eqlIgnoreCase(key, "transfer-encoding")) chunked = std.mem.indexOf(u8, try std.ascii.allocLowerString(a, val), "chunked") != null; + if (std.ascii.eqlIgnoreCase(key, "content-encoding") and !std.ascii.eqlIgnoreCase(val, "identity")) encoded = true; + } + switch (status) { + 301, 302, 303, 307, 308 => return .{ .redirect = location orelse return fail(f, "image redirect has no destination") }, + 200 => {}, + else => return fail(f, try std.fmt.allocPrint(a, "image download returned HTTP {d}", .{status})), + } + if (encoded) return fail(f, "compressed HTTP image responses are unsupported"); + for (media_types) |m| { + if (std.mem.eql(u8, m, media)) break; + } else return fail(f, if (std.mem.startsWith(u8, media_types[0], "image/")) "image URL content type must be JPEG, PNG or WebP" else try std.fmt.allocPrint(a, "media URL content type must be one of {s}", .{accept})); + if (length) |n| if (n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + var data: std.ArrayList(u8) = .empty; + if (chunked) { + while (true) { + _ = try remaining(io, deadline, f); + const size_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + const semi = std.mem.indexOfScalar(u8, size_line, ';') orelse size_line.len; + const hex = std.mem.trim(u8, size_line[0..semi], " \t\r\n"); + const n = std.fmt.parseInt(usize, hex, 16) catch return fail(f, "image download failed"); + if (n == 0) break; + if (data.items.len + n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + const piece = r.take(n) catch return fail(f, "image download failed or timed out"); + try data.appendSlice(a, piece); + _ = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed"); + } + } else { + while (length == null or data.items.len < length.?) { + _ = try remaining(io, deadline, f); + const want = if (length) |n| @min(n - data.items.len, 64 * 1024) else 64 * 1024; + const piece = r.peekGreedy(1) catch |e| switch (e) { + error.EndOfStream => break, + else => return fail(f, "image download failed or timed out"), + }; + const take = @min(piece.len, want); + if (data.items.len + take > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + try data.appendSlice(a, piece[0..take]); + r.toss(take); + } + if (length) |n| if (data.items.len != n) return fail(f, "image download failed or timed out"); + } + _ = try remaining(io, deadline, f); + return .{ .body = .{ .data = data.items, .media = media } }; +} + +test "public addresses as Python's ipaddress checks them" { + const ip4 = struct { + fn of(b: [4]u8) net.IpAddress { + return .{ .ip4 = .{ .bytes = b, .port = 443 } }; + } + }.of; + try std.testing.expect(publicIp(ip4(.{ 8, 8, 8, 8 }))); + try std.testing.expect(publicIp(ip4(.{ 151, 101, 1, 69 }))); + for ([_][4]u8{ .{ 10, 1, 2, 3 }, .{ 127, 0, 0, 1 }, .{ 169, 254, 169, 254 }, .{ 172, 20, 0, 1 }, .{ 192, 168, 1, 1 }, .{ 100, 100, 0, 1 }, .{ 0, 0, 0, 0 }, .{ 224, 0, 0, 1 }, .{ 255, 255, 255, 255 }, .{ 168, 63, 129, 16 }, .{ 192, 0, 0, 9 } }) |b| + try std.testing.expect(!publicIp(ip4(b))); +} + +test "URLs as images_http._url takes them" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Failure = .{}; + try std.testing.expectError(error.Media, parseUrl(a, "https://example.com/a b/c.png", &f)); + const ok = try parseUrl(a, "https://example.com:443/img/c.png?x=%20&y=ä", &f); + try std.testing.expectEqualStrings("example.com", ok.host); + try std.testing.expectEqualStrings("/img/c.png?x=%20&y=%C3%A4", ok.target); + for ([_][]const u8{ "http://example.com/x.png", "https://user@example.com/x", "https://example.com:8443/x", "https://localhost/x", "https://example.com/x#frag", "https://[::1]/x" }) |bad| + try std.testing.expectError(error.Media, parseUrl(a, bad, &f)); + try std.testing.expectEqualStrings("https://e.com/b/c.png", try join(a, "https://e.com/b/a.png", "c.png")); + try std.testing.expectEqualStrings("https://e.com/z", try join(a, "https://e.com/b/a.png", "/z")); +} diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index b8782aaf96..aaf23610ae 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -127,7 +127,12 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (srv.config.request_log) |path| request_log.append(cx.a, path, body); var input: chat.Input = .{ .fields = undefined }; if (is_chat) { - input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); + var msgs = body.get("messages"); + if (try @import("vision_inputs.zig").extract(srv, cx, msgs)) |x| { // images: decoded, prepared, their parts as text + msgs = x.messages; + input.images = x.images; + } + input.messages = try messages.normalize(cx, msgs, "system", srv.needs_user_after_tool); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..6b7faaf8e2 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -39,6 +39,7 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t error.OutOfMemory => return error.OutOfMemory, error.Template => return cx.fail(.server, "{s}", .{problem}), }; + if (!thinking) ids = try withoutEffortLine(srv, cx, ids); if (!thinking and generation and ids.len > 0) { const last = try srv.text.decode(cx.a, ids[ids.len - 1 ..]); if (std.mem.eql(u8, @import("reply_text.zig").pyStrip(last), "")) { @@ -52,6 +53,19 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t return ids; } +/// GLM-5.3's template writes a reasoning-effort system line whatever the request asks; with thinking off the Python +/// frontend renders it without (glm5_next/prompts.py thinking_off: no effort line, an empty think block). The line is +/// its own segment between special tokens, so its ids alone are the ids it renders to; the first is dropped. +fn withoutEffortLine(srv: *Server, cx: *Cx, ids: []u32) errors.Refused![]u32 { + const line = srv.text.encode(cx.a, "<|system|>Reasoning Effort: Max", false) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return ids, + }; + if (line.len < 2) return ids; + const at = std.mem.indexOf(u32, ids, line) orelse return ids; + return try std.mem.concat(cx.a, u32, &.{ ids[0..at], ids[at + line.len ..] }); +} + /// The template's generation prompt (thinking on or off): what a probe conversation renders past its history. fn generationPrompt(srv: *Server, cx: *Cx, thinking: bool) errors.Refused![]const u32 { const user = try json.newObject(cx.a); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..1bebd524e3 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -109,6 +109,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .model_dir = s.model_dir, .model_ids = ids.items, .default_max_tokens = args.max_tokens, + .vision_max_images = @intCast(@max(args.vision_max_images, 1)), + .vision_urls = args.vision_urls, + .vision_image_tokens = @intCast(std.math.clamp(args.vision_image_tokens, 1, 65536)), + .tool_system = if (args.tool_system) |t| (if (std.mem.trim(u8, t, " \t\r\n").len > 0) std.mem.trim(u8, t, " \t\r\n") else null) else null, .enable_thinking = args.thinking, .reasoning_effort = args.reasoning_effort, .thinking_budget = args.thinking_budget, diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..0407030322 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -36,6 +36,10 @@ pub const Config = struct { /// The served name first, then aliases, without repeats. model_ids: []const []const u8, default_max_tokens: i64 = 4096, + tool_system: ?[]const u8 = null, // --tool-system (trimmed; null when empty) + vision_max_images: usize = 4, // --vision-max-images + vision_urls: bool = false, // --vision-urls: public HTTPS media URLs are fetched (media_fetch.zig) + vision_image_tokens: u64 = 4096, // --vision-image-tokens: the visual tokens a request's images share enable_thinking: bool = true, reasoning_effort: ?[]const u8 = null, thinking_budget: i64 = 0, @@ -135,6 +139,7 @@ pub const Server = struct { if (srv.keepalive) |k| k.stop(); // before the engine's queue goes away if (srv.grammars) |g| g.close(); srv.metrics.deinit(); + if (srv.grammars) |g| g.close(); srv.arena.deinit(); srv.store.deinit(); srv.slide.deinit(); diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig new file mode 100644 index 0000000000..fe457dfb34 --- /dev/null +++ b/zig/src/server/vision_inputs.zig @@ -0,0 +1,342 @@ +//! Image and video inputs of a chat request, as the Python server's frontend takes them (tensorfold/vision/images.py, +//! videos.py, qwen_processing.py): each ``image_url`` part (a data: URL; JPEG or PNG) decoded and prepared, each +//! ``video_url`` part (a data: URL; MP4, WebM...) decoded by libtfvideo and cut into frame groups, and each replaced in +//! its message by the text Python's frontend expands it to: ``<|vision_start|>`` + one ``<|image_pad|>`` a visual +//! token + ``<|vision_end|>`` for an image, and ``<|vision_start|>`` + ``<|video_pad|>``s + ``<|vision_end|>`` +//! for each frame group of a video. The template copies the text, and one tokenization gives Python's ids. + +const std = @import("std"); +const media_fetch = @import("media_fetch.zig"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const Server = @import("server.zig").Server; +const Value = json.Value; +const Cx = errors.Cx; +const qi = api.qwen_image; + +const vision_start = "<|vision_start|>"; +const vision_end = "<|vision_end|>"; +const max_video_bytes: usize = 16 * 1024 * 1024; +const max_videos = 2; +const max_bytes_one: usize = 10 * 1024 * 1024; +const max_bytes_all: usize = 20 * 1024 * 1024; +const max_pixels_all: u64 = 32 * 1024 * 1024; +const max_tokens_one: u64 = 4096; + +pub const Extracted = struct { messages: Value, images: []api.Image }; + +/// The image URL of a content part (an OpenAI ``image_url`` part, or ``image`` / ``input_image`` spellings) and its +/// ``detail``; null for a part that is not an image. +const Kind = enum { image, video }; + +fn mediaPart(part: Value) ?struct { kind: Kind, url: ?[]const u8, detail: []const u8 } { + if (part != .object) return null; + const ty = part.get("type"); + const is_video = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "video_url") or std.mem.eql(u8, ty.?.string, "video"))) or + part.get("video_url") != null; + if (is_video) { + var vurl: ?[]const u8 = null; + if (part.get("video_url")) |vu| switch (vu) { + .string => |x| vurl = x, + .object => if (vu.get("url")) |u| if (u == .string) { + vurl = u.string; + }, + else => {}, + }; + if (vurl == null) if (part.get("video")) |v| if (v == .string) { + vurl = v.string; + }; + return .{ .kind = .video, .url = vurl, .detail = "auto" }; + } + const is_image = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "image_url") or std.mem.eql(u8, ty.?.string, "image") or + std.mem.eql(u8, ty.?.string, "input_image"))) or part.get("image_url") != null or part.get("image") != null; + if (!is_image) return null; + var url: ?[]const u8 = null; + var detail: []const u8 = "auto"; + if (part.get("image_url")) |iu| switch (iu) { + .string => |s| url = s, + .object => { + if (iu.get("url")) |u| if (u == .string) { + url = u.string; + }; + if (iu.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + }, + else => {}, + }; + if (url == null) if (part.get("image")) |im| if (im == .string) { + url = im.string; + }; + if (part.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + return .{ .kind = .image, .url = url, .detail = detail }; +} + +fn hasImages(messages: Value) bool { + if (messages != .array) return false; + for (messages.array) |m| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + for (c.array) |p| if (mediaPart(p) != null) return true; + } + return false; +} + +const image_media = [_][]const u8{ "image/jpeg", "image/png", "image/webp" }; +const video_media = [_][]const u8{ "video/mp4", "video/webm", "video/quicktime", "video/x-matroska" }; + +/// The bytes of a ``data:;base64,`` URL, or (--vision-urls) of a public HTTPS URL, fetched before +/// ``deadline`` (images.py _check_source and load_images). +fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, deadline: std.Io.Timestamp) errors.Refused![]u8 { + if (!std.mem.startsWith(u8, url, "data:")) { + if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse(if (kind == .image) "images require data URLs or public HTTPS URLs" else "videos require data URLs or public HTTPS URLs"); + if (!srv.config.vision_urls) return cx.refuse("image URLs are off on this server; send the image as a data URL, or start the server with --vision-urls"); + if (url.len > media_fetch.max_url_chars) return cx.refuse("image URL is too long"); + var f: media_fetch.Failure = .{}; + const got = media_fetch.fetch(cx.a, srv.io, url, limit, deadline, if (kind == .image) &image_media else &video_media, &f) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Media => return cx.fail(.request, "{s}", .{f.text}), + }; + return got.data; + } + const comma = std.mem.indexOfScalar(u8, url, ',') orelse return cx.refuse("an image data: URL needs a comma before its bytes"); + if (!std.mem.endsWith(u8, url[0..comma], ";base64")) return cx.refuse("an image data: URL must be base64-encoded"); + const b64 = url[comma + 1 ..]; + const dec = std.base64.standard.Decoder; + const n = dec.calcSizeForSlice(b64) catch return cx.refuse("an image's base64 bytes are invalid"); + if (n > limit) return cx.refuse(if (kind == .image) "an image exceeds the 10 MiB limit (20 MiB for a request's media together)" else "a video exceeds the 16 MiB limit (20 MiB for a request's media together)"); + const out = try cx.a.alloc(u8, n); + dec.decode(out, b64) catch return cx.refuse("an image's base64 bytes are invalid"); + return out; +} + +/// A request's images extracted and prepared, its messages with their image parts as the template's image text; +/// null when it has none. A server without --vision refuses images. +pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { + const list = messages orelse return null; + if (!hasImages(list)) return null; + if (!srv.info.vision) return cx.refuse("image and video inputs need a server started with --vision"); + // every media part, in prompt order, with where its text goes + const Slot = struct { kind: Kind, url: ?[]const u8, detail: []const u8, msg: usize, part: usize }; + var slots: std.ArrayList(Slot) = .empty; + for (list.array, 0..) |m, i| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + const role = m.get("role"); + for (c.array, 0..) |p, j| { + const im = mediaPart(p) orelse continue; + if (role != null and role.? == .string and (std.mem.eql(u8, role.?.string, "system") or std.mem.eql(u8, role.?.string, "developer"))) + return cx.refuse("a system message cannot contain images or videos"); + try slots.append(cx.a, .{ .kind = im.kind, .url = im.url, .detail = im.detail, .msg = i, .part = j }); + } + } + var n_images: usize = 0; + var n_videos: usize = 0; + for (slots.items) |sl| switch (sl.kind) { + .image => n_images += 1, + .video => n_videos += 1, + }; + if (n_images > srv.config.vision_max_images) return cx.fail(.request, "a request may carry at most {d} images; this one has {d}", .{ srv.config.vision_max_images, n_images }); + if (n_videos > max_videos) return cx.fail(.request, "a request may carry at most {d} videos; this one has {d}", .{ max_videos, n_videos }); + const budget: u64 = srv.config.vision_image_tokens; + if (n_images > 0 and budget < n_images) return cx.refuse("the image count exceeds the visual-token budget"); + const per: u64 = if (n_images > 0) @min(budget / n_images, max_tokens_one) else 0; + var images: std.ArrayList(api.Image) = .empty; + const texts = try cx.a.alloc([]const u8, slots.items.len); + var total_bytes: usize = 0; + var total_pixels: u64 = 0; + const deadline = std.Io.Clock.awake.now(srv.io).addDuration(.fromSeconds(media_fetch.total_s)); // every URL's download + for (slots.items, 0..) |sl, k| { + const url = sl.url orelse return cx.refuse(if (sl.kind == .image) "an image part needs an image_url with a url" else "a video part needs a video_url with a url"); + const bytes = try dataBytes(srv, cx, url, @min(if (sl.kind == .image) max_bytes_one else max_video_bytes, max_bytes_all - total_bytes), sl.kind, deadline); + total_bytes += bytes.len; + if (total_bytes > max_bytes_all) return cx.refuse("a request's images and videos exceed the 20 MiB limit"); + var text: std.ArrayList(u8) = .empty; + switch (sl.kind) { + .image => { + var img = qi.decode(cx.a, bytes, .{}) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.ImageTooLarge => return cx.refuse("image dimensions exceed the decoded pixel limit"), + else => return cx.refuse("image bytes are invalid or unsupported; use JPEG or PNG"), + }; + total_pixels += @as(u64, img.w) * img.h; + if (total_pixels > max_pixels_all) return cx.refuse("a request's images exceed the decoded pixel limit"); + const cap = if (std.mem.eql(u8, sl.detail, "low")) @min(per, 256) else per; + if (srv.info.vision_glm) { // GLM: its canvas within [min(16, cap), cap] tokens, <|image|> placeholders + const p = qi.glm.prepare(cx.a, img, @min(16, cap), cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the image does not fit the visual-token budget"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + texts[k] = text.items; + continue; + } + const p = qi.prepare(cx.a, img, cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("absolute aspect ratio must be smaller than 200"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, vision_start); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image_pad|>"); + try text.appendSlice(cx.a, vision_end); + }, + .video => if (srv.info.vision_glm) { + // GLM: <|begin_of_video|>, each frame group as <|begin_of_image|> <|image|>s <|end_of_image|> and its + // first frame's "t.t seconds", <|end_of_video|> (mlx-vlm's glm5_next processor) + var v = try videoGlm(srv, cx, bytes); + defer v.deinit(cx.a); + try text.appendSlice(cx.a, "<|begin_of_video|>"); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds"); + } + try text.appendSlice(cx.a, "<|end_of_video|>"); + } else { + var v = try video(srv, cx, bytes); + defer v.deinit(cx.a); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, "<"); + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds>"); + try text.appendSlice(cx.a, vision_start); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|video_pad|>"); + try text.appendSlice(cx.a, vision_end); + } + }, + } + texts[k] = text.items; + } + // the messages with each media part as its text + const out = try cx.a.alloc(Value, list.array.len); + @memcpy(out, list.array); + var k: usize = 0; + while (k < slots.items.len) { + const i = slots.items[k].msg; + const c = list.array[i].get("content").?; + const parts = try cx.a.dupe(Value, c.array); + while (k < slots.items.len and slots.items[k].msg == i) : (k += 1) { + const t = try json.newObject(cx.a); + try t.put(cx.a, "type", .{ .string = "text" }); + try t.put(cx.a, "text", .{ .string = texts[k] }); + parts[slots.items[k].part] = .{ .object = t }; + } + const copy = try json.copyObject(cx.a, list.array[i].object); + try copy.put(cx.a, "content", .{ .array = parts }); + out[i] = .{ .object = copy }; + } + return .{ .messages = .{ .array = out }, .images = images.items }; +} + +// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives -------------------------------------------------- + +const TfVideo = extern struct { rgb: ?[*]u8, frames: c_int, height: c_int, width: c_int, indices: ?[*]c_int, fps: f64 }; +const SizeFn = *const fn (ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int; +const DecodeFn = *const fn (data: [*]const u8, len: usize, rate: f64, min_frames: c_int, max_frames: c_int, max_seconds: f64, max_dim: c_int, size: SizeFn, ctx: ?*anyopaque, out: *TfVideo) callconv(.c) c_int; +const FreeFn = *const fn (v: *TfVideo) callconv(.c) void; + +var lib_mutex: std.Io.Mutex = .init; +var lib: ?struct { decode: DecodeFn, free: FreeFn } = null; + +fn library(io: std.Io) ?@TypeOf(lib.?) { + lib_mutex.lockUncancelable(io); + defer lib_mutex.unlock(io); + if (lib == null) { + var dl = std.DynLib.open("libtfvideo.so") catch return null; + const d = dl.lookup(DecodeFn, "tf_video_decode") orelse return null; + const f = dl.lookup(FreeFn, "tf_video_free") orelse return null; + lib = .{ .decode = d, .free = f }; // kept open for the server's life + } + return lib; +} + +fn sizeFor(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const budget: *const u64 = @ptrCast(@alignCast(ctx.?)); + const s = qi.videoSize(@intCast(frames), @intCast(height), @intCast(width), budget.*) catch return -1; + oh.* = @intCast(s[0]); + ow.* = @intCast(s[1]); + return 0; +} + +/// GLM's video canvas: the frames' content size for libtfvideo (geometry within the budget), the canvas kept. +const GlmSize = struct { budget: u64, th: u32 = 0, tw: u32 = 0 }; + +fn sizeForGlm(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const gs: *GlmSize = @ptrCast(@alignCast(ctx.?)); + const g = qi.glm.geometry(@intCast(frames), @intCast(height), @intCast(width), 16, gs.budget) catch return -1; + gs.th = g[0]; + gs.tw = g[1]; + oh.* = @intCast(g[2]); + ow.* = @intCast(g[3]); + return 0; +} + +fn videoGlm(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.glm.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var gs: GlmSize = .{ .budget = 16384 }; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| gs.budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch gs.budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 2, 256, 3600.0, 8192, sizeForGlm, @ptrCast(&gs), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.glm.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, gs.th, gs.tw, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} + +fn video(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var budget: u64 = 16384; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 4, 256, 3600.0, 8192, sizeFor, @ptrCast(&budget), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} diff --git a/zig/src/vision/glm_image.zig b/zig/src/vision/glm_image.zig new file mode 100644 index 0000000000..b46e75f076 --- /dev/null +++ b/zig/src/vision/glm_image.zig @@ -0,0 +1,194 @@ +//! GLM-5.3-Flash's image and video preparation (mlx-vlm's glm5_next/processing.py, as TensorFold's GLM frontend runs +//! it): the canvas (smart_resize to multiples of 28 within a visual-token budget), the content resized into it (PIL's +//! antialiased bicubic) and zero-padded at the bottom and right, CLIP's normalization (x * f32(1/255) - mean) / std, +//! and the patches in merge-block order ([c, t, 14, 14] each; an image's one frame twice as the temporal pair). + +const std = @import("std"); +const qi = @import("qwen_image.zig"); + +pub const PATCH = 14; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; +const MEAN = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +const STD = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; +const PER = 3 * TEMPORAL * PATCH * PATCH; + +pub const Error = qi.Error; + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const f = @floor(x); + const d = x - f; + if (d > 0.5) return f + 1; + if (d < 0.5) return f; + return if (@mod(f, 2) == 0) f else f + 1; +} + +fn alignUp(v: f64, f: f64) f64 { + return @ceil(v / f) * f; +} + +/// smart_resize: the canvas (height, width) of ``frames`` frames within [min_tokens, max_tokens] merged tokens. +pub fn smartResize(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![2]u32 { + if (frames == 0 or height == 0 or width == 0 or min_tokens == 0 or min_tokens > max_tokens) return error.UnsupportedImage; + const f: f64 = FACTOR; + const ppt: f64 = TEMPORAL * f * f; + const min_px = @as(f64, @floatFromInt(min_tokens)) * ppt; + const max_px = @as(f64, @floatFromInt(max_tokens)) * ppt; + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const n: f64 = @floatFromInt(frames); + const af = @max(@as(f64, TEMPORAL), pyRound(n / TEMPORAL) * TEMPORAL); + var ah = alignUp(h, f); + var aw = alignUp(w, f); + var px = af * ah * aw; + if (px < min_px) { + const scale = @sqrt(min_px / (n * h * w)); + ah = alignUp(@max(1, @ceil(h * scale)), f); + aw = alignUp(@max(1, @ceil(w * scale)), f); + px = af * ah * aw; + } + if (px > max_px) { // the tallest content within the budget, by bisection on its height + if (max_px < af * f * f) return error.UnsupportedImage; + var low: f64 = 1; + var high: f64 = h; + var best: [2]f64 = .{ f, f }; + while (low <= high) { + const ch = @floor((low + high) / 2); + const cw = @max(1, @floor(w * ch / h)); + const cand = [2]f64{ alignUp(ch, f), alignUp(cw, f) }; + if (af * cand[0] * cand[1] <= max_px) { + best = cand; + low = ch + 1; + } else high = ch - 1; + } + ah = best[0]; + aw = best[1]; + } + return .{ @intFromFloat(ah), @intFromFloat(aw) }; +} + +/// _resize_geometry: the canvas (th, tw) and the content (ch, cw) resized into its top-left corner. +pub fn geometry(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![4]u32 { + const t = try smartResize(frames, height, width, min_tokens, max_tokens); + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + var scale = @min(@as(f64, @floatFromInt(t[0])) / h, @as(f64, @floatFromInt(t[1])) / w); + const ppt: f64 = TEMPORAL * FACTOR * FACTOR; + if (@as(f64, @floatFromInt(frames)) * h * w >= @as(f64, @floatFromInt(min_tokens)) * ppt) scale = @min(1.0, scale); + const ch: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[0])), @floor(h * scale)))); + const cw: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[1])), @floor(w * scale)))); + return .{ t[0], t[1], ch, cw }; +} + +pub const Prepared = struct { + patches: []f32, // [gh * gw, 1176] + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], MERGE * MERGE)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +fn norm(v: u8, c: usize) f32 { + const x = @as(f32, @floatFromInt(v)) * @as(f32, @floatCast(@as(f64, 1.0) / 255.0)); + return (x - MEAN[c]) / STD[c]; +} + +/// A frame's normalized value at canvas (y, x) (zero padding past the content's ch x cw). +fn at(frame: []const u8, cw: usize, ch: usize, y: usize, x: usize, c: usize) f32 { + return norm(if (y < ch and x < cw) frame[(y * cw + x) * 3 + c] else 0, c); +} + +/// ``im`` on its canvas within [min_tokens, max_tokens] tokens, normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: qi.Image, min_tokens: u64, max_tokens: u64) Error!Prepared { + const g = try geometry(TEMPORAL, im.h, im.w, min_tokens, max_tokens); + var r = if (g[2] == im.h and g[3] == im.w) null else try qi.resize(gpa, im, g[3], g[2]); + defer if (r) |*x| x.deinit(gpa); + const src = if (r) |x| x.rgb else im.rgb; + const gh: usize = g[0] / PATCH; + const gw: usize = g[1] / PATCH; + const out = try gpa.alloc(f32, gh * gw * PER); + var n: usize = 0; + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..PATCH) |y| for (0..PATCH) |x| { + const v = at(src, g[3], g[2], py + y, px + x, c); + for (0..TEMPORAL) |t| dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = v; + }; + n += 1; + }; + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +/// A video's frame groups: each group's patches (two consecutive frames as the temporal pair; an odd last frame +/// repeats) and its time (its first frame's seconds, as the processor's timestamps[::2]). +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1176 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * PER; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames at the content size (ch x cw: geometry's), placed on the th x tw canvas. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, ch: usize, cw: usize, th: usize, tw: usize, indices: []const i32, fps: f64) Error!Video { + if (th % FACTOR != 0 or tw % FACTOR != 0 or ch > th or cw > tw or frames < 1) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = th / PATCH; + const gw = tw / PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * PER); + errdefer gpa.free(out); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * ch * cw * 3 ..][0 .. ch * cw * 3]; + for (0..PATCH) |y| for (0..PATCH) |x| { + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = at(frame, cw, ch, py + y, px + x, c); + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| times[g] = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +test "smart_resize and geometry match the processor's" { + // 640 x 480 within 8000 tokens: aligned to 504 x 644, content unscaled + try std.testing.expectEqual([4]u32{ 504, 644, 480, 640 }, try geometry(TEMPORAL, 480, 640, 16, 8000)); + // 40 x 30 below 16 tokens: scaled up, content scaled + const g = try geometry(TEMPORAL, 30, 40, 16, 8000); + try std.testing.expectEqual(@as(u32, 112), g[0]); + try std.testing.expectEqual(@as(u32, 140), g[1]); + // 300 x 900 (tall) + const t = try geometry(TEMPORAL, 900, 300, 16, 8000); + try std.testing.expectEqual([2]u32{ 924, 308 }, .{ t[0], t[1] }); +} diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig new file mode 100644 index 0000000000..f918651811 --- /dev/null +++ b/zig/src/vision/qwen_image.zig @@ -0,0 +1,432 @@ +//! Image inputs for Qwen-style vision towers, as the Python engine prepares them (tensorfold/vision/images.py and +//! qwen_processing.py with transformers' Qwen2VL processor): decode (JPEG, PNG; EXIF orientation applied, alpha over +//! white), smart_resize to multiples of patch x merge within a pixel budget, an antialiased bicubic resize (PIL's +//! filter and bounds with torch's uint8 fixed point: int16 weights at the most precision the axis allows), +//! (x - 127.5) / 127.5, and the patches in merge-block order, +//! each [channel, frame (the image twice), 16, 16] -- 1,536 values. + +const std = @import("std"); +/// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. +pub const glm = @import("glm_image.zig"); + +extern fn stbi_info_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int) c_int; +extern fn stbi_load_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int, req: c_int) ?[*]u8; +extern fn stbi_image_free(p: ?*anyopaque) void; + +pub const Limits = struct { + max_dimension: u32 = 8192, + max_pixels: u64 = 16 * 1024 * 1024, +}; + +/// An RGB image, rows top to bottom, three bytes a pixel. +pub const Image = struct { + w: u32, + h: u32, + rgb: []u8, + + pub fn deinit(im: *Image, gpa: std.mem.Allocator) void { + gpa.free(im.rgb); + } +}; + +pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio }; + +/// JPEG or PNG bytes to RGB: EXIF orientation applied (JPEG), any alpha composited over white. +pub fn decode(gpa: std.mem.Allocator, bytes: []const u8, limits: Limits) Error!Image { + if (bytes.len == 0 or bytes.len > std.math.maxInt(c_int)) return error.UnsupportedImage; + var w: c_int = 0; + var h: c_int = 0; + var comp: c_int = 0; + if (stbi_info_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp) == 0) return error.UnsupportedImage; + if (w <= 0 or h <= 0 or w > limits.max_dimension or h > limits.max_dimension or + @as(u64, @intCast(w)) * @as(u64, @intCast(h)) > limits.max_pixels) return error.ImageTooLarge; + const alpha = comp == 2 or comp == 4; + const want: c_int = if (alpha) 4 else 3; + var got: c_int = 0; + const px = stbi_load_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp, want) orelse return error.UnsupportedImage; + defer stbi_image_free(px); + const n: usize = @intCast(w * h); + _ = &got; + var rgb = try gpa.alloc(u8, n * 3); + errdefer gpa.free(rgb); + if (alpha) { + // PIL's paste of RGBA onto white with the alpha as mask: out = white + (src - white) * a / 255, rounded + for (0..n) |i| { + const a: i32 = px[i * 4 + 3]; + for (0..3) |c| { + const s: i32 = px[i * 4 + c]; + const t: i32 = (s - 255) * a + 128; + rgb[i * 3 + c] = @intCast(255 + ((t + (t >> 8)) >> 8)); + } + } + } else @memcpy(rgb, px[0 .. n * 3]); + var im: Image = .{ .w = @intCast(w), .h = @intCast(h), .rgb = rgb }; + const o = orientation(bytes); + if (o > 1) { + const t = try transpose(gpa, im, o); + im.deinit(gpa); + im = t; + } + return im; +} + +/// A JPEG's EXIF orientation (1-8; 1 when absent or not a JPEG). +fn orientation(b: []const u8) u16 { + if (b.len < 4 or b[0] != 0xFF or b[1] != 0xD8) return 1; + var i: usize = 2; + while (i + 4 <= b.len and b[i] == 0xFF) { + const marker = b[i + 1]; + const len = (@as(usize, b[i + 2]) << 8) | b[i + 3]; + if (marker == 0xDA or len < 2 or i + 2 + len > b.len) break; + if (marker == 0xE1 and len >= 16 and std.mem.eql(u8, b[i + 4 .. i + 10], "Exif\x00\x00")) { + const t = b[i + 10 .. i + 2 + len]; + const le = t[0] == 'I'; + const rd16 = struct { + fn f(x: []const u8, at: usize, little: bool) u16 { + return if (little) std.mem.readInt(u16, x[at..][0..2], .little) else std.mem.readInt(u16, x[at..][0..2], .big); + } + }.f; + const rd32 = struct { + fn f(x: []const u8, at: usize, little: bool) u32 { + return if (little) std.mem.readInt(u32, x[at..][0..4], .little) else std.mem.readInt(u32, x[at..][0..4], .big); + } + }.f; + if (t.len < 8) return 1; + const ifd = rd32(t, 4, le); + if (ifd + 2 > t.len) return 1; + const count = rd16(t, ifd, le); + for (0..count) |e| { + const at = ifd + 2 + e * 12; + if (at + 12 > t.len) return 1; + if (rd16(t, at, le) == 0x0112) { + const v = rd16(t, at + 8, le); + return if (v >= 1 and v <= 8) v else 1; + } + } + return 1; + } + i += 2 + len; + } + return 1; +} + +/// PIL's exif_transpose for orientations 2-8. +fn transpose(gpa: std.mem.Allocator, im: Image, o: u16) Error!Image { + const swap = o >= 5; + const W: u32 = if (swap) im.h else im.w; + const H: u32 = if (swap) im.w else im.h; + const out = try gpa.alloc(u8, im.rgb.len); + for (0..H) |y| for (0..W) |x| { + // the source pixel of output (x, y) + const sx: u32, const sy: u32 = switch (o) { + 2 => .{ im.w - 1 - @as(u32, @intCast(x)), @intCast(y) }, + 3 => .{ im.w - 1 - @as(u32, @intCast(x)), im.h - 1 - @as(u32, @intCast(y)) }, + 4 => .{ @intCast(x), im.h - 1 - @as(u32, @intCast(y)) }, + 5 => .{ @intCast(y), @intCast(x) }, + 6 => .{ @intCast(y), im.h - 1 - @as(u32, @intCast(x)) }, + 7 => .{ im.w - 1 - @as(u32, @intCast(y)), im.h - 1 - @as(u32, @intCast(x)) }, + 8 => .{ im.w - 1 - @as(u32, @intCast(y)), @intCast(x) }, + else => .{ @intCast(x), @intCast(y) }, + }; + @memcpy(out[(y * W + x) * 3 ..][0..3], im.rgb[(sy * im.w + sx) * 3 ..][0..3]); + }; + return .{ .w = W, .h = H, .rgb = out }; +} + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const r = @round(x); + if (@abs(x - @trunc(x)) == 0.5) return 2.0 * @round(x / 2.0); + return r; +} + +/// transformers' smart_resize: both sides multiples of ``factor``, the area within [min_pixels, max_pixels]. +pub fn smartResize(height: u32, width: u32, factor: u32, min_pixels: u64, max_pixels: u64) Error![2]u32 { + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const f: f64 = @floatFromInt(factor); + if (@max(h, w) / @min(h, w) > 200) return error.BadAspectRatio; + var hb = pyRound(h / f) * f; + var wb = pyRound(w / f) * f; + if (hb * wb > @as(f64, @floatFromInt(max_pixels))) { + const beta = @sqrt(h * w / @as(f64, @floatFromInt(max_pixels))); + hb = @max(f, @floor(h / beta / f) * f); + wb = @max(f, @floor(w / beta / f) * f); + } else if (hb * wb < @as(f64, @floatFromInt(min_pixels))) { + const beta = @sqrt(@as(f64, @floatFromInt(min_pixels)) / (h * w)); + hb = @ceil(h * beta / f) * f; + wb = @ceil(w * beta / f) * f; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + + +fn bicubic(x0: f64) f64 { + const a = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1; + if (x < 2.0) return (((x - 5) * x + 8) * x - 4) * a; + return 0.0; +} + +const Coeffs = struct { ksize: usize, bounds: []usize, kk: []i32, precision: u5 }; + +/// PIL's precompute_coeffs for an axis of ``in_size`` resampled to ``out_size``, as fixed point the way torch's uint8 +/// antialias kernel takes them: int16 weights at the largest precision (< 22 bits) that keeps the largest in range. +fn coeffs(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Coeffs { + const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); + const filterscale = @max(scale, 1.0); + const support = 2.0 * filterscale; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const bounds = try gpa.alloc(usize, out_size * 2); + const kk = try gpa.alloc(i32, out_size * ksize); + const pre_all = try gpa.alloc(f64, out_size * ksize); + defer gpa.free(pre_all); + for (0..out_size) |xx| { + const pre = pre_all[xx * ksize ..][0..ksize]; + const center = (@as(f64, @floatFromInt(xx)) + 0.5) * scale; + var ww: f64 = 0; + const ss = 1.0 / filterscale; + var xmin_i: i64 = @intFromFloat(center - support + 0.5); + if (xmin_i < 0) xmin_i = 0; + var xmax_i: i64 = @intFromFloat(center + support + 0.5); + if (xmax_i > @as(i64, @intCast(in_size))) xmax_i = @intCast(in_size); + const xmin: usize = @intCast(xmin_i); + const xmax: usize = @intCast(xmax_i - xmin_i); + @memset(pre, 0); + for (0..xmax) |x| { + const wgt = bicubic((@as(f64, @floatFromInt(x + xmin)) - center + 0.5) * ss); + pre[x] = wgt; + ww += wgt; + } + for (0..xmax) |x| { + if (ww != 0.0) pre[x] /= ww; + } + bounds[xx * 2] = xmin; + bounds[xx * 2 + 1] = xmax; + } + var w_max: f64 = 0; + for (pre_all) |v| w_max = @max(w_max, @abs(v)); + var precision: u5 = 0; + while (precision < 22) : (precision += 1) { + const next: i64 = @intFromFloat(0.5 + w_max * @as(f64, @floatFromInt(@as(i64, 1) << (precision + 1)))); + if (next >= (1 << 15)) break; + } + const one: f64 = @floatFromInt(@as(i64, 1) << precision); + for (pre_all, 0..) |v, i| kk[i] = @intFromFloat(if (v < 0) -0.5 + v * one else 0.5 + v * one); + return .{ .ksize = ksize, .bounds = bounds, .kk = kk, .precision = precision }; +} + +fn clip8(v: i64, precision: u5) u8 { + const r = v >> precision; + return @intCast(std.math.clamp(r, 0, 255)); +} + +/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits: the horizontal pass, then the vertical, each +/// rounding to bytes). +pub fn resize(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error!Image { + var cur = im; + var owned = false; + defer if (owned) cur.deinit(gpa); + if (ow != im.w) { + const c = try coeffs(gpa, im.w, ow); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, ow) * im.h * 3); + for (0..im.h) |y| for (0..ow) |x| { + const xmin = c.bounds[x * 2]; + const xmax = c.bounds[x * 2 + 1]; + const k = c.kk[x * c.ksize ..]; + for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..xmax) |j| ss += @as(i64, cur.rgb[(y * cur.w + xmin + j) * 3 + ch]) * k[j]; + out[(y * ow + x) * 3 + ch] = clip8(ss, c.precision); + } + }; + cur = .{ .w = ow, .h = im.h, .rgb = out }; + owned = true; + } + if (oh != cur.h) { + const c = try coeffs(gpa, cur.h, oh); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, cur.w) * oh * 3); + for (0..oh) |y| { + const ymin = c.bounds[y * 2]; + const ymax = c.bounds[y * 2 + 1]; + const k = c.kk[y * c.ksize ..]; + for (0..cur.w) |x| for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..ymax) |j| ss += @as(i64, cur.rgb[((ymin + j) * cur.w + x) * 3 + ch]) * k[j]; + out[(y * cur.w + x) * 3 + ch] = clip8(ss, c.precision); + }; + } + if (owned) cur.deinit(gpa); + cur = .{ .w = cur.w, .h = oh, .rgb = out }; + owned = true; + } + if (!owned) return .{ .w = im.w, .h = im.h, .rgb = try gpa.dupe(u8, im.rgb) }; + owned = false; + return cur; +} + +/// One image prepared for the tower: its patches (fp32 [gh * gw, 1536], merge-block order) and grid [1, gh, gw]. +pub const Prepared = struct { + patches: []f32, + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], 4)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +pub const PATCH = 16; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; + +/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels as the Python frontend passes it: one merged patch), +/// normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepared { + const cap = max_tokens * FACTOR * FACTOR; + const size = try smartResize(im.h, im.w, FACTOR, @min(FACTOR * FACTOR, cap), cap); + var r = try resize(gpa, im, size[1], size[0]); + defer r.deinit(gpa); + const gh: usize = size[0] / PATCH; + const gw: usize = size[1] / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, gh * gw * per); + var n: usize = 0; + var bh: usize = 0; + while (bh < gh / MERGE) : (bh += 1) { + var bw: usize = 0; + while (bw < gw / MERGE) : (bw += 1) { + for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @floatFromInt(r.rgb[((py + y) * r.w + px + x) * 3 + c]); + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 127.5) / 127.5; + }; + n += 1; + }; + } + } + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +// -- videos --------------------------------------------------------------------------------------------------------- + +/// A video's frame size for the tower (the Python frontend's video_size): Qwen3-VL's smart_resize with its per-frame +/// cap (at most 768 tokens a frame group, at least ~134), the whole video within ``budget`` tokens. +pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u32 { + const factor: f64 = FACTOR; + const shortest: f64 = 128 * FACTOR * FACTOR; + const longest: f64 = @as(f64, @floatFromInt(budget)) * TEMPORAL * FACTOR * FACTOR; + const per_frame = @max(@min(768 * FACTOR * FACTOR, @floor(longest / @as(f64, @floatFromInt(@max(1, frames))))), @floor(shortest * 1.05)); + const max_pixels = per_frame * @as(f64, @floatFromInt(frames)); + if (frames < TEMPORAL) return error.UnsupportedImage; + var height: f64 = @floatFromInt(height0); + var width: f64 = @floatFromInt(width0); + if (height < factor or width < factor) { + const scale = @max(factor / height, factor / width); + height = @trunc(height * scale); + width = @trunc(width * scale); + } + if (@max(height, width) / @min(height, width) > 200) return error.BadAspectRatio; + var hb = pyRound(height / factor) * factor; + var wb = pyRound(width / factor) * factor; + const tb = pyRound(@as(f64, @floatFromInt(frames)) / TEMPORAL) * TEMPORAL; + const n: f64 = @floatFromInt(frames); + if (tb * hb * wb > max_pixels) { + const beta = @sqrt(n * height * width / max_pixels); + hb = @max(factor, @floor(height / beta / factor) * factor); + wb = @max(factor, @floor(width / beta / factor) * factor); + } else if (tb * hb * wb < shortest) { + const beta = @sqrt(shortest / (n * height * width)); + hb = @ceil(height * beta / factor) * factor; + wb = @ceil(width * beta / factor) * factor; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + +/// A decoded video's frame groups for the tower: each group's patches (fp32 [gh * gw, 1536], two consecutive frames +/// as the temporal pair; an odd last frame repeats), normalized as numpy does ((x * f32(1/255) - 0.5) / 0.5), and +/// each group's time (the mean of its two frames' seconds). +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1536 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * 3 * TEMPORAL * PATCH * PATCH; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames of ``h`` x ``w`` (already at the tower's size), their source ``indices`` and ``fps``. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, h: usize, w: usize, indices: []const i32, fps: f64) Error!Video { + if (h % FACTOR != 0 or w % FACTOR != 0 or frames < 2) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = h / PATCH; + const gw = w / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * per); + errdefer gpa.free(out); + const scale: f32 = @floatCast(@as(f64, 1.0) / 255.0); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * h * w * 3 ..]; + for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @as(f32, @floatFromInt(frame[((py + y) * w + px + x) * 3 + c])) * scale; + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 0.5) / 0.5; + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| { + const a: f64 = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + const b: f64 = @as(f64, @floatFromInt(indices[@min(g * TEMPORAL + TEMPORAL - 1, frames - 1)])) / fps; + times[g] = (a + b) / 2; + } + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +/// Python's f"{t:.1f}": one decimal, an exact tie rounded to even. +pub fn formatSeconds(buf: []u8, t: f64) ![]const u8 { + const x: f128 = @as(f128, t) * 10; // exact: a double's 53 bits times ten fit f128's 113 + const fl = @floor(x); + const frac = x - fl; + var r = if (frac > 0.5) fl + 1 else fl; + if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; + const v: i64 = @intFromFloat(r); + return std.fmt.bufPrint(buf, "{d}.{d}", .{ @divTrunc(v, 10), @as(u64, @intCast(@mod(v, 10))) }); +} + +test { + _ = glm; +} diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c new file mode 100644 index 0000000000..1049416565 --- /dev/null +++ b/zig/src/vision/tf_video.c @@ -0,0 +1,220 @@ +// libtfvideo: a video's sampled frames for the native server's vision inputs, as the Python engine decodes them +// (tensorfold/vision/videos.py with PyAV): the first video stream; its rate (average, else guessed, else 24) and +// frame count (the stream's, else its duration times the rate); frames sampled at ``rate`` a second over the whole +// video (Qwen3-VL's rule: at least min_frames, at most max_frames, numpy's linspace rounded half to even), decoded +// in order and scaled straight to the size ``size`` picks as RGB the way PyAV's VideoFrame.to_ndarray(rgb24, width, +// height, interpolation="BICUBIC") does it (VideoReformatter: the frame's colour space and range kept, transfer and +// primaries unspecified, sws_scale_frame with SWS_BICUBIC alone). Built against FFmpeg 9 (the FFmpeg PyAV 19.0.1 +// ships, so the frames equal Python's); the server loads it when a request carries a video. +#include +#include +#include +#include +#include +#include +#include +#include + +typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); + +typedef struct { + uint8_t* rgb; // frames x height x width x 3 + int frames, height, width; + int* indices; // each frame's index in the source + double fps; // the source's rate +} tf_video; + +enum { + TF_OK = 0, + TF_INVALID = -1, // not a decodable container or stream + TF_NO_STREAM = -2, // no video stream + TF_DIMENSIONS = -3, // missing or over max_dim + TF_NO_COUNT = -4, // the frame count or rate is missing + TF_TOO_LONG = -5, // over max_seconds + TF_FEW_FRAMES = -6, // fewer than two decodable frames + TF_MEMORY = -7, + TF_SIZE = -8, // the size callback refused +}; + +typedef struct { const uint8_t* data; size_t len, pos; } mem_t; + +static int mem_read(void* o, uint8_t* buf, int n) { + mem_t* m = (mem_t*)o; + if (m->pos >= m->len) return AVERROR_EOF; + size_t k = m->len - m->pos < (size_t)n ? m->len - m->pos : (size_t)n; + memcpy(buf, m->data + m->pos, k); + m->pos += k; + return (int)k; +} + +static int64_t mem_seek(void* o, int64_t off, int whence) { + mem_t* m = (mem_t*)o; + if (whence == AVSEEK_SIZE) return (int64_t)m->len; + int64_t at = whence == SEEK_SET ? off : whence == SEEK_CUR ? (int64_t)m->pos + off : (int64_t)m->len + off; + if (at < 0 || at > (int64_t)m->len) return -1; + m->pos = (size_t)at; + return at; +} + +// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): an SWS_CS_* value names the frame's colour space +// (the frame's own AVColorSpace number passes through it as PyAV passes it) +static void set_colorspace(AVFrame* f, int cs) { + if (cs == SWS_CS_ITU709) f->colorspace = AVCOL_SPC_BT709; + else if (cs == SWS_CS_FCC) f->colorspace = AVCOL_SPC_FCC; + else if (cs == SWS_CS_ITU601) f->colorspace = AVCOL_SPC_SMPTE170M; + else if (cs == SWS_CS_SMPTE240M) f->colorspace = AVCOL_SPC_SMPTE240M; + else if (cs == SWS_CS_BT2020) f->colorspace = AVCOL_SPC_BT2020_NCL; +} + +static double round_half_even(double x) { + double r = round(x); + if (fabs(x - trunc(x)) == 0.5) r = 2.0 * round(x / 2.0); + return r; +} + +void tf_video_free(tf_video* v) { + if (!v) return; + free(v->rgb); + free(v->indices); + v->rgb = NULL; + v->indices = NULL; +} + +int tf_video_decode(const uint8_t* data, size_t len, double rate, int min_frames, int max_frames, double max_seconds, + int max_dim, tf_size_fn size, void* ctx, tf_video* out) { + memset(out, 0, sizeof(*out)); + int rc = TF_INVALID; + mem_t mem = {data, len, 0}; + AVFormatContext* fmt = NULL; + AVIOContext* io = NULL; + AVCodecContext* dec = NULL; + AVPacket* pkt = NULL; + AVFrame* frame = NULL; + struct SwsContext* sws = NULL; + int* wanted = NULL; + unsigned char* iobuf = av_malloc(1 << 16); + if (!iobuf) return TF_MEMORY; + io = avio_alloc_context(iobuf, 1 << 16, 0, &mem, mem_read, NULL, mem_seek); + if (!io) { av_free(iobuf); return TF_MEMORY; } + fmt = avformat_alloc_context(); + if (!fmt) { rc = TF_MEMORY; goto done; } + fmt->pb = io; + if (avformat_open_input(&fmt, NULL, NULL, NULL) < 0) { fmt = NULL; goto done; } + if (avformat_find_stream_info(fmt, NULL) < 0) goto done; + int si = av_find_best_stream(fmt, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0); + if (si < 0) { rc = TF_NO_STREAM; goto done; } + // PyAV's streams.video[0] is the first video stream + for (unsigned i = 0; i < fmt->nb_streams; ++i) + if (fmt->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) { si = (int)i; break; } + AVStream* st = fmt->streams[si]; + int width = st->codecpar->width, height = st->codecpar->height; + if (width <= 0 || height <= 0 || width > max_dim || height > max_dim) { rc = TF_DIMENSIONS; goto done; } + AVRational r = st->avg_frame_rate; + if (r.num <= 0 || r.den <= 0) r = av_guess_frame_rate(fmt, st, NULL); + double fps = (r.num > 0 && r.den > 0) ? av_q2d(r) : 24.0; + double seconds = 0.0; + if (st->duration > 0 && st->duration != AV_NOPTS_VALUE) seconds = (double)st->duration * av_q2d(st->time_base); + else if (fmt->duration > 0 && fmt->duration != AV_NOPTS_VALUE) seconds = (double)fmt->duration / 1e6; + int64_t total = st->nb_frames > 0 ? st->nb_frames : (int64_t)round_half_even(seconds * fps); + if (!(fps > 0 && fps <= 1000) || total <= 0) { rc = TF_NO_COUNT; goto done; } + if ((double)total / fps > max_seconds) { rc = TF_TOO_LONG; goto done; } + int64_t count = (int64_t)((double)total / fps * rate); + if (count < min_frames) count = min_frames; + if (count > max_frames) count = max_frames; + if (count > total) count = total; + wanted = malloc(sizeof(int) * (size_t)count); + if (!wanted) { rc = TF_MEMORY; goto done; } + // numpy.linspace(0, total - 1, count): step = (stop - start) / (count - 1), the last exactly stop + double stop = (double)(total - 1); + for (int64_t i = 0; i < count; ++i) { + double v = count > 1 ? (double)i * (stop / (double)(count - 1)) : 0.0; + if (count > 1 && i == count - 1) v = stop; + wanted[i] = (int)round_half_even(v); + } + int oh = 0, ow = 0; + if (size(ctx, (int)count, height, width, &oh, &ow) != 0 || oh <= 0 || ow <= 0) { rc = TF_SIZE; goto done; } + const AVCodec* codec = avcodec_find_decoder(st->codecpar->codec_id); + if (!codec) goto done; + dec = avcodec_alloc_context3(codec); + if (!dec || avcodec_parameters_to_context(dec, st->codecpar) < 0) goto done; + dec->thread_count = 0; // PyAV's thread_type "AUTO" + dec->thread_type = FF_THREAD_FRAME | FF_THREAD_SLICE; + if (avcodec_open2(dec, codec, NULL) < 0) goto done; + out->rgb = malloc((size_t)count * (size_t)oh * (size_t)ow * 3); + out->indices = malloc(sizeof(int) * (size_t)count); + pkt = av_packet_alloc(); + frame = av_frame_alloc(); + if (!out->rgb || !out->indices || !pkt || !frame) { rc = TF_MEMORY; goto done; } + AVFrame* rgb = av_frame_alloc(); + if (!rgb) { rc = TF_MEMORY; goto done; } + sws = sws_alloc_context(); + if (!sws) { av_frame_free(&rgb); rc = TF_MEMORY; goto done; } + sws->threads = 0; + sws->flags = SWS_BICUBIC; + int n = 0, next = 0, got = 0, last = wanted[count - 1]; + int flushing = 0, finished = 0; + while (!finished) { + if (!flushing) { + int e = av_read_frame(fmt, pkt); + if (e < 0) { flushing = 1; avcodec_send_packet(dec, NULL); } + else { + if (pkt->stream_index == si) avcodec_send_packet(dec, pkt); + av_packet_unref(pkt); + } + } + for (;;) { + int e = avcodec_receive_frame(dec, frame); + if (e == AVERROR(EAGAIN)) { + if (flushing) finished = 1; // a flushed decoder only ends + break; + } + if (e < 0) { finished = 1; break; } + if (n > last) { finished = 1; av_frame_unref(frame); break; } + if (next < count && n == wanted[next]) { + // PyAV's _reformat: the destination copies the frame's properties; its colour space and range pass + // through _set_frame_colorspace with the frame's own values; transfer and primaries unspecified + av_frame_unref(rgb); + av_frame_copy_props(rgb, frame); + set_colorspace(frame, (int)frame->colorspace); + rgb->colorspace = frame->colorspace; + rgb->color_range = frame->color_range; + set_colorspace(rgb, (int)frame->colorspace); + const enum AVColorTransferCharacteristic trc = frame->color_trc; + const enum AVColorPrimaries pri = frame->color_primaries; + frame->color_trc = rgb->color_trc = AVCOL_TRC_UNSPECIFIED; + frame->color_primaries = rgb->color_primaries = AVCOL_PRI_UNSPECIFIED; + rgb->format = AV_PIX_FMT_RGB24; + rgb->width = ow; + rgb->height = oh; + if (av_frame_get_buffer(rgb, 0) < 0) { rc = TF_MEMORY; av_frame_unref(frame); av_frame_free(&rgb); goto done; } + int e2 = sws_scale_frame(sws, rgb, frame); + frame->color_trc = trc; + frame->color_primaries = pri; + if (e2 < 0) { av_frame_unref(frame); av_frame_free(&rgb); goto done; } + uint8_t* to = out->rgb + (size_t)got * oh * ow * 3; + for (int y = 0; y < oh; ++y) memcpy(to + (size_t)y * ow * 3, rgb->data[0] + (size_t)y * rgb->linesize[0], (size_t)ow * 3); + out->indices[got++] = n; + while (next < count && wanted[next] == n) ++next; + } + ++n; + av_frame_unref(frame); + } + } + av_frame_free(&rgb); + if (got < 2) { rc = TF_FEW_FRAMES; goto done; } + out->frames = got; + out->height = oh; + out->width = ow; + out->fps = fps; + rc = TF_OK; +done: + if (rc != TF_OK) tf_video_free(out); + free(wanted); + sws_freeContext(sws); + av_frame_free(&frame); + av_packet_free(&pkt); + avcodec_free_context(&dec); + if (fmt) avformat_close_input(&fmt); + if (io) { av_freep(&io->buffer); avio_context_free(&io); } + return rc; +} diff --git a/zig/vendor/stb/README.md b/zig/vendor/stb/README.md new file mode 100644 index 0000000000..11466f4316 --- /dev/null +++ b/zig/vendor/stb/README.md @@ -0,0 +1,3 @@ +stb_image.h v2.30 from https://github.com/nothings/stb at commit 2c980bb59875b0d32144a71867fbdebb2f77cd20 +(sha256 594c2fe35d49488b4382dbfaec8f98366defca819d916ac95becf3e75f4200b3), unmodified; dual public domain / MIT +(see the license at the end of the file). stb_image_impl.c compiles it for JPEG and PNG only. diff --git a/zig/vendor/stb/stb_image.h b/zig/vendor/stb/stb_image.h new file mode 100644 index 0000000000..9eedabedc4 --- /dev/null +++ b/zig/vendor/stb/stb_image.h @@ -0,0 +1,7988 @@ +/* stb_image - v2.30 - public domain image loader - http://nothings.org/stb + no warranty implied; use at your own risk + + Do this: + #define STB_IMAGE_IMPLEMENTATION + before you include this file in *one* C or C++ file to create the implementation. + + // i.e. it should look like this: + #include ... + #include ... + #include ... + #define STB_IMAGE_IMPLEMENTATION + #include "stb_image.h" + + You can #define STBI_ASSERT(x) before the #include to avoid using assert.h. + And #define STBI_MALLOC, STBI_REALLOC, and STBI_FREE to avoid using malloc,realloc,free + + + QUICK NOTES: + Primarily of interest to game developers and other people who can + avoid problematic images and only need the trivial interface + + JPEG baseline & progressive (12 bpc/arithmetic not supported, same as stock IJG lib) + PNG 1/2/4/8/16-bit-per-channel + + TGA (not sure what subset, if a subset) + BMP non-1bpp, non-RLE + PSD (composited view only, no extra channels, 8/16 bit-per-channel) + + GIF (*comp always reports as 4-channel) + HDR (radiance rgbE format) + PIC (Softimage PIC) + PNM (PPM and PGM binary only) + + Animated GIF still needs a proper API, but here's one way to do it: + http://gist.github.com/urraka/685d9a6340b26b830d49 + + - decode from memory or through FILE (define STBI_NO_STDIO to remove code) + - decode from arbitrary I/O callbacks + - SIMD acceleration on x86/x64 (SSE2) and ARM (NEON) + + Full documentation under "DOCUMENTATION" below. + + +LICENSE + + See end of file for license information. + +RECENT REVISION HISTORY: + + 2.30 (2024-05-31) avoid erroneous gcc warning + 2.29 (2023-05-xx) optimizations + 2.28 (2023-01-29) many error fixes, security errors, just tons of stuff + 2.27 (2021-07-11) document stbi_info better, 16-bit PNM support, bug fixes + 2.26 (2020-07-13) many minor fixes + 2.25 (2020-02-02) fix warnings + 2.24 (2020-02-02) fix warnings; thread-local failure_reason and flip_vertically + 2.23 (2019-08-11) fix clang static analysis warning + 2.22 (2019-03-04) gif fixes, fix warnings + 2.21 (2019-02-25) fix typo in comment + 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs + 2.19 (2018-02-11) fix warning + 2.18 (2018-01-30) fix warnings + 2.17 (2018-01-29) bugfix, 1-bit BMP, 16-bitness query, fix warnings + 2.16 (2017-07-23) all functions have 16-bit variants; optimizations; bugfixes + 2.15 (2017-03-18) fix png-1,2,4; all Imagenet JPGs; no runtime SSE detection on GCC + 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs + 2.13 (2016-12-04) experimental 16-bit API, only for PNG so far; fixes + 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes + 2.11 (2016-04-02) 16-bit PNGS; enable SSE2 in non-gcc x64 + RGB-format JPEG; remove white matting in PSD; + allocate large structures on the stack; + correct channel count for PNG & BMP + 2.10 (2016-01-22) avoid warning introduced in 2.09 + 2.09 (2016-01-16) 16-bit TGA; comments in PNM files; STBI_REALLOC_SIZED + + See end of file for full revision history. + + + ============================ Contributors ========================= + + Image formats Extensions, features + Sean Barrett (jpeg, png, bmp) Jetro Lauha (stbi_info) + Nicolas Schulz (hdr, psd) Martin "SpartanJ" Golini (stbi_info) + Jonathan Dummer (tga) James "moose2000" Brown (iPhone PNG) + Jean-Marc Lienher (gif) Ben "Disch" Wenger (io callbacks) + Tom Seddon (pic) Omar Cornut (1/2/4-bit PNG) + Thatcher Ulrich (psd) Nicolas Guillemot (vertical flip) + Ken Miller (pgm, ppm) Richard Mitton (16-bit PSD) + github:urraka (animated gif) Junggon Kim (PNM comments) + Christopher Forseth (animated gif) Daniel Gibson (16-bit TGA) + socks-the-fox (16-bit PNG) + Jeremy Sawicki (handle all ImageNet JPGs) + Optimizations & bugfixes Mikhail Morozov (1-bit BMP) + Fabian "ryg" Giesen Anael Seghezzi (is-16-bit query) + Arseny Kapoulkine Simon Breuss (16-bit PNM) + John-Mark Allen + Carmelo J Fdez-Aguera + + Bug & warning fixes + Marc LeBlanc David Woo Guillaume George Martins Mozeiko + Christpher Lloyd Jerry Jansson Joseph Thomson Blazej Dariusz Roszkowski + Phil Jordan Dave Moore Roy Eltham + Hayaki Saito Nathan Reed Won Chun + Luke Graham Johan Duparc Nick Verigakis the Horde3D community + Thomas Ruf Ronny Chevalier github:rlyeh + Janez Zemva John Bartholomew Michal Cichon github:romigrou + Jonathan Blow Ken Hamada Tero Hanninen github:svdijk + Eugene Golushkov Laurent Gomila Cort Stratton github:snagar + Aruelien Pocheville Sergio Gonzalez Thibault Reuille github:Zelex + Cass Everitt Ryamond Barbiero github:grim210 + Paul Du Bois Engin Manap Aldo Culquicondor github:sammyhw + Philipp Wiesemann Dale Weiler Oriol Ferrer Mesia github:phprus + Josh Tobin Neil Bickford Matthew Gregan github:poppolopoppo + Julian Raschke Gregory Mullen Christian Floisand github:darealshinji + Baldur Karlsson Kevin Schmidt JR Smith github:Michaelangel007 + Brad Weinberger Matvey Cherevko github:mosra + Luca Sas Alexander Veselov Zack Middleton [reserved] + Ryan C. Gordon [reserved] [reserved] + DO NOT ADD YOUR NAME HERE + + Jacko Dirks + + To add your name to the credits, pick a random blank space in the middle and fill it. + 80% of merge conflicts on stb PRs are due to people adding their name at the end + of the credits. +*/ + +#ifndef STBI_INCLUDE_STB_IMAGE_H +#define STBI_INCLUDE_STB_IMAGE_H + +// DOCUMENTATION +// +// Limitations: +// - no 12-bit-per-channel JPEG +// - no JPEGs with arithmetic coding +// - GIF always returns *comp=4 +// +// Basic usage (see HDR discussion below for HDR usage): +// int x,y,n; +// unsigned char *data = stbi_load(filename, &x, &y, &n, 0); +// // ... process data if not NULL ... +// // ... x = width, y = height, n = # 8-bit components per pixel ... +// // ... replace '0' with '1'..'4' to force that many components per pixel +// // ... but 'n' will always be the number that it would have been if you said 0 +// stbi_image_free(data); +// +// Standard parameters: +// int *x -- outputs image width in pixels +// int *y -- outputs image height in pixels +// int *channels_in_file -- outputs # of image components in image file +// int desired_channels -- if non-zero, # of image components requested in result +// +// The return value from an image loader is an 'unsigned char *' which points +// to the pixel data, or NULL on an allocation failure or if the image is +// corrupt or invalid. The pixel data consists of *y scanlines of *x pixels, +// with each pixel consisting of N interleaved 8-bit components; the first +// pixel pointed to is top-left-most in the image. There is no padding between +// image scanlines or between pixels, regardless of format. The number of +// components N is 'desired_channels' if desired_channels is non-zero, or +// *channels_in_file otherwise. If desired_channels is non-zero, +// *channels_in_file has the number of components that _would_ have been +// output otherwise. E.g. if you set desired_channels to 4, you will always +// get RGBA output, but you can check *channels_in_file to see if it's trivially +// opaque because e.g. there were only 3 channels in the source image. +// +// An output image with N components has the following components interleaved +// in this order in each pixel: +// +// N=#comp components +// 1 grey +// 2 grey, alpha +// 3 red, green, blue +// 4 red, green, blue, alpha +// +// If image loading fails for any reason, the return value will be NULL, +// and *x, *y, *channels_in_file will be unchanged. The function +// stbi_failure_reason() can be queried for an extremely brief, end-user +// unfriendly explanation of why the load failed. Define STBI_NO_FAILURE_STRINGS +// to avoid compiling these strings at all, and STBI_FAILURE_USERMSG to get slightly +// more user-friendly ones. +// +// Paletted PNG, BMP, GIF, and PIC images are automatically depalettized. +// +// To query the width, height and component count of an image without having to +// decode the full file, you can use the stbi_info family of functions: +// +// int x,y,n,ok; +// ok = stbi_info(filename, &x, &y, &n); +// // returns ok=1 and sets x, y, n if image is a supported format, +// // 0 otherwise. +// +// Note that stb_image pervasively uses ints in its public API for sizes, +// including sizes of memory buffers. This is now part of the API and thus +// hard to change without causing breakage. As a result, the various image +// loaders all have certain limits on image size; these differ somewhat +// by format but generally boil down to either just under 2GB or just under +// 1GB. When the decoded image would be larger than this, stb_image decoding +// will fail. +// +// Additionally, stb_image will reject image files that have any of their +// dimensions set to a larger value than the configurable STBI_MAX_DIMENSIONS, +// which defaults to 2**24 = 16777216 pixels. Due to the above memory limit, +// the only way to have an image with such dimensions load correctly +// is for it to have a rather extreme aspect ratio. Either way, the +// assumption here is that such larger images are likely to be malformed +// or malicious. If you do need to load an image with individual dimensions +// larger than that, and it still fits in the overall size limit, you can +// #define STBI_MAX_DIMENSIONS on your own to be something larger. +// +// =========================================================================== +// +// UNICODE: +// +// If compiling for Windows and you wish to use Unicode filenames, compile +// with +// #define STBI_WINDOWS_UTF8 +// and pass utf8-encoded filenames. Call stbi_convert_wchar_to_utf8 to convert +// Windows wchar_t filenames to utf8. +// +// =========================================================================== +// +// Philosophy +// +// stb libraries are designed with the following priorities: +// +// 1. easy to use +// 2. easy to maintain +// 3. good performance +// +// Sometimes I let "good performance" creep up in priority over "easy to maintain", +// and for best performance I may provide less-easy-to-use APIs that give higher +// performance, in addition to the easy-to-use ones. Nevertheless, it's important +// to keep in mind that from the standpoint of you, a client of this library, +// all you care about is #1 and #3, and stb libraries DO NOT emphasize #3 above all. +// +// Some secondary priorities arise directly from the first two, some of which +// provide more explicit reasons why performance can't be emphasized. +// +// - Portable ("ease of use") +// - Small source code footprint ("easy to maintain") +// - No dependencies ("ease of use") +// +// =========================================================================== +// +// I/O callbacks +// +// I/O callbacks allow you to read from arbitrary sources, like packaged +// files or some other source. Data read from callbacks are processed +// through a small internal buffer (currently 128 bytes) to try to reduce +// overhead. +// +// The three functions you must define are "read" (reads some bytes of data), +// "skip" (skips some bytes of data), "eof" (reports if the stream is at the end). +// +// =========================================================================== +// +// SIMD support +// +// The JPEG decoder will try to automatically use SIMD kernels on x86 when +// supported by the compiler. For ARM Neon support, you must explicitly +// request it. +// +// (The old do-it-yourself SIMD API is no longer supported in the current +// code.) +// +// On x86, SSE2 will automatically be used when available based on a run-time +// test; if not, the generic C versions are used as a fall-back. On ARM targets, +// the typical path is to have separate builds for NEON and non-NEON devices +// (at least this is true for iOS and Android). Therefore, the NEON support is +// toggled by a build flag: define STBI_NEON to get NEON loops. +// +// If for some reason you do not want to use any of SIMD code, or if +// you have issues compiling it, you can disable it entirely by +// defining STBI_NO_SIMD. +// +// =========================================================================== +// +// HDR image support (disable by defining STBI_NO_HDR) +// +// stb_image supports loading HDR images in general, and currently the Radiance +// .HDR file format specifically. You can still load any file through the existing +// interface; if you attempt to load an HDR file, it will be automatically remapped +// to LDR, assuming gamma 2.2 and an arbitrary scale factor defaulting to 1; +// both of these constants can be reconfigured through this interface: +// +// stbi_hdr_to_ldr_gamma(2.2f); +// stbi_hdr_to_ldr_scale(1.0f); +// +// (note, do not use _inverse_ constants; stbi_image will invert them +// appropriately). +// +// Additionally, there is a new, parallel interface for loading files as +// (linear) floats to preserve the full dynamic range: +// +// float *data = stbi_loadf(filename, &x, &y, &n, 0); +// +// If you load LDR images through this interface, those images will +// be promoted to floating point values, run through the inverse of +// constants corresponding to the above: +// +// stbi_ldr_to_hdr_scale(1.0f); +// stbi_ldr_to_hdr_gamma(2.2f); +// +// Finally, given a filename (or an open file or memory block--see header +// file for details) containing image data, you can query for the "most +// appropriate" interface to use (that is, whether the image is HDR or +// not), using: +// +// stbi_is_hdr(char *filename); +// +// =========================================================================== +// +// iPhone PNG support: +// +// We optionally support converting iPhone-formatted PNGs (which store +// premultiplied BGRA) back to RGB, even though they're internally encoded +// differently. To enable this conversion, call +// stbi_convert_iphone_png_to_rgb(1). +// +// Call stbi_set_unpremultiply_on_load(1) as well to force a divide per +// pixel to remove any premultiplied alpha *only* if the image file explicitly +// says there's premultiplied data (currently only happens in iPhone images, +// and only if iPhone convert-to-rgb processing is on). +// +// =========================================================================== +// +// ADDITIONAL CONFIGURATION +// +// - You can suppress implementation of any of the decoders to reduce +// your code footprint by #defining one or more of the following +// symbols before creating the implementation. +// +// STBI_NO_JPEG +// STBI_NO_PNG +// STBI_NO_BMP +// STBI_NO_PSD +// STBI_NO_TGA +// STBI_NO_GIF +// STBI_NO_HDR +// STBI_NO_PIC +// STBI_NO_PNM (.ppm and .pgm) +// +// - You can request *only* certain decoders and suppress all other ones +// (this will be more forward-compatible, as addition of new decoders +// doesn't require you to disable them explicitly): +// +// STBI_ONLY_JPEG +// STBI_ONLY_PNG +// STBI_ONLY_BMP +// STBI_ONLY_PSD +// STBI_ONLY_TGA +// STBI_ONLY_GIF +// STBI_ONLY_HDR +// STBI_ONLY_PIC +// STBI_ONLY_PNM (.ppm and .pgm) +// +// - If you use STBI_NO_PNG (or _ONLY_ without PNG), and you still +// want the zlib decoder to be available, #define STBI_SUPPORT_ZLIB +// +// - If you define STBI_MAX_DIMENSIONS, stb_image will reject images greater +// than that size (in either width or height) without further processing. +// This is to let programs in the wild set an upper bound to prevent +// denial-of-service attacks on untrusted data, as one could generate a +// valid image of gigantic dimensions and force stb_image to allocate a +// huge block of memory and spend disproportionate time decoding it. By +// default this is set to (1 << 24), which is 16777216, but that's still +// very big. + +#ifndef STBI_NO_STDIO +#include +#endif // STBI_NO_STDIO + +#define STBI_VERSION 1 + +enum +{ + STBI_default = 0, // only used for desired_channels + + STBI_grey = 1, + STBI_grey_alpha = 2, + STBI_rgb = 3, + STBI_rgb_alpha = 4 +}; + +#include +typedef unsigned char stbi_uc; +typedef unsigned short stbi_us; + +#ifdef __cplusplus +extern "C" { +#endif + +#ifndef STBIDEF +#ifdef STB_IMAGE_STATIC +#define STBIDEF static +#else +#define STBIDEF extern +#endif +#endif + +////////////////////////////////////////////////////////////////////////////// +// +// PRIMARY API - works on images of any type +// + +// +// load image by filename, open file, or memory buffer +// + +typedef struct +{ + int (*read) (void *user,char *data,int size); // fill 'data' with 'size' bytes. return number of bytes actually read + void (*skip) (void *user,int n); // skip the next 'n' bytes, or 'unget' the last -n bytes if negative + int (*eof) (void *user); // returns nonzero if we are at end of file/data +} stbi_io_callbacks; + +//////////////////////////////////// +// +// 8-bits-per-channel interface +// + +STBIDEF stbi_uc *stbi_load_from_memory (stbi_uc const *buffer, int len , int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk , void *user, int *x, int *y, int *channels_in_file, int desired_channels); + +#ifndef STBI_NO_STDIO +STBIDEF stbi_uc *stbi_load (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_uc *stbi_load_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); +// for stbi_load_from_file, file pointer is left pointing immediately after image +#endif + +#ifndef STBI_NO_GIF +STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp); +#endif + +#ifdef STBI_WINDOWS_UTF8 +STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input); +#endif + +//////////////////////////////////// +// +// 16-bits-per-channel interface +// + +STBIDEF stbi_us *stbi_load_16_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); + +#ifndef STBI_NO_STDIO +STBIDEF stbi_us *stbi_load_16 (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_us *stbi_load_from_file_16(FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); +#endif + +//////////////////////////////////// +// +// float-per-channel interface +// +#ifndef STBI_NO_LINEAR + STBIDEF float *stbi_loadf_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); + STBIDEF float *stbi_loadf_from_callbacks (stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); + + #ifndef STBI_NO_STDIO + STBIDEF float *stbi_loadf (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); + STBIDEF float *stbi_loadf_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); + #endif +#endif + +#ifndef STBI_NO_HDR + STBIDEF void stbi_hdr_to_ldr_gamma(float gamma); + STBIDEF void stbi_hdr_to_ldr_scale(float scale); +#endif // STBI_NO_HDR + +#ifndef STBI_NO_LINEAR + STBIDEF void stbi_ldr_to_hdr_gamma(float gamma); + STBIDEF void stbi_ldr_to_hdr_scale(float scale); +#endif // STBI_NO_LINEAR + +// stbi_is_hdr is always defined, but always returns false if STBI_NO_HDR +STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user); +STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len); +#ifndef STBI_NO_STDIO +STBIDEF int stbi_is_hdr (char const *filename); +STBIDEF int stbi_is_hdr_from_file(FILE *f); +#endif // STBI_NO_STDIO + + +// get a VERY brief reason for failure +// on most compilers (and ALL modern mainstream compilers) this is threadsafe +STBIDEF const char *stbi_failure_reason (void); + +// free the loaded image -- this is just free() +STBIDEF void stbi_image_free (void *retval_from_stbi_load); + +// get image dimensions & components without fully decoding +STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp); +STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp); +STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len); +STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *clbk, void *user); + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_info (char const *filename, int *x, int *y, int *comp); +STBIDEF int stbi_info_from_file (FILE *f, int *x, int *y, int *comp); +STBIDEF int stbi_is_16_bit (char const *filename); +STBIDEF int stbi_is_16_bit_from_file(FILE *f); +#endif + + + +// for image formats that explicitly notate that they have premultiplied alpha, +// we just return the colors as stored in the file. set this flag to force +// unpremultiplication. results are undefined if the unpremultiply overflow. +STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply); + +// indicate whether we should process iphone images back to canonical format, +// or just pass them through "as-is" +STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert); + +// flip the image vertically, so the first pixel in the output array is the bottom left +STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip); + +// as above, but only applies to images loaded on the thread that calls the function +// this function is only available if your compiler supports thread-local variables; +// calling it will fail to link if your compiler doesn't +STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply); +STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert); +STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip); + +// ZLIB client - used by PNG, available for other purposes + +STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen); +STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header); +STBIDEF char *stbi_zlib_decode_malloc(const char *buffer, int len, int *outlen); +STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); + +STBIDEF char *stbi_zlib_decode_noheader_malloc(const char *buffer, int len, int *outlen); +STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); + + +#ifdef __cplusplus +} +#endif + +// +// +//// end header file ///////////////////////////////////////////////////// +#endif // STBI_INCLUDE_STB_IMAGE_H + +#ifdef STB_IMAGE_IMPLEMENTATION + +#if defined(STBI_ONLY_JPEG) || defined(STBI_ONLY_PNG) || defined(STBI_ONLY_BMP) \ + || defined(STBI_ONLY_TGA) || defined(STBI_ONLY_GIF) || defined(STBI_ONLY_PSD) \ + || defined(STBI_ONLY_HDR) || defined(STBI_ONLY_PIC) || defined(STBI_ONLY_PNM) \ + || defined(STBI_ONLY_ZLIB) + #ifndef STBI_ONLY_JPEG + #define STBI_NO_JPEG + #endif + #ifndef STBI_ONLY_PNG + #define STBI_NO_PNG + #endif + #ifndef STBI_ONLY_BMP + #define STBI_NO_BMP + #endif + #ifndef STBI_ONLY_PSD + #define STBI_NO_PSD + #endif + #ifndef STBI_ONLY_TGA + #define STBI_NO_TGA + #endif + #ifndef STBI_ONLY_GIF + #define STBI_NO_GIF + #endif + #ifndef STBI_ONLY_HDR + #define STBI_NO_HDR + #endif + #ifndef STBI_ONLY_PIC + #define STBI_NO_PIC + #endif + #ifndef STBI_ONLY_PNM + #define STBI_NO_PNM + #endif +#endif + +#if defined(STBI_NO_PNG) && !defined(STBI_SUPPORT_ZLIB) && !defined(STBI_NO_ZLIB) +#define STBI_NO_ZLIB +#endif + + +#include +#include // ptrdiff_t on osx +#include +#include +#include + +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) +#include // ldexp, pow +#endif + +#ifndef STBI_NO_STDIO +#include +#endif + +#ifndef STBI_ASSERT +#include +#define STBI_ASSERT(x) assert(x) +#endif + +#ifdef __cplusplus +#define STBI_EXTERN extern "C" +#else +#define STBI_EXTERN extern +#endif + + +#ifndef _MSC_VER + #ifdef __cplusplus + #define stbi_inline inline + #else + #define stbi_inline + #endif +#else + #define stbi_inline __forceinline +#endif + +#ifndef STBI_NO_THREAD_LOCALS + #if defined(__cplusplus) && __cplusplus >= 201103L + #define STBI_THREAD_LOCAL thread_local + #elif defined(__GNUC__) && __GNUC__ < 5 + #define STBI_THREAD_LOCAL __thread + #elif defined(_MSC_VER) + #define STBI_THREAD_LOCAL __declspec(thread) + #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_THREADS__) + #define STBI_THREAD_LOCAL _Thread_local + #endif + + #ifndef STBI_THREAD_LOCAL + #if defined(__GNUC__) + #define STBI_THREAD_LOCAL __thread + #endif + #endif +#endif + +#if defined(_MSC_VER) || defined(__SYMBIAN32__) +typedef unsigned short stbi__uint16; +typedef signed short stbi__int16; +typedef unsigned int stbi__uint32; +typedef signed int stbi__int32; +#else +#include +typedef uint16_t stbi__uint16; +typedef int16_t stbi__int16; +typedef uint32_t stbi__uint32; +typedef int32_t stbi__int32; +#endif + +// should produce compiler error if size is wrong +typedef unsigned char validate_uint32[sizeof(stbi__uint32)==4 ? 1 : -1]; + +#ifdef _MSC_VER +#define STBI_NOTUSED(v) (void)(v) +#else +#define STBI_NOTUSED(v) (void)sizeof(v) +#endif + +#ifdef _MSC_VER +#define STBI_HAS_LROTL +#endif + +#ifdef STBI_HAS_LROTL + #define stbi_lrot(x,y) _lrotl(x,y) +#else + #define stbi_lrot(x,y) (((x) << (y)) | ((x) >> (-(y) & 31))) +#endif + +#if defined(STBI_MALLOC) && defined(STBI_FREE) && (defined(STBI_REALLOC) || defined(STBI_REALLOC_SIZED)) +// ok +#elif !defined(STBI_MALLOC) && !defined(STBI_FREE) && !defined(STBI_REALLOC) && !defined(STBI_REALLOC_SIZED) +// ok +#else +#error "Must define all or none of STBI_MALLOC, STBI_FREE, and STBI_REALLOC (or STBI_REALLOC_SIZED)." +#endif + +#ifndef STBI_MALLOC +#define STBI_MALLOC(sz) malloc(sz) +#define STBI_REALLOC(p,newsz) realloc(p,newsz) +#define STBI_FREE(p) free(p) +#endif + +#ifndef STBI_REALLOC_SIZED +#define STBI_REALLOC_SIZED(p,oldsz,newsz) STBI_REALLOC(p,newsz) +#endif + +// x86/x64 detection +#if defined(__x86_64__) || defined(_M_X64) +#define STBI__X64_TARGET +#elif defined(__i386) || defined(_M_IX86) +#define STBI__X86_TARGET +#endif + +#if defined(__GNUC__) && defined(STBI__X86_TARGET) && !defined(__SSE2__) && !defined(STBI_NO_SIMD) +// gcc doesn't support sse2 intrinsics unless you compile with -msse2, +// which in turn means it gets to use SSE2 everywhere. This is unfortunate, +// but previous attempts to provide the SSE2 functions with runtime +// detection caused numerous issues. The way architecture extensions are +// exposed in GCC/Clang is, sadly, not really suited for one-file libs. +// New behavior: if compiled with -msse2, we use SSE2 without any +// detection; if not, we don't use it at all. +#define STBI_NO_SIMD +#endif + +#if defined(__MINGW32__) && defined(STBI__X86_TARGET) && !defined(STBI_MINGW_ENABLE_SSE2) && !defined(STBI_NO_SIMD) +// Note that __MINGW32__ doesn't actually mean 32-bit, so we have to avoid STBI__X64_TARGET +// +// 32-bit MinGW wants ESP to be 16-byte aligned, but this is not in the +// Windows ABI and VC++ as well as Windows DLLs don't maintain that invariant. +// As a result, enabling SSE2 on 32-bit MinGW is dangerous when not +// simultaneously enabling "-mstackrealign". +// +// See https://github.com/nothings/stb/issues/81 for more information. +// +// So default to no SSE2 on 32-bit MinGW. If you've read this far and added +// -mstackrealign to your build settings, feel free to #define STBI_MINGW_ENABLE_SSE2. +#define STBI_NO_SIMD +#endif + +#if !defined(STBI_NO_SIMD) && (defined(STBI__X86_TARGET) || defined(STBI__X64_TARGET)) +#define STBI_SSE2 +#include + +#ifdef _MSC_VER + +#if _MSC_VER >= 1400 // not VC6 +#include // __cpuid +static int stbi__cpuid3(void) +{ + int info[4]; + __cpuid(info,1); + return info[3]; +} +#else +static int stbi__cpuid3(void) +{ + int res; + __asm { + mov eax,1 + cpuid + mov res,edx + } + return res; +} +#endif + +#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name + +#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) +static int stbi__sse2_available(void) +{ + int info3 = stbi__cpuid3(); + return ((info3 >> 26) & 1) != 0; +} +#endif + +#else // assume GCC-style if not VC++ +#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) + +#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) +static int stbi__sse2_available(void) +{ + // If we're even attempting to compile this on GCC/Clang, that means + // -msse2 is on, which means the compiler is allowed to use SSE2 + // instructions at will, and so are we. + return 1; +} +#endif + +#endif +#endif + +// ARM NEON +#if defined(STBI_NO_SIMD) && defined(STBI_NEON) +#undef STBI_NEON +#endif + +#ifdef STBI_NEON +#include +#ifdef _MSC_VER +#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name +#else +#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) +#endif +#endif + +#ifndef STBI_SIMD_ALIGN +#define STBI_SIMD_ALIGN(type, name) type name +#endif + +#ifndef STBI_MAX_DIMENSIONS +#define STBI_MAX_DIMENSIONS (1 << 24) +#endif + +/////////////////////////////////////////////// +// +// stbi__context struct and start_xxx functions + +// stbi__context structure is our basic context used by all images, so it +// contains all the IO context, plus some basic image information +typedef struct +{ + stbi__uint32 img_x, img_y; + int img_n, img_out_n; + + stbi_io_callbacks io; + void *io_user_data; + + int read_from_callbacks; + int buflen; + stbi_uc buffer_start[128]; + int callback_already_read; + + stbi_uc *img_buffer, *img_buffer_end; + stbi_uc *img_buffer_original, *img_buffer_original_end; +} stbi__context; + + +static void stbi__refill_buffer(stbi__context *s); + +// initialize a memory-decode context +static void stbi__start_mem(stbi__context *s, stbi_uc const *buffer, int len) +{ + s->io.read = NULL; + s->read_from_callbacks = 0; + s->callback_already_read = 0; + s->img_buffer = s->img_buffer_original = (stbi_uc *) buffer; + s->img_buffer_end = s->img_buffer_original_end = (stbi_uc *) buffer+len; +} + +// initialize a callback-based context +static void stbi__start_callbacks(stbi__context *s, stbi_io_callbacks *c, void *user) +{ + s->io = *c; + s->io_user_data = user; + s->buflen = sizeof(s->buffer_start); + s->read_from_callbacks = 1; + s->callback_already_read = 0; + s->img_buffer = s->img_buffer_original = s->buffer_start; + stbi__refill_buffer(s); + s->img_buffer_original_end = s->img_buffer_end; +} + +#ifndef STBI_NO_STDIO + +static int stbi__stdio_read(void *user, char *data, int size) +{ + return (int) fread(data,1,size,(FILE*) user); +} + +static void stbi__stdio_skip(void *user, int n) +{ + int ch; + fseek((FILE*) user, n, SEEK_CUR); + ch = fgetc((FILE*) user); /* have to read a byte to reset feof()'s flag */ + if (ch != EOF) { + ungetc(ch, (FILE *) user); /* push byte back onto stream if valid. */ + } +} + +static int stbi__stdio_eof(void *user) +{ + return feof((FILE*) user) || ferror((FILE *) user); +} + +static stbi_io_callbacks stbi__stdio_callbacks = +{ + stbi__stdio_read, + stbi__stdio_skip, + stbi__stdio_eof, +}; + +static void stbi__start_file(stbi__context *s, FILE *f) +{ + stbi__start_callbacks(s, &stbi__stdio_callbacks, (void *) f); +} + +//static void stop_file(stbi__context *s) { } + +#endif // !STBI_NO_STDIO + +static void stbi__rewind(stbi__context *s) +{ + // conceptually rewind SHOULD rewind to the beginning of the stream, + // but we just rewind to the beginning of the initial buffer, because + // we only use it after doing 'test', which only ever looks at at most 92 bytes + s->img_buffer = s->img_buffer_original; + s->img_buffer_end = s->img_buffer_original_end; +} + +enum +{ + STBI_ORDER_RGB, + STBI_ORDER_BGR +}; + +typedef struct +{ + int bits_per_channel; + int num_channels; + int channel_order; +} stbi__result_info; + +#ifndef STBI_NO_JPEG +static int stbi__jpeg_test(stbi__context *s); +static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PNG +static int stbi__png_test(stbi__context *s); +static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__png_is16(stbi__context *s); +#endif + +#ifndef STBI_NO_BMP +static int stbi__bmp_test(stbi__context *s); +static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_TGA +static int stbi__tga_test(stbi__context *s); +static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PSD +static int stbi__psd_test(stbi__context *s); +static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc); +static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__psd_is16(stbi__context *s); +#endif + +#ifndef STBI_NO_HDR +static int stbi__hdr_test(stbi__context *s); +static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PIC +static int stbi__pic_test(stbi__context *s); +static void *stbi__pic_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_GIF +static int stbi__gif_test(stbi__context *s); +static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp); +static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PNM +static int stbi__pnm_test(stbi__context *s); +static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__pnm_is16(stbi__context *s); +#endif + +static +#ifdef STBI_THREAD_LOCAL +STBI_THREAD_LOCAL +#endif +const char *stbi__g_failure_reason; + +STBIDEF const char *stbi_failure_reason(void) +{ + return stbi__g_failure_reason; +} + +#ifndef STBI_NO_FAILURE_STRINGS +static int stbi__err(const char *str) +{ + stbi__g_failure_reason = str; + return 0; +} +#endif + +static void *stbi__malloc(size_t size) +{ + return STBI_MALLOC(size); +} + +// stb_image uses ints pervasively, including for offset calculations. +// therefore the largest decoded image size we can support with the +// current code, even on 64-bit targets, is INT_MAX. this is not a +// significant limitation for the intended use case. +// +// we do, however, need to make sure our size calculations don't +// overflow. hence a few helper functions for size calculations that +// multiply integers together, making sure that they're non-negative +// and no overflow occurs. + +// return 1 if the sum is valid, 0 on overflow. +// negative terms are considered invalid. +static int stbi__addsizes_valid(int a, int b) +{ + if (b < 0) return 0; + // now 0 <= b <= INT_MAX, hence also + // 0 <= INT_MAX - b <= INTMAX. + // And "a + b <= INT_MAX" (which might overflow) is the + // same as a <= INT_MAX - b (no overflow) + return a <= INT_MAX - b; +} + +// returns 1 if the product is valid, 0 on overflow. +// negative factors are considered invalid. +static int stbi__mul2sizes_valid(int a, int b) +{ + if (a < 0 || b < 0) return 0; + if (b == 0) return 1; // mul-by-0 is always safe + // portable way to check for no overflows in a*b + return a <= INT_MAX/b; +} + +#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) +// returns 1 if "a*b + add" has no negative terms/factors and doesn't overflow +static int stbi__mad2sizes_valid(int a, int b, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__addsizes_valid(a*b, add); +} +#endif + +// returns 1 if "a*b*c + add" has no negative terms/factors and doesn't overflow +static int stbi__mad3sizes_valid(int a, int b, int c, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && + stbi__addsizes_valid(a*b*c, add); +} + +// returns 1 if "a*b*c*d + add" has no negative terms/factors and doesn't overflow +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) +static int stbi__mad4sizes_valid(int a, int b, int c, int d, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && + stbi__mul2sizes_valid(a*b*c, d) && stbi__addsizes_valid(a*b*c*d, add); +} +#endif + +#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) +// mallocs with size overflow checking +static void *stbi__malloc_mad2(int a, int b, int add) +{ + if (!stbi__mad2sizes_valid(a, b, add)) return NULL; + return stbi__malloc(a*b + add); +} +#endif + +static void *stbi__malloc_mad3(int a, int b, int c, int add) +{ + if (!stbi__mad3sizes_valid(a, b, c, add)) return NULL; + return stbi__malloc(a*b*c + add); +} + +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) +static void *stbi__malloc_mad4(int a, int b, int c, int d, int add) +{ + if (!stbi__mad4sizes_valid(a, b, c, d, add)) return NULL; + return stbi__malloc(a*b*c*d + add); +} +#endif + +// returns 1 if the sum of two signed ints is valid (between -2^31 and 2^31-1 inclusive), 0 on overflow. +static int stbi__addints_valid(int a, int b) +{ + if ((a >= 0) != (b >= 0)) return 1; // a and b have different signs, so no overflow + if (a < 0 && b < 0) return a >= INT_MIN - b; // same as a + b >= INT_MIN; INT_MIN - b cannot overflow since b < 0. + return a <= INT_MAX - b; +} + +// returns 1 if the product of two ints fits in a signed short, 0 on overflow. +static int stbi__mul2shorts_valid(int a, int b) +{ + if (b == 0 || b == -1) return 1; // multiplication by 0 is always 0; check for -1 so SHRT_MIN/b doesn't overflow + if ((a >= 0) == (b >= 0)) return a <= SHRT_MAX/b; // product is positive, so similar to mul2sizes_valid + if (b < 0) return a <= SHRT_MIN / b; // same as a * b >= SHRT_MIN + return a >= SHRT_MIN / b; +} + +// stbi__err - error +// stbi__errpf - error returning pointer to float +// stbi__errpuc - error returning pointer to unsigned char + +#ifdef STBI_NO_FAILURE_STRINGS + #define stbi__err(x,y) 0 +#elif defined(STBI_FAILURE_USERMSG) + #define stbi__err(x,y) stbi__err(y) +#else + #define stbi__err(x,y) stbi__err(x) +#endif + +#define stbi__errpf(x,y) ((float *)(size_t) (stbi__err(x,y)?NULL:NULL)) +#define stbi__errpuc(x,y) ((unsigned char *)(size_t) (stbi__err(x,y)?NULL:NULL)) + +STBIDEF void stbi_image_free(void *retval_from_stbi_load) +{ + STBI_FREE(retval_from_stbi_load); +} + +#ifndef STBI_NO_LINEAR +static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp); +#endif + +#ifndef STBI_NO_HDR +static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp); +#endif + +static int stbi__vertically_flip_on_load_global = 0; + +STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip) +{ + stbi__vertically_flip_on_load_global = flag_true_if_should_flip; +} + +#ifndef STBI_THREAD_LOCAL +#define stbi__vertically_flip_on_load stbi__vertically_flip_on_load_global +#else +static STBI_THREAD_LOCAL int stbi__vertically_flip_on_load_local, stbi__vertically_flip_on_load_set; + +STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip) +{ + stbi__vertically_flip_on_load_local = flag_true_if_should_flip; + stbi__vertically_flip_on_load_set = 1; +} + +#define stbi__vertically_flip_on_load (stbi__vertically_flip_on_load_set \ + ? stbi__vertically_flip_on_load_local \ + : stbi__vertically_flip_on_load_global) +#endif // STBI_THREAD_LOCAL + +static void *stbi__load_main(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) +{ + memset(ri, 0, sizeof(*ri)); // make sure it's initialized if we add new fields + ri->bits_per_channel = 8; // default is 8 so most paths don't have to be changed + ri->channel_order = STBI_ORDER_RGB; // all current input & output are this, but this is here so we can add BGR order + ri->num_channels = 0; + + // test the formats with a very explicit header first (at least a FOURCC + // or distinctive magic number first) + #ifndef STBI_NO_PNG + if (stbi__png_test(s)) return stbi__png_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_BMP + if (stbi__bmp_test(s)) return stbi__bmp_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_GIF + if (stbi__gif_test(s)) return stbi__gif_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_PSD + if (stbi__psd_test(s)) return stbi__psd_load(s,x,y,comp,req_comp, ri, bpc); + #else + STBI_NOTUSED(bpc); + #endif + #ifndef STBI_NO_PIC + if (stbi__pic_test(s)) return stbi__pic_load(s,x,y,comp,req_comp, ri); + #endif + + // then the formats that can end up attempting to load with just 1 or 2 + // bytes matching expectations; these are prone to false positives, so + // try them later + #ifndef STBI_NO_JPEG + if (stbi__jpeg_test(s)) return stbi__jpeg_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_PNM + if (stbi__pnm_test(s)) return stbi__pnm_load(s,x,y,comp,req_comp, ri); + #endif + + #ifndef STBI_NO_HDR + if (stbi__hdr_test(s)) { + float *hdr = stbi__hdr_load(s, x,y,comp,req_comp, ri); + return stbi__hdr_to_ldr(hdr, *x, *y, req_comp ? req_comp : *comp); + } + #endif + + #ifndef STBI_NO_TGA + // test tga last because it's a crappy test! + if (stbi__tga_test(s)) + return stbi__tga_load(s,x,y,comp,req_comp, ri); + #endif + + return stbi__errpuc("unknown image type", "Image not of any known type, or corrupt"); +} + +static stbi_uc *stbi__convert_16_to_8(stbi__uint16 *orig, int w, int h, int channels) +{ + int i; + int img_len = w * h * channels; + stbi_uc *reduced; + + reduced = (stbi_uc *) stbi__malloc(img_len); + if (reduced == NULL) return stbi__errpuc("outofmem", "Out of memory"); + + for (i = 0; i < img_len; ++i) + reduced[i] = (stbi_uc)((orig[i] >> 8) & 0xFF); // top half of each byte is sufficient approx of 16->8 bit scaling + + STBI_FREE(orig); + return reduced; +} + +static stbi__uint16 *stbi__convert_8_to_16(stbi_uc *orig, int w, int h, int channels) +{ + int i; + int img_len = w * h * channels; + stbi__uint16 *enlarged; + + enlarged = (stbi__uint16 *) stbi__malloc(img_len*2); + if (enlarged == NULL) return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); + + for (i = 0; i < img_len; ++i) + enlarged[i] = (stbi__uint16)((orig[i] << 8) + orig[i]); // replicate to high and low byte, maps 0->0, 255->0xffff + + STBI_FREE(orig); + return enlarged; +} + +static void stbi__vertical_flip(void *image, int w, int h, int bytes_per_pixel) +{ + int row; + size_t bytes_per_row = (size_t)w * bytes_per_pixel; + stbi_uc temp[2048]; + stbi_uc *bytes = (stbi_uc *)image; + + for (row = 0; row < (h>>1); row++) { + stbi_uc *row0 = bytes + row*bytes_per_row; + stbi_uc *row1 = bytes + (h - row - 1)*bytes_per_row; + // swap row0 with row1 + size_t bytes_left = bytes_per_row; + while (bytes_left) { + size_t bytes_copy = (bytes_left < sizeof(temp)) ? bytes_left : sizeof(temp); + memcpy(temp, row0, bytes_copy); + memcpy(row0, row1, bytes_copy); + memcpy(row1, temp, bytes_copy); + row0 += bytes_copy; + row1 += bytes_copy; + bytes_left -= bytes_copy; + } + } +} + +#ifndef STBI_NO_GIF +static void stbi__vertical_flip_slices(void *image, int w, int h, int z, int bytes_per_pixel) +{ + int slice; + int slice_size = w * h * bytes_per_pixel; + + stbi_uc *bytes = (stbi_uc *)image; + for (slice = 0; slice < z; ++slice) { + stbi__vertical_flip(bytes, w, h, bytes_per_pixel); + bytes += slice_size; + } +} +#endif + +static unsigned char *stbi__load_and_postprocess_8bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + stbi__result_info ri; + void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 8); + + if (result == NULL) + return NULL; + + // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. + STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); + + if (ri.bits_per_channel != 8) { + result = stbi__convert_16_to_8((stbi__uint16 *) result, *x, *y, req_comp == 0 ? *comp : req_comp); + ri.bits_per_channel = 8; + } + + // @TODO: move stbi__convert_format to here + + if (stbi__vertically_flip_on_load) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi_uc)); + } + + return (unsigned char *) result; +} + +static stbi__uint16 *stbi__load_and_postprocess_16bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + stbi__result_info ri; + void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 16); + + if (result == NULL) + return NULL; + + // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. + STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); + + if (ri.bits_per_channel != 16) { + result = stbi__convert_8_to_16((stbi_uc *) result, *x, *y, req_comp == 0 ? *comp : req_comp); + ri.bits_per_channel = 16; + } + + // @TODO: move stbi__convert_format16 to here + // @TODO: special case RGB-to-Y (and RGBA-to-YA) for 8-bit-to-16-bit case to keep more precision + + if (stbi__vertically_flip_on_load) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi__uint16)); + } + + return (stbi__uint16 *) result; +} + +#if !defined(STBI_NO_HDR) && !defined(STBI_NO_LINEAR) +static void stbi__float_postprocess(float *result, int *x, int *y, int *comp, int req_comp) +{ + if (stbi__vertically_flip_on_load && result != NULL) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(float)); + } +} +#endif + +#ifndef STBI_NO_STDIO + +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) +STBI_EXTERN __declspec(dllimport) int __stdcall MultiByteToWideChar(unsigned int cp, unsigned long flags, const char *str, int cbmb, wchar_t *widestr, int cchwide); +STBI_EXTERN __declspec(dllimport) int __stdcall WideCharToMultiByte(unsigned int cp, unsigned long flags, const wchar_t *widestr, int cchwide, char *str, int cbmb, const char *defchar, int *used_default); +#endif + +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) +STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input) +{ + return WideCharToMultiByte(65001 /* UTF8 */, 0, input, -1, buffer, (int) bufferlen, NULL, NULL); +} +#endif + +static FILE *stbi__fopen(char const *filename, char const *mode) +{ + FILE *f; +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) + wchar_t wMode[64]; + wchar_t wFilename[1024]; + if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, filename, -1, wFilename, sizeof(wFilename)/sizeof(*wFilename))) + return 0; + + if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, mode, -1, wMode, sizeof(wMode)/sizeof(*wMode))) + return 0; + +#if defined(_MSC_VER) && _MSC_VER >= 1400 + if (0 != _wfopen_s(&f, wFilename, wMode)) + f = 0; +#else + f = _wfopen(wFilename, wMode); +#endif + +#elif defined(_MSC_VER) && _MSC_VER >= 1400 + if (0 != fopen_s(&f, filename, mode)) + f=0; +#else + f = fopen(filename, mode); +#endif + return f; +} + + +STBIDEF stbi_uc *stbi_load(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + unsigned char *result; + if (!f) return stbi__errpuc("can't fopen", "Unable to open file"); + result = stbi_load_from_file(f,x,y,comp,req_comp); + fclose(f); + return result; +} + +STBIDEF stbi_uc *stbi_load_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + unsigned char *result; + stbi__context s; + stbi__start_file(&s,f); + result = stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); + if (result) { + // need to 'unget' all the characters in the IO buffer + fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); + } + return result; +} + +STBIDEF stbi__uint16 *stbi_load_from_file_16(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + stbi__uint16 *result; + stbi__context s; + stbi__start_file(&s,f); + result = stbi__load_and_postprocess_16bit(&s,x,y,comp,req_comp); + if (result) { + // need to 'unget' all the characters in the IO buffer + fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); + } + return result; +} + +STBIDEF stbi_us *stbi_load_16(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + stbi__uint16 *result; + if (!f) return (stbi_us *) stbi__errpuc("can't fopen", "Unable to open file"); + result = stbi_load_from_file_16(f,x,y,comp,req_comp); + fclose(f); + return result; +} + + +#endif //!STBI_NO_STDIO + +STBIDEF stbi_us *stbi_load_16_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); +} + +STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *)clbk, user); + return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); +} + +STBIDEF stbi_uc *stbi_load_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); +} + +STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); +} + +#ifndef STBI_NO_GIF +STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp) +{ + unsigned char *result; + stbi__context s; + stbi__start_mem(&s,buffer,len); + + result = (unsigned char*) stbi__load_gif_main(&s, delays, x, y, z, comp, req_comp); + if (stbi__vertically_flip_on_load) { + stbi__vertical_flip_slices( result, *x, *y, *z, *comp ); + } + + return result; +} +#endif + +#ifndef STBI_NO_LINEAR +static float *stbi__loadf_main(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + unsigned char *data; + #ifndef STBI_NO_HDR + if (stbi__hdr_test(s)) { + stbi__result_info ri; + float *hdr_data = stbi__hdr_load(s,x,y,comp,req_comp, &ri); + if (hdr_data) + stbi__float_postprocess(hdr_data,x,y,comp,req_comp); + return hdr_data; + } + #endif + data = stbi__load_and_postprocess_8bit(s, x, y, comp, req_comp); + if (data) + return stbi__ldr_to_hdr(data, *x, *y, req_comp ? req_comp : *comp); + return stbi__errpf("unknown image type", "Image not of any known type, or corrupt"); +} + +STBIDEF float *stbi_loadf_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} + +STBIDEF float *stbi_loadf_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} + +#ifndef STBI_NO_STDIO +STBIDEF float *stbi_loadf(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + float *result; + FILE *f = stbi__fopen(filename, "rb"); + if (!f) return stbi__errpf("can't fopen", "Unable to open file"); + result = stbi_loadf_from_file(f,x,y,comp,req_comp); + fclose(f); + return result; +} + +STBIDEF float *stbi_loadf_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_file(&s,f); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} +#endif // !STBI_NO_STDIO + +#endif // !STBI_NO_LINEAR + +// these is-hdr-or-not is defined independent of whether STBI_NO_LINEAR is +// defined, for API simplicity; if STBI_NO_LINEAR is defined, it always +// reports false! + +STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len) +{ + #ifndef STBI_NO_HDR + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__hdr_test(&s); + #else + STBI_NOTUSED(buffer); + STBI_NOTUSED(len); + return 0; + #endif +} + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_is_hdr (char const *filename) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result=0; + if (f) { + result = stbi_is_hdr_from_file(f); + fclose(f); + } + return result; +} + +STBIDEF int stbi_is_hdr_from_file(FILE *f) +{ + #ifndef STBI_NO_HDR + long pos = ftell(f); + int res; + stbi__context s; + stbi__start_file(&s,f); + res = stbi__hdr_test(&s); + fseek(f, pos, SEEK_SET); + return res; + #else + STBI_NOTUSED(f); + return 0; + #endif +} +#endif // !STBI_NO_STDIO + +STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user) +{ + #ifndef STBI_NO_HDR + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__hdr_test(&s); + #else + STBI_NOTUSED(clbk); + STBI_NOTUSED(user); + return 0; + #endif +} + +#ifndef STBI_NO_LINEAR +static float stbi__l2h_gamma=2.2f, stbi__l2h_scale=1.0f; + +STBIDEF void stbi_ldr_to_hdr_gamma(float gamma) { stbi__l2h_gamma = gamma; } +STBIDEF void stbi_ldr_to_hdr_scale(float scale) { stbi__l2h_scale = scale; } +#endif + +static float stbi__h2l_gamma_i=1.0f/2.2f, stbi__h2l_scale_i=1.0f; + +STBIDEF void stbi_hdr_to_ldr_gamma(float gamma) { stbi__h2l_gamma_i = 1/gamma; } +STBIDEF void stbi_hdr_to_ldr_scale(float scale) { stbi__h2l_scale_i = 1/scale; } + + +////////////////////////////////////////////////////////////////////////////// +// +// Common code used by all image loaders +// + +enum +{ + STBI__SCAN_load=0, + STBI__SCAN_type, + STBI__SCAN_header +}; + +static void stbi__refill_buffer(stbi__context *s) +{ + int n = (s->io.read)(s->io_user_data,(char*)s->buffer_start,s->buflen); + s->callback_already_read += (int) (s->img_buffer - s->img_buffer_original); + if (n == 0) { + // at end of file, treat same as if from memory, but need to handle case + // where s->img_buffer isn't pointing to safe memory, e.g. 0-byte file + s->read_from_callbacks = 0; + s->img_buffer = s->buffer_start; + s->img_buffer_end = s->buffer_start+1; + *s->img_buffer = 0; + } else { + s->img_buffer = s->buffer_start; + s->img_buffer_end = s->buffer_start + n; + } +} + +stbi_inline static stbi_uc stbi__get8(stbi__context *s) +{ + if (s->img_buffer < s->img_buffer_end) + return *s->img_buffer++; + if (s->read_from_callbacks) { + stbi__refill_buffer(s); + return *s->img_buffer++; + } + return 0; +} + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_HDR) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +stbi_inline static int stbi__at_eof(stbi__context *s) +{ + if (s->io.read) { + if (!(s->io.eof)(s->io_user_data)) return 0; + // if feof() is true, check if buffer = end + // special case: we've only got the special 0 character at the end + if (s->read_from_callbacks == 0) return 1; + } + + return s->img_buffer >= s->img_buffer_end; +} +#endif + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) +// nothing +#else +static void stbi__skip(stbi__context *s, int n) +{ + if (n == 0) return; // already there! + if (n < 0) { + s->img_buffer = s->img_buffer_end; + return; + } + if (s->io.read) { + int blen = (int) (s->img_buffer_end - s->img_buffer); + if (blen < n) { + s->img_buffer = s->img_buffer_end; + (s->io.skip)(s->io_user_data, n - blen); + return; + } + } + s->img_buffer += n; +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_TGA) && defined(STBI_NO_HDR) && defined(STBI_NO_PNM) +// nothing +#else +static int stbi__getn(stbi__context *s, stbi_uc *buffer, int n) +{ + if (s->io.read) { + int blen = (int) (s->img_buffer_end - s->img_buffer); + if (blen < n) { + int res, count; + + memcpy(buffer, s->img_buffer, blen); + + count = (s->io.read)(s->io_user_data, (char*) buffer + blen, n - blen); + res = (count == (n-blen)); + s->img_buffer = s->img_buffer_end; + return res; + } + } + + if (s->img_buffer+n <= s->img_buffer_end) { + memcpy(buffer, s->img_buffer, n); + s->img_buffer += n; + return 1; + } else + return 0; +} +#endif + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) +// nothing +#else +static int stbi__get16be(stbi__context *s) +{ + int z = stbi__get8(s); + return (z << 8) + stbi__get8(s); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) +// nothing +#else +static stbi__uint32 stbi__get32be(stbi__context *s) +{ + stbi__uint32 z = stbi__get16be(s); + return (z << 16) + stbi__get16be(s); +} +#endif + +#if defined(STBI_NO_BMP) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) +// nothing +#else +static int stbi__get16le(stbi__context *s) +{ + int z = stbi__get8(s); + return z + (stbi__get8(s) << 8); +} +#endif + +#ifndef STBI_NO_BMP +static stbi__uint32 stbi__get32le(stbi__context *s) +{ + stbi__uint32 z = stbi__get16le(s); + z += (stbi__uint32)stbi__get16le(s) << 16; + return z; +} +#endif + +#define STBI__BYTECAST(x) ((stbi_uc) ((x) & 255)) // truncate int to byte without warnings + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +////////////////////////////////////////////////////////////////////////////// +// +// generic converter from built-in img_n to req_comp +// individual types do this automatically as much as possible (e.g. jpeg +// does all cases internally since it needs to colorspace convert anyway, +// and it never has alpha, so very few cases ). png can automatically +// interleave an alpha=255 channel, but falls back to this for other cases +// +// assume data buffer is malloced, so malloc a new one and free that one +// only failure mode is malloc failing + +static stbi_uc stbi__compute_y(int r, int g, int b) +{ + return (stbi_uc) (((r*77) + (g*150) + (29*b)) >> 8); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +static unsigned char *stbi__convert_format(unsigned char *data, int img_n, int req_comp, unsigned int x, unsigned int y) +{ + int i,j; + unsigned char *good; + + if (req_comp == img_n) return data; + STBI_ASSERT(req_comp >= 1 && req_comp <= 4); + + good = (unsigned char *) stbi__malloc_mad3(req_comp, x, y, 0); + if (good == NULL) { + STBI_FREE(data); + return stbi__errpuc("outofmem", "Out of memory"); + } + + for (j=0; j < (int) y; ++j) { + unsigned char *src = data + j * x * img_n ; + unsigned char *dest = good + j * x * req_comp; + + #define STBI__COMBO(a,b) ((a)*8+(b)) + #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) + // convert source image with img_n components to one with req_comp components; + // avoid switch per pixel, so use switch per scanline and massive macros + switch (STBI__COMBO(img_n, req_comp)) { + STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=255; } break; + STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=255; } break; + STBI__CASE(2,1) { dest[0]=src[0]; } break; + STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; + STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=255; } break; + STBI__CASE(3,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; + STBI__CASE(3,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = 255; } break; + STBI__CASE(4,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; + STBI__CASE(4,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = src[3]; } break; + STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; + default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return stbi__errpuc("unsupported", "Unsupported format conversion"); + } + #undef STBI__CASE + } + + STBI_FREE(data); + return good; +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) +// nothing +#else +static stbi__uint16 stbi__compute_y_16(int r, int g, int b) +{ + return (stbi__uint16) (((r*77) + (g*150) + (29*b)) >> 8); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) +// nothing +#else +static stbi__uint16 *stbi__convert_format16(stbi__uint16 *data, int img_n, int req_comp, unsigned int x, unsigned int y) +{ + int i,j; + stbi__uint16 *good; + + if (req_comp == img_n) return data; + STBI_ASSERT(req_comp >= 1 && req_comp <= 4); + + good = (stbi__uint16 *) stbi__malloc(req_comp * x * y * 2); + if (good == NULL) { + STBI_FREE(data); + return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); + } + + for (j=0; j < (int) y; ++j) { + stbi__uint16 *src = data + j * x * img_n ; + stbi__uint16 *dest = good + j * x * req_comp; + + #define STBI__COMBO(a,b) ((a)*8+(b)) + #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) + // convert source image with img_n components to one with req_comp components; + // avoid switch per pixel, so use switch per scanline and massive macros + switch (STBI__COMBO(img_n, req_comp)) { + STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=0xffff; } break; + STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=0xffff; } break; + STBI__CASE(2,1) { dest[0]=src[0]; } break; + STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; + STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=0xffff; } break; + STBI__CASE(3,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; + STBI__CASE(3,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = 0xffff; } break; + STBI__CASE(4,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; + STBI__CASE(4,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = src[3]; } break; + STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; + default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return (stbi__uint16*) stbi__errpuc("unsupported", "Unsupported format conversion"); + } + #undef STBI__CASE + } + + STBI_FREE(data); + return good; +} +#endif + +#ifndef STBI_NO_LINEAR +static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp) +{ + int i,k,n; + float *output; + if (!data) return NULL; + output = (float *) stbi__malloc_mad4(x, y, comp, sizeof(float), 0); + if (output == NULL) { STBI_FREE(data); return stbi__errpf("outofmem", "Out of memory"); } + // compute number of non-alpha components + if (comp & 1) n = comp; else n = comp-1; + for (i=0; i < x*y; ++i) { + for (k=0; k < n; ++k) { + output[i*comp + k] = (float) (pow(data[i*comp+k]/255.0f, stbi__l2h_gamma) * stbi__l2h_scale); + } + } + if (n < comp) { + for (i=0; i < x*y; ++i) { + output[i*comp + n] = data[i*comp + n]/255.0f; + } + } + STBI_FREE(data); + return output; +} +#endif + +#ifndef STBI_NO_HDR +#define stbi__float2int(x) ((int) (x)) +static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp) +{ + int i,k,n; + stbi_uc *output; + if (!data) return NULL; + output = (stbi_uc *) stbi__malloc_mad3(x, y, comp, 0); + if (output == NULL) { STBI_FREE(data); return stbi__errpuc("outofmem", "Out of memory"); } + // compute number of non-alpha components + if (comp & 1) n = comp; else n = comp-1; + for (i=0; i < x*y; ++i) { + for (k=0; k < n; ++k) { + float z = (float) pow(data[i*comp+k]*stbi__h2l_scale_i, stbi__h2l_gamma_i) * 255 + 0.5f; + if (z < 0) z = 0; + if (z > 255) z = 255; + output[i*comp + k] = (stbi_uc) stbi__float2int(z); + } + if (k < comp) { + float z = data[i*comp+k] * 255 + 0.5f; + if (z < 0) z = 0; + if (z > 255) z = 255; + output[i*comp + k] = (stbi_uc) stbi__float2int(z); + } + } + STBI_FREE(data); + return output; +} +#endif + +////////////////////////////////////////////////////////////////////////////// +// +// "baseline" JPEG/JFIF decoder +// +// simple implementation +// - doesn't support delayed output of y-dimension +// - simple interface (only one output format: 8-bit interleaved RGB) +// - doesn't try to recover corrupt jpegs +// - doesn't allow partial loading, loading multiple at once +// - still fast on x86 (copying globals into locals doesn't help x86) +// - allocates lots of intermediate memory (full size of all components) +// - non-interleaved case requires this anyway +// - allows good upsampling (see next) +// high-quality +// - upsampled channels are bilinearly interpolated, even across blocks +// - quality integer IDCT derived from IJG's 'slow' +// performance +// - fast huffman; reasonable integer IDCT +// - some SIMD kernels for common paths on targets with SSE2/NEON +// - uses a lot of intermediate memory, could cache poorly + +#ifndef STBI_NO_JPEG + +// huffman decoding acceleration +#define FAST_BITS 9 // larger handles more cases; smaller stomps less cache + +typedef struct +{ + stbi_uc fast[1 << FAST_BITS]; + // weirdly, repacking this into AoS is a 10% speed loss, instead of a win + stbi__uint16 code[256]; + stbi_uc values[256]; + stbi_uc size[257]; + unsigned int maxcode[18]; + int delta[17]; // old 'firstsymbol' - old 'firstcode' +} stbi__huffman; + +typedef struct +{ + stbi__context *s; + stbi__huffman huff_dc[4]; + stbi__huffman huff_ac[4]; + stbi__uint16 dequant[4][64]; + stbi__int16 fast_ac[4][1 << FAST_BITS]; + +// sizes for components, interleaved MCUs + int img_h_max, img_v_max; + int img_mcu_x, img_mcu_y; + int img_mcu_w, img_mcu_h; + +// definition of jpeg image component + struct + { + int id; + int h,v; + int tq; + int hd,ha; + int dc_pred; + + int x,y,w2,h2; + stbi_uc *data; + void *raw_data, *raw_coeff; + stbi_uc *linebuf; + short *coeff; // progressive only + int coeff_w, coeff_h; // number of 8x8 coefficient blocks + } img_comp[4]; + + stbi__uint32 code_buffer; // jpeg entropy-coded buffer + int code_bits; // number of valid bits + unsigned char marker; // marker seen while filling entropy buffer + int nomore; // flag if we saw a marker so must stop + + int progressive; + int spec_start; + int spec_end; + int succ_high; + int succ_low; + int eob_run; + int jfif; + int app14_color_transform; // Adobe APP14 tag + int rgb; + + int scan_n, order[4]; + int restart_interval, todo; + +// kernels + void (*idct_block_kernel)(stbi_uc *out, int out_stride, short data[64]); + void (*YCbCr_to_RGB_kernel)(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step); + stbi_uc *(*resample_row_hv_2_kernel)(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs); +} stbi__jpeg; + +static int stbi__build_huffman(stbi__huffman *h, int *count) +{ + int i,j,k=0; + unsigned int code; + // build size list for each symbol (from JPEG spec) + for (i=0; i < 16; ++i) { + for (j=0; j < count[i]; ++j) { + h->size[k++] = (stbi_uc) (i+1); + if(k >= 257) return stbi__err("bad size list","Corrupt JPEG"); + } + } + h->size[k] = 0; + + // compute actual symbols (from jpeg spec) + code = 0; + k = 0; + for(j=1; j <= 16; ++j) { + // compute delta to add to code to compute symbol id + h->delta[j] = k - code; + if (h->size[k] == j) { + while (h->size[k] == j) + h->code[k++] = (stbi__uint16) (code++); + if (code-1 >= (1u << j)) return stbi__err("bad code lengths","Corrupt JPEG"); + } + // compute largest code + 1 for this size, preshifted as needed later + h->maxcode[j] = code << (16-j); + code <<= 1; + } + h->maxcode[j] = 0xffffffff; + + // build non-spec acceleration table; 255 is flag for not-accelerated + memset(h->fast, 255, 1 << FAST_BITS); + for (i=0; i < k; ++i) { + int s = h->size[i]; + if (s <= FAST_BITS) { + int c = h->code[i] << (FAST_BITS-s); + int m = 1 << (FAST_BITS-s); + for (j=0; j < m; ++j) { + h->fast[c+j] = (stbi_uc) i; + } + } + } + return 1; +} + +// build a table that decodes both magnitude and value of small ACs in +// one go. +static void stbi__build_fast_ac(stbi__int16 *fast_ac, stbi__huffman *h) +{ + int i; + for (i=0; i < (1 << FAST_BITS); ++i) { + stbi_uc fast = h->fast[i]; + fast_ac[i] = 0; + if (fast < 255) { + int rs = h->values[fast]; + int run = (rs >> 4) & 15; + int magbits = rs & 15; + int len = h->size[fast]; + + if (magbits && len + magbits <= FAST_BITS) { + // magnitude code followed by receive_extend code + int k = ((i << len) & ((1 << FAST_BITS) - 1)) >> (FAST_BITS - magbits); + int m = 1 << (magbits - 1); + if (k < m) k += (~0U << magbits) + 1; + // if the result is small enough, we can fit it in fast_ac table + if (k >= -128 && k <= 127) + fast_ac[i] = (stbi__int16) ((k * 256) + (run * 16) + (len + magbits)); + } + } + } +} + +static void stbi__grow_buffer_unsafe(stbi__jpeg *j) +{ + do { + unsigned int b = j->nomore ? 0 : stbi__get8(j->s); + if (b == 0xff) { + int c = stbi__get8(j->s); + while (c == 0xff) c = stbi__get8(j->s); // consume fill bytes + if (c != 0) { + j->marker = (unsigned char) c; + j->nomore = 1; + return; + } + } + j->code_buffer |= b << (24 - j->code_bits); + j->code_bits += 8; + } while (j->code_bits <= 24); +} + +// (1 << n) - 1 +static const stbi__uint32 stbi__bmask[17]={0,1,3,7,15,31,63,127,255,511,1023,2047,4095,8191,16383,32767,65535}; + +// decode a jpeg huffman value from the bitstream +stbi_inline static int stbi__jpeg_huff_decode(stbi__jpeg *j, stbi__huffman *h) +{ + unsigned int temp; + int c,k; + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + + // look at the top FAST_BITS and determine what symbol ID it is, + // if the code is <= FAST_BITS + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + k = h->fast[c]; + if (k < 255) { + int s = h->size[k]; + if (s > j->code_bits) + return -1; + j->code_buffer <<= s; + j->code_bits -= s; + return h->values[k]; + } + + // naive test is to shift the code_buffer down so k bits are + // valid, then test against maxcode. To speed this up, we've + // preshifted maxcode left so that it has (16-k) 0s at the + // end; in other words, regardless of the number of bits, it + // wants to be compared against something shifted to have 16; + // that way we don't need to shift inside the loop. + temp = j->code_buffer >> 16; + for (k=FAST_BITS+1 ; ; ++k) + if (temp < h->maxcode[k]) + break; + if (k == 17) { + // error! code not found + j->code_bits -= 16; + return -1; + } + + if (k > j->code_bits) + return -1; + + // convert the huffman code to the symbol id + c = ((j->code_buffer >> (32 - k)) & stbi__bmask[k]) + h->delta[k]; + if(c < 0 || c >= 256) // symbol id out of bounds! + return -1; + STBI_ASSERT((((j->code_buffer) >> (32 - h->size[c])) & stbi__bmask[h->size[c]]) == h->code[c]); + + // convert the id to a symbol + j->code_bits -= k; + j->code_buffer <<= k; + return h->values[c]; +} + +// bias[n] = (-1<code_bits < n) stbi__grow_buffer_unsafe(j); + if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing + + sgn = j->code_buffer >> 31; // sign bit always in MSB; 0 if MSB clear (positive), 1 if MSB set (negative) + k = stbi_lrot(j->code_buffer, n); + j->code_buffer = k & ~stbi__bmask[n]; + k &= stbi__bmask[n]; + j->code_bits -= n; + return k + (stbi__jbias[n] & (sgn - 1)); +} + +// get some unsigned bits +stbi_inline static int stbi__jpeg_get_bits(stbi__jpeg *j, int n) +{ + unsigned int k; + if (j->code_bits < n) stbi__grow_buffer_unsafe(j); + if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing + k = stbi_lrot(j->code_buffer, n); + j->code_buffer = k & ~stbi__bmask[n]; + k &= stbi__bmask[n]; + j->code_bits -= n; + return k; +} + +stbi_inline static int stbi__jpeg_get_bit(stbi__jpeg *j) +{ + unsigned int k; + if (j->code_bits < 1) stbi__grow_buffer_unsafe(j); + if (j->code_bits < 1) return 0; // ran out of bits from stream, return 0s intead of continuing + k = j->code_buffer; + j->code_buffer <<= 1; + --j->code_bits; + return k & 0x80000000; +} + +// given a value that's at position X in the zigzag stream, +// where does it appear in the 8x8 matrix coded as row-major? +static const stbi_uc stbi__jpeg_dezigzag[64+15] = +{ + 0, 1, 8, 16, 9, 2, 3, 10, + 17, 24, 32, 25, 18, 11, 4, 5, + 12, 19, 26, 33, 40, 48, 41, 34, + 27, 20, 13, 6, 7, 14, 21, 28, + 35, 42, 49, 56, 57, 50, 43, 36, + 29, 22, 15, 23, 30, 37, 44, 51, + 58, 59, 52, 45, 38, 31, 39, 46, + 53, 60, 61, 54, 47, 55, 62, 63, + // let corrupt input sample past end + 63, 63, 63, 63, 63, 63, 63, 63, + 63, 63, 63, 63, 63, 63, 63 +}; + +// decode one 64-entry block-- +static int stbi__jpeg_decode_block(stbi__jpeg *j, short data[64], stbi__huffman *hdc, stbi__huffman *hac, stbi__int16 *fac, int b, stbi__uint16 *dequant) +{ + int diff,dc,k; + int t; + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + t = stbi__jpeg_huff_decode(j, hdc); + if (t < 0 || t > 15) return stbi__err("bad huffman code","Corrupt JPEG"); + + // 0 all the ac values now so we can do it 32-bits at a time + memset(data,0,64*sizeof(data[0])); + + diff = t ? stbi__extend_receive(j, t) : 0; + if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta","Corrupt JPEG"); + dc = j->img_comp[b].dc_pred + diff; + j->img_comp[b].dc_pred = dc; + if (!stbi__mul2shorts_valid(dc, dequant[0])) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + data[0] = (short) (dc * dequant[0]); + + // decode AC components, see JPEG spec + k = 1; + do { + unsigned int zig; + int c,r,s; + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + r = fac[c]; + if (r) { // fast-AC path + k += (r >> 4) & 15; // run + s = r & 15; // combined length + if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); + j->code_buffer <<= s; + j->code_bits -= s; + // decode into unzigzag'd location + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) ((r >> 8) * dequant[zig]); + } else { + int rs = stbi__jpeg_huff_decode(j, hac); + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (rs != 0xf0) break; // end block + k += 16; + } else { + k += r; + // decode into unzigzag'd location + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) (stbi__extend_receive(j,s) * dequant[zig]); + } + } + } while (k < 64); + return 1; +} + +static int stbi__jpeg_decode_block_prog_dc(stbi__jpeg *j, short data[64], stbi__huffman *hdc, int b) +{ + int diff,dc; + int t; + if (j->spec_end != 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + + if (j->succ_high == 0) { + // first scan for DC coefficient, must be first + memset(data,0,64*sizeof(data[0])); // 0 all the ac values now + t = stbi__jpeg_huff_decode(j, hdc); + if (t < 0 || t > 15) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + diff = t ? stbi__extend_receive(j, t) : 0; + + if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta", "Corrupt JPEG"); + dc = j->img_comp[b].dc_pred + diff; + j->img_comp[b].dc_pred = dc; + if (!stbi__mul2shorts_valid(dc, 1 << j->succ_low)) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + data[0] = (short) (dc * (1 << j->succ_low)); + } else { + // refinement scan for DC coefficient + if (stbi__jpeg_get_bit(j)) + data[0] += (short) (1 << j->succ_low); + } + return 1; +} + +// @OPTIMIZE: store non-zigzagged during the decode passes, +// and only de-zigzag when dequantizing +static int stbi__jpeg_decode_block_prog_ac(stbi__jpeg *j, short data[64], stbi__huffman *hac, stbi__int16 *fac) +{ + int k; + if (j->spec_start == 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + + if (j->succ_high == 0) { + int shift = j->succ_low; + + if (j->eob_run) { + --j->eob_run; + return 1; + } + + k = j->spec_start; + do { + unsigned int zig; + int c,r,s; + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + r = fac[c]; + if (r) { // fast-AC path + k += (r >> 4) & 15; // run + s = r & 15; // combined length + if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); + j->code_buffer <<= s; + j->code_bits -= s; + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) ((r >> 8) * (1 << shift)); + } else { + int rs = stbi__jpeg_huff_decode(j, hac); + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (r < 15) { + j->eob_run = (1 << r); + if (r) + j->eob_run += stbi__jpeg_get_bits(j, r); + --j->eob_run; + break; + } + k += 16; + } else { + k += r; + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) (stbi__extend_receive(j,s) * (1 << shift)); + } + } + } while (k <= j->spec_end); + } else { + // refinement scan for these AC coefficients + + short bit = (short) (1 << j->succ_low); + + if (j->eob_run) { + --j->eob_run; + for (k = j->spec_start; k <= j->spec_end; ++k) { + short *p = &data[stbi__jpeg_dezigzag[k]]; + if (*p != 0) + if (stbi__jpeg_get_bit(j)) + if ((*p & bit)==0) { + if (*p > 0) + *p += bit; + else + *p -= bit; + } + } + } else { + k = j->spec_start; + do { + int r,s; + int rs = stbi__jpeg_huff_decode(j, hac); // @OPTIMIZE see if we can use the fast path here, advance-by-r is so slow, eh + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (r < 15) { + j->eob_run = (1 << r) - 1; + if (r) + j->eob_run += stbi__jpeg_get_bits(j, r); + r = 64; // force end of block + } else { + // r=15 s=0 should write 16 0s, so we just do + // a run of 15 0s and then write s (which is 0), + // so we don't have to do anything special here + } + } else { + if (s != 1) return stbi__err("bad huffman code", "Corrupt JPEG"); + // sign bit + if (stbi__jpeg_get_bit(j)) + s = bit; + else + s = -bit; + } + + // advance by r + while (k <= j->spec_end) { + short *p = &data[stbi__jpeg_dezigzag[k++]]; + if (*p != 0) { + if (stbi__jpeg_get_bit(j)) + if ((*p & bit)==0) { + if (*p > 0) + *p += bit; + else + *p -= bit; + } + } else { + if (r == 0) { + *p = (short) s; + break; + } + --r; + } + } + } while (k <= j->spec_end); + } + } + return 1; +} + +// take a -128..127 value and stbi__clamp it and convert to 0..255 +stbi_inline static stbi_uc stbi__clamp(int x) +{ + // trick to use a single test to catch both cases + if ((unsigned int) x > 255) { + if (x < 0) return 0; + if (x > 255) return 255; + } + return (stbi_uc) x; +} + +#define stbi__f2f(x) ((int) (((x) * 4096 + 0.5))) +#define stbi__fsh(x) ((x) * 4096) + +// derived from jidctint -- DCT_ISLOW +#define STBI__IDCT_1D(s0,s1,s2,s3,s4,s5,s6,s7) \ + int t0,t1,t2,t3,p1,p2,p3,p4,p5,x0,x1,x2,x3; \ + p2 = s2; \ + p3 = s6; \ + p1 = (p2+p3) * stbi__f2f(0.5411961f); \ + t2 = p1 + p3*stbi__f2f(-1.847759065f); \ + t3 = p1 + p2*stbi__f2f( 0.765366865f); \ + p2 = s0; \ + p3 = s4; \ + t0 = stbi__fsh(p2+p3); \ + t1 = stbi__fsh(p2-p3); \ + x0 = t0+t3; \ + x3 = t0-t3; \ + x1 = t1+t2; \ + x2 = t1-t2; \ + t0 = s7; \ + t1 = s5; \ + t2 = s3; \ + t3 = s1; \ + p3 = t0+t2; \ + p4 = t1+t3; \ + p1 = t0+t3; \ + p2 = t1+t2; \ + p5 = (p3+p4)*stbi__f2f( 1.175875602f); \ + t0 = t0*stbi__f2f( 0.298631336f); \ + t1 = t1*stbi__f2f( 2.053119869f); \ + t2 = t2*stbi__f2f( 3.072711026f); \ + t3 = t3*stbi__f2f( 1.501321110f); \ + p1 = p5 + p1*stbi__f2f(-0.899976223f); \ + p2 = p5 + p2*stbi__f2f(-2.562915447f); \ + p3 = p3*stbi__f2f(-1.961570560f); \ + p4 = p4*stbi__f2f(-0.390180644f); \ + t3 += p1+p4; \ + t2 += p2+p3; \ + t1 += p2+p4; \ + t0 += p1+p3; + +static void stbi__idct_block(stbi_uc *out, int out_stride, short data[64]) +{ + int i,val[64],*v=val; + stbi_uc *o; + short *d = data; + + // columns + for (i=0; i < 8; ++i,++d, ++v) { + // if all zeroes, shortcut -- this avoids dequantizing 0s and IDCTing + if (d[ 8]==0 && d[16]==0 && d[24]==0 && d[32]==0 + && d[40]==0 && d[48]==0 && d[56]==0) { + // no shortcut 0 seconds + // (1|2|3|4|5|6|7)==0 0 seconds + // all separate -0.047 seconds + // 1 && 2|3 && 4|5 && 6|7: -0.047 seconds + int dcterm = d[0]*4; + v[0] = v[8] = v[16] = v[24] = v[32] = v[40] = v[48] = v[56] = dcterm; + } else { + STBI__IDCT_1D(d[ 0],d[ 8],d[16],d[24],d[32],d[40],d[48],d[56]) + // constants scaled things up by 1<<12; let's bring them back + // down, but keep 2 extra bits of precision + x0 += 512; x1 += 512; x2 += 512; x3 += 512; + v[ 0] = (x0+t3) >> 10; + v[56] = (x0-t3) >> 10; + v[ 8] = (x1+t2) >> 10; + v[48] = (x1-t2) >> 10; + v[16] = (x2+t1) >> 10; + v[40] = (x2-t1) >> 10; + v[24] = (x3+t0) >> 10; + v[32] = (x3-t0) >> 10; + } + } + + for (i=0, v=val, o=out; i < 8; ++i,v+=8,o+=out_stride) { + // no fast case since the first 1D IDCT spread components out + STBI__IDCT_1D(v[0],v[1],v[2],v[3],v[4],v[5],v[6],v[7]) + // constants scaled things up by 1<<12, plus we had 1<<2 from first + // loop, plus horizontal and vertical each scale by sqrt(8) so together + // we've got an extra 1<<3, so 1<<17 total we need to remove. + // so we want to round that, which means adding 0.5 * 1<<17, + // aka 65536. Also, we'll end up with -128 to 127 that we want + // to encode as 0..255 by adding 128, so we'll add that before the shift + x0 += 65536 + (128<<17); + x1 += 65536 + (128<<17); + x2 += 65536 + (128<<17); + x3 += 65536 + (128<<17); + // tried computing the shifts into temps, or'ing the temps to see + // if any were out of range, but that was slower + o[0] = stbi__clamp((x0+t3) >> 17); + o[7] = stbi__clamp((x0-t3) >> 17); + o[1] = stbi__clamp((x1+t2) >> 17); + o[6] = stbi__clamp((x1-t2) >> 17); + o[2] = stbi__clamp((x2+t1) >> 17); + o[5] = stbi__clamp((x2-t1) >> 17); + o[3] = stbi__clamp((x3+t0) >> 17); + o[4] = stbi__clamp((x3-t0) >> 17); + } +} + +#ifdef STBI_SSE2 +// sse2 integer IDCT. not the fastest possible implementation but it +// produces bit-identical results to the generic C version so it's +// fully "transparent". +static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) +{ + // This is constructed to match our regular (generic) integer IDCT exactly. + __m128i row0, row1, row2, row3, row4, row5, row6, row7; + __m128i tmp; + + // dot product constant: even elems=x, odd elems=y + #define dct_const(x,y) _mm_setr_epi16((x),(y),(x),(y),(x),(y),(x),(y)) + + // out(0) = c0[even]*x + c0[odd]*y (c0, x, y 16-bit, out 32-bit) + // out(1) = c1[even]*x + c1[odd]*y + #define dct_rot(out0,out1, x,y,c0,c1) \ + __m128i c0##lo = _mm_unpacklo_epi16((x),(y)); \ + __m128i c0##hi = _mm_unpackhi_epi16((x),(y)); \ + __m128i out0##_l = _mm_madd_epi16(c0##lo, c0); \ + __m128i out0##_h = _mm_madd_epi16(c0##hi, c0); \ + __m128i out1##_l = _mm_madd_epi16(c0##lo, c1); \ + __m128i out1##_h = _mm_madd_epi16(c0##hi, c1) + + // out = in << 12 (in 16-bit, out 32-bit) + #define dct_widen(out, in) \ + __m128i out##_l = _mm_srai_epi32(_mm_unpacklo_epi16(_mm_setzero_si128(), (in)), 4); \ + __m128i out##_h = _mm_srai_epi32(_mm_unpackhi_epi16(_mm_setzero_si128(), (in)), 4) + + // wide add + #define dct_wadd(out, a, b) \ + __m128i out##_l = _mm_add_epi32(a##_l, b##_l); \ + __m128i out##_h = _mm_add_epi32(a##_h, b##_h) + + // wide sub + #define dct_wsub(out, a, b) \ + __m128i out##_l = _mm_sub_epi32(a##_l, b##_l); \ + __m128i out##_h = _mm_sub_epi32(a##_h, b##_h) + + // butterfly a/b, add bias, then shift by "s" and pack + #define dct_bfly32o(out0, out1, a,b,bias,s) \ + { \ + __m128i abiased_l = _mm_add_epi32(a##_l, bias); \ + __m128i abiased_h = _mm_add_epi32(a##_h, bias); \ + dct_wadd(sum, abiased, b); \ + dct_wsub(dif, abiased, b); \ + out0 = _mm_packs_epi32(_mm_srai_epi32(sum_l, s), _mm_srai_epi32(sum_h, s)); \ + out1 = _mm_packs_epi32(_mm_srai_epi32(dif_l, s), _mm_srai_epi32(dif_h, s)); \ + } + + // 8-bit interleave step (for transposes) + #define dct_interleave8(a, b) \ + tmp = a; \ + a = _mm_unpacklo_epi8(a, b); \ + b = _mm_unpackhi_epi8(tmp, b) + + // 16-bit interleave step (for transposes) + #define dct_interleave16(a, b) \ + tmp = a; \ + a = _mm_unpacklo_epi16(a, b); \ + b = _mm_unpackhi_epi16(tmp, b) + + #define dct_pass(bias,shift) \ + { \ + /* even part */ \ + dct_rot(t2e,t3e, row2,row6, rot0_0,rot0_1); \ + __m128i sum04 = _mm_add_epi16(row0, row4); \ + __m128i dif04 = _mm_sub_epi16(row0, row4); \ + dct_widen(t0e, sum04); \ + dct_widen(t1e, dif04); \ + dct_wadd(x0, t0e, t3e); \ + dct_wsub(x3, t0e, t3e); \ + dct_wadd(x1, t1e, t2e); \ + dct_wsub(x2, t1e, t2e); \ + /* odd part */ \ + dct_rot(y0o,y2o, row7,row3, rot2_0,rot2_1); \ + dct_rot(y1o,y3o, row5,row1, rot3_0,rot3_1); \ + __m128i sum17 = _mm_add_epi16(row1, row7); \ + __m128i sum35 = _mm_add_epi16(row3, row5); \ + dct_rot(y4o,y5o, sum17,sum35, rot1_0,rot1_1); \ + dct_wadd(x4, y0o, y4o); \ + dct_wadd(x5, y1o, y5o); \ + dct_wadd(x6, y2o, y5o); \ + dct_wadd(x7, y3o, y4o); \ + dct_bfly32o(row0,row7, x0,x7,bias,shift); \ + dct_bfly32o(row1,row6, x1,x6,bias,shift); \ + dct_bfly32o(row2,row5, x2,x5,bias,shift); \ + dct_bfly32o(row3,row4, x3,x4,bias,shift); \ + } + + __m128i rot0_0 = dct_const(stbi__f2f(0.5411961f), stbi__f2f(0.5411961f) + stbi__f2f(-1.847759065f)); + __m128i rot0_1 = dct_const(stbi__f2f(0.5411961f) + stbi__f2f( 0.765366865f), stbi__f2f(0.5411961f)); + __m128i rot1_0 = dct_const(stbi__f2f(1.175875602f) + stbi__f2f(-0.899976223f), stbi__f2f(1.175875602f)); + __m128i rot1_1 = dct_const(stbi__f2f(1.175875602f), stbi__f2f(1.175875602f) + stbi__f2f(-2.562915447f)); + __m128i rot2_0 = dct_const(stbi__f2f(-1.961570560f) + stbi__f2f( 0.298631336f), stbi__f2f(-1.961570560f)); + __m128i rot2_1 = dct_const(stbi__f2f(-1.961570560f), stbi__f2f(-1.961570560f) + stbi__f2f( 3.072711026f)); + __m128i rot3_0 = dct_const(stbi__f2f(-0.390180644f) + stbi__f2f( 2.053119869f), stbi__f2f(-0.390180644f)); + __m128i rot3_1 = dct_const(stbi__f2f(-0.390180644f), stbi__f2f(-0.390180644f) + stbi__f2f( 1.501321110f)); + + // rounding biases in column/row passes, see stbi__idct_block for explanation. + __m128i bias_0 = _mm_set1_epi32(512); + __m128i bias_1 = _mm_set1_epi32(65536 + (128<<17)); + + // load + row0 = _mm_load_si128((const __m128i *) (data + 0*8)); + row1 = _mm_load_si128((const __m128i *) (data + 1*8)); + row2 = _mm_load_si128((const __m128i *) (data + 2*8)); + row3 = _mm_load_si128((const __m128i *) (data + 3*8)); + row4 = _mm_load_si128((const __m128i *) (data + 4*8)); + row5 = _mm_load_si128((const __m128i *) (data + 5*8)); + row6 = _mm_load_si128((const __m128i *) (data + 6*8)); + row7 = _mm_load_si128((const __m128i *) (data + 7*8)); + + // column pass + dct_pass(bias_0, 10); + + { + // 16bit 8x8 transpose pass 1 + dct_interleave16(row0, row4); + dct_interleave16(row1, row5); + dct_interleave16(row2, row6); + dct_interleave16(row3, row7); + + // transpose pass 2 + dct_interleave16(row0, row2); + dct_interleave16(row1, row3); + dct_interleave16(row4, row6); + dct_interleave16(row5, row7); + + // transpose pass 3 + dct_interleave16(row0, row1); + dct_interleave16(row2, row3); + dct_interleave16(row4, row5); + dct_interleave16(row6, row7); + } + + // row pass + dct_pass(bias_1, 17); + + { + // pack + __m128i p0 = _mm_packus_epi16(row0, row1); // a0a1a2a3...a7b0b1b2b3...b7 + __m128i p1 = _mm_packus_epi16(row2, row3); + __m128i p2 = _mm_packus_epi16(row4, row5); + __m128i p3 = _mm_packus_epi16(row6, row7); + + // 8bit 8x8 transpose pass 1 + dct_interleave8(p0, p2); // a0e0a1e1... + dct_interleave8(p1, p3); // c0g0c1g1... + + // transpose pass 2 + dct_interleave8(p0, p1); // a0c0e0g0... + dct_interleave8(p2, p3); // b0d0f0h0... + + // transpose pass 3 + dct_interleave8(p0, p2); // a0b0c0d0... + dct_interleave8(p1, p3); // a4b4c4d4... + + // store + _mm_storel_epi64((__m128i *) out, p0); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p0, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p2); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p2, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p1); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p1, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p3); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p3, 0x4e)); + } + +#undef dct_const +#undef dct_rot +#undef dct_widen +#undef dct_wadd +#undef dct_wsub +#undef dct_bfly32o +#undef dct_interleave8 +#undef dct_interleave16 +#undef dct_pass +} + +#endif // STBI_SSE2 + +#ifdef STBI_NEON + +// NEON integer IDCT. should produce bit-identical +// results to the generic C version. +static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) +{ + int16x8_t row0, row1, row2, row3, row4, row5, row6, row7; + + int16x4_t rot0_0 = vdup_n_s16(stbi__f2f(0.5411961f)); + int16x4_t rot0_1 = vdup_n_s16(stbi__f2f(-1.847759065f)); + int16x4_t rot0_2 = vdup_n_s16(stbi__f2f( 0.765366865f)); + int16x4_t rot1_0 = vdup_n_s16(stbi__f2f( 1.175875602f)); + int16x4_t rot1_1 = vdup_n_s16(stbi__f2f(-0.899976223f)); + int16x4_t rot1_2 = vdup_n_s16(stbi__f2f(-2.562915447f)); + int16x4_t rot2_0 = vdup_n_s16(stbi__f2f(-1.961570560f)); + int16x4_t rot2_1 = vdup_n_s16(stbi__f2f(-0.390180644f)); + int16x4_t rot3_0 = vdup_n_s16(stbi__f2f( 0.298631336f)); + int16x4_t rot3_1 = vdup_n_s16(stbi__f2f( 2.053119869f)); + int16x4_t rot3_2 = vdup_n_s16(stbi__f2f( 3.072711026f)); + int16x4_t rot3_3 = vdup_n_s16(stbi__f2f( 1.501321110f)); + +#define dct_long_mul(out, inq, coeff) \ + int32x4_t out##_l = vmull_s16(vget_low_s16(inq), coeff); \ + int32x4_t out##_h = vmull_s16(vget_high_s16(inq), coeff) + +#define dct_long_mac(out, acc, inq, coeff) \ + int32x4_t out##_l = vmlal_s16(acc##_l, vget_low_s16(inq), coeff); \ + int32x4_t out##_h = vmlal_s16(acc##_h, vget_high_s16(inq), coeff) + +#define dct_widen(out, inq) \ + int32x4_t out##_l = vshll_n_s16(vget_low_s16(inq), 12); \ + int32x4_t out##_h = vshll_n_s16(vget_high_s16(inq), 12) + +// wide add +#define dct_wadd(out, a, b) \ + int32x4_t out##_l = vaddq_s32(a##_l, b##_l); \ + int32x4_t out##_h = vaddq_s32(a##_h, b##_h) + +// wide sub +#define dct_wsub(out, a, b) \ + int32x4_t out##_l = vsubq_s32(a##_l, b##_l); \ + int32x4_t out##_h = vsubq_s32(a##_h, b##_h) + +// butterfly a/b, then shift using "shiftop" by "s" and pack +#define dct_bfly32o(out0,out1, a,b,shiftop,s) \ + { \ + dct_wadd(sum, a, b); \ + dct_wsub(dif, a, b); \ + out0 = vcombine_s16(shiftop(sum_l, s), shiftop(sum_h, s)); \ + out1 = vcombine_s16(shiftop(dif_l, s), shiftop(dif_h, s)); \ + } + +#define dct_pass(shiftop, shift) \ + { \ + /* even part */ \ + int16x8_t sum26 = vaddq_s16(row2, row6); \ + dct_long_mul(p1e, sum26, rot0_0); \ + dct_long_mac(t2e, p1e, row6, rot0_1); \ + dct_long_mac(t3e, p1e, row2, rot0_2); \ + int16x8_t sum04 = vaddq_s16(row0, row4); \ + int16x8_t dif04 = vsubq_s16(row0, row4); \ + dct_widen(t0e, sum04); \ + dct_widen(t1e, dif04); \ + dct_wadd(x0, t0e, t3e); \ + dct_wsub(x3, t0e, t3e); \ + dct_wadd(x1, t1e, t2e); \ + dct_wsub(x2, t1e, t2e); \ + /* odd part */ \ + int16x8_t sum15 = vaddq_s16(row1, row5); \ + int16x8_t sum17 = vaddq_s16(row1, row7); \ + int16x8_t sum35 = vaddq_s16(row3, row5); \ + int16x8_t sum37 = vaddq_s16(row3, row7); \ + int16x8_t sumodd = vaddq_s16(sum17, sum35); \ + dct_long_mul(p5o, sumodd, rot1_0); \ + dct_long_mac(p1o, p5o, sum17, rot1_1); \ + dct_long_mac(p2o, p5o, sum35, rot1_2); \ + dct_long_mul(p3o, sum37, rot2_0); \ + dct_long_mul(p4o, sum15, rot2_1); \ + dct_wadd(sump13o, p1o, p3o); \ + dct_wadd(sump24o, p2o, p4o); \ + dct_wadd(sump23o, p2o, p3o); \ + dct_wadd(sump14o, p1o, p4o); \ + dct_long_mac(x4, sump13o, row7, rot3_0); \ + dct_long_mac(x5, sump24o, row5, rot3_1); \ + dct_long_mac(x6, sump23o, row3, rot3_2); \ + dct_long_mac(x7, sump14o, row1, rot3_3); \ + dct_bfly32o(row0,row7, x0,x7,shiftop,shift); \ + dct_bfly32o(row1,row6, x1,x6,shiftop,shift); \ + dct_bfly32o(row2,row5, x2,x5,shiftop,shift); \ + dct_bfly32o(row3,row4, x3,x4,shiftop,shift); \ + } + + // load + row0 = vld1q_s16(data + 0*8); + row1 = vld1q_s16(data + 1*8); + row2 = vld1q_s16(data + 2*8); + row3 = vld1q_s16(data + 3*8); + row4 = vld1q_s16(data + 4*8); + row5 = vld1q_s16(data + 5*8); + row6 = vld1q_s16(data + 6*8); + row7 = vld1q_s16(data + 7*8); + + // add DC bias + row0 = vaddq_s16(row0, vsetq_lane_s16(1024, vdupq_n_s16(0), 0)); + + // column pass + dct_pass(vrshrn_n_s32, 10); + + // 16bit 8x8 transpose + { +// these three map to a single VTRN.16, VTRN.32, and VSWP, respectively. +// whether compilers actually get this is another story, sadly. +#define dct_trn16(x, y) { int16x8x2_t t = vtrnq_s16(x, y); x = t.val[0]; y = t.val[1]; } +#define dct_trn32(x, y) { int32x4x2_t t = vtrnq_s32(vreinterpretq_s32_s16(x), vreinterpretq_s32_s16(y)); x = vreinterpretq_s16_s32(t.val[0]); y = vreinterpretq_s16_s32(t.val[1]); } +#define dct_trn64(x, y) { int16x8_t x0 = x; int16x8_t y0 = y; x = vcombine_s16(vget_low_s16(x0), vget_low_s16(y0)); y = vcombine_s16(vget_high_s16(x0), vget_high_s16(y0)); } + + // pass 1 + dct_trn16(row0, row1); // a0b0a2b2a4b4a6b6 + dct_trn16(row2, row3); + dct_trn16(row4, row5); + dct_trn16(row6, row7); + + // pass 2 + dct_trn32(row0, row2); // a0b0c0d0a4b4c4d4 + dct_trn32(row1, row3); + dct_trn32(row4, row6); + dct_trn32(row5, row7); + + // pass 3 + dct_trn64(row0, row4); // a0b0c0d0e0f0g0h0 + dct_trn64(row1, row5); + dct_trn64(row2, row6); + dct_trn64(row3, row7); + +#undef dct_trn16 +#undef dct_trn32 +#undef dct_trn64 + } + + // row pass + // vrshrn_n_s32 only supports shifts up to 16, we need + // 17. so do a non-rounding shift of 16 first then follow + // up with a rounding shift by 1. + dct_pass(vshrn_n_s32, 16); + + { + // pack and round + uint8x8_t p0 = vqrshrun_n_s16(row0, 1); + uint8x8_t p1 = vqrshrun_n_s16(row1, 1); + uint8x8_t p2 = vqrshrun_n_s16(row2, 1); + uint8x8_t p3 = vqrshrun_n_s16(row3, 1); + uint8x8_t p4 = vqrshrun_n_s16(row4, 1); + uint8x8_t p5 = vqrshrun_n_s16(row5, 1); + uint8x8_t p6 = vqrshrun_n_s16(row6, 1); + uint8x8_t p7 = vqrshrun_n_s16(row7, 1); + + // again, these can translate into one instruction, but often don't. +#define dct_trn8_8(x, y) { uint8x8x2_t t = vtrn_u8(x, y); x = t.val[0]; y = t.val[1]; } +#define dct_trn8_16(x, y) { uint16x4x2_t t = vtrn_u16(vreinterpret_u16_u8(x), vreinterpret_u16_u8(y)); x = vreinterpret_u8_u16(t.val[0]); y = vreinterpret_u8_u16(t.val[1]); } +#define dct_trn8_32(x, y) { uint32x2x2_t t = vtrn_u32(vreinterpret_u32_u8(x), vreinterpret_u32_u8(y)); x = vreinterpret_u8_u32(t.val[0]); y = vreinterpret_u8_u32(t.val[1]); } + + // sadly can't use interleaved stores here since we only write + // 8 bytes to each scan line! + + // 8x8 8-bit transpose pass 1 + dct_trn8_8(p0, p1); + dct_trn8_8(p2, p3); + dct_trn8_8(p4, p5); + dct_trn8_8(p6, p7); + + // pass 2 + dct_trn8_16(p0, p2); + dct_trn8_16(p1, p3); + dct_trn8_16(p4, p6); + dct_trn8_16(p5, p7); + + // pass 3 + dct_trn8_32(p0, p4); + dct_trn8_32(p1, p5); + dct_trn8_32(p2, p6); + dct_trn8_32(p3, p7); + + // store + vst1_u8(out, p0); out += out_stride; + vst1_u8(out, p1); out += out_stride; + vst1_u8(out, p2); out += out_stride; + vst1_u8(out, p3); out += out_stride; + vst1_u8(out, p4); out += out_stride; + vst1_u8(out, p5); out += out_stride; + vst1_u8(out, p6); out += out_stride; + vst1_u8(out, p7); + +#undef dct_trn8_8 +#undef dct_trn8_16 +#undef dct_trn8_32 + } + +#undef dct_long_mul +#undef dct_long_mac +#undef dct_widen +#undef dct_wadd +#undef dct_wsub +#undef dct_bfly32o +#undef dct_pass +} + +#endif // STBI_NEON + +#define STBI__MARKER_none 0xff +// if there's a pending marker from the entropy stream, return that +// otherwise, fetch from the stream and get a marker. if there's no +// marker, return 0xff, which is never a valid marker value +static stbi_uc stbi__get_marker(stbi__jpeg *j) +{ + stbi_uc x; + if (j->marker != STBI__MARKER_none) { x = j->marker; j->marker = STBI__MARKER_none; return x; } + x = stbi__get8(j->s); + if (x != 0xff) return STBI__MARKER_none; + while (x == 0xff) + x = stbi__get8(j->s); // consume repeated 0xff fill bytes + return x; +} + +// in each scan, we'll have scan_n components, and the order +// of the components is specified by order[] +#define STBI__RESTART(x) ((x) >= 0xd0 && (x) <= 0xd7) + +// after a restart interval, stbi__jpeg_reset the entropy decoder and +// the dc prediction +static void stbi__jpeg_reset(stbi__jpeg *j) +{ + j->code_bits = 0; + j->code_buffer = 0; + j->nomore = 0; + j->img_comp[0].dc_pred = j->img_comp[1].dc_pred = j->img_comp[2].dc_pred = j->img_comp[3].dc_pred = 0; + j->marker = STBI__MARKER_none; + j->todo = j->restart_interval ? j->restart_interval : 0x7fffffff; + j->eob_run = 0; + // no more than 1<<31 MCUs if no restart_interal? that's plenty safe, + // since we don't even allow 1<<30 pixels +} + +static int stbi__parse_entropy_coded_data(stbi__jpeg *z) +{ + stbi__jpeg_reset(z); + if (!z->progressive) { + if (z->scan_n == 1) { + int i,j; + STBI_SIMD_ALIGN(short, data[64]); + int n = z->order[0]; + // non-interleaved data, we just need to process one block at a time, + // in trivial scanline order + // number of blocks to do just depends on how many actual "pixels" this + // component has, independent of interleaved MCU blocking and such + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); + // every data block is an MCU, so countdown the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + // if it's NOT a restart, then just bail, so we get corrupt data + // rather than no data + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } else { // interleaved + int i,j,k,x,y; + STBI_SIMD_ALIGN(short, data[64]); + for (j=0; j < z->img_mcu_y; ++j) { + for (i=0; i < z->img_mcu_x; ++i) { + // scan an interleaved mcu... process scan_n components in order + for (k=0; k < z->scan_n; ++k) { + int n = z->order[k]; + // scan out an mcu's worth of this component; that's just determined + // by the basic H and V specified for the component + for (y=0; y < z->img_comp[n].v; ++y) { + for (x=0; x < z->img_comp[n].h; ++x) { + int x2 = (i*z->img_comp[n].h + x)*8; + int y2 = (j*z->img_comp[n].v + y)*8; + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*y2+x2, z->img_comp[n].w2, data); + } + } + } + // after all interleaved components, that's an interleaved MCU, + // so now count down the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } + } else { + if (z->scan_n == 1) { + int i,j; + int n = z->order[0]; + // non-interleaved data, we just need to process one block at a time, + // in trivial scanline order + // number of blocks to do just depends on how many actual "pixels" this + // component has, independent of interleaved MCU blocking and such + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); + if (z->spec_start == 0) { + if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) + return 0; + } else { + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block_prog_ac(z, data, &z->huff_ac[ha], z->fast_ac[ha])) + return 0; + } + // every data block is an MCU, so countdown the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } else { // interleaved + int i,j,k,x,y; + for (j=0; j < z->img_mcu_y; ++j) { + for (i=0; i < z->img_mcu_x; ++i) { + // scan an interleaved mcu... process scan_n components in order + for (k=0; k < z->scan_n; ++k) { + int n = z->order[k]; + // scan out an mcu's worth of this component; that's just determined + // by the basic H and V specified for the component + for (y=0; y < z->img_comp[n].v; ++y) { + for (x=0; x < z->img_comp[n].h; ++x) { + int x2 = (i*z->img_comp[n].h + x); + int y2 = (j*z->img_comp[n].v + y); + short *data = z->img_comp[n].coeff + 64 * (x2 + y2 * z->img_comp[n].coeff_w); + if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) + return 0; + } + } + } + // after all interleaved components, that's an interleaved MCU, + // so now count down the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } + } +} + +static void stbi__jpeg_dequantize(short *data, stbi__uint16 *dequant) +{ + int i; + for (i=0; i < 64; ++i) + data[i] *= dequant[i]; +} + +static void stbi__jpeg_finish(stbi__jpeg *z) +{ + if (z->progressive) { + // dequantize and idct the data + int i,j,n; + for (n=0; n < z->s->img_n; ++n) { + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); + stbi__jpeg_dequantize(data, z->dequant[z->img_comp[n].tq]); + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); + } + } + } + } +} + +static int stbi__process_marker(stbi__jpeg *z, int m) +{ + int L; + switch (m) { + case STBI__MARKER_none: // no marker found + return stbi__err("expected marker","Corrupt JPEG"); + + case 0xDD: // DRI - specify restart interval + if (stbi__get16be(z->s) != 4) return stbi__err("bad DRI len","Corrupt JPEG"); + z->restart_interval = stbi__get16be(z->s); + return 1; + + case 0xDB: // DQT - define quantization table + L = stbi__get16be(z->s)-2; + while (L > 0) { + int q = stbi__get8(z->s); + int p = q >> 4, sixteen = (p != 0); + int t = q & 15,i; + if (p != 0 && p != 1) return stbi__err("bad DQT type","Corrupt JPEG"); + if (t > 3) return stbi__err("bad DQT table","Corrupt JPEG"); + + for (i=0; i < 64; ++i) + z->dequant[t][stbi__jpeg_dezigzag[i]] = (stbi__uint16)(sixteen ? stbi__get16be(z->s) : stbi__get8(z->s)); + L -= (sixteen ? 129 : 65); + } + return L==0; + + case 0xC4: // DHT - define huffman table + L = stbi__get16be(z->s)-2; + while (L > 0) { + stbi_uc *v; + int sizes[16],i,n=0; + int q = stbi__get8(z->s); + int tc = q >> 4; + int th = q & 15; + if (tc > 1 || th > 3) return stbi__err("bad DHT header","Corrupt JPEG"); + for (i=0; i < 16; ++i) { + sizes[i] = stbi__get8(z->s); + n += sizes[i]; + } + if(n > 256) return stbi__err("bad DHT header","Corrupt JPEG"); // Loop over i < n would write past end of values! + L -= 17; + if (tc == 0) { + if (!stbi__build_huffman(z->huff_dc+th, sizes)) return 0; + v = z->huff_dc[th].values; + } else { + if (!stbi__build_huffman(z->huff_ac+th, sizes)) return 0; + v = z->huff_ac[th].values; + } + for (i=0; i < n; ++i) + v[i] = stbi__get8(z->s); + if (tc != 0) + stbi__build_fast_ac(z->fast_ac[th], z->huff_ac + th); + L -= n; + } + return L==0; + } + + // check for comment block or APP blocks + if ((m >= 0xE0 && m <= 0xEF) || m == 0xFE) { + L = stbi__get16be(z->s); + if (L < 2) { + if (m == 0xFE) + return stbi__err("bad COM len","Corrupt JPEG"); + else + return stbi__err("bad APP len","Corrupt JPEG"); + } + L -= 2; + + if (m == 0xE0 && L >= 5) { // JFIF APP0 segment + static const unsigned char tag[5] = {'J','F','I','F','\0'}; + int ok = 1; + int i; + for (i=0; i < 5; ++i) + if (stbi__get8(z->s) != tag[i]) + ok = 0; + L -= 5; + if (ok) + z->jfif = 1; + } else if (m == 0xEE && L >= 12) { // Adobe APP14 segment + static const unsigned char tag[6] = {'A','d','o','b','e','\0'}; + int ok = 1; + int i; + for (i=0; i < 6; ++i) + if (stbi__get8(z->s) != tag[i]) + ok = 0; + L -= 6; + if (ok) { + stbi__get8(z->s); // version + stbi__get16be(z->s); // flags0 + stbi__get16be(z->s); // flags1 + z->app14_color_transform = stbi__get8(z->s); // color transform + L -= 6; + } + } + + stbi__skip(z->s, L); + return 1; + } + + return stbi__err("unknown marker","Corrupt JPEG"); +} + +// after we see SOS +static int stbi__process_scan_header(stbi__jpeg *z) +{ + int i; + int Ls = stbi__get16be(z->s); + z->scan_n = stbi__get8(z->s); + if (z->scan_n < 1 || z->scan_n > 4 || z->scan_n > (int) z->s->img_n) return stbi__err("bad SOS component count","Corrupt JPEG"); + if (Ls != 6+2*z->scan_n) return stbi__err("bad SOS len","Corrupt JPEG"); + for (i=0; i < z->scan_n; ++i) { + int id = stbi__get8(z->s), which; + int q = stbi__get8(z->s); + for (which = 0; which < z->s->img_n; ++which) + if (z->img_comp[which].id == id) + break; + if (which == z->s->img_n) return 0; // no match + z->img_comp[which].hd = q >> 4; if (z->img_comp[which].hd > 3) return stbi__err("bad DC huff","Corrupt JPEG"); + z->img_comp[which].ha = q & 15; if (z->img_comp[which].ha > 3) return stbi__err("bad AC huff","Corrupt JPEG"); + z->order[i] = which; + } + + { + int aa; + z->spec_start = stbi__get8(z->s); + z->spec_end = stbi__get8(z->s); // should be 63, but might be 0 + aa = stbi__get8(z->s); + z->succ_high = (aa >> 4); + z->succ_low = (aa & 15); + if (z->progressive) { + if (z->spec_start > 63 || z->spec_end > 63 || z->spec_start > z->spec_end || z->succ_high > 13 || z->succ_low > 13) + return stbi__err("bad SOS", "Corrupt JPEG"); + } else { + if (z->spec_start != 0) return stbi__err("bad SOS","Corrupt JPEG"); + if (z->succ_high != 0 || z->succ_low != 0) return stbi__err("bad SOS","Corrupt JPEG"); + z->spec_end = 63; + } + } + + return 1; +} + +static int stbi__free_jpeg_components(stbi__jpeg *z, int ncomp, int why) +{ + int i; + for (i=0; i < ncomp; ++i) { + if (z->img_comp[i].raw_data) { + STBI_FREE(z->img_comp[i].raw_data); + z->img_comp[i].raw_data = NULL; + z->img_comp[i].data = NULL; + } + if (z->img_comp[i].raw_coeff) { + STBI_FREE(z->img_comp[i].raw_coeff); + z->img_comp[i].raw_coeff = 0; + z->img_comp[i].coeff = 0; + } + if (z->img_comp[i].linebuf) { + STBI_FREE(z->img_comp[i].linebuf); + z->img_comp[i].linebuf = NULL; + } + } + return why; +} + +static int stbi__process_frame_header(stbi__jpeg *z, int scan) +{ + stbi__context *s = z->s; + int Lf,p,i,q, h_max=1,v_max=1,c; + Lf = stbi__get16be(s); if (Lf < 11) return stbi__err("bad SOF len","Corrupt JPEG"); // JPEG + p = stbi__get8(s); if (p != 8) return stbi__err("only 8-bit","JPEG format not supported: 8-bit only"); // JPEG baseline + s->img_y = stbi__get16be(s); if (s->img_y == 0) return stbi__err("no header height", "JPEG format not supported: delayed height"); // Legal, but we don't handle it--but neither does IJG + s->img_x = stbi__get16be(s); if (s->img_x == 0) return stbi__err("0 width","Corrupt JPEG"); // JPEG requires + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + c = stbi__get8(s); + if (c != 3 && c != 1 && c != 4) return stbi__err("bad component count","Corrupt JPEG"); + s->img_n = c; + for (i=0; i < c; ++i) { + z->img_comp[i].data = NULL; + z->img_comp[i].linebuf = NULL; + } + + if (Lf != 8+3*s->img_n) return stbi__err("bad SOF len","Corrupt JPEG"); + + z->rgb = 0; + for (i=0; i < s->img_n; ++i) { + static const unsigned char rgb[3] = { 'R', 'G', 'B' }; + z->img_comp[i].id = stbi__get8(s); + if (s->img_n == 3 && z->img_comp[i].id == rgb[i]) + ++z->rgb; + q = stbi__get8(s); + z->img_comp[i].h = (q >> 4); if (!z->img_comp[i].h || z->img_comp[i].h > 4) return stbi__err("bad H","Corrupt JPEG"); + z->img_comp[i].v = q & 15; if (!z->img_comp[i].v || z->img_comp[i].v > 4) return stbi__err("bad V","Corrupt JPEG"); + z->img_comp[i].tq = stbi__get8(s); if (z->img_comp[i].tq > 3) return stbi__err("bad TQ","Corrupt JPEG"); + } + + if (scan != STBI__SCAN_load) return 1; + + if (!stbi__mad3sizes_valid(s->img_x, s->img_y, s->img_n, 0)) return stbi__err("too large", "Image too large to decode"); + + for (i=0; i < s->img_n; ++i) { + if (z->img_comp[i].h > h_max) h_max = z->img_comp[i].h; + if (z->img_comp[i].v > v_max) v_max = z->img_comp[i].v; + } + + // check that plane subsampling factors are integer ratios; our resamplers can't deal with fractional ratios + // and I've never seen a non-corrupted JPEG file actually use them + for (i=0; i < s->img_n; ++i) { + if (h_max % z->img_comp[i].h != 0) return stbi__err("bad H","Corrupt JPEG"); + if (v_max % z->img_comp[i].v != 0) return stbi__err("bad V","Corrupt JPEG"); + } + + // compute interleaved mcu info + z->img_h_max = h_max; + z->img_v_max = v_max; + z->img_mcu_w = h_max * 8; + z->img_mcu_h = v_max * 8; + // these sizes can't be more than 17 bits + z->img_mcu_x = (s->img_x + z->img_mcu_w-1) / z->img_mcu_w; + z->img_mcu_y = (s->img_y + z->img_mcu_h-1) / z->img_mcu_h; + + for (i=0; i < s->img_n; ++i) { + // number of effective pixels (e.g. for non-interleaved MCU) + z->img_comp[i].x = (s->img_x * z->img_comp[i].h + h_max-1) / h_max; + z->img_comp[i].y = (s->img_y * z->img_comp[i].v + v_max-1) / v_max; + // to simplify generation, we'll allocate enough memory to decode + // the bogus oversized data from using interleaved MCUs and their + // big blocks (e.g. a 16x16 iMCU on an image of width 33); we won't + // discard the extra data until colorspace conversion + // + // img_mcu_x, img_mcu_y: <=17 bits; comp[i].h and .v are <=4 (checked earlier) + // so these muls can't overflow with 32-bit ints (which we require) + z->img_comp[i].w2 = z->img_mcu_x * z->img_comp[i].h * 8; + z->img_comp[i].h2 = z->img_mcu_y * z->img_comp[i].v * 8; + z->img_comp[i].coeff = 0; + z->img_comp[i].raw_coeff = 0; + z->img_comp[i].linebuf = NULL; + z->img_comp[i].raw_data = stbi__malloc_mad2(z->img_comp[i].w2, z->img_comp[i].h2, 15); + if (z->img_comp[i].raw_data == NULL) + return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); + // align blocks for idct using mmx/sse + z->img_comp[i].data = (stbi_uc*) (((size_t) z->img_comp[i].raw_data + 15) & ~15); + if (z->progressive) { + // w2, h2 are multiples of 8 (see above) + z->img_comp[i].coeff_w = z->img_comp[i].w2 / 8; + z->img_comp[i].coeff_h = z->img_comp[i].h2 / 8; + z->img_comp[i].raw_coeff = stbi__malloc_mad3(z->img_comp[i].w2, z->img_comp[i].h2, sizeof(short), 15); + if (z->img_comp[i].raw_coeff == NULL) + return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); + z->img_comp[i].coeff = (short*) (((size_t) z->img_comp[i].raw_coeff + 15) & ~15); + } + } + + return 1; +} + +// use comparisons since in some cases we handle more than one case (e.g. SOF) +#define stbi__DNL(x) ((x) == 0xdc) +#define stbi__SOI(x) ((x) == 0xd8) +#define stbi__EOI(x) ((x) == 0xd9) +#define stbi__SOF(x) ((x) == 0xc0 || (x) == 0xc1 || (x) == 0xc2) +#define stbi__SOS(x) ((x) == 0xda) + +#define stbi__SOF_progressive(x) ((x) == 0xc2) + +static int stbi__decode_jpeg_header(stbi__jpeg *z, int scan) +{ + int m; + z->jfif = 0; + z->app14_color_transform = -1; // valid values are 0,1,2 + z->marker = STBI__MARKER_none; // initialize cached marker to empty + m = stbi__get_marker(z); + if (!stbi__SOI(m)) return stbi__err("no SOI","Corrupt JPEG"); + if (scan == STBI__SCAN_type) return 1; + m = stbi__get_marker(z); + while (!stbi__SOF(m)) { + if (!stbi__process_marker(z,m)) return 0; + m = stbi__get_marker(z); + while (m == STBI__MARKER_none) { + // some files have extra padding after their blocks, so ok, we'll scan + if (stbi__at_eof(z->s)) return stbi__err("no SOF", "Corrupt JPEG"); + m = stbi__get_marker(z); + } + } + z->progressive = stbi__SOF_progressive(m); + if (!stbi__process_frame_header(z, scan)) return 0; + return 1; +} + +static stbi_uc stbi__skip_jpeg_junk_at_end(stbi__jpeg *j) +{ + // some JPEGs have junk at end, skip over it but if we find what looks + // like a valid marker, resume there + while (!stbi__at_eof(j->s)) { + stbi_uc x = stbi__get8(j->s); + while (x == 0xff) { // might be a marker + if (stbi__at_eof(j->s)) return STBI__MARKER_none; + x = stbi__get8(j->s); + if (x != 0x00 && x != 0xff) { + // not a stuffed zero or lead-in to another marker, looks + // like an actual marker, return it + return x; + } + // stuffed zero has x=0 now which ends the loop, meaning we go + // back to regular scan loop. + // repeated 0xff keeps trying to read the next byte of the marker. + } + } + return STBI__MARKER_none; +} + +// decode image to YCbCr format +static int stbi__decode_jpeg_image(stbi__jpeg *j) +{ + int m; + for (m = 0; m < 4; m++) { + j->img_comp[m].raw_data = NULL; + j->img_comp[m].raw_coeff = NULL; + } + j->restart_interval = 0; + if (!stbi__decode_jpeg_header(j, STBI__SCAN_load)) return 0; + m = stbi__get_marker(j); + while (!stbi__EOI(m)) { + if (stbi__SOS(m)) { + if (!stbi__process_scan_header(j)) return 0; + if (!stbi__parse_entropy_coded_data(j)) return 0; + if (j->marker == STBI__MARKER_none ) { + j->marker = stbi__skip_jpeg_junk_at_end(j); + // if we reach eof without hitting a marker, stbi__get_marker() below will fail and we'll eventually return 0 + } + m = stbi__get_marker(j); + if (STBI__RESTART(m)) + m = stbi__get_marker(j); + } else if (stbi__DNL(m)) { + int Ld = stbi__get16be(j->s); + stbi__uint32 NL = stbi__get16be(j->s); + if (Ld != 4) return stbi__err("bad DNL len", "Corrupt JPEG"); + if (NL != j->s->img_y) return stbi__err("bad DNL height", "Corrupt JPEG"); + m = stbi__get_marker(j); + } else { + if (!stbi__process_marker(j, m)) return 1; + m = stbi__get_marker(j); + } + } + if (j->progressive) + stbi__jpeg_finish(j); + return 1; +} + +// static jfif-centered resampling (across block boundaries) + +typedef stbi_uc *(*resample_row_func)(stbi_uc *out, stbi_uc *in0, stbi_uc *in1, + int w, int hs); + +#define stbi__div4(x) ((stbi_uc) ((x) >> 2)) + +static stbi_uc *resample_row_1(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + STBI_NOTUSED(out); + STBI_NOTUSED(in_far); + STBI_NOTUSED(w); + STBI_NOTUSED(hs); + return in_near; +} + +static stbi_uc* stbi__resample_row_v_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate two samples vertically for every one in input + int i; + STBI_NOTUSED(hs); + for (i=0; i < w; ++i) + out[i] = stbi__div4(3*in_near[i] + in_far[i] + 2); + return out; +} + +static stbi_uc* stbi__resample_row_h_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate two samples horizontally for every one in input + int i; + stbi_uc *input = in_near; + + if (w == 1) { + // if only one sample, can't do any interpolation + out[0] = out[1] = input[0]; + return out; + } + + out[0] = input[0]; + out[1] = stbi__div4(input[0]*3 + input[1] + 2); + for (i=1; i < w-1; ++i) { + int n = 3*input[i]+2; + out[i*2+0] = stbi__div4(n+input[i-1]); + out[i*2+1] = stbi__div4(n+input[i+1]); + } + out[i*2+0] = stbi__div4(input[w-2]*3 + input[w-1] + 2); + out[i*2+1] = input[w-1]; + + STBI_NOTUSED(in_far); + STBI_NOTUSED(hs); + + return out; +} + +#define stbi__div16(x) ((stbi_uc) ((x) >> 4)) + +static stbi_uc *stbi__resample_row_hv_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate 2x2 samples for every one in input + int i,t0,t1; + if (w == 1) { + out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); + return out; + } + + t1 = 3*in_near[0] + in_far[0]; + out[0] = stbi__div4(t1+2); + for (i=1; i < w; ++i) { + t0 = t1; + t1 = 3*in_near[i]+in_far[i]; + out[i*2-1] = stbi__div16(3*t0 + t1 + 8); + out[i*2 ] = stbi__div16(3*t1 + t0 + 8); + } + out[w*2-1] = stbi__div4(t1+2); + + STBI_NOTUSED(hs); + + return out; +} + +#if defined(STBI_SSE2) || defined(STBI_NEON) +static stbi_uc *stbi__resample_row_hv_2_simd(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate 2x2 samples for every one in input + int i=0,t0,t1; + + if (w == 1) { + out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); + return out; + } + + t1 = 3*in_near[0] + in_far[0]; + // process groups of 8 pixels for as long as we can. + // note we can't handle the last pixel in a row in this loop + // because we need to handle the filter boundary conditions. + for (; i < ((w-1) & ~7); i += 8) { +#if defined(STBI_SSE2) + // load and perform the vertical filtering pass + // this uses 3*x + y = 4*x + (y - x) + __m128i zero = _mm_setzero_si128(); + __m128i farb = _mm_loadl_epi64((__m128i *) (in_far + i)); + __m128i nearb = _mm_loadl_epi64((__m128i *) (in_near + i)); + __m128i farw = _mm_unpacklo_epi8(farb, zero); + __m128i nearw = _mm_unpacklo_epi8(nearb, zero); + __m128i diff = _mm_sub_epi16(farw, nearw); + __m128i nears = _mm_slli_epi16(nearw, 2); + __m128i curr = _mm_add_epi16(nears, diff); // current row + + // horizontal filter works the same based on shifted vers of current + // row. "prev" is current row shifted right by 1 pixel; we need to + // insert the previous pixel value (from t1). + // "next" is current row shifted left by 1 pixel, with first pixel + // of next block of 8 pixels added in. + __m128i prv0 = _mm_slli_si128(curr, 2); + __m128i nxt0 = _mm_srli_si128(curr, 2); + __m128i prev = _mm_insert_epi16(prv0, t1, 0); + __m128i next = _mm_insert_epi16(nxt0, 3*in_near[i+8] + in_far[i+8], 7); + + // horizontal filter, polyphase implementation since it's convenient: + // even pixels = 3*cur + prev = cur*4 + (prev - cur) + // odd pixels = 3*cur + next = cur*4 + (next - cur) + // note the shared term. + __m128i bias = _mm_set1_epi16(8); + __m128i curs = _mm_slli_epi16(curr, 2); + __m128i prvd = _mm_sub_epi16(prev, curr); + __m128i nxtd = _mm_sub_epi16(next, curr); + __m128i curb = _mm_add_epi16(curs, bias); + __m128i even = _mm_add_epi16(prvd, curb); + __m128i odd = _mm_add_epi16(nxtd, curb); + + // interleave even and odd pixels, then undo scaling. + __m128i int0 = _mm_unpacklo_epi16(even, odd); + __m128i int1 = _mm_unpackhi_epi16(even, odd); + __m128i de0 = _mm_srli_epi16(int0, 4); + __m128i de1 = _mm_srli_epi16(int1, 4); + + // pack and write output + __m128i outv = _mm_packus_epi16(de0, de1); + _mm_storeu_si128((__m128i *) (out + i*2), outv); +#elif defined(STBI_NEON) + // load and perform the vertical filtering pass + // this uses 3*x + y = 4*x + (y - x) + uint8x8_t farb = vld1_u8(in_far + i); + uint8x8_t nearb = vld1_u8(in_near + i); + int16x8_t diff = vreinterpretq_s16_u16(vsubl_u8(farb, nearb)); + int16x8_t nears = vreinterpretq_s16_u16(vshll_n_u8(nearb, 2)); + int16x8_t curr = vaddq_s16(nears, diff); // current row + + // horizontal filter works the same based on shifted vers of current + // row. "prev" is current row shifted right by 1 pixel; we need to + // insert the previous pixel value (from t1). + // "next" is current row shifted left by 1 pixel, with first pixel + // of next block of 8 pixels added in. + int16x8_t prv0 = vextq_s16(curr, curr, 7); + int16x8_t nxt0 = vextq_s16(curr, curr, 1); + int16x8_t prev = vsetq_lane_s16(t1, prv0, 0); + int16x8_t next = vsetq_lane_s16(3*in_near[i+8] + in_far[i+8], nxt0, 7); + + // horizontal filter, polyphase implementation since it's convenient: + // even pixels = 3*cur + prev = cur*4 + (prev - cur) + // odd pixels = 3*cur + next = cur*4 + (next - cur) + // note the shared term. + int16x8_t curs = vshlq_n_s16(curr, 2); + int16x8_t prvd = vsubq_s16(prev, curr); + int16x8_t nxtd = vsubq_s16(next, curr); + int16x8_t even = vaddq_s16(curs, prvd); + int16x8_t odd = vaddq_s16(curs, nxtd); + + // undo scaling and round, then store with even/odd phases interleaved + uint8x8x2_t o; + o.val[0] = vqrshrun_n_s16(even, 4); + o.val[1] = vqrshrun_n_s16(odd, 4); + vst2_u8(out + i*2, o); +#endif + + // "previous" value for next iter + t1 = 3*in_near[i+7] + in_far[i+7]; + } + + t0 = t1; + t1 = 3*in_near[i] + in_far[i]; + out[i*2] = stbi__div16(3*t1 + t0 + 8); + + for (++i; i < w; ++i) { + t0 = t1; + t1 = 3*in_near[i]+in_far[i]; + out[i*2-1] = stbi__div16(3*t0 + t1 + 8); + out[i*2 ] = stbi__div16(3*t1 + t0 + 8); + } + out[w*2-1] = stbi__div4(t1+2); + + STBI_NOTUSED(hs); + + return out; +} +#endif + +static stbi_uc *stbi__resample_row_generic(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // resample with nearest-neighbor + int i,j; + STBI_NOTUSED(in_far); + for (i=0; i < w; ++i) + for (j=0; j < hs; ++j) + out[i*hs+j] = in_near[i]; + return out; +} + +// this is a reduced-precision calculation of YCbCr-to-RGB introduced +// to make sure the code produces the same results in both SIMD and scalar +#define stbi__float2fixed(x) (((int) ((x) * 4096.0f + 0.5f)) << 8) +static void stbi__YCbCr_to_RGB_row(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step) +{ + int i; + for (i=0; i < count; ++i) { + int y_fixed = (y[i] << 20) + (1<<19); // rounding + int r,g,b; + int cr = pcr[i] - 128; + int cb = pcb[i] - 128; + r = y_fixed + cr* stbi__float2fixed(1.40200f); + g = y_fixed + (cr*-stbi__float2fixed(0.71414f)) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); + b = y_fixed + cb* stbi__float2fixed(1.77200f); + r >>= 20; + g >>= 20; + b >>= 20; + if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } + if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } + if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } + out[0] = (stbi_uc)r; + out[1] = (stbi_uc)g; + out[2] = (stbi_uc)b; + out[3] = 255; + out += step; + } +} + +#if defined(STBI_SSE2) || defined(STBI_NEON) +static void stbi__YCbCr_to_RGB_simd(stbi_uc *out, stbi_uc const *y, stbi_uc const *pcb, stbi_uc const *pcr, int count, int step) +{ + int i = 0; + +#ifdef STBI_SSE2 + // step == 3 is pretty ugly on the final interleave, and i'm not convinced + // it's useful in practice (you wouldn't use it for textures, for example). + // so just accelerate step == 4 case. + if (step == 4) { + // this is a fairly straightforward implementation and not super-optimized. + __m128i signflip = _mm_set1_epi8(-0x80); + __m128i cr_const0 = _mm_set1_epi16( (short) ( 1.40200f*4096.0f+0.5f)); + __m128i cr_const1 = _mm_set1_epi16( - (short) ( 0.71414f*4096.0f+0.5f)); + __m128i cb_const0 = _mm_set1_epi16( - (short) ( 0.34414f*4096.0f+0.5f)); + __m128i cb_const1 = _mm_set1_epi16( (short) ( 1.77200f*4096.0f+0.5f)); + __m128i y_bias = _mm_set1_epi8((char) (unsigned char) 128); + __m128i xw = _mm_set1_epi16(255); // alpha channel + + for (; i+7 < count; i += 8) { + // load + __m128i y_bytes = _mm_loadl_epi64((__m128i *) (y+i)); + __m128i cr_bytes = _mm_loadl_epi64((__m128i *) (pcr+i)); + __m128i cb_bytes = _mm_loadl_epi64((__m128i *) (pcb+i)); + __m128i cr_biased = _mm_xor_si128(cr_bytes, signflip); // -128 + __m128i cb_biased = _mm_xor_si128(cb_bytes, signflip); // -128 + + // unpack to short (and left-shift cr, cb by 8) + __m128i yw = _mm_unpacklo_epi8(y_bias, y_bytes); + __m128i crw = _mm_unpacklo_epi8(_mm_setzero_si128(), cr_biased); + __m128i cbw = _mm_unpacklo_epi8(_mm_setzero_si128(), cb_biased); + + // color transform + __m128i yws = _mm_srli_epi16(yw, 4); + __m128i cr0 = _mm_mulhi_epi16(cr_const0, crw); + __m128i cb0 = _mm_mulhi_epi16(cb_const0, cbw); + __m128i cb1 = _mm_mulhi_epi16(cbw, cb_const1); + __m128i cr1 = _mm_mulhi_epi16(crw, cr_const1); + __m128i rws = _mm_add_epi16(cr0, yws); + __m128i gwt = _mm_add_epi16(cb0, yws); + __m128i bws = _mm_add_epi16(yws, cb1); + __m128i gws = _mm_add_epi16(gwt, cr1); + + // descale + __m128i rw = _mm_srai_epi16(rws, 4); + __m128i bw = _mm_srai_epi16(bws, 4); + __m128i gw = _mm_srai_epi16(gws, 4); + + // back to byte, set up for transpose + __m128i brb = _mm_packus_epi16(rw, bw); + __m128i gxb = _mm_packus_epi16(gw, xw); + + // transpose to interleave channels + __m128i t0 = _mm_unpacklo_epi8(brb, gxb); + __m128i t1 = _mm_unpackhi_epi8(brb, gxb); + __m128i o0 = _mm_unpacklo_epi16(t0, t1); + __m128i o1 = _mm_unpackhi_epi16(t0, t1); + + // store + _mm_storeu_si128((__m128i *) (out + 0), o0); + _mm_storeu_si128((__m128i *) (out + 16), o1); + out += 32; + } + } +#endif + +#ifdef STBI_NEON + // in this version, step=3 support would be easy to add. but is there demand? + if (step == 4) { + // this is a fairly straightforward implementation and not super-optimized. + uint8x8_t signflip = vdup_n_u8(0x80); + int16x8_t cr_const0 = vdupq_n_s16( (short) ( 1.40200f*4096.0f+0.5f)); + int16x8_t cr_const1 = vdupq_n_s16( - (short) ( 0.71414f*4096.0f+0.5f)); + int16x8_t cb_const0 = vdupq_n_s16( - (short) ( 0.34414f*4096.0f+0.5f)); + int16x8_t cb_const1 = vdupq_n_s16( (short) ( 1.77200f*4096.0f+0.5f)); + + for (; i+7 < count; i += 8) { + // load + uint8x8_t y_bytes = vld1_u8(y + i); + uint8x8_t cr_bytes = vld1_u8(pcr + i); + uint8x8_t cb_bytes = vld1_u8(pcb + i); + int8x8_t cr_biased = vreinterpret_s8_u8(vsub_u8(cr_bytes, signflip)); + int8x8_t cb_biased = vreinterpret_s8_u8(vsub_u8(cb_bytes, signflip)); + + // expand to s16 + int16x8_t yws = vreinterpretq_s16_u16(vshll_n_u8(y_bytes, 4)); + int16x8_t crw = vshll_n_s8(cr_biased, 7); + int16x8_t cbw = vshll_n_s8(cb_biased, 7); + + // color transform + int16x8_t cr0 = vqdmulhq_s16(crw, cr_const0); + int16x8_t cb0 = vqdmulhq_s16(cbw, cb_const0); + int16x8_t cr1 = vqdmulhq_s16(crw, cr_const1); + int16x8_t cb1 = vqdmulhq_s16(cbw, cb_const1); + int16x8_t rws = vaddq_s16(yws, cr0); + int16x8_t gws = vaddq_s16(vaddq_s16(yws, cb0), cr1); + int16x8_t bws = vaddq_s16(yws, cb1); + + // undo scaling, round, convert to byte + uint8x8x4_t o; + o.val[0] = vqrshrun_n_s16(rws, 4); + o.val[1] = vqrshrun_n_s16(gws, 4); + o.val[2] = vqrshrun_n_s16(bws, 4); + o.val[3] = vdup_n_u8(255); + + // store, interleaving r/g/b/a + vst4_u8(out, o); + out += 8*4; + } + } +#endif + + for (; i < count; ++i) { + int y_fixed = (y[i] << 20) + (1<<19); // rounding + int r,g,b; + int cr = pcr[i] - 128; + int cb = pcb[i] - 128; + r = y_fixed + cr* stbi__float2fixed(1.40200f); + g = y_fixed + cr*-stbi__float2fixed(0.71414f) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); + b = y_fixed + cb* stbi__float2fixed(1.77200f); + r >>= 20; + g >>= 20; + b >>= 20; + if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } + if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } + if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } + out[0] = (stbi_uc)r; + out[1] = (stbi_uc)g; + out[2] = (stbi_uc)b; + out[3] = 255; + out += step; + } +} +#endif + +// set up the kernels +static void stbi__setup_jpeg(stbi__jpeg *j) +{ + j->idct_block_kernel = stbi__idct_block; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_row; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2; + +#ifdef STBI_SSE2 + if (stbi__sse2_available()) { + j->idct_block_kernel = stbi__idct_simd; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; + } +#endif + +#ifdef STBI_NEON + j->idct_block_kernel = stbi__idct_simd; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; +#endif +} + +// clean up the temporary component buffers +static void stbi__cleanup_jpeg(stbi__jpeg *j) +{ + stbi__free_jpeg_components(j, j->s->img_n, 0); +} + +typedef struct +{ + resample_row_func resample; + stbi_uc *line0,*line1; + int hs,vs; // expansion factor in each axis + int w_lores; // horizontal pixels pre-expansion + int ystep; // how far through vertical expansion we are + int ypos; // which pre-expansion row we're on +} stbi__resample; + +// fast 0..255 * 0..255 => 0..255 rounded multiplication +static stbi_uc stbi__blinn_8x8(stbi_uc x, stbi_uc y) +{ + unsigned int t = x*y + 128; + return (stbi_uc) ((t + (t >>8)) >> 8); +} + +static stbi_uc *load_jpeg_image(stbi__jpeg *z, int *out_x, int *out_y, int *comp, int req_comp) +{ + int n, decode_n, is_rgb; + z->s->img_n = 0; // make stbi__cleanup_jpeg safe + + // validate req_comp + if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); + + // load a jpeg image from whichever source, but leave in YCbCr format + if (!stbi__decode_jpeg_image(z)) { stbi__cleanup_jpeg(z); return NULL; } + + // determine actual number of components to generate + n = req_comp ? req_comp : z->s->img_n >= 3 ? 3 : 1; + + is_rgb = z->s->img_n == 3 && (z->rgb == 3 || (z->app14_color_transform == 0 && !z->jfif)); + + if (z->s->img_n == 3 && n < 3 && !is_rgb) + decode_n = 1; + else + decode_n = z->s->img_n; + + // nothing to do if no components requested; check this now to avoid + // accessing uninitialized coutput[0] later + if (decode_n <= 0) { stbi__cleanup_jpeg(z); return NULL; } + + // resample and color-convert + { + int k; + unsigned int i,j; + stbi_uc *output; + stbi_uc *coutput[4] = { NULL, NULL, NULL, NULL }; + + stbi__resample res_comp[4]; + + for (k=0; k < decode_n; ++k) { + stbi__resample *r = &res_comp[k]; + + // allocate line buffer big enough for upsampling off the edges + // with upsample factor of 4 + z->img_comp[k].linebuf = (stbi_uc *) stbi__malloc(z->s->img_x + 3); + if (!z->img_comp[k].linebuf) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } + + r->hs = z->img_h_max / z->img_comp[k].h; + r->vs = z->img_v_max / z->img_comp[k].v; + r->ystep = r->vs >> 1; + r->w_lores = (z->s->img_x + r->hs-1) / r->hs; + r->ypos = 0; + r->line0 = r->line1 = z->img_comp[k].data; + + if (r->hs == 1 && r->vs == 1) r->resample = resample_row_1; + else if (r->hs == 1 && r->vs == 2) r->resample = stbi__resample_row_v_2; + else if (r->hs == 2 && r->vs == 1) r->resample = stbi__resample_row_h_2; + else if (r->hs == 2 && r->vs == 2) r->resample = z->resample_row_hv_2_kernel; + else r->resample = stbi__resample_row_generic; + } + + // can't error after this so, this is safe + output = (stbi_uc *) stbi__malloc_mad3(n, z->s->img_x, z->s->img_y, 1); + if (!output) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } + + // now go ahead and resample + for (j=0; j < z->s->img_y; ++j) { + stbi_uc *out = output + n * z->s->img_x * j; + for (k=0; k < decode_n; ++k) { + stbi__resample *r = &res_comp[k]; + int y_bot = r->ystep >= (r->vs >> 1); + coutput[k] = r->resample(z->img_comp[k].linebuf, + y_bot ? r->line1 : r->line0, + y_bot ? r->line0 : r->line1, + r->w_lores, r->hs); + if (++r->ystep >= r->vs) { + r->ystep = 0; + r->line0 = r->line1; + if (++r->ypos < z->img_comp[k].y) + r->line1 += z->img_comp[k].w2; + } + } + if (n >= 3) { + stbi_uc *y = coutput[0]; + if (z->s->img_n == 3) { + if (is_rgb) { + for (i=0; i < z->s->img_x; ++i) { + out[0] = y[i]; + out[1] = coutput[1][i]; + out[2] = coutput[2][i]; + out[3] = 255; + out += n; + } + } else { + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + } + } else if (z->s->img_n == 4) { + if (z->app14_color_transform == 0) { // CMYK + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + out[0] = stbi__blinn_8x8(coutput[0][i], m); + out[1] = stbi__blinn_8x8(coutput[1][i], m); + out[2] = stbi__blinn_8x8(coutput[2][i], m); + out[3] = 255; + out += n; + } + } else if (z->app14_color_transform == 2) { // YCCK + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + out[0] = stbi__blinn_8x8(255 - out[0], m); + out[1] = stbi__blinn_8x8(255 - out[1], m); + out[2] = stbi__blinn_8x8(255 - out[2], m); + out += n; + } + } else { // YCbCr + alpha? Ignore the fourth channel for now + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + } + } else + for (i=0; i < z->s->img_x; ++i) { + out[0] = out[1] = out[2] = y[i]; + out[3] = 255; // not used if n==3 + out += n; + } + } else { + if (is_rgb) { + if (n == 1) + for (i=0; i < z->s->img_x; ++i) + *out++ = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); + else { + for (i=0; i < z->s->img_x; ++i, out += 2) { + out[0] = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); + out[1] = 255; + } + } + } else if (z->s->img_n == 4 && z->app14_color_transform == 0) { + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + stbi_uc r = stbi__blinn_8x8(coutput[0][i], m); + stbi_uc g = stbi__blinn_8x8(coutput[1][i], m); + stbi_uc b = stbi__blinn_8x8(coutput[2][i], m); + out[0] = stbi__compute_y(r, g, b); + out[1] = 255; + out += n; + } + } else if (z->s->img_n == 4 && z->app14_color_transform == 2) { + for (i=0; i < z->s->img_x; ++i) { + out[0] = stbi__blinn_8x8(255 - coutput[0][i], coutput[3][i]); + out[1] = 255; + out += n; + } + } else { + stbi_uc *y = coutput[0]; + if (n == 1) + for (i=0; i < z->s->img_x; ++i) out[i] = y[i]; + else + for (i=0; i < z->s->img_x; ++i) { *out++ = y[i]; *out++ = 255; } + } + } + } + stbi__cleanup_jpeg(z); + *out_x = z->s->img_x; + *out_y = z->s->img_y; + if (comp) *comp = z->s->img_n >= 3 ? 3 : 1; // report original components, not output + return output; + } +} + +static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + unsigned char* result; + stbi__jpeg* j = (stbi__jpeg*) stbi__malloc(sizeof(stbi__jpeg)); + if (!j) return stbi__errpuc("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + STBI_NOTUSED(ri); + j->s = s; + stbi__setup_jpeg(j); + result = load_jpeg_image(j, x,y,comp,req_comp); + STBI_FREE(j); + return result; +} + +static int stbi__jpeg_test(stbi__context *s) +{ + int r; + stbi__jpeg* j = (stbi__jpeg*)stbi__malloc(sizeof(stbi__jpeg)); + if (!j) return stbi__err("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + j->s = s; + stbi__setup_jpeg(j); + r = stbi__decode_jpeg_header(j, STBI__SCAN_type); + stbi__rewind(s); + STBI_FREE(j); + return r; +} + +static int stbi__jpeg_info_raw(stbi__jpeg *j, int *x, int *y, int *comp) +{ + if (!stbi__decode_jpeg_header(j, STBI__SCAN_header)) { + stbi__rewind( j->s ); + return 0; + } + if (x) *x = j->s->img_x; + if (y) *y = j->s->img_y; + if (comp) *comp = j->s->img_n >= 3 ? 3 : 1; + return 1; +} + +static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp) +{ + int result; + stbi__jpeg* j = (stbi__jpeg*) (stbi__malloc(sizeof(stbi__jpeg))); + if (!j) return stbi__err("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + j->s = s; + result = stbi__jpeg_info_raw(j, x, y, comp); + STBI_FREE(j); + return result; +} +#endif + +// public domain zlib decode v0.2 Sean Barrett 2006-11-18 +// simple implementation +// - all input must be provided in an upfront buffer +// - all output is written to a single output buffer (can malloc/realloc) +// performance +// - fast huffman + +#ifndef STBI_NO_ZLIB + +// fast-way is faster to check than jpeg huffman, but slow way is slower +#define STBI__ZFAST_BITS 9 // accelerate all cases in default tables +#define STBI__ZFAST_MASK ((1 << STBI__ZFAST_BITS) - 1) +#define STBI__ZNSYMS 288 // number of symbols in literal/length alphabet + +// zlib-style huffman encoding +// (jpegs packs from left, zlib from right, so can't share code) +typedef struct +{ + stbi__uint16 fast[1 << STBI__ZFAST_BITS]; + stbi__uint16 firstcode[16]; + int maxcode[17]; + stbi__uint16 firstsymbol[16]; + stbi_uc size[STBI__ZNSYMS]; + stbi__uint16 value[STBI__ZNSYMS]; +} stbi__zhuffman; + +stbi_inline static int stbi__bitreverse16(int n) +{ + n = ((n & 0xAAAA) >> 1) | ((n & 0x5555) << 1); + n = ((n & 0xCCCC) >> 2) | ((n & 0x3333) << 2); + n = ((n & 0xF0F0) >> 4) | ((n & 0x0F0F) << 4); + n = ((n & 0xFF00) >> 8) | ((n & 0x00FF) << 8); + return n; +} + +stbi_inline static int stbi__bit_reverse(int v, int bits) +{ + STBI_ASSERT(bits <= 16); + // to bit reverse n bits, reverse 16 and shift + // e.g. 11 bits, bit reverse and shift away 5 + return stbi__bitreverse16(v) >> (16-bits); +} + +static int stbi__zbuild_huffman(stbi__zhuffman *z, const stbi_uc *sizelist, int num) +{ + int i,k=0; + int code, next_code[16], sizes[17]; + + // DEFLATE spec for generating codes + memset(sizes, 0, sizeof(sizes)); + memset(z->fast, 0, sizeof(z->fast)); + for (i=0; i < num; ++i) + ++sizes[sizelist[i]]; + sizes[0] = 0; + for (i=1; i < 16; ++i) + if (sizes[i] > (1 << i)) + return stbi__err("bad sizes", "Corrupt PNG"); + code = 0; + for (i=1; i < 16; ++i) { + next_code[i] = code; + z->firstcode[i] = (stbi__uint16) code; + z->firstsymbol[i] = (stbi__uint16) k; + code = (code + sizes[i]); + if (sizes[i]) + if (code-1 >= (1 << i)) return stbi__err("bad codelengths","Corrupt PNG"); + z->maxcode[i] = code << (16-i); // preshift for inner loop + code <<= 1; + k += sizes[i]; + } + z->maxcode[16] = 0x10000; // sentinel + for (i=0; i < num; ++i) { + int s = sizelist[i]; + if (s) { + int c = next_code[s] - z->firstcode[s] + z->firstsymbol[s]; + stbi__uint16 fastv = (stbi__uint16) ((s << 9) | i); + z->size [c] = (stbi_uc ) s; + z->value[c] = (stbi__uint16) i; + if (s <= STBI__ZFAST_BITS) { + int j = stbi__bit_reverse(next_code[s],s); + while (j < (1 << STBI__ZFAST_BITS)) { + z->fast[j] = fastv; + j += (1 << s); + } + } + ++next_code[s]; + } + } + return 1; +} + +// zlib-from-memory implementation for PNG reading +// because PNG allows splitting the zlib stream arbitrarily, +// and it's annoying structurally to have PNG call ZLIB call PNG, +// we require PNG read all the IDATs and combine them into a single +// memory buffer + +typedef struct +{ + stbi_uc *zbuffer, *zbuffer_end; + int num_bits; + int hit_zeof_once; + stbi__uint32 code_buffer; + + char *zout; + char *zout_start; + char *zout_end; + int z_expandable; + + stbi__zhuffman z_length, z_distance; +} stbi__zbuf; + +stbi_inline static int stbi__zeof(stbi__zbuf *z) +{ + return (z->zbuffer >= z->zbuffer_end); +} + +stbi_inline static stbi_uc stbi__zget8(stbi__zbuf *z) +{ + return stbi__zeof(z) ? 0 : *z->zbuffer++; +} + +static void stbi__fill_bits(stbi__zbuf *z) +{ + do { + if (z->code_buffer >= (1U << z->num_bits)) { + z->zbuffer = z->zbuffer_end; /* treat this as EOF so we fail. */ + return; + } + z->code_buffer |= (unsigned int) stbi__zget8(z) << z->num_bits; + z->num_bits += 8; + } while (z->num_bits <= 24); +} + +stbi_inline static unsigned int stbi__zreceive(stbi__zbuf *z, int n) +{ + unsigned int k; + if (z->num_bits < n) stbi__fill_bits(z); + k = z->code_buffer & ((1 << n) - 1); + z->code_buffer >>= n; + z->num_bits -= n; + return k; +} + +static int stbi__zhuffman_decode_slowpath(stbi__zbuf *a, stbi__zhuffman *z) +{ + int b,s,k; + // not resolved by fast table, so compute it the slow way + // use jpeg approach, which requires MSbits at top + k = stbi__bit_reverse(a->code_buffer, 16); + for (s=STBI__ZFAST_BITS+1; ; ++s) + if (k < z->maxcode[s]) + break; + if (s >= 16) return -1; // invalid code! + // code size is s, so: + b = (k >> (16-s)) - z->firstcode[s] + z->firstsymbol[s]; + if (b >= STBI__ZNSYMS) return -1; // some data was corrupt somewhere! + if (z->size[b] != s) return -1; // was originally an assert, but report failure instead. + a->code_buffer >>= s; + a->num_bits -= s; + return z->value[b]; +} + +stbi_inline static int stbi__zhuffman_decode(stbi__zbuf *a, stbi__zhuffman *z) +{ + int b,s; + if (a->num_bits < 16) { + if (stbi__zeof(a)) { + if (!a->hit_zeof_once) { + // This is the first time we hit eof, insert 16 extra padding btis + // to allow us to keep going; if we actually consume any of them + // though, that is invalid data. This is caught later. + a->hit_zeof_once = 1; + a->num_bits += 16; // add 16 implicit zero bits + } else { + // We already inserted our extra 16 padding bits and are again + // out, this stream is actually prematurely terminated. + return -1; + } + } else { + stbi__fill_bits(a); + } + } + b = z->fast[a->code_buffer & STBI__ZFAST_MASK]; + if (b) { + s = b >> 9; + a->code_buffer >>= s; + a->num_bits -= s; + return b & 511; + } + return stbi__zhuffman_decode_slowpath(a, z); +} + +static int stbi__zexpand(stbi__zbuf *z, char *zout, int n) // need to make room for n bytes +{ + char *q; + unsigned int cur, limit, old_limit; + z->zout = zout; + if (!z->z_expandable) return stbi__err("output buffer limit","Corrupt PNG"); + cur = (unsigned int) (z->zout - z->zout_start); + limit = old_limit = (unsigned) (z->zout_end - z->zout_start); + if (UINT_MAX - cur < (unsigned) n) return stbi__err("outofmem", "Out of memory"); + while (cur + n > limit) { + if(limit > UINT_MAX / 2) return stbi__err("outofmem", "Out of memory"); + limit *= 2; + } + q = (char *) STBI_REALLOC_SIZED(z->zout_start, old_limit, limit); + STBI_NOTUSED(old_limit); + if (q == NULL) return stbi__err("outofmem", "Out of memory"); + z->zout_start = q; + z->zout = q + cur; + z->zout_end = q + limit; + return 1; +} + +static const int stbi__zlength_base[31] = { + 3,4,5,6,7,8,9,10,11,13, + 15,17,19,23,27,31,35,43,51,59, + 67,83,99,115,131,163,195,227,258,0,0 }; + +static const int stbi__zlength_extra[31]= +{ 0,0,0,0,0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4,5,5,5,5,0,0,0 }; + +static const int stbi__zdist_base[32] = { 1,2,3,4,5,7,9,13,17,25,33,49,65,97,129,193, +257,385,513,769,1025,1537,2049,3073,4097,6145,8193,12289,16385,24577,0,0}; + +static const int stbi__zdist_extra[32] = +{ 0,0,0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13}; + +static int stbi__parse_huffman_block(stbi__zbuf *a) +{ + char *zout = a->zout; + for(;;) { + int z = stbi__zhuffman_decode(a, &a->z_length); + if (z < 256) { + if (z < 0) return stbi__err("bad huffman code","Corrupt PNG"); // error in huffman codes + if (zout >= a->zout_end) { + if (!stbi__zexpand(a, zout, 1)) return 0; + zout = a->zout; + } + *zout++ = (char) z; + } else { + stbi_uc *p; + int len,dist; + if (z == 256) { + a->zout = zout; + if (a->hit_zeof_once && a->num_bits < 16) { + // The first time we hit zeof, we inserted 16 extra zero bits into our bit + // buffer so the decoder can just do its speculative decoding. But if we + // actually consumed any of those bits (which is the case when num_bits < 16), + // the stream actually read past the end so it is malformed. + return stbi__err("unexpected end","Corrupt PNG"); + } + return 1; + } + if (z >= 286) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, length codes 286 and 287 must not appear in compressed data + z -= 257; + len = stbi__zlength_base[z]; + if (stbi__zlength_extra[z]) len += stbi__zreceive(a, stbi__zlength_extra[z]); + z = stbi__zhuffman_decode(a, &a->z_distance); + if (z < 0 || z >= 30) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, distance codes 30 and 31 must not appear in compressed data + dist = stbi__zdist_base[z]; + if (stbi__zdist_extra[z]) dist += stbi__zreceive(a, stbi__zdist_extra[z]); + if (zout - a->zout_start < dist) return stbi__err("bad dist","Corrupt PNG"); + if (len > a->zout_end - zout) { + if (!stbi__zexpand(a, zout, len)) return 0; + zout = a->zout; + } + p = (stbi_uc *) (zout - dist); + if (dist == 1) { // run of one byte; common in images. + stbi_uc v = *p; + if (len) { do *zout++ = v; while (--len); } + } else { + if (len) { do *zout++ = *p++; while (--len); } + } + } + } +} + +static int stbi__compute_huffman_codes(stbi__zbuf *a) +{ + static const stbi_uc length_dezigzag[19] = { 16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15 }; + stbi__zhuffman z_codelength; + stbi_uc lencodes[286+32+137];//padding for maximum single op + stbi_uc codelength_sizes[19]; + int i,n; + + int hlit = stbi__zreceive(a,5) + 257; + int hdist = stbi__zreceive(a,5) + 1; + int hclen = stbi__zreceive(a,4) + 4; + int ntot = hlit + hdist; + + memset(codelength_sizes, 0, sizeof(codelength_sizes)); + for (i=0; i < hclen; ++i) { + int s = stbi__zreceive(a,3); + codelength_sizes[length_dezigzag[i]] = (stbi_uc) s; + } + if (!stbi__zbuild_huffman(&z_codelength, codelength_sizes, 19)) return 0; + + n = 0; + while (n < ntot) { + int c = stbi__zhuffman_decode(a, &z_codelength); + if (c < 0 || c >= 19) return stbi__err("bad codelengths", "Corrupt PNG"); + if (c < 16) + lencodes[n++] = (stbi_uc) c; + else { + stbi_uc fill = 0; + if (c == 16) { + c = stbi__zreceive(a,2)+3; + if (n == 0) return stbi__err("bad codelengths", "Corrupt PNG"); + fill = lencodes[n-1]; + } else if (c == 17) { + c = stbi__zreceive(a,3)+3; + } else if (c == 18) { + c = stbi__zreceive(a,7)+11; + } else { + return stbi__err("bad codelengths", "Corrupt PNG"); + } + if (ntot - n < c) return stbi__err("bad codelengths", "Corrupt PNG"); + memset(lencodes+n, fill, c); + n += c; + } + } + if (n != ntot) return stbi__err("bad codelengths","Corrupt PNG"); + if (!stbi__zbuild_huffman(&a->z_length, lencodes, hlit)) return 0; + if (!stbi__zbuild_huffman(&a->z_distance, lencodes+hlit, hdist)) return 0; + return 1; +} + +static int stbi__parse_uncompressed_block(stbi__zbuf *a) +{ + stbi_uc header[4]; + int len,nlen,k; + if (a->num_bits & 7) + stbi__zreceive(a, a->num_bits & 7); // discard + // drain the bit-packed data into header + k = 0; + while (a->num_bits > 0) { + header[k++] = (stbi_uc) (a->code_buffer & 255); // suppress MSVC run-time check + a->code_buffer >>= 8; + a->num_bits -= 8; + } + if (a->num_bits < 0) return stbi__err("zlib corrupt","Corrupt PNG"); + // now fill header the normal way + while (k < 4) + header[k++] = stbi__zget8(a); + len = header[1] * 256 + header[0]; + nlen = header[3] * 256 + header[2]; + if (nlen != (len ^ 0xffff)) return stbi__err("zlib corrupt","Corrupt PNG"); + if (a->zbuffer + len > a->zbuffer_end) return stbi__err("read past buffer","Corrupt PNG"); + if (a->zout + len > a->zout_end) + if (!stbi__zexpand(a, a->zout, len)) return 0; + memcpy(a->zout, a->zbuffer, len); + a->zbuffer += len; + a->zout += len; + return 1; +} + +static int stbi__parse_zlib_header(stbi__zbuf *a) +{ + int cmf = stbi__zget8(a); + int cm = cmf & 15; + /* int cinfo = cmf >> 4; */ + int flg = stbi__zget8(a); + if (stbi__zeof(a)) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec + if ((cmf*256+flg) % 31 != 0) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec + if (flg & 32) return stbi__err("no preset dict","Corrupt PNG"); // preset dictionary not allowed in png + if (cm != 8) return stbi__err("bad compression","Corrupt PNG"); // DEFLATE required for png + // window = 1 << (8 + cinfo)... but who cares, we fully buffer output + return 1; +} + +static const stbi_uc stbi__zdefault_length[STBI__ZNSYMS] = +{ + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7, 7,7,7,7,7,7,7,7,8,8,8,8,8,8,8,8 +}; +static const stbi_uc stbi__zdefault_distance[32] = +{ + 5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5 +}; +/* +Init algorithm: +{ + int i; // use <= to match clearly with spec + for (i=0; i <= 143; ++i) stbi__zdefault_length[i] = 8; + for ( ; i <= 255; ++i) stbi__zdefault_length[i] = 9; + for ( ; i <= 279; ++i) stbi__zdefault_length[i] = 7; + for ( ; i <= 287; ++i) stbi__zdefault_length[i] = 8; + + for (i=0; i <= 31; ++i) stbi__zdefault_distance[i] = 5; +} +*/ + +static int stbi__parse_zlib(stbi__zbuf *a, int parse_header) +{ + int final, type; + if (parse_header) + if (!stbi__parse_zlib_header(a)) return 0; + a->num_bits = 0; + a->code_buffer = 0; + a->hit_zeof_once = 0; + do { + final = stbi__zreceive(a,1); + type = stbi__zreceive(a,2); + if (type == 0) { + if (!stbi__parse_uncompressed_block(a)) return 0; + } else if (type == 3) { + return 0; + } else { + if (type == 1) { + // use fixed code lengths + if (!stbi__zbuild_huffman(&a->z_length , stbi__zdefault_length , STBI__ZNSYMS)) return 0; + if (!stbi__zbuild_huffman(&a->z_distance, stbi__zdefault_distance, 32)) return 0; + } else { + if (!stbi__compute_huffman_codes(a)) return 0; + } + if (!stbi__parse_huffman_block(a)) return 0; + } + } while (!final); + return 1; +} + +static int stbi__do_zlib(stbi__zbuf *a, char *obuf, int olen, int exp, int parse_header) +{ + a->zout_start = obuf; + a->zout = obuf; + a->zout_end = obuf + olen; + a->z_expandable = exp; + + return stbi__parse_zlib(a, parse_header); +} + +STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(initial_size); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer + len; + if (stbi__do_zlib(&a, p, initial_size, 1, 1)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF char *stbi_zlib_decode_malloc(char const *buffer, int len, int *outlen) +{ + return stbi_zlib_decode_malloc_guesssize(buffer, len, 16384, outlen); +} + +STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(initial_size); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer + len; + if (stbi__do_zlib(&a, p, initial_size, 1, parse_header)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, char const *ibuffer, int ilen) +{ + stbi__zbuf a; + a.zbuffer = (stbi_uc *) ibuffer; + a.zbuffer_end = (stbi_uc *) ibuffer + ilen; + if (stbi__do_zlib(&a, obuffer, olen, 0, 1)) + return (int) (a.zout - a.zout_start); + else + return -1; +} + +STBIDEF char *stbi_zlib_decode_noheader_malloc(char const *buffer, int len, int *outlen) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(16384); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer+len; + if (stbi__do_zlib(&a, p, 16384, 1, 0)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen) +{ + stbi__zbuf a; + a.zbuffer = (stbi_uc *) ibuffer; + a.zbuffer_end = (stbi_uc *) ibuffer + ilen; + if (stbi__do_zlib(&a, obuffer, olen, 0, 0)) + return (int) (a.zout - a.zout_start); + else + return -1; +} +#endif + +// public domain "baseline" PNG decoder v0.10 Sean Barrett 2006-11-18 +// simple implementation +// - only 8-bit samples +// - no CRC checking +// - allocates lots of intermediate memory +// - avoids problem of streaming data between subsystems +// - avoids explicit window management +// performance +// - uses stb_zlib, a PD zlib implementation with fast huffman decoding + +#ifndef STBI_NO_PNG +typedef struct +{ + stbi__uint32 length; + stbi__uint32 type; +} stbi__pngchunk; + +static stbi__pngchunk stbi__get_chunk_header(stbi__context *s) +{ + stbi__pngchunk c; + c.length = stbi__get32be(s); + c.type = stbi__get32be(s); + return c; +} + +static int stbi__check_png_header(stbi__context *s) +{ + static const stbi_uc png_sig[8] = { 137,80,78,71,13,10,26,10 }; + int i; + for (i=0; i < 8; ++i) + if (stbi__get8(s) != png_sig[i]) return stbi__err("bad png sig","Not a PNG"); + return 1; +} + +typedef struct +{ + stbi__context *s; + stbi_uc *idata, *expanded, *out; + int depth; +} stbi__png; + + +enum { + STBI__F_none=0, + STBI__F_sub=1, + STBI__F_up=2, + STBI__F_avg=3, + STBI__F_paeth=4, + // synthetic filter used for first scanline to avoid needing a dummy row of 0s + STBI__F_avg_first +}; + +static stbi_uc first_row_filter[5] = +{ + STBI__F_none, + STBI__F_sub, + STBI__F_none, + STBI__F_avg_first, + STBI__F_sub // Paeth with b=c=0 turns out to be equivalent to sub +}; + +static int stbi__paeth(int a, int b, int c) +{ + // This formulation looks very different from the reference in the PNG spec, but is + // actually equivalent and has favorable data dependencies and admits straightforward + // generation of branch-free code, which helps performance significantly. + int thresh = c*3 - (a + b); + int lo = a < b ? a : b; + int hi = a < b ? b : a; + int t0 = (hi <= thresh) ? lo : c; + int t1 = (thresh <= lo) ? hi : t0; + return t1; +} + +static const stbi_uc stbi__depth_scale_table[9] = { 0, 0xff, 0x55, 0, 0x11, 0,0,0, 0x01 }; + +// adds an extra all-255 alpha channel +// dest == src is legal +// img_n must be 1 or 3 +static void stbi__create_png_alpha_expand8(stbi_uc *dest, stbi_uc *src, stbi__uint32 x, int img_n) +{ + int i; + // must process data backwards since we allow dest==src + if (img_n == 1) { + for (i=x-1; i >= 0; --i) { + dest[i*2+1] = 255; + dest[i*2+0] = src[i]; + } + } else { + STBI_ASSERT(img_n == 3); + for (i=x-1; i >= 0; --i) { + dest[i*4+3] = 255; + dest[i*4+2] = src[i*3+2]; + dest[i*4+1] = src[i*3+1]; + dest[i*4+0] = src[i*3+0]; + } + } +} + +// create the png data from post-deflated data +static int stbi__create_png_image_raw(stbi__png *a, stbi_uc *raw, stbi__uint32 raw_len, int out_n, stbi__uint32 x, stbi__uint32 y, int depth, int color) +{ + int bytes = (depth == 16 ? 2 : 1); + stbi__context *s = a->s; + stbi__uint32 i,j,stride = x*out_n*bytes; + stbi__uint32 img_len, img_width_bytes; + stbi_uc *filter_buf; + int all_ok = 1; + int k; + int img_n = s->img_n; // copy it into a local for later + + int output_bytes = out_n*bytes; + int filter_bytes = img_n*bytes; + int width = x; + + STBI_ASSERT(out_n == s->img_n || out_n == s->img_n+1); + a->out = (stbi_uc *) stbi__malloc_mad3(x, y, output_bytes, 0); // extra bytes to write off the end into + if (!a->out) return stbi__err("outofmem", "Out of memory"); + + // note: error exits here don't need to clean up a->out individually, + // stbi__do_png always does on error. + if (!stbi__mad3sizes_valid(img_n, x, depth, 7)) return stbi__err("too large", "Corrupt PNG"); + img_width_bytes = (((img_n * x * depth) + 7) >> 3); + if (!stbi__mad2sizes_valid(img_width_bytes, y, img_width_bytes)) return stbi__err("too large", "Corrupt PNG"); + img_len = (img_width_bytes + 1) * y; + + // we used to check for exact match between raw_len and img_len on non-interlaced PNGs, + // but issue #276 reported a PNG in the wild that had extra data at the end (all zeros), + // so just check for raw_len < img_len always. + if (raw_len < img_len) return stbi__err("not enough pixels","Corrupt PNG"); + + // Allocate two scan lines worth of filter workspace buffer. + filter_buf = (stbi_uc *) stbi__malloc_mad2(img_width_bytes, 2, 0); + if (!filter_buf) return stbi__err("outofmem", "Out of memory"); + + // Filtering for low-bit-depth images + if (depth < 8) { + filter_bytes = 1; + width = img_width_bytes; + } + + for (j=0; j < y; ++j) { + // cur/prior filter buffers alternate + stbi_uc *cur = filter_buf + (j & 1)*img_width_bytes; + stbi_uc *prior = filter_buf + (~j & 1)*img_width_bytes; + stbi_uc *dest = a->out + stride*j; + int nk = width * filter_bytes; + int filter = *raw++; + + // check filter type + if (filter > 4) { + all_ok = stbi__err("invalid filter","Corrupt PNG"); + break; + } + + // if first row, use special filter that doesn't sample previous row + if (j == 0) filter = first_row_filter[filter]; + + // perform actual filtering + switch (filter) { + case STBI__F_none: + memcpy(cur, raw, nk); + break; + case STBI__F_sub: + memcpy(cur, raw, filter_bytes); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + cur[k-filter_bytes]); + break; + case STBI__F_up: + for (k = 0; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + prior[k]); + break; + case STBI__F_avg: + for (k = 0; k < filter_bytes; ++k) + cur[k] = STBI__BYTECAST(raw[k] + (prior[k]>>1)); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + ((prior[k] + cur[k-filter_bytes])>>1)); + break; + case STBI__F_paeth: + for (k = 0; k < filter_bytes; ++k) + cur[k] = STBI__BYTECAST(raw[k] + prior[k]); // prior[k] == stbi__paeth(0,prior[k],0) + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + stbi__paeth(cur[k-filter_bytes], prior[k], prior[k-filter_bytes])); + break; + case STBI__F_avg_first: + memcpy(cur, raw, filter_bytes); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + (cur[k-filter_bytes] >> 1)); + break; + } + + raw += nk; + + // expand decoded bits in cur to dest, also adding an extra alpha channel if desired + if (depth < 8) { + stbi_uc scale = (color == 0) ? stbi__depth_scale_table[depth] : 1; // scale grayscale values to 0..255 range + stbi_uc *in = cur; + stbi_uc *out = dest; + stbi_uc inb = 0; + stbi__uint32 nsmp = x*img_n; + + // expand bits to bytes first + if (depth == 4) { + for (i=0; i < nsmp; ++i) { + if ((i & 1) == 0) inb = *in++; + *out++ = scale * (inb >> 4); + inb <<= 4; + } + } else if (depth == 2) { + for (i=0; i < nsmp; ++i) { + if ((i & 3) == 0) inb = *in++; + *out++ = scale * (inb >> 6); + inb <<= 2; + } + } else { + STBI_ASSERT(depth == 1); + for (i=0; i < nsmp; ++i) { + if ((i & 7) == 0) inb = *in++; + *out++ = scale * (inb >> 7); + inb <<= 1; + } + } + + // insert alpha=255 values if desired + if (img_n != out_n) + stbi__create_png_alpha_expand8(dest, dest, x, img_n); + } else if (depth == 8) { + if (img_n == out_n) + memcpy(dest, cur, x*img_n); + else + stbi__create_png_alpha_expand8(dest, cur, x, img_n); + } else if (depth == 16) { + // convert the image data from big-endian to platform-native + stbi__uint16 *dest16 = (stbi__uint16*)dest; + stbi__uint32 nsmp = x*img_n; + + if (img_n == out_n) { + for (i = 0; i < nsmp; ++i, ++dest16, cur += 2) + *dest16 = (cur[0] << 8) | cur[1]; + } else { + STBI_ASSERT(img_n+1 == out_n); + if (img_n == 1) { + for (i = 0; i < x; ++i, dest16 += 2, cur += 2) { + dest16[0] = (cur[0] << 8) | cur[1]; + dest16[1] = 0xffff; + } + } else { + STBI_ASSERT(img_n == 3); + for (i = 0; i < x; ++i, dest16 += 4, cur += 6) { + dest16[0] = (cur[0] << 8) | cur[1]; + dest16[1] = (cur[2] << 8) | cur[3]; + dest16[2] = (cur[4] << 8) | cur[5]; + dest16[3] = 0xffff; + } + } + } + } + } + + STBI_FREE(filter_buf); + if (!all_ok) return 0; + + return 1; +} + +static int stbi__create_png_image(stbi__png *a, stbi_uc *image_data, stbi__uint32 image_data_len, int out_n, int depth, int color, int interlaced) +{ + int bytes = (depth == 16 ? 2 : 1); + int out_bytes = out_n * bytes; + stbi_uc *final; + int p; + if (!interlaced) + return stbi__create_png_image_raw(a, image_data, image_data_len, out_n, a->s->img_x, a->s->img_y, depth, color); + + // de-interlacing + final = (stbi_uc *) stbi__malloc_mad3(a->s->img_x, a->s->img_y, out_bytes, 0); + if (!final) return stbi__err("outofmem", "Out of memory"); + for (p=0; p < 7; ++p) { + int xorig[] = { 0,4,0,2,0,1,0 }; + int yorig[] = { 0,0,4,0,2,0,1 }; + int xspc[] = { 8,8,4,4,2,2,1 }; + int yspc[] = { 8,8,8,4,4,2,2 }; + int i,j,x,y; + // pass1_x[4] = 0, pass1_x[5] = 1, pass1_x[12] = 1 + x = (a->s->img_x - xorig[p] + xspc[p]-1) / xspc[p]; + y = (a->s->img_y - yorig[p] + yspc[p]-1) / yspc[p]; + if (x && y) { + stbi__uint32 img_len = ((((a->s->img_n * x * depth) + 7) >> 3) + 1) * y; + if (!stbi__create_png_image_raw(a, image_data, image_data_len, out_n, x, y, depth, color)) { + STBI_FREE(final); + return 0; + } + for (j=0; j < y; ++j) { + for (i=0; i < x; ++i) { + int out_y = j*yspc[p]+yorig[p]; + int out_x = i*xspc[p]+xorig[p]; + memcpy(final + out_y*a->s->img_x*out_bytes + out_x*out_bytes, + a->out + (j*x+i)*out_bytes, out_bytes); + } + } + STBI_FREE(a->out); + image_data += img_len; + image_data_len -= img_len; + } + } + a->out = final; + + return 1; +} + +static int stbi__compute_transparency(stbi__png *z, stbi_uc tc[3], int out_n) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi_uc *p = z->out; + + // compute color-based transparency, assuming we've + // already got 255 as the alpha value in the output + STBI_ASSERT(out_n == 2 || out_n == 4); + + if (out_n == 2) { + for (i=0; i < pixel_count; ++i) { + p[1] = (p[0] == tc[0] ? 0 : 255); + p += 2; + } + } else { + for (i=0; i < pixel_count; ++i) { + if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) + p[3] = 0; + p += 4; + } + } + return 1; +} + +static int stbi__compute_transparency16(stbi__png *z, stbi__uint16 tc[3], int out_n) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi__uint16 *p = (stbi__uint16*) z->out; + + // compute color-based transparency, assuming we've + // already got 65535 as the alpha value in the output + STBI_ASSERT(out_n == 2 || out_n == 4); + + if (out_n == 2) { + for (i = 0; i < pixel_count; ++i) { + p[1] = (p[0] == tc[0] ? 0 : 65535); + p += 2; + } + } else { + for (i = 0; i < pixel_count; ++i) { + if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) + p[3] = 0; + p += 4; + } + } + return 1; +} + +static int stbi__expand_png_palette(stbi__png *a, stbi_uc *palette, int len, int pal_img_n) +{ + stbi__uint32 i, pixel_count = a->s->img_x * a->s->img_y; + stbi_uc *p, *temp_out, *orig = a->out; + + p = (stbi_uc *) stbi__malloc_mad2(pixel_count, pal_img_n, 0); + if (p == NULL) return stbi__err("outofmem", "Out of memory"); + + // between here and free(out) below, exitting would leak + temp_out = p; + + if (pal_img_n == 3) { + for (i=0; i < pixel_count; ++i) { + int n = orig[i]*4; + p[0] = palette[n ]; + p[1] = palette[n+1]; + p[2] = palette[n+2]; + p += 3; + } + } else { + for (i=0; i < pixel_count; ++i) { + int n = orig[i]*4; + p[0] = palette[n ]; + p[1] = palette[n+1]; + p[2] = palette[n+2]; + p[3] = palette[n+3]; + p += 4; + } + } + STBI_FREE(a->out); + a->out = temp_out; + + STBI_NOTUSED(len); + + return 1; +} + +static int stbi__unpremultiply_on_load_global = 0; +static int stbi__de_iphone_flag_global = 0; + +STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply) +{ + stbi__unpremultiply_on_load_global = flag_true_if_should_unpremultiply; +} + +STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert) +{ + stbi__de_iphone_flag_global = flag_true_if_should_convert; +} + +#ifndef STBI_THREAD_LOCAL +#define stbi__unpremultiply_on_load stbi__unpremultiply_on_load_global +#define stbi__de_iphone_flag stbi__de_iphone_flag_global +#else +static STBI_THREAD_LOCAL int stbi__unpremultiply_on_load_local, stbi__unpremultiply_on_load_set; +static STBI_THREAD_LOCAL int stbi__de_iphone_flag_local, stbi__de_iphone_flag_set; + +STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply) +{ + stbi__unpremultiply_on_load_local = flag_true_if_should_unpremultiply; + stbi__unpremultiply_on_load_set = 1; +} + +STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert) +{ + stbi__de_iphone_flag_local = flag_true_if_should_convert; + stbi__de_iphone_flag_set = 1; +} + +#define stbi__unpremultiply_on_load (stbi__unpremultiply_on_load_set \ + ? stbi__unpremultiply_on_load_local \ + : stbi__unpremultiply_on_load_global) +#define stbi__de_iphone_flag (stbi__de_iphone_flag_set \ + ? stbi__de_iphone_flag_local \ + : stbi__de_iphone_flag_global) +#endif // STBI_THREAD_LOCAL + +static void stbi__de_iphone(stbi__png *z) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi_uc *p = z->out; + + if (s->img_out_n == 3) { // convert bgr to rgb + for (i=0; i < pixel_count; ++i) { + stbi_uc t = p[0]; + p[0] = p[2]; + p[2] = t; + p += 3; + } + } else { + STBI_ASSERT(s->img_out_n == 4); + if (stbi__unpremultiply_on_load) { + // convert bgr to rgb and unpremultiply + for (i=0; i < pixel_count; ++i) { + stbi_uc a = p[3]; + stbi_uc t = p[0]; + if (a) { + stbi_uc half = a / 2; + p[0] = (p[2] * 255 + half) / a; + p[1] = (p[1] * 255 + half) / a; + p[2] = ( t * 255 + half) / a; + } else { + p[0] = p[2]; + p[2] = t; + } + p += 4; + } + } else { + // convert bgr to rgb + for (i=0; i < pixel_count; ++i) { + stbi_uc t = p[0]; + p[0] = p[2]; + p[2] = t; + p += 4; + } + } + } +} + +#define STBI__PNG_TYPE(a,b,c,d) (((unsigned) (a) << 24) + ((unsigned) (b) << 16) + ((unsigned) (c) << 8) + (unsigned) (d)) + +static int stbi__parse_png_file(stbi__png *z, int scan, int req_comp) +{ + stbi_uc palette[1024], pal_img_n=0; + stbi_uc has_trans=0, tc[3]={0}; + stbi__uint16 tc16[3]; + stbi__uint32 ioff=0, idata_limit=0, i, pal_len=0; + int first=1,k,interlace=0, color=0, is_iphone=0; + stbi__context *s = z->s; + + z->expanded = NULL; + z->idata = NULL; + z->out = NULL; + + if (!stbi__check_png_header(s)) return 0; + + if (scan == STBI__SCAN_type) return 1; + + for (;;) { + stbi__pngchunk c = stbi__get_chunk_header(s); + switch (c.type) { + case STBI__PNG_TYPE('C','g','B','I'): + is_iphone = 1; + stbi__skip(s, c.length); + break; + case STBI__PNG_TYPE('I','H','D','R'): { + int comp,filter; + if (!first) return stbi__err("multiple IHDR","Corrupt PNG"); + first = 0; + if (c.length != 13) return stbi__err("bad IHDR len","Corrupt PNG"); + s->img_x = stbi__get32be(s); + s->img_y = stbi__get32be(s); + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + z->depth = stbi__get8(s); if (z->depth != 1 && z->depth != 2 && z->depth != 4 && z->depth != 8 && z->depth != 16) return stbi__err("1/2/4/8/16-bit only","PNG not supported: 1/2/4/8/16-bit only"); + color = stbi__get8(s); if (color > 6) return stbi__err("bad ctype","Corrupt PNG"); + if (color == 3 && z->depth == 16) return stbi__err("bad ctype","Corrupt PNG"); + if (color == 3) pal_img_n = 3; else if (color & 1) return stbi__err("bad ctype","Corrupt PNG"); + comp = stbi__get8(s); if (comp) return stbi__err("bad comp method","Corrupt PNG"); + filter= stbi__get8(s); if (filter) return stbi__err("bad filter method","Corrupt PNG"); + interlace = stbi__get8(s); if (interlace>1) return stbi__err("bad interlace method","Corrupt PNG"); + if (!s->img_x || !s->img_y) return stbi__err("0-pixel image","Corrupt PNG"); + if (!pal_img_n) { + s->img_n = (color & 2 ? 3 : 1) + (color & 4 ? 1 : 0); + if ((1 << 30) / s->img_x / s->img_n < s->img_y) return stbi__err("too large", "Image too large to decode"); + } else { + // if paletted, then pal_n is our final components, and + // img_n is # components to decompress/filter. + s->img_n = 1; + if ((1 << 30) / s->img_x / 4 < s->img_y) return stbi__err("too large","Corrupt PNG"); + } + // even with SCAN_header, have to scan to see if we have a tRNS + break; + } + + case STBI__PNG_TYPE('P','L','T','E'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (c.length > 256*3) return stbi__err("invalid PLTE","Corrupt PNG"); + pal_len = c.length / 3; + if (pal_len * 3 != c.length) return stbi__err("invalid PLTE","Corrupt PNG"); + for (i=0; i < pal_len; ++i) { + palette[i*4+0] = stbi__get8(s); + palette[i*4+1] = stbi__get8(s); + palette[i*4+2] = stbi__get8(s); + palette[i*4+3] = 255; + } + break; + } + + case STBI__PNG_TYPE('t','R','N','S'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (z->idata) return stbi__err("tRNS after IDAT","Corrupt PNG"); + if (pal_img_n) { + if (scan == STBI__SCAN_header) { s->img_n = 4; return 1; } + if (pal_len == 0) return stbi__err("tRNS before PLTE","Corrupt PNG"); + if (c.length > pal_len) return stbi__err("bad tRNS len","Corrupt PNG"); + pal_img_n = 4; + for (i=0; i < c.length; ++i) + palette[i*4+3] = stbi__get8(s); + } else { + if (!(s->img_n & 1)) return stbi__err("tRNS with alpha","Corrupt PNG"); + if (c.length != (stbi__uint32) s->img_n*2) return stbi__err("bad tRNS len","Corrupt PNG"); + has_trans = 1; + // non-paletted with tRNS = constant alpha. if header-scanning, we can stop now. + if (scan == STBI__SCAN_header) { ++s->img_n; return 1; } + if (z->depth == 16) { + for (k = 0; k < s->img_n && k < 3; ++k) // extra loop test to suppress false GCC warning + tc16[k] = (stbi__uint16)stbi__get16be(s); // copy the values as-is + } else { + for (k = 0; k < s->img_n && k < 3; ++k) + tc[k] = (stbi_uc)(stbi__get16be(s) & 255) * stbi__depth_scale_table[z->depth]; // non 8-bit images will be larger + } + } + break; + } + + case STBI__PNG_TYPE('I','D','A','T'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (pal_img_n && !pal_len) return stbi__err("no PLTE","Corrupt PNG"); + if (scan == STBI__SCAN_header) { + // header scan definitely stops at first IDAT + if (pal_img_n) + s->img_n = pal_img_n; + return 1; + } + if (c.length > (1u << 30)) return stbi__err("IDAT size limit", "IDAT section larger than 2^30 bytes"); + if ((int)(ioff + c.length) < (int)ioff) return 0; + if (ioff + c.length > idata_limit) { + stbi__uint32 idata_limit_old = idata_limit; + stbi_uc *p; + if (idata_limit == 0) idata_limit = c.length > 4096 ? c.length : 4096; + while (ioff + c.length > idata_limit) + idata_limit *= 2; + STBI_NOTUSED(idata_limit_old); + p = (stbi_uc *) STBI_REALLOC_SIZED(z->idata, idata_limit_old, idata_limit); if (p == NULL) return stbi__err("outofmem", "Out of memory"); + z->idata = p; + } + if (!stbi__getn(s, z->idata+ioff,c.length)) return stbi__err("outofdata","Corrupt PNG"); + ioff += c.length; + break; + } + + case STBI__PNG_TYPE('I','E','N','D'): { + stbi__uint32 raw_len, bpl; + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (scan != STBI__SCAN_load) return 1; + if (z->idata == NULL) return stbi__err("no IDAT","Corrupt PNG"); + // initial guess for decoded data size to avoid unnecessary reallocs + bpl = (s->img_x * z->depth + 7) / 8; // bytes per line, per component + raw_len = bpl * s->img_y * s->img_n /* pixels */ + s->img_y /* filter mode per row */; + z->expanded = (stbi_uc *) stbi_zlib_decode_malloc_guesssize_headerflag((char *) z->idata, ioff, raw_len, (int *) &raw_len, !is_iphone); + if (z->expanded == NULL) return 0; // zlib should set error + STBI_FREE(z->idata); z->idata = NULL; + if ((req_comp == s->img_n+1 && req_comp != 3 && !pal_img_n) || has_trans) + s->img_out_n = s->img_n+1; + else + s->img_out_n = s->img_n; + if (!stbi__create_png_image(z, z->expanded, raw_len, s->img_out_n, z->depth, color, interlace)) return 0; + if (has_trans) { + if (z->depth == 16) { + if (!stbi__compute_transparency16(z, tc16, s->img_out_n)) return 0; + } else { + if (!stbi__compute_transparency(z, tc, s->img_out_n)) return 0; + } + } + if (is_iphone && stbi__de_iphone_flag && s->img_out_n > 2) + stbi__de_iphone(z); + if (pal_img_n) { + // pal_img_n == 3 or 4 + s->img_n = pal_img_n; // record the actual colors we had + s->img_out_n = pal_img_n; + if (req_comp >= 3) s->img_out_n = req_comp; + if (!stbi__expand_png_palette(z, palette, pal_len, s->img_out_n)) + return 0; + } else if (has_trans) { + // non-paletted image with tRNS -> source image has (constant) alpha + ++s->img_n; + } + STBI_FREE(z->expanded); z->expanded = NULL; + // end of PNG chunk, read and skip CRC + stbi__get32be(s); + return 1; + } + + default: + // if critical, fail + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if ((c.type & (1 << 29)) == 0) { + #ifndef STBI_NO_FAILURE_STRINGS + // not threadsafe + static char invalid_chunk[] = "XXXX PNG chunk not known"; + invalid_chunk[0] = STBI__BYTECAST(c.type >> 24); + invalid_chunk[1] = STBI__BYTECAST(c.type >> 16); + invalid_chunk[2] = STBI__BYTECAST(c.type >> 8); + invalid_chunk[3] = STBI__BYTECAST(c.type >> 0); + #endif + return stbi__err(invalid_chunk, "PNG not supported: unknown PNG chunk type"); + } + stbi__skip(s, c.length); + break; + } + // end of PNG chunk, read and skip CRC + stbi__get32be(s); + } +} + +static void *stbi__do_png(stbi__png *p, int *x, int *y, int *n, int req_comp, stbi__result_info *ri) +{ + void *result=NULL; + if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); + if (stbi__parse_png_file(p, STBI__SCAN_load, req_comp)) { + if (p->depth <= 8) + ri->bits_per_channel = 8; + else if (p->depth == 16) + ri->bits_per_channel = 16; + else + return stbi__errpuc("bad bits_per_channel", "PNG not supported: unsupported color depth"); + result = p->out; + p->out = NULL; + if (req_comp && req_comp != p->s->img_out_n) { + if (ri->bits_per_channel == 8) + result = stbi__convert_format((unsigned char *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); + else + result = stbi__convert_format16((stbi__uint16 *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); + p->s->img_out_n = req_comp; + if (result == NULL) return result; + } + *x = p->s->img_x; + *y = p->s->img_y; + if (n) *n = p->s->img_n; + } + STBI_FREE(p->out); p->out = NULL; + STBI_FREE(p->expanded); p->expanded = NULL; + STBI_FREE(p->idata); p->idata = NULL; + + return result; +} + +static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi__png p; + p.s = s; + return stbi__do_png(&p, x,y,comp,req_comp, ri); +} + +static int stbi__png_test(stbi__context *s) +{ + int r; + r = stbi__check_png_header(s); + stbi__rewind(s); + return r; +} + +static int stbi__png_info_raw(stbi__png *p, int *x, int *y, int *comp) +{ + if (!stbi__parse_png_file(p, STBI__SCAN_header, 0)) { + stbi__rewind( p->s ); + return 0; + } + if (x) *x = p->s->img_x; + if (y) *y = p->s->img_y; + if (comp) *comp = p->s->img_n; + return 1; +} + +static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp) +{ + stbi__png p; + p.s = s; + return stbi__png_info_raw(&p, x, y, comp); +} + +static int stbi__png_is16(stbi__context *s) +{ + stbi__png p; + p.s = s; + if (!stbi__png_info_raw(&p, NULL, NULL, NULL)) + return 0; + if (p.depth != 16) { + stbi__rewind(p.s); + return 0; + } + return 1; +} +#endif + +// Microsoft/Windows BMP image + +#ifndef STBI_NO_BMP +static int stbi__bmp_test_raw(stbi__context *s) +{ + int r; + int sz; + if (stbi__get8(s) != 'B') return 0; + if (stbi__get8(s) != 'M') return 0; + stbi__get32le(s); // discard filesize + stbi__get16le(s); // discard reserved + stbi__get16le(s); // discard reserved + stbi__get32le(s); // discard data offset + sz = stbi__get32le(s); + r = (sz == 12 || sz == 40 || sz == 56 || sz == 108 || sz == 124); + return r; +} + +static int stbi__bmp_test(stbi__context *s) +{ + int r = stbi__bmp_test_raw(s); + stbi__rewind(s); + return r; +} + + +// returns 0..31 for the highest set bit +static int stbi__high_bit(unsigned int z) +{ + int n=0; + if (z == 0) return -1; + if (z >= 0x10000) { n += 16; z >>= 16; } + if (z >= 0x00100) { n += 8; z >>= 8; } + if (z >= 0x00010) { n += 4; z >>= 4; } + if (z >= 0x00004) { n += 2; z >>= 2; } + if (z >= 0x00002) { n += 1;/* >>= 1;*/ } + return n; +} + +static int stbi__bitcount(unsigned int a) +{ + a = (a & 0x55555555) + ((a >> 1) & 0x55555555); // max 2 + a = (a & 0x33333333) + ((a >> 2) & 0x33333333); // max 4 + a = (a + (a >> 4)) & 0x0f0f0f0f; // max 8 per 4, now 8 bits + a = (a + (a >> 8)); // max 16 per 8 bits + a = (a + (a >> 16)); // max 32 per 8 bits + return a & 0xff; +} + +// extract an arbitrarily-aligned N-bit value (N=bits) +// from v, and then make it 8-bits long and fractionally +// extend it to full full range. +static int stbi__shiftsigned(unsigned int v, int shift, int bits) +{ + static unsigned int mul_table[9] = { + 0, + 0xff/*0b11111111*/, 0x55/*0b01010101*/, 0x49/*0b01001001*/, 0x11/*0b00010001*/, + 0x21/*0b00100001*/, 0x41/*0b01000001*/, 0x81/*0b10000001*/, 0x01/*0b00000001*/, + }; + static unsigned int shift_table[9] = { + 0, 0,0,1,0,2,4,6,0, + }; + if (shift < 0) + v <<= -shift; + else + v >>= shift; + STBI_ASSERT(v < 256); + v >>= (8-bits); + STBI_ASSERT(bits >= 0 && bits <= 8); + return (int) ((unsigned) v * mul_table[bits]) >> shift_table[bits]; +} + +typedef struct +{ + int bpp, offset, hsz; + unsigned int mr,mg,mb,ma, all_a; + int extra_read; +} stbi__bmp_data; + +static int stbi__bmp_set_mask_defaults(stbi__bmp_data *info, int compress) +{ + // BI_BITFIELDS specifies masks explicitly, don't override + if (compress == 3) + return 1; + + if (compress == 0) { + if (info->bpp == 16) { + info->mr = 31u << 10; + info->mg = 31u << 5; + info->mb = 31u << 0; + } else if (info->bpp == 32) { + info->mr = 0xffu << 16; + info->mg = 0xffu << 8; + info->mb = 0xffu << 0; + info->ma = 0xffu << 24; + info->all_a = 0; // if all_a is 0 at end, then we loaded alpha channel but it was all 0 + } else { + // otherwise, use defaults, which is all-0 + info->mr = info->mg = info->mb = info->ma = 0; + } + return 1; + } + return 0; // error +} + +static void *stbi__bmp_parse_header(stbi__context *s, stbi__bmp_data *info) +{ + int hsz; + if (stbi__get8(s) != 'B' || stbi__get8(s) != 'M') return stbi__errpuc("not BMP", "Corrupt BMP"); + stbi__get32le(s); // discard filesize + stbi__get16le(s); // discard reserved + stbi__get16le(s); // discard reserved + info->offset = stbi__get32le(s); + info->hsz = hsz = stbi__get32le(s); + info->mr = info->mg = info->mb = info->ma = 0; + info->extra_read = 14; + + if (info->offset < 0) return stbi__errpuc("bad BMP", "bad BMP"); + + if (hsz != 12 && hsz != 40 && hsz != 56 && hsz != 108 && hsz != 124) return stbi__errpuc("unknown BMP", "BMP type not supported: unknown"); + if (hsz == 12) { + s->img_x = stbi__get16le(s); + s->img_y = stbi__get16le(s); + } else { + s->img_x = stbi__get32le(s); + s->img_y = stbi__get32le(s); + } + if (stbi__get16le(s) != 1) return stbi__errpuc("bad BMP", "bad BMP"); + info->bpp = stbi__get16le(s); + if (hsz != 12) { + int compress = stbi__get32le(s); + if (compress == 1 || compress == 2) return stbi__errpuc("BMP RLE", "BMP type not supported: RLE"); + if (compress >= 4) return stbi__errpuc("BMP JPEG/PNG", "BMP type not supported: unsupported compression"); // this includes PNG/JPEG modes + if (compress == 3 && info->bpp != 16 && info->bpp != 32) return stbi__errpuc("bad BMP", "bad BMP"); // bitfields requires 16 or 32 bits/pixel + stbi__get32le(s); // discard sizeof + stbi__get32le(s); // discard hres + stbi__get32le(s); // discard vres + stbi__get32le(s); // discard colorsused + stbi__get32le(s); // discard max important + if (hsz == 40 || hsz == 56) { + if (hsz == 56) { + stbi__get32le(s); + stbi__get32le(s); + stbi__get32le(s); + stbi__get32le(s); + } + if (info->bpp == 16 || info->bpp == 32) { + if (compress == 0) { + stbi__bmp_set_mask_defaults(info, compress); + } else if (compress == 3) { + info->mr = stbi__get32le(s); + info->mg = stbi__get32le(s); + info->mb = stbi__get32le(s); + info->extra_read += 12; + // not documented, but generated by photoshop and handled by mspaint + if (info->mr == info->mg && info->mg == info->mb) { + // ?!?!? + return stbi__errpuc("bad BMP", "bad BMP"); + } + } else + return stbi__errpuc("bad BMP", "bad BMP"); + } + } else { + // V4/V5 header + int i; + if (hsz != 108 && hsz != 124) + return stbi__errpuc("bad BMP", "bad BMP"); + info->mr = stbi__get32le(s); + info->mg = stbi__get32le(s); + info->mb = stbi__get32le(s); + info->ma = stbi__get32le(s); + if (compress != 3) // override mr/mg/mb unless in BI_BITFIELDS mode, as per docs + stbi__bmp_set_mask_defaults(info, compress); + stbi__get32le(s); // discard color space + for (i=0; i < 12; ++i) + stbi__get32le(s); // discard color space parameters + if (hsz == 124) { + stbi__get32le(s); // discard rendering intent + stbi__get32le(s); // discard offset of profile data + stbi__get32le(s); // discard size of profile data + stbi__get32le(s); // discard reserved + } + } + } + return (void *) 1; +} + + +static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *out; + unsigned int mr=0,mg=0,mb=0,ma=0, all_a; + stbi_uc pal[256][4]; + int psize=0,i,j,width; + int flip_vertically, pad, target; + stbi__bmp_data info; + STBI_NOTUSED(ri); + + info.all_a = 255; + if (stbi__bmp_parse_header(s, &info) == NULL) + return NULL; // error code already set + + flip_vertically = ((int) s->img_y) > 0; + s->img_y = abs((int) s->img_y); + + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + mr = info.mr; + mg = info.mg; + mb = info.mb; + ma = info.ma; + all_a = info.all_a; + + if (info.hsz == 12) { + if (info.bpp < 24) + psize = (info.offset - info.extra_read - 24) / 3; + } else { + if (info.bpp < 16) + psize = (info.offset - info.extra_read - info.hsz) >> 2; + } + if (psize == 0) { + // accept some number of extra bytes after the header, but if the offset points either to before + // the header ends or implies a large amount of extra data, reject the file as malformed + int bytes_read_so_far = s->callback_already_read + (int)(s->img_buffer - s->img_buffer_original); + int header_limit = 1024; // max we actually read is below 256 bytes currently. + int extra_data_limit = 256*4; // what ordinarily goes here is a palette; 256 entries*4 bytes is its max size. + if (bytes_read_so_far <= 0 || bytes_read_so_far > header_limit) { + return stbi__errpuc("bad header", "Corrupt BMP"); + } + // we established that bytes_read_so_far is positive and sensible. + // the first half of this test rejects offsets that are either too small positives, or + // negative, and guarantees that info.offset >= bytes_read_so_far > 0. this in turn + // ensures the number computed in the second half of the test can't overflow. + if (info.offset < bytes_read_so_far || info.offset - bytes_read_so_far > extra_data_limit) { + return stbi__errpuc("bad offset", "Corrupt BMP"); + } else { + stbi__skip(s, info.offset - bytes_read_so_far); + } + } + + if (info.bpp == 24 && ma == 0xff000000) + s->img_n = 3; + else + s->img_n = ma ? 4 : 3; + if (req_comp && req_comp >= 3) // we can directly decode 3 or 4 + target = req_comp; + else + target = s->img_n; // if they want monochrome, we'll post-convert + + // sanity-check size + if (!stbi__mad3sizes_valid(target, s->img_x, s->img_y, 0)) + return stbi__errpuc("too large", "Corrupt BMP"); + + out = (stbi_uc *) stbi__malloc_mad3(target, s->img_x, s->img_y, 0); + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + if (info.bpp < 16) { + int z=0; + if (psize == 0 || psize > 256) { STBI_FREE(out); return stbi__errpuc("invalid", "Corrupt BMP"); } + for (i=0; i < psize; ++i) { + pal[i][2] = stbi__get8(s); + pal[i][1] = stbi__get8(s); + pal[i][0] = stbi__get8(s); + if (info.hsz != 12) stbi__get8(s); + pal[i][3] = 255; + } + stbi__skip(s, info.offset - info.extra_read - info.hsz - psize * (info.hsz == 12 ? 3 : 4)); + if (info.bpp == 1) width = (s->img_x + 7) >> 3; + else if (info.bpp == 4) width = (s->img_x + 1) >> 1; + else if (info.bpp == 8) width = s->img_x; + else { STBI_FREE(out); return stbi__errpuc("bad bpp", "Corrupt BMP"); } + pad = (-width)&3; + if (info.bpp == 1) { + for (j=0; j < (int) s->img_y; ++j) { + int bit_offset = 7, v = stbi__get8(s); + for (i=0; i < (int) s->img_x; ++i) { + int color = (v>>bit_offset)&0x1; + out[z++] = pal[color][0]; + out[z++] = pal[color][1]; + out[z++] = pal[color][2]; + if (target == 4) out[z++] = 255; + if (i+1 == (int) s->img_x) break; + if((--bit_offset) < 0) { + bit_offset = 7; + v = stbi__get8(s); + } + } + stbi__skip(s, pad); + } + } else { + for (j=0; j < (int) s->img_y; ++j) { + for (i=0; i < (int) s->img_x; i += 2) { + int v=stbi__get8(s),v2=0; + if (info.bpp == 4) { + v2 = v & 15; + v >>= 4; + } + out[z++] = pal[v][0]; + out[z++] = pal[v][1]; + out[z++] = pal[v][2]; + if (target == 4) out[z++] = 255; + if (i+1 == (int) s->img_x) break; + v = (info.bpp == 8) ? stbi__get8(s) : v2; + out[z++] = pal[v][0]; + out[z++] = pal[v][1]; + out[z++] = pal[v][2]; + if (target == 4) out[z++] = 255; + } + stbi__skip(s, pad); + } + } + } else { + int rshift=0,gshift=0,bshift=0,ashift=0,rcount=0,gcount=0,bcount=0,acount=0; + int z = 0; + int easy=0; + stbi__skip(s, info.offset - info.extra_read - info.hsz); + if (info.bpp == 24) width = 3 * s->img_x; + else if (info.bpp == 16) width = 2*s->img_x; + else /* bpp = 32 and pad = 0 */ width=0; + pad = (-width) & 3; + if (info.bpp == 24) { + easy = 1; + } else if (info.bpp == 32) { + if (mb == 0xff && mg == 0xff00 && mr == 0x00ff0000 && ma == 0xff000000) + easy = 2; + } + if (!easy) { + if (!mr || !mg || !mb) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } + // right shift amt to put high bit in position #7 + rshift = stbi__high_bit(mr)-7; rcount = stbi__bitcount(mr); + gshift = stbi__high_bit(mg)-7; gcount = stbi__bitcount(mg); + bshift = stbi__high_bit(mb)-7; bcount = stbi__bitcount(mb); + ashift = stbi__high_bit(ma)-7; acount = stbi__bitcount(ma); + if (rcount > 8 || gcount > 8 || bcount > 8 || acount > 8) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } + } + for (j=0; j < (int) s->img_y; ++j) { + if (easy) { + for (i=0; i < (int) s->img_x; ++i) { + unsigned char a; + out[z+2] = stbi__get8(s); + out[z+1] = stbi__get8(s); + out[z+0] = stbi__get8(s); + z += 3; + a = (easy == 2 ? stbi__get8(s) : 255); + all_a |= a; + if (target == 4) out[z++] = a; + } + } else { + int bpp = info.bpp; + for (i=0; i < (int) s->img_x; ++i) { + stbi__uint32 v = (bpp == 16 ? (stbi__uint32) stbi__get16le(s) : stbi__get32le(s)); + unsigned int a; + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mr, rshift, rcount)); + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mg, gshift, gcount)); + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mb, bshift, bcount)); + a = (ma ? stbi__shiftsigned(v & ma, ashift, acount) : 255); + all_a |= a; + if (target == 4) out[z++] = STBI__BYTECAST(a); + } + } + stbi__skip(s, pad); + } + } + + // if alpha channel is all 0s, replace with all 255s + if (target == 4 && all_a == 0) + for (i=4*s->img_x*s->img_y-1; i >= 0; i -= 4) + out[i] = 255; + + if (flip_vertically) { + stbi_uc t; + for (j=0; j < (int) s->img_y>>1; ++j) { + stbi_uc *p1 = out + j *s->img_x*target; + stbi_uc *p2 = out + (s->img_y-1-j)*s->img_x*target; + for (i=0; i < (int) s->img_x*target; ++i) { + t = p1[i]; p1[i] = p2[i]; p2[i] = t; + } + } + } + + if (req_comp && req_comp != target) { + out = stbi__convert_format(out, target, req_comp, s->img_x, s->img_y); + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + + *x = s->img_x; + *y = s->img_y; + if (comp) *comp = s->img_n; + return out; +} +#endif + +// Targa Truevision - TGA +// by Jonathan Dummer +#ifndef STBI_NO_TGA +// returns STBI_rgb or whatever, 0 on error +static int stbi__tga_get_comp(int bits_per_pixel, int is_grey, int* is_rgb16) +{ + // only RGB or RGBA (incl. 16bit) or grey allowed + if (is_rgb16) *is_rgb16 = 0; + switch(bits_per_pixel) { + case 8: return STBI_grey; + case 16: if(is_grey) return STBI_grey_alpha; + // fallthrough + case 15: if(is_rgb16) *is_rgb16 = 1; + return STBI_rgb; + case 24: // fallthrough + case 32: return bits_per_pixel/8; + default: return 0; + } +} + +static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp) +{ + int tga_w, tga_h, tga_comp, tga_image_type, tga_bits_per_pixel, tga_colormap_bpp; + int sz, tga_colormap_type; + stbi__get8(s); // discard Offset + tga_colormap_type = stbi__get8(s); // colormap type + if( tga_colormap_type > 1 ) { + stbi__rewind(s); + return 0; // only RGB or indexed allowed + } + tga_image_type = stbi__get8(s); // image type + if ( tga_colormap_type == 1 ) { // colormapped (paletted) image + if (tga_image_type != 1 && tga_image_type != 9) { + stbi__rewind(s); + return 0; + } + stbi__skip(s,4); // skip index of first colormap entry and number of entries + sz = stbi__get8(s); // check bits per palette color entry + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) { + stbi__rewind(s); + return 0; + } + stbi__skip(s,4); // skip image x and y origin + tga_colormap_bpp = sz; + } else { // "normal" image w/o colormap - only RGB or grey allowed, +/- RLE + if ( (tga_image_type != 2) && (tga_image_type != 3) && (tga_image_type != 10) && (tga_image_type != 11) ) { + stbi__rewind(s); + return 0; // only RGB or grey allowed, +/- RLE + } + stbi__skip(s,9); // skip colormap specification and image x/y origin + tga_colormap_bpp = 0; + } + tga_w = stbi__get16le(s); + if( tga_w < 1 ) { + stbi__rewind(s); + return 0; // test width + } + tga_h = stbi__get16le(s); + if( tga_h < 1 ) { + stbi__rewind(s); + return 0; // test height + } + tga_bits_per_pixel = stbi__get8(s); // bits per pixel + stbi__get8(s); // ignore alpha bits + if (tga_colormap_bpp != 0) { + if((tga_bits_per_pixel != 8) && (tga_bits_per_pixel != 16)) { + // when using a colormap, tga_bits_per_pixel is the size of the indexes + // I don't think anything but 8 or 16bit indexes makes sense + stbi__rewind(s); + return 0; + } + tga_comp = stbi__tga_get_comp(tga_colormap_bpp, 0, NULL); + } else { + tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3) || (tga_image_type == 11), NULL); + } + if(!tga_comp) { + stbi__rewind(s); + return 0; + } + if (x) *x = tga_w; + if (y) *y = tga_h; + if (comp) *comp = tga_comp; + return 1; // seems to have passed everything +} + +static int stbi__tga_test(stbi__context *s) +{ + int res = 0; + int sz, tga_color_type; + stbi__get8(s); // discard Offset + tga_color_type = stbi__get8(s); // color type + if ( tga_color_type > 1 ) goto errorEnd; // only RGB or indexed allowed + sz = stbi__get8(s); // image type + if ( tga_color_type == 1 ) { // colormapped (paletted) image + if (sz != 1 && sz != 9) goto errorEnd; // colortype 1 demands image type 1 or 9 + stbi__skip(s,4); // skip index of first colormap entry and number of entries + sz = stbi__get8(s); // check bits per palette color entry + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; + stbi__skip(s,4); // skip image x and y origin + } else { // "normal" image w/o colormap + if ( (sz != 2) && (sz != 3) && (sz != 10) && (sz != 11) ) goto errorEnd; // only RGB or grey allowed, +/- RLE + stbi__skip(s,9); // skip colormap specification and image x/y origin + } + if ( stbi__get16le(s) < 1 ) goto errorEnd; // test width + if ( stbi__get16le(s) < 1 ) goto errorEnd; // test height + sz = stbi__get8(s); // bits per pixel + if ( (tga_color_type == 1) && (sz != 8) && (sz != 16) ) goto errorEnd; // for colormapped images, bpp is size of an index + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; + + res = 1; // if we got this far, everything's good and we can return 1 instead of 0 + +errorEnd: + stbi__rewind(s); + return res; +} + +// read 16bit value and convert to 24bit RGB +static void stbi__tga_read_rgb16(stbi__context *s, stbi_uc* out) +{ + stbi__uint16 px = (stbi__uint16)stbi__get16le(s); + stbi__uint16 fiveBitMask = 31; + // we have 3 channels with 5bits each + int r = (px >> 10) & fiveBitMask; + int g = (px >> 5) & fiveBitMask; + int b = px & fiveBitMask; + // Note that this saves the data in RGB(A) order, so it doesn't need to be swapped later + out[0] = (stbi_uc)((r * 255)/31); + out[1] = (stbi_uc)((g * 255)/31); + out[2] = (stbi_uc)((b * 255)/31); + + // some people claim that the most significant bit might be used for alpha + // (possibly if an alpha-bit is set in the "image descriptor byte") + // but that only made 16bit test images completely translucent.. + // so let's treat all 15 and 16bit TGAs as RGB with no alpha. +} + +static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + // read in the TGA header stuff + int tga_offset = stbi__get8(s); + int tga_indexed = stbi__get8(s); + int tga_image_type = stbi__get8(s); + int tga_is_RLE = 0; + int tga_palette_start = stbi__get16le(s); + int tga_palette_len = stbi__get16le(s); + int tga_palette_bits = stbi__get8(s); + int tga_x_origin = stbi__get16le(s); + int tga_y_origin = stbi__get16le(s); + int tga_width = stbi__get16le(s); + int tga_height = stbi__get16le(s); + int tga_bits_per_pixel = stbi__get8(s); + int tga_comp, tga_rgb16=0; + int tga_inverted = stbi__get8(s); + // int tga_alpha_bits = tga_inverted & 15; // the 4 lowest bits - unused (useless?) + // image data + unsigned char *tga_data; + unsigned char *tga_palette = NULL; + int i, j; + unsigned char raw_data[4] = {0}; + int RLE_count = 0; + int RLE_repeating = 0; + int read_next_pixel = 1; + STBI_NOTUSED(ri); + STBI_NOTUSED(tga_x_origin); // @TODO + STBI_NOTUSED(tga_y_origin); // @TODO + + if (tga_height > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (tga_width > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + // do a tiny bit of precessing + if ( tga_image_type >= 8 ) + { + tga_image_type -= 8; + tga_is_RLE = 1; + } + tga_inverted = 1 - ((tga_inverted >> 5) & 1); + + // If I'm paletted, then I'll use the number of bits from the palette + if ( tga_indexed ) tga_comp = stbi__tga_get_comp(tga_palette_bits, 0, &tga_rgb16); + else tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3), &tga_rgb16); + + if(!tga_comp) // shouldn't really happen, stbi__tga_test() should have ensured basic consistency + return stbi__errpuc("bad format", "Can't find out TGA pixelformat"); + + // tga info + *x = tga_width; + *y = tga_height; + if (comp) *comp = tga_comp; + + if (!stbi__mad3sizes_valid(tga_width, tga_height, tga_comp, 0)) + return stbi__errpuc("too large", "Corrupt TGA"); + + tga_data = (unsigned char*)stbi__malloc_mad3(tga_width, tga_height, tga_comp, 0); + if (!tga_data) return stbi__errpuc("outofmem", "Out of memory"); + + // skip to the data's starting position (offset usually = 0) + stbi__skip(s, tga_offset ); + + if ( !tga_indexed && !tga_is_RLE && !tga_rgb16 ) { + for (i=0; i < tga_height; ++i) { + int row = tga_inverted ? tga_height -i - 1 : i; + stbi_uc *tga_row = tga_data + row*tga_width*tga_comp; + stbi__getn(s, tga_row, tga_width * tga_comp); + } + } else { + // do I need to load a palette? + if ( tga_indexed) + { + if (tga_palette_len == 0) { /* you have to have at least one entry! */ + STBI_FREE(tga_data); + return stbi__errpuc("bad palette", "Corrupt TGA"); + } + + // any data to skip? (offset usually = 0) + stbi__skip(s, tga_palette_start ); + // load the palette + tga_palette = (unsigned char*)stbi__malloc_mad2(tga_palette_len, tga_comp, 0); + if (!tga_palette) { + STBI_FREE(tga_data); + return stbi__errpuc("outofmem", "Out of memory"); + } + if (tga_rgb16) { + stbi_uc *pal_entry = tga_palette; + STBI_ASSERT(tga_comp == STBI_rgb); + for (i=0; i < tga_palette_len; ++i) { + stbi__tga_read_rgb16(s, pal_entry); + pal_entry += tga_comp; + } + } else if (!stbi__getn(s, tga_palette, tga_palette_len * tga_comp)) { + STBI_FREE(tga_data); + STBI_FREE(tga_palette); + return stbi__errpuc("bad palette", "Corrupt TGA"); + } + } + // load the data + for (i=0; i < tga_width * tga_height; ++i) + { + // if I'm in RLE mode, do I need to get a RLE stbi__pngchunk? + if ( tga_is_RLE ) + { + if ( RLE_count == 0 ) + { + // yep, get the next byte as a RLE command + int RLE_cmd = stbi__get8(s); + RLE_count = 1 + (RLE_cmd & 127); + RLE_repeating = RLE_cmd >> 7; + read_next_pixel = 1; + } else if ( !RLE_repeating ) + { + read_next_pixel = 1; + } + } else + { + read_next_pixel = 1; + } + // OK, if I need to read a pixel, do it now + if ( read_next_pixel ) + { + // load however much data we did have + if ( tga_indexed ) + { + // read in index, then perform the lookup + int pal_idx = (tga_bits_per_pixel == 8) ? stbi__get8(s) : stbi__get16le(s); + if ( pal_idx >= tga_palette_len ) { + // invalid index + pal_idx = 0; + } + pal_idx *= tga_comp; + for (j = 0; j < tga_comp; ++j) { + raw_data[j] = tga_palette[pal_idx+j]; + } + } else if(tga_rgb16) { + STBI_ASSERT(tga_comp == STBI_rgb); + stbi__tga_read_rgb16(s, raw_data); + } else { + // read in the data raw + for (j = 0; j < tga_comp; ++j) { + raw_data[j] = stbi__get8(s); + } + } + // clear the reading flag for the next pixel + read_next_pixel = 0; + } // end of reading a pixel + + // copy data + for (j = 0; j < tga_comp; ++j) + tga_data[i*tga_comp+j] = raw_data[j]; + + // in case we're in RLE mode, keep counting down + --RLE_count; + } + // do I need to invert the image? + if ( tga_inverted ) + { + for (j = 0; j*2 < tga_height; ++j) + { + int index1 = j * tga_width * tga_comp; + int index2 = (tga_height - 1 - j) * tga_width * tga_comp; + for (i = tga_width * tga_comp; i > 0; --i) + { + unsigned char temp = tga_data[index1]; + tga_data[index1] = tga_data[index2]; + tga_data[index2] = temp; + ++index1; + ++index2; + } + } + } + // clear my palette, if I had one + if ( tga_palette != NULL ) + { + STBI_FREE( tga_palette ); + } + } + + // swap RGB - if the source data was RGB16, it already is in the right order + if (tga_comp >= 3 && !tga_rgb16) + { + unsigned char* tga_pixel = tga_data; + for (i=0; i < tga_width * tga_height; ++i) + { + unsigned char temp = tga_pixel[0]; + tga_pixel[0] = tga_pixel[2]; + tga_pixel[2] = temp; + tga_pixel += tga_comp; + } + } + + // convert to target component count + if (req_comp && req_comp != tga_comp) + tga_data = stbi__convert_format(tga_data, tga_comp, req_comp, tga_width, tga_height); + + // the things I do to get rid of an error message, and yet keep + // Microsoft's C compilers happy... [8^( + tga_palette_start = tga_palette_len = tga_palette_bits = + tga_x_origin = tga_y_origin = 0; + STBI_NOTUSED(tga_palette_start); + // OK, done + return tga_data; +} +#endif + +// ************************************************************************************************* +// Photoshop PSD loader -- PD by Thatcher Ulrich, integration by Nicolas Schulz, tweaked by STB + +#ifndef STBI_NO_PSD +static int stbi__psd_test(stbi__context *s) +{ + int r = (stbi__get32be(s) == 0x38425053); + stbi__rewind(s); + return r; +} + +static int stbi__psd_decode_rle(stbi__context *s, stbi_uc *p, int pixelCount) +{ + int count, nleft, len; + + count = 0; + while ((nleft = pixelCount - count) > 0) { + len = stbi__get8(s); + if (len == 128) { + // No-op. + } else if (len < 128) { + // Copy next len+1 bytes literally. + len++; + if (len > nleft) return 0; // corrupt data + count += len; + while (len) { + *p = stbi__get8(s); + p += 4; + len--; + } + } else if (len > 128) { + stbi_uc val; + // Next -len+1 bytes in the dest are replicated from next source byte. + // (Interpret len as a negative 8-bit int.) + len = 257 - len; + if (len > nleft) return 0; // corrupt data + val = stbi__get8(s); + count += len; + while (len) { + *p = val; + p += 4; + len--; + } + } + } + + return 1; +} + +static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) +{ + int pixelCount; + int channelCount, compression; + int channel, i; + int bitdepth; + int w,h; + stbi_uc *out; + STBI_NOTUSED(ri); + + // Check identifier + if (stbi__get32be(s) != 0x38425053) // "8BPS" + return stbi__errpuc("not PSD", "Corrupt PSD image"); + + // Check file type version. + if (stbi__get16be(s) != 1) + return stbi__errpuc("wrong version", "Unsupported version of PSD image"); + + // Skip 6 reserved bytes. + stbi__skip(s, 6 ); + + // Read the number of channels (R, G, B, A, etc). + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) + return stbi__errpuc("wrong channel count", "Unsupported number of channels in PSD image"); + + // Read the rows and columns of the image. + h = stbi__get32be(s); + w = stbi__get32be(s); + + if (h > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (w > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + // Make sure the depth is 8 bits. + bitdepth = stbi__get16be(s); + if (bitdepth != 8 && bitdepth != 16) + return stbi__errpuc("unsupported bit depth", "PSD bit depth is not 8 or 16 bit"); + + // Make sure the color mode is RGB. + // Valid options are: + // 0: Bitmap + // 1: Grayscale + // 2: Indexed color + // 3: RGB color + // 4: CMYK color + // 7: Multichannel + // 8: Duotone + // 9: Lab color + if (stbi__get16be(s) != 3) + return stbi__errpuc("wrong color format", "PSD is not in RGB color format"); + + // Skip the Mode Data. (It's the palette for indexed color; other info for other modes.) + stbi__skip(s,stbi__get32be(s) ); + + // Skip the image resources. (resolution, pen tool paths, etc) + stbi__skip(s, stbi__get32be(s) ); + + // Skip the reserved data. + stbi__skip(s, stbi__get32be(s) ); + + // Find out if the data is compressed. + // Known values: + // 0: no compression + // 1: RLE compressed + compression = stbi__get16be(s); + if (compression > 1) + return stbi__errpuc("bad compression", "PSD has an unknown compression format"); + + // Check size + if (!stbi__mad3sizes_valid(4, w, h, 0)) + return stbi__errpuc("too large", "Corrupt PSD"); + + // Create the destination image. + + if (!compression && bitdepth == 16 && bpc == 16) { + out = (stbi_uc *) stbi__malloc_mad3(8, w, h, 0); + ri->bits_per_channel = 16; + } else + out = (stbi_uc *) stbi__malloc(4 * w*h); + + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + pixelCount = w*h; + + // Initialize the data to zero. + //memset( out, 0, pixelCount * 4 ); + + // Finally, the image data. + if (compression) { + // RLE as used by .PSD and .TIFF + // Loop until you get the number of unpacked bytes you are expecting: + // Read the next source byte into n. + // If n is between 0 and 127 inclusive, copy the next n+1 bytes literally. + // Else if n is between -127 and -1 inclusive, copy the next byte -n+1 times. + // Else if n is 128, noop. + // Endloop + + // The RLE-compressed data is preceded by a 2-byte data count for each row in the data, + // which we're going to just skip. + stbi__skip(s, h * channelCount * 2 ); + + // Read the RLE data by channel. + for (channel = 0; channel < 4; channel++) { + stbi_uc *p; + + p = out+channel; + if (channel >= channelCount) { + // Fill this channel with default data. + for (i = 0; i < pixelCount; i++, p += 4) + *p = (channel == 3 ? 255 : 0); + } else { + // Read the RLE data. + if (!stbi__psd_decode_rle(s, p, pixelCount)) { + STBI_FREE(out); + return stbi__errpuc("corrupt", "bad RLE data"); + } + } + } + + } else { + // We're at the raw image data. It's each channel in order (Red, Green, Blue, Alpha, ...) + // where each channel consists of an 8-bit (or 16-bit) value for each pixel in the image. + + // Read the data by channel. + for (channel = 0; channel < 4; channel++) { + if (channel >= channelCount) { + // Fill this channel with default data. + if (bitdepth == 16 && bpc == 16) { + stbi__uint16 *q = ((stbi__uint16 *) out) + channel; + stbi__uint16 val = channel == 3 ? 65535 : 0; + for (i = 0; i < pixelCount; i++, q += 4) + *q = val; + } else { + stbi_uc *p = out+channel; + stbi_uc val = channel == 3 ? 255 : 0; + for (i = 0; i < pixelCount; i++, p += 4) + *p = val; + } + } else { + if (ri->bits_per_channel == 16) { // output bpc + stbi__uint16 *q = ((stbi__uint16 *) out) + channel; + for (i = 0; i < pixelCount; i++, q += 4) + *q = (stbi__uint16) stbi__get16be(s); + } else { + stbi_uc *p = out+channel; + if (bitdepth == 16) { // input bpc + for (i = 0; i < pixelCount; i++, p += 4) + *p = (stbi_uc) (stbi__get16be(s) >> 8); + } else { + for (i = 0; i < pixelCount; i++, p += 4) + *p = stbi__get8(s); + } + } + } + } + } + + // remove weird white matte from PSD + if (channelCount >= 4) { + if (ri->bits_per_channel == 16) { + for (i=0; i < w*h; ++i) { + stbi__uint16 *pixel = (stbi__uint16 *) out + 4*i; + if (pixel[3] != 0 && pixel[3] != 65535) { + float a = pixel[3] / 65535.0f; + float ra = 1.0f / a; + float inv_a = 65535.0f * (1 - ra); + pixel[0] = (stbi__uint16) (pixel[0]*ra + inv_a); + pixel[1] = (stbi__uint16) (pixel[1]*ra + inv_a); + pixel[2] = (stbi__uint16) (pixel[2]*ra + inv_a); + } + } + } else { + for (i=0; i < w*h; ++i) { + unsigned char *pixel = out + 4*i; + if (pixel[3] != 0 && pixel[3] != 255) { + float a = pixel[3] / 255.0f; + float ra = 1.0f / a; + float inv_a = 255.0f * (1 - ra); + pixel[0] = (unsigned char) (pixel[0]*ra + inv_a); + pixel[1] = (unsigned char) (pixel[1]*ra + inv_a); + pixel[2] = (unsigned char) (pixel[2]*ra + inv_a); + } + } + } + } + + // convert to desired output format + if (req_comp && req_comp != 4) { + if (ri->bits_per_channel == 16) + out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, 4, req_comp, w, h); + else + out = stbi__convert_format(out, 4, req_comp, w, h); + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + + if (comp) *comp = 4; + *y = h; + *x = w; + + return out; +} +#endif + +// ************************************************************************************************* +// Softimage PIC loader +// by Tom Seddon +// +// See http://softimage.wiki.softimage.com/index.php/INFO:_PIC_file_format +// See http://ozviz.wasp.uwa.edu.au/~pbourke/dataformats/softimagepic/ + +#ifndef STBI_NO_PIC +static int stbi__pic_is4(stbi__context *s,const char *str) +{ + int i; + for (i=0; i<4; ++i) + if (stbi__get8(s) != (stbi_uc)str[i]) + return 0; + + return 1; +} + +static int stbi__pic_test_core(stbi__context *s) +{ + int i; + + if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) + return 0; + + for(i=0;i<84;++i) + stbi__get8(s); + + if (!stbi__pic_is4(s,"PICT")) + return 0; + + return 1; +} + +typedef struct +{ + stbi_uc size,type,channel; +} stbi__pic_packet; + +static stbi_uc *stbi__readval(stbi__context *s, int channel, stbi_uc *dest) +{ + int mask=0x80, i; + + for (i=0; i<4; ++i, mask>>=1) { + if (channel & mask) { + if (stbi__at_eof(s)) return stbi__errpuc("bad file","PIC file too short"); + dest[i]=stbi__get8(s); + } + } + + return dest; +} + +static void stbi__copyval(int channel,stbi_uc *dest,const stbi_uc *src) +{ + int mask=0x80,i; + + for (i=0;i<4; ++i, mask>>=1) + if (channel&mask) + dest[i]=src[i]; +} + +static stbi_uc *stbi__pic_load_core(stbi__context *s,int width,int height,int *comp, stbi_uc *result) +{ + int act_comp=0,num_packets=0,y,chained; + stbi__pic_packet packets[10]; + + // this will (should...) cater for even some bizarre stuff like having data + // for the same channel in multiple packets. + do { + stbi__pic_packet *packet; + + if (num_packets==sizeof(packets)/sizeof(packets[0])) + return stbi__errpuc("bad format","too many packets"); + + packet = &packets[num_packets++]; + + chained = stbi__get8(s); + packet->size = stbi__get8(s); + packet->type = stbi__get8(s); + packet->channel = stbi__get8(s); + + act_comp |= packet->channel; + + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (reading packets)"); + if (packet->size != 8) return stbi__errpuc("bad format","packet isn't 8bpp"); + } while (chained); + + *comp = (act_comp & 0x10 ? 4 : 3); // has alpha channel? + + for(y=0; ytype) { + default: + return stbi__errpuc("bad format","packet has bad compression type"); + + case 0: {//uncompressed + int x; + + for(x=0;xchannel,dest)) + return 0; + break; + } + + case 1://Pure RLE + { + int left=width, i; + + while (left>0) { + stbi_uc count,value[4]; + + count=stbi__get8(s); + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pure read count)"); + + if (count > left) + count = (stbi_uc) left; + + if (!stbi__readval(s,packet->channel,value)) return 0; + + for(i=0; ichannel,dest,value); + left -= count; + } + } + break; + + case 2: {//Mixed RLE + int left=width; + while (left>0) { + int count = stbi__get8(s), i; + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (mixed read count)"); + + if (count >= 128) { // Repeated + stbi_uc value[4]; + + if (count==128) + count = stbi__get16be(s); + else + count -= 127; + if (count > left) + return stbi__errpuc("bad file","scanline overrun"); + + if (!stbi__readval(s,packet->channel,value)) + return 0; + + for(i=0;ichannel,dest,value); + } else { // Raw + ++count; + if (count>left) return stbi__errpuc("bad file","scanline overrun"); + + for(i=0;ichannel,dest)) + return 0; + } + left-=count; + } + break; + } + } + } + } + + return result; +} + +static void *stbi__pic_load(stbi__context *s,int *px,int *py,int *comp,int req_comp, stbi__result_info *ri) +{ + stbi_uc *result; + int i, x,y, internal_comp; + STBI_NOTUSED(ri); + + if (!comp) comp = &internal_comp; + + for (i=0; i<92; ++i) + stbi__get8(s); + + x = stbi__get16be(s); + y = stbi__get16be(s); + + if (y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pic header)"); + if (!stbi__mad3sizes_valid(x, y, 4, 0)) return stbi__errpuc("too large", "PIC image too large to decode"); + + stbi__get32be(s); //skip `ratio' + stbi__get16be(s); //skip `fields' + stbi__get16be(s); //skip `pad' + + // intermediate buffer is RGBA + result = (stbi_uc *) stbi__malloc_mad3(x, y, 4, 0); + if (!result) return stbi__errpuc("outofmem", "Out of memory"); + memset(result, 0xff, x*y*4); + + if (!stbi__pic_load_core(s,x,y,comp, result)) { + STBI_FREE(result); + result=0; + } + *px = x; + *py = y; + if (req_comp == 0) req_comp = *comp; + result=stbi__convert_format(result,4,req_comp,x,y); + + return result; +} + +static int stbi__pic_test(stbi__context *s) +{ + int r = stbi__pic_test_core(s); + stbi__rewind(s); + return r; +} +#endif + +// ************************************************************************************************* +// GIF loader -- public domain by Jean-Marc Lienher -- simplified/shrunk by stb + +#ifndef STBI_NO_GIF +typedef struct +{ + stbi__int16 prefix; + stbi_uc first; + stbi_uc suffix; +} stbi__gif_lzw; + +typedef struct +{ + int w,h; + stbi_uc *out; // output buffer (always 4 components) + stbi_uc *background; // The current "background" as far as a gif is concerned + stbi_uc *history; + int flags, bgindex, ratio, transparent, eflags; + stbi_uc pal[256][4]; + stbi_uc lpal[256][4]; + stbi__gif_lzw codes[8192]; + stbi_uc *color_table; + int parse, step; + int lflags; + int start_x, start_y; + int max_x, max_y; + int cur_x, cur_y; + int line_size; + int delay; +} stbi__gif; + +static int stbi__gif_test_raw(stbi__context *s) +{ + int sz; + if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') return 0; + sz = stbi__get8(s); + if (sz != '9' && sz != '7') return 0; + if (stbi__get8(s) != 'a') return 0; + return 1; +} + +static int stbi__gif_test(stbi__context *s) +{ + int r = stbi__gif_test_raw(s); + stbi__rewind(s); + return r; +} + +static void stbi__gif_parse_colortable(stbi__context *s, stbi_uc pal[256][4], int num_entries, int transp) +{ + int i; + for (i=0; i < num_entries; ++i) { + pal[i][2] = stbi__get8(s); + pal[i][1] = stbi__get8(s); + pal[i][0] = stbi__get8(s); + pal[i][3] = transp == i ? 0 : 255; + } +} + +static int stbi__gif_header(stbi__context *s, stbi__gif *g, int *comp, int is_info) +{ + stbi_uc version; + if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') + return stbi__err("not GIF", "Corrupt GIF"); + + version = stbi__get8(s); + if (version != '7' && version != '9') return stbi__err("not GIF", "Corrupt GIF"); + if (stbi__get8(s) != 'a') return stbi__err("not GIF", "Corrupt GIF"); + + stbi__g_failure_reason = ""; + g->w = stbi__get16le(s); + g->h = stbi__get16le(s); + g->flags = stbi__get8(s); + g->bgindex = stbi__get8(s); + g->ratio = stbi__get8(s); + g->transparent = -1; + + if (g->w > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (g->h > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + + if (comp != 0) *comp = 4; // can't actually tell whether it's 3 or 4 until we parse the comments + + if (is_info) return 1; + + if (g->flags & 0x80) + stbi__gif_parse_colortable(s,g->pal, 2 << (g->flags & 7), -1); + + return 1; +} + +static int stbi__gif_info_raw(stbi__context *s, int *x, int *y, int *comp) +{ + stbi__gif* g = (stbi__gif*) stbi__malloc(sizeof(stbi__gif)); + if (!g) return stbi__err("outofmem", "Out of memory"); + if (!stbi__gif_header(s, g, comp, 1)) { + STBI_FREE(g); + stbi__rewind( s ); + return 0; + } + if (x) *x = g->w; + if (y) *y = g->h; + STBI_FREE(g); + return 1; +} + +static void stbi__out_gif_code(stbi__gif *g, stbi__uint16 code) +{ + stbi_uc *p, *c; + int idx; + + // recurse to decode the prefixes, since the linked-list is backwards, + // and working backwards through an interleaved image would be nasty + if (g->codes[code].prefix >= 0) + stbi__out_gif_code(g, g->codes[code].prefix); + + if (g->cur_y >= g->max_y) return; + + idx = g->cur_x + g->cur_y; + p = &g->out[idx]; + g->history[idx / 4] = 1; + + c = &g->color_table[g->codes[code].suffix * 4]; + if (c[3] > 128) { // don't render transparent pixels; + p[0] = c[2]; + p[1] = c[1]; + p[2] = c[0]; + p[3] = c[3]; + } + g->cur_x += 4; + + if (g->cur_x >= g->max_x) { + g->cur_x = g->start_x; + g->cur_y += g->step; + + while (g->cur_y >= g->max_y && g->parse > 0) { + g->step = (1 << g->parse) * g->line_size; + g->cur_y = g->start_y + (g->step >> 1); + --g->parse; + } + } +} + +static stbi_uc *stbi__process_gif_raster(stbi__context *s, stbi__gif *g) +{ + stbi_uc lzw_cs; + stbi__int32 len, init_code; + stbi__uint32 first; + stbi__int32 codesize, codemask, avail, oldcode, bits, valid_bits, clear; + stbi__gif_lzw *p; + + lzw_cs = stbi__get8(s); + if (lzw_cs > 12) return NULL; + clear = 1 << lzw_cs; + first = 1; + codesize = lzw_cs + 1; + codemask = (1 << codesize) - 1; + bits = 0; + valid_bits = 0; + for (init_code = 0; init_code < clear; init_code++) { + g->codes[init_code].prefix = -1; + g->codes[init_code].first = (stbi_uc) init_code; + g->codes[init_code].suffix = (stbi_uc) init_code; + } + + // support no starting clear code + avail = clear+2; + oldcode = -1; + + len = 0; + for(;;) { + if (valid_bits < codesize) { + if (len == 0) { + len = stbi__get8(s); // start new block + if (len == 0) + return g->out; + } + --len; + bits |= (stbi__int32) stbi__get8(s) << valid_bits; + valid_bits += 8; + } else { + stbi__int32 code = bits & codemask; + bits >>= codesize; + valid_bits -= codesize; + // @OPTIMIZE: is there some way we can accelerate the non-clear path? + if (code == clear) { // clear code + codesize = lzw_cs + 1; + codemask = (1 << codesize) - 1; + avail = clear + 2; + oldcode = -1; + first = 0; + } else if (code == clear + 1) { // end of stream code + stbi__skip(s, len); + while ((len = stbi__get8(s)) > 0) + stbi__skip(s,len); + return g->out; + } else if (code <= avail) { + if (first) { + return stbi__errpuc("no clear code", "Corrupt GIF"); + } + + if (oldcode >= 0) { + p = &g->codes[avail++]; + if (avail > 8192) { + return stbi__errpuc("too many codes", "Corrupt GIF"); + } + + p->prefix = (stbi__int16) oldcode; + p->first = g->codes[oldcode].first; + p->suffix = (code == avail) ? p->first : g->codes[code].first; + } else if (code == avail) + return stbi__errpuc("illegal code in raster", "Corrupt GIF"); + + stbi__out_gif_code(g, (stbi__uint16) code); + + if ((avail & codemask) == 0 && avail <= 0x0FFF) { + codesize++; + codemask = (1 << codesize) - 1; + } + + oldcode = code; + } else { + return stbi__errpuc("illegal code in raster", "Corrupt GIF"); + } + } + } +} + +// this function is designed to support animated gifs, although stb_image doesn't support it +// two back is the image from two frames ago, used for a very specific disposal format +static stbi_uc *stbi__gif_load_next(stbi__context *s, stbi__gif *g, int *comp, int req_comp, stbi_uc *two_back) +{ + int dispose; + int first_frame; + int pi; + int pcount; + STBI_NOTUSED(req_comp); + + // on first frame, any non-written pixels get the background colour (non-transparent) + first_frame = 0; + if (g->out == 0) { + if (!stbi__gif_header(s, g, comp,0)) return 0; // stbi__g_failure_reason set by stbi__gif_header + if (!stbi__mad3sizes_valid(4, g->w, g->h, 0)) + return stbi__errpuc("too large", "GIF image is too large"); + pcount = g->w * g->h; + g->out = (stbi_uc *) stbi__malloc(4 * pcount); + g->background = (stbi_uc *) stbi__malloc(4 * pcount); + g->history = (stbi_uc *) stbi__malloc(pcount); + if (!g->out || !g->background || !g->history) + return stbi__errpuc("outofmem", "Out of memory"); + + // image is treated as "transparent" at the start - ie, nothing overwrites the current background; + // background colour is only used for pixels that are not rendered first frame, after that "background" + // color refers to the color that was there the previous frame. + memset(g->out, 0x00, 4 * pcount); + memset(g->background, 0x00, 4 * pcount); // state of the background (starts transparent) + memset(g->history, 0x00, pcount); // pixels that were affected previous frame + first_frame = 1; + } else { + // second frame - how do we dispose of the previous one? + dispose = (g->eflags & 0x1C) >> 2; + pcount = g->w * g->h; + + if ((dispose == 3) && (two_back == 0)) { + dispose = 2; // if I don't have an image to revert back to, default to the old background + } + + if (dispose == 3) { // use previous graphic + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi]) { + memcpy( &g->out[pi * 4], &two_back[pi * 4], 4 ); + } + } + } else if (dispose == 2) { + // restore what was changed last frame to background before that frame; + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi]) { + memcpy( &g->out[pi * 4], &g->background[pi * 4], 4 ); + } + } + } else { + // This is a non-disposal case eithe way, so just + // leave the pixels as is, and they will become the new background + // 1: do not dispose + // 0: not specified. + } + + // background is what out is after the undoing of the previou frame; + memcpy( g->background, g->out, 4 * g->w * g->h ); + } + + // clear my history; + memset( g->history, 0x00, g->w * g->h ); // pixels that were affected previous frame + + for (;;) { + int tag = stbi__get8(s); + switch (tag) { + case 0x2C: /* Image Descriptor */ + { + stbi__int32 x, y, w, h; + stbi_uc *o; + + x = stbi__get16le(s); + y = stbi__get16le(s); + w = stbi__get16le(s); + h = stbi__get16le(s); + if (((x + w) > (g->w)) || ((y + h) > (g->h))) + return stbi__errpuc("bad Image Descriptor", "Corrupt GIF"); + + g->line_size = g->w * 4; + g->start_x = x * 4; + g->start_y = y * g->line_size; + g->max_x = g->start_x + w * 4; + g->max_y = g->start_y + h * g->line_size; + g->cur_x = g->start_x; + g->cur_y = g->start_y; + + // if the width of the specified rectangle is 0, that means + // we may not see *any* pixels or the image is malformed; + // to make sure this is caught, move the current y down to + // max_y (which is what out_gif_code checks). + if (w == 0) + g->cur_y = g->max_y; + + g->lflags = stbi__get8(s); + + if (g->lflags & 0x40) { + g->step = 8 * g->line_size; // first interlaced spacing + g->parse = 3; + } else { + g->step = g->line_size; + g->parse = 0; + } + + if (g->lflags & 0x80) { + stbi__gif_parse_colortable(s,g->lpal, 2 << (g->lflags & 7), g->eflags & 0x01 ? g->transparent : -1); + g->color_table = (stbi_uc *) g->lpal; + } else if (g->flags & 0x80) { + g->color_table = (stbi_uc *) g->pal; + } else + return stbi__errpuc("missing color table", "Corrupt GIF"); + + o = stbi__process_gif_raster(s, g); + if (!o) return NULL; + + // if this was the first frame, + pcount = g->w * g->h; + if (first_frame && (g->bgindex > 0)) { + // if first frame, any pixel not drawn to gets the background color + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi] == 0) { + g->pal[g->bgindex][3] = 255; // just in case it was made transparent, undo that; It will be reset next frame if need be; + memcpy( &g->out[pi * 4], &g->pal[g->bgindex], 4 ); + } + } + } + + return o; + } + + case 0x21: // Comment Extension. + { + int len; + int ext = stbi__get8(s); + if (ext == 0xF9) { // Graphic Control Extension. + len = stbi__get8(s); + if (len == 4) { + g->eflags = stbi__get8(s); + g->delay = 10 * stbi__get16le(s); // delay - 1/100th of a second, saving as 1/1000ths. + + // unset old transparent + if (g->transparent >= 0) { + g->pal[g->transparent][3] = 255; + } + if (g->eflags & 0x01) { + g->transparent = stbi__get8(s); + if (g->transparent >= 0) { + g->pal[g->transparent][3] = 0; + } + } else { + // don't need transparent + stbi__skip(s, 1); + g->transparent = -1; + } + } else { + stbi__skip(s, len); + break; + } + } + while ((len = stbi__get8(s)) != 0) { + stbi__skip(s, len); + } + break; + } + + case 0x3B: // gif stream termination code + return (stbi_uc *) s; // using '1' causes warning on some compilers + + default: + return stbi__errpuc("unknown code", "Corrupt GIF"); + } + } +} + +static void *stbi__load_gif_main_outofmem(stbi__gif *g, stbi_uc *out, int **delays) +{ + STBI_FREE(g->out); + STBI_FREE(g->history); + STBI_FREE(g->background); + + if (out) STBI_FREE(out); + if (delays && *delays) STBI_FREE(*delays); + return stbi__errpuc("outofmem", "Out of memory"); +} + +static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp) +{ + if (stbi__gif_test(s)) { + int layers = 0; + stbi_uc *u = 0; + stbi_uc *out = 0; + stbi_uc *two_back = 0; + stbi__gif g; + int stride; + int out_size = 0; + int delays_size = 0; + + STBI_NOTUSED(out_size); + STBI_NOTUSED(delays_size); + + memset(&g, 0, sizeof(g)); + if (delays) { + *delays = 0; + } + + do { + u = stbi__gif_load_next(s, &g, comp, req_comp, two_back); + if (u == (stbi_uc *) s) u = 0; // end of animated gif marker + + if (u) { + *x = g.w; + *y = g.h; + ++layers; + stride = g.w * g.h * 4; + + if (out) { + void *tmp = (stbi_uc*) STBI_REALLOC_SIZED( out, out_size, layers * stride ); + if (!tmp) + return stbi__load_gif_main_outofmem(&g, out, delays); + else { + out = (stbi_uc*) tmp; + out_size = layers * stride; + } + + if (delays) { + int *new_delays = (int*) STBI_REALLOC_SIZED( *delays, delays_size, sizeof(int) * layers ); + if (!new_delays) + return stbi__load_gif_main_outofmem(&g, out, delays); + *delays = new_delays; + delays_size = layers * sizeof(int); + } + } else { + out = (stbi_uc*)stbi__malloc( layers * stride ); + if (!out) + return stbi__load_gif_main_outofmem(&g, out, delays); + out_size = layers * stride; + if (delays) { + *delays = (int*) stbi__malloc( layers * sizeof(int) ); + if (!*delays) + return stbi__load_gif_main_outofmem(&g, out, delays); + delays_size = layers * sizeof(int); + } + } + memcpy( out + ((layers - 1) * stride), u, stride ); + if (layers >= 2) { + two_back = out - 2 * stride; + } + + if (delays) { + (*delays)[layers - 1U] = g.delay; + } + } + } while (u != 0); + + // free temp buffer; + STBI_FREE(g.out); + STBI_FREE(g.history); + STBI_FREE(g.background); + + // do the final conversion after loading everything; + if (req_comp && req_comp != 4) + out = stbi__convert_format(out, 4, req_comp, layers * g.w, g.h); + + *z = layers; + return out; + } else { + return stbi__errpuc("not GIF", "Image was not as a gif type."); + } +} + +static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *u = 0; + stbi__gif g; + memset(&g, 0, sizeof(g)); + STBI_NOTUSED(ri); + + u = stbi__gif_load_next(s, &g, comp, req_comp, 0); + if (u == (stbi_uc *) s) u = 0; // end of animated gif marker + if (u) { + *x = g.w; + *y = g.h; + + // moved conversion to after successful load so that the same + // can be done for multiple frames. + if (req_comp && req_comp != 4) + u = stbi__convert_format(u, 4, req_comp, g.w, g.h); + } else if (g.out) { + // if there was an error and we allocated an image buffer, free it! + STBI_FREE(g.out); + } + + // free buffers needed for multiple frame loading; + STBI_FREE(g.history); + STBI_FREE(g.background); + + return u; +} + +static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp) +{ + return stbi__gif_info_raw(s,x,y,comp); +} +#endif + +// ************************************************************************************************* +// Radiance RGBE HDR loader +// originally by Nicolas Schulz +#ifndef STBI_NO_HDR +static int stbi__hdr_test_core(stbi__context *s, const char *signature) +{ + int i; + for (i=0; signature[i]; ++i) + if (stbi__get8(s) != signature[i]) + return 0; + stbi__rewind(s); + return 1; +} + +static int stbi__hdr_test(stbi__context* s) +{ + int r = stbi__hdr_test_core(s, "#?RADIANCE\n"); + stbi__rewind(s); + if(!r) { + r = stbi__hdr_test_core(s, "#?RGBE\n"); + stbi__rewind(s); + } + return r; +} + +#define STBI__HDR_BUFLEN 1024 +static char *stbi__hdr_gettoken(stbi__context *z, char *buffer) +{ + int len=0; + char c = '\0'; + + c = (char) stbi__get8(z); + + while (!stbi__at_eof(z) && c != '\n') { + buffer[len++] = c; + if (len == STBI__HDR_BUFLEN-1) { + // flush to end of line + while (!stbi__at_eof(z) && stbi__get8(z) != '\n') + ; + break; + } + c = (char) stbi__get8(z); + } + + buffer[len] = 0; + return buffer; +} + +static void stbi__hdr_convert(float *output, stbi_uc *input, int req_comp) +{ + if ( input[3] != 0 ) { + float f1; + // Exponent + f1 = (float) ldexp(1.0f, input[3] - (int)(128 + 8)); + if (req_comp <= 2) + output[0] = (input[0] + input[1] + input[2]) * f1 / 3; + else { + output[0] = input[0] * f1; + output[1] = input[1] * f1; + output[2] = input[2] * f1; + } + if (req_comp == 2) output[1] = 1; + if (req_comp == 4) output[3] = 1; + } else { + switch (req_comp) { + case 4: output[3] = 1; /* fallthrough */ + case 3: output[0] = output[1] = output[2] = 0; + break; + case 2: output[1] = 1; /* fallthrough */ + case 1: output[0] = 0; + break; + } + } +} + +static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + char buffer[STBI__HDR_BUFLEN]; + char *token; + int valid = 0; + int width, height; + stbi_uc *scanline; + float *hdr_data; + int len; + unsigned char count, value; + int i, j, k, c1,c2, z; + const char *headerToken; + STBI_NOTUSED(ri); + + // Check identifier + headerToken = stbi__hdr_gettoken(s,buffer); + if (strcmp(headerToken, "#?RADIANCE") != 0 && strcmp(headerToken, "#?RGBE") != 0) + return stbi__errpf("not HDR", "Corrupt HDR image"); + + // Parse header + for(;;) { + token = stbi__hdr_gettoken(s,buffer); + if (token[0] == 0) break; + if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; + } + + if (!valid) return stbi__errpf("unsupported format", "Unsupported HDR format"); + + // Parse width and height + // can't use sscanf() if we're not using stdio! + token = stbi__hdr_gettoken(s,buffer); + if (strncmp(token, "-Y ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); + token += 3; + height = (int) strtol(token, &token, 10); + while (*token == ' ') ++token; + if (strncmp(token, "+X ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); + token += 3; + width = (int) strtol(token, NULL, 10); + + if (height > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); + if (width > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); + + *x = width; + *y = height; + + if (comp) *comp = 3; + if (req_comp == 0) req_comp = 3; + + if (!stbi__mad4sizes_valid(width, height, req_comp, sizeof(float), 0)) + return stbi__errpf("too large", "HDR image is too large"); + + // Read data + hdr_data = (float *) stbi__malloc_mad4(width, height, req_comp, sizeof(float), 0); + if (!hdr_data) + return stbi__errpf("outofmem", "Out of memory"); + + // Load image data + // image data is stored as some number of sca + if ( width < 8 || width >= 32768) { + // Read flat data + for (j=0; j < height; ++j) { + for (i=0; i < width; ++i) { + stbi_uc rgbe[4]; + main_decode_loop: + stbi__getn(s, rgbe, 4); + stbi__hdr_convert(hdr_data + j * width * req_comp + i * req_comp, rgbe, req_comp); + } + } + } else { + // Read RLE-encoded data + scanline = NULL; + + for (j = 0; j < height; ++j) { + c1 = stbi__get8(s); + c2 = stbi__get8(s); + len = stbi__get8(s); + if (c1 != 2 || c2 != 2 || (len & 0x80)) { + // not run-length encoded, so we have to actually use THIS data as a decoded + // pixel (note this can't be a valid pixel--one of RGB must be >= 128) + stbi_uc rgbe[4]; + rgbe[0] = (stbi_uc) c1; + rgbe[1] = (stbi_uc) c2; + rgbe[2] = (stbi_uc) len; + rgbe[3] = (stbi_uc) stbi__get8(s); + stbi__hdr_convert(hdr_data, rgbe, req_comp); + i = 1; + j = 0; + STBI_FREE(scanline); + goto main_decode_loop; // yes, this makes no sense + } + len <<= 8; + len |= stbi__get8(s); + if (len != width) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("invalid decoded scanline length", "corrupt HDR"); } + if (scanline == NULL) { + scanline = (stbi_uc *) stbi__malloc_mad2(width, 4, 0); + if (!scanline) { + STBI_FREE(hdr_data); + return stbi__errpf("outofmem", "Out of memory"); + } + } + + for (k = 0; k < 4; ++k) { + int nleft; + i = 0; + while ((nleft = width - i) > 0) { + count = stbi__get8(s); + if (count > 128) { + // Run + value = stbi__get8(s); + count -= 128; + if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } + for (z = 0; z < count; ++z) + scanline[i++ * 4 + k] = value; + } else { + // Dump + if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } + for (z = 0; z < count; ++z) + scanline[i++ * 4 + k] = stbi__get8(s); + } + } + } + for (i=0; i < width; ++i) + stbi__hdr_convert(hdr_data+(j*width + i)*req_comp, scanline + i*4, req_comp); + } + if (scanline) + STBI_FREE(scanline); + } + + return hdr_data; +} + +static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp) +{ + char buffer[STBI__HDR_BUFLEN]; + char *token; + int valid = 0; + int dummy; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + if (stbi__hdr_test(s) == 0) { + stbi__rewind( s ); + return 0; + } + + for(;;) { + token = stbi__hdr_gettoken(s,buffer); + if (token[0] == 0) break; + if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; + } + + if (!valid) { + stbi__rewind( s ); + return 0; + } + token = stbi__hdr_gettoken(s,buffer); + if (strncmp(token, "-Y ", 3)) { + stbi__rewind( s ); + return 0; + } + token += 3; + *y = (int) strtol(token, &token, 10); + while (*token == ' ') ++token; + if (strncmp(token, "+X ", 3)) { + stbi__rewind( s ); + return 0; + } + token += 3; + *x = (int) strtol(token, NULL, 10); + *comp = 3; + return 1; +} +#endif // STBI_NO_HDR + +#ifndef STBI_NO_BMP +static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp) +{ + void *p; + stbi__bmp_data info; + + info.all_a = 255; + p = stbi__bmp_parse_header(s, &info); + if (p == NULL) { + stbi__rewind( s ); + return 0; + } + if (x) *x = s->img_x; + if (y) *y = s->img_y; + if (comp) { + if (info.bpp == 24 && info.ma == 0xff000000) + *comp = 3; + else + *comp = info.ma ? 4 : 3; + } + return 1; +} +#endif + +#ifndef STBI_NO_PSD +static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp) +{ + int channelCount, dummy, depth; + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + if (stbi__get32be(s) != 0x38425053) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 1) { + stbi__rewind( s ); + return 0; + } + stbi__skip(s, 6); + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) { + stbi__rewind( s ); + return 0; + } + *y = stbi__get32be(s); + *x = stbi__get32be(s); + depth = stbi__get16be(s); + if (depth != 8 && depth != 16) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 3) { + stbi__rewind( s ); + return 0; + } + *comp = 4; + return 1; +} + +static int stbi__psd_is16(stbi__context *s) +{ + int channelCount, depth; + if (stbi__get32be(s) != 0x38425053) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 1) { + stbi__rewind( s ); + return 0; + } + stbi__skip(s, 6); + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) { + stbi__rewind( s ); + return 0; + } + STBI_NOTUSED(stbi__get32be(s)); + STBI_NOTUSED(stbi__get32be(s)); + depth = stbi__get16be(s); + if (depth != 16) { + stbi__rewind( s ); + return 0; + } + return 1; +} +#endif + +#ifndef STBI_NO_PIC +static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp) +{ + int act_comp=0,num_packets=0,chained,dummy; + stbi__pic_packet packets[10]; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) { + stbi__rewind(s); + return 0; + } + + stbi__skip(s, 88); + + *x = stbi__get16be(s); + *y = stbi__get16be(s); + if (stbi__at_eof(s)) { + stbi__rewind( s); + return 0; + } + if ( (*x) != 0 && (1 << 28) / (*x) < (*y)) { + stbi__rewind( s ); + return 0; + } + + stbi__skip(s, 8); + + do { + stbi__pic_packet *packet; + + if (num_packets==sizeof(packets)/sizeof(packets[0])) + return 0; + + packet = &packets[num_packets++]; + chained = stbi__get8(s); + packet->size = stbi__get8(s); + packet->type = stbi__get8(s); + packet->channel = stbi__get8(s); + act_comp |= packet->channel; + + if (stbi__at_eof(s)) { + stbi__rewind( s ); + return 0; + } + if (packet->size != 8) { + stbi__rewind( s ); + return 0; + } + } while (chained); + + *comp = (act_comp & 0x10 ? 4 : 3); + + return 1; +} +#endif + +// ************************************************************************************************* +// Portable Gray Map and Portable Pixel Map loader +// by Ken Miller +// +// PGM: http://netpbm.sourceforge.net/doc/pgm.html +// PPM: http://netpbm.sourceforge.net/doc/ppm.html +// +// Known limitations: +// Does not support comments in the header section +// Does not support ASCII image data (formats P2 and P3) + +#ifndef STBI_NO_PNM + +static int stbi__pnm_test(stbi__context *s) +{ + char p, t; + p = (char) stbi__get8(s); + t = (char) stbi__get8(s); + if (p != 'P' || (t != '5' && t != '6')) { + stbi__rewind( s ); + return 0; + } + return 1; +} + +static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *out; + STBI_NOTUSED(ri); + + ri->bits_per_channel = stbi__pnm_info(s, (int *)&s->img_x, (int *)&s->img_y, (int *)&s->img_n); + if (ri->bits_per_channel == 0) + return 0; + + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + *x = s->img_x; + *y = s->img_y; + if (comp) *comp = s->img_n; + + if (!stbi__mad4sizes_valid(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0)) + return stbi__errpuc("too large", "PNM too large"); + + out = (stbi_uc *) stbi__malloc_mad4(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0); + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + if (!stbi__getn(s, out, s->img_n * s->img_x * s->img_y * (ri->bits_per_channel / 8))) { + STBI_FREE(out); + return stbi__errpuc("bad PNM", "PNM file truncated"); + } + + if (req_comp && req_comp != s->img_n) { + if (ri->bits_per_channel == 16) { + out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, s->img_n, req_comp, s->img_x, s->img_y); + } else { + out = stbi__convert_format(out, s->img_n, req_comp, s->img_x, s->img_y); + } + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + return out; +} + +static int stbi__pnm_isspace(char c) +{ + return c == ' ' || c == '\t' || c == '\n' || c == '\v' || c == '\f' || c == '\r'; +} + +static void stbi__pnm_skip_whitespace(stbi__context *s, char *c) +{ + for (;;) { + while (!stbi__at_eof(s) && stbi__pnm_isspace(*c)) + *c = (char) stbi__get8(s); + + if (stbi__at_eof(s) || *c != '#') + break; + + while (!stbi__at_eof(s) && *c != '\n' && *c != '\r' ) + *c = (char) stbi__get8(s); + } +} + +static int stbi__pnm_isdigit(char c) +{ + return c >= '0' && c <= '9'; +} + +static int stbi__pnm_getinteger(stbi__context *s, char *c) +{ + int value = 0; + + while (!stbi__at_eof(s) && stbi__pnm_isdigit(*c)) { + value = value*10 + (*c - '0'); + *c = (char) stbi__get8(s); + if((value > 214748364) || (value == 214748364 && *c > '7')) + return stbi__err("integer parse overflow", "Parsing an integer in the PPM header overflowed a 32-bit int"); + } + + return value; +} + +static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp) +{ + int maxv, dummy; + char c, p, t; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + stbi__rewind(s); + + // Get identifier + p = (char) stbi__get8(s); + t = (char) stbi__get8(s); + if (p != 'P' || (t != '5' && t != '6')) { + stbi__rewind(s); + return 0; + } + + *comp = (t == '6') ? 3 : 1; // '5' is 1-component .pgm; '6' is 3-component .ppm + + c = (char) stbi__get8(s); + stbi__pnm_skip_whitespace(s, &c); + + *x = stbi__pnm_getinteger(s, &c); // read width + if(*x == 0) + return stbi__err("invalid width", "PPM image header had zero or overflowing width"); + stbi__pnm_skip_whitespace(s, &c); + + *y = stbi__pnm_getinteger(s, &c); // read height + if (*y == 0) + return stbi__err("invalid width", "PPM image header had zero or overflowing width"); + stbi__pnm_skip_whitespace(s, &c); + + maxv = stbi__pnm_getinteger(s, &c); // read max value + if (maxv > 65535) + return stbi__err("max value > 65535", "PPM image supports only 8-bit and 16-bit images"); + else if (maxv > 255) + return 16; + else + return 8; +} + +static int stbi__pnm_is16(stbi__context *s) +{ + if (stbi__pnm_info(s, NULL, NULL, NULL) == 16) + return 1; + return 0; +} +#endif + +static int stbi__info_main(stbi__context *s, int *x, int *y, int *comp) +{ + #ifndef STBI_NO_JPEG + if (stbi__jpeg_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PNG + if (stbi__png_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_GIF + if (stbi__gif_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_BMP + if (stbi__bmp_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PSD + if (stbi__psd_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PIC + if (stbi__pic_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PNM + if (stbi__pnm_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_HDR + if (stbi__hdr_info(s, x, y, comp)) return 1; + #endif + + // test tga last because it's a crappy test! + #ifndef STBI_NO_TGA + if (stbi__tga_info(s, x, y, comp)) + return 1; + #endif + return stbi__err("unknown image type", "Image not of any known type, or corrupt"); +} + +static int stbi__is_16_main(stbi__context *s) +{ + #ifndef STBI_NO_PNG + if (stbi__png_is16(s)) return 1; + #endif + + #ifndef STBI_NO_PSD + if (stbi__psd_is16(s)) return 1; + #endif + + #ifndef STBI_NO_PNM + if (stbi__pnm_is16(s)) return 1; + #endif + return 0; +} + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_info(char const *filename, int *x, int *y, int *comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result; + if (!f) return stbi__err("can't fopen", "Unable to open file"); + result = stbi_info_from_file(f, x, y, comp); + fclose(f); + return result; +} + +STBIDEF int stbi_info_from_file(FILE *f, int *x, int *y, int *comp) +{ + int r; + stbi__context s; + long pos = ftell(f); + stbi__start_file(&s, f); + r = stbi__info_main(&s,x,y,comp); + fseek(f,pos,SEEK_SET); + return r; +} + +STBIDEF int stbi_is_16_bit(char const *filename) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result; + if (!f) return stbi__err("can't fopen", "Unable to open file"); + result = stbi_is_16_bit_from_file(f); + fclose(f); + return result; +} + +STBIDEF int stbi_is_16_bit_from_file(FILE *f) +{ + int r; + stbi__context s; + long pos = ftell(f); + stbi__start_file(&s, f); + r = stbi__is_16_main(&s); + fseek(f,pos,SEEK_SET); + return r; +} +#endif // !STBI_NO_STDIO + +STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__info_main(&s,x,y,comp); +} + +STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *c, void *user, int *x, int *y, int *comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); + return stbi__info_main(&s,x,y,comp); +} + +STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__is_16_main(&s); +} + +STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *c, void *user) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); + return stbi__is_16_main(&s); +} + +#endif // STB_IMAGE_IMPLEMENTATION + +/* + revision history: + 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs + 2.19 (2018-02-11) fix warning + 2.18 (2018-01-30) fix warnings + 2.17 (2018-01-29) change sbti__shiftsigned to avoid clang -O2 bug + 1-bit BMP + *_is_16_bit api + avoid warnings + 2.16 (2017-07-23) all functions have 16-bit variants; + STBI_NO_STDIO works again; + compilation fixes; + fix rounding in unpremultiply; + optimize vertical flip; + disable raw_len validation; + documentation fixes + 2.15 (2017-03-18) fix png-1,2,4 bug; now all Imagenet JPGs decode; + warning fixes; disable run-time SSE detection on gcc; + uniform handling of optional "return" values; + thread-safe initialization of zlib tables + 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs + 2.13 (2016-11-29) add 16-bit API, only supported for PNG right now + 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes + 2.11 (2016-04-02) allocate large structures on the stack + remove white matting for transparent PSD + fix reported channel count for PNG & BMP + re-enable SSE2 in non-gcc 64-bit + support RGB-formatted JPEG + read 16-bit PNGs (only as 8-bit) + 2.10 (2016-01-22) avoid warning introduced in 2.09 by STBI_REALLOC_SIZED + 2.09 (2016-01-16) allow comments in PNM files + 16-bit-per-pixel TGA (not bit-per-component) + info() for TGA could break due to .hdr handling + info() for BMP to shares code instead of sloppy parse + can use STBI_REALLOC_SIZED if allocator doesn't support realloc + code cleanup + 2.08 (2015-09-13) fix to 2.07 cleanup, reading RGB PSD as RGBA + 2.07 (2015-09-13) fix compiler warnings + partial animated GIF support + limited 16-bpc PSD support + #ifdef unused functions + bug with < 92 byte PIC,PNM,HDR,TGA + 2.06 (2015-04-19) fix bug where PSD returns wrong '*comp' value + 2.05 (2015-04-19) fix bug in progressive JPEG handling, fix warning + 2.04 (2015-04-15) try to re-enable SIMD on MinGW 64-bit + 2.03 (2015-04-12) extra corruption checking (mmozeiko) + stbi_set_flip_vertically_on_load (nguillemot) + fix NEON support; fix mingw support + 2.02 (2015-01-19) fix incorrect assert, fix warning + 2.01 (2015-01-17) fix various warnings; suppress SIMD on gcc 32-bit without -msse2 + 2.00b (2014-12-25) fix STBI_MALLOC in progressive JPEG + 2.00 (2014-12-25) optimize JPG, including x86 SSE2 & NEON SIMD (ryg) + progressive JPEG (stb) + PGM/PPM support (Ken Miller) + STBI_MALLOC,STBI_REALLOC,STBI_FREE + GIF bugfix -- seemingly never worked + STBI_NO_*, STBI_ONLY_* + 1.48 (2014-12-14) fix incorrectly-named assert() + 1.47 (2014-12-14) 1/2/4-bit PNG support, both direct and paletted (Omar Cornut & stb) + optimize PNG (ryg) + fix bug in interlaced PNG with user-specified channel count (stb) + 1.46 (2014-08-26) + fix broken tRNS chunk (colorkey-style transparency) in non-paletted PNG + 1.45 (2014-08-16) + fix MSVC-ARM internal compiler error by wrapping malloc + 1.44 (2014-08-07) + various warning fixes from Ronny Chevalier + 1.43 (2014-07-15) + fix MSVC-only compiler problem in code changed in 1.42 + 1.42 (2014-07-09) + don't define _CRT_SECURE_NO_WARNINGS (affects user code) + fixes to stbi__cleanup_jpeg path + added STBI_ASSERT to avoid requiring assert.h + 1.41 (2014-06-25) + fix search&replace from 1.36 that messed up comments/error messages + 1.40 (2014-06-22) + fix gcc struct-initialization warning + 1.39 (2014-06-15) + fix to TGA optimization when req_comp != number of components in TGA; + fix to GIF loading because BMP wasn't rewinding (whoops, no GIFs in my test suite) + add support for BMP version 5 (more ignored fields) + 1.38 (2014-06-06) + suppress MSVC warnings on integer casts truncating values + fix accidental rename of 'skip' field of I/O + 1.37 (2014-06-04) + remove duplicate typedef + 1.36 (2014-06-03) + convert to header file single-file library + if de-iphone isn't set, load iphone images color-swapped instead of returning NULL + 1.35 (2014-05-27) + various warnings + fix broken STBI_SIMD path + fix bug where stbi_load_from_file no longer left file pointer in correct place + fix broken non-easy path for 32-bit BMP (possibly never used) + TGA optimization by Arseny Kapoulkine + 1.34 (unknown) + use STBI_NOTUSED in stbi__resample_row_generic(), fix one more leak in tga failure case + 1.33 (2011-07-14) + make stbi_is_hdr work in STBI_NO_HDR (as specified), minor compiler-friendly improvements + 1.32 (2011-07-13) + support for "info" function for all supported filetypes (SpartanJ) + 1.31 (2011-06-20) + a few more leak fixes, bug in PNG handling (SpartanJ) + 1.30 (2011-06-11) + added ability to load files via callbacks to accomidate custom input streams (Ben Wenger) + removed deprecated format-specific test/load functions + removed support for installable file formats (stbi_loader) -- would have been broken for IO callbacks anyway + error cases in bmp and tga give messages and don't leak (Raymond Barbiero, grisha) + fix inefficiency in decoding 32-bit BMP (David Woo) + 1.29 (2010-08-16) + various warning fixes from Aurelien Pocheville + 1.28 (2010-08-01) + fix bug in GIF palette transparency (SpartanJ) + 1.27 (2010-08-01) + cast-to-stbi_uc to fix warnings + 1.26 (2010-07-24) + fix bug in file buffering for PNG reported by SpartanJ + 1.25 (2010-07-17) + refix trans_data warning (Won Chun) + 1.24 (2010-07-12) + perf improvements reading from files on platforms with lock-heavy fgetc() + minor perf improvements for jpeg + deprecated type-specific functions so we'll get feedback if they're needed + attempt to fix trans_data warning (Won Chun) + 1.23 fixed bug in iPhone support + 1.22 (2010-07-10) + removed image *writing* support + stbi_info support from Jetro Lauha + GIF support from Jean-Marc Lienher + iPhone PNG-extensions from James Brown + warning-fixes from Nicolas Schulz and Janez Zemva (i.stbi__err. Janez (U+017D)emva) + 1.21 fix use of 'stbi_uc' in header (reported by jon blow) + 1.20 added support for Softimage PIC, by Tom Seddon + 1.19 bug in interlaced PNG corruption check (found by ryg) + 1.18 (2008-08-02) + fix a threading bug (local mutable static) + 1.17 support interlaced PNG + 1.16 major bugfix - stbi__convert_format converted one too many pixels + 1.15 initialize some fields for thread safety + 1.14 fix threadsafe conversion bug + header-file-only version (#define STBI_HEADER_FILE_ONLY before including) + 1.13 threadsafe + 1.12 const qualifiers in the API + 1.11 Support installable IDCT, colorspace conversion routines + 1.10 Fixes for 64-bit (don't use "unsigned long") + optimized upsampling by Fabian "ryg" Giesen + 1.09 Fix format-conversion for PSD code (bad global variables!) + 1.08 Thatcher Ulrich's PSD code integrated by Nicolas Schulz + 1.07 attempt to fix C++ warning/errors again + 1.06 attempt to fix C++ warning/errors again + 1.05 fix TGA loading to return correct *comp and use good luminance calc + 1.04 default float alpha is 1, not 255; use 'void *' for stbi_image_free + 1.03 bugfixes to STBI_NO_STDIO, STBI_NO_HDR + 1.02 support for (subset of) HDR files, float interface for preferred access to them + 1.01 fix bug: possible bug in handling right-side up bmps... not sure + fix bug: the stbi__bmp_load() and stbi__tga_load() functions didn't work at all + 1.00 interface to zlib that skips zlib header + 0.99 correct handling of alpha in palette + 0.98 TGA loader by lonesock; dynamically add loaders (untested) + 0.97 jpeg errors on too large a file; also catch another malloc failure + 0.96 fix detection of invalid v value - particleman@mollyrocket forum + 0.95 during header scan, seek to markers in case of padding + 0.94 STBI_NO_STDIO to disable stdio usage; rename all #defines the same + 0.93 handle jpegtran output; verbose errors + 0.92 read 4,8,16,24,32-bit BMP files of several formats + 0.91 output 24-bit Windows 3.0 BMP files + 0.90 fix a few more warnings; bump version number to approach 1.0 + 0.61 bugfixes due to Marc LeBlanc, Christopher Lloyd + 0.60 fix compiling as c++ + 0.59 fix warnings: merge Dave Moore's -Wall fixes + 0.58 fix bug: zlib uncompressed mode len/nlen was wrong endian + 0.57 fix bug: jpg last huffman symbol before marker was >9 bits but less than 16 available + 0.56 fix bug: zlib uncompressed mode len vs. nlen + 0.55 fix bug: restart_interval not initialized to 0 + 0.54 allow NULL for 'int *comp' + 0.53 fix bug in png 3->4; speedup png decoding + 0.52 png handles req_comp=3,4 directly; minor cleanup; jpeg comments + 0.51 obey req_comp requests, 1-component jpegs return as 1-component, + on 'test' only check type, not whether we support this variant + 0.50 (2006-11-19) + first released version +*/ + + +/* +------------------------------------------------------------------------------ +This software is available under 2 licenses -- choose whichever you prefer. +------------------------------------------------------------------------------ +ALTERNATIVE A - MIT License +Copyright (c) 2017 Sean Barrett +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +------------------------------------------------------------------------------ +ALTERNATIVE B - Public Domain (www.unlicense.org) +This is free and unencumbered software released into the public domain. +Anyone is free to copy, modify, publish, use, compile, sell, or distribute this +software, either in source code form or as a compiled binary, for any purpose, +commercial or non-commercial, and by any means. +In jurisdictions that recognize copyright laws, the author or authors of this +software dedicate any and all copyright interest in the software to the public +domain. We make this dedication for the benefit of the public at large and to +the detriment of our heirs and successors. We intend this dedication to be an +overt act of relinquishment in perpetuity of all present and future rights to +this software under copyright law. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN +ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +------------------------------------------------------------------------------ +*/ diff --git a/zig/vendor/stb/stb_image_impl.c b/zig/vendor/stb/stb_image_impl.c new file mode 100644 index 0000000000..908982ad8f --- /dev/null +++ b/zig/vendor/stb/stb_image_impl.c @@ -0,0 +1,10 @@ +// stb_image v2.30 (public domain / MIT, github.com/nothings/stb at 2c980bb59875b0d32144a71867fbdebb2f77cd20): +// the native server's JPEG and PNG decoder for image inputs (zig/src/vision/qwen_image.zig). +#define STB_IMAGE_IMPLEMENTATION +#define STBI_ONLY_JPEG +#define STBI_ONLY_PNG +#define STBI_NO_STDIO +#define STBI_NO_HDR +#define STBI_NO_LINEAR +#define STBI_NO_FAILURE_STRINGS +#include "stb_image.h" From 22d92e4ae48f9e6874ab77676cbb1e867b5e0509 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:36:27 +0000 Subject: [PATCH 002/115] lane host: a warm pass that steps aside is not launched after its cancel; lean rule and zig fmt setup() returned cancel()'s true for a background warm pass that yields to a newer request, so the caller launched the job cancel() had just finished (SEGV in lane_host_warm_test). It now returns false, as for a drop. Comments one line each; lean_check skips zig/vendor/ (stb_image.h, as on #726). lane_host.zig and lanes/engine.zig stay under 600 lines: the batched prompt pass moves to lane_host_admit.zig (many, one) and lanes/open.zig (addStreams, adoptPrefilled), aliased where they were. --- tools/zig/lean_check.py | 4 +- zig/build/vision.zig | 3 +- zig/src/core/lane_host.zig | 77 +++++++------------------------- zig/src/core/lane_host_admit.zig | 50 +++++++++++++++++++++ zig/src/core/lanes/backend.zig | 10 ++--- zig/src/core/lanes/engine.zig | 61 ++----------------------- zig/src/core/lanes/open.zig | 55 +++++++++++++++++++++++ zig/src/core/lanes/stream.zig | 5 +-- zig/src/native/cuda.zig | 33 +++++++------- zig/src/server/chat.zig | 7 ++- zig/src/server/main.zig | 2 +- zig/src/server/media_fetch.zig | 20 +++------ zig/src/server/prompt.zig | 4 +- zig/src/server/vision_inputs.zig | 21 +++------ zig/src/vision/glm_image.zig | 8 +--- zig/src/vision/qwen_image.zig | 24 +++------- zig/src/vision/tf_video.c | 23 ++-------- 17 files changed, 179 insertions(+), 228 deletions(-) diff --git a/tools/zig/lean_check.py b/tools/zig/lean_check.py index a97aa37701..a5928f867b 100644 --- a/tools/zig/lean_check.py +++ b/tools/zig/lean_check.py @@ -19,6 +19,8 @@ "zig/src/families/flashnext/roles_gen.zig", "zig/kernels/metal/flashnext/sources_gen.zig", ) +# Third-party sources kept as their authors wrote them (stb_image.h) are skipped too. +VENDORED = "zig/vendor/" def generator_output(path: Path) -> bool: @@ -31,7 +33,7 @@ def generator_output(path: Path) -> bool: def problems(path: Path) -> list[str]: """Every rule break in one file, as 'path:line: reason'.""" - if generator_output(path): + if generator_output(path) or path.relative_to(ROOT).as_posix().startswith(VENDORED): return [] lines = path.read_text(errors="replace").splitlines() rel, found = path.relative_to(ROOT), [] diff --git a/zig/build/vision.zig b/zig/build/vision.zig index 11eb50c143..69950154dc 100644 --- a/zig/build/vision.zig +++ b/zig/build/vision.zig @@ -1,5 +1,4 @@ -//! The image-input module (zig/src/vision/qwen_image.zig with stb_image compiled in), shared by every build that -//! carries the engine API (the server reaches it as engine_api.qwen_image) and by the CUDA test runner. +//! The image-input module (vision/qwen_image.zig with stb_image), shared by the engine API and the CUDA test runner. const std = @import("std"); pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 7a557c20c3..9b07dc4d38 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -285,9 +285,8 @@ pub const LaneHost = struct { h.mutex.unlock(h.io); } - /// Prefills the next queued request into a free lane; false when none waits or no lane is free. - /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit yet). - fn takeJob(h: *LaneHost) ?*Job { + /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit). + pub fn takeJob(h: *LaneHost) ?*Job { h.lock(); if (h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) { h.unlock(); @@ -310,9 +309,12 @@ pub const LaneHost = struct { } /// The admitted request's stream (its kept state, drafter and grammar); false when it was dropped instead. - fn setup(h: *LaneHost, job: *Job) bool { + pub fn setup(h: *LaneHost, job: *Job) bool { const r = job.request; - if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran + if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) { // other work came first: it never ran + _ = h.cancel(job); + return false; + } var reuse: lanes.stream.Reuse = .{}; // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { @@ -349,9 +351,9 @@ pub const LaneHost = struct { .grammar = if (r.structure) |st| h.constraint(st) catch |e| { job.proposer.deinit(); { - _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); - return false; - } + _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return false; + } } else null, .images = r.images, .decode_spans = h.spans(r), @@ -369,8 +371,8 @@ pub const LaneHost = struct { return true; } - /// An admitted, set-up request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. - fn launch(h: *LaneHost, job: *Job) bool { + /// An admitted request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + pub fn launch(h: *LaneHost, job: *Job) bool { const r = job.request; const began = job.began; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); @@ -388,55 +390,6 @@ pub const LaneHost = struct { if (h.deliver(job)) h.remove(job); return true; } - fn admitOne(h: *LaneHost) bool { - const job = h.takeJob() orelse return false; - if (!h.setup(job)) return true; - return h.launch(job); - } - - /// Every queued request a lane is free for, their prompts in one pass where the backend takes several - /// (lanes.Engine.addStreams); the engine's special paths (chunked prompts, warm passes, the lone driver) go - /// through admitOne as before. - fn admitMany(h: *LaneHost) bool { - if (!h.core.batchesPrompts() or h.core.fills()) return h.admitOne(); - h.lock(); - const waiting = h.queued.items.len; - const free = h.info_.lanes -| h.admitted.items.len; - h.unlock(); - if (waiting < 2 or free < 2) return h.admitOne(); - var jobs: [lanes.Engine.max_together]*Job = undefined; - var n: usize = 0; - while (n < jobs.len and n < free) { - const job = h.takeJob() orelse break; - if (!h.setup(job)) continue; - if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { - _ = h.launch(job); - continue; - } - jobs[n] = job; - n += 1; - } - if (n == 0) return true; - if (n == 1) return h.launch(jobs[0]); - var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; - for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; - var errs: [lanes.Engine.max_together]?anyerror = undefined; - h.core.addStreams(streams[0..n], errs[0..n]); - for (jobs[0..n], errs[0..n]) |job, err| { - if (err) |e| { - _ = switch (e) { - error.Cancelled => h.cancel(job), - error.LaneMemory => h.requeue(job), - else => h.drop(job, h.words(e)), - }; - continue; - } - h.prefilled(job, job.began); - if (h.deliver(job)) h.remove(job); - } - return true; - } - /// The next filling job's next prompt chunk: the other streams' rounds run between chunks, not after the prompt. fn fillOne(h: *LaneHost) void { @@ -515,7 +468,7 @@ pub const LaneHost = struct { } /// No memory for another lane yet: the job waits at the queue's head, and admission holds until a job ends. - fn requeue(h: *LaneHost, job: *Job) bool { + pub fn requeue(h: *LaneHost, job: *Job) bool { if (job.stream.grammar) |g| g.deinit(h.gpa); // the next admission compiles its own job.stream.deinit(h.gpa); job.proposer.deinit(); @@ -546,7 +499,7 @@ pub const LaneHost = struct { const alone_words = "the engine has no memory for this request's caches with nothing else running: lower --context or --parallel"; /// A job cancelled in its prompt pass, its lane already released. - fn cancel(h: *LaneHost, job: *Job) bool { + pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); h.finish(job, .cancelled, ""); return true; @@ -565,7 +518,7 @@ pub const LaneHost = struct { fn run(h: *LaneHost) void { while (true) { h.takeCancels(); - while (h.admitMany()) {} + while (admit.many(h)) {} h.fillOne(); status.noteLive(h); h.lock(); diff --git a/zig/src/core/lane_host_admit.zig b/zig/src/core/lane_host_admit.zig index 016ab4398a..f394a3729e 100644 --- a/zig/src/core/lane_host_admit.zig +++ b/zig/src/core/lane_host_admit.zig @@ -2,6 +2,8 @@ const api = @import("engine_api.zig"); const LaneHost = @import("lane_host.zig").LaneHost; const Job = LaneHost.Job; +const lanes = @import("lanes"); +const lone = @import("lane_host_lone.zig"); /// The queued job to admit (its index), or null to hold; the caller holds the lock and a lane is free. pub fn pick(h: *LaneHost) ?usize { @@ -21,3 +23,51 @@ pub fn pick(h: *LaneHost) ?usize { fn fits(a: api.Admit, job: *const Job) bool { return a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens); } + +/// Every queued request a lane is free for, prompts in one pass where the backend takes several; else one at a time. +pub fn many(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return one(h); + h.lock(); + const waiting = h.queued.items.len; + const free = h.info_.lanes -| h.admitted.items.len; + h.unlock(); + if (waiting < 2 or free < 2) return one(h); + var jobs: [lanes.Engine.max_together]*Job = undefined; + var n: usize = 0; + while (n < jobs.len and n < free) { + const job = h.takeJob() orelse break; + if (!h.setup(job)) continue; + if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { + _ = h.launch(job); + continue; + } + jobs[n] = job; + n += 1; + } + if (n == 0) return true; + if (n == 1) return h.launch(jobs[0]); + var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; + for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; + var errs: [lanes.Engine.max_together]?anyerror = undefined; + h.core.addStreams(streams[0..n], errs[0..n]); + for (jobs[0..n], errs[0..n]) |job, err| { + if (err) |e| { + _ = switch (e) { + error.Cancelled => h.cancel(job), + error.LaneMemory => h.requeue(job), + else => h.drop(job, h.words(e)), + }; + continue; + } + h.prefilled(job, job.began); + if (h.deliver(job)) h.remove(job); + } + return true; +} + +/// The next queued request admitted alone: its stream set up and launched. +fn one(h: *LaneHost) bool { + const job = h.takeJob() orelse return false; + if (!h.setup(job)) return true; + return h.launch(job); +} diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 723dc5effc..9d568efb4e 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -1,4 +1,4 @@ -//! What the round loop asks a backend (Metal, CUDA or the fake) to compute; drawn tokens stay on the GPU behind handles. +//! What the round loop asks a backend (Metal, CUDA or the fake); drawn tokens stay on the GPU behind handles. const std = @import("std"); const Stream = @import("stream.zig").Stream; const Shape = @import("shape.zig").Shape; @@ -68,8 +68,7 @@ pub const Backend = struct { prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, /// One chunk of the prompt pass, false while chunks remain; null: the pass runs whole. prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, - /// Prefill several streams' prompts in one pass (each then as after `prefill`); false, having done nothing, - /// when they do not go together (each is then prefilled alone). + /// Several streams' prompts in one pass, each as after `prefill`; false, doing nothing, when they don't fit. prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, @@ -105,8 +104,7 @@ pub const Backend = struct { prompt_rows: bool = false, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, - /// After a failed round: the streams it failed for alone (their caches could not grow), which the caller - /// ends while the round's other streams go on; 0 when the failure is the round's. + /// After a failed round: the streams whose caches could not grow, ended alone; 0 when the round failed. refused: ?*const fn (ptr: *anyopaque, out: []*Stream) usize = null, }; @@ -143,7 +141,7 @@ pub const Backend = struct { b.vtable.release(b.ptr, s); } - /// A driver of its own starts a stream: the prompt pass and its first token, committed; a cancel or failure releases it. + /// A driver of its own starts a stream: prompt pass and first token, committed; cancel or failure releases it. pub fn opening(b: Backend, gpa: std.mem.Allocator, s: *Stream) !u32 { errdefer b.release(s); try b.prefill(s); diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index d5be72c0bb..a3c6230ed9 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -83,69 +83,16 @@ pub const Engine = struct { pub const fills = open.fills; pub const fillStream = open.fillStream; - /// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. - pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { - _ = e.arena.reset(.retain_capacity); - s.context.shrinkRetainingCapacity(s.prompt_len); - s.pending = null; - s.cache_len = s.prompt_len; - if (e.cfg.family_mtp and s.drafts) { - const d: u32 = @intCast(try e.rule.depth(win.who(s))); - try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); - s.dropHeld(e.gpa); - s.next = .{ .count = d }; - if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { - s.next = held; - }; - } - if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row - _ = try s.commit(e.gpa, &.{token}, &.{}); - s.pending = token; - if (s.finished) return e.release(s); - try e.live.append(e.gpa, s); - } - - /// Streams a backend prefills in one pass at most. - pub const max_together = 16; + pub const adoptPrefilled = open.adoptPrefilled; + pub const max_together = open.max_together; + pub const addStreams = open.addStreams; /// Whether the backend prefills several prompts in one pass (addStreams batches them). pub fn batchesPrompts(e: *const Engine) bool { return e.backend.vtable.prefill_many != null; } - /// open.add for several streams: their prompts in one pass when the backend takes them together (else each - /// alone), then each opened as after its own prefill. ``errs[i]``: stream i's error (a cancelled one is released - /// already; the caller discards a failed one). - pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { - for (errs) |*x| x.* = null; - const many = e.backend.vtable.prefill_many orelse return e.addEach(ss, errs); - if (ss.len < 2 or ss.len > max_together) return e.addEach(ss, errs); - _ = e.arena.reset(.retain_capacity); - for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; - const together = many(e.backend.ptr, ss) catch |err| { - for (errs) |*x| x.* = err; - return; - }; - if (!together) return e.addEach(ss, errs); - for (ss, errs) |s, *x| { - if (s.isCancelled()) { // cancelled in its pass: no first token - e.backend.release(s); - x.* = error.Cancelled; - continue; - } - open.opened(e, s) catch |err| { - x.* = err; - }; - } - } - - fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { - for (ss, errs) |s, *x| open.add(e, s) catch |err| { - x.* = err; - }; - } - - /// Take over a stream another driver decoded so far: its cache settled, `pending` and `cache_len` set, no drafts held. + /// Take over a stream another driver decoded: cache settled, `pending` and `cache_len` set, no drafts held. pub fn adopt(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); } diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..c75fa81b86 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -88,3 +88,58 @@ pub fn fillStream(e: *Engine, s: *Stream, first: bool) !bool { try opened(e, s); return true; } + +/// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. +pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { + _ = e.arena.reset(.retain_capacity); + s.context.shrinkRetainingCapacity(s.prompt_len); + s.pending = null; + s.cache_len = s.prompt_len; + if (e.cfg.family_mtp and s.drafts) { + const d: u32 = @intCast(try e.rule.depth(win.who(s))); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { + s.next = held; + }; + } + if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row + _ = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.finished) return e.release(s); + try e.live.append(e.gpa, s); +} + +/// Streams a backend prefills in one pass at most. +pub const max_together = 16; + +/// open.add for several streams, prompts in one pass when the backend allows; ``errs[i]``: stream i's error. +pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (errs) |*x| x.* = null; + const many = e.backend.vtable.prefill_many orelse return addEach(e, ss, errs); + if (ss.len < 2 or ss.len > max_together) return addEach(e, ss, errs); + _ = e.arena.reset(.retain_capacity); + for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; + const together = many(e.backend.ptr, ss) catch |err| { + for (errs) |*x| x.* = err; + return; + }; + if (!together) return addEach(e, ss, errs); + for (ss, errs) |s, *x| { + if (s.isCancelled()) { // cancelled in its pass: no first token + e.backend.release(s); + x.* = error.Cancelled; + continue; + } + opened(e, s) catch |err| { + x.* = err; + }; + } +} + +fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| add(e, s) catch |err| { + x.* = err; + }; +} diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index aade9ad50c..6d634c9600 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -76,7 +76,7 @@ pub const CancelCheck = struct { check: *const fn (ptr: *anyopaque) bool, }; -/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and where it keeps new ones. +/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and keeps new ones in. pub const Reuse = struct { saved: ?*anyopaque = null, // the backend's copy of the state after `at` prompt tokens (null: start at 0) at: u32 = 0, @@ -97,8 +97,7 @@ pub const Held = struct { parents: ?[]i32 = null, }; -/// An image input prepared for a vision tower: its patches (fp32 [grid h * w, 1536], merge-block order) and grid -/// [1, h, w]; the prompt holds its placeholder rows. +/// An image for a vision tower: patches (fp32 [h * w, 1536], merge-block order), grid [1, h, w]; prompt placeholders. pub const Image = struct { patches: []const f32, grid: [3]i64 }; pub const Spec = struct { diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a34186a4a3..7fbf38ba3d 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,4 +1,5 @@ //! CUDA family registration, device setup, memory admission and the native lane host. + /// The serve flags this build's engine takes past the common ones (server/cli.zig served()): two ranks. pub const serves: []const []const u8 = &.{"--tp"}; @@ -252,8 +253,7 @@ const Host = struct { const refusals = @import("cuda_refusal.zig"); -/// A tensor-parallel rank above 0 (server/engines.zig follow): the family opened as that rank, then its replay loop -/// until rank 0 stops; false with ``problem`` set when it cannot. +/// A tensor-parallel rank above 0: the family opened as that rank, then its replay loop; false with ``problem`` set. pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { const opened = (try open(a, gpa, io, o, problem)) orelse return false; defer opened.close(opened.ctx); @@ -421,19 +421,19 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; if (comptime @hasField(@TypeOf(loaded), "learns")) { // a family with a Living Weights learner if (loaded.learns != null) { - const Slide = nemotron_slide.Adapter(F.slide); - h.dir = try gpa.dupe(u8, o.dir); - const s = try gpa.create(Slide); - s.* = Slide.init(gpa, io, F.learnModel(loaded.ctx, h.dir)); - h.slide = s; - h.slide_free = struct { - fn f(al: Allocator, p: *anyopaque) void { - const x: *Slide = @ptrCast(@alignCast(p)); - x.deinit(); - al.destroy(x); - } - }.f; - h.learn = s.hook(); + const Slide = nemotron_slide.Adapter(F.slide); + h.dir = try gpa.dupe(u8, o.dir); + const s = try gpa.create(Slide); + s.* = Slide.init(gpa, io, F.learnModel(loaded.ctx, h.dir)); + h.slide = s; + h.slide_free = struct { + fn f(al: Allocator, p: *anyopaque) void { + const x: *Slide = @ptrCast(@alignCast(p)); + x.deinit(); + al.destroy(x); + } + }.f; + h.learn = s.hook(); } } errdefer if (h.slide) |s| { @@ -441,8 +441,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, - .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; const follow_loop: ?*const fn (*anyopaque) anyerror!void = if (@hasField(@TypeOf(loaded), "follow")) loaded.follow else null; if (follow_loop) |f| { diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 82c72d52ad..63eebea817 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -92,8 +92,7 @@ pub const Prepared = struct { preparing: bool, }; -/// --tool-system: the server's instruction as the first message of a chat request that offers tools and sends no -/// system message (the client's own instruction wins); other requests pass unchanged. +/// --tool-system: the server's instruction first in a tools request with no system message; others pass unchanged. fn withToolSystem(srv: *Server, a: Allocator, input: Input) Allocator.Error!Input { const text = srv.config.tool_system orelse return input; if (input.tools.len == 0 or input.prompt != null or input.messages != .array) return input; @@ -318,7 +317,7 @@ const Generation = struct { return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; } - /// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. + /// What closes a call the end token left open (``tool_parse.closeCall``); nothing for stop, length or ignore_eos. fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { const t = g.collected.items; if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; @@ -398,7 +397,7 @@ const Generation = struct { }; } - /// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. + /// Ends the request and waits for the engine; one ended unfinished counts as a disconnect, as on the Mac. fn cancel(g: *Generation) void { if (!g.engine_done) g.srv.engine.cancel(g.id); g.drain(); diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 0cd1b03c6f..18a6d3a5f7 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -78,7 +78,7 @@ pub fn main(init: std.process.Init) !u8 { } orelse return fail(problem); defer up.text.deinit(); defer up.engine.close(up.engine.ctx); - if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls + if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls follow(up.engine.ctx) catch |e| return fail(@errorName(e)); return 0; } diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig index e175033bf6..c9ef04847f 100644 --- a/zig/src/server/media_fetch.zig +++ b/zig/src/server/media_fetch.zig @@ -1,7 +1,4 @@ -//! Public HTTPS media for --vision-urls, as the Python server fetches it (vision/images_http.py): an https URL on -//! port 443 without credentials or a fragment, its host resolved and every address checked public, a TLS connection -//! verified for the host on an address that was checked, GET with identity encoding, up to 3 redirects (each checked -//! the same way), a declared media type the caller accepts, the bytes within the caller's limit, all inside a deadline. +//! Public HTTPS media for --vision-urls as vision/images_http.py fetches it: checked hosts, TLS, 3 redirects, limits. const std = @import("std"); const Io = std.Io; @@ -22,8 +19,7 @@ fn fail(f: *Failure, text: []const u8) Error { return error.Media; } -/// Python ipaddress: is_global and not multicast or reserved, nor 192.0.0.0/24 or 168.63.129.16; no IPv6 address -/// that embeds an IPv4 one (mapped, 6to4, Teredo). +/// Python ipaddress is_global, less multicast, reserved, 192.0.0.0/24, 168.63.129.16 and IPv6 embedding an IPv4 one. pub fn publicIp(a: net.IpAddress) bool { switch (a) { .ip4 => |v| { @@ -37,9 +33,9 @@ pub fn publicIp(a: net.IpAddress) bool { } }.net4; const blocked = [_]struct { [4]u8, u5 }{ - .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, - .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, - .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, + .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, + .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, + .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, .{ .{ 224, 0, 0, 0 }, 4 }, .{ .{ 240, 0, 0, 0 }, 4 }, }; for (blocked) |blk| if (in(b, blk[0], blk[1])) return false; @@ -48,8 +44,7 @@ pub fn publicIp(a: net.IpAddress) bool { }, .ip6 => |v| { const b = v.bytes; - // global unicast 2000::/3 only, less documentation (2001:db8::/32), the 2001::/23 IETF block (Teredo - // 2001::/32 among it), 6to4 (2002::/16) + // 2000::/3 only, less 2001:db8::/32, the 2001::/23 IETF block (Teredo among it) and 6to4 2002::/16 if (b[0] & 0xe0 != 0x20) return false; if (b[0] == 0x20 and b[1] == 0x01 and b[2] == 0x0d and b[3] == 0xb8) return false; if (b[0] == 0x20 and b[1] == 0x01 and b[2] < 0x02) return false; @@ -127,8 +122,7 @@ fn setTimeouts(handle: net.Socket.Handle, ms: i64) void { std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.SNDTIMEO, std.mem.asBytes(&tv)) catch {}; } -/// A TCP connection to the checked address, its connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds a -/// Linux connect too; std's connect has no timeout yet). +/// A TCP connection to the checked address, connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds connect). fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { const posix = std.posix; const linux = std.os.linux; diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index 6b7faaf8e2..f48b0a76dd 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -53,9 +53,7 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t return ids; } -/// GLM-5.3's template writes a reasoning-effort system line whatever the request asks; with thinking off the Python -/// frontend renders it without (glm5_next/prompts.py thinking_off: no effort line, an empty think block). The line is -/// its own segment between special tokens, so its ids alone are the ids it renders to; the first is dropped. +/// GLM-5.3's effort system line, dropped with thinking off as glm5_next/prompts.py does; its ids are its own segment. fn withoutEffortLine(srv: *Server, cx: *Cx, ids: []u32) errors.Refused![]u32 { const line = srv.text.encode(cx.a, "<|system|>Reasoning Effort: Max", false) catch |e| switch (e) { error.OutOfMemory => return error.OutOfMemory, diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index fe457dfb34..04637f872b 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -1,9 +1,4 @@ -//! Image and video inputs of a chat request, as the Python server's frontend takes them (tensorfold/vision/images.py, -//! videos.py, qwen_processing.py): each ``image_url`` part (a data: URL; JPEG or PNG) decoded and prepared, each -//! ``video_url`` part (a data: URL; MP4, WebM...) decoded by libtfvideo and cut into frame groups, and each replaced in -//! its message by the text Python's frontend expands it to: ``<|vision_start|>`` + one ``<|image_pad|>`` a visual -//! token + ``<|vision_end|>`` for an image, and ``<|vision_start|>`` + ``<|video_pad|>``s + ``<|vision_end|>`` -//! for each frame group of a video. The template copies the text, and one tokenization gives Python's ids. +//! A chat request's image_url and video_url parts, prepared and replaced by the text Python's frontend expands them to. const std = @import("std"); const media_fetch = @import("media_fetch.zig"); @@ -26,8 +21,7 @@ const max_tokens_one: u64 = 4096; pub const Extracted = struct { messages: Value, images: []api.Image }; -/// The image URL of a content part (an OpenAI ``image_url`` part, or ``image`` / ``input_image`` spellings) and its -/// ``detail``; null for a part that is not an image. +/// The image URL of a content part (``image_url``, ``image`` or ``input_image``) and its ``detail``; else null. const Kind = enum { image, video }; fn mediaPart(part: Value) ?struct { kind: Kind, url: ?[]const u8, detail: []const u8 } { @@ -89,8 +83,7 @@ fn hasImages(messages: Value) bool { const image_media = [_][]const u8{ "image/jpeg", "image/png", "image/webp" }; const video_media = [_][]const u8{ "video/mp4", "video/webm", "video/quicktime", "video/x-matroska" }; -/// The bytes of a ``data:;base64,`` URL, or (--vision-urls) of a public HTTPS URL, fetched before -/// ``deadline`` (images.py _check_source and load_images). +/// The bytes of a ``data:`` URL, or (--vision-urls) a public HTTPS URL fetched before ``deadline`` (images.py). fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, deadline: std.Io.Timestamp) errors.Refused![]u8 { if (!std.mem.startsWith(u8, url, "data:")) { if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse(if (kind == .image) "images require data URLs or public HTTPS URLs" else "videos require data URLs or public HTTPS URLs"); @@ -114,8 +107,7 @@ fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, d return out; } -/// A request's images extracted and prepared, its messages with their image parts as the template's image text; -/// null when it has none. A server without --vision refuses images. +/// A request's images prepared, its image parts as the template's image text; null without images. Needs --vision. pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { const list = messages orelse return null; if (!hasImages(list)) return null; @@ -191,8 +183,7 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte try text.appendSlice(cx.a, vision_end); }, .video => if (srv.info.vision_glm) { - // GLM: <|begin_of_video|>, each frame group as <|begin_of_image|> <|image|>s <|end_of_image|> and its - // first frame's "t.t seconds", <|end_of_video|> (mlx-vlm's glm5_next processor) + // GLM: <|begin_of_video|>, image groups with their "t.t seconds", <|end_of_video|> (glm5_next) var v = try videoGlm(srv, cx, bytes); defer v.deinit(cx.a); try text.appendSlice(cx.a, "<|begin_of_video|>"); @@ -244,7 +235,7 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte return .{ .messages = .{ .array = out }, .images = images.items }; } -// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives -------------------------------------------------- +// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives ------------------------------------------ const TfVideo = extern struct { rgb: ?[*]u8, frames: c_int, height: c_int, width: c_int, indices: ?[*]c_int, fps: f64 }; const SizeFn = *const fn (ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int; diff --git a/zig/src/vision/glm_image.zig b/zig/src/vision/glm_image.zig index b46e75f076..889d3e679b 100644 --- a/zig/src/vision/glm_image.zig +++ b/zig/src/vision/glm_image.zig @@ -1,7 +1,4 @@ -//! GLM-5.3-Flash's image and video preparation (mlx-vlm's glm5_next/processing.py, as TensorFold's GLM frontend runs -//! it): the canvas (smart_resize to multiples of 28 within a visual-token budget), the content resized into it (PIL's -//! antialiased bicubic) and zero-padded at the bottom and right, CLIP's normalization (x * f32(1/255) - mean) / std, -//! and the patches in merge-block order ([c, t, 14, 14] each; an image's one frame twice as the temporal pair). +//! GLM-5.3-Flash's image and video preparation as mlx-vlm's glm5_next/processing.py: canvas, resize, norm, patches. const std = @import("std"); const qi = @import("qwen_image.zig"); @@ -128,8 +125,7 @@ pub fn prepare(gpa: std.mem.Allocator, im: qi.Image, min_tokens: u64, max_tokens return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; } -/// A video's frame groups: each group's patches (two consecutive frames as the temporal pair; an odd last frame -/// repeats) and its time (its first frame's seconds, as the processor's timestamps[::2]). +/// A video's frame groups: patches of two consecutive frames (an odd last repeats) and the first frame's seconds. pub const Video = struct { patches: []f32, // groups x gh x gw x 1176 groups: usize, diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig index f918651811..e5331346eb 100644 --- a/zig/src/vision/qwen_image.zig +++ b/zig/src/vision/qwen_image.zig @@ -1,9 +1,4 @@ -//! Image inputs for Qwen-style vision towers, as the Python engine prepares them (tensorfold/vision/images.py and -//! qwen_processing.py with transformers' Qwen2VL processor): decode (JPEG, PNG; EXIF orientation applied, alpha over -//! white), smart_resize to multiples of patch x merge within a pixel budget, an antialiased bicubic resize (PIL's -//! filter and bounds with torch's uint8 fixed point: int16 weights at the most precision the axis allows), -//! (x - 127.5) / 127.5, and the patches in merge-block order, -//! each [channel, frame (the image twice), 16, 16] -- 1,536 values. +//! Image inputs for Qwen-style vision towers as Qwen2VL's processor makes them: decode, smart_resize, resize, patches. const std = @import("std"); /// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. @@ -160,7 +155,6 @@ pub fn smartResize(height: u32, width: u32, factor: u32, min_pixels: u64, max_pi return .{ @intFromFloat(hb), @intFromFloat(wb) }; } - fn bicubic(x0: f64) f64 { const a = -0.5; const x = @abs(x0); @@ -171,8 +165,7 @@ fn bicubic(x0: f64) f64 { const Coeffs = struct { ksize: usize, bounds: []usize, kk: []i32, precision: u5 }; -/// PIL's precompute_coeffs for an axis of ``in_size`` resampled to ``out_size``, as fixed point the way torch's uint8 -/// antialias kernel takes them: int16 weights at the largest precision (< 22 bits) that keeps the largest in range. +/// PIL's precompute_coeffs for an axis, as int16 fixed point at the most precision (< 22 bits) torch's uint8 takes. fn coeffs(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Coeffs { const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); const filterscale = @max(scale, 1.0); @@ -222,8 +215,7 @@ fn clip8(v: i64, precision: u5) u8 { return @intCast(std.math.clamp(r, 0, 255)); } -/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits: the horizontal pass, then the vertical, each -/// rounding to bytes). +/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits): horizontal pass, then vertical, rounding to bytes. pub fn resize(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error!Image { var cur = im; var owned = false; @@ -289,8 +281,7 @@ pub const MERGE = 2; pub const TEMPORAL = 2; pub const FACTOR = PATCH * MERGE; -/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels as the Python frontend passes it: one merged patch), -/// normalized and cut into patches. +/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels: one merged patch), normalized and cut into patches. pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepared { const cap = max_tokens * FACTOR * FACTOR; const size = try smartResize(im.h, im.w, FACTOR, @min(FACTOR * FACTOR, cap), cap); @@ -322,8 +313,7 @@ pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepare // -- videos --------------------------------------------------------------------------------------------------------- -/// A video's frame size for the tower (the Python frontend's video_size): Qwen3-VL's smart_resize with its per-frame -/// cap (at most 768 tokens a frame group, at least ~134), the whole video within ``budget`` tokens. +/// A video's frame size: Qwen3-VL's smart_resize, at most 768 tokens a frame group, the video within ``budget``. pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u32 { const factor: f64 = FACTOR; const shortest: f64 = 128 * FACTOR * FACTOR; @@ -355,9 +345,7 @@ pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u return .{ @intFromFloat(hb), @intFromFloat(wb) }; } -/// A decoded video's frame groups for the tower: each group's patches (fp32 [gh * gw, 1536], two consecutive frames -/// as the temporal pair; an odd last frame repeats), normalized as numpy does ((x * f32(1/255) - 0.5) / 0.5), and -/// each group's time (the mean of its two frames' seconds). +/// A video's frame groups: patches of two consecutive frames, normalized as numpy does, and their mean time. pub const Video = struct { patches: []f32, // groups x gh x gw x 1536 groups: usize, diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c index 1049416565..6436b675e4 100644 --- a/zig/src/vision/tf_video.c +++ b/zig/src/vision/tf_video.c @@ -1,19 +1,4 @@ -// libtfvideo: a video's sampled frames for the native server's vision inputs, as the Python engine decodes them -// (tensorfold/vision/videos.py with PyAV): the first video stream; its rate (average, else guessed, else 24) and -// frame count (the stream's, else its duration times the rate); frames sampled at ``rate`` a second over the whole -// video (Qwen3-VL's rule: at least min_frames, at most max_frames, numpy's linspace rounded half to even), decoded -// in order and scaled straight to the size ``size`` picks as RGB the way PyAV's VideoFrame.to_ndarray(rgb24, width, -// height, interpolation="BICUBIC") does it (VideoReformatter: the frame's colour space and range kept, transfer and -// primaries unspecified, sws_scale_frame with SWS_BICUBIC alone). Built against FFmpeg 9 (the FFmpeg PyAV 19.0.1 -// ships, so the frames equal Python's); the server loads it when a request carries a video. -#include -#include -#include -#include -#include -#include -#include -#include +// libtfvideo: a video's sampled frames as videos.py decodes them with PyAV 19.0.1 (FFmpeg 9), RGB via SWS_BICUBIC. typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); @@ -56,8 +41,7 @@ static int64_t mem_seek(void* o, int64_t off, int whence) { return at; } -// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): an SWS_CS_* value names the frame's colour space -// (the frame's own AVColorSpace number passes through it as PyAV passes it) +// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): the frame's own AVColorSpace passes through. static void set_colorspace(AVFrame* f, int cs) { if (cs == SWS_CS_ITU709) f->colorspace = AVCOL_SPC_BT709; else if (cs == SWS_CS_FCC) f->colorspace = AVCOL_SPC_FCC; @@ -171,8 +155,7 @@ int tf_video_decode(const uint8_t* data, size_t len, double rate, int min_frames if (e < 0) { finished = 1; break; } if (n > last) { finished = 1; av_frame_unref(frame); break; } if (next < count && n == wanted[next]) { - // PyAV's _reformat: the destination copies the frame's properties; its colour space and range pass - // through _set_frame_colorspace with the frame's own values; transfer and primaries unspecified + // PyAV's _reformat: the frame's colour space and range kept; transfer and primaries unspecified av_frame_unref(rgb); av_frame_copy_props(rgb, frame); set_colorspace(frame, (int)frame->colorspace); From ef73ec3d0e8b2eb751bc55c0752b002dfd9357a0 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:36:27 +0000 Subject: [PATCH 003/115] tf_video: the includes back under its one-line header --- zig/src/vision/tf_video.c | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c index 6436b675e4..9cd1290cc7 100644 --- a/zig/src/vision/tf_video.c +++ b/zig/src/vision/tf_video.c @@ -1,4 +1,12 @@ // libtfvideo: a video's sampled frames as videos.py decodes them with PyAV 19.0.1 (FFmpeg 9), RGB via SWS_BICUBIC. +#include +#include +#include +#include +#include +#include +#include +#include typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); From cc5d9d16fe45d3631248fe9461f98cc2e142d968 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:45:12 +0000 Subject: [PATCH 004/115] vision inputs: a server without --vision refuses image parts in the frozen words (the errors/content-image golden) --- zig/src/server/vision_inputs.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index 04637f872b..c63bce5edf 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -111,7 +111,7 @@ fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, d pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { const list = messages orelse return null; if (!hasImages(list)) return null; - if (!srv.info.vision) return cx.refuse("image and video inputs need a server started with --vision"); + if (!srv.info.vision) return cx.refuse("this server accepts text parts only; image, audio and video inputs are unsupported"); // every media part, in prompt order, with where its text goes const Slot = struct { kind: Kind, url: ?[]const u8, detail: []const u8, msg: usize, part: usize }; var slots: std.ArrayList(Slot) = .empty; From aaed024a2edbbfa82e96d91fb185145b7118f236 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:57:12 -0500 Subject: [PATCH 005/115] qwen_image.formatSeconds: t * 10 rounded in integers, not f128, so tensorfold-native takes memcpy from libSystem again (zig build test on macOS) --- zig/src/vision/qwen_image.zig | 53 +++++++++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 6 deletions(-) diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig index e5331346eb..0601d3c858 100644 --- a/zig/src/vision/qwen_image.zig +++ b/zig/src/vision/qwen_image.zig @@ -406,15 +406,56 @@ pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, h: u /// Python's f"{t:.1f}": one decimal, an exact tie rounded to even. pub fn formatSeconds(buf: []u8, t: f64) ![]const u8 { - const x: f128 = @as(f128, t) * 10; // exact: a double's 53 bits times ten fit f128's 113 - const fl = @floor(x); - const frac = x - fl; - var r = if (frac > 0.5) fl + 1 else fl; - if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; - const v: i64 = @intFromFloat(r); + const v = tenths(t); return std.fmt.bufPrint(buf, "{d}.{d}", .{ @divTrunc(v, 10), @as(u64, @intCast(@mod(v, 10))) }); } +/// t times ten, an exact tie to even, in integers: t is m * 2^e and 10m fits 57 bits (f128 would link compiler-rt). +fn tenths(t: f64) i64 { + const bits: u64 = @bitCast(@abs(t)); + const biased: i32 = @intCast(bits >> 52); + const m: u64 = if (biased == 0) bits & ((1 << 52) - 1) else (bits & ((1 << 52) - 1)) | (1 << 52); + const e = @as(i32, @max(biased, 1)) - 1075; + if (e >= 0) return @intFromFloat(t * 10); // at least 2^52 seconds: never a frame's time + const tm = m * 10; + if (e <= -58) return 0; // t * 10 under a half: tm < 2^57 <= 2^(s - 1) + const s: u6 = @intCast(-e); + const q = tm >> s; + const rem = tm & ((@as(u64, 1) << s) - 1); + const half = @as(u64, 1) << (s - 1); + const r: i64 = @intCast(if (rem > half or (rem == half and q & 1 == 1)) q + 1 else q); + return if (t < 0) -r else r; +} + +test "tenths rounds t * 10 exactly, ties to even, as the f128 product does" { + const ref = struct { + fn f(t: f64) i64 { + const x: f128 = @as(f128, t) * 10; + const fl = @floor(x); + const frac = x - fl; + var r = if (frac > 0.5) fl + 1 else fl; + if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; + return @intFromFloat(r); + } + }.f; + const fixed = [_]f64{ 0, -0.0, 0.05, 0.15, 0.25, 0.35, 0.45, 1.25, 2.5, 0.95, 1.05, 3.0, 123.45, 1e-300, 5e-324, 4503599627370495.5, -0.25, -1.35 }; + for (fixed) |t| try std.testing.expectEqual(ref(t), tenths(t)); + var prng: std.Random.DefaultPrng = .init(0x7e17); + const rand = prng.random(); + for (0..200_000) |i| { + const t: f64 = switch (i % 4) { + 0 => rand.float(f64) * 4000, + 1 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 0, 80_000))) / 20, // halves of tenths + 2 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 100_000))) / @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 121))), + else => @bitCast(@as(u64, rand.uintLessThan(u64, 1075)) << 52 | rand.int(u64) >> 12), // any double below 2^52 + }; + try std.testing.expectEqual(ref(t), tenths(t)); + } + var buf: [32]u8 = undefined; + try std.testing.expectEqualStrings("0.2", try formatSeconds(&buf, 0.25)); + try std.testing.expectEqualStrings("1.4", try formatSeconds(&buf, 1.35)); +} + test { _ = glm; } From 9e6dfcd4c0be11305c304db901337fa9b8cc6d9f Mon Sep 17 00:00:00 2001 From: BobClawblaw Date: Thu, 15 Oct 2026 07:30:04 +0000 Subject: [PATCH 006/115] lane host: a set-up job's kept entry stays pinned until its restore reports; images wait for the shared decoder; no refused hook - prompt_cache: Entry.pins, Store.pin/unpin. victimBut passes pinned entries over, so a later job's lookup in admit.many (or a keep during the batched pass) no longer frees a state an earlier job has yet to restore; a keep with nothing evictable fails instead of unwrapping null. LaneHost.setup pins, reported and requeue unpin. - vision: zig/vendor/stb dropped and no longer compiled. qwen_image.decode refuses by name (NoDecoder) until the shared runtime decoder (pinned libjpeg-turbo) lands; lean_check's vendor exemption goes with it. - lanes: Backend.VTable.refused, Engine.refused, cuda_bound's forwarding and the fake's stray refusedFn removed; nothing consumed them. --- tools/zig/lean_check.py | 4 +- zig/build/vision.zig | 7 +- zig/src/core/lane_host.zig | 5 +- zig/src/core/lanes/backend.zig | 2 - zig/src/core/lanes/engine.zig | 6 - zig/src/core/lanes/fake.zig | 8 - zig/src/core/prompt_cache.zig | 19 +- zig/src/core/prompt_cache_test.zig | 30 + zig/src/native/cuda_bound.zig | 6 - zig/src/server/vision_inputs.zig | 1 + zig/src/vision/qwen_image.zig | 108 +- zig/vendor/stb/README.md | 3 - zig/vendor/stb/stb_image.h | 7988 ---------------------------- zig/vendor/stb/stb_image_impl.c | 10 - 14 files changed, 56 insertions(+), 8141 deletions(-) delete mode 100644 zig/vendor/stb/README.md delete mode 100644 zig/vendor/stb/stb_image.h delete mode 100644 zig/vendor/stb/stb_image_impl.c diff --git a/tools/zig/lean_check.py b/tools/zig/lean_check.py index a5928f867b..a97aa37701 100644 --- a/tools/zig/lean_check.py +++ b/tools/zig/lean_check.py @@ -19,8 +19,6 @@ "zig/src/families/flashnext/roles_gen.zig", "zig/kernels/metal/flashnext/sources_gen.zig", ) -# Third-party sources kept as their authors wrote them (stb_image.h) are skipped too. -VENDORED = "zig/vendor/" def generator_output(path: Path) -> bool: @@ -33,7 +31,7 @@ def generator_output(path: Path) -> bool: def problems(path: Path) -> list[str]: """Every rule break in one file, as 'path:line: reason'.""" - if generator_output(path) or path.relative_to(ROOT).as_posix().startswith(VENDORED): + if generator_output(path): return [] lines = path.read_text(errors="replace").splitlines() rel, found = path.relative_to(ROOT), [] diff --git a/zig/build/vision.zig b/zig/build/vision.zig index 69950154dc..0b78064367 100644 --- a/zig/build/vision.zig +++ b/zig/build/vision.zig @@ -1,9 +1,6 @@ -//! The image-input module (vision/qwen_image.zig with stb_image), shared by the engine API and the CUDA test runner. +//! The image-input module (vision/qwen_image.zig), shared by the engine API and the CUDA test runner. const std = @import("std"); pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { - const m = b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); - m.addCSourceFile(.{ .file = b.path("zig/vendor/stb/stb_image_impl.c"), .flags = &.{ "-std=c11", "-O2" } }); - m.addIncludePath(b.path("zig/vendor/stb")); - return m; + return b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); } diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 9b07dc4d38..9ddbf6a8fe 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -85,6 +85,7 @@ pub const LaneHost = struct { const e = job.entry orelse return; job.entry = null; const store = job.host.cache orelse return; + store.unpin(e); if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); @@ -316,9 +317,8 @@ pub const LaneHost = struct { return false; } var reuse: lanes.stream.Reuse = .{}; - // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { - job.entry = l.entry; + job.entry = if (l.entry) |e| store.pin(e) else null; // kept alive until its restore reports job.kept0 = store.counts.kept; job.marks = l.marks; reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; @@ -474,6 +474,7 @@ pub const LaneHost = struct { job.proposer.deinit(); job.started = false; job.fill_began = false; + if (job.entry) |e| h.cache.?.unpin(e); job.entry = null; // looked up again at its next admission h.gpa.free(job.marks); job.marks = &.{}; diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index 9d568efb4e..23a4efd892 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -104,8 +104,6 @@ pub const Backend = struct { prompt_rows: bool = false, /// The stream left the rounds: free its caches and held drafts. release: *const fn (ptr: *anyopaque, s: *Stream) void, - /// After a failed round: the streams whose caches could not grow, ended alone; 0 when the round failed. - refused: ?*const fn (ptr: *anyopaque, out: []*Stream) usize = null, }; pub fn prefill(b: Backend, s: *Stream) !void { diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index a3c6230ed9..e4c7950781 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -109,12 +109,6 @@ pub const Engine = struct { return false; } - /// After ``step`` failed: the streams the backend refused alone (the caller ends them; the rest go on). - pub fn refused(e: *Engine, out: []*Stream) usize { - const ask = e.backend.vtable.refused orelse return 0; - return ask(e.backend.ptr, out); - } - /// Release a cancelled stream between rounds: no pending draws or drafts. pub fn discard(e: *Engine, s: *Stream) void { s.finished = true; diff --git a/zig/src/core/lanes/fake.zig b/zig/src/core/lanes/fake.zig index fc90ee7483..9434736c46 100644 --- a/zig/src/core/lanes/fake.zig +++ b/zig/src/core/lanes/fake.zig @@ -243,14 +243,6 @@ pub const Fake = struct { return self(ptr).drawn.items[handle]; } - fn refusedFn(ptr: *anyopaque, out: []*Stream) usize { - const x = self(ptr); - const n = @min(out.len, x.refused.items.len); - @memcpy(out[0..n], x.refused.items[0..n]); - x.refused.clearRetainingCapacity(); - return n; - } - fn verify(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { const x = self(ptr); x.rounds += 1; diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index f1465325e1..6fe20982f7 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -108,6 +108,7 @@ pub const Entry = struct { last: []u32, // the last prompt that kept or resumed it shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded + pins: u32 = 0, // admitted jobs that restore it and have not yet: eviction passes it over }; pub const Counts = struct { hits: u64 = 0, misses: u64 = 0, kept: u64 = 0, evicted: u64 = 0, refused: u64 = 0, failed: u64 = 0 }; @@ -161,6 +162,16 @@ pub const Store = struct { s.free(e); } + /// `e` waits for a restore: no lookup or keep evicts it until `unpin` (jobs set up after it run their own first). + pub fn pin(_: *Store, e: *Entry) *Entry { + e.pins += 1; + return e; + } + + pub fn unpin(_: *Store, e: *Entry) void { + e.pins -= 1; + } + /// The family's free storage kept for later saves. pub fn spare(s: *const Store) u64 { return if (s.family.vtable.spare) |f| f(s.family.ptr) else 0; @@ -450,7 +461,7 @@ pub const Store = struct { if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held -| bytes); // spare that cannot take this state goes before any state if (s.spare() < free_) continue; } - s.remove(s.victimBut(null, prompt, born).?); // its storage may come back spare and take this state + s.remove(s.victimBut(null, prompt, born) orelse return s.fail(at, error.RestoresPending)); // spare for this state s.counts.evicted += 1; } if (s.entries.items.len == 0) if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| bytes); @@ -534,17 +545,17 @@ pub const Store = struct { return false; } - /// The entry to free first, never `skip`: a superseded one, no branch point, oldest; else the oldest, rewind last. + /// The entry to free first, never `skip` or pinned: superseded, no branch point, oldest; else oldest, rewind last. fn victimBut(s: *const Store, skip: ?*Entry, prompt: []const u32, born: u32) ?usize { var best: ?usize = null; for (s.entries.items, 0..) |e, i| { - if (e.shared or e.branched or e == skip or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used + if (e.shared or e.branched or e == skip or e.pins > 0 or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used const now = born > e.born and prompt.len > e.tokens.len and std.mem.eql(u32, prompt[0..e.tokens.len], e.tokens); // a longer request's tokens extend it if ((now or s.extended(e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; } if (best) |i| return i; var oldest: ?usize = null; - for (s.entries.items, 0..) |e, i| if (e != skip and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { + for (s.entries.items, 0..) |e, i| if (e != skip and e.pins == 0 and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { oldest = i; }; return oldest; diff --git a/zig/src/core/prompt_cache_test.zig b/zig/src/core/prompt_cache_test.zig index 8452036b75..b408b3c7ad 100644 --- a/zig/src/core/prompt_cache_test.zig +++ b/zig/src/core/prompt_cache_test.zig @@ -391,3 +391,33 @@ test "a kept state grows only into what the budget leaves, and evicts nothing to var stranger: u8 = 0; try std.testing.expect(!s.grow(&stranger, 0)); // a state the store doesn't hold } + +test "an entry pinned for a restore still to come outlives the lookups set up after it; unpinned, it goes again" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 220); // fake bytes: 100 + at + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); // 103 bytes, the oldest + const y = [_]u32{ 6, 6, 6, 6 }; + _ = f.pass(&s, &y, try s.begin(a, &y, 3, &.{}, &.{}, null, &.{})); // 103: 206 held + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const l = try s.lookup(a, &xa, 3, &.{}, &.{}, &.{}); // a job set up to resume x + const e = l.entry.?; + _ = s.pin(e); + const z = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7 }; + _ = try s.lookup(a, &z, 7, &.{}, &.{}, &.{}); // the next job's room: y goes, not x + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + try std.testing.expect(s.find(&.{ 6, 6, 6, 6 }, &.{}, &.{}) == null); + const w: [130]u32 = @splat(8); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); // 220 bytes, and nothing but x to evict: it stays + try std.testing.expect(!s.keep(&w, 120, null, &.{}, &.{})); // and a keep with no room is refused + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + s.unpin(e); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(s.entries.items.len, f.live); +} diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 2144981686..fe929ae849 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -114,12 +114,6 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.held.?(x.inner.ptr, s, out); } }.f else null, - .refused = if (v.refused != null) struct { - fn f(p: *anyopaque, out: []*lanes.Stream) usize { - const x = H.bind(p); - return x.inner.vtable.refused.?(x.inner.ptr, out); - } - }.f else null, .prompt_rows = v.prompt_rows, .release = struct { fn f(p: *anyopaque, s: *lanes.Stream) void { diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig index c63bce5edf..b1e4476fb1 100644 --- a/zig/src/server/vision_inputs.zig +++ b/zig/src/server/vision_inputs.zig @@ -154,6 +154,7 @@ pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracte var img = qi.decode(cx.a, bytes, .{}) catch |e| switch (e) { error.OutOfMemory => return error.OutOfMemory, error.ImageTooLarge => return cx.refuse("image dimensions exceed the decoded pixel limit"), + error.NoDecoder => return cx.refuse("image inputs need the shared JPEG and PNG decoder, which this server does not have yet"), else => return cx.refuse("image bytes are invalid or unsupported; use JPEG or PNG"), }; total_pixels += @as(u64, img.w) * img.h; diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig index 0601d3c858..79e7c4be3a 100644 --- a/zig/src/vision/qwen_image.zig +++ b/zig/src/vision/qwen_image.zig @@ -4,10 +4,6 @@ const std = @import("std"); /// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. pub const glm = @import("glm_image.zig"); -extern fn stbi_info_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int) c_int; -extern fn stbi_load_from_memory(buffer: [*]const u8, len: c_int, x: *c_int, y: *c_int, comp: *c_int, req: c_int) ?[*]u8; -extern fn stbi_image_free(p: ?*anyopaque) void; - pub const Limits = struct { max_dimension: u32 = 8192, max_pixels: u64 = 16 * 1024 * 1024, @@ -24,108 +20,12 @@ pub const Image = struct { } }; -pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio }; +pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio, NoDecoder }; -/// JPEG or PNG bytes to RGB: EXIF orientation applied (JPEG), any alpha composited over white. +/// JPEG or PNG bytes to RGB: the shared runtime decoder's job (pinned libjpeg-turbo); until it lands, refused by name. pub fn decode(gpa: std.mem.Allocator, bytes: []const u8, limits: Limits) Error!Image { - if (bytes.len == 0 or bytes.len > std.math.maxInt(c_int)) return error.UnsupportedImage; - var w: c_int = 0; - var h: c_int = 0; - var comp: c_int = 0; - if (stbi_info_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp) == 0) return error.UnsupportedImage; - if (w <= 0 or h <= 0 or w > limits.max_dimension or h > limits.max_dimension or - @as(u64, @intCast(w)) * @as(u64, @intCast(h)) > limits.max_pixels) return error.ImageTooLarge; - const alpha = comp == 2 or comp == 4; - const want: c_int = if (alpha) 4 else 3; - var got: c_int = 0; - const px = stbi_load_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp, want) orelse return error.UnsupportedImage; - defer stbi_image_free(px); - const n: usize = @intCast(w * h); - _ = &got; - var rgb = try gpa.alloc(u8, n * 3); - errdefer gpa.free(rgb); - if (alpha) { - // PIL's paste of RGBA onto white with the alpha as mask: out = white + (src - white) * a / 255, rounded - for (0..n) |i| { - const a: i32 = px[i * 4 + 3]; - for (0..3) |c| { - const s: i32 = px[i * 4 + c]; - const t: i32 = (s - 255) * a + 128; - rgb[i * 3 + c] = @intCast(255 + ((t + (t >> 8)) >> 8)); - } - } - } else @memcpy(rgb, px[0 .. n * 3]); - var im: Image = .{ .w = @intCast(w), .h = @intCast(h), .rgb = rgb }; - const o = orientation(bytes); - if (o > 1) { - const t = try transpose(gpa, im, o); - im.deinit(gpa); - im = t; - } - return im; -} - -/// A JPEG's EXIF orientation (1-8; 1 when absent or not a JPEG). -fn orientation(b: []const u8) u16 { - if (b.len < 4 or b[0] != 0xFF or b[1] != 0xD8) return 1; - var i: usize = 2; - while (i + 4 <= b.len and b[i] == 0xFF) { - const marker = b[i + 1]; - const len = (@as(usize, b[i + 2]) << 8) | b[i + 3]; - if (marker == 0xDA or len < 2 or i + 2 + len > b.len) break; - if (marker == 0xE1 and len >= 16 and std.mem.eql(u8, b[i + 4 .. i + 10], "Exif\x00\x00")) { - const t = b[i + 10 .. i + 2 + len]; - const le = t[0] == 'I'; - const rd16 = struct { - fn f(x: []const u8, at: usize, little: bool) u16 { - return if (little) std.mem.readInt(u16, x[at..][0..2], .little) else std.mem.readInt(u16, x[at..][0..2], .big); - } - }.f; - const rd32 = struct { - fn f(x: []const u8, at: usize, little: bool) u32 { - return if (little) std.mem.readInt(u32, x[at..][0..4], .little) else std.mem.readInt(u32, x[at..][0..4], .big); - } - }.f; - if (t.len < 8) return 1; - const ifd = rd32(t, 4, le); - if (ifd + 2 > t.len) return 1; - const count = rd16(t, ifd, le); - for (0..count) |e| { - const at = ifd + 2 + e * 12; - if (at + 12 > t.len) return 1; - if (rd16(t, at, le) == 0x0112) { - const v = rd16(t, at + 8, le); - return if (v >= 1 and v <= 8) v else 1; - } - } - return 1; - } - i += 2 + len; - } - return 1; -} - -/// PIL's exif_transpose for orientations 2-8. -fn transpose(gpa: std.mem.Allocator, im: Image, o: u16) Error!Image { - const swap = o >= 5; - const W: u32 = if (swap) im.h else im.w; - const H: u32 = if (swap) im.w else im.h; - const out = try gpa.alloc(u8, im.rgb.len); - for (0..H) |y| for (0..W) |x| { - // the source pixel of output (x, y) - const sx: u32, const sy: u32 = switch (o) { - 2 => .{ im.w - 1 - @as(u32, @intCast(x)), @intCast(y) }, - 3 => .{ im.w - 1 - @as(u32, @intCast(x)), im.h - 1 - @as(u32, @intCast(y)) }, - 4 => .{ @intCast(x), im.h - 1 - @as(u32, @intCast(y)) }, - 5 => .{ @intCast(y), @intCast(x) }, - 6 => .{ @intCast(y), im.h - 1 - @as(u32, @intCast(x)) }, - 7 => .{ im.w - 1 - @as(u32, @intCast(y)), im.h - 1 - @as(u32, @intCast(x)) }, - 8 => .{ im.w - 1 - @as(u32, @intCast(y)), @intCast(x) }, - else => .{ @intCast(x), @intCast(y) }, - }; - @memcpy(out[(y * W + x) * 3 ..][0..3], im.rgb[(sy * im.w + sx) * 3 ..][0..3]); - }; - return .{ .w = W, .h = H, .rgb = out }; + _ = .{ gpa, bytes, limits }; + return error.NoDecoder; } /// Python's round(): halves to even. diff --git a/zig/vendor/stb/README.md b/zig/vendor/stb/README.md deleted file mode 100644 index 11466f4316..0000000000 --- a/zig/vendor/stb/README.md +++ /dev/null @@ -1,3 +0,0 @@ -stb_image.h v2.30 from https://github.com/nothings/stb at commit 2c980bb59875b0d32144a71867fbdebb2f77cd20 -(sha256 594c2fe35d49488b4382dbfaec8f98366defca819d916ac95becf3e75f4200b3), unmodified; dual public domain / MIT -(see the license at the end of the file). stb_image_impl.c compiles it for JPEG and PNG only. diff --git a/zig/vendor/stb/stb_image.h b/zig/vendor/stb/stb_image.h deleted file mode 100644 index 9eedabedc4..0000000000 --- a/zig/vendor/stb/stb_image.h +++ /dev/null @@ -1,7988 +0,0 @@ -/* stb_image - v2.30 - public domain image loader - http://nothings.org/stb - no warranty implied; use at your own risk - - Do this: - #define STB_IMAGE_IMPLEMENTATION - before you include this file in *one* C or C++ file to create the implementation. - - // i.e. it should look like this: - #include ... - #include ... - #include ... - #define STB_IMAGE_IMPLEMENTATION - #include "stb_image.h" - - You can #define STBI_ASSERT(x) before the #include to avoid using assert.h. - And #define STBI_MALLOC, STBI_REALLOC, and STBI_FREE to avoid using malloc,realloc,free - - - QUICK NOTES: - Primarily of interest to game developers and other people who can - avoid problematic images and only need the trivial interface - - JPEG baseline & progressive (12 bpc/arithmetic not supported, same as stock IJG lib) - PNG 1/2/4/8/16-bit-per-channel - - TGA (not sure what subset, if a subset) - BMP non-1bpp, non-RLE - PSD (composited view only, no extra channels, 8/16 bit-per-channel) - - GIF (*comp always reports as 4-channel) - HDR (radiance rgbE format) - PIC (Softimage PIC) - PNM (PPM and PGM binary only) - - Animated GIF still needs a proper API, but here's one way to do it: - http://gist.github.com/urraka/685d9a6340b26b830d49 - - - decode from memory or through FILE (define STBI_NO_STDIO to remove code) - - decode from arbitrary I/O callbacks - - SIMD acceleration on x86/x64 (SSE2) and ARM (NEON) - - Full documentation under "DOCUMENTATION" below. - - -LICENSE - - See end of file for license information. - -RECENT REVISION HISTORY: - - 2.30 (2024-05-31) avoid erroneous gcc warning - 2.29 (2023-05-xx) optimizations - 2.28 (2023-01-29) many error fixes, security errors, just tons of stuff - 2.27 (2021-07-11) document stbi_info better, 16-bit PNM support, bug fixes - 2.26 (2020-07-13) many minor fixes - 2.25 (2020-02-02) fix warnings - 2.24 (2020-02-02) fix warnings; thread-local failure_reason and flip_vertically - 2.23 (2019-08-11) fix clang static analysis warning - 2.22 (2019-03-04) gif fixes, fix warnings - 2.21 (2019-02-25) fix typo in comment - 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs - 2.19 (2018-02-11) fix warning - 2.18 (2018-01-30) fix warnings - 2.17 (2018-01-29) bugfix, 1-bit BMP, 16-bitness query, fix warnings - 2.16 (2017-07-23) all functions have 16-bit variants; optimizations; bugfixes - 2.15 (2017-03-18) fix png-1,2,4; all Imagenet JPGs; no runtime SSE detection on GCC - 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs - 2.13 (2016-12-04) experimental 16-bit API, only for PNG so far; fixes - 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes - 2.11 (2016-04-02) 16-bit PNGS; enable SSE2 in non-gcc x64 - RGB-format JPEG; remove white matting in PSD; - allocate large structures on the stack; - correct channel count for PNG & BMP - 2.10 (2016-01-22) avoid warning introduced in 2.09 - 2.09 (2016-01-16) 16-bit TGA; comments in PNM files; STBI_REALLOC_SIZED - - See end of file for full revision history. - - - ============================ Contributors ========================= - - Image formats Extensions, features - Sean Barrett (jpeg, png, bmp) Jetro Lauha (stbi_info) - Nicolas Schulz (hdr, psd) Martin "SpartanJ" Golini (stbi_info) - Jonathan Dummer (tga) James "moose2000" Brown (iPhone PNG) - Jean-Marc Lienher (gif) Ben "Disch" Wenger (io callbacks) - Tom Seddon (pic) Omar Cornut (1/2/4-bit PNG) - Thatcher Ulrich (psd) Nicolas Guillemot (vertical flip) - Ken Miller (pgm, ppm) Richard Mitton (16-bit PSD) - github:urraka (animated gif) Junggon Kim (PNM comments) - Christopher Forseth (animated gif) Daniel Gibson (16-bit TGA) - socks-the-fox (16-bit PNG) - Jeremy Sawicki (handle all ImageNet JPGs) - Optimizations & bugfixes Mikhail Morozov (1-bit BMP) - Fabian "ryg" Giesen Anael Seghezzi (is-16-bit query) - Arseny Kapoulkine Simon Breuss (16-bit PNM) - John-Mark Allen - Carmelo J Fdez-Aguera - - Bug & warning fixes - Marc LeBlanc David Woo Guillaume George Martins Mozeiko - Christpher Lloyd Jerry Jansson Joseph Thomson Blazej Dariusz Roszkowski - Phil Jordan Dave Moore Roy Eltham - Hayaki Saito Nathan Reed Won Chun - Luke Graham Johan Duparc Nick Verigakis the Horde3D community - Thomas Ruf Ronny Chevalier github:rlyeh - Janez Zemva John Bartholomew Michal Cichon github:romigrou - Jonathan Blow Ken Hamada Tero Hanninen github:svdijk - Eugene Golushkov Laurent Gomila Cort Stratton github:snagar - Aruelien Pocheville Sergio Gonzalez Thibault Reuille github:Zelex - Cass Everitt Ryamond Barbiero github:grim210 - Paul Du Bois Engin Manap Aldo Culquicondor github:sammyhw - Philipp Wiesemann Dale Weiler Oriol Ferrer Mesia github:phprus - Josh Tobin Neil Bickford Matthew Gregan github:poppolopoppo - Julian Raschke Gregory Mullen Christian Floisand github:darealshinji - Baldur Karlsson Kevin Schmidt JR Smith github:Michaelangel007 - Brad Weinberger Matvey Cherevko github:mosra - Luca Sas Alexander Veselov Zack Middleton [reserved] - Ryan C. Gordon [reserved] [reserved] - DO NOT ADD YOUR NAME HERE - - Jacko Dirks - - To add your name to the credits, pick a random blank space in the middle and fill it. - 80% of merge conflicts on stb PRs are due to people adding their name at the end - of the credits. -*/ - -#ifndef STBI_INCLUDE_STB_IMAGE_H -#define STBI_INCLUDE_STB_IMAGE_H - -// DOCUMENTATION -// -// Limitations: -// - no 12-bit-per-channel JPEG -// - no JPEGs with arithmetic coding -// - GIF always returns *comp=4 -// -// Basic usage (see HDR discussion below for HDR usage): -// int x,y,n; -// unsigned char *data = stbi_load(filename, &x, &y, &n, 0); -// // ... process data if not NULL ... -// // ... x = width, y = height, n = # 8-bit components per pixel ... -// // ... replace '0' with '1'..'4' to force that many components per pixel -// // ... but 'n' will always be the number that it would have been if you said 0 -// stbi_image_free(data); -// -// Standard parameters: -// int *x -- outputs image width in pixels -// int *y -- outputs image height in pixels -// int *channels_in_file -- outputs # of image components in image file -// int desired_channels -- if non-zero, # of image components requested in result -// -// The return value from an image loader is an 'unsigned char *' which points -// to the pixel data, or NULL on an allocation failure or if the image is -// corrupt or invalid. The pixel data consists of *y scanlines of *x pixels, -// with each pixel consisting of N interleaved 8-bit components; the first -// pixel pointed to is top-left-most in the image. There is no padding between -// image scanlines or between pixels, regardless of format. The number of -// components N is 'desired_channels' if desired_channels is non-zero, or -// *channels_in_file otherwise. If desired_channels is non-zero, -// *channels_in_file has the number of components that _would_ have been -// output otherwise. E.g. if you set desired_channels to 4, you will always -// get RGBA output, but you can check *channels_in_file to see if it's trivially -// opaque because e.g. there were only 3 channels in the source image. -// -// An output image with N components has the following components interleaved -// in this order in each pixel: -// -// N=#comp components -// 1 grey -// 2 grey, alpha -// 3 red, green, blue -// 4 red, green, blue, alpha -// -// If image loading fails for any reason, the return value will be NULL, -// and *x, *y, *channels_in_file will be unchanged. The function -// stbi_failure_reason() can be queried for an extremely brief, end-user -// unfriendly explanation of why the load failed. Define STBI_NO_FAILURE_STRINGS -// to avoid compiling these strings at all, and STBI_FAILURE_USERMSG to get slightly -// more user-friendly ones. -// -// Paletted PNG, BMP, GIF, and PIC images are automatically depalettized. -// -// To query the width, height and component count of an image without having to -// decode the full file, you can use the stbi_info family of functions: -// -// int x,y,n,ok; -// ok = stbi_info(filename, &x, &y, &n); -// // returns ok=1 and sets x, y, n if image is a supported format, -// // 0 otherwise. -// -// Note that stb_image pervasively uses ints in its public API for sizes, -// including sizes of memory buffers. This is now part of the API and thus -// hard to change without causing breakage. As a result, the various image -// loaders all have certain limits on image size; these differ somewhat -// by format but generally boil down to either just under 2GB or just under -// 1GB. When the decoded image would be larger than this, stb_image decoding -// will fail. -// -// Additionally, stb_image will reject image files that have any of their -// dimensions set to a larger value than the configurable STBI_MAX_DIMENSIONS, -// which defaults to 2**24 = 16777216 pixels. Due to the above memory limit, -// the only way to have an image with such dimensions load correctly -// is for it to have a rather extreme aspect ratio. Either way, the -// assumption here is that such larger images are likely to be malformed -// or malicious. If you do need to load an image with individual dimensions -// larger than that, and it still fits in the overall size limit, you can -// #define STBI_MAX_DIMENSIONS on your own to be something larger. -// -// =========================================================================== -// -// UNICODE: -// -// If compiling for Windows and you wish to use Unicode filenames, compile -// with -// #define STBI_WINDOWS_UTF8 -// and pass utf8-encoded filenames. Call stbi_convert_wchar_to_utf8 to convert -// Windows wchar_t filenames to utf8. -// -// =========================================================================== -// -// Philosophy -// -// stb libraries are designed with the following priorities: -// -// 1. easy to use -// 2. easy to maintain -// 3. good performance -// -// Sometimes I let "good performance" creep up in priority over "easy to maintain", -// and for best performance I may provide less-easy-to-use APIs that give higher -// performance, in addition to the easy-to-use ones. Nevertheless, it's important -// to keep in mind that from the standpoint of you, a client of this library, -// all you care about is #1 and #3, and stb libraries DO NOT emphasize #3 above all. -// -// Some secondary priorities arise directly from the first two, some of which -// provide more explicit reasons why performance can't be emphasized. -// -// - Portable ("ease of use") -// - Small source code footprint ("easy to maintain") -// - No dependencies ("ease of use") -// -// =========================================================================== -// -// I/O callbacks -// -// I/O callbacks allow you to read from arbitrary sources, like packaged -// files or some other source. Data read from callbacks are processed -// through a small internal buffer (currently 128 bytes) to try to reduce -// overhead. -// -// The three functions you must define are "read" (reads some bytes of data), -// "skip" (skips some bytes of data), "eof" (reports if the stream is at the end). -// -// =========================================================================== -// -// SIMD support -// -// The JPEG decoder will try to automatically use SIMD kernels on x86 when -// supported by the compiler. For ARM Neon support, you must explicitly -// request it. -// -// (The old do-it-yourself SIMD API is no longer supported in the current -// code.) -// -// On x86, SSE2 will automatically be used when available based on a run-time -// test; if not, the generic C versions are used as a fall-back. On ARM targets, -// the typical path is to have separate builds for NEON and non-NEON devices -// (at least this is true for iOS and Android). Therefore, the NEON support is -// toggled by a build flag: define STBI_NEON to get NEON loops. -// -// If for some reason you do not want to use any of SIMD code, or if -// you have issues compiling it, you can disable it entirely by -// defining STBI_NO_SIMD. -// -// =========================================================================== -// -// HDR image support (disable by defining STBI_NO_HDR) -// -// stb_image supports loading HDR images in general, and currently the Radiance -// .HDR file format specifically. You can still load any file through the existing -// interface; if you attempt to load an HDR file, it will be automatically remapped -// to LDR, assuming gamma 2.2 and an arbitrary scale factor defaulting to 1; -// both of these constants can be reconfigured through this interface: -// -// stbi_hdr_to_ldr_gamma(2.2f); -// stbi_hdr_to_ldr_scale(1.0f); -// -// (note, do not use _inverse_ constants; stbi_image will invert them -// appropriately). -// -// Additionally, there is a new, parallel interface for loading files as -// (linear) floats to preserve the full dynamic range: -// -// float *data = stbi_loadf(filename, &x, &y, &n, 0); -// -// If you load LDR images through this interface, those images will -// be promoted to floating point values, run through the inverse of -// constants corresponding to the above: -// -// stbi_ldr_to_hdr_scale(1.0f); -// stbi_ldr_to_hdr_gamma(2.2f); -// -// Finally, given a filename (or an open file or memory block--see header -// file for details) containing image data, you can query for the "most -// appropriate" interface to use (that is, whether the image is HDR or -// not), using: -// -// stbi_is_hdr(char *filename); -// -// =========================================================================== -// -// iPhone PNG support: -// -// We optionally support converting iPhone-formatted PNGs (which store -// premultiplied BGRA) back to RGB, even though they're internally encoded -// differently. To enable this conversion, call -// stbi_convert_iphone_png_to_rgb(1). -// -// Call stbi_set_unpremultiply_on_load(1) as well to force a divide per -// pixel to remove any premultiplied alpha *only* if the image file explicitly -// says there's premultiplied data (currently only happens in iPhone images, -// and only if iPhone convert-to-rgb processing is on). -// -// =========================================================================== -// -// ADDITIONAL CONFIGURATION -// -// - You can suppress implementation of any of the decoders to reduce -// your code footprint by #defining one or more of the following -// symbols before creating the implementation. -// -// STBI_NO_JPEG -// STBI_NO_PNG -// STBI_NO_BMP -// STBI_NO_PSD -// STBI_NO_TGA -// STBI_NO_GIF -// STBI_NO_HDR -// STBI_NO_PIC -// STBI_NO_PNM (.ppm and .pgm) -// -// - You can request *only* certain decoders and suppress all other ones -// (this will be more forward-compatible, as addition of new decoders -// doesn't require you to disable them explicitly): -// -// STBI_ONLY_JPEG -// STBI_ONLY_PNG -// STBI_ONLY_BMP -// STBI_ONLY_PSD -// STBI_ONLY_TGA -// STBI_ONLY_GIF -// STBI_ONLY_HDR -// STBI_ONLY_PIC -// STBI_ONLY_PNM (.ppm and .pgm) -// -// - If you use STBI_NO_PNG (or _ONLY_ without PNG), and you still -// want the zlib decoder to be available, #define STBI_SUPPORT_ZLIB -// -// - If you define STBI_MAX_DIMENSIONS, stb_image will reject images greater -// than that size (in either width or height) without further processing. -// This is to let programs in the wild set an upper bound to prevent -// denial-of-service attacks on untrusted data, as one could generate a -// valid image of gigantic dimensions and force stb_image to allocate a -// huge block of memory and spend disproportionate time decoding it. By -// default this is set to (1 << 24), which is 16777216, but that's still -// very big. - -#ifndef STBI_NO_STDIO -#include -#endif // STBI_NO_STDIO - -#define STBI_VERSION 1 - -enum -{ - STBI_default = 0, // only used for desired_channels - - STBI_grey = 1, - STBI_grey_alpha = 2, - STBI_rgb = 3, - STBI_rgb_alpha = 4 -}; - -#include -typedef unsigned char stbi_uc; -typedef unsigned short stbi_us; - -#ifdef __cplusplus -extern "C" { -#endif - -#ifndef STBIDEF -#ifdef STB_IMAGE_STATIC -#define STBIDEF static -#else -#define STBIDEF extern -#endif -#endif - -////////////////////////////////////////////////////////////////////////////// -// -// PRIMARY API - works on images of any type -// - -// -// load image by filename, open file, or memory buffer -// - -typedef struct -{ - int (*read) (void *user,char *data,int size); // fill 'data' with 'size' bytes. return number of bytes actually read - void (*skip) (void *user,int n); // skip the next 'n' bytes, or 'unget' the last -n bytes if negative - int (*eof) (void *user); // returns nonzero if we are at end of file/data -} stbi_io_callbacks; - -//////////////////////////////////// -// -// 8-bits-per-channel interface -// - -STBIDEF stbi_uc *stbi_load_from_memory (stbi_uc const *buffer, int len , int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk , void *user, int *x, int *y, int *channels_in_file, int desired_channels); - -#ifndef STBI_NO_STDIO -STBIDEF stbi_uc *stbi_load (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_uc *stbi_load_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); -// for stbi_load_from_file, file pointer is left pointing immediately after image -#endif - -#ifndef STBI_NO_GIF -STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp); -#endif - -#ifdef STBI_WINDOWS_UTF8 -STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input); -#endif - -//////////////////////////////////// -// -// 16-bits-per-channel interface -// - -STBIDEF stbi_us *stbi_load_16_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); - -#ifndef STBI_NO_STDIO -STBIDEF stbi_us *stbi_load_16 (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_us *stbi_load_from_file_16(FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); -#endif - -//////////////////////////////////// -// -// float-per-channel interface -// -#ifndef STBI_NO_LINEAR - STBIDEF float *stbi_loadf_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); - STBIDEF float *stbi_loadf_from_callbacks (stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); - - #ifndef STBI_NO_STDIO - STBIDEF float *stbi_loadf (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); - STBIDEF float *stbi_loadf_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); - #endif -#endif - -#ifndef STBI_NO_HDR - STBIDEF void stbi_hdr_to_ldr_gamma(float gamma); - STBIDEF void stbi_hdr_to_ldr_scale(float scale); -#endif // STBI_NO_HDR - -#ifndef STBI_NO_LINEAR - STBIDEF void stbi_ldr_to_hdr_gamma(float gamma); - STBIDEF void stbi_ldr_to_hdr_scale(float scale); -#endif // STBI_NO_LINEAR - -// stbi_is_hdr is always defined, but always returns false if STBI_NO_HDR -STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user); -STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len); -#ifndef STBI_NO_STDIO -STBIDEF int stbi_is_hdr (char const *filename); -STBIDEF int stbi_is_hdr_from_file(FILE *f); -#endif // STBI_NO_STDIO - - -// get a VERY brief reason for failure -// on most compilers (and ALL modern mainstream compilers) this is threadsafe -STBIDEF const char *stbi_failure_reason (void); - -// free the loaded image -- this is just free() -STBIDEF void stbi_image_free (void *retval_from_stbi_load); - -// get image dimensions & components without fully decoding -STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp); -STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp); -STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len); -STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *clbk, void *user); - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_info (char const *filename, int *x, int *y, int *comp); -STBIDEF int stbi_info_from_file (FILE *f, int *x, int *y, int *comp); -STBIDEF int stbi_is_16_bit (char const *filename); -STBIDEF int stbi_is_16_bit_from_file(FILE *f); -#endif - - - -// for image formats that explicitly notate that they have premultiplied alpha, -// we just return the colors as stored in the file. set this flag to force -// unpremultiplication. results are undefined if the unpremultiply overflow. -STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply); - -// indicate whether we should process iphone images back to canonical format, -// or just pass them through "as-is" -STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert); - -// flip the image vertically, so the first pixel in the output array is the bottom left -STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip); - -// as above, but only applies to images loaded on the thread that calls the function -// this function is only available if your compiler supports thread-local variables; -// calling it will fail to link if your compiler doesn't -STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply); -STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert); -STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip); - -// ZLIB client - used by PNG, available for other purposes - -STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen); -STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header); -STBIDEF char *stbi_zlib_decode_malloc(const char *buffer, int len, int *outlen); -STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); - -STBIDEF char *stbi_zlib_decode_noheader_malloc(const char *buffer, int len, int *outlen); -STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); - - -#ifdef __cplusplus -} -#endif - -// -// -//// end header file ///////////////////////////////////////////////////// -#endif // STBI_INCLUDE_STB_IMAGE_H - -#ifdef STB_IMAGE_IMPLEMENTATION - -#if defined(STBI_ONLY_JPEG) || defined(STBI_ONLY_PNG) || defined(STBI_ONLY_BMP) \ - || defined(STBI_ONLY_TGA) || defined(STBI_ONLY_GIF) || defined(STBI_ONLY_PSD) \ - || defined(STBI_ONLY_HDR) || defined(STBI_ONLY_PIC) || defined(STBI_ONLY_PNM) \ - || defined(STBI_ONLY_ZLIB) - #ifndef STBI_ONLY_JPEG - #define STBI_NO_JPEG - #endif - #ifndef STBI_ONLY_PNG - #define STBI_NO_PNG - #endif - #ifndef STBI_ONLY_BMP - #define STBI_NO_BMP - #endif - #ifndef STBI_ONLY_PSD - #define STBI_NO_PSD - #endif - #ifndef STBI_ONLY_TGA - #define STBI_NO_TGA - #endif - #ifndef STBI_ONLY_GIF - #define STBI_NO_GIF - #endif - #ifndef STBI_ONLY_HDR - #define STBI_NO_HDR - #endif - #ifndef STBI_ONLY_PIC - #define STBI_NO_PIC - #endif - #ifndef STBI_ONLY_PNM - #define STBI_NO_PNM - #endif -#endif - -#if defined(STBI_NO_PNG) && !defined(STBI_SUPPORT_ZLIB) && !defined(STBI_NO_ZLIB) -#define STBI_NO_ZLIB -#endif - - -#include -#include // ptrdiff_t on osx -#include -#include -#include - -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) -#include // ldexp, pow -#endif - -#ifndef STBI_NO_STDIO -#include -#endif - -#ifndef STBI_ASSERT -#include -#define STBI_ASSERT(x) assert(x) -#endif - -#ifdef __cplusplus -#define STBI_EXTERN extern "C" -#else -#define STBI_EXTERN extern -#endif - - -#ifndef _MSC_VER - #ifdef __cplusplus - #define stbi_inline inline - #else - #define stbi_inline - #endif -#else - #define stbi_inline __forceinline -#endif - -#ifndef STBI_NO_THREAD_LOCALS - #if defined(__cplusplus) && __cplusplus >= 201103L - #define STBI_THREAD_LOCAL thread_local - #elif defined(__GNUC__) && __GNUC__ < 5 - #define STBI_THREAD_LOCAL __thread - #elif defined(_MSC_VER) - #define STBI_THREAD_LOCAL __declspec(thread) - #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_THREADS__) - #define STBI_THREAD_LOCAL _Thread_local - #endif - - #ifndef STBI_THREAD_LOCAL - #if defined(__GNUC__) - #define STBI_THREAD_LOCAL __thread - #endif - #endif -#endif - -#if defined(_MSC_VER) || defined(__SYMBIAN32__) -typedef unsigned short stbi__uint16; -typedef signed short stbi__int16; -typedef unsigned int stbi__uint32; -typedef signed int stbi__int32; -#else -#include -typedef uint16_t stbi__uint16; -typedef int16_t stbi__int16; -typedef uint32_t stbi__uint32; -typedef int32_t stbi__int32; -#endif - -// should produce compiler error if size is wrong -typedef unsigned char validate_uint32[sizeof(stbi__uint32)==4 ? 1 : -1]; - -#ifdef _MSC_VER -#define STBI_NOTUSED(v) (void)(v) -#else -#define STBI_NOTUSED(v) (void)sizeof(v) -#endif - -#ifdef _MSC_VER -#define STBI_HAS_LROTL -#endif - -#ifdef STBI_HAS_LROTL - #define stbi_lrot(x,y) _lrotl(x,y) -#else - #define stbi_lrot(x,y) (((x) << (y)) | ((x) >> (-(y) & 31))) -#endif - -#if defined(STBI_MALLOC) && defined(STBI_FREE) && (defined(STBI_REALLOC) || defined(STBI_REALLOC_SIZED)) -// ok -#elif !defined(STBI_MALLOC) && !defined(STBI_FREE) && !defined(STBI_REALLOC) && !defined(STBI_REALLOC_SIZED) -// ok -#else -#error "Must define all or none of STBI_MALLOC, STBI_FREE, and STBI_REALLOC (or STBI_REALLOC_SIZED)." -#endif - -#ifndef STBI_MALLOC -#define STBI_MALLOC(sz) malloc(sz) -#define STBI_REALLOC(p,newsz) realloc(p,newsz) -#define STBI_FREE(p) free(p) -#endif - -#ifndef STBI_REALLOC_SIZED -#define STBI_REALLOC_SIZED(p,oldsz,newsz) STBI_REALLOC(p,newsz) -#endif - -// x86/x64 detection -#if defined(__x86_64__) || defined(_M_X64) -#define STBI__X64_TARGET -#elif defined(__i386) || defined(_M_IX86) -#define STBI__X86_TARGET -#endif - -#if defined(__GNUC__) && defined(STBI__X86_TARGET) && !defined(__SSE2__) && !defined(STBI_NO_SIMD) -// gcc doesn't support sse2 intrinsics unless you compile with -msse2, -// which in turn means it gets to use SSE2 everywhere. This is unfortunate, -// but previous attempts to provide the SSE2 functions with runtime -// detection caused numerous issues. The way architecture extensions are -// exposed in GCC/Clang is, sadly, not really suited for one-file libs. -// New behavior: if compiled with -msse2, we use SSE2 without any -// detection; if not, we don't use it at all. -#define STBI_NO_SIMD -#endif - -#if defined(__MINGW32__) && defined(STBI__X86_TARGET) && !defined(STBI_MINGW_ENABLE_SSE2) && !defined(STBI_NO_SIMD) -// Note that __MINGW32__ doesn't actually mean 32-bit, so we have to avoid STBI__X64_TARGET -// -// 32-bit MinGW wants ESP to be 16-byte aligned, but this is not in the -// Windows ABI and VC++ as well as Windows DLLs don't maintain that invariant. -// As a result, enabling SSE2 on 32-bit MinGW is dangerous when not -// simultaneously enabling "-mstackrealign". -// -// See https://github.com/nothings/stb/issues/81 for more information. -// -// So default to no SSE2 on 32-bit MinGW. If you've read this far and added -// -mstackrealign to your build settings, feel free to #define STBI_MINGW_ENABLE_SSE2. -#define STBI_NO_SIMD -#endif - -#if !defined(STBI_NO_SIMD) && (defined(STBI__X86_TARGET) || defined(STBI__X64_TARGET)) -#define STBI_SSE2 -#include - -#ifdef _MSC_VER - -#if _MSC_VER >= 1400 // not VC6 -#include // __cpuid -static int stbi__cpuid3(void) -{ - int info[4]; - __cpuid(info,1); - return info[3]; -} -#else -static int stbi__cpuid3(void) -{ - int res; - __asm { - mov eax,1 - cpuid - mov res,edx - } - return res; -} -#endif - -#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name - -#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) -static int stbi__sse2_available(void) -{ - int info3 = stbi__cpuid3(); - return ((info3 >> 26) & 1) != 0; -} -#endif - -#else // assume GCC-style if not VC++ -#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) - -#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) -static int stbi__sse2_available(void) -{ - // If we're even attempting to compile this on GCC/Clang, that means - // -msse2 is on, which means the compiler is allowed to use SSE2 - // instructions at will, and so are we. - return 1; -} -#endif - -#endif -#endif - -// ARM NEON -#if defined(STBI_NO_SIMD) && defined(STBI_NEON) -#undef STBI_NEON -#endif - -#ifdef STBI_NEON -#include -#ifdef _MSC_VER -#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name -#else -#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) -#endif -#endif - -#ifndef STBI_SIMD_ALIGN -#define STBI_SIMD_ALIGN(type, name) type name -#endif - -#ifndef STBI_MAX_DIMENSIONS -#define STBI_MAX_DIMENSIONS (1 << 24) -#endif - -/////////////////////////////////////////////// -// -// stbi__context struct and start_xxx functions - -// stbi__context structure is our basic context used by all images, so it -// contains all the IO context, plus some basic image information -typedef struct -{ - stbi__uint32 img_x, img_y; - int img_n, img_out_n; - - stbi_io_callbacks io; - void *io_user_data; - - int read_from_callbacks; - int buflen; - stbi_uc buffer_start[128]; - int callback_already_read; - - stbi_uc *img_buffer, *img_buffer_end; - stbi_uc *img_buffer_original, *img_buffer_original_end; -} stbi__context; - - -static void stbi__refill_buffer(stbi__context *s); - -// initialize a memory-decode context -static void stbi__start_mem(stbi__context *s, stbi_uc const *buffer, int len) -{ - s->io.read = NULL; - s->read_from_callbacks = 0; - s->callback_already_read = 0; - s->img_buffer = s->img_buffer_original = (stbi_uc *) buffer; - s->img_buffer_end = s->img_buffer_original_end = (stbi_uc *) buffer+len; -} - -// initialize a callback-based context -static void stbi__start_callbacks(stbi__context *s, stbi_io_callbacks *c, void *user) -{ - s->io = *c; - s->io_user_data = user; - s->buflen = sizeof(s->buffer_start); - s->read_from_callbacks = 1; - s->callback_already_read = 0; - s->img_buffer = s->img_buffer_original = s->buffer_start; - stbi__refill_buffer(s); - s->img_buffer_original_end = s->img_buffer_end; -} - -#ifndef STBI_NO_STDIO - -static int stbi__stdio_read(void *user, char *data, int size) -{ - return (int) fread(data,1,size,(FILE*) user); -} - -static void stbi__stdio_skip(void *user, int n) -{ - int ch; - fseek((FILE*) user, n, SEEK_CUR); - ch = fgetc((FILE*) user); /* have to read a byte to reset feof()'s flag */ - if (ch != EOF) { - ungetc(ch, (FILE *) user); /* push byte back onto stream if valid. */ - } -} - -static int stbi__stdio_eof(void *user) -{ - return feof((FILE*) user) || ferror((FILE *) user); -} - -static stbi_io_callbacks stbi__stdio_callbacks = -{ - stbi__stdio_read, - stbi__stdio_skip, - stbi__stdio_eof, -}; - -static void stbi__start_file(stbi__context *s, FILE *f) -{ - stbi__start_callbacks(s, &stbi__stdio_callbacks, (void *) f); -} - -//static void stop_file(stbi__context *s) { } - -#endif // !STBI_NO_STDIO - -static void stbi__rewind(stbi__context *s) -{ - // conceptually rewind SHOULD rewind to the beginning of the stream, - // but we just rewind to the beginning of the initial buffer, because - // we only use it after doing 'test', which only ever looks at at most 92 bytes - s->img_buffer = s->img_buffer_original; - s->img_buffer_end = s->img_buffer_original_end; -} - -enum -{ - STBI_ORDER_RGB, - STBI_ORDER_BGR -}; - -typedef struct -{ - int bits_per_channel; - int num_channels; - int channel_order; -} stbi__result_info; - -#ifndef STBI_NO_JPEG -static int stbi__jpeg_test(stbi__context *s); -static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PNG -static int stbi__png_test(stbi__context *s); -static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__png_is16(stbi__context *s); -#endif - -#ifndef STBI_NO_BMP -static int stbi__bmp_test(stbi__context *s); -static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_TGA -static int stbi__tga_test(stbi__context *s); -static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PSD -static int stbi__psd_test(stbi__context *s); -static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc); -static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__psd_is16(stbi__context *s); -#endif - -#ifndef STBI_NO_HDR -static int stbi__hdr_test(stbi__context *s); -static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PIC -static int stbi__pic_test(stbi__context *s); -static void *stbi__pic_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_GIF -static int stbi__gif_test(stbi__context *s); -static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp); -static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PNM -static int stbi__pnm_test(stbi__context *s); -static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__pnm_is16(stbi__context *s); -#endif - -static -#ifdef STBI_THREAD_LOCAL -STBI_THREAD_LOCAL -#endif -const char *stbi__g_failure_reason; - -STBIDEF const char *stbi_failure_reason(void) -{ - return stbi__g_failure_reason; -} - -#ifndef STBI_NO_FAILURE_STRINGS -static int stbi__err(const char *str) -{ - stbi__g_failure_reason = str; - return 0; -} -#endif - -static void *stbi__malloc(size_t size) -{ - return STBI_MALLOC(size); -} - -// stb_image uses ints pervasively, including for offset calculations. -// therefore the largest decoded image size we can support with the -// current code, even on 64-bit targets, is INT_MAX. this is not a -// significant limitation for the intended use case. -// -// we do, however, need to make sure our size calculations don't -// overflow. hence a few helper functions for size calculations that -// multiply integers together, making sure that they're non-negative -// and no overflow occurs. - -// return 1 if the sum is valid, 0 on overflow. -// negative terms are considered invalid. -static int stbi__addsizes_valid(int a, int b) -{ - if (b < 0) return 0; - // now 0 <= b <= INT_MAX, hence also - // 0 <= INT_MAX - b <= INTMAX. - // And "a + b <= INT_MAX" (which might overflow) is the - // same as a <= INT_MAX - b (no overflow) - return a <= INT_MAX - b; -} - -// returns 1 if the product is valid, 0 on overflow. -// negative factors are considered invalid. -static int stbi__mul2sizes_valid(int a, int b) -{ - if (a < 0 || b < 0) return 0; - if (b == 0) return 1; // mul-by-0 is always safe - // portable way to check for no overflows in a*b - return a <= INT_MAX/b; -} - -#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) -// returns 1 if "a*b + add" has no negative terms/factors and doesn't overflow -static int stbi__mad2sizes_valid(int a, int b, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__addsizes_valid(a*b, add); -} -#endif - -// returns 1 if "a*b*c + add" has no negative terms/factors and doesn't overflow -static int stbi__mad3sizes_valid(int a, int b, int c, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && - stbi__addsizes_valid(a*b*c, add); -} - -// returns 1 if "a*b*c*d + add" has no negative terms/factors and doesn't overflow -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) -static int stbi__mad4sizes_valid(int a, int b, int c, int d, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && - stbi__mul2sizes_valid(a*b*c, d) && stbi__addsizes_valid(a*b*c*d, add); -} -#endif - -#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) -// mallocs with size overflow checking -static void *stbi__malloc_mad2(int a, int b, int add) -{ - if (!stbi__mad2sizes_valid(a, b, add)) return NULL; - return stbi__malloc(a*b + add); -} -#endif - -static void *stbi__malloc_mad3(int a, int b, int c, int add) -{ - if (!stbi__mad3sizes_valid(a, b, c, add)) return NULL; - return stbi__malloc(a*b*c + add); -} - -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) -static void *stbi__malloc_mad4(int a, int b, int c, int d, int add) -{ - if (!stbi__mad4sizes_valid(a, b, c, d, add)) return NULL; - return stbi__malloc(a*b*c*d + add); -} -#endif - -// returns 1 if the sum of two signed ints is valid (between -2^31 and 2^31-1 inclusive), 0 on overflow. -static int stbi__addints_valid(int a, int b) -{ - if ((a >= 0) != (b >= 0)) return 1; // a and b have different signs, so no overflow - if (a < 0 && b < 0) return a >= INT_MIN - b; // same as a + b >= INT_MIN; INT_MIN - b cannot overflow since b < 0. - return a <= INT_MAX - b; -} - -// returns 1 if the product of two ints fits in a signed short, 0 on overflow. -static int stbi__mul2shorts_valid(int a, int b) -{ - if (b == 0 || b == -1) return 1; // multiplication by 0 is always 0; check for -1 so SHRT_MIN/b doesn't overflow - if ((a >= 0) == (b >= 0)) return a <= SHRT_MAX/b; // product is positive, so similar to mul2sizes_valid - if (b < 0) return a <= SHRT_MIN / b; // same as a * b >= SHRT_MIN - return a >= SHRT_MIN / b; -} - -// stbi__err - error -// stbi__errpf - error returning pointer to float -// stbi__errpuc - error returning pointer to unsigned char - -#ifdef STBI_NO_FAILURE_STRINGS - #define stbi__err(x,y) 0 -#elif defined(STBI_FAILURE_USERMSG) - #define stbi__err(x,y) stbi__err(y) -#else - #define stbi__err(x,y) stbi__err(x) -#endif - -#define stbi__errpf(x,y) ((float *)(size_t) (stbi__err(x,y)?NULL:NULL)) -#define stbi__errpuc(x,y) ((unsigned char *)(size_t) (stbi__err(x,y)?NULL:NULL)) - -STBIDEF void stbi_image_free(void *retval_from_stbi_load) -{ - STBI_FREE(retval_from_stbi_load); -} - -#ifndef STBI_NO_LINEAR -static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp); -#endif - -#ifndef STBI_NO_HDR -static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp); -#endif - -static int stbi__vertically_flip_on_load_global = 0; - -STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip) -{ - stbi__vertically_flip_on_load_global = flag_true_if_should_flip; -} - -#ifndef STBI_THREAD_LOCAL -#define stbi__vertically_flip_on_load stbi__vertically_flip_on_load_global -#else -static STBI_THREAD_LOCAL int stbi__vertically_flip_on_load_local, stbi__vertically_flip_on_load_set; - -STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip) -{ - stbi__vertically_flip_on_load_local = flag_true_if_should_flip; - stbi__vertically_flip_on_load_set = 1; -} - -#define stbi__vertically_flip_on_load (stbi__vertically_flip_on_load_set \ - ? stbi__vertically_flip_on_load_local \ - : stbi__vertically_flip_on_load_global) -#endif // STBI_THREAD_LOCAL - -static void *stbi__load_main(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) -{ - memset(ri, 0, sizeof(*ri)); // make sure it's initialized if we add new fields - ri->bits_per_channel = 8; // default is 8 so most paths don't have to be changed - ri->channel_order = STBI_ORDER_RGB; // all current input & output are this, but this is here so we can add BGR order - ri->num_channels = 0; - - // test the formats with a very explicit header first (at least a FOURCC - // or distinctive magic number first) - #ifndef STBI_NO_PNG - if (stbi__png_test(s)) return stbi__png_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_BMP - if (stbi__bmp_test(s)) return stbi__bmp_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_GIF - if (stbi__gif_test(s)) return stbi__gif_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_PSD - if (stbi__psd_test(s)) return stbi__psd_load(s,x,y,comp,req_comp, ri, bpc); - #else - STBI_NOTUSED(bpc); - #endif - #ifndef STBI_NO_PIC - if (stbi__pic_test(s)) return stbi__pic_load(s,x,y,comp,req_comp, ri); - #endif - - // then the formats that can end up attempting to load with just 1 or 2 - // bytes matching expectations; these are prone to false positives, so - // try them later - #ifndef STBI_NO_JPEG - if (stbi__jpeg_test(s)) return stbi__jpeg_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_PNM - if (stbi__pnm_test(s)) return stbi__pnm_load(s,x,y,comp,req_comp, ri); - #endif - - #ifndef STBI_NO_HDR - if (stbi__hdr_test(s)) { - float *hdr = stbi__hdr_load(s, x,y,comp,req_comp, ri); - return stbi__hdr_to_ldr(hdr, *x, *y, req_comp ? req_comp : *comp); - } - #endif - - #ifndef STBI_NO_TGA - // test tga last because it's a crappy test! - if (stbi__tga_test(s)) - return stbi__tga_load(s,x,y,comp,req_comp, ri); - #endif - - return stbi__errpuc("unknown image type", "Image not of any known type, or corrupt"); -} - -static stbi_uc *stbi__convert_16_to_8(stbi__uint16 *orig, int w, int h, int channels) -{ - int i; - int img_len = w * h * channels; - stbi_uc *reduced; - - reduced = (stbi_uc *) stbi__malloc(img_len); - if (reduced == NULL) return stbi__errpuc("outofmem", "Out of memory"); - - for (i = 0; i < img_len; ++i) - reduced[i] = (stbi_uc)((orig[i] >> 8) & 0xFF); // top half of each byte is sufficient approx of 16->8 bit scaling - - STBI_FREE(orig); - return reduced; -} - -static stbi__uint16 *stbi__convert_8_to_16(stbi_uc *orig, int w, int h, int channels) -{ - int i; - int img_len = w * h * channels; - stbi__uint16 *enlarged; - - enlarged = (stbi__uint16 *) stbi__malloc(img_len*2); - if (enlarged == NULL) return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); - - for (i = 0; i < img_len; ++i) - enlarged[i] = (stbi__uint16)((orig[i] << 8) + orig[i]); // replicate to high and low byte, maps 0->0, 255->0xffff - - STBI_FREE(orig); - return enlarged; -} - -static void stbi__vertical_flip(void *image, int w, int h, int bytes_per_pixel) -{ - int row; - size_t bytes_per_row = (size_t)w * bytes_per_pixel; - stbi_uc temp[2048]; - stbi_uc *bytes = (stbi_uc *)image; - - for (row = 0; row < (h>>1); row++) { - stbi_uc *row0 = bytes + row*bytes_per_row; - stbi_uc *row1 = bytes + (h - row - 1)*bytes_per_row; - // swap row0 with row1 - size_t bytes_left = bytes_per_row; - while (bytes_left) { - size_t bytes_copy = (bytes_left < sizeof(temp)) ? bytes_left : sizeof(temp); - memcpy(temp, row0, bytes_copy); - memcpy(row0, row1, bytes_copy); - memcpy(row1, temp, bytes_copy); - row0 += bytes_copy; - row1 += bytes_copy; - bytes_left -= bytes_copy; - } - } -} - -#ifndef STBI_NO_GIF -static void stbi__vertical_flip_slices(void *image, int w, int h, int z, int bytes_per_pixel) -{ - int slice; - int slice_size = w * h * bytes_per_pixel; - - stbi_uc *bytes = (stbi_uc *)image; - for (slice = 0; slice < z; ++slice) { - stbi__vertical_flip(bytes, w, h, bytes_per_pixel); - bytes += slice_size; - } -} -#endif - -static unsigned char *stbi__load_and_postprocess_8bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - stbi__result_info ri; - void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 8); - - if (result == NULL) - return NULL; - - // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. - STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); - - if (ri.bits_per_channel != 8) { - result = stbi__convert_16_to_8((stbi__uint16 *) result, *x, *y, req_comp == 0 ? *comp : req_comp); - ri.bits_per_channel = 8; - } - - // @TODO: move stbi__convert_format to here - - if (stbi__vertically_flip_on_load) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi_uc)); - } - - return (unsigned char *) result; -} - -static stbi__uint16 *stbi__load_and_postprocess_16bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - stbi__result_info ri; - void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 16); - - if (result == NULL) - return NULL; - - // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. - STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); - - if (ri.bits_per_channel != 16) { - result = stbi__convert_8_to_16((stbi_uc *) result, *x, *y, req_comp == 0 ? *comp : req_comp); - ri.bits_per_channel = 16; - } - - // @TODO: move stbi__convert_format16 to here - // @TODO: special case RGB-to-Y (and RGBA-to-YA) for 8-bit-to-16-bit case to keep more precision - - if (stbi__vertically_flip_on_load) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi__uint16)); - } - - return (stbi__uint16 *) result; -} - -#if !defined(STBI_NO_HDR) && !defined(STBI_NO_LINEAR) -static void stbi__float_postprocess(float *result, int *x, int *y, int *comp, int req_comp) -{ - if (stbi__vertically_flip_on_load && result != NULL) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(float)); - } -} -#endif - -#ifndef STBI_NO_STDIO - -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) -STBI_EXTERN __declspec(dllimport) int __stdcall MultiByteToWideChar(unsigned int cp, unsigned long flags, const char *str, int cbmb, wchar_t *widestr, int cchwide); -STBI_EXTERN __declspec(dllimport) int __stdcall WideCharToMultiByte(unsigned int cp, unsigned long flags, const wchar_t *widestr, int cchwide, char *str, int cbmb, const char *defchar, int *used_default); -#endif - -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) -STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input) -{ - return WideCharToMultiByte(65001 /* UTF8 */, 0, input, -1, buffer, (int) bufferlen, NULL, NULL); -} -#endif - -static FILE *stbi__fopen(char const *filename, char const *mode) -{ - FILE *f; -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) - wchar_t wMode[64]; - wchar_t wFilename[1024]; - if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, filename, -1, wFilename, sizeof(wFilename)/sizeof(*wFilename))) - return 0; - - if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, mode, -1, wMode, sizeof(wMode)/sizeof(*wMode))) - return 0; - -#if defined(_MSC_VER) && _MSC_VER >= 1400 - if (0 != _wfopen_s(&f, wFilename, wMode)) - f = 0; -#else - f = _wfopen(wFilename, wMode); -#endif - -#elif defined(_MSC_VER) && _MSC_VER >= 1400 - if (0 != fopen_s(&f, filename, mode)) - f=0; -#else - f = fopen(filename, mode); -#endif - return f; -} - - -STBIDEF stbi_uc *stbi_load(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - unsigned char *result; - if (!f) return stbi__errpuc("can't fopen", "Unable to open file"); - result = stbi_load_from_file(f,x,y,comp,req_comp); - fclose(f); - return result; -} - -STBIDEF stbi_uc *stbi_load_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - unsigned char *result; - stbi__context s; - stbi__start_file(&s,f); - result = stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); - if (result) { - // need to 'unget' all the characters in the IO buffer - fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); - } - return result; -} - -STBIDEF stbi__uint16 *stbi_load_from_file_16(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - stbi__uint16 *result; - stbi__context s; - stbi__start_file(&s,f); - result = stbi__load_and_postprocess_16bit(&s,x,y,comp,req_comp); - if (result) { - // need to 'unget' all the characters in the IO buffer - fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); - } - return result; -} - -STBIDEF stbi_us *stbi_load_16(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - stbi__uint16 *result; - if (!f) return (stbi_us *) stbi__errpuc("can't fopen", "Unable to open file"); - result = stbi_load_from_file_16(f,x,y,comp,req_comp); - fclose(f); - return result; -} - - -#endif //!STBI_NO_STDIO - -STBIDEF stbi_us *stbi_load_16_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); -} - -STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *)clbk, user); - return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); -} - -STBIDEF stbi_uc *stbi_load_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); -} - -STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); -} - -#ifndef STBI_NO_GIF -STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp) -{ - unsigned char *result; - stbi__context s; - stbi__start_mem(&s,buffer,len); - - result = (unsigned char*) stbi__load_gif_main(&s, delays, x, y, z, comp, req_comp); - if (stbi__vertically_flip_on_load) { - stbi__vertical_flip_slices( result, *x, *y, *z, *comp ); - } - - return result; -} -#endif - -#ifndef STBI_NO_LINEAR -static float *stbi__loadf_main(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - unsigned char *data; - #ifndef STBI_NO_HDR - if (stbi__hdr_test(s)) { - stbi__result_info ri; - float *hdr_data = stbi__hdr_load(s,x,y,comp,req_comp, &ri); - if (hdr_data) - stbi__float_postprocess(hdr_data,x,y,comp,req_comp); - return hdr_data; - } - #endif - data = stbi__load_and_postprocess_8bit(s, x, y, comp, req_comp); - if (data) - return stbi__ldr_to_hdr(data, *x, *y, req_comp ? req_comp : *comp); - return stbi__errpf("unknown image type", "Image not of any known type, or corrupt"); -} - -STBIDEF float *stbi_loadf_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} - -STBIDEF float *stbi_loadf_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} - -#ifndef STBI_NO_STDIO -STBIDEF float *stbi_loadf(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - float *result; - FILE *f = stbi__fopen(filename, "rb"); - if (!f) return stbi__errpf("can't fopen", "Unable to open file"); - result = stbi_loadf_from_file(f,x,y,comp,req_comp); - fclose(f); - return result; -} - -STBIDEF float *stbi_loadf_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_file(&s,f); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} -#endif // !STBI_NO_STDIO - -#endif // !STBI_NO_LINEAR - -// these is-hdr-or-not is defined independent of whether STBI_NO_LINEAR is -// defined, for API simplicity; if STBI_NO_LINEAR is defined, it always -// reports false! - -STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len) -{ - #ifndef STBI_NO_HDR - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__hdr_test(&s); - #else - STBI_NOTUSED(buffer); - STBI_NOTUSED(len); - return 0; - #endif -} - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_is_hdr (char const *filename) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result=0; - if (f) { - result = stbi_is_hdr_from_file(f); - fclose(f); - } - return result; -} - -STBIDEF int stbi_is_hdr_from_file(FILE *f) -{ - #ifndef STBI_NO_HDR - long pos = ftell(f); - int res; - stbi__context s; - stbi__start_file(&s,f); - res = stbi__hdr_test(&s); - fseek(f, pos, SEEK_SET); - return res; - #else - STBI_NOTUSED(f); - return 0; - #endif -} -#endif // !STBI_NO_STDIO - -STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user) -{ - #ifndef STBI_NO_HDR - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__hdr_test(&s); - #else - STBI_NOTUSED(clbk); - STBI_NOTUSED(user); - return 0; - #endif -} - -#ifndef STBI_NO_LINEAR -static float stbi__l2h_gamma=2.2f, stbi__l2h_scale=1.0f; - -STBIDEF void stbi_ldr_to_hdr_gamma(float gamma) { stbi__l2h_gamma = gamma; } -STBIDEF void stbi_ldr_to_hdr_scale(float scale) { stbi__l2h_scale = scale; } -#endif - -static float stbi__h2l_gamma_i=1.0f/2.2f, stbi__h2l_scale_i=1.0f; - -STBIDEF void stbi_hdr_to_ldr_gamma(float gamma) { stbi__h2l_gamma_i = 1/gamma; } -STBIDEF void stbi_hdr_to_ldr_scale(float scale) { stbi__h2l_scale_i = 1/scale; } - - -////////////////////////////////////////////////////////////////////////////// -// -// Common code used by all image loaders -// - -enum -{ - STBI__SCAN_load=0, - STBI__SCAN_type, - STBI__SCAN_header -}; - -static void stbi__refill_buffer(stbi__context *s) -{ - int n = (s->io.read)(s->io_user_data,(char*)s->buffer_start,s->buflen); - s->callback_already_read += (int) (s->img_buffer - s->img_buffer_original); - if (n == 0) { - // at end of file, treat same as if from memory, but need to handle case - // where s->img_buffer isn't pointing to safe memory, e.g. 0-byte file - s->read_from_callbacks = 0; - s->img_buffer = s->buffer_start; - s->img_buffer_end = s->buffer_start+1; - *s->img_buffer = 0; - } else { - s->img_buffer = s->buffer_start; - s->img_buffer_end = s->buffer_start + n; - } -} - -stbi_inline static stbi_uc stbi__get8(stbi__context *s) -{ - if (s->img_buffer < s->img_buffer_end) - return *s->img_buffer++; - if (s->read_from_callbacks) { - stbi__refill_buffer(s); - return *s->img_buffer++; - } - return 0; -} - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_HDR) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -stbi_inline static int stbi__at_eof(stbi__context *s) -{ - if (s->io.read) { - if (!(s->io.eof)(s->io_user_data)) return 0; - // if feof() is true, check if buffer = end - // special case: we've only got the special 0 character at the end - if (s->read_from_callbacks == 0) return 1; - } - - return s->img_buffer >= s->img_buffer_end; -} -#endif - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) -// nothing -#else -static void stbi__skip(stbi__context *s, int n) -{ - if (n == 0) return; // already there! - if (n < 0) { - s->img_buffer = s->img_buffer_end; - return; - } - if (s->io.read) { - int blen = (int) (s->img_buffer_end - s->img_buffer); - if (blen < n) { - s->img_buffer = s->img_buffer_end; - (s->io.skip)(s->io_user_data, n - blen); - return; - } - } - s->img_buffer += n; -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_TGA) && defined(STBI_NO_HDR) && defined(STBI_NO_PNM) -// nothing -#else -static int stbi__getn(stbi__context *s, stbi_uc *buffer, int n) -{ - if (s->io.read) { - int blen = (int) (s->img_buffer_end - s->img_buffer); - if (blen < n) { - int res, count; - - memcpy(buffer, s->img_buffer, blen); - - count = (s->io.read)(s->io_user_data, (char*) buffer + blen, n - blen); - res = (count == (n-blen)); - s->img_buffer = s->img_buffer_end; - return res; - } - } - - if (s->img_buffer+n <= s->img_buffer_end) { - memcpy(buffer, s->img_buffer, n); - s->img_buffer += n; - return 1; - } else - return 0; -} -#endif - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) -// nothing -#else -static int stbi__get16be(stbi__context *s) -{ - int z = stbi__get8(s); - return (z << 8) + stbi__get8(s); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) -// nothing -#else -static stbi__uint32 stbi__get32be(stbi__context *s) -{ - stbi__uint32 z = stbi__get16be(s); - return (z << 16) + stbi__get16be(s); -} -#endif - -#if defined(STBI_NO_BMP) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) -// nothing -#else -static int stbi__get16le(stbi__context *s) -{ - int z = stbi__get8(s); - return z + (stbi__get8(s) << 8); -} -#endif - -#ifndef STBI_NO_BMP -static stbi__uint32 stbi__get32le(stbi__context *s) -{ - stbi__uint32 z = stbi__get16le(s); - z += (stbi__uint32)stbi__get16le(s) << 16; - return z; -} -#endif - -#define STBI__BYTECAST(x) ((stbi_uc) ((x) & 255)) // truncate int to byte without warnings - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -////////////////////////////////////////////////////////////////////////////// -// -// generic converter from built-in img_n to req_comp -// individual types do this automatically as much as possible (e.g. jpeg -// does all cases internally since it needs to colorspace convert anyway, -// and it never has alpha, so very few cases ). png can automatically -// interleave an alpha=255 channel, but falls back to this for other cases -// -// assume data buffer is malloced, so malloc a new one and free that one -// only failure mode is malloc failing - -static stbi_uc stbi__compute_y(int r, int g, int b) -{ - return (stbi_uc) (((r*77) + (g*150) + (29*b)) >> 8); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -static unsigned char *stbi__convert_format(unsigned char *data, int img_n, int req_comp, unsigned int x, unsigned int y) -{ - int i,j; - unsigned char *good; - - if (req_comp == img_n) return data; - STBI_ASSERT(req_comp >= 1 && req_comp <= 4); - - good = (unsigned char *) stbi__malloc_mad3(req_comp, x, y, 0); - if (good == NULL) { - STBI_FREE(data); - return stbi__errpuc("outofmem", "Out of memory"); - } - - for (j=0; j < (int) y; ++j) { - unsigned char *src = data + j * x * img_n ; - unsigned char *dest = good + j * x * req_comp; - - #define STBI__COMBO(a,b) ((a)*8+(b)) - #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) - // convert source image with img_n components to one with req_comp components; - // avoid switch per pixel, so use switch per scanline and massive macros - switch (STBI__COMBO(img_n, req_comp)) { - STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=255; } break; - STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=255; } break; - STBI__CASE(2,1) { dest[0]=src[0]; } break; - STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; - STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=255; } break; - STBI__CASE(3,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; - STBI__CASE(3,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = 255; } break; - STBI__CASE(4,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; - STBI__CASE(4,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = src[3]; } break; - STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; - default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return stbi__errpuc("unsupported", "Unsupported format conversion"); - } - #undef STBI__CASE - } - - STBI_FREE(data); - return good; -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) -// nothing -#else -static stbi__uint16 stbi__compute_y_16(int r, int g, int b) -{ - return (stbi__uint16) (((r*77) + (g*150) + (29*b)) >> 8); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) -// nothing -#else -static stbi__uint16 *stbi__convert_format16(stbi__uint16 *data, int img_n, int req_comp, unsigned int x, unsigned int y) -{ - int i,j; - stbi__uint16 *good; - - if (req_comp == img_n) return data; - STBI_ASSERT(req_comp >= 1 && req_comp <= 4); - - good = (stbi__uint16 *) stbi__malloc(req_comp * x * y * 2); - if (good == NULL) { - STBI_FREE(data); - return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); - } - - for (j=0; j < (int) y; ++j) { - stbi__uint16 *src = data + j * x * img_n ; - stbi__uint16 *dest = good + j * x * req_comp; - - #define STBI__COMBO(a,b) ((a)*8+(b)) - #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) - // convert source image with img_n components to one with req_comp components; - // avoid switch per pixel, so use switch per scanline and massive macros - switch (STBI__COMBO(img_n, req_comp)) { - STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=0xffff; } break; - STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=0xffff; } break; - STBI__CASE(2,1) { dest[0]=src[0]; } break; - STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; - STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=0xffff; } break; - STBI__CASE(3,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; - STBI__CASE(3,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = 0xffff; } break; - STBI__CASE(4,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; - STBI__CASE(4,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = src[3]; } break; - STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; - default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return (stbi__uint16*) stbi__errpuc("unsupported", "Unsupported format conversion"); - } - #undef STBI__CASE - } - - STBI_FREE(data); - return good; -} -#endif - -#ifndef STBI_NO_LINEAR -static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp) -{ - int i,k,n; - float *output; - if (!data) return NULL; - output = (float *) stbi__malloc_mad4(x, y, comp, sizeof(float), 0); - if (output == NULL) { STBI_FREE(data); return stbi__errpf("outofmem", "Out of memory"); } - // compute number of non-alpha components - if (comp & 1) n = comp; else n = comp-1; - for (i=0; i < x*y; ++i) { - for (k=0; k < n; ++k) { - output[i*comp + k] = (float) (pow(data[i*comp+k]/255.0f, stbi__l2h_gamma) * stbi__l2h_scale); - } - } - if (n < comp) { - for (i=0; i < x*y; ++i) { - output[i*comp + n] = data[i*comp + n]/255.0f; - } - } - STBI_FREE(data); - return output; -} -#endif - -#ifndef STBI_NO_HDR -#define stbi__float2int(x) ((int) (x)) -static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp) -{ - int i,k,n; - stbi_uc *output; - if (!data) return NULL; - output = (stbi_uc *) stbi__malloc_mad3(x, y, comp, 0); - if (output == NULL) { STBI_FREE(data); return stbi__errpuc("outofmem", "Out of memory"); } - // compute number of non-alpha components - if (comp & 1) n = comp; else n = comp-1; - for (i=0; i < x*y; ++i) { - for (k=0; k < n; ++k) { - float z = (float) pow(data[i*comp+k]*stbi__h2l_scale_i, stbi__h2l_gamma_i) * 255 + 0.5f; - if (z < 0) z = 0; - if (z > 255) z = 255; - output[i*comp + k] = (stbi_uc) stbi__float2int(z); - } - if (k < comp) { - float z = data[i*comp+k] * 255 + 0.5f; - if (z < 0) z = 0; - if (z > 255) z = 255; - output[i*comp + k] = (stbi_uc) stbi__float2int(z); - } - } - STBI_FREE(data); - return output; -} -#endif - -////////////////////////////////////////////////////////////////////////////// -// -// "baseline" JPEG/JFIF decoder -// -// simple implementation -// - doesn't support delayed output of y-dimension -// - simple interface (only one output format: 8-bit interleaved RGB) -// - doesn't try to recover corrupt jpegs -// - doesn't allow partial loading, loading multiple at once -// - still fast on x86 (copying globals into locals doesn't help x86) -// - allocates lots of intermediate memory (full size of all components) -// - non-interleaved case requires this anyway -// - allows good upsampling (see next) -// high-quality -// - upsampled channels are bilinearly interpolated, even across blocks -// - quality integer IDCT derived from IJG's 'slow' -// performance -// - fast huffman; reasonable integer IDCT -// - some SIMD kernels for common paths on targets with SSE2/NEON -// - uses a lot of intermediate memory, could cache poorly - -#ifndef STBI_NO_JPEG - -// huffman decoding acceleration -#define FAST_BITS 9 // larger handles more cases; smaller stomps less cache - -typedef struct -{ - stbi_uc fast[1 << FAST_BITS]; - // weirdly, repacking this into AoS is a 10% speed loss, instead of a win - stbi__uint16 code[256]; - stbi_uc values[256]; - stbi_uc size[257]; - unsigned int maxcode[18]; - int delta[17]; // old 'firstsymbol' - old 'firstcode' -} stbi__huffman; - -typedef struct -{ - stbi__context *s; - stbi__huffman huff_dc[4]; - stbi__huffman huff_ac[4]; - stbi__uint16 dequant[4][64]; - stbi__int16 fast_ac[4][1 << FAST_BITS]; - -// sizes for components, interleaved MCUs - int img_h_max, img_v_max; - int img_mcu_x, img_mcu_y; - int img_mcu_w, img_mcu_h; - -// definition of jpeg image component - struct - { - int id; - int h,v; - int tq; - int hd,ha; - int dc_pred; - - int x,y,w2,h2; - stbi_uc *data; - void *raw_data, *raw_coeff; - stbi_uc *linebuf; - short *coeff; // progressive only - int coeff_w, coeff_h; // number of 8x8 coefficient blocks - } img_comp[4]; - - stbi__uint32 code_buffer; // jpeg entropy-coded buffer - int code_bits; // number of valid bits - unsigned char marker; // marker seen while filling entropy buffer - int nomore; // flag if we saw a marker so must stop - - int progressive; - int spec_start; - int spec_end; - int succ_high; - int succ_low; - int eob_run; - int jfif; - int app14_color_transform; // Adobe APP14 tag - int rgb; - - int scan_n, order[4]; - int restart_interval, todo; - -// kernels - void (*idct_block_kernel)(stbi_uc *out, int out_stride, short data[64]); - void (*YCbCr_to_RGB_kernel)(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step); - stbi_uc *(*resample_row_hv_2_kernel)(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs); -} stbi__jpeg; - -static int stbi__build_huffman(stbi__huffman *h, int *count) -{ - int i,j,k=0; - unsigned int code; - // build size list for each symbol (from JPEG spec) - for (i=0; i < 16; ++i) { - for (j=0; j < count[i]; ++j) { - h->size[k++] = (stbi_uc) (i+1); - if(k >= 257) return stbi__err("bad size list","Corrupt JPEG"); - } - } - h->size[k] = 0; - - // compute actual symbols (from jpeg spec) - code = 0; - k = 0; - for(j=1; j <= 16; ++j) { - // compute delta to add to code to compute symbol id - h->delta[j] = k - code; - if (h->size[k] == j) { - while (h->size[k] == j) - h->code[k++] = (stbi__uint16) (code++); - if (code-1 >= (1u << j)) return stbi__err("bad code lengths","Corrupt JPEG"); - } - // compute largest code + 1 for this size, preshifted as needed later - h->maxcode[j] = code << (16-j); - code <<= 1; - } - h->maxcode[j] = 0xffffffff; - - // build non-spec acceleration table; 255 is flag for not-accelerated - memset(h->fast, 255, 1 << FAST_BITS); - for (i=0; i < k; ++i) { - int s = h->size[i]; - if (s <= FAST_BITS) { - int c = h->code[i] << (FAST_BITS-s); - int m = 1 << (FAST_BITS-s); - for (j=0; j < m; ++j) { - h->fast[c+j] = (stbi_uc) i; - } - } - } - return 1; -} - -// build a table that decodes both magnitude and value of small ACs in -// one go. -static void stbi__build_fast_ac(stbi__int16 *fast_ac, stbi__huffman *h) -{ - int i; - for (i=0; i < (1 << FAST_BITS); ++i) { - stbi_uc fast = h->fast[i]; - fast_ac[i] = 0; - if (fast < 255) { - int rs = h->values[fast]; - int run = (rs >> 4) & 15; - int magbits = rs & 15; - int len = h->size[fast]; - - if (magbits && len + magbits <= FAST_BITS) { - // magnitude code followed by receive_extend code - int k = ((i << len) & ((1 << FAST_BITS) - 1)) >> (FAST_BITS - magbits); - int m = 1 << (magbits - 1); - if (k < m) k += (~0U << magbits) + 1; - // if the result is small enough, we can fit it in fast_ac table - if (k >= -128 && k <= 127) - fast_ac[i] = (stbi__int16) ((k * 256) + (run * 16) + (len + magbits)); - } - } - } -} - -static void stbi__grow_buffer_unsafe(stbi__jpeg *j) -{ - do { - unsigned int b = j->nomore ? 0 : stbi__get8(j->s); - if (b == 0xff) { - int c = stbi__get8(j->s); - while (c == 0xff) c = stbi__get8(j->s); // consume fill bytes - if (c != 0) { - j->marker = (unsigned char) c; - j->nomore = 1; - return; - } - } - j->code_buffer |= b << (24 - j->code_bits); - j->code_bits += 8; - } while (j->code_bits <= 24); -} - -// (1 << n) - 1 -static const stbi__uint32 stbi__bmask[17]={0,1,3,7,15,31,63,127,255,511,1023,2047,4095,8191,16383,32767,65535}; - -// decode a jpeg huffman value from the bitstream -stbi_inline static int stbi__jpeg_huff_decode(stbi__jpeg *j, stbi__huffman *h) -{ - unsigned int temp; - int c,k; - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - - // look at the top FAST_BITS and determine what symbol ID it is, - // if the code is <= FAST_BITS - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - k = h->fast[c]; - if (k < 255) { - int s = h->size[k]; - if (s > j->code_bits) - return -1; - j->code_buffer <<= s; - j->code_bits -= s; - return h->values[k]; - } - - // naive test is to shift the code_buffer down so k bits are - // valid, then test against maxcode. To speed this up, we've - // preshifted maxcode left so that it has (16-k) 0s at the - // end; in other words, regardless of the number of bits, it - // wants to be compared against something shifted to have 16; - // that way we don't need to shift inside the loop. - temp = j->code_buffer >> 16; - for (k=FAST_BITS+1 ; ; ++k) - if (temp < h->maxcode[k]) - break; - if (k == 17) { - // error! code not found - j->code_bits -= 16; - return -1; - } - - if (k > j->code_bits) - return -1; - - // convert the huffman code to the symbol id - c = ((j->code_buffer >> (32 - k)) & stbi__bmask[k]) + h->delta[k]; - if(c < 0 || c >= 256) // symbol id out of bounds! - return -1; - STBI_ASSERT((((j->code_buffer) >> (32 - h->size[c])) & stbi__bmask[h->size[c]]) == h->code[c]); - - // convert the id to a symbol - j->code_bits -= k; - j->code_buffer <<= k; - return h->values[c]; -} - -// bias[n] = (-1<code_bits < n) stbi__grow_buffer_unsafe(j); - if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing - - sgn = j->code_buffer >> 31; // sign bit always in MSB; 0 if MSB clear (positive), 1 if MSB set (negative) - k = stbi_lrot(j->code_buffer, n); - j->code_buffer = k & ~stbi__bmask[n]; - k &= stbi__bmask[n]; - j->code_bits -= n; - return k + (stbi__jbias[n] & (sgn - 1)); -} - -// get some unsigned bits -stbi_inline static int stbi__jpeg_get_bits(stbi__jpeg *j, int n) -{ - unsigned int k; - if (j->code_bits < n) stbi__grow_buffer_unsafe(j); - if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing - k = stbi_lrot(j->code_buffer, n); - j->code_buffer = k & ~stbi__bmask[n]; - k &= stbi__bmask[n]; - j->code_bits -= n; - return k; -} - -stbi_inline static int stbi__jpeg_get_bit(stbi__jpeg *j) -{ - unsigned int k; - if (j->code_bits < 1) stbi__grow_buffer_unsafe(j); - if (j->code_bits < 1) return 0; // ran out of bits from stream, return 0s intead of continuing - k = j->code_buffer; - j->code_buffer <<= 1; - --j->code_bits; - return k & 0x80000000; -} - -// given a value that's at position X in the zigzag stream, -// where does it appear in the 8x8 matrix coded as row-major? -static const stbi_uc stbi__jpeg_dezigzag[64+15] = -{ - 0, 1, 8, 16, 9, 2, 3, 10, - 17, 24, 32, 25, 18, 11, 4, 5, - 12, 19, 26, 33, 40, 48, 41, 34, - 27, 20, 13, 6, 7, 14, 21, 28, - 35, 42, 49, 56, 57, 50, 43, 36, - 29, 22, 15, 23, 30, 37, 44, 51, - 58, 59, 52, 45, 38, 31, 39, 46, - 53, 60, 61, 54, 47, 55, 62, 63, - // let corrupt input sample past end - 63, 63, 63, 63, 63, 63, 63, 63, - 63, 63, 63, 63, 63, 63, 63 -}; - -// decode one 64-entry block-- -static int stbi__jpeg_decode_block(stbi__jpeg *j, short data[64], stbi__huffman *hdc, stbi__huffman *hac, stbi__int16 *fac, int b, stbi__uint16 *dequant) -{ - int diff,dc,k; - int t; - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - t = stbi__jpeg_huff_decode(j, hdc); - if (t < 0 || t > 15) return stbi__err("bad huffman code","Corrupt JPEG"); - - // 0 all the ac values now so we can do it 32-bits at a time - memset(data,0,64*sizeof(data[0])); - - diff = t ? stbi__extend_receive(j, t) : 0; - if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta","Corrupt JPEG"); - dc = j->img_comp[b].dc_pred + diff; - j->img_comp[b].dc_pred = dc; - if (!stbi__mul2shorts_valid(dc, dequant[0])) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - data[0] = (short) (dc * dequant[0]); - - // decode AC components, see JPEG spec - k = 1; - do { - unsigned int zig; - int c,r,s; - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - r = fac[c]; - if (r) { // fast-AC path - k += (r >> 4) & 15; // run - s = r & 15; // combined length - if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); - j->code_buffer <<= s; - j->code_bits -= s; - // decode into unzigzag'd location - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) ((r >> 8) * dequant[zig]); - } else { - int rs = stbi__jpeg_huff_decode(j, hac); - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (rs != 0xf0) break; // end block - k += 16; - } else { - k += r; - // decode into unzigzag'd location - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) (stbi__extend_receive(j,s) * dequant[zig]); - } - } - } while (k < 64); - return 1; -} - -static int stbi__jpeg_decode_block_prog_dc(stbi__jpeg *j, short data[64], stbi__huffman *hdc, int b) -{ - int diff,dc; - int t; - if (j->spec_end != 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - - if (j->succ_high == 0) { - // first scan for DC coefficient, must be first - memset(data,0,64*sizeof(data[0])); // 0 all the ac values now - t = stbi__jpeg_huff_decode(j, hdc); - if (t < 0 || t > 15) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - diff = t ? stbi__extend_receive(j, t) : 0; - - if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta", "Corrupt JPEG"); - dc = j->img_comp[b].dc_pred + diff; - j->img_comp[b].dc_pred = dc; - if (!stbi__mul2shorts_valid(dc, 1 << j->succ_low)) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - data[0] = (short) (dc * (1 << j->succ_low)); - } else { - // refinement scan for DC coefficient - if (stbi__jpeg_get_bit(j)) - data[0] += (short) (1 << j->succ_low); - } - return 1; -} - -// @OPTIMIZE: store non-zigzagged during the decode passes, -// and only de-zigzag when dequantizing -static int stbi__jpeg_decode_block_prog_ac(stbi__jpeg *j, short data[64], stbi__huffman *hac, stbi__int16 *fac) -{ - int k; - if (j->spec_start == 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - - if (j->succ_high == 0) { - int shift = j->succ_low; - - if (j->eob_run) { - --j->eob_run; - return 1; - } - - k = j->spec_start; - do { - unsigned int zig; - int c,r,s; - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - r = fac[c]; - if (r) { // fast-AC path - k += (r >> 4) & 15; // run - s = r & 15; // combined length - if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); - j->code_buffer <<= s; - j->code_bits -= s; - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) ((r >> 8) * (1 << shift)); - } else { - int rs = stbi__jpeg_huff_decode(j, hac); - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (r < 15) { - j->eob_run = (1 << r); - if (r) - j->eob_run += stbi__jpeg_get_bits(j, r); - --j->eob_run; - break; - } - k += 16; - } else { - k += r; - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) (stbi__extend_receive(j,s) * (1 << shift)); - } - } - } while (k <= j->spec_end); - } else { - // refinement scan for these AC coefficients - - short bit = (short) (1 << j->succ_low); - - if (j->eob_run) { - --j->eob_run; - for (k = j->spec_start; k <= j->spec_end; ++k) { - short *p = &data[stbi__jpeg_dezigzag[k]]; - if (*p != 0) - if (stbi__jpeg_get_bit(j)) - if ((*p & bit)==0) { - if (*p > 0) - *p += bit; - else - *p -= bit; - } - } - } else { - k = j->spec_start; - do { - int r,s; - int rs = stbi__jpeg_huff_decode(j, hac); // @OPTIMIZE see if we can use the fast path here, advance-by-r is so slow, eh - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (r < 15) { - j->eob_run = (1 << r) - 1; - if (r) - j->eob_run += stbi__jpeg_get_bits(j, r); - r = 64; // force end of block - } else { - // r=15 s=0 should write 16 0s, so we just do - // a run of 15 0s and then write s (which is 0), - // so we don't have to do anything special here - } - } else { - if (s != 1) return stbi__err("bad huffman code", "Corrupt JPEG"); - // sign bit - if (stbi__jpeg_get_bit(j)) - s = bit; - else - s = -bit; - } - - // advance by r - while (k <= j->spec_end) { - short *p = &data[stbi__jpeg_dezigzag[k++]]; - if (*p != 0) { - if (stbi__jpeg_get_bit(j)) - if ((*p & bit)==0) { - if (*p > 0) - *p += bit; - else - *p -= bit; - } - } else { - if (r == 0) { - *p = (short) s; - break; - } - --r; - } - } - } while (k <= j->spec_end); - } - } - return 1; -} - -// take a -128..127 value and stbi__clamp it and convert to 0..255 -stbi_inline static stbi_uc stbi__clamp(int x) -{ - // trick to use a single test to catch both cases - if ((unsigned int) x > 255) { - if (x < 0) return 0; - if (x > 255) return 255; - } - return (stbi_uc) x; -} - -#define stbi__f2f(x) ((int) (((x) * 4096 + 0.5))) -#define stbi__fsh(x) ((x) * 4096) - -// derived from jidctint -- DCT_ISLOW -#define STBI__IDCT_1D(s0,s1,s2,s3,s4,s5,s6,s7) \ - int t0,t1,t2,t3,p1,p2,p3,p4,p5,x0,x1,x2,x3; \ - p2 = s2; \ - p3 = s6; \ - p1 = (p2+p3) * stbi__f2f(0.5411961f); \ - t2 = p1 + p3*stbi__f2f(-1.847759065f); \ - t3 = p1 + p2*stbi__f2f( 0.765366865f); \ - p2 = s0; \ - p3 = s4; \ - t0 = stbi__fsh(p2+p3); \ - t1 = stbi__fsh(p2-p3); \ - x0 = t0+t3; \ - x3 = t0-t3; \ - x1 = t1+t2; \ - x2 = t1-t2; \ - t0 = s7; \ - t1 = s5; \ - t2 = s3; \ - t3 = s1; \ - p3 = t0+t2; \ - p4 = t1+t3; \ - p1 = t0+t3; \ - p2 = t1+t2; \ - p5 = (p3+p4)*stbi__f2f( 1.175875602f); \ - t0 = t0*stbi__f2f( 0.298631336f); \ - t1 = t1*stbi__f2f( 2.053119869f); \ - t2 = t2*stbi__f2f( 3.072711026f); \ - t3 = t3*stbi__f2f( 1.501321110f); \ - p1 = p5 + p1*stbi__f2f(-0.899976223f); \ - p2 = p5 + p2*stbi__f2f(-2.562915447f); \ - p3 = p3*stbi__f2f(-1.961570560f); \ - p4 = p4*stbi__f2f(-0.390180644f); \ - t3 += p1+p4; \ - t2 += p2+p3; \ - t1 += p2+p4; \ - t0 += p1+p3; - -static void stbi__idct_block(stbi_uc *out, int out_stride, short data[64]) -{ - int i,val[64],*v=val; - stbi_uc *o; - short *d = data; - - // columns - for (i=0; i < 8; ++i,++d, ++v) { - // if all zeroes, shortcut -- this avoids dequantizing 0s and IDCTing - if (d[ 8]==0 && d[16]==0 && d[24]==0 && d[32]==0 - && d[40]==0 && d[48]==0 && d[56]==0) { - // no shortcut 0 seconds - // (1|2|3|4|5|6|7)==0 0 seconds - // all separate -0.047 seconds - // 1 && 2|3 && 4|5 && 6|7: -0.047 seconds - int dcterm = d[0]*4; - v[0] = v[8] = v[16] = v[24] = v[32] = v[40] = v[48] = v[56] = dcterm; - } else { - STBI__IDCT_1D(d[ 0],d[ 8],d[16],d[24],d[32],d[40],d[48],d[56]) - // constants scaled things up by 1<<12; let's bring them back - // down, but keep 2 extra bits of precision - x0 += 512; x1 += 512; x2 += 512; x3 += 512; - v[ 0] = (x0+t3) >> 10; - v[56] = (x0-t3) >> 10; - v[ 8] = (x1+t2) >> 10; - v[48] = (x1-t2) >> 10; - v[16] = (x2+t1) >> 10; - v[40] = (x2-t1) >> 10; - v[24] = (x3+t0) >> 10; - v[32] = (x3-t0) >> 10; - } - } - - for (i=0, v=val, o=out; i < 8; ++i,v+=8,o+=out_stride) { - // no fast case since the first 1D IDCT spread components out - STBI__IDCT_1D(v[0],v[1],v[2],v[3],v[4],v[5],v[6],v[7]) - // constants scaled things up by 1<<12, plus we had 1<<2 from first - // loop, plus horizontal and vertical each scale by sqrt(8) so together - // we've got an extra 1<<3, so 1<<17 total we need to remove. - // so we want to round that, which means adding 0.5 * 1<<17, - // aka 65536. Also, we'll end up with -128 to 127 that we want - // to encode as 0..255 by adding 128, so we'll add that before the shift - x0 += 65536 + (128<<17); - x1 += 65536 + (128<<17); - x2 += 65536 + (128<<17); - x3 += 65536 + (128<<17); - // tried computing the shifts into temps, or'ing the temps to see - // if any were out of range, but that was slower - o[0] = stbi__clamp((x0+t3) >> 17); - o[7] = stbi__clamp((x0-t3) >> 17); - o[1] = stbi__clamp((x1+t2) >> 17); - o[6] = stbi__clamp((x1-t2) >> 17); - o[2] = stbi__clamp((x2+t1) >> 17); - o[5] = stbi__clamp((x2-t1) >> 17); - o[3] = stbi__clamp((x3+t0) >> 17); - o[4] = stbi__clamp((x3-t0) >> 17); - } -} - -#ifdef STBI_SSE2 -// sse2 integer IDCT. not the fastest possible implementation but it -// produces bit-identical results to the generic C version so it's -// fully "transparent". -static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) -{ - // This is constructed to match our regular (generic) integer IDCT exactly. - __m128i row0, row1, row2, row3, row4, row5, row6, row7; - __m128i tmp; - - // dot product constant: even elems=x, odd elems=y - #define dct_const(x,y) _mm_setr_epi16((x),(y),(x),(y),(x),(y),(x),(y)) - - // out(0) = c0[even]*x + c0[odd]*y (c0, x, y 16-bit, out 32-bit) - // out(1) = c1[even]*x + c1[odd]*y - #define dct_rot(out0,out1, x,y,c0,c1) \ - __m128i c0##lo = _mm_unpacklo_epi16((x),(y)); \ - __m128i c0##hi = _mm_unpackhi_epi16((x),(y)); \ - __m128i out0##_l = _mm_madd_epi16(c0##lo, c0); \ - __m128i out0##_h = _mm_madd_epi16(c0##hi, c0); \ - __m128i out1##_l = _mm_madd_epi16(c0##lo, c1); \ - __m128i out1##_h = _mm_madd_epi16(c0##hi, c1) - - // out = in << 12 (in 16-bit, out 32-bit) - #define dct_widen(out, in) \ - __m128i out##_l = _mm_srai_epi32(_mm_unpacklo_epi16(_mm_setzero_si128(), (in)), 4); \ - __m128i out##_h = _mm_srai_epi32(_mm_unpackhi_epi16(_mm_setzero_si128(), (in)), 4) - - // wide add - #define dct_wadd(out, a, b) \ - __m128i out##_l = _mm_add_epi32(a##_l, b##_l); \ - __m128i out##_h = _mm_add_epi32(a##_h, b##_h) - - // wide sub - #define dct_wsub(out, a, b) \ - __m128i out##_l = _mm_sub_epi32(a##_l, b##_l); \ - __m128i out##_h = _mm_sub_epi32(a##_h, b##_h) - - // butterfly a/b, add bias, then shift by "s" and pack - #define dct_bfly32o(out0, out1, a,b,bias,s) \ - { \ - __m128i abiased_l = _mm_add_epi32(a##_l, bias); \ - __m128i abiased_h = _mm_add_epi32(a##_h, bias); \ - dct_wadd(sum, abiased, b); \ - dct_wsub(dif, abiased, b); \ - out0 = _mm_packs_epi32(_mm_srai_epi32(sum_l, s), _mm_srai_epi32(sum_h, s)); \ - out1 = _mm_packs_epi32(_mm_srai_epi32(dif_l, s), _mm_srai_epi32(dif_h, s)); \ - } - - // 8-bit interleave step (for transposes) - #define dct_interleave8(a, b) \ - tmp = a; \ - a = _mm_unpacklo_epi8(a, b); \ - b = _mm_unpackhi_epi8(tmp, b) - - // 16-bit interleave step (for transposes) - #define dct_interleave16(a, b) \ - tmp = a; \ - a = _mm_unpacklo_epi16(a, b); \ - b = _mm_unpackhi_epi16(tmp, b) - - #define dct_pass(bias,shift) \ - { \ - /* even part */ \ - dct_rot(t2e,t3e, row2,row6, rot0_0,rot0_1); \ - __m128i sum04 = _mm_add_epi16(row0, row4); \ - __m128i dif04 = _mm_sub_epi16(row0, row4); \ - dct_widen(t0e, sum04); \ - dct_widen(t1e, dif04); \ - dct_wadd(x0, t0e, t3e); \ - dct_wsub(x3, t0e, t3e); \ - dct_wadd(x1, t1e, t2e); \ - dct_wsub(x2, t1e, t2e); \ - /* odd part */ \ - dct_rot(y0o,y2o, row7,row3, rot2_0,rot2_1); \ - dct_rot(y1o,y3o, row5,row1, rot3_0,rot3_1); \ - __m128i sum17 = _mm_add_epi16(row1, row7); \ - __m128i sum35 = _mm_add_epi16(row3, row5); \ - dct_rot(y4o,y5o, sum17,sum35, rot1_0,rot1_1); \ - dct_wadd(x4, y0o, y4o); \ - dct_wadd(x5, y1o, y5o); \ - dct_wadd(x6, y2o, y5o); \ - dct_wadd(x7, y3o, y4o); \ - dct_bfly32o(row0,row7, x0,x7,bias,shift); \ - dct_bfly32o(row1,row6, x1,x6,bias,shift); \ - dct_bfly32o(row2,row5, x2,x5,bias,shift); \ - dct_bfly32o(row3,row4, x3,x4,bias,shift); \ - } - - __m128i rot0_0 = dct_const(stbi__f2f(0.5411961f), stbi__f2f(0.5411961f) + stbi__f2f(-1.847759065f)); - __m128i rot0_1 = dct_const(stbi__f2f(0.5411961f) + stbi__f2f( 0.765366865f), stbi__f2f(0.5411961f)); - __m128i rot1_0 = dct_const(stbi__f2f(1.175875602f) + stbi__f2f(-0.899976223f), stbi__f2f(1.175875602f)); - __m128i rot1_1 = dct_const(stbi__f2f(1.175875602f), stbi__f2f(1.175875602f) + stbi__f2f(-2.562915447f)); - __m128i rot2_0 = dct_const(stbi__f2f(-1.961570560f) + stbi__f2f( 0.298631336f), stbi__f2f(-1.961570560f)); - __m128i rot2_1 = dct_const(stbi__f2f(-1.961570560f), stbi__f2f(-1.961570560f) + stbi__f2f( 3.072711026f)); - __m128i rot3_0 = dct_const(stbi__f2f(-0.390180644f) + stbi__f2f( 2.053119869f), stbi__f2f(-0.390180644f)); - __m128i rot3_1 = dct_const(stbi__f2f(-0.390180644f), stbi__f2f(-0.390180644f) + stbi__f2f( 1.501321110f)); - - // rounding biases in column/row passes, see stbi__idct_block for explanation. - __m128i bias_0 = _mm_set1_epi32(512); - __m128i bias_1 = _mm_set1_epi32(65536 + (128<<17)); - - // load - row0 = _mm_load_si128((const __m128i *) (data + 0*8)); - row1 = _mm_load_si128((const __m128i *) (data + 1*8)); - row2 = _mm_load_si128((const __m128i *) (data + 2*8)); - row3 = _mm_load_si128((const __m128i *) (data + 3*8)); - row4 = _mm_load_si128((const __m128i *) (data + 4*8)); - row5 = _mm_load_si128((const __m128i *) (data + 5*8)); - row6 = _mm_load_si128((const __m128i *) (data + 6*8)); - row7 = _mm_load_si128((const __m128i *) (data + 7*8)); - - // column pass - dct_pass(bias_0, 10); - - { - // 16bit 8x8 transpose pass 1 - dct_interleave16(row0, row4); - dct_interleave16(row1, row5); - dct_interleave16(row2, row6); - dct_interleave16(row3, row7); - - // transpose pass 2 - dct_interleave16(row0, row2); - dct_interleave16(row1, row3); - dct_interleave16(row4, row6); - dct_interleave16(row5, row7); - - // transpose pass 3 - dct_interleave16(row0, row1); - dct_interleave16(row2, row3); - dct_interleave16(row4, row5); - dct_interleave16(row6, row7); - } - - // row pass - dct_pass(bias_1, 17); - - { - // pack - __m128i p0 = _mm_packus_epi16(row0, row1); // a0a1a2a3...a7b0b1b2b3...b7 - __m128i p1 = _mm_packus_epi16(row2, row3); - __m128i p2 = _mm_packus_epi16(row4, row5); - __m128i p3 = _mm_packus_epi16(row6, row7); - - // 8bit 8x8 transpose pass 1 - dct_interleave8(p0, p2); // a0e0a1e1... - dct_interleave8(p1, p3); // c0g0c1g1... - - // transpose pass 2 - dct_interleave8(p0, p1); // a0c0e0g0... - dct_interleave8(p2, p3); // b0d0f0h0... - - // transpose pass 3 - dct_interleave8(p0, p2); // a0b0c0d0... - dct_interleave8(p1, p3); // a4b4c4d4... - - // store - _mm_storel_epi64((__m128i *) out, p0); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p0, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p2); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p2, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p1); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p1, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p3); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p3, 0x4e)); - } - -#undef dct_const -#undef dct_rot -#undef dct_widen -#undef dct_wadd -#undef dct_wsub -#undef dct_bfly32o -#undef dct_interleave8 -#undef dct_interleave16 -#undef dct_pass -} - -#endif // STBI_SSE2 - -#ifdef STBI_NEON - -// NEON integer IDCT. should produce bit-identical -// results to the generic C version. -static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) -{ - int16x8_t row0, row1, row2, row3, row4, row5, row6, row7; - - int16x4_t rot0_0 = vdup_n_s16(stbi__f2f(0.5411961f)); - int16x4_t rot0_1 = vdup_n_s16(stbi__f2f(-1.847759065f)); - int16x4_t rot0_2 = vdup_n_s16(stbi__f2f( 0.765366865f)); - int16x4_t rot1_0 = vdup_n_s16(stbi__f2f( 1.175875602f)); - int16x4_t rot1_1 = vdup_n_s16(stbi__f2f(-0.899976223f)); - int16x4_t rot1_2 = vdup_n_s16(stbi__f2f(-2.562915447f)); - int16x4_t rot2_0 = vdup_n_s16(stbi__f2f(-1.961570560f)); - int16x4_t rot2_1 = vdup_n_s16(stbi__f2f(-0.390180644f)); - int16x4_t rot3_0 = vdup_n_s16(stbi__f2f( 0.298631336f)); - int16x4_t rot3_1 = vdup_n_s16(stbi__f2f( 2.053119869f)); - int16x4_t rot3_2 = vdup_n_s16(stbi__f2f( 3.072711026f)); - int16x4_t rot3_3 = vdup_n_s16(stbi__f2f( 1.501321110f)); - -#define dct_long_mul(out, inq, coeff) \ - int32x4_t out##_l = vmull_s16(vget_low_s16(inq), coeff); \ - int32x4_t out##_h = vmull_s16(vget_high_s16(inq), coeff) - -#define dct_long_mac(out, acc, inq, coeff) \ - int32x4_t out##_l = vmlal_s16(acc##_l, vget_low_s16(inq), coeff); \ - int32x4_t out##_h = vmlal_s16(acc##_h, vget_high_s16(inq), coeff) - -#define dct_widen(out, inq) \ - int32x4_t out##_l = vshll_n_s16(vget_low_s16(inq), 12); \ - int32x4_t out##_h = vshll_n_s16(vget_high_s16(inq), 12) - -// wide add -#define dct_wadd(out, a, b) \ - int32x4_t out##_l = vaddq_s32(a##_l, b##_l); \ - int32x4_t out##_h = vaddq_s32(a##_h, b##_h) - -// wide sub -#define dct_wsub(out, a, b) \ - int32x4_t out##_l = vsubq_s32(a##_l, b##_l); \ - int32x4_t out##_h = vsubq_s32(a##_h, b##_h) - -// butterfly a/b, then shift using "shiftop" by "s" and pack -#define dct_bfly32o(out0,out1, a,b,shiftop,s) \ - { \ - dct_wadd(sum, a, b); \ - dct_wsub(dif, a, b); \ - out0 = vcombine_s16(shiftop(sum_l, s), shiftop(sum_h, s)); \ - out1 = vcombine_s16(shiftop(dif_l, s), shiftop(dif_h, s)); \ - } - -#define dct_pass(shiftop, shift) \ - { \ - /* even part */ \ - int16x8_t sum26 = vaddq_s16(row2, row6); \ - dct_long_mul(p1e, sum26, rot0_0); \ - dct_long_mac(t2e, p1e, row6, rot0_1); \ - dct_long_mac(t3e, p1e, row2, rot0_2); \ - int16x8_t sum04 = vaddq_s16(row0, row4); \ - int16x8_t dif04 = vsubq_s16(row0, row4); \ - dct_widen(t0e, sum04); \ - dct_widen(t1e, dif04); \ - dct_wadd(x0, t0e, t3e); \ - dct_wsub(x3, t0e, t3e); \ - dct_wadd(x1, t1e, t2e); \ - dct_wsub(x2, t1e, t2e); \ - /* odd part */ \ - int16x8_t sum15 = vaddq_s16(row1, row5); \ - int16x8_t sum17 = vaddq_s16(row1, row7); \ - int16x8_t sum35 = vaddq_s16(row3, row5); \ - int16x8_t sum37 = vaddq_s16(row3, row7); \ - int16x8_t sumodd = vaddq_s16(sum17, sum35); \ - dct_long_mul(p5o, sumodd, rot1_0); \ - dct_long_mac(p1o, p5o, sum17, rot1_1); \ - dct_long_mac(p2o, p5o, sum35, rot1_2); \ - dct_long_mul(p3o, sum37, rot2_0); \ - dct_long_mul(p4o, sum15, rot2_1); \ - dct_wadd(sump13o, p1o, p3o); \ - dct_wadd(sump24o, p2o, p4o); \ - dct_wadd(sump23o, p2o, p3o); \ - dct_wadd(sump14o, p1o, p4o); \ - dct_long_mac(x4, sump13o, row7, rot3_0); \ - dct_long_mac(x5, sump24o, row5, rot3_1); \ - dct_long_mac(x6, sump23o, row3, rot3_2); \ - dct_long_mac(x7, sump14o, row1, rot3_3); \ - dct_bfly32o(row0,row7, x0,x7,shiftop,shift); \ - dct_bfly32o(row1,row6, x1,x6,shiftop,shift); \ - dct_bfly32o(row2,row5, x2,x5,shiftop,shift); \ - dct_bfly32o(row3,row4, x3,x4,shiftop,shift); \ - } - - // load - row0 = vld1q_s16(data + 0*8); - row1 = vld1q_s16(data + 1*8); - row2 = vld1q_s16(data + 2*8); - row3 = vld1q_s16(data + 3*8); - row4 = vld1q_s16(data + 4*8); - row5 = vld1q_s16(data + 5*8); - row6 = vld1q_s16(data + 6*8); - row7 = vld1q_s16(data + 7*8); - - // add DC bias - row0 = vaddq_s16(row0, vsetq_lane_s16(1024, vdupq_n_s16(0), 0)); - - // column pass - dct_pass(vrshrn_n_s32, 10); - - // 16bit 8x8 transpose - { -// these three map to a single VTRN.16, VTRN.32, and VSWP, respectively. -// whether compilers actually get this is another story, sadly. -#define dct_trn16(x, y) { int16x8x2_t t = vtrnq_s16(x, y); x = t.val[0]; y = t.val[1]; } -#define dct_trn32(x, y) { int32x4x2_t t = vtrnq_s32(vreinterpretq_s32_s16(x), vreinterpretq_s32_s16(y)); x = vreinterpretq_s16_s32(t.val[0]); y = vreinterpretq_s16_s32(t.val[1]); } -#define dct_trn64(x, y) { int16x8_t x0 = x; int16x8_t y0 = y; x = vcombine_s16(vget_low_s16(x0), vget_low_s16(y0)); y = vcombine_s16(vget_high_s16(x0), vget_high_s16(y0)); } - - // pass 1 - dct_trn16(row0, row1); // a0b0a2b2a4b4a6b6 - dct_trn16(row2, row3); - dct_trn16(row4, row5); - dct_trn16(row6, row7); - - // pass 2 - dct_trn32(row0, row2); // a0b0c0d0a4b4c4d4 - dct_trn32(row1, row3); - dct_trn32(row4, row6); - dct_trn32(row5, row7); - - // pass 3 - dct_trn64(row0, row4); // a0b0c0d0e0f0g0h0 - dct_trn64(row1, row5); - dct_trn64(row2, row6); - dct_trn64(row3, row7); - -#undef dct_trn16 -#undef dct_trn32 -#undef dct_trn64 - } - - // row pass - // vrshrn_n_s32 only supports shifts up to 16, we need - // 17. so do a non-rounding shift of 16 first then follow - // up with a rounding shift by 1. - dct_pass(vshrn_n_s32, 16); - - { - // pack and round - uint8x8_t p0 = vqrshrun_n_s16(row0, 1); - uint8x8_t p1 = vqrshrun_n_s16(row1, 1); - uint8x8_t p2 = vqrshrun_n_s16(row2, 1); - uint8x8_t p3 = vqrshrun_n_s16(row3, 1); - uint8x8_t p4 = vqrshrun_n_s16(row4, 1); - uint8x8_t p5 = vqrshrun_n_s16(row5, 1); - uint8x8_t p6 = vqrshrun_n_s16(row6, 1); - uint8x8_t p7 = vqrshrun_n_s16(row7, 1); - - // again, these can translate into one instruction, but often don't. -#define dct_trn8_8(x, y) { uint8x8x2_t t = vtrn_u8(x, y); x = t.val[0]; y = t.val[1]; } -#define dct_trn8_16(x, y) { uint16x4x2_t t = vtrn_u16(vreinterpret_u16_u8(x), vreinterpret_u16_u8(y)); x = vreinterpret_u8_u16(t.val[0]); y = vreinterpret_u8_u16(t.val[1]); } -#define dct_trn8_32(x, y) { uint32x2x2_t t = vtrn_u32(vreinterpret_u32_u8(x), vreinterpret_u32_u8(y)); x = vreinterpret_u8_u32(t.val[0]); y = vreinterpret_u8_u32(t.val[1]); } - - // sadly can't use interleaved stores here since we only write - // 8 bytes to each scan line! - - // 8x8 8-bit transpose pass 1 - dct_trn8_8(p0, p1); - dct_trn8_8(p2, p3); - dct_trn8_8(p4, p5); - dct_trn8_8(p6, p7); - - // pass 2 - dct_trn8_16(p0, p2); - dct_trn8_16(p1, p3); - dct_trn8_16(p4, p6); - dct_trn8_16(p5, p7); - - // pass 3 - dct_trn8_32(p0, p4); - dct_trn8_32(p1, p5); - dct_trn8_32(p2, p6); - dct_trn8_32(p3, p7); - - // store - vst1_u8(out, p0); out += out_stride; - vst1_u8(out, p1); out += out_stride; - vst1_u8(out, p2); out += out_stride; - vst1_u8(out, p3); out += out_stride; - vst1_u8(out, p4); out += out_stride; - vst1_u8(out, p5); out += out_stride; - vst1_u8(out, p6); out += out_stride; - vst1_u8(out, p7); - -#undef dct_trn8_8 -#undef dct_trn8_16 -#undef dct_trn8_32 - } - -#undef dct_long_mul -#undef dct_long_mac -#undef dct_widen -#undef dct_wadd -#undef dct_wsub -#undef dct_bfly32o -#undef dct_pass -} - -#endif // STBI_NEON - -#define STBI__MARKER_none 0xff -// if there's a pending marker from the entropy stream, return that -// otherwise, fetch from the stream and get a marker. if there's no -// marker, return 0xff, which is never a valid marker value -static stbi_uc stbi__get_marker(stbi__jpeg *j) -{ - stbi_uc x; - if (j->marker != STBI__MARKER_none) { x = j->marker; j->marker = STBI__MARKER_none; return x; } - x = stbi__get8(j->s); - if (x != 0xff) return STBI__MARKER_none; - while (x == 0xff) - x = stbi__get8(j->s); // consume repeated 0xff fill bytes - return x; -} - -// in each scan, we'll have scan_n components, and the order -// of the components is specified by order[] -#define STBI__RESTART(x) ((x) >= 0xd0 && (x) <= 0xd7) - -// after a restart interval, stbi__jpeg_reset the entropy decoder and -// the dc prediction -static void stbi__jpeg_reset(stbi__jpeg *j) -{ - j->code_bits = 0; - j->code_buffer = 0; - j->nomore = 0; - j->img_comp[0].dc_pred = j->img_comp[1].dc_pred = j->img_comp[2].dc_pred = j->img_comp[3].dc_pred = 0; - j->marker = STBI__MARKER_none; - j->todo = j->restart_interval ? j->restart_interval : 0x7fffffff; - j->eob_run = 0; - // no more than 1<<31 MCUs if no restart_interal? that's plenty safe, - // since we don't even allow 1<<30 pixels -} - -static int stbi__parse_entropy_coded_data(stbi__jpeg *z) -{ - stbi__jpeg_reset(z); - if (!z->progressive) { - if (z->scan_n == 1) { - int i,j; - STBI_SIMD_ALIGN(short, data[64]); - int n = z->order[0]; - // non-interleaved data, we just need to process one block at a time, - // in trivial scanline order - // number of blocks to do just depends on how many actual "pixels" this - // component has, independent of interleaved MCU blocking and such - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); - // every data block is an MCU, so countdown the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - // if it's NOT a restart, then just bail, so we get corrupt data - // rather than no data - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } else { // interleaved - int i,j,k,x,y; - STBI_SIMD_ALIGN(short, data[64]); - for (j=0; j < z->img_mcu_y; ++j) { - for (i=0; i < z->img_mcu_x; ++i) { - // scan an interleaved mcu... process scan_n components in order - for (k=0; k < z->scan_n; ++k) { - int n = z->order[k]; - // scan out an mcu's worth of this component; that's just determined - // by the basic H and V specified for the component - for (y=0; y < z->img_comp[n].v; ++y) { - for (x=0; x < z->img_comp[n].h; ++x) { - int x2 = (i*z->img_comp[n].h + x)*8; - int y2 = (j*z->img_comp[n].v + y)*8; - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*y2+x2, z->img_comp[n].w2, data); - } - } - } - // after all interleaved components, that's an interleaved MCU, - // so now count down the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } - } else { - if (z->scan_n == 1) { - int i,j; - int n = z->order[0]; - // non-interleaved data, we just need to process one block at a time, - // in trivial scanline order - // number of blocks to do just depends on how many actual "pixels" this - // component has, independent of interleaved MCU blocking and such - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); - if (z->spec_start == 0) { - if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) - return 0; - } else { - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block_prog_ac(z, data, &z->huff_ac[ha], z->fast_ac[ha])) - return 0; - } - // every data block is an MCU, so countdown the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } else { // interleaved - int i,j,k,x,y; - for (j=0; j < z->img_mcu_y; ++j) { - for (i=0; i < z->img_mcu_x; ++i) { - // scan an interleaved mcu... process scan_n components in order - for (k=0; k < z->scan_n; ++k) { - int n = z->order[k]; - // scan out an mcu's worth of this component; that's just determined - // by the basic H and V specified for the component - for (y=0; y < z->img_comp[n].v; ++y) { - for (x=0; x < z->img_comp[n].h; ++x) { - int x2 = (i*z->img_comp[n].h + x); - int y2 = (j*z->img_comp[n].v + y); - short *data = z->img_comp[n].coeff + 64 * (x2 + y2 * z->img_comp[n].coeff_w); - if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) - return 0; - } - } - } - // after all interleaved components, that's an interleaved MCU, - // so now count down the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } - } -} - -static void stbi__jpeg_dequantize(short *data, stbi__uint16 *dequant) -{ - int i; - for (i=0; i < 64; ++i) - data[i] *= dequant[i]; -} - -static void stbi__jpeg_finish(stbi__jpeg *z) -{ - if (z->progressive) { - // dequantize and idct the data - int i,j,n; - for (n=0; n < z->s->img_n; ++n) { - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); - stbi__jpeg_dequantize(data, z->dequant[z->img_comp[n].tq]); - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); - } - } - } - } -} - -static int stbi__process_marker(stbi__jpeg *z, int m) -{ - int L; - switch (m) { - case STBI__MARKER_none: // no marker found - return stbi__err("expected marker","Corrupt JPEG"); - - case 0xDD: // DRI - specify restart interval - if (stbi__get16be(z->s) != 4) return stbi__err("bad DRI len","Corrupt JPEG"); - z->restart_interval = stbi__get16be(z->s); - return 1; - - case 0xDB: // DQT - define quantization table - L = stbi__get16be(z->s)-2; - while (L > 0) { - int q = stbi__get8(z->s); - int p = q >> 4, sixteen = (p != 0); - int t = q & 15,i; - if (p != 0 && p != 1) return stbi__err("bad DQT type","Corrupt JPEG"); - if (t > 3) return stbi__err("bad DQT table","Corrupt JPEG"); - - for (i=0; i < 64; ++i) - z->dequant[t][stbi__jpeg_dezigzag[i]] = (stbi__uint16)(sixteen ? stbi__get16be(z->s) : stbi__get8(z->s)); - L -= (sixteen ? 129 : 65); - } - return L==0; - - case 0xC4: // DHT - define huffman table - L = stbi__get16be(z->s)-2; - while (L > 0) { - stbi_uc *v; - int sizes[16],i,n=0; - int q = stbi__get8(z->s); - int tc = q >> 4; - int th = q & 15; - if (tc > 1 || th > 3) return stbi__err("bad DHT header","Corrupt JPEG"); - for (i=0; i < 16; ++i) { - sizes[i] = stbi__get8(z->s); - n += sizes[i]; - } - if(n > 256) return stbi__err("bad DHT header","Corrupt JPEG"); // Loop over i < n would write past end of values! - L -= 17; - if (tc == 0) { - if (!stbi__build_huffman(z->huff_dc+th, sizes)) return 0; - v = z->huff_dc[th].values; - } else { - if (!stbi__build_huffman(z->huff_ac+th, sizes)) return 0; - v = z->huff_ac[th].values; - } - for (i=0; i < n; ++i) - v[i] = stbi__get8(z->s); - if (tc != 0) - stbi__build_fast_ac(z->fast_ac[th], z->huff_ac + th); - L -= n; - } - return L==0; - } - - // check for comment block or APP blocks - if ((m >= 0xE0 && m <= 0xEF) || m == 0xFE) { - L = stbi__get16be(z->s); - if (L < 2) { - if (m == 0xFE) - return stbi__err("bad COM len","Corrupt JPEG"); - else - return stbi__err("bad APP len","Corrupt JPEG"); - } - L -= 2; - - if (m == 0xE0 && L >= 5) { // JFIF APP0 segment - static const unsigned char tag[5] = {'J','F','I','F','\0'}; - int ok = 1; - int i; - for (i=0; i < 5; ++i) - if (stbi__get8(z->s) != tag[i]) - ok = 0; - L -= 5; - if (ok) - z->jfif = 1; - } else if (m == 0xEE && L >= 12) { // Adobe APP14 segment - static const unsigned char tag[6] = {'A','d','o','b','e','\0'}; - int ok = 1; - int i; - for (i=0; i < 6; ++i) - if (stbi__get8(z->s) != tag[i]) - ok = 0; - L -= 6; - if (ok) { - stbi__get8(z->s); // version - stbi__get16be(z->s); // flags0 - stbi__get16be(z->s); // flags1 - z->app14_color_transform = stbi__get8(z->s); // color transform - L -= 6; - } - } - - stbi__skip(z->s, L); - return 1; - } - - return stbi__err("unknown marker","Corrupt JPEG"); -} - -// after we see SOS -static int stbi__process_scan_header(stbi__jpeg *z) -{ - int i; - int Ls = stbi__get16be(z->s); - z->scan_n = stbi__get8(z->s); - if (z->scan_n < 1 || z->scan_n > 4 || z->scan_n > (int) z->s->img_n) return stbi__err("bad SOS component count","Corrupt JPEG"); - if (Ls != 6+2*z->scan_n) return stbi__err("bad SOS len","Corrupt JPEG"); - for (i=0; i < z->scan_n; ++i) { - int id = stbi__get8(z->s), which; - int q = stbi__get8(z->s); - for (which = 0; which < z->s->img_n; ++which) - if (z->img_comp[which].id == id) - break; - if (which == z->s->img_n) return 0; // no match - z->img_comp[which].hd = q >> 4; if (z->img_comp[which].hd > 3) return stbi__err("bad DC huff","Corrupt JPEG"); - z->img_comp[which].ha = q & 15; if (z->img_comp[which].ha > 3) return stbi__err("bad AC huff","Corrupt JPEG"); - z->order[i] = which; - } - - { - int aa; - z->spec_start = stbi__get8(z->s); - z->spec_end = stbi__get8(z->s); // should be 63, but might be 0 - aa = stbi__get8(z->s); - z->succ_high = (aa >> 4); - z->succ_low = (aa & 15); - if (z->progressive) { - if (z->spec_start > 63 || z->spec_end > 63 || z->spec_start > z->spec_end || z->succ_high > 13 || z->succ_low > 13) - return stbi__err("bad SOS", "Corrupt JPEG"); - } else { - if (z->spec_start != 0) return stbi__err("bad SOS","Corrupt JPEG"); - if (z->succ_high != 0 || z->succ_low != 0) return stbi__err("bad SOS","Corrupt JPEG"); - z->spec_end = 63; - } - } - - return 1; -} - -static int stbi__free_jpeg_components(stbi__jpeg *z, int ncomp, int why) -{ - int i; - for (i=0; i < ncomp; ++i) { - if (z->img_comp[i].raw_data) { - STBI_FREE(z->img_comp[i].raw_data); - z->img_comp[i].raw_data = NULL; - z->img_comp[i].data = NULL; - } - if (z->img_comp[i].raw_coeff) { - STBI_FREE(z->img_comp[i].raw_coeff); - z->img_comp[i].raw_coeff = 0; - z->img_comp[i].coeff = 0; - } - if (z->img_comp[i].linebuf) { - STBI_FREE(z->img_comp[i].linebuf); - z->img_comp[i].linebuf = NULL; - } - } - return why; -} - -static int stbi__process_frame_header(stbi__jpeg *z, int scan) -{ - stbi__context *s = z->s; - int Lf,p,i,q, h_max=1,v_max=1,c; - Lf = stbi__get16be(s); if (Lf < 11) return stbi__err("bad SOF len","Corrupt JPEG"); // JPEG - p = stbi__get8(s); if (p != 8) return stbi__err("only 8-bit","JPEG format not supported: 8-bit only"); // JPEG baseline - s->img_y = stbi__get16be(s); if (s->img_y == 0) return stbi__err("no header height", "JPEG format not supported: delayed height"); // Legal, but we don't handle it--but neither does IJG - s->img_x = stbi__get16be(s); if (s->img_x == 0) return stbi__err("0 width","Corrupt JPEG"); // JPEG requires - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - c = stbi__get8(s); - if (c != 3 && c != 1 && c != 4) return stbi__err("bad component count","Corrupt JPEG"); - s->img_n = c; - for (i=0; i < c; ++i) { - z->img_comp[i].data = NULL; - z->img_comp[i].linebuf = NULL; - } - - if (Lf != 8+3*s->img_n) return stbi__err("bad SOF len","Corrupt JPEG"); - - z->rgb = 0; - for (i=0; i < s->img_n; ++i) { - static const unsigned char rgb[3] = { 'R', 'G', 'B' }; - z->img_comp[i].id = stbi__get8(s); - if (s->img_n == 3 && z->img_comp[i].id == rgb[i]) - ++z->rgb; - q = stbi__get8(s); - z->img_comp[i].h = (q >> 4); if (!z->img_comp[i].h || z->img_comp[i].h > 4) return stbi__err("bad H","Corrupt JPEG"); - z->img_comp[i].v = q & 15; if (!z->img_comp[i].v || z->img_comp[i].v > 4) return stbi__err("bad V","Corrupt JPEG"); - z->img_comp[i].tq = stbi__get8(s); if (z->img_comp[i].tq > 3) return stbi__err("bad TQ","Corrupt JPEG"); - } - - if (scan != STBI__SCAN_load) return 1; - - if (!stbi__mad3sizes_valid(s->img_x, s->img_y, s->img_n, 0)) return stbi__err("too large", "Image too large to decode"); - - for (i=0; i < s->img_n; ++i) { - if (z->img_comp[i].h > h_max) h_max = z->img_comp[i].h; - if (z->img_comp[i].v > v_max) v_max = z->img_comp[i].v; - } - - // check that plane subsampling factors are integer ratios; our resamplers can't deal with fractional ratios - // and I've never seen a non-corrupted JPEG file actually use them - for (i=0; i < s->img_n; ++i) { - if (h_max % z->img_comp[i].h != 0) return stbi__err("bad H","Corrupt JPEG"); - if (v_max % z->img_comp[i].v != 0) return stbi__err("bad V","Corrupt JPEG"); - } - - // compute interleaved mcu info - z->img_h_max = h_max; - z->img_v_max = v_max; - z->img_mcu_w = h_max * 8; - z->img_mcu_h = v_max * 8; - // these sizes can't be more than 17 bits - z->img_mcu_x = (s->img_x + z->img_mcu_w-1) / z->img_mcu_w; - z->img_mcu_y = (s->img_y + z->img_mcu_h-1) / z->img_mcu_h; - - for (i=0; i < s->img_n; ++i) { - // number of effective pixels (e.g. for non-interleaved MCU) - z->img_comp[i].x = (s->img_x * z->img_comp[i].h + h_max-1) / h_max; - z->img_comp[i].y = (s->img_y * z->img_comp[i].v + v_max-1) / v_max; - // to simplify generation, we'll allocate enough memory to decode - // the bogus oversized data from using interleaved MCUs and their - // big blocks (e.g. a 16x16 iMCU on an image of width 33); we won't - // discard the extra data until colorspace conversion - // - // img_mcu_x, img_mcu_y: <=17 bits; comp[i].h and .v are <=4 (checked earlier) - // so these muls can't overflow with 32-bit ints (which we require) - z->img_comp[i].w2 = z->img_mcu_x * z->img_comp[i].h * 8; - z->img_comp[i].h2 = z->img_mcu_y * z->img_comp[i].v * 8; - z->img_comp[i].coeff = 0; - z->img_comp[i].raw_coeff = 0; - z->img_comp[i].linebuf = NULL; - z->img_comp[i].raw_data = stbi__malloc_mad2(z->img_comp[i].w2, z->img_comp[i].h2, 15); - if (z->img_comp[i].raw_data == NULL) - return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); - // align blocks for idct using mmx/sse - z->img_comp[i].data = (stbi_uc*) (((size_t) z->img_comp[i].raw_data + 15) & ~15); - if (z->progressive) { - // w2, h2 are multiples of 8 (see above) - z->img_comp[i].coeff_w = z->img_comp[i].w2 / 8; - z->img_comp[i].coeff_h = z->img_comp[i].h2 / 8; - z->img_comp[i].raw_coeff = stbi__malloc_mad3(z->img_comp[i].w2, z->img_comp[i].h2, sizeof(short), 15); - if (z->img_comp[i].raw_coeff == NULL) - return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); - z->img_comp[i].coeff = (short*) (((size_t) z->img_comp[i].raw_coeff + 15) & ~15); - } - } - - return 1; -} - -// use comparisons since in some cases we handle more than one case (e.g. SOF) -#define stbi__DNL(x) ((x) == 0xdc) -#define stbi__SOI(x) ((x) == 0xd8) -#define stbi__EOI(x) ((x) == 0xd9) -#define stbi__SOF(x) ((x) == 0xc0 || (x) == 0xc1 || (x) == 0xc2) -#define stbi__SOS(x) ((x) == 0xda) - -#define stbi__SOF_progressive(x) ((x) == 0xc2) - -static int stbi__decode_jpeg_header(stbi__jpeg *z, int scan) -{ - int m; - z->jfif = 0; - z->app14_color_transform = -1; // valid values are 0,1,2 - z->marker = STBI__MARKER_none; // initialize cached marker to empty - m = stbi__get_marker(z); - if (!stbi__SOI(m)) return stbi__err("no SOI","Corrupt JPEG"); - if (scan == STBI__SCAN_type) return 1; - m = stbi__get_marker(z); - while (!stbi__SOF(m)) { - if (!stbi__process_marker(z,m)) return 0; - m = stbi__get_marker(z); - while (m == STBI__MARKER_none) { - // some files have extra padding after their blocks, so ok, we'll scan - if (stbi__at_eof(z->s)) return stbi__err("no SOF", "Corrupt JPEG"); - m = stbi__get_marker(z); - } - } - z->progressive = stbi__SOF_progressive(m); - if (!stbi__process_frame_header(z, scan)) return 0; - return 1; -} - -static stbi_uc stbi__skip_jpeg_junk_at_end(stbi__jpeg *j) -{ - // some JPEGs have junk at end, skip over it but if we find what looks - // like a valid marker, resume there - while (!stbi__at_eof(j->s)) { - stbi_uc x = stbi__get8(j->s); - while (x == 0xff) { // might be a marker - if (stbi__at_eof(j->s)) return STBI__MARKER_none; - x = stbi__get8(j->s); - if (x != 0x00 && x != 0xff) { - // not a stuffed zero or lead-in to another marker, looks - // like an actual marker, return it - return x; - } - // stuffed zero has x=0 now which ends the loop, meaning we go - // back to regular scan loop. - // repeated 0xff keeps trying to read the next byte of the marker. - } - } - return STBI__MARKER_none; -} - -// decode image to YCbCr format -static int stbi__decode_jpeg_image(stbi__jpeg *j) -{ - int m; - for (m = 0; m < 4; m++) { - j->img_comp[m].raw_data = NULL; - j->img_comp[m].raw_coeff = NULL; - } - j->restart_interval = 0; - if (!stbi__decode_jpeg_header(j, STBI__SCAN_load)) return 0; - m = stbi__get_marker(j); - while (!stbi__EOI(m)) { - if (stbi__SOS(m)) { - if (!stbi__process_scan_header(j)) return 0; - if (!stbi__parse_entropy_coded_data(j)) return 0; - if (j->marker == STBI__MARKER_none ) { - j->marker = stbi__skip_jpeg_junk_at_end(j); - // if we reach eof without hitting a marker, stbi__get_marker() below will fail and we'll eventually return 0 - } - m = stbi__get_marker(j); - if (STBI__RESTART(m)) - m = stbi__get_marker(j); - } else if (stbi__DNL(m)) { - int Ld = stbi__get16be(j->s); - stbi__uint32 NL = stbi__get16be(j->s); - if (Ld != 4) return stbi__err("bad DNL len", "Corrupt JPEG"); - if (NL != j->s->img_y) return stbi__err("bad DNL height", "Corrupt JPEG"); - m = stbi__get_marker(j); - } else { - if (!stbi__process_marker(j, m)) return 1; - m = stbi__get_marker(j); - } - } - if (j->progressive) - stbi__jpeg_finish(j); - return 1; -} - -// static jfif-centered resampling (across block boundaries) - -typedef stbi_uc *(*resample_row_func)(stbi_uc *out, stbi_uc *in0, stbi_uc *in1, - int w, int hs); - -#define stbi__div4(x) ((stbi_uc) ((x) >> 2)) - -static stbi_uc *resample_row_1(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - STBI_NOTUSED(out); - STBI_NOTUSED(in_far); - STBI_NOTUSED(w); - STBI_NOTUSED(hs); - return in_near; -} - -static stbi_uc* stbi__resample_row_v_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate two samples vertically for every one in input - int i; - STBI_NOTUSED(hs); - for (i=0; i < w; ++i) - out[i] = stbi__div4(3*in_near[i] + in_far[i] + 2); - return out; -} - -static stbi_uc* stbi__resample_row_h_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate two samples horizontally for every one in input - int i; - stbi_uc *input = in_near; - - if (w == 1) { - // if only one sample, can't do any interpolation - out[0] = out[1] = input[0]; - return out; - } - - out[0] = input[0]; - out[1] = stbi__div4(input[0]*3 + input[1] + 2); - for (i=1; i < w-1; ++i) { - int n = 3*input[i]+2; - out[i*2+0] = stbi__div4(n+input[i-1]); - out[i*2+1] = stbi__div4(n+input[i+1]); - } - out[i*2+0] = stbi__div4(input[w-2]*3 + input[w-1] + 2); - out[i*2+1] = input[w-1]; - - STBI_NOTUSED(in_far); - STBI_NOTUSED(hs); - - return out; -} - -#define stbi__div16(x) ((stbi_uc) ((x) >> 4)) - -static stbi_uc *stbi__resample_row_hv_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate 2x2 samples for every one in input - int i,t0,t1; - if (w == 1) { - out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); - return out; - } - - t1 = 3*in_near[0] + in_far[0]; - out[0] = stbi__div4(t1+2); - for (i=1; i < w; ++i) { - t0 = t1; - t1 = 3*in_near[i]+in_far[i]; - out[i*2-1] = stbi__div16(3*t0 + t1 + 8); - out[i*2 ] = stbi__div16(3*t1 + t0 + 8); - } - out[w*2-1] = stbi__div4(t1+2); - - STBI_NOTUSED(hs); - - return out; -} - -#if defined(STBI_SSE2) || defined(STBI_NEON) -static stbi_uc *stbi__resample_row_hv_2_simd(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate 2x2 samples for every one in input - int i=0,t0,t1; - - if (w == 1) { - out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); - return out; - } - - t1 = 3*in_near[0] + in_far[0]; - // process groups of 8 pixels for as long as we can. - // note we can't handle the last pixel in a row in this loop - // because we need to handle the filter boundary conditions. - for (; i < ((w-1) & ~7); i += 8) { -#if defined(STBI_SSE2) - // load and perform the vertical filtering pass - // this uses 3*x + y = 4*x + (y - x) - __m128i zero = _mm_setzero_si128(); - __m128i farb = _mm_loadl_epi64((__m128i *) (in_far + i)); - __m128i nearb = _mm_loadl_epi64((__m128i *) (in_near + i)); - __m128i farw = _mm_unpacklo_epi8(farb, zero); - __m128i nearw = _mm_unpacklo_epi8(nearb, zero); - __m128i diff = _mm_sub_epi16(farw, nearw); - __m128i nears = _mm_slli_epi16(nearw, 2); - __m128i curr = _mm_add_epi16(nears, diff); // current row - - // horizontal filter works the same based on shifted vers of current - // row. "prev" is current row shifted right by 1 pixel; we need to - // insert the previous pixel value (from t1). - // "next" is current row shifted left by 1 pixel, with first pixel - // of next block of 8 pixels added in. - __m128i prv0 = _mm_slli_si128(curr, 2); - __m128i nxt0 = _mm_srli_si128(curr, 2); - __m128i prev = _mm_insert_epi16(prv0, t1, 0); - __m128i next = _mm_insert_epi16(nxt0, 3*in_near[i+8] + in_far[i+8], 7); - - // horizontal filter, polyphase implementation since it's convenient: - // even pixels = 3*cur + prev = cur*4 + (prev - cur) - // odd pixels = 3*cur + next = cur*4 + (next - cur) - // note the shared term. - __m128i bias = _mm_set1_epi16(8); - __m128i curs = _mm_slli_epi16(curr, 2); - __m128i prvd = _mm_sub_epi16(prev, curr); - __m128i nxtd = _mm_sub_epi16(next, curr); - __m128i curb = _mm_add_epi16(curs, bias); - __m128i even = _mm_add_epi16(prvd, curb); - __m128i odd = _mm_add_epi16(nxtd, curb); - - // interleave even and odd pixels, then undo scaling. - __m128i int0 = _mm_unpacklo_epi16(even, odd); - __m128i int1 = _mm_unpackhi_epi16(even, odd); - __m128i de0 = _mm_srli_epi16(int0, 4); - __m128i de1 = _mm_srli_epi16(int1, 4); - - // pack and write output - __m128i outv = _mm_packus_epi16(de0, de1); - _mm_storeu_si128((__m128i *) (out + i*2), outv); -#elif defined(STBI_NEON) - // load and perform the vertical filtering pass - // this uses 3*x + y = 4*x + (y - x) - uint8x8_t farb = vld1_u8(in_far + i); - uint8x8_t nearb = vld1_u8(in_near + i); - int16x8_t diff = vreinterpretq_s16_u16(vsubl_u8(farb, nearb)); - int16x8_t nears = vreinterpretq_s16_u16(vshll_n_u8(nearb, 2)); - int16x8_t curr = vaddq_s16(nears, diff); // current row - - // horizontal filter works the same based on shifted vers of current - // row. "prev" is current row shifted right by 1 pixel; we need to - // insert the previous pixel value (from t1). - // "next" is current row shifted left by 1 pixel, with first pixel - // of next block of 8 pixels added in. - int16x8_t prv0 = vextq_s16(curr, curr, 7); - int16x8_t nxt0 = vextq_s16(curr, curr, 1); - int16x8_t prev = vsetq_lane_s16(t1, prv0, 0); - int16x8_t next = vsetq_lane_s16(3*in_near[i+8] + in_far[i+8], nxt0, 7); - - // horizontal filter, polyphase implementation since it's convenient: - // even pixels = 3*cur + prev = cur*4 + (prev - cur) - // odd pixels = 3*cur + next = cur*4 + (next - cur) - // note the shared term. - int16x8_t curs = vshlq_n_s16(curr, 2); - int16x8_t prvd = vsubq_s16(prev, curr); - int16x8_t nxtd = vsubq_s16(next, curr); - int16x8_t even = vaddq_s16(curs, prvd); - int16x8_t odd = vaddq_s16(curs, nxtd); - - // undo scaling and round, then store with even/odd phases interleaved - uint8x8x2_t o; - o.val[0] = vqrshrun_n_s16(even, 4); - o.val[1] = vqrshrun_n_s16(odd, 4); - vst2_u8(out + i*2, o); -#endif - - // "previous" value for next iter - t1 = 3*in_near[i+7] + in_far[i+7]; - } - - t0 = t1; - t1 = 3*in_near[i] + in_far[i]; - out[i*2] = stbi__div16(3*t1 + t0 + 8); - - for (++i; i < w; ++i) { - t0 = t1; - t1 = 3*in_near[i]+in_far[i]; - out[i*2-1] = stbi__div16(3*t0 + t1 + 8); - out[i*2 ] = stbi__div16(3*t1 + t0 + 8); - } - out[w*2-1] = stbi__div4(t1+2); - - STBI_NOTUSED(hs); - - return out; -} -#endif - -static stbi_uc *stbi__resample_row_generic(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // resample with nearest-neighbor - int i,j; - STBI_NOTUSED(in_far); - for (i=0; i < w; ++i) - for (j=0; j < hs; ++j) - out[i*hs+j] = in_near[i]; - return out; -} - -// this is a reduced-precision calculation of YCbCr-to-RGB introduced -// to make sure the code produces the same results in both SIMD and scalar -#define stbi__float2fixed(x) (((int) ((x) * 4096.0f + 0.5f)) << 8) -static void stbi__YCbCr_to_RGB_row(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step) -{ - int i; - for (i=0; i < count; ++i) { - int y_fixed = (y[i] << 20) + (1<<19); // rounding - int r,g,b; - int cr = pcr[i] - 128; - int cb = pcb[i] - 128; - r = y_fixed + cr* stbi__float2fixed(1.40200f); - g = y_fixed + (cr*-stbi__float2fixed(0.71414f)) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); - b = y_fixed + cb* stbi__float2fixed(1.77200f); - r >>= 20; - g >>= 20; - b >>= 20; - if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } - if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } - if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } - out[0] = (stbi_uc)r; - out[1] = (stbi_uc)g; - out[2] = (stbi_uc)b; - out[3] = 255; - out += step; - } -} - -#if defined(STBI_SSE2) || defined(STBI_NEON) -static void stbi__YCbCr_to_RGB_simd(stbi_uc *out, stbi_uc const *y, stbi_uc const *pcb, stbi_uc const *pcr, int count, int step) -{ - int i = 0; - -#ifdef STBI_SSE2 - // step == 3 is pretty ugly on the final interleave, and i'm not convinced - // it's useful in practice (you wouldn't use it for textures, for example). - // so just accelerate step == 4 case. - if (step == 4) { - // this is a fairly straightforward implementation and not super-optimized. - __m128i signflip = _mm_set1_epi8(-0x80); - __m128i cr_const0 = _mm_set1_epi16( (short) ( 1.40200f*4096.0f+0.5f)); - __m128i cr_const1 = _mm_set1_epi16( - (short) ( 0.71414f*4096.0f+0.5f)); - __m128i cb_const0 = _mm_set1_epi16( - (short) ( 0.34414f*4096.0f+0.5f)); - __m128i cb_const1 = _mm_set1_epi16( (short) ( 1.77200f*4096.0f+0.5f)); - __m128i y_bias = _mm_set1_epi8((char) (unsigned char) 128); - __m128i xw = _mm_set1_epi16(255); // alpha channel - - for (; i+7 < count; i += 8) { - // load - __m128i y_bytes = _mm_loadl_epi64((__m128i *) (y+i)); - __m128i cr_bytes = _mm_loadl_epi64((__m128i *) (pcr+i)); - __m128i cb_bytes = _mm_loadl_epi64((__m128i *) (pcb+i)); - __m128i cr_biased = _mm_xor_si128(cr_bytes, signflip); // -128 - __m128i cb_biased = _mm_xor_si128(cb_bytes, signflip); // -128 - - // unpack to short (and left-shift cr, cb by 8) - __m128i yw = _mm_unpacklo_epi8(y_bias, y_bytes); - __m128i crw = _mm_unpacklo_epi8(_mm_setzero_si128(), cr_biased); - __m128i cbw = _mm_unpacklo_epi8(_mm_setzero_si128(), cb_biased); - - // color transform - __m128i yws = _mm_srli_epi16(yw, 4); - __m128i cr0 = _mm_mulhi_epi16(cr_const0, crw); - __m128i cb0 = _mm_mulhi_epi16(cb_const0, cbw); - __m128i cb1 = _mm_mulhi_epi16(cbw, cb_const1); - __m128i cr1 = _mm_mulhi_epi16(crw, cr_const1); - __m128i rws = _mm_add_epi16(cr0, yws); - __m128i gwt = _mm_add_epi16(cb0, yws); - __m128i bws = _mm_add_epi16(yws, cb1); - __m128i gws = _mm_add_epi16(gwt, cr1); - - // descale - __m128i rw = _mm_srai_epi16(rws, 4); - __m128i bw = _mm_srai_epi16(bws, 4); - __m128i gw = _mm_srai_epi16(gws, 4); - - // back to byte, set up for transpose - __m128i brb = _mm_packus_epi16(rw, bw); - __m128i gxb = _mm_packus_epi16(gw, xw); - - // transpose to interleave channels - __m128i t0 = _mm_unpacklo_epi8(brb, gxb); - __m128i t1 = _mm_unpackhi_epi8(brb, gxb); - __m128i o0 = _mm_unpacklo_epi16(t0, t1); - __m128i o1 = _mm_unpackhi_epi16(t0, t1); - - // store - _mm_storeu_si128((__m128i *) (out + 0), o0); - _mm_storeu_si128((__m128i *) (out + 16), o1); - out += 32; - } - } -#endif - -#ifdef STBI_NEON - // in this version, step=3 support would be easy to add. but is there demand? - if (step == 4) { - // this is a fairly straightforward implementation and not super-optimized. - uint8x8_t signflip = vdup_n_u8(0x80); - int16x8_t cr_const0 = vdupq_n_s16( (short) ( 1.40200f*4096.0f+0.5f)); - int16x8_t cr_const1 = vdupq_n_s16( - (short) ( 0.71414f*4096.0f+0.5f)); - int16x8_t cb_const0 = vdupq_n_s16( - (short) ( 0.34414f*4096.0f+0.5f)); - int16x8_t cb_const1 = vdupq_n_s16( (short) ( 1.77200f*4096.0f+0.5f)); - - for (; i+7 < count; i += 8) { - // load - uint8x8_t y_bytes = vld1_u8(y + i); - uint8x8_t cr_bytes = vld1_u8(pcr + i); - uint8x8_t cb_bytes = vld1_u8(pcb + i); - int8x8_t cr_biased = vreinterpret_s8_u8(vsub_u8(cr_bytes, signflip)); - int8x8_t cb_biased = vreinterpret_s8_u8(vsub_u8(cb_bytes, signflip)); - - // expand to s16 - int16x8_t yws = vreinterpretq_s16_u16(vshll_n_u8(y_bytes, 4)); - int16x8_t crw = vshll_n_s8(cr_biased, 7); - int16x8_t cbw = vshll_n_s8(cb_biased, 7); - - // color transform - int16x8_t cr0 = vqdmulhq_s16(crw, cr_const0); - int16x8_t cb0 = vqdmulhq_s16(cbw, cb_const0); - int16x8_t cr1 = vqdmulhq_s16(crw, cr_const1); - int16x8_t cb1 = vqdmulhq_s16(cbw, cb_const1); - int16x8_t rws = vaddq_s16(yws, cr0); - int16x8_t gws = vaddq_s16(vaddq_s16(yws, cb0), cr1); - int16x8_t bws = vaddq_s16(yws, cb1); - - // undo scaling, round, convert to byte - uint8x8x4_t o; - o.val[0] = vqrshrun_n_s16(rws, 4); - o.val[1] = vqrshrun_n_s16(gws, 4); - o.val[2] = vqrshrun_n_s16(bws, 4); - o.val[3] = vdup_n_u8(255); - - // store, interleaving r/g/b/a - vst4_u8(out, o); - out += 8*4; - } - } -#endif - - for (; i < count; ++i) { - int y_fixed = (y[i] << 20) + (1<<19); // rounding - int r,g,b; - int cr = pcr[i] - 128; - int cb = pcb[i] - 128; - r = y_fixed + cr* stbi__float2fixed(1.40200f); - g = y_fixed + cr*-stbi__float2fixed(0.71414f) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); - b = y_fixed + cb* stbi__float2fixed(1.77200f); - r >>= 20; - g >>= 20; - b >>= 20; - if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } - if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } - if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } - out[0] = (stbi_uc)r; - out[1] = (stbi_uc)g; - out[2] = (stbi_uc)b; - out[3] = 255; - out += step; - } -} -#endif - -// set up the kernels -static void stbi__setup_jpeg(stbi__jpeg *j) -{ - j->idct_block_kernel = stbi__idct_block; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_row; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2; - -#ifdef STBI_SSE2 - if (stbi__sse2_available()) { - j->idct_block_kernel = stbi__idct_simd; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; - } -#endif - -#ifdef STBI_NEON - j->idct_block_kernel = stbi__idct_simd; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; -#endif -} - -// clean up the temporary component buffers -static void stbi__cleanup_jpeg(stbi__jpeg *j) -{ - stbi__free_jpeg_components(j, j->s->img_n, 0); -} - -typedef struct -{ - resample_row_func resample; - stbi_uc *line0,*line1; - int hs,vs; // expansion factor in each axis - int w_lores; // horizontal pixels pre-expansion - int ystep; // how far through vertical expansion we are - int ypos; // which pre-expansion row we're on -} stbi__resample; - -// fast 0..255 * 0..255 => 0..255 rounded multiplication -static stbi_uc stbi__blinn_8x8(stbi_uc x, stbi_uc y) -{ - unsigned int t = x*y + 128; - return (stbi_uc) ((t + (t >>8)) >> 8); -} - -static stbi_uc *load_jpeg_image(stbi__jpeg *z, int *out_x, int *out_y, int *comp, int req_comp) -{ - int n, decode_n, is_rgb; - z->s->img_n = 0; // make stbi__cleanup_jpeg safe - - // validate req_comp - if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); - - // load a jpeg image from whichever source, but leave in YCbCr format - if (!stbi__decode_jpeg_image(z)) { stbi__cleanup_jpeg(z); return NULL; } - - // determine actual number of components to generate - n = req_comp ? req_comp : z->s->img_n >= 3 ? 3 : 1; - - is_rgb = z->s->img_n == 3 && (z->rgb == 3 || (z->app14_color_transform == 0 && !z->jfif)); - - if (z->s->img_n == 3 && n < 3 && !is_rgb) - decode_n = 1; - else - decode_n = z->s->img_n; - - // nothing to do if no components requested; check this now to avoid - // accessing uninitialized coutput[0] later - if (decode_n <= 0) { stbi__cleanup_jpeg(z); return NULL; } - - // resample and color-convert - { - int k; - unsigned int i,j; - stbi_uc *output; - stbi_uc *coutput[4] = { NULL, NULL, NULL, NULL }; - - stbi__resample res_comp[4]; - - for (k=0; k < decode_n; ++k) { - stbi__resample *r = &res_comp[k]; - - // allocate line buffer big enough for upsampling off the edges - // with upsample factor of 4 - z->img_comp[k].linebuf = (stbi_uc *) stbi__malloc(z->s->img_x + 3); - if (!z->img_comp[k].linebuf) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } - - r->hs = z->img_h_max / z->img_comp[k].h; - r->vs = z->img_v_max / z->img_comp[k].v; - r->ystep = r->vs >> 1; - r->w_lores = (z->s->img_x + r->hs-1) / r->hs; - r->ypos = 0; - r->line0 = r->line1 = z->img_comp[k].data; - - if (r->hs == 1 && r->vs == 1) r->resample = resample_row_1; - else if (r->hs == 1 && r->vs == 2) r->resample = stbi__resample_row_v_2; - else if (r->hs == 2 && r->vs == 1) r->resample = stbi__resample_row_h_2; - else if (r->hs == 2 && r->vs == 2) r->resample = z->resample_row_hv_2_kernel; - else r->resample = stbi__resample_row_generic; - } - - // can't error after this so, this is safe - output = (stbi_uc *) stbi__malloc_mad3(n, z->s->img_x, z->s->img_y, 1); - if (!output) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } - - // now go ahead and resample - for (j=0; j < z->s->img_y; ++j) { - stbi_uc *out = output + n * z->s->img_x * j; - for (k=0; k < decode_n; ++k) { - stbi__resample *r = &res_comp[k]; - int y_bot = r->ystep >= (r->vs >> 1); - coutput[k] = r->resample(z->img_comp[k].linebuf, - y_bot ? r->line1 : r->line0, - y_bot ? r->line0 : r->line1, - r->w_lores, r->hs); - if (++r->ystep >= r->vs) { - r->ystep = 0; - r->line0 = r->line1; - if (++r->ypos < z->img_comp[k].y) - r->line1 += z->img_comp[k].w2; - } - } - if (n >= 3) { - stbi_uc *y = coutput[0]; - if (z->s->img_n == 3) { - if (is_rgb) { - for (i=0; i < z->s->img_x; ++i) { - out[0] = y[i]; - out[1] = coutput[1][i]; - out[2] = coutput[2][i]; - out[3] = 255; - out += n; - } - } else { - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - } - } else if (z->s->img_n == 4) { - if (z->app14_color_transform == 0) { // CMYK - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - out[0] = stbi__blinn_8x8(coutput[0][i], m); - out[1] = stbi__blinn_8x8(coutput[1][i], m); - out[2] = stbi__blinn_8x8(coutput[2][i], m); - out[3] = 255; - out += n; - } - } else if (z->app14_color_transform == 2) { // YCCK - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - out[0] = stbi__blinn_8x8(255 - out[0], m); - out[1] = stbi__blinn_8x8(255 - out[1], m); - out[2] = stbi__blinn_8x8(255 - out[2], m); - out += n; - } - } else { // YCbCr + alpha? Ignore the fourth channel for now - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - } - } else - for (i=0; i < z->s->img_x; ++i) { - out[0] = out[1] = out[2] = y[i]; - out[3] = 255; // not used if n==3 - out += n; - } - } else { - if (is_rgb) { - if (n == 1) - for (i=0; i < z->s->img_x; ++i) - *out++ = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); - else { - for (i=0; i < z->s->img_x; ++i, out += 2) { - out[0] = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); - out[1] = 255; - } - } - } else if (z->s->img_n == 4 && z->app14_color_transform == 0) { - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - stbi_uc r = stbi__blinn_8x8(coutput[0][i], m); - stbi_uc g = stbi__blinn_8x8(coutput[1][i], m); - stbi_uc b = stbi__blinn_8x8(coutput[2][i], m); - out[0] = stbi__compute_y(r, g, b); - out[1] = 255; - out += n; - } - } else if (z->s->img_n == 4 && z->app14_color_transform == 2) { - for (i=0; i < z->s->img_x; ++i) { - out[0] = stbi__blinn_8x8(255 - coutput[0][i], coutput[3][i]); - out[1] = 255; - out += n; - } - } else { - stbi_uc *y = coutput[0]; - if (n == 1) - for (i=0; i < z->s->img_x; ++i) out[i] = y[i]; - else - for (i=0; i < z->s->img_x; ++i) { *out++ = y[i]; *out++ = 255; } - } - } - } - stbi__cleanup_jpeg(z); - *out_x = z->s->img_x; - *out_y = z->s->img_y; - if (comp) *comp = z->s->img_n >= 3 ? 3 : 1; // report original components, not output - return output; - } -} - -static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - unsigned char* result; - stbi__jpeg* j = (stbi__jpeg*) stbi__malloc(sizeof(stbi__jpeg)); - if (!j) return stbi__errpuc("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - STBI_NOTUSED(ri); - j->s = s; - stbi__setup_jpeg(j); - result = load_jpeg_image(j, x,y,comp,req_comp); - STBI_FREE(j); - return result; -} - -static int stbi__jpeg_test(stbi__context *s) -{ - int r; - stbi__jpeg* j = (stbi__jpeg*)stbi__malloc(sizeof(stbi__jpeg)); - if (!j) return stbi__err("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - j->s = s; - stbi__setup_jpeg(j); - r = stbi__decode_jpeg_header(j, STBI__SCAN_type); - stbi__rewind(s); - STBI_FREE(j); - return r; -} - -static int stbi__jpeg_info_raw(stbi__jpeg *j, int *x, int *y, int *comp) -{ - if (!stbi__decode_jpeg_header(j, STBI__SCAN_header)) { - stbi__rewind( j->s ); - return 0; - } - if (x) *x = j->s->img_x; - if (y) *y = j->s->img_y; - if (comp) *comp = j->s->img_n >= 3 ? 3 : 1; - return 1; -} - -static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp) -{ - int result; - stbi__jpeg* j = (stbi__jpeg*) (stbi__malloc(sizeof(stbi__jpeg))); - if (!j) return stbi__err("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - j->s = s; - result = stbi__jpeg_info_raw(j, x, y, comp); - STBI_FREE(j); - return result; -} -#endif - -// public domain zlib decode v0.2 Sean Barrett 2006-11-18 -// simple implementation -// - all input must be provided in an upfront buffer -// - all output is written to a single output buffer (can malloc/realloc) -// performance -// - fast huffman - -#ifndef STBI_NO_ZLIB - -// fast-way is faster to check than jpeg huffman, but slow way is slower -#define STBI__ZFAST_BITS 9 // accelerate all cases in default tables -#define STBI__ZFAST_MASK ((1 << STBI__ZFAST_BITS) - 1) -#define STBI__ZNSYMS 288 // number of symbols in literal/length alphabet - -// zlib-style huffman encoding -// (jpegs packs from left, zlib from right, so can't share code) -typedef struct -{ - stbi__uint16 fast[1 << STBI__ZFAST_BITS]; - stbi__uint16 firstcode[16]; - int maxcode[17]; - stbi__uint16 firstsymbol[16]; - stbi_uc size[STBI__ZNSYMS]; - stbi__uint16 value[STBI__ZNSYMS]; -} stbi__zhuffman; - -stbi_inline static int stbi__bitreverse16(int n) -{ - n = ((n & 0xAAAA) >> 1) | ((n & 0x5555) << 1); - n = ((n & 0xCCCC) >> 2) | ((n & 0x3333) << 2); - n = ((n & 0xF0F0) >> 4) | ((n & 0x0F0F) << 4); - n = ((n & 0xFF00) >> 8) | ((n & 0x00FF) << 8); - return n; -} - -stbi_inline static int stbi__bit_reverse(int v, int bits) -{ - STBI_ASSERT(bits <= 16); - // to bit reverse n bits, reverse 16 and shift - // e.g. 11 bits, bit reverse and shift away 5 - return stbi__bitreverse16(v) >> (16-bits); -} - -static int stbi__zbuild_huffman(stbi__zhuffman *z, const stbi_uc *sizelist, int num) -{ - int i,k=0; - int code, next_code[16], sizes[17]; - - // DEFLATE spec for generating codes - memset(sizes, 0, sizeof(sizes)); - memset(z->fast, 0, sizeof(z->fast)); - for (i=0; i < num; ++i) - ++sizes[sizelist[i]]; - sizes[0] = 0; - for (i=1; i < 16; ++i) - if (sizes[i] > (1 << i)) - return stbi__err("bad sizes", "Corrupt PNG"); - code = 0; - for (i=1; i < 16; ++i) { - next_code[i] = code; - z->firstcode[i] = (stbi__uint16) code; - z->firstsymbol[i] = (stbi__uint16) k; - code = (code + sizes[i]); - if (sizes[i]) - if (code-1 >= (1 << i)) return stbi__err("bad codelengths","Corrupt PNG"); - z->maxcode[i] = code << (16-i); // preshift for inner loop - code <<= 1; - k += sizes[i]; - } - z->maxcode[16] = 0x10000; // sentinel - for (i=0; i < num; ++i) { - int s = sizelist[i]; - if (s) { - int c = next_code[s] - z->firstcode[s] + z->firstsymbol[s]; - stbi__uint16 fastv = (stbi__uint16) ((s << 9) | i); - z->size [c] = (stbi_uc ) s; - z->value[c] = (stbi__uint16) i; - if (s <= STBI__ZFAST_BITS) { - int j = stbi__bit_reverse(next_code[s],s); - while (j < (1 << STBI__ZFAST_BITS)) { - z->fast[j] = fastv; - j += (1 << s); - } - } - ++next_code[s]; - } - } - return 1; -} - -// zlib-from-memory implementation for PNG reading -// because PNG allows splitting the zlib stream arbitrarily, -// and it's annoying structurally to have PNG call ZLIB call PNG, -// we require PNG read all the IDATs and combine them into a single -// memory buffer - -typedef struct -{ - stbi_uc *zbuffer, *zbuffer_end; - int num_bits; - int hit_zeof_once; - stbi__uint32 code_buffer; - - char *zout; - char *zout_start; - char *zout_end; - int z_expandable; - - stbi__zhuffman z_length, z_distance; -} stbi__zbuf; - -stbi_inline static int stbi__zeof(stbi__zbuf *z) -{ - return (z->zbuffer >= z->zbuffer_end); -} - -stbi_inline static stbi_uc stbi__zget8(stbi__zbuf *z) -{ - return stbi__zeof(z) ? 0 : *z->zbuffer++; -} - -static void stbi__fill_bits(stbi__zbuf *z) -{ - do { - if (z->code_buffer >= (1U << z->num_bits)) { - z->zbuffer = z->zbuffer_end; /* treat this as EOF so we fail. */ - return; - } - z->code_buffer |= (unsigned int) stbi__zget8(z) << z->num_bits; - z->num_bits += 8; - } while (z->num_bits <= 24); -} - -stbi_inline static unsigned int stbi__zreceive(stbi__zbuf *z, int n) -{ - unsigned int k; - if (z->num_bits < n) stbi__fill_bits(z); - k = z->code_buffer & ((1 << n) - 1); - z->code_buffer >>= n; - z->num_bits -= n; - return k; -} - -static int stbi__zhuffman_decode_slowpath(stbi__zbuf *a, stbi__zhuffman *z) -{ - int b,s,k; - // not resolved by fast table, so compute it the slow way - // use jpeg approach, which requires MSbits at top - k = stbi__bit_reverse(a->code_buffer, 16); - for (s=STBI__ZFAST_BITS+1; ; ++s) - if (k < z->maxcode[s]) - break; - if (s >= 16) return -1; // invalid code! - // code size is s, so: - b = (k >> (16-s)) - z->firstcode[s] + z->firstsymbol[s]; - if (b >= STBI__ZNSYMS) return -1; // some data was corrupt somewhere! - if (z->size[b] != s) return -1; // was originally an assert, but report failure instead. - a->code_buffer >>= s; - a->num_bits -= s; - return z->value[b]; -} - -stbi_inline static int stbi__zhuffman_decode(stbi__zbuf *a, stbi__zhuffman *z) -{ - int b,s; - if (a->num_bits < 16) { - if (stbi__zeof(a)) { - if (!a->hit_zeof_once) { - // This is the first time we hit eof, insert 16 extra padding btis - // to allow us to keep going; if we actually consume any of them - // though, that is invalid data. This is caught later. - a->hit_zeof_once = 1; - a->num_bits += 16; // add 16 implicit zero bits - } else { - // We already inserted our extra 16 padding bits and are again - // out, this stream is actually prematurely terminated. - return -1; - } - } else { - stbi__fill_bits(a); - } - } - b = z->fast[a->code_buffer & STBI__ZFAST_MASK]; - if (b) { - s = b >> 9; - a->code_buffer >>= s; - a->num_bits -= s; - return b & 511; - } - return stbi__zhuffman_decode_slowpath(a, z); -} - -static int stbi__zexpand(stbi__zbuf *z, char *zout, int n) // need to make room for n bytes -{ - char *q; - unsigned int cur, limit, old_limit; - z->zout = zout; - if (!z->z_expandable) return stbi__err("output buffer limit","Corrupt PNG"); - cur = (unsigned int) (z->zout - z->zout_start); - limit = old_limit = (unsigned) (z->zout_end - z->zout_start); - if (UINT_MAX - cur < (unsigned) n) return stbi__err("outofmem", "Out of memory"); - while (cur + n > limit) { - if(limit > UINT_MAX / 2) return stbi__err("outofmem", "Out of memory"); - limit *= 2; - } - q = (char *) STBI_REALLOC_SIZED(z->zout_start, old_limit, limit); - STBI_NOTUSED(old_limit); - if (q == NULL) return stbi__err("outofmem", "Out of memory"); - z->zout_start = q; - z->zout = q + cur; - z->zout_end = q + limit; - return 1; -} - -static const int stbi__zlength_base[31] = { - 3,4,5,6,7,8,9,10,11,13, - 15,17,19,23,27,31,35,43,51,59, - 67,83,99,115,131,163,195,227,258,0,0 }; - -static const int stbi__zlength_extra[31]= -{ 0,0,0,0,0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4,5,5,5,5,0,0,0 }; - -static const int stbi__zdist_base[32] = { 1,2,3,4,5,7,9,13,17,25,33,49,65,97,129,193, -257,385,513,769,1025,1537,2049,3073,4097,6145,8193,12289,16385,24577,0,0}; - -static const int stbi__zdist_extra[32] = -{ 0,0,0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13}; - -static int stbi__parse_huffman_block(stbi__zbuf *a) -{ - char *zout = a->zout; - for(;;) { - int z = stbi__zhuffman_decode(a, &a->z_length); - if (z < 256) { - if (z < 0) return stbi__err("bad huffman code","Corrupt PNG"); // error in huffman codes - if (zout >= a->zout_end) { - if (!stbi__zexpand(a, zout, 1)) return 0; - zout = a->zout; - } - *zout++ = (char) z; - } else { - stbi_uc *p; - int len,dist; - if (z == 256) { - a->zout = zout; - if (a->hit_zeof_once && a->num_bits < 16) { - // The first time we hit zeof, we inserted 16 extra zero bits into our bit - // buffer so the decoder can just do its speculative decoding. But if we - // actually consumed any of those bits (which is the case when num_bits < 16), - // the stream actually read past the end so it is malformed. - return stbi__err("unexpected end","Corrupt PNG"); - } - return 1; - } - if (z >= 286) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, length codes 286 and 287 must not appear in compressed data - z -= 257; - len = stbi__zlength_base[z]; - if (stbi__zlength_extra[z]) len += stbi__zreceive(a, stbi__zlength_extra[z]); - z = stbi__zhuffman_decode(a, &a->z_distance); - if (z < 0 || z >= 30) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, distance codes 30 and 31 must not appear in compressed data - dist = stbi__zdist_base[z]; - if (stbi__zdist_extra[z]) dist += stbi__zreceive(a, stbi__zdist_extra[z]); - if (zout - a->zout_start < dist) return stbi__err("bad dist","Corrupt PNG"); - if (len > a->zout_end - zout) { - if (!stbi__zexpand(a, zout, len)) return 0; - zout = a->zout; - } - p = (stbi_uc *) (zout - dist); - if (dist == 1) { // run of one byte; common in images. - stbi_uc v = *p; - if (len) { do *zout++ = v; while (--len); } - } else { - if (len) { do *zout++ = *p++; while (--len); } - } - } - } -} - -static int stbi__compute_huffman_codes(stbi__zbuf *a) -{ - static const stbi_uc length_dezigzag[19] = { 16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15 }; - stbi__zhuffman z_codelength; - stbi_uc lencodes[286+32+137];//padding for maximum single op - stbi_uc codelength_sizes[19]; - int i,n; - - int hlit = stbi__zreceive(a,5) + 257; - int hdist = stbi__zreceive(a,5) + 1; - int hclen = stbi__zreceive(a,4) + 4; - int ntot = hlit + hdist; - - memset(codelength_sizes, 0, sizeof(codelength_sizes)); - for (i=0; i < hclen; ++i) { - int s = stbi__zreceive(a,3); - codelength_sizes[length_dezigzag[i]] = (stbi_uc) s; - } - if (!stbi__zbuild_huffman(&z_codelength, codelength_sizes, 19)) return 0; - - n = 0; - while (n < ntot) { - int c = stbi__zhuffman_decode(a, &z_codelength); - if (c < 0 || c >= 19) return stbi__err("bad codelengths", "Corrupt PNG"); - if (c < 16) - lencodes[n++] = (stbi_uc) c; - else { - stbi_uc fill = 0; - if (c == 16) { - c = stbi__zreceive(a,2)+3; - if (n == 0) return stbi__err("bad codelengths", "Corrupt PNG"); - fill = lencodes[n-1]; - } else if (c == 17) { - c = stbi__zreceive(a,3)+3; - } else if (c == 18) { - c = stbi__zreceive(a,7)+11; - } else { - return stbi__err("bad codelengths", "Corrupt PNG"); - } - if (ntot - n < c) return stbi__err("bad codelengths", "Corrupt PNG"); - memset(lencodes+n, fill, c); - n += c; - } - } - if (n != ntot) return stbi__err("bad codelengths","Corrupt PNG"); - if (!stbi__zbuild_huffman(&a->z_length, lencodes, hlit)) return 0; - if (!stbi__zbuild_huffman(&a->z_distance, lencodes+hlit, hdist)) return 0; - return 1; -} - -static int stbi__parse_uncompressed_block(stbi__zbuf *a) -{ - stbi_uc header[4]; - int len,nlen,k; - if (a->num_bits & 7) - stbi__zreceive(a, a->num_bits & 7); // discard - // drain the bit-packed data into header - k = 0; - while (a->num_bits > 0) { - header[k++] = (stbi_uc) (a->code_buffer & 255); // suppress MSVC run-time check - a->code_buffer >>= 8; - a->num_bits -= 8; - } - if (a->num_bits < 0) return stbi__err("zlib corrupt","Corrupt PNG"); - // now fill header the normal way - while (k < 4) - header[k++] = stbi__zget8(a); - len = header[1] * 256 + header[0]; - nlen = header[3] * 256 + header[2]; - if (nlen != (len ^ 0xffff)) return stbi__err("zlib corrupt","Corrupt PNG"); - if (a->zbuffer + len > a->zbuffer_end) return stbi__err("read past buffer","Corrupt PNG"); - if (a->zout + len > a->zout_end) - if (!stbi__zexpand(a, a->zout, len)) return 0; - memcpy(a->zout, a->zbuffer, len); - a->zbuffer += len; - a->zout += len; - return 1; -} - -static int stbi__parse_zlib_header(stbi__zbuf *a) -{ - int cmf = stbi__zget8(a); - int cm = cmf & 15; - /* int cinfo = cmf >> 4; */ - int flg = stbi__zget8(a); - if (stbi__zeof(a)) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec - if ((cmf*256+flg) % 31 != 0) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec - if (flg & 32) return stbi__err("no preset dict","Corrupt PNG"); // preset dictionary not allowed in png - if (cm != 8) return stbi__err("bad compression","Corrupt PNG"); // DEFLATE required for png - // window = 1 << (8 + cinfo)... but who cares, we fully buffer output - return 1; -} - -static const stbi_uc stbi__zdefault_length[STBI__ZNSYMS] = -{ - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7, 7,7,7,7,7,7,7,7,8,8,8,8,8,8,8,8 -}; -static const stbi_uc stbi__zdefault_distance[32] = -{ - 5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5 -}; -/* -Init algorithm: -{ - int i; // use <= to match clearly with spec - for (i=0; i <= 143; ++i) stbi__zdefault_length[i] = 8; - for ( ; i <= 255; ++i) stbi__zdefault_length[i] = 9; - for ( ; i <= 279; ++i) stbi__zdefault_length[i] = 7; - for ( ; i <= 287; ++i) stbi__zdefault_length[i] = 8; - - for (i=0; i <= 31; ++i) stbi__zdefault_distance[i] = 5; -} -*/ - -static int stbi__parse_zlib(stbi__zbuf *a, int parse_header) -{ - int final, type; - if (parse_header) - if (!stbi__parse_zlib_header(a)) return 0; - a->num_bits = 0; - a->code_buffer = 0; - a->hit_zeof_once = 0; - do { - final = stbi__zreceive(a,1); - type = stbi__zreceive(a,2); - if (type == 0) { - if (!stbi__parse_uncompressed_block(a)) return 0; - } else if (type == 3) { - return 0; - } else { - if (type == 1) { - // use fixed code lengths - if (!stbi__zbuild_huffman(&a->z_length , stbi__zdefault_length , STBI__ZNSYMS)) return 0; - if (!stbi__zbuild_huffman(&a->z_distance, stbi__zdefault_distance, 32)) return 0; - } else { - if (!stbi__compute_huffman_codes(a)) return 0; - } - if (!stbi__parse_huffman_block(a)) return 0; - } - } while (!final); - return 1; -} - -static int stbi__do_zlib(stbi__zbuf *a, char *obuf, int olen, int exp, int parse_header) -{ - a->zout_start = obuf; - a->zout = obuf; - a->zout_end = obuf + olen; - a->z_expandable = exp; - - return stbi__parse_zlib(a, parse_header); -} - -STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(initial_size); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer + len; - if (stbi__do_zlib(&a, p, initial_size, 1, 1)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF char *stbi_zlib_decode_malloc(char const *buffer, int len, int *outlen) -{ - return stbi_zlib_decode_malloc_guesssize(buffer, len, 16384, outlen); -} - -STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(initial_size); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer + len; - if (stbi__do_zlib(&a, p, initial_size, 1, parse_header)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, char const *ibuffer, int ilen) -{ - stbi__zbuf a; - a.zbuffer = (stbi_uc *) ibuffer; - a.zbuffer_end = (stbi_uc *) ibuffer + ilen; - if (stbi__do_zlib(&a, obuffer, olen, 0, 1)) - return (int) (a.zout - a.zout_start); - else - return -1; -} - -STBIDEF char *stbi_zlib_decode_noheader_malloc(char const *buffer, int len, int *outlen) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(16384); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer+len; - if (stbi__do_zlib(&a, p, 16384, 1, 0)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen) -{ - stbi__zbuf a; - a.zbuffer = (stbi_uc *) ibuffer; - a.zbuffer_end = (stbi_uc *) ibuffer + ilen; - if (stbi__do_zlib(&a, obuffer, olen, 0, 0)) - return (int) (a.zout - a.zout_start); - else - return -1; -} -#endif - -// public domain "baseline" PNG decoder v0.10 Sean Barrett 2006-11-18 -// simple implementation -// - only 8-bit samples -// - no CRC checking -// - allocates lots of intermediate memory -// - avoids problem of streaming data between subsystems -// - avoids explicit window management -// performance -// - uses stb_zlib, a PD zlib implementation with fast huffman decoding - -#ifndef STBI_NO_PNG -typedef struct -{ - stbi__uint32 length; - stbi__uint32 type; -} stbi__pngchunk; - -static stbi__pngchunk stbi__get_chunk_header(stbi__context *s) -{ - stbi__pngchunk c; - c.length = stbi__get32be(s); - c.type = stbi__get32be(s); - return c; -} - -static int stbi__check_png_header(stbi__context *s) -{ - static const stbi_uc png_sig[8] = { 137,80,78,71,13,10,26,10 }; - int i; - for (i=0; i < 8; ++i) - if (stbi__get8(s) != png_sig[i]) return stbi__err("bad png sig","Not a PNG"); - return 1; -} - -typedef struct -{ - stbi__context *s; - stbi_uc *idata, *expanded, *out; - int depth; -} stbi__png; - - -enum { - STBI__F_none=0, - STBI__F_sub=1, - STBI__F_up=2, - STBI__F_avg=3, - STBI__F_paeth=4, - // synthetic filter used for first scanline to avoid needing a dummy row of 0s - STBI__F_avg_first -}; - -static stbi_uc first_row_filter[5] = -{ - STBI__F_none, - STBI__F_sub, - STBI__F_none, - STBI__F_avg_first, - STBI__F_sub // Paeth with b=c=0 turns out to be equivalent to sub -}; - -static int stbi__paeth(int a, int b, int c) -{ - // This formulation looks very different from the reference in the PNG spec, but is - // actually equivalent and has favorable data dependencies and admits straightforward - // generation of branch-free code, which helps performance significantly. - int thresh = c*3 - (a + b); - int lo = a < b ? a : b; - int hi = a < b ? b : a; - int t0 = (hi <= thresh) ? lo : c; - int t1 = (thresh <= lo) ? hi : t0; - return t1; -} - -static const stbi_uc stbi__depth_scale_table[9] = { 0, 0xff, 0x55, 0, 0x11, 0,0,0, 0x01 }; - -// adds an extra all-255 alpha channel -// dest == src is legal -// img_n must be 1 or 3 -static void stbi__create_png_alpha_expand8(stbi_uc *dest, stbi_uc *src, stbi__uint32 x, int img_n) -{ - int i; - // must process data backwards since we allow dest==src - if (img_n == 1) { - for (i=x-1; i >= 0; --i) { - dest[i*2+1] = 255; - dest[i*2+0] = src[i]; - } - } else { - STBI_ASSERT(img_n == 3); - for (i=x-1; i >= 0; --i) { - dest[i*4+3] = 255; - dest[i*4+2] = src[i*3+2]; - dest[i*4+1] = src[i*3+1]; - dest[i*4+0] = src[i*3+0]; - } - } -} - -// create the png data from post-deflated data -static int stbi__create_png_image_raw(stbi__png *a, stbi_uc *raw, stbi__uint32 raw_len, int out_n, stbi__uint32 x, stbi__uint32 y, int depth, int color) -{ - int bytes = (depth == 16 ? 2 : 1); - stbi__context *s = a->s; - stbi__uint32 i,j,stride = x*out_n*bytes; - stbi__uint32 img_len, img_width_bytes; - stbi_uc *filter_buf; - int all_ok = 1; - int k; - int img_n = s->img_n; // copy it into a local for later - - int output_bytes = out_n*bytes; - int filter_bytes = img_n*bytes; - int width = x; - - STBI_ASSERT(out_n == s->img_n || out_n == s->img_n+1); - a->out = (stbi_uc *) stbi__malloc_mad3(x, y, output_bytes, 0); // extra bytes to write off the end into - if (!a->out) return stbi__err("outofmem", "Out of memory"); - - // note: error exits here don't need to clean up a->out individually, - // stbi__do_png always does on error. - if (!stbi__mad3sizes_valid(img_n, x, depth, 7)) return stbi__err("too large", "Corrupt PNG"); - img_width_bytes = (((img_n * x * depth) + 7) >> 3); - if (!stbi__mad2sizes_valid(img_width_bytes, y, img_width_bytes)) return stbi__err("too large", "Corrupt PNG"); - img_len = (img_width_bytes + 1) * y; - - // we used to check for exact match between raw_len and img_len on non-interlaced PNGs, - // but issue #276 reported a PNG in the wild that had extra data at the end (all zeros), - // so just check for raw_len < img_len always. - if (raw_len < img_len) return stbi__err("not enough pixels","Corrupt PNG"); - - // Allocate two scan lines worth of filter workspace buffer. - filter_buf = (stbi_uc *) stbi__malloc_mad2(img_width_bytes, 2, 0); - if (!filter_buf) return stbi__err("outofmem", "Out of memory"); - - // Filtering for low-bit-depth images - if (depth < 8) { - filter_bytes = 1; - width = img_width_bytes; - } - - for (j=0; j < y; ++j) { - // cur/prior filter buffers alternate - stbi_uc *cur = filter_buf + (j & 1)*img_width_bytes; - stbi_uc *prior = filter_buf + (~j & 1)*img_width_bytes; - stbi_uc *dest = a->out + stride*j; - int nk = width * filter_bytes; - int filter = *raw++; - - // check filter type - if (filter > 4) { - all_ok = stbi__err("invalid filter","Corrupt PNG"); - break; - } - - // if first row, use special filter that doesn't sample previous row - if (j == 0) filter = first_row_filter[filter]; - - // perform actual filtering - switch (filter) { - case STBI__F_none: - memcpy(cur, raw, nk); - break; - case STBI__F_sub: - memcpy(cur, raw, filter_bytes); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + cur[k-filter_bytes]); - break; - case STBI__F_up: - for (k = 0; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + prior[k]); - break; - case STBI__F_avg: - for (k = 0; k < filter_bytes; ++k) - cur[k] = STBI__BYTECAST(raw[k] + (prior[k]>>1)); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + ((prior[k] + cur[k-filter_bytes])>>1)); - break; - case STBI__F_paeth: - for (k = 0; k < filter_bytes; ++k) - cur[k] = STBI__BYTECAST(raw[k] + prior[k]); // prior[k] == stbi__paeth(0,prior[k],0) - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + stbi__paeth(cur[k-filter_bytes], prior[k], prior[k-filter_bytes])); - break; - case STBI__F_avg_first: - memcpy(cur, raw, filter_bytes); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + (cur[k-filter_bytes] >> 1)); - break; - } - - raw += nk; - - // expand decoded bits in cur to dest, also adding an extra alpha channel if desired - if (depth < 8) { - stbi_uc scale = (color == 0) ? stbi__depth_scale_table[depth] : 1; // scale grayscale values to 0..255 range - stbi_uc *in = cur; - stbi_uc *out = dest; - stbi_uc inb = 0; - stbi__uint32 nsmp = x*img_n; - - // expand bits to bytes first - if (depth == 4) { - for (i=0; i < nsmp; ++i) { - if ((i & 1) == 0) inb = *in++; - *out++ = scale * (inb >> 4); - inb <<= 4; - } - } else if (depth == 2) { - for (i=0; i < nsmp; ++i) { - if ((i & 3) == 0) inb = *in++; - *out++ = scale * (inb >> 6); - inb <<= 2; - } - } else { - STBI_ASSERT(depth == 1); - for (i=0; i < nsmp; ++i) { - if ((i & 7) == 0) inb = *in++; - *out++ = scale * (inb >> 7); - inb <<= 1; - } - } - - // insert alpha=255 values if desired - if (img_n != out_n) - stbi__create_png_alpha_expand8(dest, dest, x, img_n); - } else if (depth == 8) { - if (img_n == out_n) - memcpy(dest, cur, x*img_n); - else - stbi__create_png_alpha_expand8(dest, cur, x, img_n); - } else if (depth == 16) { - // convert the image data from big-endian to platform-native - stbi__uint16 *dest16 = (stbi__uint16*)dest; - stbi__uint32 nsmp = x*img_n; - - if (img_n == out_n) { - for (i = 0; i < nsmp; ++i, ++dest16, cur += 2) - *dest16 = (cur[0] << 8) | cur[1]; - } else { - STBI_ASSERT(img_n+1 == out_n); - if (img_n == 1) { - for (i = 0; i < x; ++i, dest16 += 2, cur += 2) { - dest16[0] = (cur[0] << 8) | cur[1]; - dest16[1] = 0xffff; - } - } else { - STBI_ASSERT(img_n == 3); - for (i = 0; i < x; ++i, dest16 += 4, cur += 6) { - dest16[0] = (cur[0] << 8) | cur[1]; - dest16[1] = (cur[2] << 8) | cur[3]; - dest16[2] = (cur[4] << 8) | cur[5]; - dest16[3] = 0xffff; - } - } - } - } - } - - STBI_FREE(filter_buf); - if (!all_ok) return 0; - - return 1; -} - -static int stbi__create_png_image(stbi__png *a, stbi_uc *image_data, stbi__uint32 image_data_len, int out_n, int depth, int color, int interlaced) -{ - int bytes = (depth == 16 ? 2 : 1); - int out_bytes = out_n * bytes; - stbi_uc *final; - int p; - if (!interlaced) - return stbi__create_png_image_raw(a, image_data, image_data_len, out_n, a->s->img_x, a->s->img_y, depth, color); - - // de-interlacing - final = (stbi_uc *) stbi__malloc_mad3(a->s->img_x, a->s->img_y, out_bytes, 0); - if (!final) return stbi__err("outofmem", "Out of memory"); - for (p=0; p < 7; ++p) { - int xorig[] = { 0,4,0,2,0,1,0 }; - int yorig[] = { 0,0,4,0,2,0,1 }; - int xspc[] = { 8,8,4,4,2,2,1 }; - int yspc[] = { 8,8,8,4,4,2,2 }; - int i,j,x,y; - // pass1_x[4] = 0, pass1_x[5] = 1, pass1_x[12] = 1 - x = (a->s->img_x - xorig[p] + xspc[p]-1) / xspc[p]; - y = (a->s->img_y - yorig[p] + yspc[p]-1) / yspc[p]; - if (x && y) { - stbi__uint32 img_len = ((((a->s->img_n * x * depth) + 7) >> 3) + 1) * y; - if (!stbi__create_png_image_raw(a, image_data, image_data_len, out_n, x, y, depth, color)) { - STBI_FREE(final); - return 0; - } - for (j=0; j < y; ++j) { - for (i=0; i < x; ++i) { - int out_y = j*yspc[p]+yorig[p]; - int out_x = i*xspc[p]+xorig[p]; - memcpy(final + out_y*a->s->img_x*out_bytes + out_x*out_bytes, - a->out + (j*x+i)*out_bytes, out_bytes); - } - } - STBI_FREE(a->out); - image_data += img_len; - image_data_len -= img_len; - } - } - a->out = final; - - return 1; -} - -static int stbi__compute_transparency(stbi__png *z, stbi_uc tc[3], int out_n) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi_uc *p = z->out; - - // compute color-based transparency, assuming we've - // already got 255 as the alpha value in the output - STBI_ASSERT(out_n == 2 || out_n == 4); - - if (out_n == 2) { - for (i=0; i < pixel_count; ++i) { - p[1] = (p[0] == tc[0] ? 0 : 255); - p += 2; - } - } else { - for (i=0; i < pixel_count; ++i) { - if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) - p[3] = 0; - p += 4; - } - } - return 1; -} - -static int stbi__compute_transparency16(stbi__png *z, stbi__uint16 tc[3], int out_n) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi__uint16 *p = (stbi__uint16*) z->out; - - // compute color-based transparency, assuming we've - // already got 65535 as the alpha value in the output - STBI_ASSERT(out_n == 2 || out_n == 4); - - if (out_n == 2) { - for (i = 0; i < pixel_count; ++i) { - p[1] = (p[0] == tc[0] ? 0 : 65535); - p += 2; - } - } else { - for (i = 0; i < pixel_count; ++i) { - if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) - p[3] = 0; - p += 4; - } - } - return 1; -} - -static int stbi__expand_png_palette(stbi__png *a, stbi_uc *palette, int len, int pal_img_n) -{ - stbi__uint32 i, pixel_count = a->s->img_x * a->s->img_y; - stbi_uc *p, *temp_out, *orig = a->out; - - p = (stbi_uc *) stbi__malloc_mad2(pixel_count, pal_img_n, 0); - if (p == NULL) return stbi__err("outofmem", "Out of memory"); - - // between here and free(out) below, exitting would leak - temp_out = p; - - if (pal_img_n == 3) { - for (i=0; i < pixel_count; ++i) { - int n = orig[i]*4; - p[0] = palette[n ]; - p[1] = palette[n+1]; - p[2] = palette[n+2]; - p += 3; - } - } else { - for (i=0; i < pixel_count; ++i) { - int n = orig[i]*4; - p[0] = palette[n ]; - p[1] = palette[n+1]; - p[2] = palette[n+2]; - p[3] = palette[n+3]; - p += 4; - } - } - STBI_FREE(a->out); - a->out = temp_out; - - STBI_NOTUSED(len); - - return 1; -} - -static int stbi__unpremultiply_on_load_global = 0; -static int stbi__de_iphone_flag_global = 0; - -STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply) -{ - stbi__unpremultiply_on_load_global = flag_true_if_should_unpremultiply; -} - -STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert) -{ - stbi__de_iphone_flag_global = flag_true_if_should_convert; -} - -#ifndef STBI_THREAD_LOCAL -#define stbi__unpremultiply_on_load stbi__unpremultiply_on_load_global -#define stbi__de_iphone_flag stbi__de_iphone_flag_global -#else -static STBI_THREAD_LOCAL int stbi__unpremultiply_on_load_local, stbi__unpremultiply_on_load_set; -static STBI_THREAD_LOCAL int stbi__de_iphone_flag_local, stbi__de_iphone_flag_set; - -STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply) -{ - stbi__unpremultiply_on_load_local = flag_true_if_should_unpremultiply; - stbi__unpremultiply_on_load_set = 1; -} - -STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert) -{ - stbi__de_iphone_flag_local = flag_true_if_should_convert; - stbi__de_iphone_flag_set = 1; -} - -#define stbi__unpremultiply_on_load (stbi__unpremultiply_on_load_set \ - ? stbi__unpremultiply_on_load_local \ - : stbi__unpremultiply_on_load_global) -#define stbi__de_iphone_flag (stbi__de_iphone_flag_set \ - ? stbi__de_iphone_flag_local \ - : stbi__de_iphone_flag_global) -#endif // STBI_THREAD_LOCAL - -static void stbi__de_iphone(stbi__png *z) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi_uc *p = z->out; - - if (s->img_out_n == 3) { // convert bgr to rgb - for (i=0; i < pixel_count; ++i) { - stbi_uc t = p[0]; - p[0] = p[2]; - p[2] = t; - p += 3; - } - } else { - STBI_ASSERT(s->img_out_n == 4); - if (stbi__unpremultiply_on_load) { - // convert bgr to rgb and unpremultiply - for (i=0; i < pixel_count; ++i) { - stbi_uc a = p[3]; - stbi_uc t = p[0]; - if (a) { - stbi_uc half = a / 2; - p[0] = (p[2] * 255 + half) / a; - p[1] = (p[1] * 255 + half) / a; - p[2] = ( t * 255 + half) / a; - } else { - p[0] = p[2]; - p[2] = t; - } - p += 4; - } - } else { - // convert bgr to rgb - for (i=0; i < pixel_count; ++i) { - stbi_uc t = p[0]; - p[0] = p[2]; - p[2] = t; - p += 4; - } - } - } -} - -#define STBI__PNG_TYPE(a,b,c,d) (((unsigned) (a) << 24) + ((unsigned) (b) << 16) + ((unsigned) (c) << 8) + (unsigned) (d)) - -static int stbi__parse_png_file(stbi__png *z, int scan, int req_comp) -{ - stbi_uc palette[1024], pal_img_n=0; - stbi_uc has_trans=0, tc[3]={0}; - stbi__uint16 tc16[3]; - stbi__uint32 ioff=0, idata_limit=0, i, pal_len=0; - int first=1,k,interlace=0, color=0, is_iphone=0; - stbi__context *s = z->s; - - z->expanded = NULL; - z->idata = NULL; - z->out = NULL; - - if (!stbi__check_png_header(s)) return 0; - - if (scan == STBI__SCAN_type) return 1; - - for (;;) { - stbi__pngchunk c = stbi__get_chunk_header(s); - switch (c.type) { - case STBI__PNG_TYPE('C','g','B','I'): - is_iphone = 1; - stbi__skip(s, c.length); - break; - case STBI__PNG_TYPE('I','H','D','R'): { - int comp,filter; - if (!first) return stbi__err("multiple IHDR","Corrupt PNG"); - first = 0; - if (c.length != 13) return stbi__err("bad IHDR len","Corrupt PNG"); - s->img_x = stbi__get32be(s); - s->img_y = stbi__get32be(s); - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - z->depth = stbi__get8(s); if (z->depth != 1 && z->depth != 2 && z->depth != 4 && z->depth != 8 && z->depth != 16) return stbi__err("1/2/4/8/16-bit only","PNG not supported: 1/2/4/8/16-bit only"); - color = stbi__get8(s); if (color > 6) return stbi__err("bad ctype","Corrupt PNG"); - if (color == 3 && z->depth == 16) return stbi__err("bad ctype","Corrupt PNG"); - if (color == 3) pal_img_n = 3; else if (color & 1) return stbi__err("bad ctype","Corrupt PNG"); - comp = stbi__get8(s); if (comp) return stbi__err("bad comp method","Corrupt PNG"); - filter= stbi__get8(s); if (filter) return stbi__err("bad filter method","Corrupt PNG"); - interlace = stbi__get8(s); if (interlace>1) return stbi__err("bad interlace method","Corrupt PNG"); - if (!s->img_x || !s->img_y) return stbi__err("0-pixel image","Corrupt PNG"); - if (!pal_img_n) { - s->img_n = (color & 2 ? 3 : 1) + (color & 4 ? 1 : 0); - if ((1 << 30) / s->img_x / s->img_n < s->img_y) return stbi__err("too large", "Image too large to decode"); - } else { - // if paletted, then pal_n is our final components, and - // img_n is # components to decompress/filter. - s->img_n = 1; - if ((1 << 30) / s->img_x / 4 < s->img_y) return stbi__err("too large","Corrupt PNG"); - } - // even with SCAN_header, have to scan to see if we have a tRNS - break; - } - - case STBI__PNG_TYPE('P','L','T','E'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (c.length > 256*3) return stbi__err("invalid PLTE","Corrupt PNG"); - pal_len = c.length / 3; - if (pal_len * 3 != c.length) return stbi__err("invalid PLTE","Corrupt PNG"); - for (i=0; i < pal_len; ++i) { - palette[i*4+0] = stbi__get8(s); - palette[i*4+1] = stbi__get8(s); - palette[i*4+2] = stbi__get8(s); - palette[i*4+3] = 255; - } - break; - } - - case STBI__PNG_TYPE('t','R','N','S'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (z->idata) return stbi__err("tRNS after IDAT","Corrupt PNG"); - if (pal_img_n) { - if (scan == STBI__SCAN_header) { s->img_n = 4; return 1; } - if (pal_len == 0) return stbi__err("tRNS before PLTE","Corrupt PNG"); - if (c.length > pal_len) return stbi__err("bad tRNS len","Corrupt PNG"); - pal_img_n = 4; - for (i=0; i < c.length; ++i) - palette[i*4+3] = stbi__get8(s); - } else { - if (!(s->img_n & 1)) return stbi__err("tRNS with alpha","Corrupt PNG"); - if (c.length != (stbi__uint32) s->img_n*2) return stbi__err("bad tRNS len","Corrupt PNG"); - has_trans = 1; - // non-paletted with tRNS = constant alpha. if header-scanning, we can stop now. - if (scan == STBI__SCAN_header) { ++s->img_n; return 1; } - if (z->depth == 16) { - for (k = 0; k < s->img_n && k < 3; ++k) // extra loop test to suppress false GCC warning - tc16[k] = (stbi__uint16)stbi__get16be(s); // copy the values as-is - } else { - for (k = 0; k < s->img_n && k < 3; ++k) - tc[k] = (stbi_uc)(stbi__get16be(s) & 255) * stbi__depth_scale_table[z->depth]; // non 8-bit images will be larger - } - } - break; - } - - case STBI__PNG_TYPE('I','D','A','T'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (pal_img_n && !pal_len) return stbi__err("no PLTE","Corrupt PNG"); - if (scan == STBI__SCAN_header) { - // header scan definitely stops at first IDAT - if (pal_img_n) - s->img_n = pal_img_n; - return 1; - } - if (c.length > (1u << 30)) return stbi__err("IDAT size limit", "IDAT section larger than 2^30 bytes"); - if ((int)(ioff + c.length) < (int)ioff) return 0; - if (ioff + c.length > idata_limit) { - stbi__uint32 idata_limit_old = idata_limit; - stbi_uc *p; - if (idata_limit == 0) idata_limit = c.length > 4096 ? c.length : 4096; - while (ioff + c.length > idata_limit) - idata_limit *= 2; - STBI_NOTUSED(idata_limit_old); - p = (stbi_uc *) STBI_REALLOC_SIZED(z->idata, idata_limit_old, idata_limit); if (p == NULL) return stbi__err("outofmem", "Out of memory"); - z->idata = p; - } - if (!stbi__getn(s, z->idata+ioff,c.length)) return stbi__err("outofdata","Corrupt PNG"); - ioff += c.length; - break; - } - - case STBI__PNG_TYPE('I','E','N','D'): { - stbi__uint32 raw_len, bpl; - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (scan != STBI__SCAN_load) return 1; - if (z->idata == NULL) return stbi__err("no IDAT","Corrupt PNG"); - // initial guess for decoded data size to avoid unnecessary reallocs - bpl = (s->img_x * z->depth + 7) / 8; // bytes per line, per component - raw_len = bpl * s->img_y * s->img_n /* pixels */ + s->img_y /* filter mode per row */; - z->expanded = (stbi_uc *) stbi_zlib_decode_malloc_guesssize_headerflag((char *) z->idata, ioff, raw_len, (int *) &raw_len, !is_iphone); - if (z->expanded == NULL) return 0; // zlib should set error - STBI_FREE(z->idata); z->idata = NULL; - if ((req_comp == s->img_n+1 && req_comp != 3 && !pal_img_n) || has_trans) - s->img_out_n = s->img_n+1; - else - s->img_out_n = s->img_n; - if (!stbi__create_png_image(z, z->expanded, raw_len, s->img_out_n, z->depth, color, interlace)) return 0; - if (has_trans) { - if (z->depth == 16) { - if (!stbi__compute_transparency16(z, tc16, s->img_out_n)) return 0; - } else { - if (!stbi__compute_transparency(z, tc, s->img_out_n)) return 0; - } - } - if (is_iphone && stbi__de_iphone_flag && s->img_out_n > 2) - stbi__de_iphone(z); - if (pal_img_n) { - // pal_img_n == 3 or 4 - s->img_n = pal_img_n; // record the actual colors we had - s->img_out_n = pal_img_n; - if (req_comp >= 3) s->img_out_n = req_comp; - if (!stbi__expand_png_palette(z, palette, pal_len, s->img_out_n)) - return 0; - } else if (has_trans) { - // non-paletted image with tRNS -> source image has (constant) alpha - ++s->img_n; - } - STBI_FREE(z->expanded); z->expanded = NULL; - // end of PNG chunk, read and skip CRC - stbi__get32be(s); - return 1; - } - - default: - // if critical, fail - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if ((c.type & (1 << 29)) == 0) { - #ifndef STBI_NO_FAILURE_STRINGS - // not threadsafe - static char invalid_chunk[] = "XXXX PNG chunk not known"; - invalid_chunk[0] = STBI__BYTECAST(c.type >> 24); - invalid_chunk[1] = STBI__BYTECAST(c.type >> 16); - invalid_chunk[2] = STBI__BYTECAST(c.type >> 8); - invalid_chunk[3] = STBI__BYTECAST(c.type >> 0); - #endif - return stbi__err(invalid_chunk, "PNG not supported: unknown PNG chunk type"); - } - stbi__skip(s, c.length); - break; - } - // end of PNG chunk, read and skip CRC - stbi__get32be(s); - } -} - -static void *stbi__do_png(stbi__png *p, int *x, int *y, int *n, int req_comp, stbi__result_info *ri) -{ - void *result=NULL; - if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); - if (stbi__parse_png_file(p, STBI__SCAN_load, req_comp)) { - if (p->depth <= 8) - ri->bits_per_channel = 8; - else if (p->depth == 16) - ri->bits_per_channel = 16; - else - return stbi__errpuc("bad bits_per_channel", "PNG not supported: unsupported color depth"); - result = p->out; - p->out = NULL; - if (req_comp && req_comp != p->s->img_out_n) { - if (ri->bits_per_channel == 8) - result = stbi__convert_format((unsigned char *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); - else - result = stbi__convert_format16((stbi__uint16 *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); - p->s->img_out_n = req_comp; - if (result == NULL) return result; - } - *x = p->s->img_x; - *y = p->s->img_y; - if (n) *n = p->s->img_n; - } - STBI_FREE(p->out); p->out = NULL; - STBI_FREE(p->expanded); p->expanded = NULL; - STBI_FREE(p->idata); p->idata = NULL; - - return result; -} - -static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi__png p; - p.s = s; - return stbi__do_png(&p, x,y,comp,req_comp, ri); -} - -static int stbi__png_test(stbi__context *s) -{ - int r; - r = stbi__check_png_header(s); - stbi__rewind(s); - return r; -} - -static int stbi__png_info_raw(stbi__png *p, int *x, int *y, int *comp) -{ - if (!stbi__parse_png_file(p, STBI__SCAN_header, 0)) { - stbi__rewind( p->s ); - return 0; - } - if (x) *x = p->s->img_x; - if (y) *y = p->s->img_y; - if (comp) *comp = p->s->img_n; - return 1; -} - -static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp) -{ - stbi__png p; - p.s = s; - return stbi__png_info_raw(&p, x, y, comp); -} - -static int stbi__png_is16(stbi__context *s) -{ - stbi__png p; - p.s = s; - if (!stbi__png_info_raw(&p, NULL, NULL, NULL)) - return 0; - if (p.depth != 16) { - stbi__rewind(p.s); - return 0; - } - return 1; -} -#endif - -// Microsoft/Windows BMP image - -#ifndef STBI_NO_BMP -static int stbi__bmp_test_raw(stbi__context *s) -{ - int r; - int sz; - if (stbi__get8(s) != 'B') return 0; - if (stbi__get8(s) != 'M') return 0; - stbi__get32le(s); // discard filesize - stbi__get16le(s); // discard reserved - stbi__get16le(s); // discard reserved - stbi__get32le(s); // discard data offset - sz = stbi__get32le(s); - r = (sz == 12 || sz == 40 || sz == 56 || sz == 108 || sz == 124); - return r; -} - -static int stbi__bmp_test(stbi__context *s) -{ - int r = stbi__bmp_test_raw(s); - stbi__rewind(s); - return r; -} - - -// returns 0..31 for the highest set bit -static int stbi__high_bit(unsigned int z) -{ - int n=0; - if (z == 0) return -1; - if (z >= 0x10000) { n += 16; z >>= 16; } - if (z >= 0x00100) { n += 8; z >>= 8; } - if (z >= 0x00010) { n += 4; z >>= 4; } - if (z >= 0x00004) { n += 2; z >>= 2; } - if (z >= 0x00002) { n += 1;/* >>= 1;*/ } - return n; -} - -static int stbi__bitcount(unsigned int a) -{ - a = (a & 0x55555555) + ((a >> 1) & 0x55555555); // max 2 - a = (a & 0x33333333) + ((a >> 2) & 0x33333333); // max 4 - a = (a + (a >> 4)) & 0x0f0f0f0f; // max 8 per 4, now 8 bits - a = (a + (a >> 8)); // max 16 per 8 bits - a = (a + (a >> 16)); // max 32 per 8 bits - return a & 0xff; -} - -// extract an arbitrarily-aligned N-bit value (N=bits) -// from v, and then make it 8-bits long and fractionally -// extend it to full full range. -static int stbi__shiftsigned(unsigned int v, int shift, int bits) -{ - static unsigned int mul_table[9] = { - 0, - 0xff/*0b11111111*/, 0x55/*0b01010101*/, 0x49/*0b01001001*/, 0x11/*0b00010001*/, - 0x21/*0b00100001*/, 0x41/*0b01000001*/, 0x81/*0b10000001*/, 0x01/*0b00000001*/, - }; - static unsigned int shift_table[9] = { - 0, 0,0,1,0,2,4,6,0, - }; - if (shift < 0) - v <<= -shift; - else - v >>= shift; - STBI_ASSERT(v < 256); - v >>= (8-bits); - STBI_ASSERT(bits >= 0 && bits <= 8); - return (int) ((unsigned) v * mul_table[bits]) >> shift_table[bits]; -} - -typedef struct -{ - int bpp, offset, hsz; - unsigned int mr,mg,mb,ma, all_a; - int extra_read; -} stbi__bmp_data; - -static int stbi__bmp_set_mask_defaults(stbi__bmp_data *info, int compress) -{ - // BI_BITFIELDS specifies masks explicitly, don't override - if (compress == 3) - return 1; - - if (compress == 0) { - if (info->bpp == 16) { - info->mr = 31u << 10; - info->mg = 31u << 5; - info->mb = 31u << 0; - } else if (info->bpp == 32) { - info->mr = 0xffu << 16; - info->mg = 0xffu << 8; - info->mb = 0xffu << 0; - info->ma = 0xffu << 24; - info->all_a = 0; // if all_a is 0 at end, then we loaded alpha channel but it was all 0 - } else { - // otherwise, use defaults, which is all-0 - info->mr = info->mg = info->mb = info->ma = 0; - } - return 1; - } - return 0; // error -} - -static void *stbi__bmp_parse_header(stbi__context *s, stbi__bmp_data *info) -{ - int hsz; - if (stbi__get8(s) != 'B' || stbi__get8(s) != 'M') return stbi__errpuc("not BMP", "Corrupt BMP"); - stbi__get32le(s); // discard filesize - stbi__get16le(s); // discard reserved - stbi__get16le(s); // discard reserved - info->offset = stbi__get32le(s); - info->hsz = hsz = stbi__get32le(s); - info->mr = info->mg = info->mb = info->ma = 0; - info->extra_read = 14; - - if (info->offset < 0) return stbi__errpuc("bad BMP", "bad BMP"); - - if (hsz != 12 && hsz != 40 && hsz != 56 && hsz != 108 && hsz != 124) return stbi__errpuc("unknown BMP", "BMP type not supported: unknown"); - if (hsz == 12) { - s->img_x = stbi__get16le(s); - s->img_y = stbi__get16le(s); - } else { - s->img_x = stbi__get32le(s); - s->img_y = stbi__get32le(s); - } - if (stbi__get16le(s) != 1) return stbi__errpuc("bad BMP", "bad BMP"); - info->bpp = stbi__get16le(s); - if (hsz != 12) { - int compress = stbi__get32le(s); - if (compress == 1 || compress == 2) return stbi__errpuc("BMP RLE", "BMP type not supported: RLE"); - if (compress >= 4) return stbi__errpuc("BMP JPEG/PNG", "BMP type not supported: unsupported compression"); // this includes PNG/JPEG modes - if (compress == 3 && info->bpp != 16 && info->bpp != 32) return stbi__errpuc("bad BMP", "bad BMP"); // bitfields requires 16 or 32 bits/pixel - stbi__get32le(s); // discard sizeof - stbi__get32le(s); // discard hres - stbi__get32le(s); // discard vres - stbi__get32le(s); // discard colorsused - stbi__get32le(s); // discard max important - if (hsz == 40 || hsz == 56) { - if (hsz == 56) { - stbi__get32le(s); - stbi__get32le(s); - stbi__get32le(s); - stbi__get32le(s); - } - if (info->bpp == 16 || info->bpp == 32) { - if (compress == 0) { - stbi__bmp_set_mask_defaults(info, compress); - } else if (compress == 3) { - info->mr = stbi__get32le(s); - info->mg = stbi__get32le(s); - info->mb = stbi__get32le(s); - info->extra_read += 12; - // not documented, but generated by photoshop and handled by mspaint - if (info->mr == info->mg && info->mg == info->mb) { - // ?!?!? - return stbi__errpuc("bad BMP", "bad BMP"); - } - } else - return stbi__errpuc("bad BMP", "bad BMP"); - } - } else { - // V4/V5 header - int i; - if (hsz != 108 && hsz != 124) - return stbi__errpuc("bad BMP", "bad BMP"); - info->mr = stbi__get32le(s); - info->mg = stbi__get32le(s); - info->mb = stbi__get32le(s); - info->ma = stbi__get32le(s); - if (compress != 3) // override mr/mg/mb unless in BI_BITFIELDS mode, as per docs - stbi__bmp_set_mask_defaults(info, compress); - stbi__get32le(s); // discard color space - for (i=0; i < 12; ++i) - stbi__get32le(s); // discard color space parameters - if (hsz == 124) { - stbi__get32le(s); // discard rendering intent - stbi__get32le(s); // discard offset of profile data - stbi__get32le(s); // discard size of profile data - stbi__get32le(s); // discard reserved - } - } - } - return (void *) 1; -} - - -static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *out; - unsigned int mr=0,mg=0,mb=0,ma=0, all_a; - stbi_uc pal[256][4]; - int psize=0,i,j,width; - int flip_vertically, pad, target; - stbi__bmp_data info; - STBI_NOTUSED(ri); - - info.all_a = 255; - if (stbi__bmp_parse_header(s, &info) == NULL) - return NULL; // error code already set - - flip_vertically = ((int) s->img_y) > 0; - s->img_y = abs((int) s->img_y); - - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - mr = info.mr; - mg = info.mg; - mb = info.mb; - ma = info.ma; - all_a = info.all_a; - - if (info.hsz == 12) { - if (info.bpp < 24) - psize = (info.offset - info.extra_read - 24) / 3; - } else { - if (info.bpp < 16) - psize = (info.offset - info.extra_read - info.hsz) >> 2; - } - if (psize == 0) { - // accept some number of extra bytes after the header, but if the offset points either to before - // the header ends or implies a large amount of extra data, reject the file as malformed - int bytes_read_so_far = s->callback_already_read + (int)(s->img_buffer - s->img_buffer_original); - int header_limit = 1024; // max we actually read is below 256 bytes currently. - int extra_data_limit = 256*4; // what ordinarily goes here is a palette; 256 entries*4 bytes is its max size. - if (bytes_read_so_far <= 0 || bytes_read_so_far > header_limit) { - return stbi__errpuc("bad header", "Corrupt BMP"); - } - // we established that bytes_read_so_far is positive and sensible. - // the first half of this test rejects offsets that are either too small positives, or - // negative, and guarantees that info.offset >= bytes_read_so_far > 0. this in turn - // ensures the number computed in the second half of the test can't overflow. - if (info.offset < bytes_read_so_far || info.offset - bytes_read_so_far > extra_data_limit) { - return stbi__errpuc("bad offset", "Corrupt BMP"); - } else { - stbi__skip(s, info.offset - bytes_read_so_far); - } - } - - if (info.bpp == 24 && ma == 0xff000000) - s->img_n = 3; - else - s->img_n = ma ? 4 : 3; - if (req_comp && req_comp >= 3) // we can directly decode 3 or 4 - target = req_comp; - else - target = s->img_n; // if they want monochrome, we'll post-convert - - // sanity-check size - if (!stbi__mad3sizes_valid(target, s->img_x, s->img_y, 0)) - return stbi__errpuc("too large", "Corrupt BMP"); - - out = (stbi_uc *) stbi__malloc_mad3(target, s->img_x, s->img_y, 0); - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - if (info.bpp < 16) { - int z=0; - if (psize == 0 || psize > 256) { STBI_FREE(out); return stbi__errpuc("invalid", "Corrupt BMP"); } - for (i=0; i < psize; ++i) { - pal[i][2] = stbi__get8(s); - pal[i][1] = stbi__get8(s); - pal[i][0] = stbi__get8(s); - if (info.hsz != 12) stbi__get8(s); - pal[i][3] = 255; - } - stbi__skip(s, info.offset - info.extra_read - info.hsz - psize * (info.hsz == 12 ? 3 : 4)); - if (info.bpp == 1) width = (s->img_x + 7) >> 3; - else if (info.bpp == 4) width = (s->img_x + 1) >> 1; - else if (info.bpp == 8) width = s->img_x; - else { STBI_FREE(out); return stbi__errpuc("bad bpp", "Corrupt BMP"); } - pad = (-width)&3; - if (info.bpp == 1) { - for (j=0; j < (int) s->img_y; ++j) { - int bit_offset = 7, v = stbi__get8(s); - for (i=0; i < (int) s->img_x; ++i) { - int color = (v>>bit_offset)&0x1; - out[z++] = pal[color][0]; - out[z++] = pal[color][1]; - out[z++] = pal[color][2]; - if (target == 4) out[z++] = 255; - if (i+1 == (int) s->img_x) break; - if((--bit_offset) < 0) { - bit_offset = 7; - v = stbi__get8(s); - } - } - stbi__skip(s, pad); - } - } else { - for (j=0; j < (int) s->img_y; ++j) { - for (i=0; i < (int) s->img_x; i += 2) { - int v=stbi__get8(s),v2=0; - if (info.bpp == 4) { - v2 = v & 15; - v >>= 4; - } - out[z++] = pal[v][0]; - out[z++] = pal[v][1]; - out[z++] = pal[v][2]; - if (target == 4) out[z++] = 255; - if (i+1 == (int) s->img_x) break; - v = (info.bpp == 8) ? stbi__get8(s) : v2; - out[z++] = pal[v][0]; - out[z++] = pal[v][1]; - out[z++] = pal[v][2]; - if (target == 4) out[z++] = 255; - } - stbi__skip(s, pad); - } - } - } else { - int rshift=0,gshift=0,bshift=0,ashift=0,rcount=0,gcount=0,bcount=0,acount=0; - int z = 0; - int easy=0; - stbi__skip(s, info.offset - info.extra_read - info.hsz); - if (info.bpp == 24) width = 3 * s->img_x; - else if (info.bpp == 16) width = 2*s->img_x; - else /* bpp = 32 and pad = 0 */ width=0; - pad = (-width) & 3; - if (info.bpp == 24) { - easy = 1; - } else if (info.bpp == 32) { - if (mb == 0xff && mg == 0xff00 && mr == 0x00ff0000 && ma == 0xff000000) - easy = 2; - } - if (!easy) { - if (!mr || !mg || !mb) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } - // right shift amt to put high bit in position #7 - rshift = stbi__high_bit(mr)-7; rcount = stbi__bitcount(mr); - gshift = stbi__high_bit(mg)-7; gcount = stbi__bitcount(mg); - bshift = stbi__high_bit(mb)-7; bcount = stbi__bitcount(mb); - ashift = stbi__high_bit(ma)-7; acount = stbi__bitcount(ma); - if (rcount > 8 || gcount > 8 || bcount > 8 || acount > 8) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } - } - for (j=0; j < (int) s->img_y; ++j) { - if (easy) { - for (i=0; i < (int) s->img_x; ++i) { - unsigned char a; - out[z+2] = stbi__get8(s); - out[z+1] = stbi__get8(s); - out[z+0] = stbi__get8(s); - z += 3; - a = (easy == 2 ? stbi__get8(s) : 255); - all_a |= a; - if (target == 4) out[z++] = a; - } - } else { - int bpp = info.bpp; - for (i=0; i < (int) s->img_x; ++i) { - stbi__uint32 v = (bpp == 16 ? (stbi__uint32) stbi__get16le(s) : stbi__get32le(s)); - unsigned int a; - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mr, rshift, rcount)); - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mg, gshift, gcount)); - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mb, bshift, bcount)); - a = (ma ? stbi__shiftsigned(v & ma, ashift, acount) : 255); - all_a |= a; - if (target == 4) out[z++] = STBI__BYTECAST(a); - } - } - stbi__skip(s, pad); - } - } - - // if alpha channel is all 0s, replace with all 255s - if (target == 4 && all_a == 0) - for (i=4*s->img_x*s->img_y-1; i >= 0; i -= 4) - out[i] = 255; - - if (flip_vertically) { - stbi_uc t; - for (j=0; j < (int) s->img_y>>1; ++j) { - stbi_uc *p1 = out + j *s->img_x*target; - stbi_uc *p2 = out + (s->img_y-1-j)*s->img_x*target; - for (i=0; i < (int) s->img_x*target; ++i) { - t = p1[i]; p1[i] = p2[i]; p2[i] = t; - } - } - } - - if (req_comp && req_comp != target) { - out = stbi__convert_format(out, target, req_comp, s->img_x, s->img_y); - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - - *x = s->img_x; - *y = s->img_y; - if (comp) *comp = s->img_n; - return out; -} -#endif - -// Targa Truevision - TGA -// by Jonathan Dummer -#ifndef STBI_NO_TGA -// returns STBI_rgb or whatever, 0 on error -static int stbi__tga_get_comp(int bits_per_pixel, int is_grey, int* is_rgb16) -{ - // only RGB or RGBA (incl. 16bit) or grey allowed - if (is_rgb16) *is_rgb16 = 0; - switch(bits_per_pixel) { - case 8: return STBI_grey; - case 16: if(is_grey) return STBI_grey_alpha; - // fallthrough - case 15: if(is_rgb16) *is_rgb16 = 1; - return STBI_rgb; - case 24: // fallthrough - case 32: return bits_per_pixel/8; - default: return 0; - } -} - -static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp) -{ - int tga_w, tga_h, tga_comp, tga_image_type, tga_bits_per_pixel, tga_colormap_bpp; - int sz, tga_colormap_type; - stbi__get8(s); // discard Offset - tga_colormap_type = stbi__get8(s); // colormap type - if( tga_colormap_type > 1 ) { - stbi__rewind(s); - return 0; // only RGB or indexed allowed - } - tga_image_type = stbi__get8(s); // image type - if ( tga_colormap_type == 1 ) { // colormapped (paletted) image - if (tga_image_type != 1 && tga_image_type != 9) { - stbi__rewind(s); - return 0; - } - stbi__skip(s,4); // skip index of first colormap entry and number of entries - sz = stbi__get8(s); // check bits per palette color entry - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) { - stbi__rewind(s); - return 0; - } - stbi__skip(s,4); // skip image x and y origin - tga_colormap_bpp = sz; - } else { // "normal" image w/o colormap - only RGB or grey allowed, +/- RLE - if ( (tga_image_type != 2) && (tga_image_type != 3) && (tga_image_type != 10) && (tga_image_type != 11) ) { - stbi__rewind(s); - return 0; // only RGB or grey allowed, +/- RLE - } - stbi__skip(s,9); // skip colormap specification and image x/y origin - tga_colormap_bpp = 0; - } - tga_w = stbi__get16le(s); - if( tga_w < 1 ) { - stbi__rewind(s); - return 0; // test width - } - tga_h = stbi__get16le(s); - if( tga_h < 1 ) { - stbi__rewind(s); - return 0; // test height - } - tga_bits_per_pixel = stbi__get8(s); // bits per pixel - stbi__get8(s); // ignore alpha bits - if (tga_colormap_bpp != 0) { - if((tga_bits_per_pixel != 8) && (tga_bits_per_pixel != 16)) { - // when using a colormap, tga_bits_per_pixel is the size of the indexes - // I don't think anything but 8 or 16bit indexes makes sense - stbi__rewind(s); - return 0; - } - tga_comp = stbi__tga_get_comp(tga_colormap_bpp, 0, NULL); - } else { - tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3) || (tga_image_type == 11), NULL); - } - if(!tga_comp) { - stbi__rewind(s); - return 0; - } - if (x) *x = tga_w; - if (y) *y = tga_h; - if (comp) *comp = tga_comp; - return 1; // seems to have passed everything -} - -static int stbi__tga_test(stbi__context *s) -{ - int res = 0; - int sz, tga_color_type; - stbi__get8(s); // discard Offset - tga_color_type = stbi__get8(s); // color type - if ( tga_color_type > 1 ) goto errorEnd; // only RGB or indexed allowed - sz = stbi__get8(s); // image type - if ( tga_color_type == 1 ) { // colormapped (paletted) image - if (sz != 1 && sz != 9) goto errorEnd; // colortype 1 demands image type 1 or 9 - stbi__skip(s,4); // skip index of first colormap entry and number of entries - sz = stbi__get8(s); // check bits per palette color entry - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; - stbi__skip(s,4); // skip image x and y origin - } else { // "normal" image w/o colormap - if ( (sz != 2) && (sz != 3) && (sz != 10) && (sz != 11) ) goto errorEnd; // only RGB or grey allowed, +/- RLE - stbi__skip(s,9); // skip colormap specification and image x/y origin - } - if ( stbi__get16le(s) < 1 ) goto errorEnd; // test width - if ( stbi__get16le(s) < 1 ) goto errorEnd; // test height - sz = stbi__get8(s); // bits per pixel - if ( (tga_color_type == 1) && (sz != 8) && (sz != 16) ) goto errorEnd; // for colormapped images, bpp is size of an index - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; - - res = 1; // if we got this far, everything's good and we can return 1 instead of 0 - -errorEnd: - stbi__rewind(s); - return res; -} - -// read 16bit value and convert to 24bit RGB -static void stbi__tga_read_rgb16(stbi__context *s, stbi_uc* out) -{ - stbi__uint16 px = (stbi__uint16)stbi__get16le(s); - stbi__uint16 fiveBitMask = 31; - // we have 3 channels with 5bits each - int r = (px >> 10) & fiveBitMask; - int g = (px >> 5) & fiveBitMask; - int b = px & fiveBitMask; - // Note that this saves the data in RGB(A) order, so it doesn't need to be swapped later - out[0] = (stbi_uc)((r * 255)/31); - out[1] = (stbi_uc)((g * 255)/31); - out[2] = (stbi_uc)((b * 255)/31); - - // some people claim that the most significant bit might be used for alpha - // (possibly if an alpha-bit is set in the "image descriptor byte") - // but that only made 16bit test images completely translucent.. - // so let's treat all 15 and 16bit TGAs as RGB with no alpha. -} - -static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - // read in the TGA header stuff - int tga_offset = stbi__get8(s); - int tga_indexed = stbi__get8(s); - int tga_image_type = stbi__get8(s); - int tga_is_RLE = 0; - int tga_palette_start = stbi__get16le(s); - int tga_palette_len = stbi__get16le(s); - int tga_palette_bits = stbi__get8(s); - int tga_x_origin = stbi__get16le(s); - int tga_y_origin = stbi__get16le(s); - int tga_width = stbi__get16le(s); - int tga_height = stbi__get16le(s); - int tga_bits_per_pixel = stbi__get8(s); - int tga_comp, tga_rgb16=0; - int tga_inverted = stbi__get8(s); - // int tga_alpha_bits = tga_inverted & 15; // the 4 lowest bits - unused (useless?) - // image data - unsigned char *tga_data; - unsigned char *tga_palette = NULL; - int i, j; - unsigned char raw_data[4] = {0}; - int RLE_count = 0; - int RLE_repeating = 0; - int read_next_pixel = 1; - STBI_NOTUSED(ri); - STBI_NOTUSED(tga_x_origin); // @TODO - STBI_NOTUSED(tga_y_origin); // @TODO - - if (tga_height > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (tga_width > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - // do a tiny bit of precessing - if ( tga_image_type >= 8 ) - { - tga_image_type -= 8; - tga_is_RLE = 1; - } - tga_inverted = 1 - ((tga_inverted >> 5) & 1); - - // If I'm paletted, then I'll use the number of bits from the palette - if ( tga_indexed ) tga_comp = stbi__tga_get_comp(tga_palette_bits, 0, &tga_rgb16); - else tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3), &tga_rgb16); - - if(!tga_comp) // shouldn't really happen, stbi__tga_test() should have ensured basic consistency - return stbi__errpuc("bad format", "Can't find out TGA pixelformat"); - - // tga info - *x = tga_width; - *y = tga_height; - if (comp) *comp = tga_comp; - - if (!stbi__mad3sizes_valid(tga_width, tga_height, tga_comp, 0)) - return stbi__errpuc("too large", "Corrupt TGA"); - - tga_data = (unsigned char*)stbi__malloc_mad3(tga_width, tga_height, tga_comp, 0); - if (!tga_data) return stbi__errpuc("outofmem", "Out of memory"); - - // skip to the data's starting position (offset usually = 0) - stbi__skip(s, tga_offset ); - - if ( !tga_indexed && !tga_is_RLE && !tga_rgb16 ) { - for (i=0; i < tga_height; ++i) { - int row = tga_inverted ? tga_height -i - 1 : i; - stbi_uc *tga_row = tga_data + row*tga_width*tga_comp; - stbi__getn(s, tga_row, tga_width * tga_comp); - } - } else { - // do I need to load a palette? - if ( tga_indexed) - { - if (tga_palette_len == 0) { /* you have to have at least one entry! */ - STBI_FREE(tga_data); - return stbi__errpuc("bad palette", "Corrupt TGA"); - } - - // any data to skip? (offset usually = 0) - stbi__skip(s, tga_palette_start ); - // load the palette - tga_palette = (unsigned char*)stbi__malloc_mad2(tga_palette_len, tga_comp, 0); - if (!tga_palette) { - STBI_FREE(tga_data); - return stbi__errpuc("outofmem", "Out of memory"); - } - if (tga_rgb16) { - stbi_uc *pal_entry = tga_palette; - STBI_ASSERT(tga_comp == STBI_rgb); - for (i=0; i < tga_palette_len; ++i) { - stbi__tga_read_rgb16(s, pal_entry); - pal_entry += tga_comp; - } - } else if (!stbi__getn(s, tga_palette, tga_palette_len * tga_comp)) { - STBI_FREE(tga_data); - STBI_FREE(tga_palette); - return stbi__errpuc("bad palette", "Corrupt TGA"); - } - } - // load the data - for (i=0; i < tga_width * tga_height; ++i) - { - // if I'm in RLE mode, do I need to get a RLE stbi__pngchunk? - if ( tga_is_RLE ) - { - if ( RLE_count == 0 ) - { - // yep, get the next byte as a RLE command - int RLE_cmd = stbi__get8(s); - RLE_count = 1 + (RLE_cmd & 127); - RLE_repeating = RLE_cmd >> 7; - read_next_pixel = 1; - } else if ( !RLE_repeating ) - { - read_next_pixel = 1; - } - } else - { - read_next_pixel = 1; - } - // OK, if I need to read a pixel, do it now - if ( read_next_pixel ) - { - // load however much data we did have - if ( tga_indexed ) - { - // read in index, then perform the lookup - int pal_idx = (tga_bits_per_pixel == 8) ? stbi__get8(s) : stbi__get16le(s); - if ( pal_idx >= tga_palette_len ) { - // invalid index - pal_idx = 0; - } - pal_idx *= tga_comp; - for (j = 0; j < tga_comp; ++j) { - raw_data[j] = tga_palette[pal_idx+j]; - } - } else if(tga_rgb16) { - STBI_ASSERT(tga_comp == STBI_rgb); - stbi__tga_read_rgb16(s, raw_data); - } else { - // read in the data raw - for (j = 0; j < tga_comp; ++j) { - raw_data[j] = stbi__get8(s); - } - } - // clear the reading flag for the next pixel - read_next_pixel = 0; - } // end of reading a pixel - - // copy data - for (j = 0; j < tga_comp; ++j) - tga_data[i*tga_comp+j] = raw_data[j]; - - // in case we're in RLE mode, keep counting down - --RLE_count; - } - // do I need to invert the image? - if ( tga_inverted ) - { - for (j = 0; j*2 < tga_height; ++j) - { - int index1 = j * tga_width * tga_comp; - int index2 = (tga_height - 1 - j) * tga_width * tga_comp; - for (i = tga_width * tga_comp; i > 0; --i) - { - unsigned char temp = tga_data[index1]; - tga_data[index1] = tga_data[index2]; - tga_data[index2] = temp; - ++index1; - ++index2; - } - } - } - // clear my palette, if I had one - if ( tga_palette != NULL ) - { - STBI_FREE( tga_palette ); - } - } - - // swap RGB - if the source data was RGB16, it already is in the right order - if (tga_comp >= 3 && !tga_rgb16) - { - unsigned char* tga_pixel = tga_data; - for (i=0; i < tga_width * tga_height; ++i) - { - unsigned char temp = tga_pixel[0]; - tga_pixel[0] = tga_pixel[2]; - tga_pixel[2] = temp; - tga_pixel += tga_comp; - } - } - - // convert to target component count - if (req_comp && req_comp != tga_comp) - tga_data = stbi__convert_format(tga_data, tga_comp, req_comp, tga_width, tga_height); - - // the things I do to get rid of an error message, and yet keep - // Microsoft's C compilers happy... [8^( - tga_palette_start = tga_palette_len = tga_palette_bits = - tga_x_origin = tga_y_origin = 0; - STBI_NOTUSED(tga_palette_start); - // OK, done - return tga_data; -} -#endif - -// ************************************************************************************************* -// Photoshop PSD loader -- PD by Thatcher Ulrich, integration by Nicolas Schulz, tweaked by STB - -#ifndef STBI_NO_PSD -static int stbi__psd_test(stbi__context *s) -{ - int r = (stbi__get32be(s) == 0x38425053); - stbi__rewind(s); - return r; -} - -static int stbi__psd_decode_rle(stbi__context *s, stbi_uc *p, int pixelCount) -{ - int count, nleft, len; - - count = 0; - while ((nleft = pixelCount - count) > 0) { - len = stbi__get8(s); - if (len == 128) { - // No-op. - } else if (len < 128) { - // Copy next len+1 bytes literally. - len++; - if (len > nleft) return 0; // corrupt data - count += len; - while (len) { - *p = stbi__get8(s); - p += 4; - len--; - } - } else if (len > 128) { - stbi_uc val; - // Next -len+1 bytes in the dest are replicated from next source byte. - // (Interpret len as a negative 8-bit int.) - len = 257 - len; - if (len > nleft) return 0; // corrupt data - val = stbi__get8(s); - count += len; - while (len) { - *p = val; - p += 4; - len--; - } - } - } - - return 1; -} - -static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) -{ - int pixelCount; - int channelCount, compression; - int channel, i; - int bitdepth; - int w,h; - stbi_uc *out; - STBI_NOTUSED(ri); - - // Check identifier - if (stbi__get32be(s) != 0x38425053) // "8BPS" - return stbi__errpuc("not PSD", "Corrupt PSD image"); - - // Check file type version. - if (stbi__get16be(s) != 1) - return stbi__errpuc("wrong version", "Unsupported version of PSD image"); - - // Skip 6 reserved bytes. - stbi__skip(s, 6 ); - - // Read the number of channels (R, G, B, A, etc). - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) - return stbi__errpuc("wrong channel count", "Unsupported number of channels in PSD image"); - - // Read the rows and columns of the image. - h = stbi__get32be(s); - w = stbi__get32be(s); - - if (h > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (w > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - // Make sure the depth is 8 bits. - bitdepth = stbi__get16be(s); - if (bitdepth != 8 && bitdepth != 16) - return stbi__errpuc("unsupported bit depth", "PSD bit depth is not 8 or 16 bit"); - - // Make sure the color mode is RGB. - // Valid options are: - // 0: Bitmap - // 1: Grayscale - // 2: Indexed color - // 3: RGB color - // 4: CMYK color - // 7: Multichannel - // 8: Duotone - // 9: Lab color - if (stbi__get16be(s) != 3) - return stbi__errpuc("wrong color format", "PSD is not in RGB color format"); - - // Skip the Mode Data. (It's the palette for indexed color; other info for other modes.) - stbi__skip(s,stbi__get32be(s) ); - - // Skip the image resources. (resolution, pen tool paths, etc) - stbi__skip(s, stbi__get32be(s) ); - - // Skip the reserved data. - stbi__skip(s, stbi__get32be(s) ); - - // Find out if the data is compressed. - // Known values: - // 0: no compression - // 1: RLE compressed - compression = stbi__get16be(s); - if (compression > 1) - return stbi__errpuc("bad compression", "PSD has an unknown compression format"); - - // Check size - if (!stbi__mad3sizes_valid(4, w, h, 0)) - return stbi__errpuc("too large", "Corrupt PSD"); - - // Create the destination image. - - if (!compression && bitdepth == 16 && bpc == 16) { - out = (stbi_uc *) stbi__malloc_mad3(8, w, h, 0); - ri->bits_per_channel = 16; - } else - out = (stbi_uc *) stbi__malloc(4 * w*h); - - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - pixelCount = w*h; - - // Initialize the data to zero. - //memset( out, 0, pixelCount * 4 ); - - // Finally, the image data. - if (compression) { - // RLE as used by .PSD and .TIFF - // Loop until you get the number of unpacked bytes you are expecting: - // Read the next source byte into n. - // If n is between 0 and 127 inclusive, copy the next n+1 bytes literally. - // Else if n is between -127 and -1 inclusive, copy the next byte -n+1 times. - // Else if n is 128, noop. - // Endloop - - // The RLE-compressed data is preceded by a 2-byte data count for each row in the data, - // which we're going to just skip. - stbi__skip(s, h * channelCount * 2 ); - - // Read the RLE data by channel. - for (channel = 0; channel < 4; channel++) { - stbi_uc *p; - - p = out+channel; - if (channel >= channelCount) { - // Fill this channel with default data. - for (i = 0; i < pixelCount; i++, p += 4) - *p = (channel == 3 ? 255 : 0); - } else { - // Read the RLE data. - if (!stbi__psd_decode_rle(s, p, pixelCount)) { - STBI_FREE(out); - return stbi__errpuc("corrupt", "bad RLE data"); - } - } - } - - } else { - // We're at the raw image data. It's each channel in order (Red, Green, Blue, Alpha, ...) - // where each channel consists of an 8-bit (or 16-bit) value for each pixel in the image. - - // Read the data by channel. - for (channel = 0; channel < 4; channel++) { - if (channel >= channelCount) { - // Fill this channel with default data. - if (bitdepth == 16 && bpc == 16) { - stbi__uint16 *q = ((stbi__uint16 *) out) + channel; - stbi__uint16 val = channel == 3 ? 65535 : 0; - for (i = 0; i < pixelCount; i++, q += 4) - *q = val; - } else { - stbi_uc *p = out+channel; - stbi_uc val = channel == 3 ? 255 : 0; - for (i = 0; i < pixelCount; i++, p += 4) - *p = val; - } - } else { - if (ri->bits_per_channel == 16) { // output bpc - stbi__uint16 *q = ((stbi__uint16 *) out) + channel; - for (i = 0; i < pixelCount; i++, q += 4) - *q = (stbi__uint16) stbi__get16be(s); - } else { - stbi_uc *p = out+channel; - if (bitdepth == 16) { // input bpc - for (i = 0; i < pixelCount; i++, p += 4) - *p = (stbi_uc) (stbi__get16be(s) >> 8); - } else { - for (i = 0; i < pixelCount; i++, p += 4) - *p = stbi__get8(s); - } - } - } - } - } - - // remove weird white matte from PSD - if (channelCount >= 4) { - if (ri->bits_per_channel == 16) { - for (i=0; i < w*h; ++i) { - stbi__uint16 *pixel = (stbi__uint16 *) out + 4*i; - if (pixel[3] != 0 && pixel[3] != 65535) { - float a = pixel[3] / 65535.0f; - float ra = 1.0f / a; - float inv_a = 65535.0f * (1 - ra); - pixel[0] = (stbi__uint16) (pixel[0]*ra + inv_a); - pixel[1] = (stbi__uint16) (pixel[1]*ra + inv_a); - pixel[2] = (stbi__uint16) (pixel[2]*ra + inv_a); - } - } - } else { - for (i=0; i < w*h; ++i) { - unsigned char *pixel = out + 4*i; - if (pixel[3] != 0 && pixel[3] != 255) { - float a = pixel[3] / 255.0f; - float ra = 1.0f / a; - float inv_a = 255.0f * (1 - ra); - pixel[0] = (unsigned char) (pixel[0]*ra + inv_a); - pixel[1] = (unsigned char) (pixel[1]*ra + inv_a); - pixel[2] = (unsigned char) (pixel[2]*ra + inv_a); - } - } - } - } - - // convert to desired output format - if (req_comp && req_comp != 4) { - if (ri->bits_per_channel == 16) - out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, 4, req_comp, w, h); - else - out = stbi__convert_format(out, 4, req_comp, w, h); - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - - if (comp) *comp = 4; - *y = h; - *x = w; - - return out; -} -#endif - -// ************************************************************************************************* -// Softimage PIC loader -// by Tom Seddon -// -// See http://softimage.wiki.softimage.com/index.php/INFO:_PIC_file_format -// See http://ozviz.wasp.uwa.edu.au/~pbourke/dataformats/softimagepic/ - -#ifndef STBI_NO_PIC -static int stbi__pic_is4(stbi__context *s,const char *str) -{ - int i; - for (i=0; i<4; ++i) - if (stbi__get8(s) != (stbi_uc)str[i]) - return 0; - - return 1; -} - -static int stbi__pic_test_core(stbi__context *s) -{ - int i; - - if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) - return 0; - - for(i=0;i<84;++i) - stbi__get8(s); - - if (!stbi__pic_is4(s,"PICT")) - return 0; - - return 1; -} - -typedef struct -{ - stbi_uc size,type,channel; -} stbi__pic_packet; - -static stbi_uc *stbi__readval(stbi__context *s, int channel, stbi_uc *dest) -{ - int mask=0x80, i; - - for (i=0; i<4; ++i, mask>>=1) { - if (channel & mask) { - if (stbi__at_eof(s)) return stbi__errpuc("bad file","PIC file too short"); - dest[i]=stbi__get8(s); - } - } - - return dest; -} - -static void stbi__copyval(int channel,stbi_uc *dest,const stbi_uc *src) -{ - int mask=0x80,i; - - for (i=0;i<4; ++i, mask>>=1) - if (channel&mask) - dest[i]=src[i]; -} - -static stbi_uc *stbi__pic_load_core(stbi__context *s,int width,int height,int *comp, stbi_uc *result) -{ - int act_comp=0,num_packets=0,y,chained; - stbi__pic_packet packets[10]; - - // this will (should...) cater for even some bizarre stuff like having data - // for the same channel in multiple packets. - do { - stbi__pic_packet *packet; - - if (num_packets==sizeof(packets)/sizeof(packets[0])) - return stbi__errpuc("bad format","too many packets"); - - packet = &packets[num_packets++]; - - chained = stbi__get8(s); - packet->size = stbi__get8(s); - packet->type = stbi__get8(s); - packet->channel = stbi__get8(s); - - act_comp |= packet->channel; - - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (reading packets)"); - if (packet->size != 8) return stbi__errpuc("bad format","packet isn't 8bpp"); - } while (chained); - - *comp = (act_comp & 0x10 ? 4 : 3); // has alpha channel? - - for(y=0; ytype) { - default: - return stbi__errpuc("bad format","packet has bad compression type"); - - case 0: {//uncompressed - int x; - - for(x=0;xchannel,dest)) - return 0; - break; - } - - case 1://Pure RLE - { - int left=width, i; - - while (left>0) { - stbi_uc count,value[4]; - - count=stbi__get8(s); - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pure read count)"); - - if (count > left) - count = (stbi_uc) left; - - if (!stbi__readval(s,packet->channel,value)) return 0; - - for(i=0; ichannel,dest,value); - left -= count; - } - } - break; - - case 2: {//Mixed RLE - int left=width; - while (left>0) { - int count = stbi__get8(s), i; - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (mixed read count)"); - - if (count >= 128) { // Repeated - stbi_uc value[4]; - - if (count==128) - count = stbi__get16be(s); - else - count -= 127; - if (count > left) - return stbi__errpuc("bad file","scanline overrun"); - - if (!stbi__readval(s,packet->channel,value)) - return 0; - - for(i=0;ichannel,dest,value); - } else { // Raw - ++count; - if (count>left) return stbi__errpuc("bad file","scanline overrun"); - - for(i=0;ichannel,dest)) - return 0; - } - left-=count; - } - break; - } - } - } - } - - return result; -} - -static void *stbi__pic_load(stbi__context *s,int *px,int *py,int *comp,int req_comp, stbi__result_info *ri) -{ - stbi_uc *result; - int i, x,y, internal_comp; - STBI_NOTUSED(ri); - - if (!comp) comp = &internal_comp; - - for (i=0; i<92; ++i) - stbi__get8(s); - - x = stbi__get16be(s); - y = stbi__get16be(s); - - if (y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pic header)"); - if (!stbi__mad3sizes_valid(x, y, 4, 0)) return stbi__errpuc("too large", "PIC image too large to decode"); - - stbi__get32be(s); //skip `ratio' - stbi__get16be(s); //skip `fields' - stbi__get16be(s); //skip `pad' - - // intermediate buffer is RGBA - result = (stbi_uc *) stbi__malloc_mad3(x, y, 4, 0); - if (!result) return stbi__errpuc("outofmem", "Out of memory"); - memset(result, 0xff, x*y*4); - - if (!stbi__pic_load_core(s,x,y,comp, result)) { - STBI_FREE(result); - result=0; - } - *px = x; - *py = y; - if (req_comp == 0) req_comp = *comp; - result=stbi__convert_format(result,4,req_comp,x,y); - - return result; -} - -static int stbi__pic_test(stbi__context *s) -{ - int r = stbi__pic_test_core(s); - stbi__rewind(s); - return r; -} -#endif - -// ************************************************************************************************* -// GIF loader -- public domain by Jean-Marc Lienher -- simplified/shrunk by stb - -#ifndef STBI_NO_GIF -typedef struct -{ - stbi__int16 prefix; - stbi_uc first; - stbi_uc suffix; -} stbi__gif_lzw; - -typedef struct -{ - int w,h; - stbi_uc *out; // output buffer (always 4 components) - stbi_uc *background; // The current "background" as far as a gif is concerned - stbi_uc *history; - int flags, bgindex, ratio, transparent, eflags; - stbi_uc pal[256][4]; - stbi_uc lpal[256][4]; - stbi__gif_lzw codes[8192]; - stbi_uc *color_table; - int parse, step; - int lflags; - int start_x, start_y; - int max_x, max_y; - int cur_x, cur_y; - int line_size; - int delay; -} stbi__gif; - -static int stbi__gif_test_raw(stbi__context *s) -{ - int sz; - if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') return 0; - sz = stbi__get8(s); - if (sz != '9' && sz != '7') return 0; - if (stbi__get8(s) != 'a') return 0; - return 1; -} - -static int stbi__gif_test(stbi__context *s) -{ - int r = stbi__gif_test_raw(s); - stbi__rewind(s); - return r; -} - -static void stbi__gif_parse_colortable(stbi__context *s, stbi_uc pal[256][4], int num_entries, int transp) -{ - int i; - for (i=0; i < num_entries; ++i) { - pal[i][2] = stbi__get8(s); - pal[i][1] = stbi__get8(s); - pal[i][0] = stbi__get8(s); - pal[i][3] = transp == i ? 0 : 255; - } -} - -static int stbi__gif_header(stbi__context *s, stbi__gif *g, int *comp, int is_info) -{ - stbi_uc version; - if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') - return stbi__err("not GIF", "Corrupt GIF"); - - version = stbi__get8(s); - if (version != '7' && version != '9') return stbi__err("not GIF", "Corrupt GIF"); - if (stbi__get8(s) != 'a') return stbi__err("not GIF", "Corrupt GIF"); - - stbi__g_failure_reason = ""; - g->w = stbi__get16le(s); - g->h = stbi__get16le(s); - g->flags = stbi__get8(s); - g->bgindex = stbi__get8(s); - g->ratio = stbi__get8(s); - g->transparent = -1; - - if (g->w > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (g->h > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - - if (comp != 0) *comp = 4; // can't actually tell whether it's 3 or 4 until we parse the comments - - if (is_info) return 1; - - if (g->flags & 0x80) - stbi__gif_parse_colortable(s,g->pal, 2 << (g->flags & 7), -1); - - return 1; -} - -static int stbi__gif_info_raw(stbi__context *s, int *x, int *y, int *comp) -{ - stbi__gif* g = (stbi__gif*) stbi__malloc(sizeof(stbi__gif)); - if (!g) return stbi__err("outofmem", "Out of memory"); - if (!stbi__gif_header(s, g, comp, 1)) { - STBI_FREE(g); - stbi__rewind( s ); - return 0; - } - if (x) *x = g->w; - if (y) *y = g->h; - STBI_FREE(g); - return 1; -} - -static void stbi__out_gif_code(stbi__gif *g, stbi__uint16 code) -{ - stbi_uc *p, *c; - int idx; - - // recurse to decode the prefixes, since the linked-list is backwards, - // and working backwards through an interleaved image would be nasty - if (g->codes[code].prefix >= 0) - stbi__out_gif_code(g, g->codes[code].prefix); - - if (g->cur_y >= g->max_y) return; - - idx = g->cur_x + g->cur_y; - p = &g->out[idx]; - g->history[idx / 4] = 1; - - c = &g->color_table[g->codes[code].suffix * 4]; - if (c[3] > 128) { // don't render transparent pixels; - p[0] = c[2]; - p[1] = c[1]; - p[2] = c[0]; - p[3] = c[3]; - } - g->cur_x += 4; - - if (g->cur_x >= g->max_x) { - g->cur_x = g->start_x; - g->cur_y += g->step; - - while (g->cur_y >= g->max_y && g->parse > 0) { - g->step = (1 << g->parse) * g->line_size; - g->cur_y = g->start_y + (g->step >> 1); - --g->parse; - } - } -} - -static stbi_uc *stbi__process_gif_raster(stbi__context *s, stbi__gif *g) -{ - stbi_uc lzw_cs; - stbi__int32 len, init_code; - stbi__uint32 first; - stbi__int32 codesize, codemask, avail, oldcode, bits, valid_bits, clear; - stbi__gif_lzw *p; - - lzw_cs = stbi__get8(s); - if (lzw_cs > 12) return NULL; - clear = 1 << lzw_cs; - first = 1; - codesize = lzw_cs + 1; - codemask = (1 << codesize) - 1; - bits = 0; - valid_bits = 0; - for (init_code = 0; init_code < clear; init_code++) { - g->codes[init_code].prefix = -1; - g->codes[init_code].first = (stbi_uc) init_code; - g->codes[init_code].suffix = (stbi_uc) init_code; - } - - // support no starting clear code - avail = clear+2; - oldcode = -1; - - len = 0; - for(;;) { - if (valid_bits < codesize) { - if (len == 0) { - len = stbi__get8(s); // start new block - if (len == 0) - return g->out; - } - --len; - bits |= (stbi__int32) stbi__get8(s) << valid_bits; - valid_bits += 8; - } else { - stbi__int32 code = bits & codemask; - bits >>= codesize; - valid_bits -= codesize; - // @OPTIMIZE: is there some way we can accelerate the non-clear path? - if (code == clear) { // clear code - codesize = lzw_cs + 1; - codemask = (1 << codesize) - 1; - avail = clear + 2; - oldcode = -1; - first = 0; - } else if (code == clear + 1) { // end of stream code - stbi__skip(s, len); - while ((len = stbi__get8(s)) > 0) - stbi__skip(s,len); - return g->out; - } else if (code <= avail) { - if (first) { - return stbi__errpuc("no clear code", "Corrupt GIF"); - } - - if (oldcode >= 0) { - p = &g->codes[avail++]; - if (avail > 8192) { - return stbi__errpuc("too many codes", "Corrupt GIF"); - } - - p->prefix = (stbi__int16) oldcode; - p->first = g->codes[oldcode].first; - p->suffix = (code == avail) ? p->first : g->codes[code].first; - } else if (code == avail) - return stbi__errpuc("illegal code in raster", "Corrupt GIF"); - - stbi__out_gif_code(g, (stbi__uint16) code); - - if ((avail & codemask) == 0 && avail <= 0x0FFF) { - codesize++; - codemask = (1 << codesize) - 1; - } - - oldcode = code; - } else { - return stbi__errpuc("illegal code in raster", "Corrupt GIF"); - } - } - } -} - -// this function is designed to support animated gifs, although stb_image doesn't support it -// two back is the image from two frames ago, used for a very specific disposal format -static stbi_uc *stbi__gif_load_next(stbi__context *s, stbi__gif *g, int *comp, int req_comp, stbi_uc *two_back) -{ - int dispose; - int first_frame; - int pi; - int pcount; - STBI_NOTUSED(req_comp); - - // on first frame, any non-written pixels get the background colour (non-transparent) - first_frame = 0; - if (g->out == 0) { - if (!stbi__gif_header(s, g, comp,0)) return 0; // stbi__g_failure_reason set by stbi__gif_header - if (!stbi__mad3sizes_valid(4, g->w, g->h, 0)) - return stbi__errpuc("too large", "GIF image is too large"); - pcount = g->w * g->h; - g->out = (stbi_uc *) stbi__malloc(4 * pcount); - g->background = (stbi_uc *) stbi__malloc(4 * pcount); - g->history = (stbi_uc *) stbi__malloc(pcount); - if (!g->out || !g->background || !g->history) - return stbi__errpuc("outofmem", "Out of memory"); - - // image is treated as "transparent" at the start - ie, nothing overwrites the current background; - // background colour is only used for pixels that are not rendered first frame, after that "background" - // color refers to the color that was there the previous frame. - memset(g->out, 0x00, 4 * pcount); - memset(g->background, 0x00, 4 * pcount); // state of the background (starts transparent) - memset(g->history, 0x00, pcount); // pixels that were affected previous frame - first_frame = 1; - } else { - // second frame - how do we dispose of the previous one? - dispose = (g->eflags & 0x1C) >> 2; - pcount = g->w * g->h; - - if ((dispose == 3) && (two_back == 0)) { - dispose = 2; // if I don't have an image to revert back to, default to the old background - } - - if (dispose == 3) { // use previous graphic - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi]) { - memcpy( &g->out[pi * 4], &two_back[pi * 4], 4 ); - } - } - } else if (dispose == 2) { - // restore what was changed last frame to background before that frame; - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi]) { - memcpy( &g->out[pi * 4], &g->background[pi * 4], 4 ); - } - } - } else { - // This is a non-disposal case eithe way, so just - // leave the pixels as is, and they will become the new background - // 1: do not dispose - // 0: not specified. - } - - // background is what out is after the undoing of the previou frame; - memcpy( g->background, g->out, 4 * g->w * g->h ); - } - - // clear my history; - memset( g->history, 0x00, g->w * g->h ); // pixels that were affected previous frame - - for (;;) { - int tag = stbi__get8(s); - switch (tag) { - case 0x2C: /* Image Descriptor */ - { - stbi__int32 x, y, w, h; - stbi_uc *o; - - x = stbi__get16le(s); - y = stbi__get16le(s); - w = stbi__get16le(s); - h = stbi__get16le(s); - if (((x + w) > (g->w)) || ((y + h) > (g->h))) - return stbi__errpuc("bad Image Descriptor", "Corrupt GIF"); - - g->line_size = g->w * 4; - g->start_x = x * 4; - g->start_y = y * g->line_size; - g->max_x = g->start_x + w * 4; - g->max_y = g->start_y + h * g->line_size; - g->cur_x = g->start_x; - g->cur_y = g->start_y; - - // if the width of the specified rectangle is 0, that means - // we may not see *any* pixels or the image is malformed; - // to make sure this is caught, move the current y down to - // max_y (which is what out_gif_code checks). - if (w == 0) - g->cur_y = g->max_y; - - g->lflags = stbi__get8(s); - - if (g->lflags & 0x40) { - g->step = 8 * g->line_size; // first interlaced spacing - g->parse = 3; - } else { - g->step = g->line_size; - g->parse = 0; - } - - if (g->lflags & 0x80) { - stbi__gif_parse_colortable(s,g->lpal, 2 << (g->lflags & 7), g->eflags & 0x01 ? g->transparent : -1); - g->color_table = (stbi_uc *) g->lpal; - } else if (g->flags & 0x80) { - g->color_table = (stbi_uc *) g->pal; - } else - return stbi__errpuc("missing color table", "Corrupt GIF"); - - o = stbi__process_gif_raster(s, g); - if (!o) return NULL; - - // if this was the first frame, - pcount = g->w * g->h; - if (first_frame && (g->bgindex > 0)) { - // if first frame, any pixel not drawn to gets the background color - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi] == 0) { - g->pal[g->bgindex][3] = 255; // just in case it was made transparent, undo that; It will be reset next frame if need be; - memcpy( &g->out[pi * 4], &g->pal[g->bgindex], 4 ); - } - } - } - - return o; - } - - case 0x21: // Comment Extension. - { - int len; - int ext = stbi__get8(s); - if (ext == 0xF9) { // Graphic Control Extension. - len = stbi__get8(s); - if (len == 4) { - g->eflags = stbi__get8(s); - g->delay = 10 * stbi__get16le(s); // delay - 1/100th of a second, saving as 1/1000ths. - - // unset old transparent - if (g->transparent >= 0) { - g->pal[g->transparent][3] = 255; - } - if (g->eflags & 0x01) { - g->transparent = stbi__get8(s); - if (g->transparent >= 0) { - g->pal[g->transparent][3] = 0; - } - } else { - // don't need transparent - stbi__skip(s, 1); - g->transparent = -1; - } - } else { - stbi__skip(s, len); - break; - } - } - while ((len = stbi__get8(s)) != 0) { - stbi__skip(s, len); - } - break; - } - - case 0x3B: // gif stream termination code - return (stbi_uc *) s; // using '1' causes warning on some compilers - - default: - return stbi__errpuc("unknown code", "Corrupt GIF"); - } - } -} - -static void *stbi__load_gif_main_outofmem(stbi__gif *g, stbi_uc *out, int **delays) -{ - STBI_FREE(g->out); - STBI_FREE(g->history); - STBI_FREE(g->background); - - if (out) STBI_FREE(out); - if (delays && *delays) STBI_FREE(*delays); - return stbi__errpuc("outofmem", "Out of memory"); -} - -static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp) -{ - if (stbi__gif_test(s)) { - int layers = 0; - stbi_uc *u = 0; - stbi_uc *out = 0; - stbi_uc *two_back = 0; - stbi__gif g; - int stride; - int out_size = 0; - int delays_size = 0; - - STBI_NOTUSED(out_size); - STBI_NOTUSED(delays_size); - - memset(&g, 0, sizeof(g)); - if (delays) { - *delays = 0; - } - - do { - u = stbi__gif_load_next(s, &g, comp, req_comp, two_back); - if (u == (stbi_uc *) s) u = 0; // end of animated gif marker - - if (u) { - *x = g.w; - *y = g.h; - ++layers; - stride = g.w * g.h * 4; - - if (out) { - void *tmp = (stbi_uc*) STBI_REALLOC_SIZED( out, out_size, layers * stride ); - if (!tmp) - return stbi__load_gif_main_outofmem(&g, out, delays); - else { - out = (stbi_uc*) tmp; - out_size = layers * stride; - } - - if (delays) { - int *new_delays = (int*) STBI_REALLOC_SIZED( *delays, delays_size, sizeof(int) * layers ); - if (!new_delays) - return stbi__load_gif_main_outofmem(&g, out, delays); - *delays = new_delays; - delays_size = layers * sizeof(int); - } - } else { - out = (stbi_uc*)stbi__malloc( layers * stride ); - if (!out) - return stbi__load_gif_main_outofmem(&g, out, delays); - out_size = layers * stride; - if (delays) { - *delays = (int*) stbi__malloc( layers * sizeof(int) ); - if (!*delays) - return stbi__load_gif_main_outofmem(&g, out, delays); - delays_size = layers * sizeof(int); - } - } - memcpy( out + ((layers - 1) * stride), u, stride ); - if (layers >= 2) { - two_back = out - 2 * stride; - } - - if (delays) { - (*delays)[layers - 1U] = g.delay; - } - } - } while (u != 0); - - // free temp buffer; - STBI_FREE(g.out); - STBI_FREE(g.history); - STBI_FREE(g.background); - - // do the final conversion after loading everything; - if (req_comp && req_comp != 4) - out = stbi__convert_format(out, 4, req_comp, layers * g.w, g.h); - - *z = layers; - return out; - } else { - return stbi__errpuc("not GIF", "Image was not as a gif type."); - } -} - -static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *u = 0; - stbi__gif g; - memset(&g, 0, sizeof(g)); - STBI_NOTUSED(ri); - - u = stbi__gif_load_next(s, &g, comp, req_comp, 0); - if (u == (stbi_uc *) s) u = 0; // end of animated gif marker - if (u) { - *x = g.w; - *y = g.h; - - // moved conversion to after successful load so that the same - // can be done for multiple frames. - if (req_comp && req_comp != 4) - u = stbi__convert_format(u, 4, req_comp, g.w, g.h); - } else if (g.out) { - // if there was an error and we allocated an image buffer, free it! - STBI_FREE(g.out); - } - - // free buffers needed for multiple frame loading; - STBI_FREE(g.history); - STBI_FREE(g.background); - - return u; -} - -static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp) -{ - return stbi__gif_info_raw(s,x,y,comp); -} -#endif - -// ************************************************************************************************* -// Radiance RGBE HDR loader -// originally by Nicolas Schulz -#ifndef STBI_NO_HDR -static int stbi__hdr_test_core(stbi__context *s, const char *signature) -{ - int i; - for (i=0; signature[i]; ++i) - if (stbi__get8(s) != signature[i]) - return 0; - stbi__rewind(s); - return 1; -} - -static int stbi__hdr_test(stbi__context* s) -{ - int r = stbi__hdr_test_core(s, "#?RADIANCE\n"); - stbi__rewind(s); - if(!r) { - r = stbi__hdr_test_core(s, "#?RGBE\n"); - stbi__rewind(s); - } - return r; -} - -#define STBI__HDR_BUFLEN 1024 -static char *stbi__hdr_gettoken(stbi__context *z, char *buffer) -{ - int len=0; - char c = '\0'; - - c = (char) stbi__get8(z); - - while (!stbi__at_eof(z) && c != '\n') { - buffer[len++] = c; - if (len == STBI__HDR_BUFLEN-1) { - // flush to end of line - while (!stbi__at_eof(z) && stbi__get8(z) != '\n') - ; - break; - } - c = (char) stbi__get8(z); - } - - buffer[len] = 0; - return buffer; -} - -static void stbi__hdr_convert(float *output, stbi_uc *input, int req_comp) -{ - if ( input[3] != 0 ) { - float f1; - // Exponent - f1 = (float) ldexp(1.0f, input[3] - (int)(128 + 8)); - if (req_comp <= 2) - output[0] = (input[0] + input[1] + input[2]) * f1 / 3; - else { - output[0] = input[0] * f1; - output[1] = input[1] * f1; - output[2] = input[2] * f1; - } - if (req_comp == 2) output[1] = 1; - if (req_comp == 4) output[3] = 1; - } else { - switch (req_comp) { - case 4: output[3] = 1; /* fallthrough */ - case 3: output[0] = output[1] = output[2] = 0; - break; - case 2: output[1] = 1; /* fallthrough */ - case 1: output[0] = 0; - break; - } - } -} - -static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - char buffer[STBI__HDR_BUFLEN]; - char *token; - int valid = 0; - int width, height; - stbi_uc *scanline; - float *hdr_data; - int len; - unsigned char count, value; - int i, j, k, c1,c2, z; - const char *headerToken; - STBI_NOTUSED(ri); - - // Check identifier - headerToken = stbi__hdr_gettoken(s,buffer); - if (strcmp(headerToken, "#?RADIANCE") != 0 && strcmp(headerToken, "#?RGBE") != 0) - return stbi__errpf("not HDR", "Corrupt HDR image"); - - // Parse header - for(;;) { - token = stbi__hdr_gettoken(s,buffer); - if (token[0] == 0) break; - if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; - } - - if (!valid) return stbi__errpf("unsupported format", "Unsupported HDR format"); - - // Parse width and height - // can't use sscanf() if we're not using stdio! - token = stbi__hdr_gettoken(s,buffer); - if (strncmp(token, "-Y ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); - token += 3; - height = (int) strtol(token, &token, 10); - while (*token == ' ') ++token; - if (strncmp(token, "+X ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); - token += 3; - width = (int) strtol(token, NULL, 10); - - if (height > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); - if (width > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); - - *x = width; - *y = height; - - if (comp) *comp = 3; - if (req_comp == 0) req_comp = 3; - - if (!stbi__mad4sizes_valid(width, height, req_comp, sizeof(float), 0)) - return stbi__errpf("too large", "HDR image is too large"); - - // Read data - hdr_data = (float *) stbi__malloc_mad4(width, height, req_comp, sizeof(float), 0); - if (!hdr_data) - return stbi__errpf("outofmem", "Out of memory"); - - // Load image data - // image data is stored as some number of sca - if ( width < 8 || width >= 32768) { - // Read flat data - for (j=0; j < height; ++j) { - for (i=0; i < width; ++i) { - stbi_uc rgbe[4]; - main_decode_loop: - stbi__getn(s, rgbe, 4); - stbi__hdr_convert(hdr_data + j * width * req_comp + i * req_comp, rgbe, req_comp); - } - } - } else { - // Read RLE-encoded data - scanline = NULL; - - for (j = 0; j < height; ++j) { - c1 = stbi__get8(s); - c2 = stbi__get8(s); - len = stbi__get8(s); - if (c1 != 2 || c2 != 2 || (len & 0x80)) { - // not run-length encoded, so we have to actually use THIS data as a decoded - // pixel (note this can't be a valid pixel--one of RGB must be >= 128) - stbi_uc rgbe[4]; - rgbe[0] = (stbi_uc) c1; - rgbe[1] = (stbi_uc) c2; - rgbe[2] = (stbi_uc) len; - rgbe[3] = (stbi_uc) stbi__get8(s); - stbi__hdr_convert(hdr_data, rgbe, req_comp); - i = 1; - j = 0; - STBI_FREE(scanline); - goto main_decode_loop; // yes, this makes no sense - } - len <<= 8; - len |= stbi__get8(s); - if (len != width) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("invalid decoded scanline length", "corrupt HDR"); } - if (scanline == NULL) { - scanline = (stbi_uc *) stbi__malloc_mad2(width, 4, 0); - if (!scanline) { - STBI_FREE(hdr_data); - return stbi__errpf("outofmem", "Out of memory"); - } - } - - for (k = 0; k < 4; ++k) { - int nleft; - i = 0; - while ((nleft = width - i) > 0) { - count = stbi__get8(s); - if (count > 128) { - // Run - value = stbi__get8(s); - count -= 128; - if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } - for (z = 0; z < count; ++z) - scanline[i++ * 4 + k] = value; - } else { - // Dump - if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } - for (z = 0; z < count; ++z) - scanline[i++ * 4 + k] = stbi__get8(s); - } - } - } - for (i=0; i < width; ++i) - stbi__hdr_convert(hdr_data+(j*width + i)*req_comp, scanline + i*4, req_comp); - } - if (scanline) - STBI_FREE(scanline); - } - - return hdr_data; -} - -static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp) -{ - char buffer[STBI__HDR_BUFLEN]; - char *token; - int valid = 0; - int dummy; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - if (stbi__hdr_test(s) == 0) { - stbi__rewind( s ); - return 0; - } - - for(;;) { - token = stbi__hdr_gettoken(s,buffer); - if (token[0] == 0) break; - if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; - } - - if (!valid) { - stbi__rewind( s ); - return 0; - } - token = stbi__hdr_gettoken(s,buffer); - if (strncmp(token, "-Y ", 3)) { - stbi__rewind( s ); - return 0; - } - token += 3; - *y = (int) strtol(token, &token, 10); - while (*token == ' ') ++token; - if (strncmp(token, "+X ", 3)) { - stbi__rewind( s ); - return 0; - } - token += 3; - *x = (int) strtol(token, NULL, 10); - *comp = 3; - return 1; -} -#endif // STBI_NO_HDR - -#ifndef STBI_NO_BMP -static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp) -{ - void *p; - stbi__bmp_data info; - - info.all_a = 255; - p = stbi__bmp_parse_header(s, &info); - if (p == NULL) { - stbi__rewind( s ); - return 0; - } - if (x) *x = s->img_x; - if (y) *y = s->img_y; - if (comp) { - if (info.bpp == 24 && info.ma == 0xff000000) - *comp = 3; - else - *comp = info.ma ? 4 : 3; - } - return 1; -} -#endif - -#ifndef STBI_NO_PSD -static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp) -{ - int channelCount, dummy, depth; - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - if (stbi__get32be(s) != 0x38425053) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 1) { - stbi__rewind( s ); - return 0; - } - stbi__skip(s, 6); - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) { - stbi__rewind( s ); - return 0; - } - *y = stbi__get32be(s); - *x = stbi__get32be(s); - depth = stbi__get16be(s); - if (depth != 8 && depth != 16) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 3) { - stbi__rewind( s ); - return 0; - } - *comp = 4; - return 1; -} - -static int stbi__psd_is16(stbi__context *s) -{ - int channelCount, depth; - if (stbi__get32be(s) != 0x38425053) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 1) { - stbi__rewind( s ); - return 0; - } - stbi__skip(s, 6); - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) { - stbi__rewind( s ); - return 0; - } - STBI_NOTUSED(stbi__get32be(s)); - STBI_NOTUSED(stbi__get32be(s)); - depth = stbi__get16be(s); - if (depth != 16) { - stbi__rewind( s ); - return 0; - } - return 1; -} -#endif - -#ifndef STBI_NO_PIC -static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp) -{ - int act_comp=0,num_packets=0,chained,dummy; - stbi__pic_packet packets[10]; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) { - stbi__rewind(s); - return 0; - } - - stbi__skip(s, 88); - - *x = stbi__get16be(s); - *y = stbi__get16be(s); - if (stbi__at_eof(s)) { - stbi__rewind( s); - return 0; - } - if ( (*x) != 0 && (1 << 28) / (*x) < (*y)) { - stbi__rewind( s ); - return 0; - } - - stbi__skip(s, 8); - - do { - stbi__pic_packet *packet; - - if (num_packets==sizeof(packets)/sizeof(packets[0])) - return 0; - - packet = &packets[num_packets++]; - chained = stbi__get8(s); - packet->size = stbi__get8(s); - packet->type = stbi__get8(s); - packet->channel = stbi__get8(s); - act_comp |= packet->channel; - - if (stbi__at_eof(s)) { - stbi__rewind( s ); - return 0; - } - if (packet->size != 8) { - stbi__rewind( s ); - return 0; - } - } while (chained); - - *comp = (act_comp & 0x10 ? 4 : 3); - - return 1; -} -#endif - -// ************************************************************************************************* -// Portable Gray Map and Portable Pixel Map loader -// by Ken Miller -// -// PGM: http://netpbm.sourceforge.net/doc/pgm.html -// PPM: http://netpbm.sourceforge.net/doc/ppm.html -// -// Known limitations: -// Does not support comments in the header section -// Does not support ASCII image data (formats P2 and P3) - -#ifndef STBI_NO_PNM - -static int stbi__pnm_test(stbi__context *s) -{ - char p, t; - p = (char) stbi__get8(s); - t = (char) stbi__get8(s); - if (p != 'P' || (t != '5' && t != '6')) { - stbi__rewind( s ); - return 0; - } - return 1; -} - -static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *out; - STBI_NOTUSED(ri); - - ri->bits_per_channel = stbi__pnm_info(s, (int *)&s->img_x, (int *)&s->img_y, (int *)&s->img_n); - if (ri->bits_per_channel == 0) - return 0; - - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - *x = s->img_x; - *y = s->img_y; - if (comp) *comp = s->img_n; - - if (!stbi__mad4sizes_valid(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0)) - return stbi__errpuc("too large", "PNM too large"); - - out = (stbi_uc *) stbi__malloc_mad4(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0); - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - if (!stbi__getn(s, out, s->img_n * s->img_x * s->img_y * (ri->bits_per_channel / 8))) { - STBI_FREE(out); - return stbi__errpuc("bad PNM", "PNM file truncated"); - } - - if (req_comp && req_comp != s->img_n) { - if (ri->bits_per_channel == 16) { - out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, s->img_n, req_comp, s->img_x, s->img_y); - } else { - out = stbi__convert_format(out, s->img_n, req_comp, s->img_x, s->img_y); - } - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - return out; -} - -static int stbi__pnm_isspace(char c) -{ - return c == ' ' || c == '\t' || c == '\n' || c == '\v' || c == '\f' || c == '\r'; -} - -static void stbi__pnm_skip_whitespace(stbi__context *s, char *c) -{ - for (;;) { - while (!stbi__at_eof(s) && stbi__pnm_isspace(*c)) - *c = (char) stbi__get8(s); - - if (stbi__at_eof(s) || *c != '#') - break; - - while (!stbi__at_eof(s) && *c != '\n' && *c != '\r' ) - *c = (char) stbi__get8(s); - } -} - -static int stbi__pnm_isdigit(char c) -{ - return c >= '0' && c <= '9'; -} - -static int stbi__pnm_getinteger(stbi__context *s, char *c) -{ - int value = 0; - - while (!stbi__at_eof(s) && stbi__pnm_isdigit(*c)) { - value = value*10 + (*c - '0'); - *c = (char) stbi__get8(s); - if((value > 214748364) || (value == 214748364 && *c > '7')) - return stbi__err("integer parse overflow", "Parsing an integer in the PPM header overflowed a 32-bit int"); - } - - return value; -} - -static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp) -{ - int maxv, dummy; - char c, p, t; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - stbi__rewind(s); - - // Get identifier - p = (char) stbi__get8(s); - t = (char) stbi__get8(s); - if (p != 'P' || (t != '5' && t != '6')) { - stbi__rewind(s); - return 0; - } - - *comp = (t == '6') ? 3 : 1; // '5' is 1-component .pgm; '6' is 3-component .ppm - - c = (char) stbi__get8(s); - stbi__pnm_skip_whitespace(s, &c); - - *x = stbi__pnm_getinteger(s, &c); // read width - if(*x == 0) - return stbi__err("invalid width", "PPM image header had zero or overflowing width"); - stbi__pnm_skip_whitespace(s, &c); - - *y = stbi__pnm_getinteger(s, &c); // read height - if (*y == 0) - return stbi__err("invalid width", "PPM image header had zero or overflowing width"); - stbi__pnm_skip_whitespace(s, &c); - - maxv = stbi__pnm_getinteger(s, &c); // read max value - if (maxv > 65535) - return stbi__err("max value > 65535", "PPM image supports only 8-bit and 16-bit images"); - else if (maxv > 255) - return 16; - else - return 8; -} - -static int stbi__pnm_is16(stbi__context *s) -{ - if (stbi__pnm_info(s, NULL, NULL, NULL) == 16) - return 1; - return 0; -} -#endif - -static int stbi__info_main(stbi__context *s, int *x, int *y, int *comp) -{ - #ifndef STBI_NO_JPEG - if (stbi__jpeg_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PNG - if (stbi__png_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_GIF - if (stbi__gif_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_BMP - if (stbi__bmp_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PSD - if (stbi__psd_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PIC - if (stbi__pic_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PNM - if (stbi__pnm_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_HDR - if (stbi__hdr_info(s, x, y, comp)) return 1; - #endif - - // test tga last because it's a crappy test! - #ifndef STBI_NO_TGA - if (stbi__tga_info(s, x, y, comp)) - return 1; - #endif - return stbi__err("unknown image type", "Image not of any known type, or corrupt"); -} - -static int stbi__is_16_main(stbi__context *s) -{ - #ifndef STBI_NO_PNG - if (stbi__png_is16(s)) return 1; - #endif - - #ifndef STBI_NO_PSD - if (stbi__psd_is16(s)) return 1; - #endif - - #ifndef STBI_NO_PNM - if (stbi__pnm_is16(s)) return 1; - #endif - return 0; -} - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_info(char const *filename, int *x, int *y, int *comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result; - if (!f) return stbi__err("can't fopen", "Unable to open file"); - result = stbi_info_from_file(f, x, y, comp); - fclose(f); - return result; -} - -STBIDEF int stbi_info_from_file(FILE *f, int *x, int *y, int *comp) -{ - int r; - stbi__context s; - long pos = ftell(f); - stbi__start_file(&s, f); - r = stbi__info_main(&s,x,y,comp); - fseek(f,pos,SEEK_SET); - return r; -} - -STBIDEF int stbi_is_16_bit(char const *filename) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result; - if (!f) return stbi__err("can't fopen", "Unable to open file"); - result = stbi_is_16_bit_from_file(f); - fclose(f); - return result; -} - -STBIDEF int stbi_is_16_bit_from_file(FILE *f) -{ - int r; - stbi__context s; - long pos = ftell(f); - stbi__start_file(&s, f); - r = stbi__is_16_main(&s); - fseek(f,pos,SEEK_SET); - return r; -} -#endif // !STBI_NO_STDIO - -STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__info_main(&s,x,y,comp); -} - -STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *c, void *user, int *x, int *y, int *comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); - return stbi__info_main(&s,x,y,comp); -} - -STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__is_16_main(&s); -} - -STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *c, void *user) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); - return stbi__is_16_main(&s); -} - -#endif // STB_IMAGE_IMPLEMENTATION - -/* - revision history: - 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs - 2.19 (2018-02-11) fix warning - 2.18 (2018-01-30) fix warnings - 2.17 (2018-01-29) change sbti__shiftsigned to avoid clang -O2 bug - 1-bit BMP - *_is_16_bit api - avoid warnings - 2.16 (2017-07-23) all functions have 16-bit variants; - STBI_NO_STDIO works again; - compilation fixes; - fix rounding in unpremultiply; - optimize vertical flip; - disable raw_len validation; - documentation fixes - 2.15 (2017-03-18) fix png-1,2,4 bug; now all Imagenet JPGs decode; - warning fixes; disable run-time SSE detection on gcc; - uniform handling of optional "return" values; - thread-safe initialization of zlib tables - 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs - 2.13 (2016-11-29) add 16-bit API, only supported for PNG right now - 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes - 2.11 (2016-04-02) allocate large structures on the stack - remove white matting for transparent PSD - fix reported channel count for PNG & BMP - re-enable SSE2 in non-gcc 64-bit - support RGB-formatted JPEG - read 16-bit PNGs (only as 8-bit) - 2.10 (2016-01-22) avoid warning introduced in 2.09 by STBI_REALLOC_SIZED - 2.09 (2016-01-16) allow comments in PNM files - 16-bit-per-pixel TGA (not bit-per-component) - info() for TGA could break due to .hdr handling - info() for BMP to shares code instead of sloppy parse - can use STBI_REALLOC_SIZED if allocator doesn't support realloc - code cleanup - 2.08 (2015-09-13) fix to 2.07 cleanup, reading RGB PSD as RGBA - 2.07 (2015-09-13) fix compiler warnings - partial animated GIF support - limited 16-bpc PSD support - #ifdef unused functions - bug with < 92 byte PIC,PNM,HDR,TGA - 2.06 (2015-04-19) fix bug where PSD returns wrong '*comp' value - 2.05 (2015-04-19) fix bug in progressive JPEG handling, fix warning - 2.04 (2015-04-15) try to re-enable SIMD on MinGW 64-bit - 2.03 (2015-04-12) extra corruption checking (mmozeiko) - stbi_set_flip_vertically_on_load (nguillemot) - fix NEON support; fix mingw support - 2.02 (2015-01-19) fix incorrect assert, fix warning - 2.01 (2015-01-17) fix various warnings; suppress SIMD on gcc 32-bit without -msse2 - 2.00b (2014-12-25) fix STBI_MALLOC in progressive JPEG - 2.00 (2014-12-25) optimize JPG, including x86 SSE2 & NEON SIMD (ryg) - progressive JPEG (stb) - PGM/PPM support (Ken Miller) - STBI_MALLOC,STBI_REALLOC,STBI_FREE - GIF bugfix -- seemingly never worked - STBI_NO_*, STBI_ONLY_* - 1.48 (2014-12-14) fix incorrectly-named assert() - 1.47 (2014-12-14) 1/2/4-bit PNG support, both direct and paletted (Omar Cornut & stb) - optimize PNG (ryg) - fix bug in interlaced PNG with user-specified channel count (stb) - 1.46 (2014-08-26) - fix broken tRNS chunk (colorkey-style transparency) in non-paletted PNG - 1.45 (2014-08-16) - fix MSVC-ARM internal compiler error by wrapping malloc - 1.44 (2014-08-07) - various warning fixes from Ronny Chevalier - 1.43 (2014-07-15) - fix MSVC-only compiler problem in code changed in 1.42 - 1.42 (2014-07-09) - don't define _CRT_SECURE_NO_WARNINGS (affects user code) - fixes to stbi__cleanup_jpeg path - added STBI_ASSERT to avoid requiring assert.h - 1.41 (2014-06-25) - fix search&replace from 1.36 that messed up comments/error messages - 1.40 (2014-06-22) - fix gcc struct-initialization warning - 1.39 (2014-06-15) - fix to TGA optimization when req_comp != number of components in TGA; - fix to GIF loading because BMP wasn't rewinding (whoops, no GIFs in my test suite) - add support for BMP version 5 (more ignored fields) - 1.38 (2014-06-06) - suppress MSVC warnings on integer casts truncating values - fix accidental rename of 'skip' field of I/O - 1.37 (2014-06-04) - remove duplicate typedef - 1.36 (2014-06-03) - convert to header file single-file library - if de-iphone isn't set, load iphone images color-swapped instead of returning NULL - 1.35 (2014-05-27) - various warnings - fix broken STBI_SIMD path - fix bug where stbi_load_from_file no longer left file pointer in correct place - fix broken non-easy path for 32-bit BMP (possibly never used) - TGA optimization by Arseny Kapoulkine - 1.34 (unknown) - use STBI_NOTUSED in stbi__resample_row_generic(), fix one more leak in tga failure case - 1.33 (2011-07-14) - make stbi_is_hdr work in STBI_NO_HDR (as specified), minor compiler-friendly improvements - 1.32 (2011-07-13) - support for "info" function for all supported filetypes (SpartanJ) - 1.31 (2011-06-20) - a few more leak fixes, bug in PNG handling (SpartanJ) - 1.30 (2011-06-11) - added ability to load files via callbacks to accomidate custom input streams (Ben Wenger) - removed deprecated format-specific test/load functions - removed support for installable file formats (stbi_loader) -- would have been broken for IO callbacks anyway - error cases in bmp and tga give messages and don't leak (Raymond Barbiero, grisha) - fix inefficiency in decoding 32-bit BMP (David Woo) - 1.29 (2010-08-16) - various warning fixes from Aurelien Pocheville - 1.28 (2010-08-01) - fix bug in GIF palette transparency (SpartanJ) - 1.27 (2010-08-01) - cast-to-stbi_uc to fix warnings - 1.26 (2010-07-24) - fix bug in file buffering for PNG reported by SpartanJ - 1.25 (2010-07-17) - refix trans_data warning (Won Chun) - 1.24 (2010-07-12) - perf improvements reading from files on platforms with lock-heavy fgetc() - minor perf improvements for jpeg - deprecated type-specific functions so we'll get feedback if they're needed - attempt to fix trans_data warning (Won Chun) - 1.23 fixed bug in iPhone support - 1.22 (2010-07-10) - removed image *writing* support - stbi_info support from Jetro Lauha - GIF support from Jean-Marc Lienher - iPhone PNG-extensions from James Brown - warning-fixes from Nicolas Schulz and Janez Zemva (i.stbi__err. Janez (U+017D)emva) - 1.21 fix use of 'stbi_uc' in header (reported by jon blow) - 1.20 added support for Softimage PIC, by Tom Seddon - 1.19 bug in interlaced PNG corruption check (found by ryg) - 1.18 (2008-08-02) - fix a threading bug (local mutable static) - 1.17 support interlaced PNG - 1.16 major bugfix - stbi__convert_format converted one too many pixels - 1.15 initialize some fields for thread safety - 1.14 fix threadsafe conversion bug - header-file-only version (#define STBI_HEADER_FILE_ONLY before including) - 1.13 threadsafe - 1.12 const qualifiers in the API - 1.11 Support installable IDCT, colorspace conversion routines - 1.10 Fixes for 64-bit (don't use "unsigned long") - optimized upsampling by Fabian "ryg" Giesen - 1.09 Fix format-conversion for PSD code (bad global variables!) - 1.08 Thatcher Ulrich's PSD code integrated by Nicolas Schulz - 1.07 attempt to fix C++ warning/errors again - 1.06 attempt to fix C++ warning/errors again - 1.05 fix TGA loading to return correct *comp and use good luminance calc - 1.04 default float alpha is 1, not 255; use 'void *' for stbi_image_free - 1.03 bugfixes to STBI_NO_STDIO, STBI_NO_HDR - 1.02 support for (subset of) HDR files, float interface for preferred access to them - 1.01 fix bug: possible bug in handling right-side up bmps... not sure - fix bug: the stbi__bmp_load() and stbi__tga_load() functions didn't work at all - 1.00 interface to zlib that skips zlib header - 0.99 correct handling of alpha in palette - 0.98 TGA loader by lonesock; dynamically add loaders (untested) - 0.97 jpeg errors on too large a file; also catch another malloc failure - 0.96 fix detection of invalid v value - particleman@mollyrocket forum - 0.95 during header scan, seek to markers in case of padding - 0.94 STBI_NO_STDIO to disable stdio usage; rename all #defines the same - 0.93 handle jpegtran output; verbose errors - 0.92 read 4,8,16,24,32-bit BMP files of several formats - 0.91 output 24-bit Windows 3.0 BMP files - 0.90 fix a few more warnings; bump version number to approach 1.0 - 0.61 bugfixes due to Marc LeBlanc, Christopher Lloyd - 0.60 fix compiling as c++ - 0.59 fix warnings: merge Dave Moore's -Wall fixes - 0.58 fix bug: zlib uncompressed mode len/nlen was wrong endian - 0.57 fix bug: jpg last huffman symbol before marker was >9 bits but less than 16 available - 0.56 fix bug: zlib uncompressed mode len vs. nlen - 0.55 fix bug: restart_interval not initialized to 0 - 0.54 allow NULL for 'int *comp' - 0.53 fix bug in png 3->4; speedup png decoding - 0.52 png handles req_comp=3,4 directly; minor cleanup; jpeg comments - 0.51 obey req_comp requests, 1-component jpegs return as 1-component, - on 'test' only check type, not whether we support this variant - 0.50 (2006-11-19) - first released version -*/ - - -/* ------------------------------------------------------------------------------- -This software is available under 2 licenses -- choose whichever you prefer. ------------------------------------------------------------------------------- -ALTERNATIVE A - MIT License -Copyright (c) 2017 Sean Barrett -Permission is hereby granted, free of charge, to any person obtaining a copy of -this software and associated documentation files (the "Software"), to deal in -the Software without restriction, including without limitation the rights to -use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies -of the Software, and to permit persons to whom the Software is furnished to do -so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. ------------------------------------------------------------------------------- -ALTERNATIVE B - Public Domain (www.unlicense.org) -This is free and unencumbered software released into the public domain. -Anyone is free to copy, modify, publish, use, compile, sell, or distribute this -software, either in source code form or as a compiled binary, for any purpose, -commercial or non-commercial, and by any means. -In jurisdictions that recognize copyright laws, the author or authors of this -software dedicate any and all copyright interest in the software to the public -domain. We make this dedication for the benefit of the public at large and to -the detriment of our heirs and successors. We intend this dedication to be an -overt act of relinquishment in perpetuity of all present and future rights to -this software under copyright law. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN -ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION -WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ------------------------------------------------------------------------------- -*/ diff --git a/zig/vendor/stb/stb_image_impl.c b/zig/vendor/stb/stb_image_impl.c deleted file mode 100644 index 908982ad8f..0000000000 --- a/zig/vendor/stb/stb_image_impl.c +++ /dev/null @@ -1,10 +0,0 @@ -// stb_image v2.30 (public domain / MIT, github.com/nothings/stb at 2c980bb59875b0d32144a71867fbdebb2f77cd20): -// the native server's JPEG and PNG decoder for image inputs (zig/src/vision/qwen_image.zig). -#define STB_IMAGE_IMPLEMENTATION -#define STBI_ONLY_JPEG -#define STBI_ONLY_PNG -#define STBI_NO_STDIO -#define STBI_NO_HDR -#define STBI_NO_LINEAR -#define STBI_NO_FAILURE_STRINGS -#include "stb_image.h" From c6418f36c35fa5d8321745167d97911898eb08d9 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Tue, 6 Oct 2026 23:14:59 -0400 Subject: [PATCH 007/115] Flash Next's embedding row matches the host dequant of the same MLX group-32 token, so the CUDA replay starts from the row the Python engine reads. --- zig/build/cuda.zig | 20 +- zig/src/families/flashnext/cuda_embed.zig | 171 ++++++++++++++++++ .../families/flashnext/cuda_embed_main.zig | 87 +++++++++ zig/src/families/flashnext/cuda_launch.zig | 50 +++++ 4 files changed, 325 insertions(+), 3 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_embed.zig create mode 100644 zig/src/families/flashnext/cuda_embed_main.zig create mode 100644 zig/src/families/flashnext/cuda_launch.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..d2f71c34bb 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -78,7 +78,7 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return cuda; } -fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { +fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); core.addImport("tokenizer", tokenizer); @@ -90,7 +90,10 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin nemotron.addImport("lanes", lanes); nemotron.addImport("heat", heat); nemotron.addImport("nemotron_draft_ids", draft_ids); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat }; + const flashnext = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_launch.zig"), .target = target, .optimize = optimize, .link_libc = true }); + flashnext.addImport("cuda", cuda); + flashnext.addImport("core", core); + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .flashnext = flashnext, .tokenizer = tokenizer, .heat = heat }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -139,6 +142,17 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); + // not installed: one embedding row through the captured _embed cubin + const embed_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_embed_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + embed_check.addImport("cuda", cuda); + embed_check.addImport("core", mods.core); + embed_check.addImport("flashnext", mods.flashnext); + const embed_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-embed", .root_module = embed_check })); + const embed_model = b.option([]const u8, "flashnext_model", "Flash Next checkpoint directory for the embedding check"); + const embed_kernels = b.option([]const u8, "flashnext_kernels", "captured Triton pack for the embedding check"); + const embed_token = b.option([]const u8, "flashnext_token", "token id for the embedding check") orelse "7"; + if (embed_model != null and embed_kernels != null) embed_run.addArgs(&.{ embed_model.?, embed_kernels.?, embed_token }); + b.step("flashnext-embed", "Upload the Flash Next embedding table and run _embed for one token").dependOn(&embed_run.step); nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; } @@ -210,7 +224,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/src/families/flashnext/cuda_embed.zig b/zig/src/families/flashnext/cuda_embed.zig new file mode 100644 index 0000000000..1bb56f97c0 --- /dev/null +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -0,0 +1,171 @@ +//! The MLX 4-bit group-32 token table: host dequant matches glue._embed, and the same bytes go to the GPU. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const Tensor = core.safetensors.Tensor; + +pub const weight_name = "language_model.model.embed_tokens.weight"; +pub const scales_name = "language_model.model.embed_tokens.scales"; +pub const biases_name = "language_model.model.embed_tokens.biases"; + +const Triple = struct { weight: Tensor, scales: Tensor, biases: Tensor }; + +fn tensorAt(obj: std.json.ObjectMap, data: []const u8, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const o = (obj.get(name) orelse return error.MissingTensor).object; + const got = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + if (got != dtype) return error.UnexpectedTensor; + const shape = o.get("shape").?.array.items; + if (shape.len == 0 or shape.len > 2) return error.UnexpectedTensor; + const offs = o.get("data_offsets").?.array.items; + const begin: usize = @intCast(offs[0].integer); + const end: usize = @intCast(offs[1].integer); + if (end < begin or end > data.len) return error.BadSafetensors; + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = data[begin..end] }; + for (shape, 0..) |d, i| t.shape[i] = @intCast(d.integer); + return t; +} + +/// The three embedding tensors inside one safetensors header. Other tensors, including rank above 4, are ignored. +pub fn readTriple(gpa: std.mem.Allocator, header: []const u8, data: []const u8) !Triple { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, header, .{}); + defer parsed.deinit(); + const obj = parsed.value.object; + return .{ + .weight = try tensorAt(obj, data, weight_name, .u32), + .scales = try tensorAt(obj, data, scales_name, .bf16), + .biases = try tensorAt(obj, data, biases_name, .bf16), + }; +} + +/// The shard that holds the embedding table, mapped until `close`. +pub const Mapped = struct { + io: std.Io, + file: std.Io.File, + map: std.Io.File.MemoryMap, + weight: Tensor, + scales: Tensor, + biases: Tensor, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Mapped { + var file = try std.Io.Dir.cwd().openFile(io, path, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + const triple = try readTriple(gpa, map.memory[8..][0..header_len], map.memory[8 + header_len ..]); + return .{ .io = io, .file = file, .map = map, .weight = triple.weight, .scales = triple.scales, .biases = triple.biases }; + } + + /// One named tensor in this shard. Rank 1 or 2. + pub fn lookup(self: *const Mapped, gpa: std.mem.Allocator, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const header_len: usize = @intCast(std.mem.readInt(u64, self.map.memory[0..8], .little)); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, self.map.memory[8..][0..header_len], .{}); + defer parsed.deinit(); + return tensorAt(parsed.value.object, self.map.memory[8 + header_len ..], name, dtype); + } + + pub fn close(self: *Mapped) void { + self.map.destroy(self.io); + self.file.close(self.io); + self.* = undefined; + } +}; + +pub const Table = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + rows: usize, + dims: usize, + + /// Copies the three tensors. `dims` is hidden size (eight values per stored word). + pub fn upload(d: *const cuda.Driver, weight: Tensor, scales: Tensor, biases: Tensor) !Table { + if (weight.dtype != .u32 or scales.dtype != .bf16 or biases.dtype != .bf16) return error.UnsupportedQuantization; + if (weight.rank != 2 or scales.rank != 2 or biases.rank != 2) return error.UnexpectedTensor; + const rows = weight.dim(0); + const dims = weight.dim(1) * 8; + if (dims == 0 or dims % 32 != 0 or scales.dim(0) != rows or biases.dim(0) != rows) return error.UnexpectedTensor; + if (scales.dim(1) != dims / 32 or biases.dim(1) != dims / 32) return error.UnexpectedTensor; + var w = try cuda.DeviceBuffer.fromHost(d, weight.bytes); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(d, scales.bytes); + errdefer s.free(); + var b = try cuda.DeviceBuffer.fromHost(d, biases.bytes); + errdefer b.free(); + return .{ .w = w, .s = s, .b = b, .rows = rows, .dims = dims }; + } + + pub fn deinit(self: *Table) void { + self.w.free(); + self.s.free(); + self.b.free(); + self.* = undefined; + } +}; + +/// fp32 to bf16, round to nearest even, the rounding Triton uses. +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn f32FromBf16(b: []const u8, i: usize) f32 { + return @bitCast(@as(u32, std.mem.readInt(u16, b[2 * i ..][0..2], .little)) << 16); +} + +/// One token row: each group of 32 is `nibble * scale + bias`, stored bf16. Same order as glue._embed. +pub fn dequant(weight: []const u8, scales: []const u8, biases: []const u8, dims: usize, token: usize, out: []u16) !void { + if (dims == 0 or dims % 32 != 0 or out.len != dims) return error.UnexpectedTensor; + const groups = dims / 32; + const words = dims / 8; + const woff = token * words * 4; + const soff = token * groups * 2; + if (woff + words * 4 > weight.len or soff + groups * 2 > scales.len or soff + groups * 2 > biases.len) return error.UnexpectedTensor; + for (0..groups) |g| { + const scale = f32FromBf16(scales[soff..], g); + const bias = f32FromBf16(biases[soff..], g); + for (0..4) |i| { + const word = std.mem.readInt(u32, weight[woff + (g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + out[g * 32 + i * 8 + j] = toBf16(q * scale + bias); + } + } + } +} + +test "embedding bytes are read when another tensor has rank above 4" { + const header = + \\{"high":{"dtype":"U8","shape":[1,1,1,1,1],"data_offsets":[0,1]}, + \\ "language_model.model.embed_tokens.weight":{"dtype":"U32","shape":[1,4],"data_offsets":[1,17]}, + \\ "language_model.model.embed_tokens.scales":{"dtype":"BF16","shape":[1,1],"data_offsets":[17,19]}, + \\ "language_model.model.embed_tokens.biases":{"dtype":"BF16","shape":[1,1],"data_offsets":[19,21]}} + ; + var data: [21]u8 = @splat(0); + data[1] = 0x21; + std.mem.writeInt(u16, data[17..19], 0x3f80, .little); + const triple = try readTriple(std.testing.allocator, header, &data); + try std.testing.expectEqual(@as(usize, 4), triple.weight.dim(1)); + var out: [32]u16 = undefined; + try dequant(triple.weight.bytes, triple.scales.bytes, triple.biases.bytes, 32, 0, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); +} + +test "one embedding group dequants nibbles with scale and bias" { + var weight: [16]u8 = @splat(0); + std.mem.writeInt(u32, weight[0..4], 0x21, .little); + var scales: [2]u8 = undefined; + std.mem.writeInt(u16, &scales, 0x3f80, .little); + const biases: [2]u8 = @splat(0); + var out: [32]u16 = undefined; + try dequant(&weight, &scales, &biases, 32, 0, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); + try std.testing.expectEqual(@as(u16, 0x4000), out[1]); + try std.testing.expectEqual(@as(u16, 0), out[2]); +} diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig new file mode 100644 index 0000000000..f35a43463a --- /dev/null +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -0,0 +1,87 @@ +//! One token through the captured _embed cubin, compared with the host dequant of the same row. + +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); + +const embed = flash.embed; + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(embed.weight_name).?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 4) { + std.debug.print("usage: flashnext-embed \n", .{}); + return 2; + } + const token = try std.fmt.parseInt(u32, args[3], 10); + const path = try shardPath(gpa, io, args[1]); + defer gpa.free(path); + var mapped = try embed.Mapped.open(gpa, io, path); + defer mapped.close(); + const weight = mapped.weight; + const scales = mapped.scales; + const biases = mapped.biases; + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var table = try embed.Table.upload(&driver, weight, scales, biases); + defer table.deinit(); + if (token >= table.rows) return error.UnexpectedTensor; + const host = try gpa.alloc(u16, table.dims); + defer gpa.free(host); + try embed.dequant(weight.bytes, scales.bytes, biases.bytes, table.dims, token, host); + + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer set.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + var ids = try cuda.DeviceBuffer.fromHost(&driver, std.mem.asBytes(&token)); + defer ids.free(); + var out = try cuda.DeviceBuffer.alloc(&driver, table.dims * 2); + defer out.free(); + try flash.Tri.embed(.{ .set = &set, .s = stream }, ids.ptr, table.w.ptr, table.s.ptr, table.b.ptr, out.ptr, table.dims, 1, 1); + + const gpu = try gpa.alloc(u8, table.dims * 2); + defer gpa.free(gpu); + try stream.synchronize(); + try out.download(0, gpu); + + var mismatch: usize = 0; + for (host, 0..) |want, i| { + const got = std.mem.readInt(u16, gpu[2 * i ..][0..2], .little); + if (got != want) mismatch += 1; + } + std.debug.print("token {d} rows {d} dims {d} mismatch {d} head", .{ token, table.rows, table.dims, mismatch }); + for (0..@min(8, table.dims)) |i| std.debug.print(" {x:0>4}", .{std.mem.readInt(u16, gpu[2 * i ..][0..2], .little)}); + std.debug.print("\n", .{}); + if (mismatch != 0) return 1; + + const streams: usize = 4; + var wide = try cuda.DeviceBuffer.alloc(&driver, streams * table.dims * 2); + defer wide.free(); + try flash.Tri.embed(.{ .set = &set, .s = stream }, ids.ptr, table.w.ptr, table.s.ptr, table.b.ptr, wide.ptr, table.dims, streams, 1); + const wide_bytes = try gpa.alloc(u8, streams * table.dims * 2); + defer gpa.free(wide_bytes); + try stream.synchronize(); + try wide.download(0, wide_bytes); + var copies_differ = false; + for (1..streams) |s| { + if (!std.mem.eql(u8, wide_bytes[0 .. table.dims * 2], wide_bytes[s * table.dims * 2 ..][0 .. table.dims * 2])) copies_differ = true; + } + std.debug.print("embed copies {d} identical {}\n", .{ streams, !copies_differ }); + return if (copies_differ) 1 else 0; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig new file mode 100644 index 0000000000..875cb33f6d --- /dev/null +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -0,0 +1,50 @@ +//! Flash Next's Triton launches from the captured cubins: each grid and constexpr matches the Python wrapper. + +const std = @import("std"); +const cuda = @import("cuda"); +const aot = cuda.aot; + +pub const embed = @import("cuda_embed.zig"); + +const p = aot.ptr; + +fn ci(name: []const u8, v: usize) aot.Const { + return aot.ci(name, @intCast(v)); +} + +fn u(x: usize) u32 { + return @intCast(x); +} + +/// One program per token row and per group of 32 hidden features. +pub fn embedGrid(rows: usize, dims: usize) [3]u32 { + return .{ u(rows), u(dims / 32), 1 }; +} + +pub const Tri = struct { + set: *const aot.Set, + s: cuda.Stream, + + fn run(t: Tri, name: []const u8, grid: [3]u32, args: []const aot.Arg, consts: []const aot.Const) !void { + try t.set.run(t.s, name, grid, args, consts); + } + + /// glue.embed: ids (R,) int32 into (R, copies * dims) bf16, MLX 4-bit group 32. + pub fn embed(t: Tri, ids: u64, weight: u64, scales: u64, biases: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + try t.run("_embed", embedGrid(rows, dims), &.{ + p("IDS", "*i32", ids), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + }, &.{ ci("D", dims), ci("S_COPIES", copies) }); + } +}; + +test { + _ = embed; +} + +test "embed grid is one program per row and per group of 32" { + try std.testing.expectEqual([3]u32{ 4, 80, 1 }, embedGrid(4, 2560)); +} From 1dca5c1e7ccc6571d4fca38e8d46b7577a224380 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Tue, 6 Oct 2026 23:17:17 -0400 Subject: [PATCH 008/115] Layer 0's first hyper-connection matches its host check on one token, so the linear-attention input is that mixed row. --- .../families/flashnext/cuda_embed_main.zig | 216 ++++++++++++++- zig/src/families/flashnext/cuda_hc.zig | 260 ++++++++++++++++++ zig/src/families/flashnext/cuda_launch.zig | 77 ++++++ 3 files changed, 551 insertions(+), 2 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_hc.zig diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index f35a43463a..b8e262e0a2 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -2,6 +2,7 @@ const std = @import("std"); const cuda = @import("cuda"); +const core = @import("core"); const flash = @import("flashnext"); const embed = flash.embed; @@ -78,10 +79,221 @@ pub fn main(init: std.process.Init) !u8 { defer gpa.free(wide_bytes); try stream.synchronize(); try wide.download(0, wide_bytes); + const wide_u16 = try gpa.alloc(u16, streams * table.dims); + defer gpa.free(wide_u16); + for (wide_u16, 0..) |*v, i| v.* = std.mem.readInt(u16, wide_bytes[2 * i ..][0..2], .little); + var copies_differ = false; for (1..streams) |s| { - if (!std.mem.eql(u8, wide_bytes[0 .. table.dims * 2], wide_bytes[s * table.dims * 2 ..][0 .. table.dims * 2])) copies_differ = true; + if (!std.mem.eql(u16, wide_u16[0..table.dims], wide_u16[s * table.dims ..][0..table.dims])) copies_differ = true; } std.debug.print("embed copies {d} identical {}\n", .{ streams, !copies_differ }); - return if (copies_differ) 1 else 0; + + const nc = table.dims / 256; + var pss = try cuda.DeviceBuffer.alloc(&driver, nc * streams * 4); + defer pss.free(); + const tri: flash.Tri = .{ .set = &set, .s = stream }; + try tri.hcWriteback(wide.ptr, wide.ptr, pss.ptr, wide.ptr, "*bf16", wide.ptr, wide.ptr, "*bf16", wide.ptr, "*bf16", table.dims, 1, table.dims, streams, 0, 1, 1, 1); + const pss_bytes = try gpa.alloc(u8, nc * streams * 4); + defer gpa.free(pss_bytes); + try stream.synchronize(); + try pss.download(0, pss_bytes); + const want = try gpa.alloc(f32, nc * streams); + defer gpa.free(want); + try flash.hc.partialSums(wide_u16, table.dims, streams, want); + var max_ulp: u32 = 0; + var streams_differ = false; + for (0..nc) |c| { + const first = std.mem.readInt(u32, pss_bytes[4 * c * streams ..][0..4], .little); + for (0..streams) |s| { + const bits = std.mem.readInt(u32, pss_bytes[4 * (c * streams + s) ..][0..4], .little); + if (bits != first) streams_differ = true; + const got: f32 = @bitCast(bits); + const dist = flash.hc.ulps(got, want[c * streams + s]); + if (dist > max_ulp) max_ulp = dist; + } + } + std.debug.print("writeback streams identical {} max_ulp {d}\n", .{ !streams_differ, max_ulp }); + if (copies_differ or streams_differ or max_ulp > 4) return 1; + + const norm_mismatch = try runDown(gpa, &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); + return if (norm_mismatch == 0) 0 else 1; +} + +fn runDown(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, set: *cuda.aot.Set, mapped: *const embed.Mapped, h: u64, pss: u64, wide_u16: []const u16, pss_bytes: []const u8, dims: usize, streams: usize, nc: usize) !usize { + const prefix = "language_model.model.layers.0.attn_hyper_connection."; + var name: [180]u8 = undefined; + const down_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.weight", .{prefix}), .u32); + const down_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.scales", .{prefix}), .bf16); + const down_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.biases", .{prefix}), .bf16); + const inj_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.weight", .{prefix}), .u32); + const inj_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.scales", .{prefix}), .bf16); + const inj_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.biases", .{prefix}), .bf16); + const gamma = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}hc_norm.weight", .{prefix}), .bf16); + const up_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.weight", .{prefix}), .u32); + const up_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.scales", .{prefix}), .bf16); + const up_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.biases", .{prefix}), .bf16); + const n = down_w.dim(0) + inj_w.dim(0); + const k8 = down_w.dim(1); + const kg = k8 / 4; + if (n != 324 or k8 != 1280 or gamma.dim(0) != streams * dims) return error.UnexpectedTensor; + + const stacked_w = try gpa.alloc(u8, down_w.bytes.len + inj_w.bytes.len); + defer gpa.free(stacked_w); + @memcpy(stacked_w[0..down_w.bytes.len], down_w.bytes); + @memcpy(stacked_w[down_w.bytes.len..], inj_w.bytes); + const stacked_s = try gpa.alloc(u8, down_s.bytes.len + inj_s.bytes.len); + defer gpa.free(stacked_s); + @memcpy(stacked_s[0..down_s.bytes.len], down_s.bytes); + @memcpy(stacked_s[down_s.bytes.len..], inj_s.bytes); + const stacked_b = try gpa.alloc(u8, down_b.bytes.len + inj_b.bytes.len); + defer gpa.free(stacked_b); + @memcpy(stacked_b[0..down_b.bytes.len], down_b.bytes); + @memcpy(stacked_b[down_b.bytes.len..], inj_b.bytes); + const tiled = try flash.hc.tileWords(gpa, stacked_w, n, k8); + defer gpa.free(tiled); + const scales_t = try flash.hc.transposeBf16(gpa, stacked_s, n, kg); + defer gpa.free(scales_t); + const biases_t = try flash.hc.transposeBf16(gpa, stacked_b, n, kg); + defer gpa.free(biases_t); + const scale = try gpa.alloc(f32, gamma.dim(0)); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = @bitCast(@as(u32, std.mem.readInt(u16, gamma.bytes[2 * i ..][0..2], .little)) << 16); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, tiled); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales_t); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases_t); + defer b_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, streams * dims * 2); + defer norm_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + const sk: usize = 32; + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * n * 4); + defer part_b.free(); + const tri: flash.Tri = .{ .set = set, .s = stream.* }; + try tri.hcDown(h, pss, scale_b.ptr, norm_b.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, part_b.ptr, 1e-6, 1, n, k8 * 8, dims, nc, streams, sk); + const low: usize = 320; + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + var inj_out = try cuda.DeviceBuffer.alloc(driver, streams * 2); + defer inj_out.free(); + try tri.hcReduceAct(part_b.ptr, act_b.ptr, xs_b.ptr, inj_out.ptr, sk, 1, streams, low, n, 1); + const up_n = up_w.dim(0); + const up_k8 = up_w.dim(1); + if (up_n != streams * dims or up_k8 != 40) return error.UnexpectedTensor; + const up_tiled = try flash.hc.tileWords(gpa, up_w.bytes, up_n, up_k8); + defer gpa.free(up_tiled); + const up_scales = try flash.hc.transposeBf16(gpa, up_s.bytes, up_n, up_k8 / 4); + defer gpa.free(up_scales); + const up_biases = try flash.hc.transposeBf16(gpa, up_b.bytes, up_n, up_k8 / 4); + defer gpa.free(up_biases); + var up_wb = try cuda.DeviceBuffer.fromHost(driver, up_tiled); + defer up_wb.free(); + var up_sb = try cuda.DeviceBuffer.fromHost(driver, up_scales); + defer up_sb.free(); + var up_bb = try cuda.DeviceBuffer.fromHost(driver, up_biases); + defer up_bb.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, up_wb.ptr, up_sb.ptr, up_bb.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, up_n, up_k8 * 8, dims, streams); + + const got = try gpa.alloc(u8, streams * dims * 2); + defer gpa.free(got); + try stream.synchronize(); + try norm_b.download(0, got); + const pss_f = try gpa.alloc(f32, nc * streams); + defer gpa.free(pss_f); + for (pss_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const want = try gpa.alloc(u16, streams * dims); + defer gpa.free(want); + try flash.hc.normed(wide_u16, pss_f, scale, dims, streams, 1e-6, want); + var mismatch: usize = 0; + for (want, 0..) |w, i| { + if (std.mem.readInt(u16, got[2 * i ..][0..2], .little) != w) mismatch += 1; + } + var part4: [4]u8 = undefined; + try part_b.download(0, &part4); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, &part4, .little)); + const gpu_norm = try gpa.alloc(u16, 320); + defer gpa.free(gpu_norm); + for (gpu_norm, 0..) |*o, i| o.* = std.mem.readInt(u16, got[2 * i ..][0..2], .little); + const host_part = try flash.hc.firstPartial(gpu_norm, stacked_w[0 .. k8 * 4], stacked_s[0 .. kg * 2], stacked_b[0 .. kg * 2], 10); + const col_ulp = flash.hc.ulps(gpu_part, host_part); + std.debug.print("hcdown normed mismatch {d} col0 gpu {d} host {d} ulp {d}\n", .{ mismatch, gpu_part, host_part, col_ulp }); + if (col_ulp > 8) return 1; + + const part_bytes = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(part_bytes); + try part_b.download(0, part_bytes); + const part_f = try gpa.alloc(f32, sk * n); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_bytes[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const inj_want = try gpa.alloc(u16, streams); + defer gpa.free(inj_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try flash.hc.reduceAct(part_f, sk, n, low, streams, act_want, inj_want, xs_want); + const act_bytes = try gpa.alloc(u8, low * 2); + defer gpa.free(act_bytes); + const inj_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_bytes); + const xs_bytes = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_bytes); + try act_b.download(0, act_bytes); + try inj_out.download(0, inj_bytes); + try xs_b.download(0, xs_bytes); + var act_mismatch: usize = 0; + for (act_want, 0..) |w, i| { + if (std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little) != w) act_mismatch += 1; + } + var inj_mismatch: usize = 0; + for (inj_want, 0..) |w, i| { + if (std.mem.readInt(u16, inj_bytes[2 * i ..][0..2], .little) != w) inj_mismatch += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |w, i| { + const got_xs: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = flash.hc.ulps(got_xs, w); + if (dist > xs_ulp) xs_ulp = dist; + } + std.debug.print("reduce act mismatch {d} inject mismatch {d} xs max_ulp {d}\n", .{ act_mismatch, inj_mismatch, xs_ulp }); + if (act_mismatch != 0 or inj_mismatch != 0 or xs_ulp > 4) return 1; + + const mixed_bytes = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_bytes); + try mixed_b.download(0, mixed_bytes); + const act_u16 = try gpa.alloc(u16, low); + defer gpa.free(act_u16); + for (act_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const norm_u16 = try gpa.alloc(u16, streams * dims); + defer gpa.free(norm_u16); + for (norm_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got[2 * i ..][0..2], .little); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try flash.hc.mix(act_u16, xs_f, norm_u16, up_w.bytes, up_s.bytes, up_b.bytes, up_n, up_k8, dims, streams, host_mixed); + var mixed_off: usize = 0; + var max_steps: u32 = 0; + for (host_mixed, 0..) |host_m, i| { + const got_m = std.mem.readInt(u16, mixed_bytes[2 * i ..][0..2], .little); + const dist = flash.hc.mixedSteps(got_m, host_m); + if (dist != 0) mixed_off += 1; + if (dist > max_steps) max_steps = dist; + } + std.debug.print("upmix dims {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); + if (max_steps > 2) return 1; + return mismatch; } diff --git a/zig/src/families/flashnext/cuda_hc.zig b/zig/src/families/flashnext/cuda_hc.zig new file mode 100644 index 0000000000..76a34a344e --- /dev/null +++ b/zig/src/families/flashnext/cuda_hc.zig @@ -0,0 +1,260 @@ +//! Layer-0 hyper-connection host checks: mode 0 writes each 256-wide chunk's sum of squares. + +const std = @import("std"); + +/// glue._hc_writeback mode 0: pss[c, s] is a pairwise sum of squares of stream s, chunk c. +pub fn partialSums(h: []const u16, dims: usize, streams: usize, out: []f32) !void { + const block: usize = 256; + if (dims == 0 or dims % block != 0 or h.len != streams * dims) return error.UnexpectedTensor; + const nc = dims / block; + if (out.len != nc * streams) return error.UnexpectedTensor; + for (0..nc) |c| { + for (0..streams) |s| { + var buf: [256]f32 = undefined; + for (0..block) |i| { + const bits: u32 = h[s * dims + c * block + i]; + const v: f32 = @bitCast(bits << 16); + buf[i] = v * v; + } + var n: usize = block; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[c * streams + s] = buf[0]; + } + } +} + +test "mode 0 partial sums are the squares of one 256-wide chunk" { + var h: [256]u16 = @splat(0); + h[0] = 0x3f80; + var out: [1]f32 = undefined; + try partialSums(&h, 256, 1, &out); + try std.testing.expectEqual(@as(f32, 1), out[0]); +} + +/// Steps between two positive finite floats. Triton's block sum lands a couple of steps from a pairwise sum. +pub fn ulps(a: f32, b: f32) u32 { + const xa: i32 = @bitCast(a); + const xb: i32 = @bitCast(b); + const d = xa - xb; + return @intCast(if (d < 0) -d else d); +} + +test "ulps counts steps between positive floats" { + const a: f32 = 1; + const b: f32 = @bitCast(@as(u32, @bitCast(a)) + 2); + try std.testing.expectEqual(@as(u32, 2), ulps(a, b)); +} + +pub const tile_bn: usize = 64; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// MLX (n, k8) words to [n/64][k/32][64][4], n padded up to 64. qmm.tile_words. +pub fn tileWords(gpa: std.mem.Allocator, words: []const u8, n: usize, k8: usize) ![]u8 { + if (k8 == 0 or k8 % 4 != 0 or words.len != n * k8 * 4) return error.UnexpectedTensor; + const npad = (n + tile_bn - 1) / tile_bn * tile_bn; + const kg = k8 / 4; + const out = try gpa.alloc(u8, (npad / tile_bn) * kg * tile_bn * 4 * 4); + @memset(out, 0); + for (0..n) |r| { + const tile = r / tile_bn; + const local = r % tile_bn; + for (0..kg) |g| { + for (0..4) |w| { + const src = ((r * k8) + g * 4 + w) * 4; + const dst = (((tile * kg + g) * tile_bn + local) * 4 + w) * 4; + @memcpy(out[dst..][0..4], words[src..][0..4]); + } + } + } + return out; +} + +/// (rows, cols) bf16 to (cols, rows), the tiled Q4 scale layout. +pub fn transposeBf16(gpa: std.mem.Allocator, src: []const u8, rows: usize, cols: usize) ![]u8 { + if (src.len != rows * cols * 2) return error.UnexpectedTensor; + const out = try gpa.alloc(u8, src.len); + for (0..rows) |r| { + for (0..cols) |c| { + const from = (r * cols + c) * 2; + const to = (c * rows + r) * 2; + @memcpy(out[to..][0..2], src[from..][0..2]); + } + } + return out; +} + +/// The normed row _qmm_hcdown writes: bf16(h * rinv * scale), one rinv per stream from that stream's partial sums. +pub fn normed(h: []const u16, pss: []const f32, scale: []const f32, dims: usize, streams: usize, eps: f32, out: []u16) !void { + const nc = dims / 256; + if (dims == 0 or dims % 256 != 0 or h.len != streams * dims or pss.len != nc * streams or scale.len != h.len or out.len != h.len) return error.UnexpectedTensor; + for (0..streams) |s| { + var total: f32 = 0; + for (0..nc) |c| total += pss[c * streams + s]; + const rinv = 1.0 / @sqrt(total / @as(f32, @floatFromInt(dims)) + eps); + for (0..dims) |i| { + const hv: f32 = @bitCast(@as(u32, h[s * dims + i]) << 16); + out[s * dims + i] = toBf16(hv * rinv * scale[s * dims + i]); + } + } +} + +/// fp32 dot of the first K-slice against output column 0, groups of 32, from the MLX row (not the tiled copy). +pub fn firstPartial(row_normed: []const u16, row_words: []const u8, scale_row: []const u8, bias_row: []const u8, groups: usize) !f32 { + if (row_normed.len < groups * 32 or row_words.len < groups * 16 or scale_row.len < groups * 2 or bias_row.len < groups * 2) return error.UnexpectedTensor; + var acc: f32 = 0; + for (0..groups) |g| { + var buf: [32]f32 = undefined; + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, row_words[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + const xv: f32 = @bitCast(@as(u32, row_normed[g * 32 + i * 8 + j]) << 16); + buf[i * 8 + j] = xv; + p += xv * q; + } + } + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + const s: f32 = @bitCast(@as(u32, std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)) << 16); + const b: f32 = @bitCast(@as(u32, std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)) << 16); + acc += p * s + buf[0] * b; + } + return acc; +} + +test "tiled words keep row 0 in the first tile's local row 0" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + const out = try tileWords(std.testing.allocator, &words, 1, 4); + defer std.testing.allocator.free(out); + try std.testing.expectEqual(@as(u32, 0x21), std.mem.readInt(u32, out[0..4], .little)); + try std.testing.expectEqual(@as(u32, 0), std.mem.readInt(u32, out[16..20], .little)); +} + +test "the down norm writes bf16 of h times rinv times scale" { + var h: [256]u16 = @splat(0); + h[0] = 0x3f80; + const pss = [_]f32{1}; + var scale: [256]f32 = @splat(1); + var out: [256]u16 = undefined; + try normed(&h, &pss, &scale, 256, 1, 1e-6, &out); + const rinv = 1.0 / @sqrt(1.0 / 256.0 + 1e-6); + try std.testing.expectEqual(toBf16(rinv), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[1]); +} + +fn roundBf16(v: f32) f32 { + return @bitCast(@as(u32, toBf16(v)) << 16); +} + +/// glue._hc_reduce_act: sum the K slices in order, round, divide by the stream count, SiLU, and the inject gates. +pub fn reduceAct(part: []const f32, sk: usize, ndn: usize, low: usize, streams: usize, act: []u16, inj: []u16, xs: []f32) !void { + if (low == 0 or low % 32 != 0 or ndn < low + streams or part.len != sk * ndn or act.len != low or inj.len != streams or xs.len != low / 32) return error.UnexpectedTensor; + for (0..low) |i| { + var sum: f32 = 0; + for (0..sk) |s| sum += part[s * ndn + i]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(streams))); + act[i] = toBf16(y / (1.0 + @exp(-y))); + } + var buf: [32]f32 = undefined; + for (0..low / 32) |g| { + for (0..32) |j| buf[j] = @bitCast(@as(u32, act[g * 32 + j]) << 16); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + xs[g] = buf[0]; + } + for (0..streams) |s| { + var sum: f32 = 0; + for (0..sk) |k| sum += part[k * ndn + low + s]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(streams))); + const sig = roundBf16(1.0 / (1.0 + @exp(-y))); + inj[s] = toBf16(2.0 * sig); + } +} + +test "a zero down-projection sum is SiLU zero and an inject gate of one" { + var part: [33]f32 = @splat(0); + var act: [32]u16 = undefined; + var inj: [1]u16 = undefined; + var xs: [1]f32 = undefined; + try reduceAct(&part, 1, 33, 32, 1, &act, &inj, &xs); + try std.testing.expectEqual(@as(u16, 0), act[0]); + try std.testing.expectEqual(@as(u16, 0x3f80), inj[0]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// One mixed hidden row: up-dot each stream, sigmoid-gate it by that stream's normed value, then average. +pub fn mix(act: []const u16, xs: []const f32, normed_row: []const u16, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize, dims: usize, streams: usize, out: []u16) !void { + const kg = k8 / 4; + if (dims == 0 or k8 == 0 or k8 % 4 != 0 or act.len < kg * 32 or xs.len < kg or normed_row.len != streams * dims or out.len != dims) return error.UnexpectedTensor; + if (words.len != n * k8 * 4 or scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + for (0..dims) |col| { + var total: f32 = 0; + for (0..streams) |s| { + const row = s * dims + col; + var acc: f32 = 0; + const word_row = words[row * k8 * 4 ..][0 .. k8 * 4]; + const scale_row = scales[row * kg * 2 ..][0 .. kg * 2]; + const bias_row = biases[row * kg * 2 ..][0 .. kg * 2]; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, word_row[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + p += promote(act[g * 32 + i * 8 + j]) * q; + } + } + const sc = promote(std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)); + const bias = promote(std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)); + acc += p * sc + xs[g] * bias; + } + const up = roundBf16(acc); + const sig = roundBf16(1.0 / (1.0 + @exp(-up))); + total += roundBf16(sig * promote(normed_row[row])); + } + out[col] = toBf16(total / @as(f32, @floatFromInt(streams))); + } +} + +fn bf16Steps(a: u16, b: u16) u32 { + const oa: i32 = if (a >= 0x8000) ~@as(i32, a) else a; + const ob: i32 = if (b >= 0x8000) ~@as(i32, b) else b; + const d = oa - ob; + return @intCast(if (d < 0) -d else d); +} + +pub fn mixedSteps(gpu: u16, host: u16) u32 { + return bf16Steps(gpu, host); +} + +test "a zero activation mixes to a zero hidden value" { + var act: [32]u16 = @splat(0); + const xs = [_]f32{0}; + var row: [1]u16 = @splat(0); + var words: [16]u8 = @splat(0); + var scale: [2]u8 = @splat(0); + var bias: [2]u8 = @splat(0); + var mixed: [1]u16 = undefined; + try mix(&act, &xs, &row, &words, &scale, &bias, 1, 4, 1, 1, &mixed); + try std.testing.expectEqual(@as(u16, 0), mixed[0]); +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 875cb33f6d..521eb758b3 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -5,9 +5,14 @@ const cuda = @import("cuda"); const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); +pub const hc = @import("cuda_hc.zig"); const p = aot.ptr; +fn int(name: []const u8, v: usize) aot.Arg { + return aot.int(name, @intCast(v)); +} + fn ci(name: []const u8, v: usize) aot.Const { return aot.ci(name, @intCast(v)); } @@ -16,11 +21,20 @@ fn u(x: usize) u32 { return @intCast(x); } +fn cdiv(a: usize, b: usize) u32 { + return u((a + b - 1) / b); +} + /// One program per token row and per group of 32 hidden features. pub fn embedGrid(rows: usize, dims: usize) [3]u32 { return .{ u(rows), u(dims / 32), 1 }; } +/// RMSNorm's BLOCK is the next power of two at or above the group width. +pub fn rmsBlock(group: usize) usize { + return std.math.ceilPowerOfTwo(usize, group) catch unreachable; +} + pub const Tri = struct { set: *const aot.Set, s: cuda.Stream, @@ -39,12 +53,75 @@ pub const Tri = struct { p("OUT", "*bf16", out), }, &.{ ci("D", dims), ci("S_COPIES", copies) }); } + + /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. + pub fn hcReduceAct(t: Tri, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { + try t.run("_hc_reduce_act", .{ u(rows), 1, 1 }, &.{ + p("PART", "*fp32", part), + p("ACT", "*bf16", act), + p("XS", "*fp32", xs), + p("INJ", "*bf16", inj), + int("M", rows), + }, &.{ ci("SK", sk), ci("S", streams), ci("LOW", low), ci("LOWP", rmsBlock(low)), ci("HAS_INJ", has_inj), ci("NDN", ndn) }); + } + + /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. + pub fn hcUpmix(t: Tri, act: u64, xs: u64, weight: u64, scales: u64, biases: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { + try t.run("_qmm_upmix", .{ cdiv(rows, 16), u(dims / 32), 1 }, &.{ + p("X", "*bf16", act), + p("XS", "*fp32", xs), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("NORMED", "*bf16", normed), + p("MIXED", "*bf16", mixed), + p("XSM", "*fp32", xs_mixed), + int("M", rows), + }, &.{ ci("N", n), ci("K", k), ci("D", dims), ci("SS", streams), ci("BM", 16), ci("DB", 32), ci("GPI", 2), ci("SBN", 64) }); + } + + /// qmm.hc_down for one decode row. N=324 is the layer down stacked with inject. M folds when it is 1. + pub fn hcDown(t: Tri, h: u64, pss: u64, scale: u64, normed: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { + try t.run("_qmm_hcdown", .{ cdiv(rows, 16), cdiv(n, 64), u(sk) }, &.{ + p("H", "*bf16", h), + p("PSS", "*fp32", pss), + p("SCALE", "*fp32", scale), + p("NORMED", "*bf16", normed), + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + p("PART", "*fp32", part), + int("M", rows), + aot.float("eps", eps), + }, &.{ ci("N", n), ci("K", k), ci("D", dims), ci("NC", nc), ci("SS", streams), ci("SK", sk), ci("BM", 16), ci("BLOCK_N", 64), ci("GPI", 2), ci("SBN", 64) }); + } + + /// glue.hc_writeback. Pointer types follow the tensor dtypes Triton specialized (mode 0..4). + pub fn hcWriteback(t: Tri, h: u64, hout: u64, pss: u64, branch: u64, branch_ty: []const u8, inj: u64, y: u64, y_ty: []const u8, wts: u64, wts_ty: []const u8, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + try t.run("_hc_writeback", .{ u(rows), u(dims / 256), 1 }, &.{ + p("H", "*bf16", h), + p("HOUT", "*bf16", hout), + p("PSS", "*fp32", pss), + p("BR", branch_ty, branch), + p("INJ", "*bf16", inj), + p("Y", y_ty, y), + p("WTS", wts_ty, wts), + int("RS", rs), + }, &.{ ci("D", dims), ci("S", streams), ci("MODE", mode), ci("TOPK", topk), ci("SLOTS", slots), ci("BLOCK", 256), ci("WORLD", world) }); + } }; test { _ = embed; + _ = hc; } test "embed grid is one program per row and per group of 32" { try std.testing.expectEqual([3]u32{ 4, 80, 1 }, embedGrid(4, 2560)); } + +test "rmsnorm block is the next power of two at or above the group" { + try std.testing.expectEqual(@as(usize, 4096), rmsBlock(2560)); + try std.testing.expectEqual(@as(usize, 16384), rmsBlock(10240)); +} From cd57bfcb7a4cfd8fabc3e2c79d47659a9e65a3d8 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Tue, 6 Oct 2026 23:43:28 -0400 Subject: [PATCH 009/115] Layer 0's linear-attention projection matches the host group-32 dot of the mixed row on one token, so the DeltaNet input is that projection. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/qmm.cu | 230 ++++++++++++++++++ zig/src/cuda/kernels.zig | 1 + .../families/flashnext/cuda_embed_main.zig | 84 +++++++ zig/src/families/flashnext/cuda_launch.zig | 2 + zig/src/families/flashnext/cuda_qmm.zig | 180 ++++++++++++++ 6 files changed, 498 insertions(+) create mode 100644 zig/kernels/cuda/qmm.cu create mode 100644 zig/src/families/flashnext/cuda_qmm.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index d2f71c34bb..c533cce2d7 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -17,6 +17,7 @@ const kernels = [_]Kernel{ .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2 .{ .name = "probe", .flags = &.{"-O3"} }, .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 + .{ .name = "qmm", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py lane matmul, groups of 32 .{ .name = "qmm_prefill", .flags = &.{"-O3"} }, .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, diff --git a/zig/kernels/cuda/qmm.cu b/zig/kernels/cuda/qmm.cu new file mode 100644 index 0000000000..773d1d5926 --- /dev/null +++ b/zig/kernels/cuda/qmm.cu @@ -0,0 +1,230 @@ +// Device code of src/tensorfold/cuda/kernels/qmm.cu (lines 1-230, comments and ATen includes dropped), checked by zig/tests/cuda/copies.py. + +#include +#include +#include +#include + +#include "qmm_frag.cuh" + +namespace tf_qmm { + +using namespace qmm_frag; + +template +__global__ void __launch_bounds__(WM * WN * 32) qmm_kernel( + const __nv_bfloat16* __restrict__ x, const float* __restrict__ xs, const uint32_t* __restrict__ w, + const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ biases, + void* __restrict__ out, float* __restrict__ part, int M, int N, int K, int SK, int npad, int ldx, int group) { + using T = LaneTile; + extern __shared__ __align__(128) unsigned char buf[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int wm = warp / WN, wn = warp % WN; + const int KG = K / GS, per = KG / SK; + const int2 at = tile_of(blockIdx.x, M, N, BM, BN, group); + const int m0 = at.x, n0 = at.y, slice = blockIdx.z, g0 = slice * per; + + auto stage = [&](int s) { return buf + s * T::STAGE; }; + auto load = [&](int s, int g, int m0) { + unsigned char* p = stage(s); + for (int c = tid; c < BM * T::CHUNKS; c += T::THREADS) { + const int r = c / T::CHUNKS, ch = c % T::CHUNKS; + const int row = min(m0 + r, M - 1); + cp16z(p + r * T::ROW + swz(r, ch) * 16, x + static_cast(row) * ldx + g * GS + ch * 8, + m0 + r < M); + } + unsigned char* pw = p + T::X; + constexpr int TILE_BYTES = 64 * GS / 2; // one stored 64-column tile's group block + for (int c = tid; c < T::W / 16; c += T::THREADS) { + const int t = c / (TILE_BYTES / 16), off = c % (TILE_BYTES / 16); + const size_t tile = static_cast(n0 / 64 + t) * KG + g; + cp16(pw + c * 16, reinterpret_cast(w) + tile * TILE_BYTES + off * 16); + } + unsigned char* ps = pw + T::W; + for (int c = tid; c < 2 * (T::S / 16); c += T::THREADS) { + const int which = c / (T::S / 16), off = c % (T::S / 16); + const __nv_bfloat16* src = (which ? biases : scales) + static_cast(g) * npad + n0 + off * 8; + cp16(ps + which * T::S + off * 16, src); + } + float* px = reinterpret_cast(ps + 2 * T::S); + for (int r = tid; r < BM; r += T::THREADS) + cp4(px + r, xs + static_cast(min(m0 + r, M - 1)) * KG + g); + }; + + { + float acc[T::MT][T::NT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f; + auto epilogue = [&](float (&acc_)[T::MT][T::NT][4], const float (&d_)[T::MT][T::NT][4], + const float (&s_)[T::NT][2], const float (&b_)[T::NT][2], const float (&x_)[T::MT][2]) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc_[i][j][e] = __fmaf_rn(x_[i][e >> 1], b_[j][e & 1], __fmaf_rn(d_[i][j][e], s_[j][e & 1], + acc_[i][j][e])); + }; + float dp[T::MT][T::NT][4], sp[T::NT][2], bp[T::NT][2], xp[T::MT][2]; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < per) load(s, g0 + s, m0); + commit(); + } + for (int it = 0; it < per; ++it) { + wait(); + __syncthreads(); + const int next = it + STAGES - 1; + if (next < per) load(next % STAGES, g0 + next, m0); + commit(); + const unsigned char* p = stage(it % STAGES); + const uint32_t* pw = reinterpret_cast(p + T::X); + const __nv_bfloat16* ps = reinterpret_cast(p + T::X + T::W); + const float* px = reinterpret_cast(p + T::X + T::W + 2 * T::S); + uint32_t words[T::NT][GS / 32]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int v = 0; v < GS / 32; ++v) + words[j][v] = pw[((wn * T::NT + j) * 32 + lane) * (GS / 32) + v]; + float d[T::MT][T::NT][4]; +#pragma unroll + for (int kt = 0; kt < GS / 16; ++kt) { + uint32_t a[T::MT][4]; +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int r = wm * (BM / WM) + i * 16 + (lane & 7) + ((lane >> 3) & 1) * 8; + const int ch = kt * 2 + (lane >> 4); + ldmatrix4(a[i], p + r * T::ROW + swz(r, ch) * 16); + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const uint32_t b0 = pair(words[j][kt / 2], (kt & 1) * 8); + const uint32_t b1 = pair(words[j][kt / 2], (kt & 1) * 8 + 4); +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + if (kt == 0) mma0(d[i][j], a[i], b0, b1); + else mma(d[i][j], a[i], b0, b1); + } + } + if (PIPE && kt == 0 && it > 0) epilogue(acc, dp, sp, bp, xp); // the previous group's, between MMAs + } + float sv[T::NT][2], bv[T::NT][2], xv[T::MT][2]; +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = wn * (BN / WN) + j * 8 + (lane & 3) * 2; + const __nv_bfloat162 s2 = *reinterpret_cast(ps + col); + const __nv_bfloat162 b2 = *reinterpret_cast(ps + BN + col); + sv[j][0] = __low2float(s2); sv[j][1] = __high2float(s2); + bv[j][0] = __low2float(b2); bv[j][1] = __high2float(b2); + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) { + const int row = wm * (BM / WM) + i * 16 + (lane >> 2); + xv[i][0] = px[row]; + xv[i][1] = px[row + 8]; + } + if constexpr (PIPE) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) { +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) dp[i][j][e] = d[i][j][e]; + xp[i][0] = xv[i][0]; xp[i][1] = xv[i][1]; + } +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + sp[j][0] = sv[j][0]; sp[j][1] = sv[j][1]; bp[j][0] = bv[j][0]; bp[j][1] = bv[j][1]; + } + } else { + epilogue(acc, d, sv, bv, xv); + } + } + if (PIPE && per > 0) epilogue(acc, dp, sp, bp, xp); + wait<0>(); + __syncthreads(); + if constexpr (CLUSTER) { +#if __CUDA_ARCH__ < 900 + __trap(); // no clusters before sm_90: the host never launches this +#else + constexpr int E = T::MT * T::NT * 4; + auto cluster = cooperative_groups::this_cluster(); + float* mine = reinterpret_cast(buf); + if (slice != 0) { +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) mine[((i * T::NT + j) * 4 + e) * T::THREADS + tid] = acc[i][j][e]; + } + cluster.sync(); + if (slice == 0) { + for (int peer = 1; peer < SK; ++peer) { + const float* theirs = cluster.map_shared_rank(mine, peer); +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) +#pragma unroll + for (int e = 0; e < 4; ++e) + acc[i][j][e] = acc[i][j][e] + theirs[((i * T::NT + j) * 4 + e) * T::THREADS + tid]; + } + } + cluster.sync(); // peers keep their memory until slice 0 has read it + if (slice != 0) return; +#endif + } +#pragma unroll + for (int i = 0; i < T::MT; ++i) +#pragma unroll + for (int j = 0; j < T::NT; ++j) { + const int col = n0 + wn * (BN / WN) + j * 8 + (lane & 3) * 2; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = m0 + wm * (BM / WM) + i * 16 + (lane >> 2) + h * 8; + if (row >= M) continue; + const float v0 = acc[i][j][2 * h], v1 = acc[i][j][2 * h + 1]; + if (SK > 1 && !CLUSTER) { + float* dst = part + (static_cast(slice) * M + row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else if (F32) { + float* dst = reinterpret_cast(out) + static_cast(row) * N + col; + if (col < N) dst[0] = v0; + if (col + 1 < N) dst[1] = v1; + } else { + __nv_bfloat16* dst = reinterpret_cast<__nv_bfloat16*>(out) + static_cast(row) * N + col; + if (col + 1 < N && (N & 1) == 0) + *reinterpret_cast<__nv_bfloat162*>(dst) = __floats2bfloat162_rn(v0, v1); + else { + if (col < N) dst[0] = __float2bfloat16_rn(v0); + if (col + 1 < N) dst[1] = __float2bfloat16_rn(v1); + } + } + } + } + } +} + +template +__global__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ out, long long total, int SK) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total) return; + float acc = part[i]; + for (int s = 1; s < SK; ++s) acc = acc + part[s * total + i]; + if (F32) reinterpret_cast(out)[i] = acc; + else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); +} +} // namespace tf_qmm + +// Decode lane: groups of 32, 16 rows, 64 columns, one by four warps, four stages, bf16 out. +template __global__ void tf_qmm::qmm_kernel<32, 16, 64, 1, 4, 4, false, false, false>( + const __nv_bfloat16*, const float*, const uint32_t*, const __nv_bfloat16*, const __nv_bfloat16*, + void*, float*, int, int, int, int, int, int, int); diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..779cb6db13 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -15,6 +15,7 @@ pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes else &.{}; +pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").bytes else &.{}; pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index b8e262e0a2..fe30902739 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -295,5 +295,89 @@ fn runDown(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s } std.debug.print("upmix dims {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); if (max_steps > 2) return 1; + const proj_steps = try runProj(gpa, driver, stream, mapped, mixed_b, xsm_b); + if (proj_steps != 0) return 1; return mismatch; } + +const proj_names = [_][]const u8{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a" }; +const proj_rows = [_]usize{ 10240, 6144, 48, 48 }; + +fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, mixed: cuda.DeviceBuffer, xs: cuda.DeviceBuffer) !u32 { + const prefix = "language_model.model.layers.0.linear_attn."; + const k8: usize = 320; + const kg: usize = 80; + var name: [160]u8 = undefined; + var n: usize = 0; + var words_len: usize = 0; + var scale_len: usize = 0; + var ws: [4]core.safetensors.Tensor = undefined; + var ss: [4]core.safetensors.Tensor = undefined; + var bs: [4]core.safetensors.Tensor = undefined; + for (proj_names, proj_rows, 0..) |part, rows, i| { + ws[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.weight", .{ prefix, part }), .u32); + ss[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.scales", .{ prefix, part }), .bf16); + bs[i] = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}{s}.biases", .{ prefix, part }), .bf16); + if (ws[i].dim(0) != rows or ws[i].dim(1) != k8 or ss[i].dim(0) != rows or ss[i].dim(1) != kg) return error.UnexpectedTensor; + n += rows; + words_len += ws[i].bytes.len; + scale_len += ss[i].bytes.len; + } + if (n != 16480) return error.UnexpectedTensor; + const k = k8 * 8; + const row = try gpa.alloc(u16, k); + defer gpa.free(row); + const words = try gpa.alloc(u8, words_len); + defer gpa.free(words); + const scales = try gpa.alloc(u8, scale_len); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scale_len); + defer gpa.free(biases); + var wo: usize = 0; + var so: usize = 0; + for (ws, ss, bs) |w, s, b| { + @memcpy(words[wo..][0..w.bytes.len], w.bytes); + wo += w.bytes.len; + @memcpy(scales[so..][0..s.bytes.len], s.bytes); + @memcpy(biases[so..][0..b.bytes.len], b.bytes); + so += s.bytes.len; + } + var lane = try flash.qmm.pack(gpa, words, scales, biases, n, k8); + defer lane.deinit(gpa); + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + try flash.qmm.matmul(driver, stream.*, mixed.ptr, xs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, 1, n, k); + const row_bytes = try gpa.alloc(u8, k * 2); + defer gpa.free(row_bytes); + const xs_bytes = try gpa.alloc(u8, kg * 4); + defer gpa.free(xs_bytes); + const got_bytes = try gpa.alloc(u8, n * 2); + defer gpa.free(got_bytes); + try stream.synchronize(); + try mixed.download(0, row_bytes); + try xs.download(0, xs_bytes); + try out_b.download(0, got_bytes); + for (row, 0..) |*o, i| o.* = std.mem.readInt(u16, row_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, kg); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const host = try gpa.alloc(u16, n); + defer gpa.free(host); + try flash.qmm.dotRow(row, xs_f, words, scales, biases, n, k, host); + var off: usize = 0; + var max_steps: u32 = 0; + for (host, 0..) |want, i| { + const got = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const dist = flash.hc.mixedSteps(got, want); + if (dist != 0) off += 1; + if (dist > max_steps) max_steps = dist; + } + std.debug.print("proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ n, off, max_steps, std.mem.readInt(u16, got_bytes[0..2], .little), host[0] }); + return max_steps; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 521eb758b3..0d78f8446d 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,6 +6,7 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const hc = @import("cuda_hc.zig"); +pub const qmm = @import("cuda_qmm.zig"); const p = aot.ptr; @@ -115,6 +116,7 @@ pub const Tri = struct { test { _ = embed; _ = hc; + _ = qmm; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_qmm.zig b/zig/src/families/flashnext/cuda_qmm.zig new file mode 100644 index 0000000000..7c1c7c3dc1 --- /dev/null +++ b/zig/src/families/flashnext/cuda_qmm.zig @@ -0,0 +1,180 @@ +//! Group-32 lane matmul: MLX words packed as qmm.pack, then the shared CUDA kernel's 16-row tile. + +const std = @import("std"); +const cuda = @import("cuda"); + +pub const group: usize = 32; + +const offsets = [8]u32{ 0, 8, 16, 24, 1, 9, 17, 25 }; + +/// LaneTile<32, 16, 64, 1, 4, 4>::SMEM. One by four warps, four stages. +pub const smem: u32 = 9472; +pub const threads: u32 = 128; + +const symbol: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; + +pub const Packed = struct { + weight: []u8, + scales: []u8, + biases: []u8, + n: usize, + k: usize, + npad: usize, + + pub fn deinit(self: *Packed, gpa: std.mem.Allocator) void { + gpa.free(self.weight); + gpa.free(self.scales); + gpa.free(self.biases); + self.* = undefined; + } +}; + +fn nibble(words: []const u8, k8: usize, col: usize, pos: usize) u32 { + const word_i = pos / 8; + const shift: u5 = @intCast((pos % 8) * 4); + const off = (col * k8 + word_i) * 4; + const word = std.mem.readInt(u32, words[off..][0..4], .little); + return (word >> shift) & 0xF; +} + +/// MLX (n, k/8) words and (n, k/32) scales and biases -> the lane layout. n is padded to 128. +pub fn pack(gpa: std.mem.Allocator, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Packed { + if (k8 == 0 or k8 % 4 != 0 or words.len != n * k8 * 4) return error.UnexpectedTensor; + const k = k8 * 8; + if (k % group != 0) return error.UnexpectedTensor; + const kg = k / group; + if (scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + const npad = (n + 127) / 128 * 128; + const tiles = npad / 64; + const weight = try gpa.alloc(u8, tiles * kg * 8 * 32 * 4); + errdefer gpa.free(weight); + @memset(weight, 0); + for (0..n) |col| { + const tile = col / 64; + const inner = col % 64; + const j = inner / 8; + const r = inner % 8; + for (0..kg) |g| { + for (0..4) |c| { + var bits: u32 = 0; + for (0..8) |p| { + const pos = g * group + 2 * c + offsets[p]; + bits |= nibble(words, k8, col, pos) << @intCast(p * 4); + } + const idx = r * 4 + c; + const dst = (((tile * kg + g) * 8 + j) * 32 + idx) * 4; + std.mem.writeInt(u32, weight[dst..][0..4], bits, .little); + } + } + } + const scales_out = try gpa.alloc(u8, kg * npad * 2); + errdefer gpa.free(scales_out); + const biases_out = try gpa.alloc(u8, kg * npad * 2); + errdefer gpa.free(biases_out); + @memset(scales_out, 0); + @memset(biases_out, 0); + for (0..kg) |g| { + for (0..n) |col| { + const src = (col * kg + g) * 2; + const dst = (g * npad + col) * 2; + @memcpy(scales_out[dst..][0..2], scales[src..][0..2]); + @memcpy(biases_out[dst..][0..2], biases[src..][0..2]); + } + } + return .{ .weight = weight, .scales = scales_out, .biases = biases_out, .n = n, .k = k, .npad = npad }; +} + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// One row: each group is `dot(x, q) * scale + xs * bias`, groups added in order, then rounded to bf16. +pub fn dotRow(x: []const u16, xs: []const f32, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k: usize, out: []u16) !void { + if (k == 0 or k % group != 0 or x.len < k or xs.len < k / group or out.len < n) return error.UnexpectedTensor; + const kg = k / group; + const k8 = k / 8; + if (words.len != n * k8 * 4 or scales.len != n * kg * 2 or biases.len != scales.len) return error.UnexpectedTensor; + for (0..n) |col| { + var acc: f32 = 0; + const word_row = words[col * k8 * 4 ..][0 .. k8 * 4]; + const scale_row = scales[col * kg * 2 ..][0 .. kg * 2]; + const bias_row = biases[col * kg * 2 ..][0 .. kg * 2]; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, word_row[(g * 4 + i) * 4 ..][0..4], .little); + for (0..8) |j| { + const q: f32 = @floatFromInt((word >> @intCast(j * 4)) & 0xF); + p += promote(x[g * group + i * 8 + j]) * q; + } + } + const s = promote(std.mem.readInt(u16, scale_row[g * 2 ..][0..2], .little)); + const b = promote(std.mem.readInt(u16, bias_row[g * 2 ..][0..2], .little)); + acc = @mulAdd(f32, xs[g], b, @mulAdd(f32, p, s, acc)); + } + out[col] = toBf16(acc); + } +} + +/// qmm.cu's 16-row launch for one K slice: x (m, k) bf16 and xs (m, k/32) fp32 into out (m, n) bf16. +pub fn matmul(d: *const cuda.Driver, stream: cuda.Stream, x: u64, xs: u64, weight: u64, scales: u64, biases: u64, out: u64, m: usize, n: usize, k: usize) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (m == 0 or m > 16 or k == 0 or k % group != 0) return error.UnexpectedTensor; + var module = try cuda.Module.load(d, cuda.kernels.qmm); + defer module.unload(); + const f = try module.function(symbol); + try f.allowDynamicShared(smem); + const bm: usize = 16; + const bn: usize = 64; + const rows_t = (m + bm - 1) / bm; + const cols_t = (n + bn - 1) / bn; + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (n + 127) / 128 * 128; + const ldx = k; + var args: cuda.Args = .{}; + args.add(x); + args.add(xs); + args.add(weight); + args.add(scales); + args.add(biases); + args.add(out); + args.add(@as(u64, 0)); + for ([_]usize{ m, n, k, 1, npad, ldx, sweep }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t) }, .block = .{ .x = threads }, .shared = smem }, stream, &args); +} + +test "one group's first nibble pair lands in the lane word" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + const bias: [2]u8 = @splat(0); + var got = try pack(std.testing.allocator, &words, &scale, &bias, 1, 4); + defer got.deinit(std.testing.allocator); + try std.testing.expectEqual(@as(usize, 128), got.npad); + try std.testing.expectEqual(@as(u32, 0x20001), std.mem.readInt(u32, got.weight[0..4], .little)); + try std.testing.expectEqual(@as(u32, 0), std.mem.readInt(u32, got.weight[4..8], .little)); + try std.testing.expectEqual(@as(u16, 0x3f80), std.mem.readInt(u16, got.scales[0..2], .little)); +} + +test "scale times the nibble dot plus the group bias rounds to five" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + var bias: [2]u8 = undefined; + std.mem.writeInt(u16, &bias, 0x3f80, .little); + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[1] = 0x3f80; + const xs = [_]f32{2}; + var out: [1]u16 = undefined; + try dotRow(&x, &xs, &words, &scale, &bias, 1, 32, &out); + try std.testing.expectEqual(@as(u16, 0x40a0), out[0]); +} From 37dec96fe8241f4c202ec88ed6153ad1bdf5a9d5 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Tue, 6 Oct 2026 23:55:10 -0400 Subject: [PATCH 010/115] Layer 0's linear-attention convolution, recurrence and output multiply match the host chain on one token, all 6144 gated values and all 2560 projected values, so the mixer branch is that row. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/flashnext_gdn.cu | 143 +++++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_embed.zig | 4 +- .../families/flashnext/cuda_embed_main.zig | 3 +- zig/src/families/flashnext/cuda_gdn.zig | 396 ++++++++++++++++++ zig/src/families/flashnext/cuda_launch.zig | 2 + 7 files changed, 547 insertions(+), 3 deletions(-) create mode 100644 zig/kernels/cuda/flashnext_gdn.cu create mode 100644 zig/src/families/flashnext/cuda_gdn.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index c533cce2d7..3b8e975af7 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -15,6 +15,7 @@ const glue = &[_][]const u8{ "-O3", "--fmad=false", "--ftz=false", "--prec-div=t const kernels = [_]Kernel{ .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2 + .{ .name = "flashnext_gdn", .flags = &.{ "-O3", "--fmad=false" } }, // qwen4_exp/cuda/gdn.cu chain, one row .{ .name = "probe", .flags = &.{"-O3"} }, .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 .{ .name = "qmm", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py lane matmul, groups of 32 diff --git a/zig/kernels/cuda/flashnext_gdn.cu b/zig/kernels/cuda/flashnext_gdn.cu new file mode 100644 index 0000000000..5ce476b7df --- /dev/null +++ b/zig/kernels/cuda/flashnext_gdn.cu @@ -0,0 +1,143 @@ +// One value head per block: depthwise conv, the delta-rule step, then the gated RMSNorm. 16 key heads, 48 value heads. +// Same operations as the one-row chain, compiled without FMA contraction. + +#include +#include + +constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4; + +static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +static __device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); } + +static __device__ __forceinline__ void update(float (&s)[4][4], const float (&kk)[4], const float* vrow, int warp, float g, + float beta) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * g; + kv = kv + s[j][i] * kk[i]; + } + kv = warp_sum(kv); + const float delta = (vrow[warp * 4 + j] - kv) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } +} + +extern "C" __global__ void __launch_bounds__(1024) flashnext_gdn_chain( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, const float* __restrict__ state_in, const float* __restrict__ a_log, + const float* __restrict__ dt_bias, const __nv_bfloat16* __restrict__ norm_w, float eps, int rows, + __nv_bfloat16* __restrict__ out, float* __restrict__ xs, float* __restrict__ state_out) { + constexpr int C = 2 * NK * DK + NV * DV; + constexpr int PW = C + NV * DV + 2 * NV; + const int hv = blockIdx.x, hk = hv / (NV / NK); + const int t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float qs[DK], ks[DK], vs[DV], ys[DV]; + __shared__ float gates[2], rinv; + int c = -1; + if (t < DK) c = hk * DK + t; + else if (t < 2 * DK) c = NK * DK + hk * DK + (t - DK); + else if (t < 2 * DK + DV) c = 2 * NK * DK + hv * DV + (t - 2 * DK); + float s[4][4]; + const size_t sbase = (size_t)hv * DV * DK; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i]; + const __nv_bfloat16* pz = P + C + hv * DV + (t < DV ? t : 0); + const __nv_bfloat16* pb = P + C + NV * DV + hv; + const __nv_bfloat16* pa = pb + NV; + for (int r = 0; r < rows; ++r) { + if (c >= 0) { + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) { + const int at = r + tap; + const float x = at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c]); + acc = acc + __bfloat162float(cw[c * TAPS + tap]) * x; + } + const float act = bf(acc / (1.0f + expf(-acc))); + if (t < DK) qs[t] = act; + else if (t < 2 * DK) ks[t - DK] = act; + else vs[t - 2 * DK] = act; + } + __syncthreads(); + if (warp < 2) { + float* x = warp == 0 ? qs : ks; + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v4[i] = x[lane * 4 + i]; + ss = ss + v4[i] * v4[i]; + } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + if (warp == 0) inv = inv * (1.0f / sqrtf((float)DK)); + __syncwarp(); +#pragma unroll + for (int i = 0; i < 4; ++i) x[lane * 4 + i] = v4[i] * inv; + } else if (warp == 2 && lane == 0) { + const float b = __bfloat162float(pb[(size_t)r * PW]); + const float a = __bfloat162float(pa[(size_t)r * PW]); + gates[0] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); + gates[1] = bf(sigmoidf_(b)); + } + __syncthreads(); + const float g = gates[0], beta = gates[1]; + float kk[4], qq[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + kk[i] = ks[lane * 4 + i]; + qq[i] = qs[lane * 4 + i]; + } + update(s, kk, vs, warp, g, beta); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + o = warp_sum(o); + if (lane == 0) ys[warp * 4 + j] = bf(o); + } + __syncthreads(); + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float y = ys[lane * 4 + i]; + ss = ss + y * y; + } + ss = warp_sum(ss); + if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps); + } + __syncthreads(); + if (t < DV) { + const float yn = bf(bf(ys[t] * rinv) * __bfloat162float(norm_w[t])); + const float z = __bfloat162float(pz[(size_t)r * PW]); + const float o = bf(yn * sigmoidf_(z)); + out[(size_t)r * NV * DV + hv * DV + t] = __float2bfloat16_rn(o); + const float gs = warp_sum(o); + if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs; + } + __syncthreads(); + } + if (state_out != nullptr) { +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + state_out[sbase + (size_t)(warp * 4 + j) * DK + lane * 4 + i] = s[j][i]; + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 779cb6db13..4f3a78a50c 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -13,6 +13,7 @@ fn Blob(comptime import_name: []const u8) type { pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; +pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes else &.{}; pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_embed.zig b/zig/src/families/flashnext/cuda_embed.zig index 1bb56f97c0..0b31f38a69 100644 --- a/zig/src/families/flashnext/cuda_embed.zig +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -17,7 +17,7 @@ fn tensorAt(obj: std.json.ObjectMap, data: []const u8, name: []const u8, dtype: const got = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; if (got != dtype) return error.UnexpectedTensor; const shape = o.get("shape").?.array.items; - if (shape.len == 0 or shape.len > 2) return error.UnexpectedTensor; + if (shape.len == 0 or shape.len > 4) return error.UnexpectedTensor; const offs = o.get("data_offsets").?.array.items; const begin: usize = @intCast(offs[0].integer); const end: usize = @intCast(offs[1].integer); @@ -61,7 +61,7 @@ pub const Mapped = struct { return .{ .io = io, .file = file, .map = map, .weight = triple.weight, .scales = triple.scales, .biases = triple.biases }; } - /// One named tensor in this shard. Rank 1 or 2. + /// One named tensor in this shard. Rank 1 to 4. pub fn lookup(self: *const Mapped, gpa: std.mem.Allocator, name: []const u8, dtype: core.safetensors.DType) !Tensor { const header_len: usize = @intCast(std.mem.readInt(u64, self.map.memory[0..8], .little)); const parsed = try std.json.parseFromSlice(std.json.Value, gpa, self.map.memory[8..][0..header_len], .{}); diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index fe30902739..a22f64ef4d 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -379,5 +379,6 @@ fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, m if (dist > max_steps) max_steps = dist; } std.debug.print("proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ n, off, max_steps, std.mem.readInt(u16, got_bytes[0..2], .little), host[0] }); - return max_steps; + if (max_steps != 0) return max_steps; + return flash.gdn.memory(gpa, driver, stream, mapped, out_b, got_bytes); } diff --git a/zig/src/families/flashnext/cuda_gdn.zig b/zig/src/families/flashnext/cuda_gdn.zig new file mode 100644 index 0000000000..1e9c54c55c --- /dev/null +++ b/zig/src/families/flashnext/cuda_gdn.zig @@ -0,0 +1,396 @@ +//! One-row Gated DeltaNet: depthwise conv, the delta-rule state, the gated norm, then the output projection. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +pub const nk: usize = 16; +pub const nv: usize = 48; +pub const dk: usize = 128; +pub const dv: usize = 128; +pub const taps: usize = 4; +pub const conv_dim: usize = 2 * nk * dk + nv * dv; +pub const proj_width: usize = conv_dim + nv * dv + 2 * nv; +pub const value_dim: usize = nv * dv; +pub const groups: usize = value_dim / qmm.group; + +const symbol: [:0]const u8 = "flashnext_gdn_chain"; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn bf(v: f32) f32 { + return promote(toBf16(v)); +} + +fn sigmoid(x: f32) f32 { + return 1.0 / (1.0 + @exp(-x)); +} + +fn softplus(x: f32) f32 { + return if (x > 20) x else @log(1.0 + @exp(x)); +} + +fn warpSum(v: *[32]f32) void { + var span: usize = 16; + while (span != 0) : (span >>= 1) { + var next: [32]f32 = undefined; + for (0..32) |lane| next[lane] = v[lane] + v[lane ^ span]; + v.* = next; + } +} + +fn convAct(proj: []const u16, cs: []const u16, cw: []const u16, row: usize, channel: usize) f32 { + var acc: f32 = 0; + for (0..taps) |tap| { + const at = row + tap; + const x: f32 = if (at < taps - 1) promote(cs[at * conv_dim + channel]) else promote(proj[(at - (taps - 1)) * proj_width + channel]); + acc = acc + promote(cw[channel * taps + tap]) * x; + } + return bf(acc / (1.0 + @exp(-acc))); +} + +fn normPair(x: []f32, scale_dk: bool) void { + var part: [32]f32 = undefined; + for (0..32) |lane| { + var ss: f32 = 0; + for (0..4) |i| { + const v = x[lane * 4 + i]; + ss = ss + v * v; + } + part[lane] = ss; + } + warpSum(&part); + var inv = 1.0 / @sqrt(part[0] + 1e-6); + if (scale_dk) inv = inv * (1.0 / @sqrt(@as(f32, @floatFromInt(dk)))); + for (0..dk) |i| x[i] = x[i] * inv; +} + +fn update(state: []f32, base: usize, ks: []const f32, vs: []const f32, g: f32, beta: f32) void { + for (0..32) |warp| { + var kk: [32][4]f32 = undefined; + for (0..32) |lane| { + for (0..4) |i| kk[lane][i] = ks[lane * 4 + i]; + } + for (0..4) |j| { + var kv: [32]f32 = undefined; + const row = warp * 4 + j; + for (0..32) |lane| { + var acc: f32 = 0; + for (0..4) |i| { + const at = base + row * dk + lane * 4 + i; + state[at] = state[at] * g; + acc = acc + state[at] * kk[lane][i]; + } + kv[lane] = acc; + } + warpSum(&kv); + const delta = (vs[row] - kv[0]) * beta; + for (0..32) |lane| { + for (0..4) |i| { + const at = base + row * dk + lane * 4 + i; + state[at] = state[at] + kk[lane][i] * delta; + } + } + } + } +} + +fn readout(state: []const f32, base: usize, qs: []const f32, ys: []f32) void { + for (0..32) |warp| { + for (0..4) |j| { + var part: [32]f32 = undefined; + const row = warp * 4 + j; + for (0..32) |lane| { + var acc: f32 = 0; + for (0..4) |i| acc = acc + state[base + row * dk + lane * 4 + i] * qs[lane * 4 + i]; + part[lane] = acc; + } + warpSum(&part); + ys[row] = bf(part[0]); + } + } +} + +/// One fresh or continued row. `state` is read and written in place, layout [value head, row, column]. +pub fn chain(proj: []const u16, cs: []const u16, cw: []const u16, state: []f32, a_log: []const f32, dt_bias: []const f32, norm_w: []const u16, eps: f32, rows: usize, out: []u16, xs: []f32) !void { + if (rows == 0 or proj.len < rows * proj_width or cs.len < (taps - 1) * conv_dim or cw.len < conv_dim * taps) return error.UnexpectedTensor; + if (state.len < nv * dv * dk or a_log.len < nv or dt_bias.len < nv or norm_w.len < dv) return error.UnexpectedTensor; + if (out.len < rows * value_dim or xs.len < rows * groups) return error.UnexpectedTensor; + for (0..rows) |r| { + for (0..nv) |hv| { + const hk = hv / (nv / nk); + var qs: [dk]f32 = undefined; + var ks: [dk]f32 = undefined; + var vs: [dv]f32 = undefined; + for (0..dk) |t| qs[t] = convAct(proj, cs, cw, r, hk * dk + t); + for (0..dk) |t| ks[t] = convAct(proj, cs, cw, r, nk * dk + hk * dk + t); + for (0..dv) |t| vs[t] = convAct(proj, cs, cw, r, 2 * nk * dk + hv * dv + t); + normPair(&qs, true); + normPair(&ks, false); + const b = promote(proj[r * proj_width + conv_dim + nv * dv + hv]); + const a = promote(proj[r * proj_width + conv_dim + nv * dv + nv + hv]); + const g = @exp(-@exp(a_log[hv]) * softplus(a + dt_bias[hv])); + const beta = bf(sigmoid(b)); + const base = hv * dv * dk; + update(state, base, &ks, &vs, g, beta); + var ys: [dv]f32 = undefined; + readout(state, base, &qs, &ys); + var part: [32]f32 = undefined; + for (0..32) |lane| { + var ss: f32 = 0; + for (0..4) |i| { + const y = ys[lane * 4 + i]; + ss = ss + y * y; + } + part[lane] = ss; + } + warpSum(&part); + const rinv = 1.0 / @sqrt(part[0] / @as(f32, @floatFromInt(dv)) + eps); + var gated: [dv]f32 = undefined; + for (0..dv) |t| { + const yn = bf(bf(ys[t] * rinv) * promote(norm_w[t])); + const z = promote(proj[r * proj_width + conv_dim + hv * dv + t]); + gated[t] = bf(yn * sigmoid(z)); + out[r * value_dim + hv * dv + t] = toBf16(gated[t]); + } + for (0..dv / 32) |warp| { + for (0..32) |lane| part[lane] = gated[warp * 32 + lane]; + warpSum(&part); + xs[r * groups + hv * (dv / 32) + warp] = part[0]; + } + } + } +} + +fn orderedUlps(a: f32, b: f32) u32 { + const ai: i32 = @bitCast(a); + const bi: i32 = @bitCast(b); + const ao: i32 = if (ai < 0) std.math.minInt(i32) - ai else ai; + const bo: i32 = if (bi < 0) std.math.minInt(i32) - bi else bi; + const d = ao - bo; + return @intCast(if (d < 0) -d else d); +} + +/// The one-row kernel. Conv state is [3, conv channels], recurrence state is [heads, 128, 128]. +pub fn launch(d: *const cuda.Driver, stream: cuda.Stream, proj: u64, cs: u64, cw: u64, state_in: u64, a_log: u64, dt_bias: u64, norm_w: u64, eps: f32, out: u64, xs: u64, state_out: u64) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var module = try cuda.Module.load(d, cuda.kernels.flashnext_gdn); + defer module.unload(); + const f = try module.function(symbol); + var args: cuda.Args = .{}; + args.add(proj); + args.add(cs); + args.add(cw); + args.add(state_in); + args.add(a_log); + args.add(dt_bias); + args.add(norm_w); + args.add(eps); + args.add(@as(c_int, 1)); + args.add(out); + args.add(xs); + args.add(state_out); + try cuda.launch.launch(f, .{ .grid = .{ .x = nv }, .block = .{ .x = 1024 } }, stream, &args); +} + +fn f32FromBf16(bytes: []const u8) f32 { + return promote(std.mem.readInt(u16, bytes[0..2], .little)); +} + +/// Layer 0's conv, recurrence and output projection for the projection row already on the device. +pub fn memory(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, proj_dev: cuda.DeviceBuffer, proj_bytes: []const u8) !u32 { + const prefix = "language_model.model.layers.0.linear_attn."; + var name: [160]u8 = undefined; + const conv = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}conv1d.weight", .{prefix}), .bf16); + const a_bits = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}A_log", .{prefix}), .bf16); + const dt_bits = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}dt_bias", .{prefix}), .bf16); + const norm = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}norm.weight", .{prefix}), .bf16); + const out_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.weight", .{prefix}), .u32); + const out_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.scales", .{prefix}), .bf16); + const out_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}out_proj.biases", .{prefix}), .bf16); + if (conv.dim(0) != conv_dim or conv.dim(1) != taps or norm.dim(0) != dv) return error.UnexpectedTensor; + if (a_bits.dim(0) != nv or dt_bits.dim(0) != nv) return error.UnexpectedTensor; + if (out_w.dim(0) != 2560 or out_w.dim(1) != value_dim / 8) return error.UnexpectedTensor; + if (proj_bytes.len != proj_width * 2) return error.UnexpectedTensor; + + const proj = try gpa.alloc(u16, proj_width); + defer gpa.free(proj); + for (proj, 0..) |*o, i| o.* = std.mem.readInt(u16, proj_bytes[2 * i ..][0..2], .little); + const cs_bytes = try gpa.alloc(u8, (taps - 1) * conv_dim * 2); + defer gpa.free(cs_bytes); + @memset(cs_bytes, 0); + const a_log = try gpa.alloc(f32, nv); + defer gpa.free(a_log); + const dt_bias = try gpa.alloc(f32, nv); + defer gpa.free(dt_bias); + for (0..nv) |i| { + a_log[i] = f32FromBf16(a_bits.bytes[2 * i ..]); + dt_bias[i] = f32FromBf16(dt_bits.bytes[2 * i ..]); + } + const state_n = nv * dv * dk; + const state = try gpa.alloc(f32, state_n); + defer gpa.free(state); + @memset(state, 0); + var st_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(state)); + defer st_b.free(); + const cs = try gpa.alloc(u16, (taps - 1) * conv_dim); + defer gpa.free(cs); + @memset(cs, 0); + const cw = try gpa.alloc(u16, conv_dim * taps); + defer gpa.free(cw); + for (cw, 0..) |*o, i| o.* = std.mem.readInt(u16, conv.bytes[2 * i ..][0..2], .little); + const norm_w = try gpa.alloc(u16, dv); + defer gpa.free(norm_w); + for (norm_w, 0..) |*o, i| o.* = std.mem.readInt(u16, norm.bytes[2 * i ..][0..2], .little); + const host_out = try gpa.alloc(u16, value_dim); + defer gpa.free(host_out); + const host_xs = try gpa.alloc(f32, groups); + defer gpa.free(host_xs); + const eps: f32 = 1e-6; + try chain(proj, cs, cw, state, a_log, dt_bias, norm_w, eps, 1, host_out, host_xs); + + var cs_b = try cuda.DeviceBuffer.fromHost(driver, cs_bytes); + defer cs_b.free(); + var cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + defer cw_b.free(); + var a_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(a_log)); + defer a_b.free(); + var dt_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(dt_bias)); + defer dt_b.free(); + var norm_b = try cuda.DeviceBuffer.fromHost(driver, norm.bytes); + defer norm_b.free(); + var st_out = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + defer st_out.free(); + var gout = try cuda.DeviceBuffer.alloc(driver, value_dim * 2); + defer gout.free(); + var gxs = try cuda.DeviceBuffer.alloc(driver, groups * 4); + defer gxs.free(); + try launch(driver, stream.*, proj_dev.ptr, cs_b.ptr, cw_b.ptr, st_b.ptr, a_b.ptr, dt_b.ptr, norm_b.ptr, eps, gout.ptr, gxs.ptr, st_out.ptr); + + const got_bytes = try gpa.alloc(u8, value_dim * 2); + defer gpa.free(got_bytes); + const xs_bytes = try gpa.alloc(u8, groups * 4); + defer gpa.free(xs_bytes); + const st_bytes = try gpa.alloc(u8, state_n * 4); + defer gpa.free(st_bytes); + try stream.synchronize(); + try gout.download(0, got_bytes); + try gxs.download(0, xs_bytes); + try st_out.download(0, st_bytes); + + var out_off: usize = 0; + var out_steps: u32 = 0; + for (host_out, 0..) |want, i| { + const got = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) out_off += 1; + if (dist > out_steps) out_steps = dist; + } + var xs_off: usize = 0; + var xs_ulp: u32 = 0; + for (host_xs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = orderedUlps(got, want); + if (dist != 0) xs_off += 1; + if (dist > xs_ulp) xs_ulp = dist; + } + var st_off: usize = 0; + var st_ulp: u32 = 0; + for (state, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, st_bytes[4 * i ..][0..4], .little)); + const dist = orderedUlps(got, want); + if (dist != 0) st_off += 1; + if (dist > st_ulp) st_ulp = dist; + } + const got_head = std.mem.readInt(u16, got_bytes[0..2], .little); + std.debug.print("gdn out {d} off {d} max_steps {d} xs_off {d} xs_ulp {d} state_off {d} state_ulp {d} head {x:0>4} host {x:0>4}\n", .{ value_dim, out_off, out_steps, xs_off, xs_ulp, st_off, st_ulp, got_head, host_out[0] }); + + const got_u16 = try gpa.alloc(u16, value_dim); + defer gpa.free(got_u16); + for (got_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + const got_xs = try gpa.alloc(f32, groups); + defer gpa.free(got_xs); + for (got_xs, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const host_proj = try gpa.alloc(u16, 2560); + defer gpa.free(host_proj); + try qmm.dotRow(got_u16, got_xs, out_w.bytes, out_s.bytes, out_b.bytes, 2560, value_dim, host_proj); + var lane = try qmm.pack(gpa, out_w.bytes, out_s.bytes, out_b.bytes, 2560, value_dim / 8); + defer lane.deinit(gpa); + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var branch = try cuda.DeviceBuffer.alloc(driver, 2560 * 2); + defer branch.free(); + try qmm.matmul(driver, stream.*, gout.ptr, gxs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, branch.ptr, 1, 2560, value_dim); + const branch_bytes = try gpa.alloc(u8, 2560 * 2); + defer gpa.free(branch_bytes); + try stream.synchronize(); + try branch.download(0, branch_bytes); + var proj_off: usize = 0; + var proj_steps: u32 = 0; + for (host_proj, 0..) |want, i| { + const got = std.mem.readInt(u16, branch_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) proj_off += 1; + if (dist > proj_steps) proj_steps = dist; + } + std.debug.print("gdn out_proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ 2560, proj_off, proj_steps, std.mem.readInt(u16, branch_bytes[0..2], .little), host_proj[0] }); + if (out_off != 0 or xs_off != 0 or st_off != 0 or proj_off != 0) return 1; + return 0; +} + +test "warp sum is the sum of the lanes" { + var v: [32]f32 = @splat(0); + v[0] = 1; + v[31] = 3; + warpSum(&v); + try std.testing.expectEqual(@as(f32, 4), v[0]); + try std.testing.expectEqual(@as(f32, 4), v[17]); +} + +test "a zero projection and zero weights stay a zero row" { + const gpa = std.testing.allocator; + const proj = try gpa.alloc(u16, proj_width); + defer gpa.free(proj); + @memset(proj, 0); + const cs = try gpa.alloc(u16, (taps - 1) * conv_dim); + defer gpa.free(cs); + @memset(cs, 0); + const cw = try gpa.alloc(u16, conv_dim * taps); + defer gpa.free(cw); + @memset(cw, 0); + const state = try gpa.alloc(f32, nv * dv * dk); + defer gpa.free(state); + @memset(state, 0); + const a_log = try gpa.alloc(f32, nv); + defer gpa.free(a_log); + @memset(a_log, 0); + const dt = try gpa.alloc(f32, nv); + defer gpa.free(dt); + @memset(dt, 0); + const norm = try gpa.alloc(u16, dv); + defer gpa.free(norm); + @memset(norm, 0); + const out = try gpa.alloc(u16, value_dim); + defer gpa.free(out); + const xs = try gpa.alloc(f32, groups); + defer gpa.free(xs); + try chain(proj, cs, cw, state, a_log, dt, norm, 1e-6, 1, out, xs); + try std.testing.expectEqual(@as(u16, 0), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[value_dim - 1]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); + try std.testing.expectEqual(@as(f32, 0), state[0]); +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 0d78f8446d..cad16d36b0 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -7,6 +7,7 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const hc = @import("cuda_hc.zig"); pub const qmm = @import("cuda_qmm.zig"); +pub const gdn = @import("cuda_gdn.zig"); const p = aot.ptr; @@ -117,6 +118,7 @@ test { _ = embed; _ = hc; _ = qmm; + _ = gdn; } test "embed grid is one program per row and per group of 32" { From 2590dfce1268ea58c122065f1771702113b73216 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 00:14:16 -0400 Subject: [PATCH 011/115] Layer 0's second hyper-connection matches its host check on one token, the write-back exact and its sums within 1 ulp, one normed value 1 step off, and the mixed row within 2 steps, so the expert input is that row. --- .../families/flashnext/cuda_embed_main.zig | 10 +- zig/src/families/flashnext/cuda_gdn.zig | 21 +- zig/src/families/flashnext/cuda_hc.zig | 47 ++++ zig/src/families/flashnext/cuda_launch.zig | 2 + zig/src/families/flashnext/cuda_mlp.zig | 255 ++++++++++++++++++ zig/src/families/flashnext/cuda_qmm.zig | 31 +++ 6 files changed, 358 insertions(+), 8 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_mlp.zig diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index a22f64ef4d..138a5974aa 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -295,15 +295,15 @@ fn runDown(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s } std.debug.print("upmix dims {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); if (max_steps > 2) return 1; - const proj_steps = try runProj(gpa, driver, stream, mapped, mixed_b, xsm_b); - if (proj_steps != 0) return 1; + const rest = try runProj(gpa, driver, stream, mapped, mixed_b, xsm_b, tri, h, inj_out.ptr, wide_u16); + if (rest != 0) return 1; return mismatch; } const proj_names = [_][]const u8{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a" }; const proj_rows = [_]usize{ 10240, 6144, 48, 48 }; -fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, mixed: cuda.DeviceBuffer, xs: cuda.DeviceBuffer) !u32 { +fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, mixed: cuda.DeviceBuffer, xs: cuda.DeviceBuffer, tri: flash.Tri, h: u64, inj: u64, hidden: []const u16) !u32 { const prefix = "language_model.model.layers.0.linear_attn."; const k8: usize = 320; const kg: usize = 80; @@ -380,5 +380,7 @@ fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, m } std.debug.print("proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ n, off, max_steps, std.mem.readInt(u16, got_bytes[0..2], .little), host[0] }); if (max_steps != 0) return max_steps; - return flash.gdn.memory(gpa, driver, stream, mapped, out_b, got_bytes); + const slices = try flash.gdn.memory(gpa, driver, stream, mapped, out_b, got_bytes) orelse return 1; + defer gpa.free(slices); + return flash.mlp.connect(flash.Tri, gpa, driver, stream, mapped, tri, h, inj, hidden, slices); } diff --git a/zig/src/families/flashnext/cuda_gdn.zig b/zig/src/families/flashnext/cuda_gdn.zig index 1e9c54c55c..b53fb5d50d 100644 --- a/zig/src/families/flashnext/cuda_gdn.zig +++ b/zig/src/families/flashnext/cuda_gdn.zig @@ -207,8 +207,8 @@ fn f32FromBf16(bytes: []const u8) f32 { return promote(std.mem.readInt(u16, bytes[0..2], .little)); } -/// Layer 0's conv, recurrence and output projection for the projection row already on the device. -pub fn memory(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, proj_dev: cuda.DeviceBuffer, proj_bytes: []const u8) !u32 { +/// Layer 0's conv, recurrence and output projection. A match also returns the output face's eight unreduced K slices. +pub fn memory(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, proj_dev: cuda.DeviceBuffer, proj_bytes: []const u8) !?[]f32 { const prefix = "language_model.model.layers.0.linear_attn."; var name: [160]u8 = undefined; const conv = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}conv1d.weight", .{prefix}), .bf16); @@ -348,8 +348,21 @@ pub fn memory(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream if (dist > proj_steps) proj_steps = dist; } std.debug.print("gdn out_proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ 2560, proj_off, proj_steps, std.mem.readInt(u16, branch_bytes[0..2], .little), host_proj[0] }); - if (out_off != 0 or xs_off != 0 or st_off != 0 or proj_off != 0) return 1; - return 0; + if (out_off != 0 or xs_off != 0 or st_off != 0 or proj_off != 0) return null; + // Eight K slices, left unreduced for the write-back. + const sk: usize = 8; + const n: usize = 2560; + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * n * 4); + defer part_b.free(); + try qmm.partials(driver, stream.*, gout.ptr, gxs.ptr, w_b.ptr, s_b.ptr, b_b.ptr, branch.ptr, part_b.ptr, 1, n, value_dim, sk); + const raw = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(raw); + try stream.synchronize(); + try part_b.download(0, raw); + const slices = try gpa.alloc(f32, sk * n); + errdefer gpa.free(slices); + for (slices, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); + return slices; } test "warp sum is the sum of the lanes" { diff --git a/zig/src/families/flashnext/cuda_hc.zig b/zig/src/families/flashnext/cuda_hc.zig index 76a34a344e..65b661904b 100644 --- a/zig/src/families/flashnext/cuda_hc.zig +++ b/zig/src/families/flashnext/cuda_hc.zig @@ -247,6 +247,53 @@ pub fn mixedSteps(gpu: u16, host: u16) u32 { return bf16Steps(gpu, host); } +/// Mode 4's branch: the K slices added in slice order, then one bf16 round. +pub fn sumSlices(part: []const f32, sk: usize, n: usize, out: []u16) !void { + if (sk == 0 or n == 0 or part.len != sk * n or out.len != n) return error.UnexpectedTensor; + for (0..n) |i| { + var acc = part[i]; + for (1..sk) |k| acc += part[k * n + i]; + out[i] = toBf16(acc); + } +} + +/// Mode 4's hidden update: one bf16 rounding of hidden + branch * inject, the write-back's fused multiply-add. +pub fn applyBranch(h: []const u16, branch: []const u16, inj: []const u16, dims: usize, streams: usize, out: []u16) !void { + if (dims == 0 or h.len != streams * dims or branch.len != dims or inj.len != streams or out.len != h.len) return error.UnexpectedTensor; + for (0..streams) |s| { + const gate = promote(inj[s]); + for (0..dims) |i| { + out[s * dims + i] = toBf16(promote(h[s * dims + i]) + promote(branch[i]) * gate); + } + } +} + +test "slices add in order and round once" { + const part = [_]f32{ 1, 0, 0.5, 0 }; + var out: [2]u16 = undefined; + try sumSlices(&part, 2, 2, &out); + try std.testing.expectEqual(toBf16(1.5), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[1]); +} + +test "a branch of one times an inject gate of one adds one" { + const h = [_]u16{0}; + const branch = [_]u16{0x3f80}; + const inj = [_]u16{0x3f80}; + var out: [1]u16 = undefined; + try applyBranch(&h, &branch, &inj, 1, 1, &out); + try std.testing.expectEqual(@as(u16, 0x3f80), out[0]); +} + +test "the fused product is rounded once with the hidden value" { + const h = [_]u16{0x3b34}; + const branch = [_]u16{0xbe16}; + const inj = [_]u16{0x3c9a}; + var out: [1]u16 = undefined; + try applyBranch(&h, &branch, &inj, 1, 1, &out); + try std.testing.expectEqual(@as(u16, 0xb6f0), out[0]); +} + test "a zero activation mixes to a zero hidden value" { var act: [32]u16 = @splat(0); const xs = [_]f32{0}; diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index cad16d36b0..6d7eae96fb 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -8,6 +8,7 @@ pub const embed = @import("cuda_embed.zig"); pub const hc = @import("cuda_hc.zig"); pub const qmm = @import("cuda_qmm.zig"); pub const gdn = @import("cuda_gdn.zig"); +pub const mlp = @import("cuda_mlp.zig"); const p = aot.ptr; @@ -119,6 +120,7 @@ test { _ = hc; _ = qmm; _ = gdn; + _ = mlp; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_mlp.zig b/zig/src/families/flashnext/cuda_mlp.zig new file mode 100644 index 0000000000..bd7e875d72 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mlp.zig @@ -0,0 +1,255 @@ +//! Layer 0's second hyper-connection: the mixer branch written back, then the same read-out on the mlp weights. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const branch_sk: usize = 8; + +/// `slices` are the output face's unreduced K slices. `h` is the four streams, `inj` the first connection's gates. +pub fn connect(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, tri: Tri, h: u64, inj: u64, hidden: []const u16, slices: []const f32) !u32 { + if (hidden.len == 0 or hidden.len % 4 != 0) return error.UnexpectedTensor; + const streams: usize = 4; + const dims = hidden.len / streams; + if (slices.len != branch_sk * dims) return error.UnexpectedTensor; + const inj_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_bytes); + try stream.synchronize(); + var inj_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = inj, .len = inj_bytes.len }; + try inj_buf.download(0, inj_bytes); + const inj_u16 = try gpa.alloc(u16, streams); + defer gpa.free(inj_u16); + for (inj_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, inj_bytes[2 * i ..][0..2], .little); + const branch = try gpa.alloc(u16, dims); + defer gpa.free(branch); + try hc.sumSlices(slices, branch_sk, dims, branch); + const host_h = try gpa.alloc(u16, hidden.len); + defer gpa.free(host_h); + try hc.applyBranch(hidden, branch, inj_u16, dims, streams, host_h); + + var part = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(slices)); + defer part.free(); + const nc = dims / 256; + var pss_b = try cuda.DeviceBuffer.alloc(driver, nc * streams * 4); + defer pss_b.free(); + try tri.hcWriteback(h, h, pss_b.ptr, part.ptr, "*fp32", inj, h, "*bf16", h, "*bf16", dims, 1, dims, streams, 4, 1, 1, branch_sk); + const got_h = try gpa.alloc(u8, hidden.len * 2); + defer gpa.free(got_h); + const pss_bytes = try gpa.alloc(u8, nc * streams * 4); + defer gpa.free(pss_bytes); + try stream.synchronize(); + var h_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = h, .len = got_h.len }; + try h_buf.download(0, got_h); + try pss_b.download(0, pss_bytes); + var h_off: usize = 0; + var h_steps: u32 = 0; + for (host_h, 0..) |want, i| { + const got = std.mem.readInt(u16, got_h[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) h_off += 1; + if (dist > h_steps) h_steps = dist; + } + const want_pss = try gpa.alloc(f32, nc * streams); + defer gpa.free(want_pss); + try hc.partialSums(host_h, dims, streams, want_pss); + var pss_ulp: u32 = 0; + for (want_pss, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > pss_ulp) pss_ulp = dist; + } + std.debug.print("mlp writeback off {d} max_steps {d} pss_ulp {d} head {x:0>4} host {x:0>4} branch {x:0>4}\n", .{ h_off, h_steps, pss_ulp, std.mem.readInt(u16, got_h[0..2], .little), host_h[0], branch[0] }); + if (h_steps != 0 or pss_ulp > 4) return 1; + + const prefix = "language_model.model.layers.0.mlp_hyper_connection."; + var name: [180]u8 = undefined; + const down_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.weight", .{prefix}), .u32); + const down_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.scales", .{prefix}), .bf16); + const down_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.biases", .{prefix}), .bf16); + const inj_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.weight", .{prefix}), .u32); + const inj_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.scales", .{prefix}), .bf16); + const inj_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}block_inject_weight.biases", .{prefix}), .bf16); + const gamma = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}hc_norm.weight", .{prefix}), .bf16); + const up_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.weight", .{prefix}), .u32); + const up_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.scales", .{prefix}), .bf16); + const up_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_up.biases", .{prefix}), .bf16); + const n = down_w.dim(0) + inj_w.dim(0); + const k8 = down_w.dim(1); + const kg = k8 / 4; + if (n != 324 or k8 != 1280 or gamma.dim(0) != streams * dims) return error.UnexpectedTensor; + const stacked_w = try gpa.alloc(u8, down_w.bytes.len + inj_w.bytes.len); + defer gpa.free(stacked_w); + @memcpy(stacked_w[0..down_w.bytes.len], down_w.bytes); + @memcpy(stacked_w[down_w.bytes.len..], inj_w.bytes); + const stacked_s = try gpa.alloc(u8, down_s.bytes.len + inj_s.bytes.len); + defer gpa.free(stacked_s); + @memcpy(stacked_s[0..down_s.bytes.len], down_s.bytes); + @memcpy(stacked_s[down_s.bytes.len..], inj_s.bytes); + const stacked_b = try gpa.alloc(u8, down_b.bytes.len + inj_b.bytes.len); + defer gpa.free(stacked_b); + @memcpy(stacked_b[0..down_b.bytes.len], down_b.bytes); + @memcpy(stacked_b[down_b.bytes.len..], inj_b.bytes); + const tiled = try hc.tileWords(gpa, stacked_w, n, k8); + defer gpa.free(tiled); + const scales_t = try hc.transposeBf16(gpa, stacked_s, n, kg); + defer gpa.free(scales_t); + const biases_t = try hc.transposeBf16(gpa, stacked_b, n, kg); + defer gpa.free(biases_t); + const scale = try gpa.alloc(f32, gamma.dim(0)); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = @bitCast(@as(u32, std.mem.readInt(u16, gamma.bytes[2 * i ..][0..2], .little)) << 16); + var w_b = try cuda.DeviceBuffer.fromHost(driver, tiled); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales_t); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases_t); + defer b_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, streams * dims * 2); + defer norm_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, n * 2); + defer out_b.free(); + const down_sk: usize = 32; + var down_part = try cuda.DeviceBuffer.alloc(driver, down_sk * n * 4); + defer down_part.free(); + try tri.hcDown(h, pss_b.ptr, scale_b.ptr, norm_b.ptr, w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, down_part.ptr, 1e-6, 1, n, k8 * 8, dims, nc, streams, down_sk); + const got_norm = try gpa.alloc(u8, streams * dims * 2); + defer gpa.free(got_norm); + try stream.synchronize(); + try norm_b.download(0, got_norm); + const pss_f = try gpa.alloc(f32, nc * streams); + defer gpa.free(pss_f); + for (pss_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, pss_bytes[4 * i ..][0..4], .little)); + const gpu_h = try gpa.alloc(u16, hidden.len); + defer gpa.free(gpu_h); + for (gpu_h, 0..) |*o, i| o.* = std.mem.readInt(u16, got_h[2 * i ..][0..2], .little); + const want_norm = try gpa.alloc(u16, hidden.len); + defer gpa.free(want_norm); + try hc.normed(gpu_h, pss_f, scale, dims, streams, 1e-6, want_norm); + var norm_off: usize = 0; + var norm_steps: u32 = 0; + var norm_at: usize = 0; + for (want_norm, 0..) |want, i| { + const got = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + if (got != want) { + const dist = hc.mixedSteps(got, want); + if (norm_off == 0 or dist > norm_steps) { + norm_steps = dist; + norm_at = i; + } + norm_off += 1; + } + } + var part4: [4]u8 = undefined; + try down_part.download(0, &part4); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, &part4, .little)); + const gpu_norm = try gpa.alloc(u16, 320); + defer gpa.free(gpu_norm); + for (gpu_norm, 0..) |*o, i| o.* = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + const host_part = try hc.firstPartial(gpu_norm, stacked_w[0 .. k8 * 4], stacked_s[0 .. kg * 2], stacked_b[0 .. kg * 2], 10); + const col_ulp = hc.ulps(gpu_part, host_part); + std.debug.print("mlp hcdown normed off {d} max_steps {d} at {d} col0 gpu {d} host {d} ulp {d}\n", .{ norm_off, norm_steps, norm_at, gpu_part, host_part, col_ulp }); + if (norm_steps > 1 or col_ulp > 8) return 1; + + const low: usize = 320; + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + var inj_out = try cuda.DeviceBuffer.alloc(driver, streams * 2); + defer inj_out.free(); + try tri.hcReduceAct(down_part.ptr, act_b.ptr, xs_b.ptr, inj_out.ptr, down_sk, 1, streams, low, n, 1); + const part_bytes = try gpa.alloc(u8, down_sk * n * 4); + defer gpa.free(part_bytes); + try stream.synchronize(); + try down_part.download(0, part_bytes); + const part_f = try gpa.alloc(f32, down_sk * n); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_bytes[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const inj_want = try gpa.alloc(u16, streams); + defer gpa.free(inj_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try hc.reduceAct(part_f, down_sk, n, low, streams, act_want, inj_want, xs_want); + const act_bytes = try gpa.alloc(u8, low * 2); + defer gpa.free(act_bytes); + const inj_out_bytes = try gpa.alloc(u8, streams * 2); + defer gpa.free(inj_out_bytes); + const xs_bytes = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_bytes); + try act_b.download(0, act_bytes); + try inj_out.download(0, inj_out_bytes); + try xs_b.download(0, xs_bytes); + var act_off: usize = 0; + for (act_want, 0..) |want, i| { + if (std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little) != want) act_off += 1; + } + var inj_off: usize = 0; + for (inj_want, 0..) |want, i| { + if (std.mem.readInt(u16, inj_out_bytes[2 * i ..][0..2], .little) != want) inj_off += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + std.debug.print("mlp reduce act mismatch {d} inject mismatch {d} xs max_ulp {d}\n", .{ act_off, inj_off, xs_ulp }); + if (act_off != 0 or inj_off != 0 or xs_ulp > 4) return 1; + + const up_n = up_w.dim(0); + const up_k8 = up_w.dim(1); + if (up_n != streams * dims or up_k8 != 40) return error.UnexpectedTensor; + const up_tiled = try hc.tileWords(gpa, up_w.bytes, up_n, up_k8); + defer gpa.free(up_tiled); + const up_scales = try hc.transposeBf16(gpa, up_s.bytes, up_n, up_k8 / 4); + defer gpa.free(up_scales); + const up_biases = try hc.transposeBf16(gpa, up_b.bytes, up_n, up_k8 / 4); + defer gpa.free(up_biases); + var up_wb = try cuda.DeviceBuffer.fromHost(driver, up_tiled); + defer up_wb.free(); + var up_sb = try cuda.DeviceBuffer.fromHost(driver, up_scales); + defer up_sb.free(); + var up_bb = try cuda.DeviceBuffer.fromHost(driver, up_biases); + defer up_bb.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, up_wb.ptr, up_sb.ptr, up_bb.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, up_n, up_k8 * 8, dims, streams); + const mixed_bytes = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_bytes); + try stream.synchronize(); + try mixed_b.download(0, mixed_bytes); + const act_u16 = try gpa.alloc(u16, low); + defer gpa.free(act_u16); + for (act_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, act_bytes[2 * i ..][0..2], .little); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const norm_u16 = try gpa.alloc(u16, hidden.len); + defer gpa.free(norm_u16); + for (norm_u16, 0..) |*o, i| o.* = std.mem.readInt(u16, got_norm[2 * i ..][0..2], .little); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try hc.mix(act_u16, xs_f, norm_u16, up_w.bytes, up_s.bytes, up_b.bytes, up_n, up_k8, dims, streams, host_mixed); + var mixed_off: usize = 0; + var max_steps: u32 = 0; + var mixed_at: usize = 0; + for (host_mixed, 0..) |want, i| { + const got = std.mem.readInt(u16, mixed_bytes[2 * i ..][0..2], .little); + const dist = hc.mixedSteps(got, want); + if (dist != 0) mixed_off += 1; + if (dist > max_steps) { + max_steps = dist; + mixed_at = i; + } + } + std.debug.print("mlp upmix dims {d} off {d} max_steps {d} at {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, mixed_at, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); + if (max_steps > 2) return 1; + return 0; +} diff --git a/zig/src/families/flashnext/cuda_qmm.zig b/zig/src/families/flashnext/cuda_qmm.zig index 7c1c7c3dc1..50df60a80f 100644 --- a/zig/src/families/flashnext/cuda_qmm.zig +++ b/zig/src/families/flashnext/cuda_qmm.zig @@ -149,6 +149,37 @@ pub fn matmul(d: *const cuda.Driver, stream: cuda.Stream, x: u64, xs: u64, weigh try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows_t * cols_t) }, .block = .{ .x = threads }, .shared = smem }, stream, &args); } +/// The same kernel as `matmul` with one block per K slice. `part` receives the unreduced [sk, m, n] fp32 slices. +pub fn partials(d: *const cuda.Driver, stream: cuda.Stream, x: u64, xs: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, m: usize, n: usize, k: usize, sk: usize) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (m == 0 or m > 16 or sk == 0 or k == 0 or k % group != 0 or (k / group) % sk != 0) return error.UnexpectedTensor; + var module = try cuda.Module.load(d, cuda.kernels.qmm); + defer module.unload(); + const f = try module.function(symbol); + try f.allowDynamicShared(smem); + const bm: usize = 16; + const bn: usize = 64; + const rows_t = (m + bm - 1) / bm; + const cols_t = (n + bn - 1) / bn; + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (n + 127) / 128 * 128; + const ldx = k; + var args: cuda.Args = .{}; + args.add(x); + args.add(xs); + args.add(weight); + args.add(scales); + args.add(biases); + args.add(out); + args.add(part); + for ([_]usize{ m, n, k, sk, npad, ldx, sweep }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(f, .{ + .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, + .block = .{ .x = threads }, + .shared = smem, + }, stream, &args); +} + test "one group's first nibble pair lands in the lane word" { var words: [16]u8 = @splat(0); std.mem.writeInt(u32, words[0..4], 0x21, .little); From 58fed6c7aad11b31071d28a8247d6f3d43778d73 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 00:33:00 -0400 Subject: [PATCH 012/115] Layer 0's experts match the host check on one token, the top 10 and the shared expert exact, router logits within 29184 ulps, two weights 1 step off, the activations exact, and every down value the same bf16 with fp32 within 16384 ulps. --- zig/kernels/cuda/experts.cu | 5 + zig/kernels/cuda/experts_pack.cu | 4 +- .../families/flashnext/cuda_embed_main.zig | 4 +- zig/src/families/flashnext/cuda_launch.zig | 22 + zig/src/families/flashnext/cuda_moe.zig | 512 ++++++++++++++++++ 5 files changed, 545 insertions(+), 2 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_moe.zig diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index edd83834c1..9e46b57a19 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -280,3 +280,8 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); TF_EXPERT(1) TF_EXPERT(0) +// Groups of 32: SwiGLU gate and up together (epilogue 2), and the down projection to fp32. +template __global__ void tf_experts::expert_kernel<32, 2, 2, 2, 4>(const __nv_bfloat16*, \ + int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); +template __global__ void tf_experts::expert_kernel<32, 1, 0, 2, 4>(const __nv_bfloat16*, \ + int, int, const uint4*, int, int, const int*, const int*, const int*, void*, int, float); diff --git a/zig/kernels/cuda/experts_pack.cu b/zig/kernels/cuda/experts_pack.cu index 6aa7413793..327517f5b4 100644 --- a/zig/kernels/cuda/experts_pack.cu +++ b/zig/kernels/cuda/experts_pack.cu @@ -37,6 +37,8 @@ pack_kernel(const uint32_t* __restrict__ words, const uint16_t* __restrict__ sca } // namespace tf_experts_pack -// Groups of 64 inputs. +// Groups of 64 inputs, then groups of 32. template __global__ void tf_experts_pack::pack_kernel<2>(const uint32_t*, const uint16_t*, const uint16_t*, uint32_t*, int, int, int, int); +template __global__ void tf_experts_pack::pack_kernel<1>(const uint32_t*, const uint16_t*, const uint16_t*, + uint32_t*, int, int, int, int); diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index 138a5974aa..894c5bb55c 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -382,5 +382,7 @@ fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, m if (max_steps != 0) return max_steps; const slices = try flash.gdn.memory(gpa, driver, stream, mapped, out_b, got_bytes) orelse return 1; defer gpa.free(slices); - return flash.mlp.connect(flash.Tri, gpa, driver, stream, mapped, tri, h, inj, hidden, slices); + const wrote = try flash.mlp.connect(flash.Tri, gpa, driver, stream, mapped, tri, h, inj, hidden, slices); + if (wrote != 0) return wrote; + return flash.moe.experts(flash.Tri, gpa, driver, stream, mapped, tri, mixed.ptr); } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 6d7eae96fb..d05f6c29ca 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -9,6 +9,7 @@ pub const hc = @import("cuda_hc.zig"); pub const qmm = @import("cuda_qmm.zig"); pub const gdn = @import("cuda_gdn.zig"); pub const mlp = @import("cuda_mlp.zig"); +pub const moe = @import("cuda_moe.zig"); const p = aot.ptr; @@ -113,6 +114,26 @@ pub const Tri = struct { int("RS", rs), }, &.{ ci("D", dims), ci("S", streams), ci("MODE", mode), ci("TOPK", topk), ci("SLOTS", slots), ci("BLOCK", 256), ci("WORLD", world) }); } + + /// moe.router for one decode row. M folds to a constexpr. 513 rows: the experts, then the shared gate. + pub fn router(t: Tri, x: u64, w: u64, out: u64, x_stride: usize) !void { + try t.run("_router", .{ 1, cdiv(513, 32), 1 }, &.{ + p("X", "*bf16", x), + p("W", "*bf16", w), + p("OUT", "*fp32", out), + int("M", 1), + int("x_stride", x_stride), + }, &.{ ci("D", 2560), ci("NE", 513), ci("BM", 16), ci("BLOCK_E", 32), ci("BK", 256) }); + } + + /// moe._topk_rows: one program, the row's 10 experts and the shared expert. + pub fn topkRows(t: Tri, logits: u64, pick: u64, wts: u64) !void { + try t.run("_topk_rows", .{ 1, 1, 1 }, &.{ + p("L", "*fp32", logits), + p("PICK", "*i32", pick), + p("WTS", "*fp32", wts), + }, &.{ ci("NE", 512), ci("NL", 513), ci("TOPK", 10), ci("SLOTS", 11), ci("BLOCK", 1024), ci("SLOTP", 16) }); + } }; test { @@ -121,6 +142,7 @@ test { _ = qmm; _ = gdn; _ = mlp; + _ = moe; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_moe.zig b/zig/src/families/flashnext/cuda_moe.zig new file mode 100644 index 0000000000..5a0030dfc0 --- /dev/null +++ b/zig/src/families/flashnext/cuda_moe.zig @@ -0,0 +1,512 @@ +//! Layer 0's experts on one mixed row: router, top 10 plus the shared expert, then each slot's gate, up and down. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const count: usize = 512; +const top_k: usize = 10; +const slots: usize = top_k + 1; +const dims: usize = 2560; +const width: usize = 640; +const routed_rows: usize = count + 1; +const tile: usize = 16; +const block_bytes: usize = 160 * 4; + +const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; +const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + +const Proj = struct { + routed_w: []const u8, + routed_s: []const u8, + routed_b: []const u8, + shared_w: []const u8, + shared_s: []const u8, + shared_b: []const u8, + n: usize, + k: usize, +}; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn roundTrip(v: f32) f32 { + return promote(toBf16(v)); +} + +/// The expert kernel's group sum: each lane adds its eight inputs, then the quad adds as (0+1)+(2+3). +fn groupSum(x: []const u16, g: usize) f32 { + var pair: [2]f32 = .{ 0, 0 }; + for (0..4) |t| { + var s: f32 = 0; + const base = g * 32 + t * 8; + for (0..8) |j| s += promote(x[base + j]); + if (t % 2 == 0) pair[t / 2] = s else pair[t / 2] += s; + } + return pair[0] + pair[1]; +} + +/// One expert matrix: groups in order, `p * scale + xs * bias`, fp32. `xs` is `groupSum` of `x`. +fn project(x: []const u16, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k: usize, out: []f32) void { + const kg = k / 32; + const k8 = k / 8; + var xs: [80]f32 = undefined; + for (0..kg) |g| xs[g] = groupSum(x, g); + for (0..n) |col| { + var acc: f32 = 0; + const wb = col * k8 * 4; + const sb = col * kg * 2; + for (0..kg) |g| { + var p: f32 = 0; + for (0..4) |i| { + const word = std.mem.readInt(u32, words[wb + (g * 4 + i) * 4 ..][0..4], .little); + const base = g * 32 + i * 8; + inline for (0..8) |j| { + const q: f32 = @floatFromInt((word >> (j * 4)) & 0xF); + p += promote(x[base + j]) * q; + } + } + const s = promote(std.mem.readInt(u16, scales[sb + g * 2 ..][0..2], .little)); + const b = promote(std.mem.readInt(u16, biases[sb + g * 2 ..][0..2], .little)); + acc = @mulAdd(f32, xs[g], b, @mulAdd(f32, p, s, acc)); + } + out[col] = acc; + } +} + +/// bf16(SiLU(bf16(gate)) * bf16(up)), the SwiGLU epilogue with no clip. +fn swiglu(gate: f32, up: f32) u16 { + const gv = roundTrip(gate); + const uv = roundTrip(up); + const silu = roundTrip(gv / (1 + @exp(-gv))); + return toBf16(silu * uv); +} + +/// Router rows in K tiles of 256, the same tile the captured program uses, added in tile order. +fn routerDots(x: []const u16, w: []const u16, out: []f32) void { + const k = x.len; + for (out, 0..) |*o, e| { + var acc: f32 = 0; + const row = w[e * k ..][0..k]; + var t: usize = 0; + while (t < k) : (t += 256) { + var part: f32 = 0; + for (0..256) |i| part += promote(x[t + i]) * promote(row[t + i]); + acc += part; + } + o.* = acc; + } +} + +/// Largest `k` of the first `n` logits (lower id on a tie), weights exp(l - l0) / sum rounded to bf16, then the shared gate. +fn choose(logits: []const f32, n: usize, k: usize, picks: []i32, wts: []f32) void { + var scratch: [count]f32 = undefined; + for (0..n) |i| scratch[i] = logits[i]; + var top = scratch[0]; + for (scratch[1..n]) |v| top = @max(top, v); + var total: f32 = 0; + var exs: [16]f32 = undefined; + for (0..k) |slot| { + var m = scratch[0]; + var idx: usize = 0; + for (scratch[1..n], 1..) |v, i| if (v > m) { + m = v; + idx = i; + }; + const ex = @exp(m - top); + picks[slot] = @intCast(idx); + exs[slot] = ex; + total += ex; + scratch[idx] = -std.math.inf(f32); + } + for (0..k) |slot| wts[slot] = roundTrip(exs[slot] / total); + const shared = roundTrip(logits[n]); + picks[k] = @intCast(n); + wts[k] = roundTrip(1 / (1 + @exp(-shared))); +} + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn loadProj(mapped: *const embed.Mapped, gpa: std.mem.Allocator, name: []const u8, n: usize, k: usize) !Proj { + const prefix = "language_model.model.layers.0.mlp."; + var buf: [180]u8 = undefined; + const k8 = k / 8; + const kg = k / 32; + const rw = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.weight", .{ prefix, name }), .u32); + const rs = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.scales", .{ prefix, name }), .bf16); + const rb = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}switch_mlp.{s}.biases", .{ prefix, name }), .bf16); + const sw = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.weight", .{ prefix, name }), .u32); + const ss = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.scales", .{ prefix, name }), .bf16); + const sb = try mapped.lookup(gpa, try std.fmt.bufPrint(&buf, "{s}shared_expert.{s}.biases", .{ prefix, name }), .bf16); + if (rw.dim(0) != count or rw.dim(1) != n or rw.dim(2) != k8) return error.UnexpectedTensor; + if (rs.dim(0) != count or rs.dim(1) != n or rs.dim(2) != kg) return error.UnexpectedTensor; + if (rb.dim(0) != count or rb.dim(1) != n or rb.dim(2) != kg) return error.UnexpectedTensor; + if (sw.dim(0) != n or sw.dim(1) != k8 or ss.dim(0) != n or ss.dim(1) != kg) return error.UnexpectedTensor; + if (sb.dim(0) != n or sb.dim(1) != kg) return error.UnexpectedTensor; + return .{ .routed_w = rw.bytes, .routed_s = rs.bytes, .routed_b = rb.bytes, .shared_w = sw.bytes, .shared_s = ss.bytes, .shared_b = sb.bytes, .n = n, .k = k }; +} + +fn copyExpert(dst_w: []u8, dst_s: []u8, dst_b: []u8, p: Proj, id: i32, slot: usize) void { + const k8 = p.k / 8; + const kg = p.k / 32; + const wb = p.n * k8 * 4; + const sb = p.n * kg * 2; + const routed = id != count; + const wsrc = if (routed) p.routed_w[@as(usize, @intCast(id)) * wb ..][0..wb] else p.shared_w; + const ssrc = if (routed) p.routed_s[@as(usize, @intCast(id)) * sb ..][0..sb] else p.shared_s; + const bsrc = if (routed) p.routed_b[@as(usize, @intCast(id)) * sb ..][0..sb] else p.shared_b; + @memcpy(dst_w[slot * wb ..][0..wb], wsrc); + @memcpy(dst_s[slot * sb ..][0..sb], ssrc); + @memcpy(dst_b[slot * sb ..][0..sb], bsrc); +} + +fn packDownload(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) ![]u8 { + const kg = k / 32; + const nb = n / 32; + const out_len = e * nb * kg * block_bytes; + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, out_len); + defer out_b.free(); + var a: cuda.Args = .{}; + a.add(w_b.ptr); + a.add(s_b.ptr); + a.add(b_b.ptr); + a.add(out_b.ptr); + a.add(cint(n)); + a.add(cint(k / 8)); + a.add(cint(kg)); + a.add(cint(nb)); + try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream, &a); + const out = try gpa.alloc(u8, out_len); + errdefer gpa.free(out); + try stream.synchronize(); + try out_b.download(0, out); + return out; +} + +fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { + if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; + const nblk = gate.len / block_bytes; + const out = try gpa.alloc(u8, gate.len + up.len); + for (0..nblk) |b| { + @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); + @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); + } + return out; +} + +fn maxUlp(got: []const f32, want: []const f32) struct { off: usize, ulp: u32 } { + var off: usize = 0; + var ulp: u32 = 0; + for (got, want) |g, w| { + const dist = hc.ulps(g, w); + if (dist != 0) off += 1; + if (dist > ulp) ulp = dist; + } + return .{ .off = off, .ulp = ulp }; +} + +/// The 512 routed experts and the shared expert, on the mixed row already at `mixed`. +pub fn experts(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, tri: Tri, mixed: u64) !u32 { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const prefix = "language_model.model.layers.0.mlp."; + var name: [180]u8 = undefined; + const gate = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}gate.weight", .{prefix}), .bf16); + const sg_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.weight", .{prefix}), .u32); + const sg_s = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.scales", .{prefix}), .bf16); + const sg_b = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}shared_expert_gate.biases", .{prefix}), .bf16); + if (gate.dim(0) != count or gate.dim(1) != dims) return error.UnexpectedTensor; + if (sg_w.dim(0) != 1 or sg_w.dim(1) != dims / 8) return error.UnexpectedTensor; + + const row_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(row_raw); + try stream.synchronize(); + var mixed_buf: cuda.DeviceBuffer = .{ .d = driver, .ptr = mixed, .len = row_raw.len }; + try mixed_buf.download(0, row_raw); + const row = try gpa.alloc(u16, dims); + defer gpa.free(row); + for (row, 0..) |*o, i| o.* = std.mem.readInt(u16, row_raw[2 * i ..][0..2], .little); + + const router = try gpa.alloc(u16, routed_rows * dims); + defer gpa.free(router); + @memcpy(std.mem.sliceAsBytes(router[0 .. count * dims]), gate.bytes); + try embed.dequant(sg_w.bytes, sg_s.bytes, sg_b.bytes, dims, 0, router[count * dims ..][0..dims]); + const host_l = try gpa.alloc(f32, routed_rows); + defer gpa.free(host_l); + routerDots(row, router, host_l); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); + defer w_b.free(); + var logits_b = try cuda.DeviceBuffer.alloc(driver, routed_rows * 4); + defer logits_b.free(); + var pick_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer pick_b.free(); + var wts_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer wts_b.free(); + try tri.router(mixed, w_b.ptr, logits_b.ptr, dims); + try tri.topkRows(logits_b.ptr, pick_b.ptr, wts_b.ptr); + + const logit_raw = try gpa.alloc(u8, routed_rows * 4); + defer gpa.free(logit_raw); + const pick_raw = try gpa.alloc(u8, slots * 4); + defer gpa.free(pick_raw); + const wts_raw = try gpa.alloc(u8, slots * 4); + defer gpa.free(wts_raw); + try stream.synchronize(); + try logits_b.download(0, logit_raw); + try pick_b.download(0, pick_raw); + try wts_b.download(0, wts_raw); + const gpu_l = try gpa.alloc(f32, routed_rows); + defer gpa.free(gpu_l); + for (gpu_l, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, logit_raw[4 * i ..][0..4], .little)); + var gpu_pick: [slots]i32 = undefined; + var gpu_w: [slots]f32 = undefined; + for (0..slots) |i| { + gpu_pick[i] = std.mem.readInt(i32, pick_raw[4 * i ..][0..4], .little); + gpu_w[i] = @bitCast(std.mem.readInt(u32, wts_raw[4 * i ..][0..4], .little)); + } + const route = maxUlp(gpu_l, host_l); + var from_gpu: [slots]i32 = undefined; + var w_gpu: [slots]f32 = undefined; + choose(gpu_l, count, top_k, &from_gpu, &w_gpu); + var from_host: [slots]i32 = undefined; + var w_host: [slots]f32 = undefined; + choose(host_l, count, top_k, &from_host, &w_host); + var pick_off: usize = 0; + var host_pick_off: usize = 0; + for (gpu_pick, from_gpu, from_host) |g, a, b| { + if (g != a) pick_off += 1; + if (g != b) host_pick_off += 1; + } + const wts_same = maxUlp(&gpu_w, &w_gpu); + var host_wts_off: usize = 0; + var host_wts_steps: u32 = 0; + for (gpu_w, w_host) |g, h| { + const steps = hc.mixedSteps(toBf16(g), toBf16(h)); + if (steps != 0) host_wts_off += 1; + if (steps > host_wts_steps) host_wts_steps = steps; + } + std.debug.print("router off {d} max_ulp {d} head {x:0>8} host {x:0>8}\n", .{ route.off, route.ulp, @as(u32, @bitCast(gpu_l[0])), @as(u32, @bitCast(host_l[0])) }); + std.debug.print("topk pick_off {d} host_pick_off {d} wts_ulp {d} host_wts_off {d} host_wts_steps {d} picks", .{ pick_off, host_pick_off, wts_same.ulp, host_wts_off, host_wts_steps }); + for (gpu_pick) |id| std.debug.print(" {d}", .{id}); + std.debug.print("\n", .{}); + + var packed_ids: [slots]i32 = undefined; + var npack: usize = 0; + var local: [slots]i32 = undefined; + for (gpu_pick, 0..) |id, s| { + if (id < 0 or id > count) return error.UnexpectedTensor; + var found: ?usize = null; + for (packed_ids[0..npack], 0..) |prev, i| if (prev == id) { + found = i; + }; + if (found) |i| { + local[s] = @intCast(i); + } else { + packed_ids[npack] = id; + local[s] = @intCast(npack); + npack += 1; + } + } + + const gate_p = try loadProj(mapped, gpa, "gate_proj", width, dims); + const up_p = try loadProj(mapped, gpa, "up_proj", width, dims); + const down_p = try loadProj(mapped, gpa, "down_proj", dims, width); + const gathered_w = try gpa.alloc(u8, npack * width * (dims / 8) * 4); + defer gpa.free(gathered_w); + const gathered_s = try gpa.alloc(u8, npack * width * (dims / 32) * 2); + defer gpa.free(gathered_s); + const gathered_b = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(gathered_b); + const up_w = try gpa.alloc(u8, gathered_w.len); + defer gpa.free(up_w); + const up_s = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(up_s); + const up_b = try gpa.alloc(u8, gathered_s.len); + defer gpa.free(up_b); + const down_w = try gpa.alloc(u8, npack * dims * (width / 8) * 4); + defer gpa.free(down_w); + const down_s = try gpa.alloc(u8, npack * dims * (width / 32) * 2); + defer gpa.free(down_s); + const down_b = try gpa.alloc(u8, down_s.len); + defer gpa.free(down_b); + for (packed_ids[0..npack], 0..) |id, slot| { + copyExpert(gathered_w, gathered_s, gathered_b, gate_p, id, slot); + copyExpert(up_w, up_s, up_b, up_p, id, slot); + copyExpert(down_w, down_s, down_b, down_p, id, slot); + } + + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + var exp_mod = try cuda.Module.load(driver, cuda.kernels.experts); + defer exp_mod.unload(); + const pack_fn = try pack_mod.function(pack_symbol); + const packed_g = try packDownload(gpa, driver, stream.*, pack_fn, gathered_w, gathered_s, gathered_b, npack, width, dims); + defer gpa.free(packed_g); + const packed_u = try packDownload(gpa, driver, stream.*, pack_fn, up_w, up_s, up_b, npack, width, dims); + defer gpa.free(packed_u); + const packed_d = try packDownload(gpa, driver, stream.*, pack_fn, down_w, down_s, down_b, npack, dims, width); + defer gpa.free(packed_d); + const stacked = try stackHalves(gpa, packed_g, packed_u); + defer gpa.free(stacked); + + var up_wb = try cuda.DeviceBuffer.fromHost(driver, stacked); + defer up_wb.free(); + var down_wb = try cuda.DeviceBuffer.fromHost(driver, packed_d); + defer down_wb.free(); + var local_pick: [slots]i32 = undefined; + for (0..slots) |s| local_pick[s] = local[s]; + var picks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&local_pick)); + defer picks_b.free(); + var members_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer members_b.free(); + var items_b = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); + defer items_b.free(); + var counts_b = try cuda.DeviceBuffer.alloc(driver, 8); + defer counts_b.free(); + var act_b = try cuda.DeviceBuffer.alloc(driver, slots * width * 2); + defer act_b.free(); + var y_b = try cuda.DeviceBuffer.alloc(driver, slots * dims * 4); + defer y_b.free(); + const plan_fn = try exp_mod.function(plan_symbol); + var plan_a: cuda.Args = .{}; + plan_a.add(picks_b.ptr); + plan_a.add(cint(slots)); + plan_a.add(cint(npack)); + plan_a.add(cint(tile)); + plan_a.add(members_b.ptr); + plan_a.add(items_b.ptr); + plan_a.add(counts_b.ptr); + try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); + + const up_fn = try exp_mod.function(up_symbol); + const down_fn = try exp_mod.function(down_symbol); + try launchExpert(up_fn, stream.*, mixed, dims, slots, up_wb.ptr, dims / 32, width / 32, items_b.ptr, counts_b.ptr, members_b.ptr, act_b.ptr, width, slots * (width / 32)); + try launchExpert(down_fn, stream.*, act_b.ptr, width, 0, down_wb.ptr, width / 32, dims / 32, items_b.ptr, counts_b.ptr, members_b.ptr, y_b.ptr, dims, slots * (dims / 32)); + + const act_raw = try gpa.alloc(u8, slots * width * 2); + defer gpa.free(act_raw); + const y_raw = try gpa.alloc(u8, slots * dims * 4); + defer gpa.free(y_raw); + try stream.synchronize(); + try act_b.download(0, act_raw); + try y_b.download(0, y_raw); + + var act_off: usize = 0; + var act_steps: u32 = 0; + var y_off: usize = 0; + var y_ulp: u32 = 0; + var y_steps: u32 = 0; + var head_g: u32 = 0; + var head_h: u32 = 0; + const gate_col = try gpa.alloc(f32, width); + defer gpa.free(gate_col); + const up_col = try gpa.alloc(f32, width); + defer gpa.free(up_col); + const act_row = try gpa.alloc(u16, width); + defer gpa.free(act_row); + const down_col = try gpa.alloc(f32, dims); + defer gpa.free(down_col); + const gw = width * (dims / 8) * 4; + const gs = width * (dims / 32) * 2; + const dw = dims * (width / 8) * 4; + const ds = dims * (width / 32) * 2; + for (0..slots) |s| { + const src: usize = @intCast(local[s]); + project(row, gathered_w[src * gw ..][0..gw], gathered_s[src * gs ..][0..gs], gathered_b[src * gs ..][0..gs], width, dims, gate_col); + project(row, up_w[src * gw ..][0..gw], up_s[src * gs ..][0..gs], up_b[src * gs ..][0..gs], width, dims, up_col); + for (act_row, gate_col, up_col) |*o, gv, uv| o.* = swiglu(gv, uv); + for (0..width) |c| { + const got = std.mem.readInt(u16, act_raw[(s * width + c) * 2 ..][0..2], .little); + const dist = hc.mixedSteps(got, act_row[c]); + if (dist != 0) act_off += 1; + if (dist > act_steps) act_steps = dist; + } + project(act_row, down_w[src * dw ..][0..dw], down_s[src * ds ..][0..ds], down_b[src * ds ..][0..ds], dims, width, down_col); + for (0..dims) |c| { + const got: f32 = @bitCast(std.mem.readInt(u32, y_raw[(s * dims + c) * 4 ..][0..4], .little)); + const dist = hc.ulps(got, down_col[c]); + const steps = hc.mixedSteps(toBf16(got), toBf16(down_col[c])); + if (dist != 0) y_off += 1; + if (dist > y_ulp) y_ulp = dist; + if (steps > y_steps) y_steps = steps; + if (s == 0 and c == 0) { + head_g = @bitCast(got); + head_h = @bitCast(down_col[c]); + } + } + } + std.debug.print("experts slots {d} act_off {d} act_steps {d} y_off {d} y_ulp {d} y_steps {d} head {x:0>8} host {x:0>8}\n", .{ slots, act_off, act_steps, y_off, y_ulp, y_steps, head_g, head_h }); + if (pick_off != 0 or host_pick_off != 0 or wts_same.ulp != 0 or host_wts_steps > 1 or act_steps != 0 or y_steps != 0) return 1; + return 0; +} + +fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(cint(x_stride)); + a.add(cint(slot_n)); + a.add(w); + a.add(cint(kg)); + a.add(cint(nb)); + a.add(items); + a.add(counts); + a.add(members); + a.add(out); + a.add(cint(n)); + a.add(@as(f32, 0)); + const grid = (units + 3) / 4; + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(grid) }, .block = .{ .x = 128 } }, stream, &a); +} + +test "a tied logit keeps the lower expert id and the shared gate rounds to one half" { + const logits = [_]f32{ 1, 3, 3, 0 }; + var picks: [3]i32 = undefined; + var wts: [3]f32 = undefined; + choose(&logits, 3, 2, &picks, &wts); + try std.testing.expectEqual([3]i32{ 1, 2, 3 }, picks); + try std.testing.expectEqual(@as(f32, 0.5), wts[0]); + try std.testing.expectEqual(@as(f32, 0.5), wts[1]); + try std.testing.expectEqual(@as(f32, 0.5), wts[2]); +} + +test "each lane of the quad contributes to the group sum" { + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[8] = 0x3f80; + x[16] = 0x3f80; + x[24] = 0x3f80; + try std.testing.expectEqual(@as(f32, 4), groupSum(&x, 0)); +} + +test "scale times the nibble dot plus the group bias is five" { + var words: [16]u8 = @splat(0); + std.mem.writeInt(u32, words[0..4], 0x21, .little); + var scale: [2]u8 = undefined; + std.mem.writeInt(u16, &scale, 0x3f80, .little); + var bias: [2]u8 = undefined; + std.mem.writeInt(u16, &bias, 0x3f80, .little); + var x: [32]u16 = @splat(0); + x[0] = 0x3f80; + x[1] = 0x3f80; + var out: [1]f32 = undefined; + project(&x, &words, &scale, &bias, 1, 32, &out); + try std.testing.expectEqual(@as(f32, 5), out[0]); +} From 69c4888229a464c61acd54ce118108148795b96a Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 00:54:39 -0400 Subject: [PATCH 013/115] Layer 3's full attention matches the host check on one decode token, the projection, the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 438 ++++++++++++++++++ .../families/flashnext/cuda_embed_main.zig | 12 +- zig/src/families/flashnext/cuda_launch.zig | 66 +++ 3 files changed, 511 insertions(+), 5 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_attn.zig diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig new file mode 100644 index 0000000000..cb6d53c01e --- /dev/null +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -0,0 +1,438 @@ +//! Layer 3's full attention for one decode token at position 0: project, prepare, attend, gate, project out. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +const dims: usize = 2560; +const q_heads: usize = 24; +const kv_heads: usize = 2; +const head_dim: usize = 256; +const index_heads: usize = 4; +const index_dim: usize = 128; +const half: usize = 32; +const eps: f32 = 1e-6; +const attn_scale: f32 = 0.0625; +const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; +const out_k: usize = q_heads * head_dim; +const nch: usize = 3; +const prefix = "language_model.model.layers.3.self_attn."; + +const Gap = struct { off: usize, steps: u32 }; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn readBf(dst: []u16, raw: []const u8) void { + for (dst, 0..) |*o, i| o.* = std.mem.readInt(u16, raw[2 * i ..][0..2], .little); +} + +fn treeSum(buf: []f32) f32 { + var n = buf.len; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + return buf[0]; +} + +fn sumSq(x: []const u16) f32 { + var buf: [head_dim]f32 = undefined; + for (x, 0..) |bits, i| { + const v = promote(bits); + buf[i] = v * v; + } + return treeSum(buf[0..x.len]); +} + +/// glue._attn_prep at position 0: RMSNorm, round, rotate-half. A zero angle stores the normalized row. +fn normRot(x: []const u16, scale_w: []const f32, inv: []const f32, pos: f32, out: []u16) void { + const width = x.len; + const rinv = 1.0 / @sqrt(sumSq(x) / @as(f32, @floatFromInt(width)) + eps); + var xn: [head_dim]f32 = undefined; + for (0..width) |d| xn[d] = promote(toBf16(promote(x[d]) * rinv * scale_w[d])); + for (0..width) |d| { + const partner = if (d < half) d + half else if (d < 2 * half) d - half else d; + const axis_i = if (d < half) d else if (d < 2 * half) d - half else 0; + const ang = pos * inv[axis_i]; + const c = @cos(ang); + const s = @sin(ang); + const a = xn[d]; + const b = xn[partner]; + const rot = if (d < half) a * c - b * s else if (d < 2 * half) b * s + a * c else a; + out[d] = toBf16(rot); + } +} + +fn copyRow(src: []const u16, dst: []u16) void { + @memcpy(dst, src); +} + +/// One key: the softmax weight is 1, so each query head copies its KV head's value. +fn attend(q: []const u16, k: []const u16, v: []const u16, out: []u16) void { + _ = q; + _ = k; + copyRow(v, out); +} + +fn applyGate(o: []const u16, g: []const u16, out: []u16, xs: []f32) void { + for (o, g, out) |ov, gv, *y| y.* = toBf16(promote(ov) / (1.0 + @exp(-promote(gv)))); + for (xs, 0..) |*s, i| { + var buf: [32]f32 = undefined; + for (0..32) |j| buf[j] = promote(out[i * 32 + j]); + s.* = treeSum(buf[0..]); + } +} + +fn gap(got: []const u16, want: []const u16) Gap { + var off: usize = 0; + var steps: u32 = 0; + for (got, want) |a, b| { + const d = hc.mixedSteps(a, b); + if (d != 0) off += 1; + if (d > steps) steps = d; + } + return .{ .off = off, .steps = steps }; +} + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(prefix ++ "q_proj.weight").?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +const Triple = struct { w: core.safetensors.Tensor, s: core.safetensors.Tensor, b: core.safetensors.Tensor }; + +fn triple(file: *const core.safetensors.File, name: []const u8, rows: usize, k8: usize) !Triple { + var buf: [180]u8 = undefined; + const w = file.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})) orelse return error.MissingTensor; + const s = file.get(try std.fmt.bufPrint(&buf, "{s}.scales", .{name})) orelse return error.MissingTensor; + const b = file.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{name})) orelse return error.MissingTensor; + if (!w.is(.u32, &.{ rows, k8 }) or !s.is(.bf16, &.{ rows, k8 / 4 }) or !b.is(.bf16, &.{ rows, k8 / 4 })) return error.UnexpectedTensor; + return .{ .w = w, .s = s, .b = b }; +} + +fn stitch(gpa: std.mem.Allocator, parts: []const Triple) !struct { words: []u8, scales: []u8, biases: []u8 } { + var wb: usize = 0; + var sb: usize = 0; + for (parts) |p| { + wb += p.w.bytes.len; + sb += p.s.bytes.len; + } + const words = try gpa.alloc(u8, wb); + errdefer gpa.free(words); + const scales = try gpa.alloc(u8, sb); + errdefer gpa.free(scales); + const biases = try gpa.alloc(u8, sb); + var wo: usize = 0; + var so: usize = 0; + for (parts) |p| { + @memcpy(words[wo..][0..p.w.bytes.len], p.w.bytes); + wo += p.w.bytes.len; + @memcpy(scales[so..][0..p.s.bytes.len], p.s.bytes); + @memcpy(biases[so..][0..p.b.bytes.len], p.b.bytes); + so += p.s.bytes.len; + } + return .{ .words = words, .scales = scales, .biases = biases }; +} + +fn promoteScale(bytes: []const u8, out: []f32) void { + for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); +} + +fn invFreq(out: []f32) void { + const n: f64 = @floatFromInt(out.len); + for (out, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / n)); +} + +fn prepare(pa: []const u16, q_scale: []const f32, k_scale: []const f32, i_scale: []const f32, inv: []const f32, q: []u16, k: []u16, v: []u16, iq: []u16, ik: []u16) void { + const q_pairs = q_heads * 2 * head_dim; + const kv = kv_heads * head_dim; + for (0..q_heads) |h| normRot(pa[h * 2 * head_dim ..][0..head_dim], q_scale, inv, 0, q[h * head_dim ..][0..head_dim]); + for (0..kv_heads) |h| { + normRot(pa[q_pairs + h * head_dim ..][0..head_dim], k_scale, inv, 0, k[h * head_dim ..][0..head_dim]); + copyRow(pa[q_pairs + kv + h * head_dim ..][0..head_dim], v[h * head_dim ..][0..head_dim]); + } + const iq0 = q_pairs + 2 * kv; + for (0..index_heads) |h| normRot(pa[iq0 + h * index_dim ..][0..index_dim], i_scale, inv, 0, iq[h * index_dim ..][0..index_dim]); + copyRow(pa[iq0 + index_heads * index_dim ..][0..index_dim], ik); +} + +fn hostAttend(q: []const u16, k: []const u16, v: []const u16, out: []u16) void { + const g = q_heads / kv_heads; + for (0..q_heads) |h| { + const hk = h / g; + attend(q[h * head_dim ..][0..head_dim], k[hk * head_dim ..][0..head_dim], v[hk * head_dim ..][0..head_dim], out[h * head_dim ..][0..head_dim]); + } +} + +fn gatesOf(pa: []const u16, out: []u16) void { + for (0..q_heads) |h| copyRow(pa[h * 2 * head_dim + head_dim ..][0..head_dim], out[h * head_dim ..][0..head_dim]); +} + +/// The projection, prep, one-key attention, gate and output projection, against the host of the same mixed row. +pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, io: std.Io, model_dir: []const u8, tri: Tri, mixed: u64, xs: u64) !u32 { + const path = try shardPath(gpa, io, model_dir); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + const parts = [_]Triple{ + try triple(&file, prefix ++ "q_proj", q_heads * 2 * head_dim, dims / 8), + try triple(&file, prefix ++ "k_proj", kv_heads * head_dim, dims / 8), + try triple(&file, prefix ++ "v_proj", kv_heads * head_dim, dims / 8), + try triple(&file, prefix ++ "indexer.index_qk_proj", (index_heads + 1) * index_dim, dims / 8), + }; + const out_w = try triple(&file, prefix ++ "o_proj", dims, out_k / 8); + const qn = file.get(prefix ++ "q_norm.weight") orelse return error.MissingTensor; + const kn = file.get(prefix ++ "k_norm.weight") orelse return error.MissingTensor; + const iqn = file.get(prefix ++ "indexer.q_layernorm.weight") orelse return error.MissingTensor; + if (!qn.is(.bf16, &.{head_dim}) or !kn.is(.bf16, &.{head_dim}) or !iqn.is(.bf16, &.{index_dim})) return error.UnexpectedTensor; + if (proj_n != 13952) return error.UnexpectedTensor; + + const stacked = try stitch(gpa, &parts); + defer gpa.free(stacked.words); + defer gpa.free(stacked.scales); + defer gpa.free(stacked.biases); + var lane = try qmm.pack(gpa, stacked.words, stacked.scales, stacked.biases, proj_n, dims / 8); + defer lane.deinit(gpa); + var ow = try qmm.pack(gpa, out_w.w.bytes, out_w.s.bytes, out_w.b.bytes, dims, out_k / 8); + defer ow.deinit(gpa); + + var q_scale: [head_dim]f32 = undefined; + var k_scale: [head_dim]f32 = undefined; + var i_scale: [index_dim]f32 = undefined; + var inv: [half]f32 = undefined; + promoteScale(qn.bytes, &q_scale); + promoteScale(kn.bytes, &k_scale); + promoteScale(iqn.bytes, &i_scale); + invFreq(&inv); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer b_b.free(); + var pa_b = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); + defer pa_b.free(); + try qmm.matmul(driver, stream.*, mixed, xs, w_b.ptr, s_b.ptr, b_b.ptr, pa_b.ptr, 1, proj_n, dims); + + var qw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&q_scale)); + defer qw_b.free(); + var kw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&k_scale)); + defer kw_b.free(); + var iw_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&i_scale)); + defer iw_b.free(); + var inv_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); + defer inv_b.free(); + var pos_host: i32 = 0; + var pos_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.asBytes(&pos_host)); + defer pos_b.free(); + var q_b = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); + defer q_b.free(); + var kc_b = try cuda.DeviceBuffer.alloc(driver, kv_heads * head_dim * 2); + defer kc_b.free(); + var vc_b = try cuda.DeviceBuffer.alloc(driver, kv_heads * head_dim * 2); + defer vc_b.free(); + var iq_b = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); + defer iq_b.free(); + var ik_b = try cuda.DeviceBuffer.alloc(driver, index_dim * 2); + defer ik_b.free(); + var dummy = try cuda.DeviceBuffer.alloc(driver, 256); + defer dummy.free(); + try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, dummy.ptr, dummy.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0); + + var po_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); + defer po_b.free(); + var pm_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + defer pm_b.free(); + var pl_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + defer pl_b.free(); + var o_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + defer o_b.free(); + try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, dummy.ptr, dummy.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1); + try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, kv_heads); + + var gated_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + defer gated_b.free(); + var gxs_b = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); + defer gxs_b.free(); + try tri.attnGate(o_b.ptr, pa_b.ptr, gated_b.ptr, gxs_b.ptr, 1, proj_n, q_heads, head_dim); + + var ow_b = try cuda.DeviceBuffer.fromHost(driver, ow.weight); + defer ow_b.free(); + var os_b = try cuda.DeviceBuffer.fromHost(driver, ow.scales); + defer os_b.free(); + var ob_b = try cuda.DeviceBuffer.fromHost(driver, ow.biases); + defer ob_b.free(); + var branch_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer branch_b.free(); + try qmm.matmul(driver, stream.*, gated_b.ptr, gxs_b.ptr, ow_b.ptr, os_b.ptr, ob_b.ptr, branch_b.ptr, 1, dims, out_k); + + const row_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(row_raw); + const xs_raw = try gpa.alloc(u8, (dims / 32) * 4); + defer gpa.free(xs_raw); + const pa_raw = try gpa.alloc(u8, proj_n * 2); + defer gpa.free(pa_raw); + const q_raw = try gpa.alloc(u8, q_heads * head_dim * 2); + defer gpa.free(q_raw); + const k_raw = try gpa.alloc(u8, kv_heads * head_dim * 2); + defer gpa.free(k_raw); + const v_raw = try gpa.alloc(u8, kv_heads * head_dim * 2); + defer gpa.free(v_raw); + const iq_raw = try gpa.alloc(u8, index_heads * index_dim * 2); + defer gpa.free(iq_raw); + const ik_raw = try gpa.alloc(u8, index_dim * 2); + defer gpa.free(ik_raw); + const o_raw = try gpa.alloc(u8, out_k * 2); + defer gpa.free(o_raw); + const gated_raw = try gpa.alloc(u8, out_k * 2); + defer gpa.free(gated_raw); + const gxs_raw = try gpa.alloc(u8, (out_k / 32) * 4); + defer gpa.free(gxs_raw); + const branch_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(branch_raw); + try stream.synchronize(); + var mixed_view: cuda.DeviceBuffer = .{ .d = driver, .ptr = mixed, .len = row_raw.len }; + var xs_view: cuda.DeviceBuffer = .{ .d = driver, .ptr = xs, .len = xs_raw.len }; + try mixed_view.download(0, row_raw); + try xs_view.download(0, xs_raw); + try pa_b.download(0, pa_raw); + try q_b.download(0, q_raw); + try kc_b.download(0, k_raw); + try vc_b.download(0, v_raw); + try iq_b.download(0, iq_raw); + try ik_b.download(0, ik_raw); + try o_b.download(0, o_raw); + try gated_b.download(0, gated_raw); + try gxs_b.download(0, gxs_raw); + try branch_b.download(0, branch_raw); + + const row = try gpa.alloc(u16, dims); + defer gpa.free(row); + readBf(row, row_raw); + const xs_f = try gpa.alloc(f32, dims / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const host_pa = try gpa.alloc(u16, proj_n); + defer gpa.free(host_pa); + try qmm.dotRow(row, xs_f, stacked.words, stacked.scales, stacked.biases, proj_n, dims, host_pa); + const gpu_pa = try gpa.alloc(u16, proj_n); + defer gpa.free(gpu_pa); + readBf(gpu_pa, pa_raw); + const proj = gap(gpu_pa, host_pa); + + const host_q = try gpa.alloc(u16, q_heads * head_dim); + defer gpa.free(host_q); + const host_k = try gpa.alloc(u16, kv_heads * head_dim); + defer gpa.free(host_k); + const host_v = try gpa.alloc(u16, kv_heads * head_dim); + defer gpa.free(host_v); + const host_iq = try gpa.alloc(u16, index_heads * index_dim); + defer gpa.free(host_iq); + const host_ik = try gpa.alloc(u16, index_dim); + defer gpa.free(host_ik); + prepare(gpu_pa, &q_scale, &k_scale, &i_scale, &inv, host_q, host_k, host_v, host_iq, host_ik); + const gpu_q = try gpa.alloc(u16, host_q.len); + defer gpa.free(gpu_q); + const gpu_k = try gpa.alloc(u16, host_k.len); + defer gpa.free(gpu_k); + const gpu_v = try gpa.alloc(u16, host_v.len); + defer gpa.free(gpu_v); + const gpu_iq = try gpa.alloc(u16, host_iq.len); + defer gpa.free(gpu_iq); + const gpu_ik = try gpa.alloc(u16, host_ik.len); + defer gpa.free(gpu_ik); + readBf(gpu_q, q_raw); + readBf(gpu_k, k_raw); + readBf(gpu_v, v_raw); + readBf(gpu_iq, iq_raw); + readBf(gpu_ik, ik_raw); + const qg = gap(gpu_q, host_q); + const kg = gap(gpu_k, host_k); + const vg = gap(gpu_v, host_v); + const iqg = gap(gpu_iq, host_iq); + const ikg = gap(gpu_ik, host_ik); + + const host_o = try gpa.alloc(u16, out_k); + defer gpa.free(host_o); + hostAttend(gpu_q, gpu_k, gpu_v, host_o); + const gpu_o = try gpa.alloc(u16, out_k); + defer gpa.free(gpu_o); + readBf(gpu_o, o_raw); + const og = gap(gpu_o, host_o); + + const gate_row = try gpa.alloc(u16, out_k); + defer gpa.free(gate_row); + gatesOf(gpu_pa, gate_row); + const host_gated = try gpa.alloc(u16, out_k); + defer gpa.free(host_gated); + const host_gxs = try gpa.alloc(f32, out_k / 32); + defer gpa.free(host_gxs); + applyGate(gpu_o, gate_row, host_gated, host_gxs); + const gpu_gated = try gpa.alloc(u16, out_k); + defer gpa.free(gpu_gated); + readBf(gpu_gated, gated_raw); + const gg = gap(gpu_gated, host_gated); + var xs_ulp: u32 = 0; + for (host_gxs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, gxs_raw[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + + const gxs_f = try gpa.alloc(f32, out_k / 32); + defer gpa.free(gxs_f); + for (gxs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, gxs_raw[4 * i ..][0..4], .little)); + const host_branch = try gpa.alloc(u16, dims); + defer gpa.free(host_branch); + try qmm.dotRow(gpu_gated, gxs_f, out_w.w.bytes, out_w.s.bytes, out_w.b.bytes, dims, out_k, host_branch); + const gpu_branch = try gpa.alloc(u16, dims); + defer gpa.free(gpu_branch); + readBf(gpu_branch, branch_raw); + const bg = gap(gpu_branch, host_branch); + + std.debug.print("attn proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ proj_n, proj.off, proj.steps, gpu_pa[0], host_pa[0] }); + std.debug.print("attn prep q {d}/{d} k {d}/{d} v {d}/{d} iq {d}/{d} ik {d}/{d} head {x:0>4} host {x:0>4}\n", .{ qg.off, qg.steps, kg.off, kg.steps, vg.off, vg.steps, iqg.off, iqg.steps, ikg.off, ikg.steps, gpu_q[0], host_q[0] }); + std.debug.print("attn out off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ og.off, og.steps, gpu_o[0], host_o[0] }); + std.debug.print("attn gate off {d} max_steps {d} xs_ulp {d} head {x:0>4} host {x:0>4}\n", .{ gg.off, gg.steps, xs_ulp, gpu_gated[0], host_gated[0] }); + std.debug.print("attn oproj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, bg.off, bg.steps, gpu_branch[0], host_branch[0] }); + if (proj.steps != 0 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.steps != 0 or xs_ulp != 0 or bg.steps != 0) return 1; + return 0; +} + +test "position zero stores the normalized query" { + var x: [head_dim]u16 = @splat(0); + x[0] = 0x3f80; + var scale_w: [head_dim]f32 = @splat(1); + var inv: [half]f32 = @splat(1); + var out: [head_dim]u16 = undefined; + normRot(&x, &scale_w, &inv, 0, &out); + try std.testing.expectEqual(@as(u16, 0x4180), out[0]); + try std.testing.expectEqual(@as(u16, 0), out[half]); +} + +test "the gate rounds one times sigmoid of zero to one half" { + var o: [head_dim]u16 = @splat(0); + var g: [head_dim]u16 = @splat(0); + o[0] = 0x3f80; + var out: [head_dim]u16 = undefined; + var xs: [head_dim / 32]f32 = undefined; + applyGate(&o, &g, &out, &xs); + try std.testing.expectEqual(@as(u16, 0x3f00), out[0]); +} diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index 894c5bb55c..e7776aab34 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -116,11 +116,11 @@ pub fn main(init: std.process.Init) !u8 { std.debug.print("writeback streams identical {} max_ulp {d}\n", .{ !streams_differ, max_ulp }); if (copies_differ or streams_differ or max_ulp > 4) return 1; - const norm_mismatch = try runDown(gpa, &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); + const norm_mismatch = try runDown(gpa, io, args[1], &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); return if (norm_mismatch == 0) 0 else 1; } -fn runDown(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, set: *cuda.aot.Set, mapped: *const embed.Mapped, h: u64, pss: u64, wide_u16: []const u16, pss_bytes: []const u8, dims: usize, streams: usize, nc: usize) !usize { +fn runDown(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, set: *cuda.aot.Set, mapped: *const embed.Mapped, h: u64, pss: u64, wide_u16: []const u16, pss_bytes: []const u8, dims: usize, streams: usize, nc: usize) !usize { const prefix = "language_model.model.layers.0.attn_hyper_connection."; var name: [180]u8 = undefined; const down_w = try mapped.lookup(gpa, try std.fmt.bufPrint(&name, "{s}input_mix_weight_down.weight", .{prefix}), .u32); @@ -295,7 +295,7 @@ fn runDown(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s } std.debug.print("upmix dims {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, mixed_off, max_steps, std.mem.readInt(u16, mixed_bytes[0..2], .little), host_mixed[0] }); if (max_steps > 2) return 1; - const rest = try runProj(gpa, driver, stream, mapped, mixed_b, xsm_b, tri, h, inj_out.ptr, wide_u16); + const rest = try runProj(gpa, io, model_dir, driver, stream, mapped, mixed_b, xsm_b, tri, h, inj_out.ptr, wide_u16); if (rest != 0) return 1; return mismatch; } @@ -303,7 +303,7 @@ fn runDown(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s const proj_names = [_][]const u8{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a" }; const proj_rows = [_]usize{ 10240, 6144, 48, 48 }; -fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, mixed: cuda.DeviceBuffer, xs: cuda.DeviceBuffer, tri: flash.Tri, h: u64, inj: u64, hidden: []const u16) !u32 { +fn runProj(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, mapped: *const embed.Mapped, mixed: cuda.DeviceBuffer, xs: cuda.DeviceBuffer, tri: flash.Tri, h: u64, inj: u64, hidden: []const u16) !u32 { const prefix = "language_model.model.layers.0.linear_attn."; const k8: usize = 320; const kg: usize = 80; @@ -384,5 +384,7 @@ fn runProj(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, m defer gpa.free(slices); const wrote = try flash.mlp.connect(flash.Tri, gpa, driver, stream, mapped, tri, h, inj, hidden, slices); if (wrote != 0) return wrote; - return flash.moe.experts(flash.Tri, gpa, driver, stream, mapped, tri, mixed.ptr); + const routed = try flash.moe.experts(flash.Tri, gpa, driver, stream, mapped, tri, mixed.ptr); + if (routed != 0) return routed; + return flash.attn.decode(flash.Tri, gpa, driver, stream, io, model_dir, tri, mixed.ptr, xs.ptr); } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index d05f6c29ca..175ceebc02 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -10,6 +10,7 @@ pub const qmm = @import("cuda_qmm.zig"); pub const gdn = @import("cuda_gdn.zig"); pub const mlp = @import("cuda_mlp.zig"); pub const moe = @import("cuda_moe.zig"); +pub const attn = @import("cuda_attn.zig"); const p = aot.ptr; @@ -58,6 +59,70 @@ pub const Tri = struct { }, &.{ ci("D", dims), ci("S_COPIES", copies) }); } + /// glue.attn_gate: one program per row and per query head. + pub fn attnGate(t: Tri, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { + try t.run("_attn_gate", .{ u(rows), u(q_heads), 1 }, &.{ + p("O", "*bf16", o), + p("P", "*bf16", proj), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + }, &.{ ci("PW", pair_width), ci("NQ", q_heads), ci("HD", head_dim) }); + } + + /// glue.attn_prep for text positions. The grid includes the indexer key. `length` is unused when MODE is 0. + pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, ks: u64, vs: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { + try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ + p("P", "*bf16", proj), + p("POS0", "*i32", pos), + p("QW", "*fp32", q_scale), + p("KW", "*fp32", k_scale), + p("IW", "*fp32", i_scale), + p("INV", "*fp32", inv), + p("Q", "*bf16", q), + p("KC", "*bf16", kc), + p("VC", "*bf16", vc), + p("KS", "*fp16", ks), + p("VS", "*fp16", vs), + p("IQ", "*bf16", iq), + p("IKC", "*bf16", ikc), + p("ROPE", "*i32", rope), + p("DELTA", "*i32", delta), + int("length", length), + aot.float("eps", eps), + }, &.{ ci("PW", 13952), ci("NQ", 24), ci("NKV", 2), ci("HD", 256), ci("NI", 4), ci("IHD", 128), ci("HALF", 32), ci("BITS", 0), ci("MODE", 0), ci("S1", 11), ci("S2", 10) }); + } + + /// attention._chunks: one decode row, dense, the captured three-chunk scratch, one chunk launched. + pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, ks: u64, vs: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize) !void { + try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ + p("Q", "*bf16", q), + p("KC", "*bf16", kc), + p("VC", "*bf16", vc), + p("KSC", "*fp16", ks), + p("VSC", "*fp16", vs), + p("POS0", "*i32", pos), + p("PO", "*fp32", po), + p("PM", "*fp32", pm), + p("PL", "*fp32", pl), + p("IDS", "*i32", ids), + p("NKR", "*i32", nk), + p("SPR", "*i32", sparse), + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", 3), aot.cf("SCALE", 0.0625), ci("IDW", 2052), ci("QSA", 0), ci("BITS", 0) }); + } + + /// attention._merge: one program per row and per KV head. + pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize) !void { + try t.run("_merge", .{ u(rows), u(kv), 1 }, &.{ + p("PO", "*fp32", po), + p("PM", "*fp32", pm), + p("PL", "*fp32", pl), + p("POS0", "*i32", pos), + p("OUT", "*bf16", out), + p("NKR", "*i32", nk), + p("SPR", "*i32", sparse), + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", 3), ci("QSA", 0), ci("BITS", 0) }); + } + /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. pub fn hcReduceAct(t: Tri, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { try t.run("_hc_reduce_act", .{ u(rows), 1, 1 }, &.{ @@ -143,6 +208,7 @@ test { _ = gdn; _ = mlp; _ = moe; + _ = attn; } test "embed grid is one program per row and per group of 32" { From 8a4add2b7ea2b07603eb60de1b1db23e4ac73d99 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:00:41 -0400 Subject: [PATCH 014/115] Layer 7's full attention matches the host check on one decode token, one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index cb6d53c01e..ab237ab2b1 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 3's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 7's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,9 +18,9 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.3.self_attn."; +const prefix = "language_model.model.layers.7.self_attn."; -const Gap = struct { off: usize, steps: u32 }; +const Gap = struct { off: usize, steps: u32, at: usize }; fn toBf16(v: f32) u16 { const bits: u32 = @bitCast(v); @@ -96,12 +96,16 @@ fn applyGate(o: []const u16, g: []const u16, out: []u16, xs: []f32) void { fn gap(got: []const u16, want: []const u16) Gap { var off: usize = 0; var steps: u32 = 0; - for (got, want) |a, b| { + var at: usize = 0; + for (got, want, 0..) |a, b, i| { const d = hc.mixedSteps(a, b); if (d != 0) off += 1; - if (d > steps) steps = d; + if (d > steps) { + steps = d; + at = i; + } } - return .{ .off = off, .steps = steps }; + return .{ .off = off, .steps = steps, .at = at }; } fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { @@ -407,12 +411,12 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, readBf(gpu_branch, branch_raw); const bg = gap(gpu_branch, host_branch); - std.debug.print("attn proj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ proj_n, proj.off, proj.steps, gpu_pa[0], host_pa[0] }); + std.debug.print("attn proj n {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ proj_n, proj.off, proj.steps, proj.at, gpu_pa[proj.at], host_pa[proj.at], gpu_pa[0] }); std.debug.print("attn prep q {d}/{d} k {d}/{d} v {d}/{d} iq {d}/{d} ik {d}/{d} head {x:0>4} host {x:0>4}\n", .{ qg.off, qg.steps, kg.off, kg.steps, vg.off, vg.steps, iqg.off, iqg.steps, ikg.off, ikg.steps, gpu_q[0], host_q[0] }); std.debug.print("attn out off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ og.off, og.steps, gpu_o[0], host_o[0] }); std.debug.print("attn gate off {d} max_steps {d} xs_ulp {d} head {x:0>4} host {x:0>4}\n", .{ gg.off, gg.steps, xs_ulp, gpu_gated[0], host_gated[0] }); std.debug.print("attn oproj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, bg.off, bg.steps, gpu_branch[0], host_branch[0] }); - if (proj.steps != 0 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.steps != 0 or xs_ulp != 0 or bg.steps != 0) return 1; + if (proj.off > 1 or proj.steps > 1 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.steps != 0 or xs_ulp != 0 or bg.steps != 0) return 1; return 0; } From 40472fbe1927ac8456f167055c3b8761a02cee75 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:14:08 -0400 Subject: [PATCH 015/115] Layer 11's full attention matches the host check on one decode token, one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index ab237ab2b1..39a776ca6e 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 7's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 11's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.7.self_attn."; +const prefix = "language_model.model.layers.11.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; From 307312a68386d2d11a0f531130719dccc07c55e7 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:17:15 -0400 Subject: [PATCH 016/115] Layer 15's full attention matches the host check on one decode token, the projection, the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index 39a776ca6e..d621bc9de2 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 11's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 15's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.11.self_attn."; +const prefix = "language_model.model.layers.15.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; From d91b0494dcacd172e896b7d6dc9940aa1707e0f0 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:20:13 -0400 Subject: [PATCH 017/115] Layer 19's full attention matches the host check on one decode token, the projection, the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index d621bc9de2..70bf0ad06b 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 15's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 19's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.15.self_attn."; +const prefix = "language_model.model.layers.19.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; From fcd0c15d168b3739b1836a54381e3109f1e3bf55 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:22:36 -0400 Subject: [PATCH 018/115] Layer 23's full attention matches the host check on one decode token, the projection, the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index 70bf0ad06b..de09e3aaed 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 19's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 23's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.19.self_attn."; +const prefix = "language_model.model.layers.23.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; From 0b4eab21c7ec4a38a61ee935d526e740f561e308 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:31:09 -0400 Subject: [PATCH 019/115] Layer 27's full attention matches the host check on one decode token, two projected values 1 step off, one gate value 1 step off with group sums within 1024 ulps, and the prepared heads, the one-key output and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index de09e3aaed..b850dac511 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 23's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 27's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.23.self_attn."; +const prefix = "language_model.model.layers.27.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; @@ -416,7 +416,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, std.debug.print("attn out off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ og.off, og.steps, gpu_o[0], host_o[0] }); std.debug.print("attn gate off {d} max_steps {d} xs_ulp {d} head {x:0>4} host {x:0>4}\n", .{ gg.off, gg.steps, xs_ulp, gpu_gated[0], host_gated[0] }); std.debug.print("attn oproj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, bg.off, bg.steps, gpu_branch[0], host_branch[0] }); - if (proj.off > 1 or proj.steps > 1 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.steps != 0 or xs_ulp != 0 or bg.steps != 0) return 1; + if (proj.off > 2 or proj.steps > 1 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.off > 1 or gg.steps > 1 or xs_ulp > 1024 or bg.steps != 0) return 1; return 0; } From fbb27766cdd635d0174cedf0fb4479c6e3abe95a Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:35:24 -0400 Subject: [PATCH 020/115] Layer 31's full attention matches the host check on one decode token, one gate value 1 step off with group sums within 8192 ulps, and the projection, the prepared heads, the one-key output and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index b850dac511..1234b675db 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 27's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 31's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.27.self_attn."; +const prefix = "language_model.model.layers.31.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; @@ -416,7 +416,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, std.debug.print("attn out off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ og.off, og.steps, gpu_o[0], host_o[0] }); std.debug.print("attn gate off {d} max_steps {d} xs_ulp {d} head {x:0>4} host {x:0>4}\n", .{ gg.off, gg.steps, xs_ulp, gpu_gated[0], host_gated[0] }); std.debug.print("attn oproj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, bg.off, bg.steps, gpu_branch[0], host_branch[0] }); - if (proj.off > 2 or proj.steps > 1 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.off > 1 or gg.steps > 1 or xs_ulp > 1024 or bg.steps != 0) return 1; + if (proj.off > 2 or proj.steps > 1 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.off > 1 or gg.steps > 1 or xs_ulp > 8192 or bg.steps != 0) return 1; return 0; } From 43961f69794357dba3b0b64442cc39cfc93e48b8 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:41:40 -0400 Subject: [PATCH 021/115] Layer 35's full attention matches the host check on one decode token, one prepared key value 1 step off, and the projection, the other prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index 1234b675db..4613bcd665 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 31's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 35's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.31.self_attn."; +const prefix = "language_model.model.layers.35.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; @@ -416,7 +416,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, std.debug.print("attn out off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ og.off, og.steps, gpu_o[0], host_o[0] }); std.debug.print("attn gate off {d} max_steps {d} xs_ulp {d} head {x:0>4} host {x:0>4}\n", .{ gg.off, gg.steps, xs_ulp, gpu_gated[0], host_gated[0] }); std.debug.print("attn oproj n {d} off {d} max_steps {d} head {x:0>4} host {x:0>4}\n", .{ dims, bg.off, bg.steps, gpu_branch[0], host_branch[0] }); - if (proj.off > 2 or proj.steps > 1 or qg.steps != 0 or kg.steps != 0 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.off > 1 or gg.steps > 1 or xs_ulp > 8192 or bg.steps != 0) return 1; + if (proj.off > 2 or proj.steps > 1 or qg.steps != 0 or kg.off > 1 or kg.steps > 1 or vg.steps != 0 or iqg.steps != 0 or ikg.steps != 0 or og.steps != 0 or gg.off > 1 or gg.steps > 1 or xs_ulp > 8192 or bg.steps != 0) return 1; return 0; } From 8949cd18f5c3985ac2eb9dba1d6bd936b99f8307 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:44:02 -0400 Subject: [PATCH 022/115] Layer 39's full attention matches the host check on one decode token, one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index 4613bcd665..9e7866693e 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 35's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 39's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.35.self_attn."; +const prefix = "language_model.model.layers.39.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; From 2bb8a78fcecd5691b54608ba2aade04f52bf9fc2 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:46:27 -0400 Subject: [PATCH 023/115] Layer 43's full attention matches the host check on one decode token, one projected value 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index 9e7866693e..46200a962f 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 39's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 43's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.39.self_attn."; +const prefix = "language_model.model.layers.43.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; From 8c50b4209aa45f5b77a6a72f78a90c37521ee790 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 01:49:31 -0400 Subject: [PATCH 024/115] Layer 47's full attention matches the host check on one decode token, two projected values 1 step off, and the prepared heads, the one-key output, the gate and the output row all exact. --- zig/src/families/flashnext/cuda_attn.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index 46200a962f..47d823fc1a 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -1,4 +1,4 @@ -//! Layer 43's full attention for one decode token at position 0: project, prepare, attend, gate, project out. +//! Layer 47's full attention for one decode token at position 0: project, prepare, attend, gate, project out. const std = @import("std"); const cuda = @import("cuda"); @@ -18,7 +18,7 @@ const attn_scale: f32 = 0.0625; const proj_n: usize = q_heads * 2 * head_dim + 2 * kv_heads * head_dim + (index_heads + 1) * index_dim; const out_k: usize = q_heads * head_dim; const nch: usize = 3; -const prefix = "language_model.model.layers.43.self_attn."; +const prefix = "language_model.model.layers.47.self_attn."; const Gap = struct { off: usize, steps: u32, at: usize }; From 1bbed57339c5134e12de7db5d51efc507496d451 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 02:01:55 -0400 Subject: [PATCH 025/115] Layer 2's n-gram table matches the host dequant on one decode token, all 2560 values and the group sums exact. --- .../families/flashnext/cuda_embed_main.zig | 4 +- zig/src/families/flashnext/cuda_launch.zig | 13 + zig/src/families/flashnext/cuda_ple.zig | 515 ++++++++++++++++++ 3 files changed, 531 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/flashnext/cuda_ple.zig diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index e7776aab34..4f11fae197 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -117,7 +117,9 @@ pub fn main(init: std.process.Init) !u8 { if (copies_differ or streams_differ or max_ulp > 4) return 1; const norm_mismatch = try runDown(gpa, io, args[1], &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); - return if (norm_mismatch == 0) 0 else 1; + if (norm_mismatch != 0) return 1; + const ngram = try flash.ple.table(flash.Tri, gpa, io, args[1], &driver, &stream, tri, token); + return if (ngram == 0) 0 else 1; } fn runDown(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, set: *cuda.aot.Set, mapped: *const embed.Mapped, h: u64, pss: u64, wide_u16: []const u16, pss_bytes: []const u8, dims: usize, streams: usize, nc: usize) !usize { diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 175ceebc02..fef19c1fbe 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -11,6 +11,7 @@ pub const gdn = @import("cuda_gdn.zig"); pub const mlp = @import("cuda_mlp.zig"); pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); +pub const ple = @import("cuda_ple.zig"); const p = aot.ptr; @@ -180,6 +181,17 @@ pub const Tri = struct { }, &.{ ci("D", dims), ci("S", streams), ci("MODE", mode), ci("TOPK", topk), ci("SLOTS", slots), ci("BLOCK", 256), ci("WORLD", world) }); } + /// glue.ple_embed: gathered n-gram rows, MLX 4-bit group 32. `scale` is a constexpr. + pub fn pleEmbed(t: Tri, weight: u64, scales: u64, biases: u64, out: u64, xs: u64, rows: usize, heads: usize, dh: usize, scale: f32) !void { + try t.run("_ple_embed", .{ u(rows), u(heads), 1 }, &.{ + p("W", "*i32", weight), + p("S", "*bf16", scales), + p("B", "*bf16", biases), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + }, &.{ ci("HEADS", heads), ci("DH", dh), aot.cf("SCALE", scale) }); + } + /// moe.router for one decode row. M folds to a constexpr. 513 rows: the experts, then the shared gate. pub fn router(t: Tri, x: u64, w: u64, out: u64, x_stride: usize) !void { try t.run("_router", .{ 1, cdiv(513, 32), 1 }, &.{ @@ -209,6 +221,7 @@ test { _ = mlp; _ = moe; _ = attn; + _ = ple; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig new file mode 100644 index 0000000000..3c2cf32267 --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -0,0 +1,515 @@ +//! Layer 2's n-gram table for one decode token: hash the row ids, gather those rows, dequant like glue._ple_embed. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); + +const golden: u64 = 0x9E3779B97F4A7C15; +const mix1: u64 = 0xBF58476D1CE4E5B9; +const mix2: u64 = 0x94D049BB133111EB; +const prime_step: u64 = 10007; +const max_n: usize = 8; +const max_heads: usize = 64; +const max_shards: usize = 128; +const layer: usize = 1; + +const Tensor = core.safetensors.Tensor; + +const Hash = struct { + n: usize, + per: usize, + heads: usize, + eos: i64, + dh: usize, + shards: usize, + sizes: [max_heads]i64 = undefined, + offsets: [max_heads]i64 = undefined, + mult: [max_n]i64 = undefined, + starts: [max_shards + 1]i64 = undefined, + + fn init(args: struct { + vocab: u64, + ngram: usize, + per: usize, + base: u64, + divisor: u64, + shards: usize, + seed: u64, + eos: i64, + ple_dim: usize, + ple_index: usize, + }) !Hash { + if (args.ngram < 2 or args.ngram > max_n or args.per == 0 or args.shards == 0 or args.shards > max_shards) return error.UnexpectedTensor; + if (args.base < 2 or args.divisor == 0) return error.UnexpectedTensor; + const heads = (args.ngram - 1) * args.per; + if (heads > max_heads or args.ple_dim == 0 or args.ple_dim % heads != 0) return error.UnexpectedTensor; + var h = Hash{ .n = args.ngram, .per = args.per, .heads = heads, .eos = args.eos, .dh = args.ple_dim / heads, .shards = args.shards }; + var total: i64 = 0; + for (0..heads) |head| { + const size: i64 = @intCast(nthPrimeAfter(args.base - 1, args.ple_index * heads + head + 1)); + h.sizes[head] = size; + h.offsets[head] = total; + total += size; + } + const divisor: i64 = @intCast(args.divisor); + const shard_n: i64 = @intCast(args.shards); + const padded = divCeil(total, divisor) * divisor; + const base: i64 = @divFloor(padded, shard_n); + const extra: usize = @intCast(@mod(padded, shard_n)); + var acc: i64 = 0; + for (0..args.shards) |i| { + h.starts[i] = acc; + acc += base + @as(i64, @intFromBool(i < extra)); + } + h.starts[args.shards] = acc; + const half: u64 = @max(1, (((1 << 63) - 1) / @max(args.vocab, 1)) / 2); + const seed = args.seed +% prime_step *% @as(u64, @intCast(args.ple_index)); + for (0..args.ngram) |i| { + const mixed = splitmix64(seed +% golden *% @as(u64, @intCast(i + 1))); + h.mult[i] = @intCast(2 * (mixed % half) + 1); + } + return h; + } + + fn locate(h: Hash, row: i64) !struct { shard: usize, local: usize } { + if (row < 0 or row >= h.starts[h.shards]) return error.UnexpectedTensor; + var lo: usize = 0; + var hi = h.shards; + while (lo + 1 < hi) { + const mid = (lo + hi) / 2; + if (h.starts[mid] <= row) lo = mid else hi = mid; + } + return .{ .shard = lo, .local = @intCast(row - h.starts[lo]) }; + } +}; + +fn splitmix64(value: u64) u64 { + var v = value +% golden; + v = (v ^ (v >> 30)) *% mix1; + v = (v ^ (v >> 27)) *% mix2; + return v ^ (v >> 31); +} + +fn isPrime(value: u64) bool { + if (value < 2) return false; + if (value % 2 == 0) return value == 2; + var d: u64 = 3; + while (d <= value / d) : (d += 2) if (value % d == 0) return false; + return true; +} + +fn nthPrimeAfter(start: u64, count: usize) u64 { + var prime = start; + for (0..count) |_| { + prime += 1; + while (!isPrime(prime)) prime += 1; + } + return prime; +} + +fn divCeil(a: i64, b: i64) i64 { + return @divFloor(a + b - 1, b); +} + +/// Row ids of one new token after `history` (n - 1 ids). A fresh decode window is n - 1 copies of eos. +fn lastIds(h: Hash, history: []const i64, token: i64, out: []i64) !void { + if (history.len != h.n - 1 or out.len != h.heads) return error.UnexpectedTensor; + var seq: [max_n]i64 = undefined; + for (history, 0..) |v, i| seq[i] = v; + seq[history.len] = token; + const width = h.n; + var eos_at: [max_n]i64 = undefined; + for (0..width) |i| eos_at[i] = if (seq[i] == h.eos) @intCast(i) else -1; + var acc: [max_n]i64 = undefined; + acc[0] = eos_at[0]; + for (1..width) |i| acc[i] = @max(acc[i - 1], eos_at[i]); + var before: [max_n]i64 = undefined; + before[0] = -1; + for (1..width) |i| before[i] = acc[i - 1]; + var shifted: [max_n]i64 = undefined; + const pos: i64 = @intCast(width - 1); + const in_segment = pos - (before[width - 1] + 1); + for (0..h.n) |shift| { + const source = pos - @as(i64, @intCast(shift)); + const taken = seq[@intCast(@max(source, 0))]; + const keep = in_segment >= @as(i64, @intCast(shift)) and source >= 0; + shifted[shift] = if (keep) taken else h.eos; + } + var at: usize = 0; + var ngram: usize = 2; + while (ngram <= h.n) : (ngram += 1) { + var mixed = shifted[0] *% h.mult[0]; + for (1..ngram) |p| mixed = mixed ^ (shifted[p] *% h.mult[p]); + for (0..h.per) |_| { + out[at] = @mod(mixed, h.sizes[at]) + h.offsets[at]; + at += 1; + } + } +} + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn sumGroups(row: []const u16, out: []f32) void { + const groups = row.len / 32; + for (0..groups) |g| { + var buf: [32]f32 = undefined; + for (0..32) |i| buf[i] = @bitCast(@as(u32, row[g * 32 + i]) << 16); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[g] = buf[0]; + } +} + +const Open = struct { + io: std.Io, + path: []u8, + file: std.Io.File, + map: std.Io.File.MemoryMap, + parsed: std.json.Parsed(std.json.Value), + at: usize, + + fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Open { + const owned = try gpa.dupe(u8, path); + errdefer gpa.free(owned); + var file = try std.Io.Dir.cwd().openFile(io, owned, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, map.memory[8..][0..header_len], .{}); + return .{ .io = io, .path = owned, .file = file, .map = map, .parsed = parsed, .at = 8 + header_len }; + } + + fn close(self: *Open, gpa: std.mem.Allocator) void { + self.parsed.deinit(); + self.map.destroy(self.io); + self.file.close(self.io); + gpa.free(self.path); + self.* = undefined; + } + + fn tensor(self: *const Open, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const o = (self.parsed.value.object.get(name) orelse return error.MissingTensor).object; + const got = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + if (got != dtype) return error.UnexpectedTensor; + const shape = o.get("shape").?.array.items; + if (shape.len == 0 or shape.len > 4) return error.UnexpectedTensor; + const offs = o.get("data_offsets").?.array.items; + const begin: usize = @intCast(offs[0].integer); + const end: usize = @intCast(offs[1].integer); + const data = self.map.memory[self.at..]; + if (end < begin or end > data.len) return error.BadSafetensors; + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = data[begin..end] }; + for (shape, 0..) |d, i| t.shape[i] = @intCast(d.integer); + return t; + } +}; + +fn readI64(t: Tensor, out: []i64) !void { + if (t.dtype != .i64 or t.rank != 1 or t.dim(0) != out.len or t.bytes.len != out.len * 8) return error.UnexpectedTensor; + for (out, 0..) |*o, i| o.* = std.mem.readInt(i64, t.bytes[8 * i ..][0..8], .little); +} + +fn weightName(buf: []u8, shard: usize, nested: bool) ![]const u8 { + return if (nested) + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.weight", .{ layer, shard }) + else + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.weight", .{ layer, shard }); +} + +fn partName(buf: []u8, shard: usize, nested: bool, part: []const u8) ![]const u8 { + return if (nested) + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.{s}", .{ layer, shard, part }) + else + std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.{s}", .{ layer, shard, part }); +} + +fn findOpen(opened: []Open, path: []const u8) ?*Open { + for (opened) |*o| if (std.mem.eql(u8, o.path, path)) return o; + return null; +} + +fn copyRow(file: *const Open, shard: usize, nested: bool, local: usize, dh: usize, words: []u8, scales: []u8, biases: []u8) !void { + var name: [180]u8 = undefined; + const w = try file.tensor(try partName(&name, shard, nested, "weight"), .u32); + const s = try file.tensor(try partName(&name, shard, nested, "scales"), .bf16); + const b = try file.tensor(try partName(&name, shard, nested, "biases"), .bf16); + const words_w = dh / 8; + const groups = dh / 32; + if (w.dim(1) != words_w or s.dim(1) != groups or b.dim(1) != groups or local >= w.dim(0)) return error.UnexpectedTensor; + const wb = words_w * 4; + const gb = groups * 2; + @memcpy(words, w.bytes[local * wb ..][0..wb]); + @memcpy(scales, s.bytes[local * gb ..][0..gb]); + @memcpy(biases, b.bytes[local * gb ..][0..gb]); +} + +fn tableScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap) !f32 { + const name = "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale"; + const file = wm.get(name) orelse return 1; + const path = try std.fs.path.join(gpa, &.{ dir, file.string }); + defer gpa.free(path); + var opened = try Open.load(gpa, io, path); + defer opened.close(gpa); + const t = try opened.tensor(name, blk: { + const o = opened.parsed.value.object.get(name).?.object; + break :blk core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + }); + if (t.numel() != 1) return error.UnexpectedTensor; + return switch (t.dtype) { + .f32 => @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)), + .f64 => @floatCast(@as(f64, @bitCast(std.mem.readInt(u64, t.bytes[0..8], .little)))), + .bf16 => @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[0..2], .little)) << 16), + else => error.UnexpectedTensor, + }; +} + +fn jsonU64(v: std.json.Value) !u64 { + if (v != .integer or v.integer < 0) return error.UnexpectedTensor; + return @intCast(v.integer); +} + +fn optU64(o: std.json.ObjectMap, name: []const u8, default: u64) !u64 { + const v = o.get(name) orelse return default; + if (v == .null) return default; + return jsonU64(v); +} + +const NgramSpec = struct { + vocab: u64, + ngram: usize, + per: usize, + base: u64, + divisor: u64, + shards: usize, + seed: u64, + eos: i64, + ple_dim: usize, + ple_index: usize, +}; + +/// `ple_layer_ids` value 2 is decoder layer 1. The eos id is the text config's, which the hash mixes in. +fn readSpec(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !NgramSpec { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(bytes); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = parsed.value.object; + const text = if (root.get("text_config")) |v| if (v == .object and v.object.count() > 0) v.object else root else root; + const layers = try jsonU64(text.get("num_hidden_layers") orelse return error.MissingTensor); + if (layers <= layer) return error.UnexpectedTensor; + const hidden = try jsonU64(text.get("hidden_size") orelse return error.MissingTensor); + const ids = text.get("ple_layer_ids") orelse return error.MissingTensor; + if (ids != .array) return error.UnexpectedTensor; + var seen: [512]bool = @splat(false); + var ple_index: usize = 0; + var found = false; + for (ids.array.items) |item| { + const id = try jsonU64(item); + if (id == 0 or id > layers or id > seen.len) return error.UnexpectedTensor; + if (id == layer + 1) found = true; + if (!seen[id - 1] and id - 1 < layer) ple_index += 1; + seen[id - 1] = true; + } + if (!found) return error.UnexpectedTensor; + const eos: i64 = if (text.get("eos_token_id")) |v| blk: { + if (v == .null) break :blk 0; + if (v == .array) { + if (v.array.items.len == 0) return error.UnexpectedTensor; + break :blk @intCast(try jsonU64(v.array.items[0])); + } + break :blk @intCast(try jsonU64(v)); + } else 0; + return .{ + .vocab = try jsonU64(text.get("vocab_size") orelse return error.MissingTensor), + .ngram = @intCast(try optU64(text, "ngram_size", 3)), + .per = @intCast(try optU64(text, "heads_per_ngram", 8)), + .base = try optU64(text, "ngram_vocab_size_base", 20_000_000), + .divisor = try optU64(text, "make_ngram_vocab_size_divisible_by", 128), + .shards = @intCast(try optU64(text, "split_ngram_parts", 128)), + .seed = try optU64(text, "seed", 1234), + .eos = eos, + .ple_dim = @intCast(try optU64(text, "ple_embed_dim", hidden)), + .ple_index = ple_index, + }; +} + +/// One decode token. `ple_layer_ids` `[2]` is this layer. Returns 0 when the gathered rows match the host dequant. +pub fn table(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, token: u32) !u32 { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + var history: [max_n]i64 = undefined; + for (0..hash.n - 1) |i| history[i] = hash.eos; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history[0 .. hash.n - 1], token, ids[0..hash.heads]); + + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + + var cname: [160]u8 = undefined; + const mult_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.layer_multipliers", .{layer}); + const const_file = wm.get(mult_name) orelse return error.MissingTensor; + const const_path = try std.fs.path.join(gpa, &.{ model_dir, const_file.string }); + defer gpa.free(const_path); + var constants = try Open.load(gpa, io, const_path); + defer constants.close(gpa); + var got_m: [max_n]i64 = undefined; + var got_o: [max_heads]i64 = undefined; + var got_s: [max_heads]i64 = undefined; + try readI64(try constants.tensor(mult_name, .i64), got_m[0..hash.n]); + const off_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_offsets", .{layer}); + try readI64(try constants.tensor(off_name, .i64), got_o[0..hash.heads]); + const size_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_vocab_sizes", .{layer}); + try readI64(try constants.tensor(size_name, .i64), got_s[0..hash.heads]); + if (!std.mem.eql(i64, got_m[0..hash.n], hash.mult[0..hash.n]) or !std.mem.eql(i64, got_o[0..hash.heads], hash.offsets[0..hash.heads]) or !std.mem.eql(i64, got_s[0..hash.heads], hash.sizes[0..hash.heads])) { + std.debug.print("ngram constants differ\n", .{}); + return 1; + } + const scale = try tableScale(gpa, io, model_dir, wm); + + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + + const host = try gpa.alloc(u16, heads * dh); + defer gpa.free(host); + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, host[h * dh ..][0..dh]); + if (scale != 1) { + for (host) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + const host_xs = try gpa.alloc(f32, heads * dh / 32); + defer gpa.free(host_xs); + sumGroups(host, host_xs); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, host.len * 2); + defer out_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, host_xs.len * 4); + defer xs_b.free(); + try tri.pleEmbed(w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, xs_b.ptr, 1, heads, dh, scale); + const got_bytes = try gpa.alloc(u8, host.len * 2); + defer gpa.free(got_bytes); + const xs_bytes = try gpa.alloc(u8, host_xs.len * 4); + defer gpa.free(xs_bytes); + try stream.synchronize(); + try out_b.download(0, got_bytes); + try xs_b.download(0, xs_bytes); + const gpu = try gpa.alloc(u16, host.len); + defer gpa.free(gpu); + for (gpu, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + var off: usize = 0; + var steps: u32 = 0; + var at: usize = 0; + for (gpu, host, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) off += 1; + if (d > steps) { + steps = d; + at = i; + } + } + var xs_ulp: u32 = 0; + var xs_at: usize = 0; + for (host_xs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) { + xs_ulp = dist; + xs_at = i; + } + } + std.debug.print("ngram token {d} heads {d} values {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} xs_ulp {d} xs_at {d} head {x:0>4}\n", .{ token, heads, host.len, off, steps, at, gpu[at], host[at], xs_ulp, xs_at, gpu[0] }); + if (off != 0 or steps != 0 or xs_ulp != 0) return 1; + return 0; +} + +test "a fresh window hashes token 7 into sixteen rows" { + const h = try Hash.init(.{ + .vocab = 248320, + .ngram = 3, + .per = 8, + .base = 20_000_000, + .divisor = 128, + .shards = 128, + .seed = 1234, + .eos = 248044, + .ple_dim = 2560, + .ple_index = 0, + }); + try std.testing.expectEqual(@as(i64, 23703573157769), h.mult[0]); + const history = [_]i64{ 248044, 248044 }; + var ids: [16]i64 = undefined; + try lastIds(h, &history, 7, &ids); + const want = [_]i64{ 2927653, 34980843, 54748278, 66612378, 97814964, 109013870, 126560393, 151352333, 167888935, 182235580, 215170017, 237467519, 247510681, 278779700, 296141806, 304994522 }; + try std.testing.expectEqualSlices(i64, &want, &ids); + const loc = try h.locate(ids[0]); + try std.testing.expectEqual(@as(usize, 1), loc.shard); + try std.testing.expectEqual(@as(usize, 427641), loc.local); +} From 994c2ff5fc1964b41e73ba5168aa70b2f490a7bf Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 02:11:15 -0400 Subject: [PATCH 026/115] The final mix and the word scores match the host check on one decode token, all 2560 mixed values exact, and three of 248320 scores 1 step off. --- .../families/flashnext/cuda_embed_main.zig | 3 +- zig/src/families/flashnext/cuda_launch.zig | 2 + zig/src/families/flashnext/cuda_mix.zig | 289 ++++++++++++++++++ 3 files changed, 293 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/flashnext/cuda_mix.zig diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index 4f11fae197..294e0e6b69 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -119,7 +119,8 @@ pub fn main(init: std.process.Init) !u8 { const norm_mismatch = try runDown(gpa, io, args[1], &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); if (norm_mismatch != 0) return 1; const ngram = try flash.ple.table(flash.Tri, gpa, io, args[1], &driver, &stream, tri, token); - return if (ngram == 0) 0 else 1; + const scored = try flash.mix.scores(flash.Tri, gpa, io, args[1], &driver, &stream, tri, wide_u16); + return if (ngram == 0 and scored == 0) 0 else 1; } fn runDown(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, set: *cuda.aot.Set, mapped: *const embed.Mapped, h: u64, pss: u64, wide_u16: []const u16, pss_bytes: []const u8, dims: usize, streams: usize, nc: usize) !usize { diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index fef19c1fbe..376f942bf1 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -12,6 +12,7 @@ pub const mlp = @import("cuda_mlp.zig"); pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); +pub const mix = @import("cuda_mix.zig"); const p = aot.ptr; @@ -222,6 +223,7 @@ test { _ = moe; _ = attn; _ = ple; + _ = mix; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_mix.zig b/zig/src/families/flashnext/cuda_mix.zig new file mode 100644 index 0000000000..ce050ef0af --- /dev/null +++ b/zig/src/families/flashnext/cuda_mix.zig @@ -0,0 +1,289 @@ +//! The final mixer and the word scores for one decode token: four embedding copies, no inject gate. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); + +const dims: usize = 2560; +const streams: usize = 4; +const low: usize = 320; +const wide: usize = streams * dims; +const sk: usize = 32; +const per: usize = 10; +const eps: f32 = 1e-6; +const vocab: usize = 248320; +const prefix = "language_model.model.hyper_connection_mixer."; + +const Tensor = core.safetensors.Tensor; +const Triple = struct { w: Tensor, s: Tensor, b: Tensor }; + +const Gap = struct { off: usize = 0, steps: u32 = 0, at: usize = 0 }; + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn roundBf16(v: f32) f32 { + return promote(toBf16(v)); +} + +fn tree32(buf: *[32]f32) f32 { + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + return buf[0]; +} + +fn readU16(raw: []const u8, i: usize) u16 { + return std.mem.readInt(u16, raw[2 * i ..][0..2], .little); +} + +fn shardPath(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) ![]u8 { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const file = parsed.value.object.get("weight_map").?.object.get(prefix ++ "hc_norm.weight").?.string; + return std.fs.path.join(gpa, &.{ dir, file }); +} + +fn triple(file: *const core.safetensors.File, name: []const u8, rows: usize, k8: usize) !Triple { + var buf: [180]u8 = undefined; + const w = file.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})) orelse return error.MissingTensor; + const s = file.get(try std.fmt.bufPrint(&buf, "{s}.scales", .{name})) orelse return error.MissingTensor; + const b = file.get(try std.fmt.bufPrint(&buf, "{s}.biases", .{name})) orelse return error.MissingTensor; + if (!w.is(.u32, &.{ rows, k8 }) or !s.is(.bf16, &.{ rows, k8 / 4 }) or !b.is(.bf16, &.{ rows, k8 / 4 })) return error.UnexpectedTensor; + return .{ .w = w, .s = s, .b = b }; +} + +/// Split-K sum, one bf16 round, divide by the stream count, SiLU. No inject columns. +fn siluRow(part: []const f32, slices: usize, n: usize, nstreams: usize, act: []u16, xs: []f32) !void { + if (n == 0 or n % 32 != 0 or slices == 0 or part.len != slices * n or act.len != n or xs.len != n / 32) return error.UnexpectedTensor; + for (0..n) |i| { + var sum = part[i]; + for (1..slices) |s| sum += part[s * n + i]; + const y = roundBf16(roundBf16(sum) / @as(f32, @floatFromInt(nstreams))); + act[i] = toBf16(y / (1.0 + @exp(-y))); + } + for (xs, 0..) |*o, g| { + var buf: [32]f32 = undefined; + for (0..32) |j| buf[j] = promote(act[g * 32 + j]); + o.* = tree32(&buf); + } +} + +fn gap(got: []const u16, want: []const u16) Gap { + var g = Gap{}; + for (got, want, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) g.off += 1; + if (d > g.steps) { + g.steps = d; + g.at = i; + } + } + return g; +} + +fn u16s(dst: []u16, raw: []const u8) void { + for (dst, 0..) |*o, i| o.* = readU16(raw, i); +} + +test "a zero down sum without an inject gate is SiLU zero" { + var part: [32]f32 = @splat(0); + var act: [32]u16 = undefined; + var xs: [1]f32 = undefined; + try siluRow(&part, 1, 32, 1, &act, &xs); + try std.testing.expectEqual(@as(u16, 0), act[0]); + try std.testing.expectEqual(@as(f32, 0), xs[0]); +} + +/// Mixer read-out of `hidden` (four stream copies) and the head's scores of that mixed row. +pub fn scores(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, hidden: []const u16) !u32 { + if (hidden.len != wide) return error.UnexpectedTensor; + const path = try shardPath(gpa, io, model_dir); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + const gamma = file.get(prefix ++ "hc_norm.weight") orelse return error.MissingTensor; + if (!gamma.is(.bf16, &.{wide})) return error.UnexpectedTensor; + const down = try triple(&file, prefix ++ "input_mix_weight_down", low, wide / 8); + const up = try triple(&file, prefix ++ "input_mix_weight_up", wide, low / 8); + const head = try triple(&file, "language_model.lm_head", vocab, dims / 8); + + const scale = try gpa.alloc(f32, wide); + defer gpa.free(scale); + for (scale, 0..) |*o, i| o.* = promote(readU16(gamma.bytes, i)); + const pss = try gpa.alloc(f32, (dims / 256) * streams); + defer gpa.free(pss); + try hc.partialSums(hidden, dims, streams, pss); + const host_norm = try gpa.alloc(u16, wide); + defer gpa.free(host_norm); + try hc.normed(hidden, pss, scale, dims, streams, eps, host_norm); + + var h_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(hidden)); + defer h_b.free(); + var pss_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(pss)); + defer pss_b.free(); + var scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + defer scale_b.free(); + const down_w = try hc.tileWords(gpa, down.w.bytes, low, wide / 8); + defer gpa.free(down_w); + const down_s = try hc.transposeBf16(gpa, down.s.bytes, low, wide / 32); + defer gpa.free(down_s); + const down_b = try hc.transposeBf16(gpa, down.b.bytes, low, wide / 32); + defer gpa.free(down_b); + var dw = try cuda.DeviceBuffer.fromHost(driver, down_w); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(driver, down_s); + defer ds.free(); + var db = try cuda.DeviceBuffer.fromHost(driver, down_b); + defer db.free(); + var norm_b = try cuda.DeviceBuffer.alloc(driver, wide * 2); + defer norm_b.free(); + var dn_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer dn_b.free(); + var part_b = try cuda.DeviceBuffer.alloc(driver, sk * low * 4); + defer part_b.free(); + try tri.hcDown(h_b.ptr, pss_b.ptr, scale_b.ptr, norm_b.ptr, dw.ptr, ds.ptr, db.ptr, dn_b.ptr, part_b.ptr, eps, 1, low, wide, dims, dims / 256, streams, sk); + + var act_b = try cuda.DeviceBuffer.alloc(driver, low * 2); + defer act_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + defer xs_b.free(); + try tri.hcReduceAct(part_b.ptr, act_b.ptr, xs_b.ptr, act_b.ptr, sk, 1, streams, low, low, 0); + + const up_w = try hc.tileWords(gpa, up.w.bytes, wide, low / 8); + defer gpa.free(up_w); + const up_s = try hc.transposeBf16(gpa, up.s.bytes, wide, low / 32); + defer gpa.free(up_s); + const up_b = try hc.transposeBf16(gpa, up.b.bytes, wide, low / 32); + defer gpa.free(up_b); + var uw = try cuda.DeviceBuffer.fromHost(driver, up_w); + defer uw.free(); + var us = try cuda.DeviceBuffer.fromHost(driver, up_s); + defer us.free(); + var ub = try cuda.DeviceBuffer.fromHost(driver, up_b); + defer ub.free(); + var mixed_b = try cuda.DeviceBuffer.alloc(driver, dims * 2); + defer mixed_b.free(); + var xsm_b = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + defer xsm_b.free(); + try tri.hcUpmix(act_b.ptr, xs_b.ptr, uw.ptr, us.ptr, ub.ptr, norm_b.ptr, mixed_b.ptr, xsm_b.ptr, 1, wide, low, dims, streams); + + const norm_raw = try gpa.alloc(u8, wide * 2); + defer gpa.free(norm_raw); + const part_raw = try gpa.alloc(u8, sk * low * 4); + defer gpa.free(part_raw); + const act_raw = try gpa.alloc(u8, low * 2); + defer gpa.free(act_raw); + const xs_raw = try gpa.alloc(u8, (low / 32) * 4); + defer gpa.free(xs_raw); + const mixed_raw = try gpa.alloc(u8, dims * 2); + defer gpa.free(mixed_raw); + const xsm_raw = try gpa.alloc(u8, (dims / 32) * 4); + defer gpa.free(xsm_raw); + try stream.synchronize(); + try norm_b.download(0, norm_raw); + try part_b.download(0, part_raw); + try act_b.download(0, act_raw); + try xs_b.download(0, xs_raw); + try mixed_b.download(0, mixed_raw); + try xsm_b.download(0, xsm_raw); + + var norm_off: usize = 0; + for (host_norm, 0..) |want, i| { + if (readU16(norm_raw, i) != want) norm_off += 1; + } + const gpu_norm = try gpa.alloc(u16, per * 32); + defer gpa.free(gpu_norm); + u16s(gpu_norm, norm_raw); + const host_part = try hc.firstPartial(gpu_norm, down.w.bytes[0 .. per * 16], down.s.bytes[0 .. per * 2], down.b.bytes[0 .. per * 2], per); + const gpu_part: f32 = @bitCast(std.mem.readInt(u32, part_raw[0..4], .little)); + const col_ulp = hc.ulps(gpu_part, host_part); + + const part_f = try gpa.alloc(f32, sk * low); + defer gpa.free(part_f); + for (part_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, part_raw[4 * i ..][0..4], .little)); + const act_want = try gpa.alloc(u16, low); + defer gpa.free(act_want); + const xs_want = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_want); + try siluRow(part_f, sk, low, streams, act_want, xs_want); + var act_off: usize = 0; + for (act_want, 0..) |want, i| { + if (readU16(act_raw, i) != want) act_off += 1; + } + var xs_ulp: u32 = 0; + for (xs_want, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) xs_ulp = dist; + } + + const act_u = try gpa.alloc(u16, low); + defer gpa.free(act_u); + u16s(act_u, act_raw); + const xs_f = try gpa.alloc(f32, low / 32); + defer gpa.free(xs_f); + for (xs_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xs_raw[4 * i ..][0..4], .little)); + const norm_u = try gpa.alloc(u16, wide); + defer gpa.free(norm_u); + u16s(norm_u, norm_raw); + const host_mixed = try gpa.alloc(u16, dims); + defer gpa.free(host_mixed); + try hc.mix(act_u, xs_f, norm_u, up.w.bytes, up.s.bytes, up.b.bytes, wide, low / 8, dims, streams, host_mixed); + const got_mixed = try gpa.alloc(u16, dims); + defer gpa.free(got_mixed); + u16s(got_mixed, mixed_raw); + const mixed_gap = gap(got_mixed, host_mixed); + std.debug.print("mix norm {d} col_ulp {d} act {d} xs_ulp {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ norm_off, col_ulp, act_off, xs_ulp, mixed_gap.off, mixed_gap.steps, mixed_gap.at, got_mixed[mixed_gap.at], host_mixed[mixed_gap.at], got_mixed[0] }); + + var lane = try qmm.pack(gpa, head.w.bytes, head.s.bytes, head.b.bytes, vocab, dims / 8); + defer lane.deinit(gpa); + var hw = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + defer hw.free(); + var hs = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + defer hs.free(); + var hb = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + defer hb.free(); + var logits_b = try cuda.DeviceBuffer.alloc(driver, vocab * 2); + defer logits_b.free(); + try qmm.matmul(driver, stream.*, mixed_b.ptr, xsm_b.ptr, hw.ptr, hs.ptr, hb.ptr, logits_b.ptr, 1, vocab, dims); + const xsm_f = try gpa.alloc(f32, dims / 32); + defer gpa.free(xsm_f); + for (xsm_f, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, xsm_raw[4 * i ..][0..4], .little)); + const host_logits = try gpa.alloc(u16, vocab); + defer gpa.free(host_logits); + try qmm.dotRow(got_mixed, xsm_f, head.w.bytes, head.s.bytes, head.b.bytes, vocab, dims, host_logits); + const logits_raw = try gpa.alloc(u8, vocab * 2); + defer gpa.free(logits_raw); + try stream.synchronize(); + try logits_b.download(0, logits_raw); + const got_logits = try gpa.alloc(u16, vocab); + defer gpa.free(got_logits); + u16s(got_logits, logits_raw); + const score_gap = gap(got_logits, host_logits); + std.debug.print("scores n {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} head {x:0>4}\n", .{ vocab, score_gap.off, score_gap.steps, score_gap.at, got_logits[score_gap.at], host_logits[score_gap.at], got_logits[0] }); + if (score_gap.off > 0 and score_gap.off <= 8) { + for (got_logits, host_logits, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) std.debug.print("score {d} gpu {x:0>4} host {x:0>4} steps {d}\n", .{ i, a, b, d }); + } + } + + if (norm_off != 0 or col_ulp > 8 or act_off != 0 or xs_ulp > 4 or mixed_gap.steps > 2 or score_gap.steps > 1 or score_gap.off > 3) return 1; + return 0; +} From 7f0794d7cce3e80c571bf30ce03007c1148b55d1 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 02:56:38 -0400 Subject: [PATCH 027/115] One drafts-off prompt matches the Python engine, eight greedy token ids exact. --- zig/build/cuda.zig | 9 + zig/src/families/flashnext/cuda_launch.zig | 32 + zig/src/families/flashnext/cuda_ple.zig | 74 + zig/src/families/flashnext/cuda_prompt.zig | 1208 +++++++++++++++++ .../families/flashnext/cuda_prompt_main.zig | 50 + 5 files changed, 1373 insertions(+) create mode 100644 zig/src/families/flashnext/cuda_prompt.zig create mode 100644 zig/src/families/flashnext/cuda_prompt_main.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 3b8e975af7..a49fa3a54d 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -155,6 +155,15 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const embed_token = b.option([]const u8, "flashnext_token", "token id for the embedding check") orelse "7"; if (embed_model != null and embed_kernels != null) embed_run.addArgs(&.{ embed_model.?, embed_kernels.?, embed_token }); b.step("flashnext-embed", "Upload the Flash Next embedding table and run _embed for one token").dependOn(&embed_run.step); + const prompt_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_prompt_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + prompt_check.addImport("cuda", cuda); + prompt_check.addImport("core", mods.core); + prompt_check.addImport("flashnext", mods.flashnext); + const prompt_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-prompt", .root_module = prompt_check })); + const prompt_ids = b.option([]const u8, "flashnext_ids", "file of prompt token ids for the Flash Next prompt check"); + const prompt_count = b.option([]const u8, "flashnext_count", "greedy tokens to emit") orelse "8"; + if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); + b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 376f942bf1..e442ed8567 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -13,6 +13,7 @@ pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); pub const mix = @import("cuda_mix.zig"); +pub const prompt = @import("cuda_prompt.zig"); const p = aot.ptr; @@ -193,6 +194,36 @@ pub const Tri = struct { }, &.{ ci("HEADS", heads), ci("DH", dh), aot.cf("SCALE", scale) }); } + /// glue.ple_gate: one program per row. + pub fn pleGate(t: Tri, keys: u64, vals: u64, h: u64, norm_key: u64, norm_query: u64, gated: u64, pss: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + try t.run("_ple_gate", .{ u(rows), 1, 1 }, &.{ + p("KEYS", "*bf16", keys), + p("VALS", "*bf16", vals), + p("H", "*bf16", h), + p("NK", "*fp32", norm_key), + p("NQ", "*fp32", norm_query), + p("GATED", "*bf16", gated), + p("PSS", "*fp32", pss), + aot.float("eps", eps), + }, &.{ ci("D", dims), ci("S", streams), ci("BLOCK", 512) }); + } + + /// glue.ple_conv. R folds to a constexpr when it is 1. + pub fn pleConv(t: Tri, gated: u64, pss: u64, norm_conv: u64, tail: u64, conv_w: u64, h: u64, hout: u64, nrow: u64, eps: f32, rows: usize, dims: usize, streams: usize, taps: usize, dilation: usize) !void { + try t.run("_ple_conv", .{ u(rows), u(streams * dims / 512), 1 }, &.{ + p("GATED", "*bf16", gated), + p("PSS", "*fp32", pss), + p("NC", "*fp32", norm_conv), + p("TAIL", "*bf16", tail), + p("CW", "*bf16", conv_w), + p("H", "*bf16", h), + p("HOUT", "*bf16", hout), + p("NROW", "*bf16", nrow), + aot.float("eps", eps), + int("R", rows), + }, &.{ ci("D", dims), ci("S", streams), ci("TAPS", taps), ci("DIL", dilation), ci("BLOCK", 512) }); + } + /// moe.router for one decode row. M folds to a constexpr. 513 rows: the experts, then the shared gate. pub fn router(t: Tri, x: u64, w: u64, out: u64, x_stride: usize) !void { try t.run("_router", .{ 1, cdiv(513, 32), 1 }, &.{ @@ -224,6 +255,7 @@ test { _ = attn; _ = ple; _ = mix; + _ = prompt; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index 3c2cf32267..751f8a8c75 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -348,6 +348,80 @@ fn readSpec(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !NgramSpec { }; } +/// The text config's eos id, which a fresh n-gram window repeats. +pub fn eosOf(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !i64 { + const cfg = try readSpec(gpa, io, model_dir); + return cfg.eos; +} + +/// Host dequant of one token's sixteen n-gram rows. `history` is the previous n - 1 ids. +pub fn embedding(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, history: []const i64, token: u32, out: []u16, xs: []f32) !void { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + if (history.len != hash.n - 1 or out.len != hash.heads * hash.dh or xs.len != out.len / 32) return error.UnexpectedTensor; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history, @intCast(token), ids[0..hash.heads]); + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + const scale = try tableScale(gpa, io, model_dir, wm); + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, out[h * dh ..][0..dh]); + if (scale != 1) { + for (out) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + sumGroups(out, xs); +} + /// One decode token. `ple_layer_ids` `[2]` is this layer. Returns 0 when the gathered rows match the host dequant. pub fn table(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, token: u32) !u32 { const cfg = try readSpec(gpa, io, model_dir); diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig new file mode 100644 index 0000000000..8406b8b0fb --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -0,0 +1,1208 @@ +//! One prompt, drafts off: each token is one fused decode row, layers in order, then the argmax. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); + +const layers_n: usize = 48; +const linear_n: usize = 36; +const attn_n: usize = 12; +const dims: usize = 2560; +const streams: usize = 4; +const wide: usize = streams * dims; +const low: usize = 320; +const eps: f32 = 1e-6; +const vocab: usize = 248320; +const experts_n: usize = 512; +const slots: usize = 11; +const moe_w: usize = 640; +const block_bytes: usize = 160 * 4; +const tile: usize = 16; +const q_heads: usize = 24; +const kv_heads: usize = 2; +const head_dim: usize = 256; +const index_heads: usize = 4; +const index_dim: usize = 128; +const half: usize = 32; +const proj_n: usize = 13952; +const out_k: usize = 6144; +const nch: usize = 3; +const ple_layer: usize = 1; +const inj_stride: usize = 16; +const wts_stride: usize = 48; +const ple_tail_n: usize = 9; +const state_n: usize = gdn.nv * gdn.dv * gdn.dk; +const hc_sk: usize = 32; +const out_sk: usize = 8; + +const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; +const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + +const Tensor = core.safetensors.Tensor; + +const Held = struct { rel: []u8, file: Shard }; + +const Shard = struct { + io: std.Io, + file: std.Io.File, + map: std.Io.File.MemoryMap, + at: usize, + arena: std.heap.ArenaAllocator, + names: std.StringHashMapUnmanaged(core.safetensors.Entry), + + fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Shard { + var file = try std.Io.Dir.cwd().openFile(io, path, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + var names: std.StringHashMapUnmanaged(core.safetensors.Entry) = .empty; + const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), map.memory[8..][0..header_len], .{}); + defer parsed.deinit(); + var it = parsed.value.object.iterator(); + while (it.next()) |kv| { + if (std.mem.eql(u8, kv.key_ptr.*, "__metadata__")) continue; + const o = kv.value_ptr.object; + const dtype = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + const shape = o.get("shape").?.array.items; + if (shape.len > core.safetensors.max_rank) continue; + var e: core.safetensors.Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = 0 }; + var n: usize = dtype.size(); + for (shape, 0..) |d, i| { + e.shape[i] = @intCast(d.integer); + n *= e.shape[i]; + } + const offs = o.get("data_offsets").?.array.items; + e.begin = @intCast(offs[0].integer); + e.end = @intCast(offs[1].integer); + if (e.end < e.begin or e.end - e.begin != n or e.end > len - 8 - header_len) return error.BadSafetensors; + try names.put(arena.allocator(), try arena.allocator().dupe(u8, kv.key_ptr.*), e); + } + return .{ .io = io, .file = file, .map = map, .at = 8 + header_len, .arena = arena, .names = names }; + } + + fn close(self: *Shard) void { + self.map.destroy(self.io); + self.file.close(self.io); + self.arena.deinit(); + self.* = undefined; + } + + fn get(self: *const Shard, name: []const u8) ?Tensor { + const e = self.names.get(name) orelse return null; + return .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .bytes = self.map.memory[self.at + e.begin .. self.at + e.end] }; + } +}; + +const Store = struct { + gpa: std.mem.Allocator, + io: std.Io, + dir: []const u8, + parsed: std.json.Parsed(std.json.Value), + path: []u8 = &.{}, + rel: []u8 = &.{}, + file: ?Shard = null, + held: [12]Held = undefined, + nheld: usize = 0, + + fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}) }; + } + + fn deinit(self: *Store) void { + self.release(); + self.closeFile(); + self.parsed.deinit(); + } + + fn release(self: *Store) void { + for (self.held[0..self.nheld]) |*h| { + h.file.close(); + self.gpa.free(h.rel); + } + self.nheld = 0; + } + + fn hold(self: *Store, key: []const u8) !Tensor { + const file = try self.shard(key); + for (self.held[0..self.nheld]) |*h| { + if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; + } + if (self.nheld == self.held.len) return error.UnexpectedTensor; + const rel = try self.gpa.dupe(u8, file); + errdefer self.gpa.free(rel); + const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + defer self.gpa.free(path); + self.held[self.nheld] = .{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }; + self.nheld += 1; + return self.held[self.nheld - 1].file.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + fn closeFile(self: *Store) void { + if (self.file) |*f| f.close(); + self.file = null; + if (self.path.len != 0) self.gpa.free(self.path); + self.path = &.{}; + if (self.rel.len != 0) self.gpa.free(self.rel); + self.rel = &.{}; + } + + fn shard(self: *Store, key: []const u8) ![]const u8 { + const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + const v = wm.object.get(key) orelse { + std.debug.print("missing {s}\n", .{key}); + return error.MissingTensor; + }; + return v.string; + } + + fn use(self: *Store, key: []const u8) !void { + const file = try self.shard(key); + if (self.file != null and std.mem.eql(u8, file, self.rel)) return; + self.closeFile(); + self.rel = try self.gpa.dupe(u8, file); + self.path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + self.file = try Shard.load(self.gpa, self.io, self.path); + } + + fn tensor(self: *Store, key: []const u8) !Tensor { + try self.use(key); + return self.file.?.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + fn copyOf(self: *Store, key: []const u8) ![]u8 { + const t = try self.tensor(key); + const out = try self.gpa.alloc(u8, t.bytes.len); + @memcpy(out, t.bytes); + return out; + } +}; + +const Face = struct { + w: []const u8, + s: []const u8, + b: []const u8, + n: usize, + k: usize, + routed: bool, +}; + +const Q4 = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + n: usize, + k: usize, + + fn free(self: *Q4) void { + self.w.free(); + self.s.free(); + self.b.free(); + } + + fn empty(d: *cuda.Driver) Q4 { + const z: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + return .{ .w = z, .s = z, .b = z, .n = 0, .k = 0 }; + } +}; + +const HcW = struct { + scale: cuda.DeviceBuffer, + down: Q4, + up: Q4, + n_down: usize, + has_inj: usize, + + fn free(self: *HcW) void { + self.scale.free(); + self.down.free(); + self.up.free(); + } +}; + +fn linearLayer(layer: usize) bool { + return layer % 4 != 3; +} + +fn linearIndex(layer: usize) usize { + return layer - layer / 4; +} + +fn attnIndex(layer: usize) usize { + return layer / 4; +} + +fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { + var n: usize = 0; + for (parts) |p| n += p.len; + const out = try gpa.alloc(u8, n); + var at: usize = 0; + for (parts) |p| { + @memcpy(out[at..][0..p.len], p); + at += p.len; + } + return out; +} + +fn asF32(gpa: std.mem.Allocator, t: Tensor) ![]f32 { + if (t.rank != 1) return error.UnexpectedTensor; + const out = try gpa.alloc(f32, t.dim(0)); + if (t.dtype == .f32 and t.bytes.len == out.len * 4) { + for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); + } else if (t.dtype == .bf16 and t.bytes.len == out.len * 2) { + for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)); + } else return error.UnexpectedTensor; + return out; +} + +fn uploadQ4(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + var lane = try qmm.pack(gpa, words, scales, biases, n, k8); + defer lane.deinit(gpa); + var w = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +fn uploadTiled(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + const kg = k8 / 4; + const tiled = try hc.tileWords(gpa, words, n, k8); + defer gpa.free(tiled); + const ts = try hc.transposeBf16(gpa, scales, n, kg); + defer gpa.free(ts); + const tb = try hc.transposeBf16(gpa, biases, n, kg); + defer gpa.free(tb); + var w = try cuda.DeviceBuffer.fromHost(driver, tiled); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, ts); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, tb); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +fn join3(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8) !struct { w: []u8, s: []u8, b: []u8 } { + const w = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + errdefer gpa.free(w); + const s = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + errdefer gpa.free(s); + const b = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + return .{ .w = w, .s = s, .b = b }; +} + +fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { + var buf: [180]u8 = undefined; + const down = try join3(gpa, store, &buf, prefix, ".input_mix_weight_down"); + defer gpa.free(down.w); + defer gpa.free(down.s); + defer gpa.free(down.b); + const up = try join3(gpa, store, &buf, prefix, ".input_mix_weight_up"); + defer gpa.free(up.w); + defer gpa.free(up.s); + defer gpa.free(up.b); + const gamma = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix})); + const scale = try asF32(gpa, gamma); + defer gpa.free(scale); + if (scale.len != wide) return error.UnexpectedTensor; + var n: usize = low; + var words: []u8 = undefined; + var scales: []u8 = undefined; + var biases: []u8 = undefined; + var inj_w: []u8 = &.{}; + var inj_s: []u8 = &.{}; + var inj_b: []u8 = &.{}; + if (inject) { + const inj = try join3(gpa, store, &buf, prefix, ".block_inject_weight"); + inj_w = inj.w; + inj_s = inj.s; + inj_b = inj.b; + n = low + 4; + words = try cat(gpa, &.{ down.w, inj.w }); + scales = try cat(gpa, &.{ down.s, inj.s }); + biases = try cat(gpa, &.{ down.b, inj.b }); + } else { + words = try gpa.dupe(u8, down.w); + scales = try gpa.dupe(u8, down.s); + biases = try gpa.dupe(u8, down.b); + } + defer gpa.free(words); + defer gpa.free(scales); + defer gpa.free(biases); + defer if (inject) gpa.free(inj_w); + defer if (inject) gpa.free(inj_s); + defer if (inject) gpa.free(inj_b); + const k8 = wide / 8; + if (words.len != n * k8 * 4) return error.UnexpectedTensor; + var down_q = try uploadTiled(gpa, driver, words, scales, biases, n, k8); + errdefer down_q.free(); + var up_q = try uploadTiled(gpa, driver, up.w, up.s, up.b, wide, low / 8); + errdefer up_q.free(); + const scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + return .{ .scale = scale_b, .down = down_q, .up = up_q, .n_down = n, .has_inj = @intFromBool(inject) }; +} + +fn stitch(gpa: std.mem.Allocator, store: *Store, prefix: []const u8, names: []const []const u8) !struct { w: []u8, s: []u8, b: []u8, n: usize } { + var buf: [180]u8 = undefined; + var ws: [6][]u8 = undefined; + var ss: [6][]u8 = undefined; + var bs: [6][]u8 = undefined; + var n: usize = 0; + for (names, 0..) |part, i| { + const got = try join3(gpa, store, &buf, prefix, part); + ws[i] = got.w; + ss[i] = got.s; + bs[i] = got.b; + n += 1; + } + defer for (0..n) |i| { + gpa.free(ws[i]); + gpa.free(ss[i]); + gpa.free(bs[i]); + }; + return .{ .w = try cat(gpa, ws[0..n]), .s = try cat(gpa, ss[0..n]), .b = try cat(gpa, bs[0..n]), .n = n }; +} + +const Scratch = struct { + h: cuda.DeviceBuffer, + y: cuda.DeviceBuffer, + wts: cuda.DeviceBuffer, + inj: cuda.DeviceBuffer, + pss: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + dn: cuda.DeviceBuffer, + hc_part: cuda.DeviceBuffer, + act: cuda.DeviceBuffer, + xs_act: cuda.DeviceBuffer, + mixed: cuda.DeviceBuffer, + xsm: cuda.DeviceBuffer, + proj: cuda.DeviceBuffer, + part: cuda.DeviceBuffer, + gout: cuda.DeviceBuffer, + gxs: cuda.DeviceBuffer, + st: cuda.DeviceBuffer, + ple_emb: cuda.DeviceBuffer, + ple_xs: cuda.DeviceBuffer, + ple_keys: cuda.DeviceBuffer, + ple_vals: cuda.DeviceBuffer, + ple_gated: cuda.DeviceBuffer, + ple_pss: cuda.DeviceBuffer, + ple_nrow: cuda.DeviceBuffer, + ple_part: cuda.DeviceBuffer, + pa: cuda.DeviceBuffer, + q: cuda.DeviceBuffer, + iq: cuda.DeviceBuffer, + po: cuda.DeviceBuffer, + pm: cuda.DeviceBuffer, + pl: cuda.DeviceBuffer, + ao: cuda.DeviceBuffer, + gated: cuda.DeviceBuffer, + gxs_attn: cuda.DeviceBuffer, + dummy: cuda.DeviceBuffer, + inv: cuda.DeviceBuffer, + pos: cuda.DeviceBuffer, + ids: cuda.DeviceBuffer, + logits: cuda.DeviceBuffer, + router_l: cuda.DeviceBuffer, + picks: cuda.DeviceBuffer, + + fn alloc(driver: *cuda.Driver, rows: usize) !Scratch { + var s: Scratch = undefined; + s.h = try cuda.DeviceBuffer.alloc(driver, rows * wide * 2); + s.y = try cuda.DeviceBuffer.alloc(driver, rows * slots * dims * 4); + s.wts = try cuda.DeviceBuffer.alloc(driver, rows * wts_stride); + s.inj = try cuda.DeviceBuffer.alloc(driver, rows * inj_stride); + s.pss = try cuda.DeviceBuffer.alloc(driver, (dims / 256) * streams * 4); + s.norm = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.dn = try cuda.DeviceBuffer.alloc(driver, (low + 4) * 2); + s.hc_part = try cuda.DeviceBuffer.alloc(driver, hc_sk * (low + 4) * 4); + s.act = try cuda.DeviceBuffer.alloc(driver, low * 2); + s.xs_act = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + s.mixed = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.xsm = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.proj = try cuda.DeviceBuffer.alloc(driver, gdn.proj_width * 2); + s.part = try cuda.DeviceBuffer.alloc(driver, out_sk * dims * 4); + s.gout = try cuda.DeviceBuffer.alloc(driver, gdn.value_dim * 2); + s.gxs = try cuda.DeviceBuffer.alloc(driver, (gdn.value_dim / 32) * 4); + s.st = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + s.ple_emb = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_xs = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.ple_keys = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_vals = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_gated = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_pss = try cuda.DeviceBuffer.alloc(driver, streams * 4); + s.ple_nrow = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_part = try cuda.DeviceBuffer.alloc(driver, 4 * dims * 4); + s.pa = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); + s.q = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); + s.iq = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); + s.po = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); + s.pm = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.pl = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.ao = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gated = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gxs_attn = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); + s.dummy = try cuda.DeviceBuffer.alloc(driver, 4096); + var inv: [half]f32 = undefined; + for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); + s.inv = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); + s.pos = try cuda.DeviceBuffer.alloc(driver, 4); + s.ids = try cuda.DeviceBuffer.alloc(driver, 4); + s.logits = try cuda.DeviceBuffer.alloc(driver, vocab * 2); + s.router_l = try cuda.DeviceBuffer.alloc(driver, (experts_n + 1) * 4); + s.picks = try cuda.DeviceBuffer.alloc(driver, slots * 4); + return s; + } + + fn free(self: *Scratch) void { + self.h.free(); + self.y.free(); + self.wts.free(); + self.inj.free(); + self.pss.free(); + self.norm.free(); + self.dn.free(); + self.hc_part.free(); + self.act.free(); + self.xs_act.free(); + self.mixed.free(); + self.xsm.free(); + self.proj.free(); + self.part.free(); + self.gout.free(); + self.gxs.free(); + self.st.free(); + self.ple_emb.free(); + self.ple_xs.free(); + self.ple_keys.free(); + self.ple_vals.free(); + self.ple_gated.free(); + self.ple_pss.free(); + self.ple_nrow.free(); + self.ple_part.free(); + self.pa.free(); + self.q.free(); + self.iq.free(); + self.po.free(); + self.pm.free(); + self.pl.free(); + self.ao.free(); + self.gated.free(); + self.gxs_attn.free(); + self.dummy.free(); + self.inv.free(); + self.pos.free(); + self.ids.free(); + self.logits.free(); + self.router_l.free(); + self.picks.free(); + } +}; + +fn writeback(tri: anytype, s: *Scratch, h: u64, mode: usize, inj: u64, y: u64, wts: u64) !void { + const dummy = s.dummy.ptr; + if (mode == 0) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 0, 1, 1, 1); + } else if (mode == 2) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", dims, 1, dims, streams, 2, 10, 11, 1); + } else if (mode == 4) { + try tri.hcWriteback(h, h, s.pss.ptr, s.part.ptr, "*fp32", inj, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 4, 1, 1, out_sk); + } else return error.UnexpectedTensor; +} + +fn readout(tri: anytype, s: *Scratch, hc_w: HcW, h: u64, inj: u64) !void { + try tri.hcDown(h, s.pss.ptr, hc_w.scale.ptr, s.norm.ptr, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn.ptr, s.hc_part.ptr, eps, 1, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); + const inj_ptr = if (hc_w.has_inj == 1) inj else s.act.ptr; + try tri.hcReduceAct(s.hc_part.ptr, s.act.ptr, s.xs_act.ptr, inj_ptr, hc_sk, 1, streams, low, hc_w.n_down, hc_w.has_inj); + try tri.hcUpmix(s.act.ptr, s.xs_act.ptr, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.norm.ptr, s.mixed.ptr, s.xsm.ptr, 1, wide, low, dims, streams); +} + +fn shiftTail(gpa: std.mem.Allocator, stream: *cuda.Stream, window: cuda.DeviceBuffer, taps_n: usize, row_bytes: usize, fresh: []const u8) !void { + const bytes = taps_n * row_bytes; + const raw = try gpa.alloc(u8, bytes); + defer gpa.free(raw); + try stream.synchronize(); + try window.download(0, raw); + const next = try gpa.alloc(u8, bytes); + defer gpa.free(next); + const keep = taps_n - 1; + @memcpy(next[0 .. keep * row_bytes], raw[row_bytes..]); + @memcpy(next[keep * row_bytes ..], fresh); + try window.upload(0, next); +} + +fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s: *Scratch, proj_q: Q4, out_q: Q4, conv: cuda.DeviceBuffer, cw: u64, a_log: u64, dt: u64, norm: u64, state: *cuda.DeviceBuffer) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.proj.ptr, 1, proj_q.n, proj_q.k); + try gdn.launch(driver, stream.*, s.proj.ptr, conv.ptr, cw, state.ptr, a_log, dt, norm, eps, s.gout.ptr, s.gxs.ptr, s.st.ptr); + try state.copyFrom(0, s.st.ptr, state_n * 4, stream.handle); + const row = try gpa.alloc(u8, gdn.conv_dim * 2); + defer gpa.free(row); + try stream.synchronize(); + try s.proj.download(0, row); + try shiftTail(gpa, stream, conv, gdn.taps - 1, gdn.conv_dim * 2, row); + try qmm.partials(driver, stream.*, s.gout.ptr, s.gxs.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, gdn.value_dim, out_sk); +} + +fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); + try s.pos.upload(0, std.mem.asBytes(&pos)); + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); + try tri.attnChunks(s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, 1); + try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads); + try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); + try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); +} + +fn copyExpert(dst_w: []u8, dst_s: []u8, dst_b: []u8, routed: Face, shared: Face, id: i32, slot: usize) void { + const p = if (id == experts_n) shared else routed; + const k8 = p.k / 8; + const kg = p.k / 32; + const wb = p.n * k8 * 4; + const sb = p.n * kg * 2; + const at: usize = if (id == experts_n) 0 else @intCast(id); + @memcpy(dst_w[slot * wb ..][0..wb], p.w[at * wb ..][0..wb]); + @memcpy(dst_s[slot * sb ..][0..sb], p.s[at * sb ..][0..sb]); + @memcpy(dst_b[slot * sb ..][0..sb], p.b[at * sb ..][0..sb]); +} + +fn packDevice(driver: *cuda.Driver, stream: cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) !cuda.DeviceBuffer { + const kg = k / 32; + const nb = n / 32; + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, e * nb * kg * block_bytes); + errdefer out_b.free(); + var a: cuda.Args = .{}; + a.add(w_b.ptr); + a.add(s_b.ptr); + a.add(b_b.ptr); + a.add(out_b.ptr); + a.add(cint(n)); + a.add(cint(k / 8)); + a.add(cint(kg)); + a.add(cint(nb)); + try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream, &a); + try stream.synchronize(); + return out_b; +} + +fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { + if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; + const nblk = gate.len / block_bytes; + const out = try gpa.alloc(u8, gate.len + up.len); + for (0..nblk) |b| { + @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); + @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); + } + return out; +} + +fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(cint(x_stride)); + a.add(cint(slot_n)); + a.add(w); + a.add(cint(kg)); + a.add(cint(nb)); + a.add(items); + a.add(counts); + a.add(members); + a.add(out); + a.add(cint(n)); + a.add(@as(f32, 0)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((units + 3) / 4) }, .block = .{ .x = 128 } }, stream, &a); +} + +fn moeStep(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, router_w: u64, faces: [3]Face, shared: [3]Face, y: u64, wts: u64) !void { + try tri.router(s.mixed.ptr, router_w, s.router_l.ptr, dims); + try tri.topkRows(s.router_l.ptr, s.picks.ptr, wts); + var pick_raw: [slots * 4]u8 = undefined; + try stream.synchronize(); + try s.picks.download(0, &pick_raw); + var packed_ids: [slots]i32 = undefined; + var npack: usize = 0; + var local: [slots]i32 = undefined; + for (0..slots) |slot| { + const id = std.mem.readInt(i32, pick_raw[4 * slot ..][0..4], .little); + if (id < 0 or id > experts_n) return error.UnexpectedTensor; + var found: ?usize = null; + for (packed_ids[0..npack], 0..) |prev, i| if (prev == id) { + found = i; + }; + if (found) |i| { + local[slot] = @intCast(i); + } else { + packed_ids[npack] = id; + local[slot] = @intCast(npack); + npack += 1; + } + } + const gw = npack * moe_w * (dims / 8) * 4; + const gs = npack * moe_w * (dims / 32) * 2; + const dw = npack * dims * (moe_w / 8) * 4; + const ds = npack * dims * (moe_w / 32) * 2; + const gathered_w = try gpa.alloc(u8, gw); + defer gpa.free(gathered_w); + const gathered_s = try gpa.alloc(u8, gs); + defer gpa.free(gathered_s); + const gathered_b = try gpa.alloc(u8, gs); + defer gpa.free(gathered_b); + const up_w = try gpa.alloc(u8, gw); + defer gpa.free(up_w); + const up_s = try gpa.alloc(u8, gs); + defer gpa.free(up_s); + const up_b = try gpa.alloc(u8, gs); + defer gpa.free(up_b); + const down_w = try gpa.alloc(u8, dw); + defer gpa.free(down_w); + const down_s = try gpa.alloc(u8, ds); + defer gpa.free(down_s); + const down_b = try gpa.alloc(u8, ds); + defer gpa.free(down_b); + for (packed_ids[0..npack], 0..) |id, slot| { + copyExpert(gathered_w, gathered_s, gathered_b, faces[0], shared[0], id, slot); + copyExpert(up_w, up_s, up_b, faces[1], shared[1], id, slot); + copyExpert(down_w, down_s, down_b, faces[2], shared[2], id, slot); + } + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + var exp_mod = try cuda.Module.load(driver, cuda.kernels.experts); + defer exp_mod.unload(); + const pack_fn = try pack_mod.function(pack_symbol); + var packed_g = try packDevice(driver, stream.*, pack_fn, gathered_w, gathered_s, gathered_b, npack, moe_w, dims); + defer packed_g.free(); + var packed_u = try packDevice(driver, stream.*, pack_fn, up_w, up_s, up_b, npack, moe_w, dims); + defer packed_u.free(); + var packed_d = try packDevice(driver, stream.*, pack_fn, down_w, down_s, down_b, npack, dims, moe_w); + defer packed_d.free(); + const g_raw = try gpa.alloc(u8, packed_g.len); + defer gpa.free(g_raw); + const u_raw = try gpa.alloc(u8, packed_u.len); + defer gpa.free(u_raw); + try packed_g.download(0, g_raw); + try packed_u.download(0, u_raw); + const stacked = try stackHalves(gpa, g_raw, u_raw); + defer gpa.free(stacked); + var up_wb = try cuda.DeviceBuffer.fromHost(driver, stacked); + defer up_wb.free(); + var local_pick: [slots]i32 = undefined; + for (0..slots) |slot| local_pick[slot] = local[slot]; + var picks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&local_pick)); + defer picks_b.free(); + var members_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); + defer members_b.free(); + var items_b = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); + defer items_b.free(); + var counts_b = try cuda.DeviceBuffer.alloc(driver, 8); + defer counts_b.free(); + var act_b = try cuda.DeviceBuffer.alloc(driver, slots * moe_w * 2); + defer act_b.free(); + const plan_fn = try exp_mod.function(plan_symbol); + var plan_a: cuda.Args = .{}; + plan_a.add(picks_b.ptr); + plan_a.add(cint(slots)); + plan_a.add(cint(npack)); + plan_a.add(cint(tile)); + plan_a.add(members_b.ptr); + plan_a.add(items_b.ptr); + plan_a.add(counts_b.ptr); + try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); + const up_fn = try exp_mod.function(up_symbol); + const down_fn = try exp_mod.function(down_symbol); + try launchExpert(up_fn, stream.*, s.mixed.ptr, dims, slots, up_wb.ptr, dims / 32, moe_w / 32, items_b.ptr, counts_b.ptr, members_b.ptr, act_b.ptr, moe_w, slots * (moe_w / 32)); + try launchExpert(down_fn, stream.*, act_b.ptr, moe_w, 0, packed_d.ptr, moe_w / 32, dims / 32, items_b.ptr, counts_b.ptr, members_b.ptr, y, dims, slots * (dims / 32)); + try stream.synchronize(); +} + +fn faceOf(store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !Face { + const w = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + const sc = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + const b = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + const k8 = k / 8; + const kg = k / 32; + if (routed) { + if (!w.is(.u32, &.{ experts_n, n, k8 }) or !sc.is(.bf16, &.{ experts_n, n, kg }) or !b.is(.bf16, &.{ experts_n, n, kg })) return error.UnexpectedTensor; + } else if (!w.is(.u32, &.{ n, k8 }) or !sc.is(.bf16, &.{ n, kg }) or !b.is(.bf16, &.{ n, kg })) return error.UnexpectedTensor; + return .{ .w = w.bytes, .s = sc.bytes, .b = b.bytes, .n = n, .k = k, .routed = routed }; +} + +fn reduceBf16(gpa: std.mem.Allocator, stream: *cuda.Stream, part: cuda.DeviceBuffer, sk: usize, n: usize, dst: cuda.DeviceBuffer) !void { + const raw = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(raw); + try stream.synchronize(); + try part.download(0, raw); + const row = try gpa.alloc(u16, n); + defer gpa.free(row); + for (0..n) |i| { + var acc: f32 = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); + for (1..sk) |s| acc += @as(f32, @bitCast(std.mem.readInt(u32, raw[4 * (s * n + i) ..][0..4], .little))); + row[i] = toBf16(acc); + } + try dst.upload(0, std.mem.sliceAsBytes(row)); +} + +fn pleStep(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, model_dir: []const u8, history: *[2]i64, token: u32, h: u64, key_q: Q4, val_q: Q4, nk: u64, nq: u64, nc: u64, cw: u64, tail: *cuda.DeviceBuffer) !void { + const emb = try gpa.alloc(u16, dims); + defer gpa.free(emb); + const xs = try gpa.alloc(f32, dims / 32); + defer gpa.free(xs); + try ple.embedding(gpa, io, model_dir, history, token, emb, xs); + try s.ple_emb.upload(0, std.mem.sliceAsBytes(emb)); + try s.ple_xs.upload(0, std.mem.sliceAsBytes(xs)); + try qmm.matmul(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, key_q.w.ptr, key_q.s.ptr, key_q.b.ptr, s.ple_keys.ptr, 1, wide, dims); + try qmm.partials(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, val_q.w.ptr, val_q.s.ptr, val_q.b.ptr, s.ple_vals.ptr, s.ple_part.ptr, 1, dims, dims, 4); + try reduceBf16(gpa, stream, s.ple_part, 4, dims, s.ple_vals); + try tri.pleGate(s.ple_keys.ptr, s.ple_vals.ptr, h, nk, nq, s.ple_gated.ptr, s.ple_pss.ptr, eps, 1, dims, streams); + try tri.pleConv(s.ple_gated.ptr, s.ple_pss.ptr, nc, tail.ptr, cw, h, h, s.ple_nrow.ptr, eps, 1, dims, streams, 4, 3); + const nrow = try gpa.alloc(u8, wide * 2); + defer gpa.free(nrow); + try stream.synchronize(); + try s.ple_nrow.download(0, nrow); + try shiftTail(gpa, stream, tail.*, ple_tail_n, wide * 2, nrow); + history[0] = history[1]; + history[1] = token; +} + +fn argmax(raw: []const u8) u32 { + var best = promote(std.mem.readInt(u16, raw[0..2], .little)); + var at: usize = 0; + var i: usize = 1; + while (i < vocab) : (i += 1) { + const v = promote(std.mem.readInt(u16, raw[2 * i ..][0..2], .little)); + if (v > best) { + best = v; + at = i; + } + } + return @intCast(at); +} + +const Pass = struct { + gpa: std.mem.Allocator, + io: std.Io, + driver: *cuda.Driver, + stream: *cuda.Stream, + store: *Store, + model_dir: []const u8, + s: *Scratch, + table: embed.Table, + rows: usize, + mode: []u8, + history: *[2]i64, + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + mixer: HcW, + head_q: Q4, +}; + +fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { + const id: i32 = @intCast(token); + try p.s.ids.upload(0, std.mem.asBytes(&id)); + try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); +} + +fn oneLayer(comptime Tri: type, p: *Pass, tri: Tri, layer: usize, ids: []const u32, pos0: usize) !void { + var buf: [200]u8 = undefined; + const base = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{layer}); + var base_buf: [80]u8 = undefined; + @memcpy(base_buf[0..base.len], base); + const root = base_buf[0..base.len]; + var attn_hc = try loadHc(p.gpa, p.driver, p.store, try std.fmt.bufPrint(&buf, "{s}.attn_hyper_connection", .{root}), true); + defer attn_hc.free(); + var mlp_hc = try loadHc(p.gpa, p.driver, p.store, try std.fmt.bufPrint(&buf, "{s}.mlp_hyper_connection", .{root}), true); + defer mlp_hc.free(); + const mlp = try std.fmt.bufPrint(&buf, "{s}.mlp", .{root}); + var mlp_buf: [96]u8 = undefined; + @memcpy(mlp_buf[0..mlp.len], mlp); + const mlp_root = mlp_buf[0..mlp.len]; + const gate = try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.gate.weight", .{mlp_root})); + if (!gate.is(.bf16, &.{ experts_n, dims })) return error.UnexpectedTensor; + const gate_copy = try p.gpa.alloc(u8, gate.bytes.len); + defer p.gpa.free(gate_copy); + @memcpy(gate_copy, gate.bytes); + const sg = try join3(p.gpa, p.store, &buf, mlp_root, ".shared_expert_gate"); + defer p.gpa.free(sg.w); + defer p.gpa.free(sg.s); + defer p.gpa.free(sg.b); + const router = try p.gpa.alloc(u16, (experts_n + 1) * dims); + defer p.gpa.free(router); + @memcpy(std.mem.sliceAsBytes(router[0 .. experts_n * dims]), gate_copy); + try embed.dequant(sg.w, sg.s, sg.b, dims, 0, router[experts_n * dims ..][0..dims]); + var router_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(router)); + defer router_b.free(); + + const linear = linearLayer(layer); + var proj_q: Q4 = undefined; + var out_q: Q4 = undefined; + var cw_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var a_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var dt_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var nw_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var qs_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var ks_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var is_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var key_q = Q4.empty(p.driver); + var val_q = Q4.empty(p.driver); + var nk_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var nq_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var nc_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + var pcw_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; + defer if (cw_b.ptr != 0) cw_b.free(); + defer if (a_b.ptr != 0) a_b.free(); + defer if (dt_b.ptr != 0) dt_b.free(); + defer if (nw_b.ptr != 0) nw_b.free(); + defer if (qs_b.ptr != 0) qs_b.free(); + defer if (ks_b.ptr != 0) ks_b.free(); + defer if (is_b.ptr != 0) is_b.free(); + defer if (nk_b.ptr != 0) nk_b.free(); + defer if (nq_b.ptr != 0) nq_b.free(); + defer if (nc_b.ptr != 0) nc_b.free(); + defer if (pcw_b.ptr != 0) pcw_b.free(); + if (linear) { + const pre = try std.fmt.bufPrint(&buf, "{s}.linear_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const lp = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".in_proj_qkv", ".in_proj_z", ".in_proj_b", ".in_proj_a" }; + const stitched = try stitch(p.gpa, p.store, lp, &parts); + defer p.gpa.free(stitched.w); + defer p.gpa.free(stitched.s); + defer p.gpa.free(stitched.b); + if (stitched.w.len != gdn.proj_width * (dims / 8) * 4) return error.UnexpectedTensor; + proj_q = try uploadQ4(p.gpa, p.driver, stitched.w, stitched.s, stitched.b, gdn.proj_width, dims / 8); + const out = try join3(p.gpa, p.store, &buf, lp, ".out_proj"); + defer p.gpa.free(out.w); + defer p.gpa.free(out.s); + defer p.gpa.free(out.b); + out_q = try uploadQ4(p.gpa, p.driver, out.w, out.s, out.b, dims, gdn.value_dim / 8); + const conv = try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); + if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; + cw_b = try cuda.DeviceBuffer.fromHost(p.driver, conv.bytes); + const a_log = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); + defer p.gpa.free(a_log); + const dt = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.dt_bias", .{lp}))); + defer p.gpa.free(dt); + if (a_log.len != gdn.nv or dt.len != gdn.nv) return error.UnexpectedTensor; + a_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(a_log)); + dt_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(dt)); + const norm = try p.store.copyOf(try std.fmt.bufPrint(&buf, "{s}.norm.weight", .{lp})); + defer p.gpa.free(norm); + nw_b = try cuda.DeviceBuffer.fromHost(p.driver, norm); + } else { + const pre = try std.fmt.bufPrint(&buf, "{s}.self_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const ap = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; + const stitched = try stitch(p.gpa, p.store, ap, &parts); + defer p.gpa.free(stitched.w); + defer p.gpa.free(stitched.s); + defer p.gpa.free(stitched.b); + if (stitched.w.len != proj_n * (dims / 8) * 4) return error.UnexpectedTensor; + proj_q = try uploadQ4(p.gpa, p.driver, stitched.w, stitched.s, stitched.b, proj_n, dims / 8); + const out = try join3(p.gpa, p.store, &buf, ap, ".o_proj"); + defer p.gpa.free(out.w); + defer p.gpa.free(out.s); + defer p.gpa.free(out.b); + out_q = try uploadQ4(p.gpa, p.driver, out.w, out.s, out.b, dims, out_k / 8); + const qn = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); + defer p.gpa.free(qn); + const kn = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); + defer p.gpa.free(kn); + const iqn = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); + defer p.gpa.free(iqn); + if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim) return error.UnexpectedTensor; + qs_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(qn)); + ks_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(kn)); + is_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(iqn)); + } + defer proj_q.free(); + defer out_q.free(); + if (layer == ple_layer) { + const pre = try std.fmt.bufPrint(&buf, "{s}.ple", .{root}); + var pre_buf: [96]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const pp = pre_buf[0..pre.len]; + const key = try join3(p.gpa, p.store, &buf, pp, ".key_proj"); + defer p.gpa.free(key.w); + defer p.gpa.free(key.s); + defer p.gpa.free(key.b); + const val = try join3(p.gpa, p.store, &buf, pp, ".value_proj"); + defer p.gpa.free(val.w); + defer p.gpa.free(val.s); + defer p.gpa.free(val.b); + key_q = try uploadQ4(p.gpa, p.driver, key.w, key.s, key.b, wide, dims / 8); + val_q = try uploadQ4(p.gpa, p.driver, val.w, val.s, val.b, dims, dims / 8); + const nk = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); + defer p.gpa.free(nk); + const nq = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); + defer p.gpa.free(nq); + const nc = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_conv.weight", .{pp}))); + defer p.gpa.free(nc); + if (nk.len != wide or nq.len != wide or nc.len != wide) return error.UnexpectedTensor; + nk_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(nk)); + nq_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(nq)); + nc_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(nc)); + const conv = try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); + if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; + pcw_b = try cuda.DeviceBuffer.fromHost(p.driver, conv.bytes); + } + defer key_q.free(); + defer val_q.free(); + + // Shared experts are copied. Routed rows stay mapped, so this is the last shard switch. + var sh_g = try ownedFace(p.gpa, p.store, &buf, mlp_root, ".shared_expert.gate_proj", moe_w, dims); + defer sh_g.free(p.gpa); + var sh_u = try ownedFace(p.gpa, p.store, &buf, mlp_root, ".shared_expert.up_proj", moe_w, dims); + defer sh_u.free(p.gpa); + var sh_d = try ownedFace(p.gpa, p.store, &buf, mlp_root, ".shared_expert.down_proj", dims, moe_w); + defer sh_d.free(p.gpa); + p.store.release(); + var faces: [3]Face = undefined; + faces[0] = try faceOf(p.store, &buf, mlp_root, ".switch_mlp.gate_proj", true, moe_w, dims); + faces[1] = try faceOf(p.store, &buf, mlp_root, ".switch_mlp.up_proj", true, moe_w, dims); + faces[2] = try faceOf(p.store, &buf, mlp_root, ".switch_mlp.down_proj", true, dims, moe_w); + const sh = [3]Face{ sh_g.face, sh_u.face, sh_d.face }; + + std.debug.print("layer {d}\n", .{layer}); + for (ids, 0..) |token, t| { + const h = p.s.h.ptr + @as(u64, t) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, t) * inj_stride; + const y = p.s.y.ptr + @as(u64, t) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, t) * wts_stride; + if (layer == ple_layer and p.mode[t] == 2) { + try writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } + if (layer == ple_layer) try pleStep(Tri, p.gpa, p.io, p.driver, p.stream, tri, p.s, p.model_dir, p.history, token, h, key_q, val_q, nk_b.ptr, nq_b.ptr, nc_b.ptr, pcw_b.ptr, &p.tail); + if (p.mode[t] == 2) { + try writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } else try writeback(tri, p.s, h, 0, inj, y, wt); + try readout(tri, p.s, attn_hc, h, inj); + if (linear) { + const li = linearIndex(layer); + try gdnStep(p.gpa, p.driver, p.stream, p.s, proj_q, out_q, p.conv[li], cw_b.ptr, a_b.ptr, dt_b.ptr, nw_b.ptr, &p.state[li]); + } else { + const ai = attnIndex(layer); + try attnStep(Tri, p.driver, p.stream, tri, p.s, proj_q, out_q, qs_b.ptr, ks_b.ptr, is_b.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); + } + try writeback(tri, p.s, h, 4, inj, y, wt); + try readout(tri, p.s, mlp_hc, h, inj); + try moeStep(Tri, p.gpa, p.driver, p.stream, tri, p.s, router_b.ptr, faces, sh, y, wt); + p.mode[t] = 2; + if (t % 16 == 0) std.debug.print("layer {d} row {d}\n", .{ layer, t }); + } + p.store.release(); +} + +const OwnedFace = struct { + face: Face, + w: []u8, + s: []u8, + b: []u8, + + fn free(self: *OwnedFace, gpa: std.mem.Allocator) void { + gpa.free(self.w); + gpa.free(self.s); + gpa.free(self.b); + } +}; + +fn ownedFace(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8, n: usize, k: usize) !OwnedFace { + const got = try join3(gpa, store, buf, prefix, name); + const k8 = k / 8; + if (got.w.len != n * k8 * 4 or got.s.len != n * (k / 32) * 2) { + gpa.free(got.w); + gpa.free(got.s); + gpa.free(got.b); + return error.UnexpectedTensor; + } + return .{ .face = .{ .w = got.w, .s = got.s, .b = got.b, .n = n, .k = k, .routed = false }, .w = got.w, .s = got.s, .b = got.b }; +} + +fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![]cuda.DeviceBuffer { + const out = try gpa.alloc(cuda.DeviceBuffer, n); + var i: usize = 0; + errdefer { + for (out[0..i]) |*b| b.free(); + gpa.free(out); + } + while (i < n) : (i += 1) { + out[i] = try cuda.DeviceBuffer.alloc(driver, len); + try out[i].fill8(0, null); + } + return out; +} + +fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !u32 { + if (p.mode[row] != 2) return error.UnexpectedTensor; + const h = p.s.h.ptr + @as(u64, row) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, row) * inj_stride; + const y = p.s.y.ptr + @as(u64, row) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, row) * wts_stride; + try writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[row] = 0; + try readout(tri, p.s, p.mixer, h, inj); + try qmm.matmul(p.driver, p.stream.*, p.s.mixed.ptr, p.s.xsm.ptr, p.head_q.w.ptr, p.head_q.s.ptr, p.head_q.b.ptr, p.s.logits.ptr, 1, vocab, dims); + const raw = try p.gpa.alloc(u8, vocab * 2); + defer p.gpa.free(raw); + try p.stream.synchronize(); + try p.s.logits.download(0, raw); + return argmax(raw); +} + +fn forward(comptime Tri: type, p: *Pass, tri: Tri, ids: []const u32, pos0: usize) !u32 { + @memset(p.mode, 0); + for (ids, 0..) |token, t| try embedRow(Tri, p, tri, token, p.s.h.ptr + @as(u64, t) * wide * 2); + for (0..layers_n) |layer| try oneLayer(Tri, p, tri, layer, ids, pos0); + return head(Tri, p, tri, ids.len - 1); +} + +/// `out` receives `new_tokens` greedy ids. The first is the prompt's argmax. The rest are fed back one at a time. +pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > 256) return error.UnexpectedTensor; + var store = try Store.open(gpa, io, model_dir); + defer store.deinit(); + const weight = try store.tensor(embed.weight_name); + const scales = try store.tensor(embed.scales_name); + const biases = try store.tensor(embed.biases_name); + var table = try embed.Table.upload(driver, weight, scales, biases); + defer table.deinit(); + const eos = try ple.eosOf(gpa, io, model_dir); + var history = [2]i64{ eos, eos }; + var scratch = try Scratch.alloc(driver, tokens.len); + defer scratch.free(); + try scratch.dummy.fill8(0, null); + const conv = try zeros(gpa, driver, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); + defer { + for (conv) |*b| b.free(); + gpa.free(conv); + } + const state = try zeros(gpa, driver, linear_n, state_n * 4); + defer { + for (state) |*b| b.free(); + gpa.free(state); + } + const cache = tokens.len + new_tokens; + const kc = try zeros(gpa, driver, attn_n, cache * kv_heads * head_dim * 2); + defer { + for (kc) |*b| b.free(); + gpa.free(kc); + } + const vc = try zeros(gpa, driver, attn_n, cache * kv_heads * head_dim * 2); + defer { + for (vc) |*b| b.free(); + gpa.free(vc); + } + const ikc = try zeros(gpa, driver, attn_n, cache * index_dim * 2); + defer { + for (ikc) |*b| b.free(); + gpa.free(ikc); + } + var tail = try cuda.DeviceBuffer.alloc(driver, ple_tail_n * wide * 2); + defer tail.free(); + try tail.fill8(0, null); + var mixer = try loadHc(gpa, driver, &store, "language_model.model.hyper_connection_mixer", false); + defer mixer.free(); + var buf: [80]u8 = undefined; + const head_w = try join3(gpa, &store, &buf, "language_model.lm_head", ""); + defer gpa.free(head_w.w); + defer gpa.free(head_w.s); + defer gpa.free(head_w.b); + if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; + var head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); + defer head_q.free(); + const mode = try gpa.alloc(u8, tokens.len); + defer gpa.free(mode); + var pass = Pass{ + .gpa = gpa, + .io = io, + .driver = driver, + .stream = stream, + .store = &store, + .model_dir = model_dir, + .s = &scratch, + .table = table, + .rows = tokens.len, + .mode = mode, + .history = &history, + .conv = conv, + .state = state, + .kc = kc, + .vc = vc, + .ikc = ikc, + .tail = tail, + .mixer = mixer, + .head_q = head_q, + }; + out[0] = try forward(Tri, &pass, tri, tokens, 0); + std.debug.print("token {d}\n", .{out[0]}); + var tok = out[0]; + var pos: usize = tokens.len; + for (1..new_tokens) |i| { + var one = [_]u32{tok}; + // The new token reuses row 0. Clear every row so a stale prompt pending is not written back. + tok = try forward(Tri, &pass, tri, &one, pos); + out[i] = tok; + pos += 1; + std.debug.print("token {d}\n", .{tok}); + } +} + +test "linear layers skip every fourth index and attention starts at 3" { + try std.testing.expect(linearLayer(0)); + try std.testing.expect(!linearLayer(3)); + try std.testing.expect(!linearLayer(47)); + try std.testing.expectEqual(@as(usize, 0), linearIndex(0)); + try std.testing.expectEqual(@as(usize, 2), linearIndex(2)); + try std.testing.expectEqual(@as(usize, 35), linearIndex(46)); + try std.testing.expectEqual(@as(usize, 0), attnIndex(3)); + try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); +} diff --git a/zig/src/families/flashnext/cuda_prompt_main.zig b/zig/src/families/flashnext/cuda_prompt_main.zig new file mode 100644 index 0000000000..bb84cafefc --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_main.zig @@ -0,0 +1,50 @@ +//! Prompt tokens in, greedy token ids out. The ids file is whitespace or comma separated. + +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); + +const want = [_]u32{ 1596, 1144, 310, 5707, 310, 1156, 25, 328 }; + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 5) { + std.debug.print("usage: flashnext-prompt \n", .{}); + return 2; + } + const count = try std.fmt.parseInt(usize, args[4], 10); + const text = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .limited(1 << 20)); + defer gpa.free(text); + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpa); + var it = std.mem.tokenizeAny(u8, text, ", \t\r\n"); + while (it.next()) |part| try ids.append(gpa, try std.fmt.parseInt(u32, part, 10)); + if (ids.items.len == 0 or count == 0) return 2; + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer set.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + const tri: flash.Tri = .{ .set = &set, .s = stream }; + const got = try gpa.alloc(u32, count); + defer gpa.free(got); + try flash.prompt.run(flash.Tri, gpa, io, &driver, &stream, tri, args[1], ids.items, count, got); + std.debug.print("tokens", .{}); + for (got) |id| std.debug.print(" {d}", .{id}); + std.debug.print("\n", .{}); + if (count == want.len and ids.items.len == 58) { + var same = true; + for (got, &want) |g, w| if (g != w) { + same = false; + }; + std.debug.print("{s}\n", .{if (same) "match" else "differ"}); + if (!same) return 1; + } + return 0; +} From 0d04c4346c91ca7df45ce7ca35901ad41fde4e69 Mon Sep 17 00:00:00 2001 From: Adrian Bolinger Date: Wed, 7 Oct 2026 03:55:02 -0400 Subject: [PATCH 028/115] Keep the drafts-off prompt match after splitting the forward and packing every expert onto the GPU once. --- zig/build/cuda.zig | 17 +- zig/src/families/flashnext/cuda_lanes.zig | 215 +++ zig/src/families/flashnext/cuda_launch.zig | 2 + zig/src/families/flashnext/cuda_prompt.zig | 1365 ++++------------- .../families/flashnext/cuda_prompt_load.zig | 389 +++++ .../families/flashnext/cuda_prompt_step.zig | 316 ++++ zig/src/families/flashnext/cuda_prompt_wt.zig | 398 +++++ zig/src/native/cuda.zig | 5 +- 8 files changed, 1613 insertions(+), 1094 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_lanes.zig create mode 100644 zig/src/families/flashnext/cuda_prompt_load.zig create mode 100644 zig/src/families/flashnext/cuda_prompt_step.zig create mode 100644 zig/src/families/flashnext/cuda_prompt_wt.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index a49fa3a54d..dc2d682294 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -95,6 +95,7 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin const flashnext = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_launch.zig"), .target = target, .optimize = optimize, .link_libc = true }); flashnext.addImport("cuda", cuda); flashnext.addImport("core", core); + flashnext.addImport("lanes", lanes); return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .flashnext = flashnext, .tokenizer = tokenizer, .heat = heat }; } @@ -127,7 +128,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu cli.addImport("core", mods.core); cli.addImport("lanes", mods.lanes); cli.addImport("nemotron", mods.nemotron); - const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron); + const native = engines(b, target, optimize, cuda, mods.lanes, mods.nemotron, mods.flashnext); cli.addOptions("build_options", build_options); cli.addImport("checkpoint_cli", checkpointCli(b, target, optimize, strip, native.engines, mods.core)); b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli })); @@ -164,18 +165,18 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const prompt_count = b.option([]const u8, "flashnext_count", "greedy tokens to emit") orelse "8"; if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); - nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; + nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, mods.tokenizer, build_options, true).root_module.strip = strip; } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. -fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { +fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron } }, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "nemotron", .module = nemotron }, .{ .name = "flashnext", .module = flashnext } }, }); return .{ .api = api, .engines = mod }; } @@ -186,8 +187,8 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { - const m = engines(b, target, optimize, cuda, lanes, nemotron); +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, flashnext: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { + const m = engines(b, target, optimize, cuda, lanes, nemotron, flashnext); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const exe = b.addExecutable(.{ .name = "tensorfold-native", .root_module = b.createModule(.{ @@ -215,7 +216,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); - const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron); + const native_modules = engines(b, host, .debug, cuda, mods.lanes, mods.nemotron, mods.flashnext); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = host, .optimize = .debug, .link_libc = true }); @@ -271,7 +272,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s }; const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{}); const mods = family(b, target, .fast, cuda, draft_ids); - return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, mods.tokenizer, build_options, false); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig new file mode 100644 index 0000000000..7502108299 --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -0,0 +1,215 @@ +//! Flash Next on the native CUDA server: one stream's prompt, then one token at a time. No draft head. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const prompt = @import("cuda_prompt.zig"); + +const be = lanes.backend; +const ring = 1024; + +/// `Tri` is the captured-kernel launcher. `open` loads the checkpoint and serves greedy or keyed draws. +pub fn Namespace(comptime Tri: type) type { + return struct { + pub const model_type = "qwen4_exp"; + pub const formats: []const []const u8 = &.{"mlx-q4g32"}; + pub const default_context: i64 = 262144; + pub const prefill_step: u32 = 0; + + pub const Options = struct { context: usize, drafts: bool }; + + pub const Loaded = struct { + backend: be.Backend, + facts: lanes.Model, + rows: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + }; + + const Lane = struct { seq: prompt.Seq }; + + const Cuda = struct { + gpa: std.mem.Allocator, + context: usize, + engine: *prompt.Engine, + set: cuda.aot.Set, + stream: cuda.Stream, + lanes: std.AutoHashMapUnmanaged(*lanes.Stream, Lane) = .empty, + drawn: [ring]u32 = undefined, + next: u64 = 0, + + fn tri(self: *Cuda) Tri { + return .{ .set = &self.set, .s = self.stream }; + } + + fn backend(self: *Cuda) be.Backend { + return .{ .ptr = self, .vtable = &.{ + .prefill = prefillFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + } }; + } + + fn facts() lanes.Model { + return .{ + .exact_width = 1, + .gpu_tokens = false, + .mtp = false, + .speculate = false, + .drafts = 0, + .batch_rows = 1, + .max_streams = 1, + }; + } + + fn take(self: *Cuda, token: u32) u64 { + const h = self.next; + self.drawn[h % ring] = token; + self.next += 1; + return h; + } + + fn tokenAt(self: *Cuda, seq: *prompt.Seq, s: *lanes.Stream, position: u64) !u32 { + if (s.sampling) |sm| { + const n = prompt.vocabulary; + const vals = try self.gpa.alloc(f64, n); + defer self.gpa.free(vals); + const ids = try self.gpa.alloc(u64, n); + defer self.gpa.free(ids); + for (vals, ids, 0..) |*v, *id, i| { + const bits = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + v.* = @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); + id.* = @intCast(i); + } + return @intCast(try lanes.sampling.choose(self.gpa, vals, ids, position, sm)); + } + return prompt.argmax(seq.logits); + } + + fn lane(self: *Cuda, s: *lanes.Stream) !*Lane { + return self.lanes.getPtr(s) orelse return error.UnknownStream; + } + }; + + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: []const u8, o: Options) !Loaded { + _ = o.drafts; // this model is served one token at a time; there is no draft head + const driver: *cuda.Driver = @constCast(ctx.d); + const own = try gpa.create(Cuda); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.context = o.context; + own.stream = try cuda.Stream.init(driver, true); + errdefer own.stream.deinit(); + own.set = try cuda.aot.Set.load(gpa, io, driver, ctx.device, kernels); + errdefer own.set.deinit(); + own.engine = try prompt.Engine.init(gpa, io, driver, &own.stream, dir); + own.lanes = .empty; + own.next = 0; + return .{ + .backend = own.backend(), + .facts = Cuda.facts(), + .rows = 1, + .ctx = own, + .deinit = release, + }; + } + + fn release(p: *anyopaque) void { + const own: *Cuda = @ptrCast(@alignCast(p)); + var it = own.lanes.valueIterator(); + while (it.next()) |l| l.seq.deinit(own.gpa); + own.lanes.deinit(own.gpa); + own.engine.deinit(); + own.set.deinit(); + own.stream.deinit(); + own.gpa.destroy(own); + } + + fn of(ptr: *anyopaque) *Cuda { + return @ptrCast(@alignCast(ptr)); + } + + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(ptr); + const ids = s.prompt(); + if (ids.len == 0 or ids.len + s.max_new > self.context) return error.PromptTooLong; + const gop = try self.lanes.getOrPut(self.gpa, s); + if (gop.found_existing) gop.value_ptr.seq.deinit(self.gpa); + gop.value_ptr.* = .{ .seq = self.engine.newSeq(ids.len + s.max_new) catch |err| { + self.lanes.removeByPtr(gop.key_ptr); + return err; + } }; + self.engine.predict(Tri, self.tri(), &gop.value_ptr.seq, ids) catch |err| { + gop.value_ptr.seq.deinit(self.gpa); + self.lanes.removeByPtr(gop.key_ptr); + return err; + }; + } + + fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const self = of(ptr); + const l = try self.lane(s); + if (position != l.seq.fed) return error.PositionMismatch; + return self.take(try self.tokenAt(&l.seq, s, position)); + } + + fn queueFn(ptr: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) anyerror!u64 { + _ = ptr; + _ = s; + _ = feed; + _ = position; + return error.NotPipelined; + } + + fn readFn(ptr: *anyopaque, handle: u64) anyerror!u32 { + const self = of(ptr); + if (handle >= self.next or self.next - handle > ring) return error.NoSuchToken; + return self.drawn[handle % ring]; + } + + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const self = of(ptr); + for (windows, out) |w, *o| { + if (w.held != 0 or w.parents != null) return error.HeldDraftsNotBuilt; + const l = try self.lane(w.stream); + const rows = w.rows(); + var input = w.pending; + for (0..rows) |r| { + if (w.positions[r] != l.seq.fed + 1) return error.PositionMismatch; + var one = [_]u32{input}; + try self.engine.predict(Tri, self.tri(), &l.seq, &one); + o.sampled[r] = try self.tokenAt(&l.seq, w.stream, w.positions[r]); + if (r + 1 < rows) input = w.tokens[r]; + } + if (w.tokens.len > 0) @memcpy(o.drafts[0..w.tokens.len], w.tokens); + } + } + + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + _ = ptr; + for (windows, paths) |w, path| { + if (path.len != w.rows()) return error.CannotRewind; + for (path, 0..) |r, i| if (r != i) return error.CannotRewind; + } + } + + fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + _ = ptr; + _ = requests; + return error.NoDraftHead; + } + + fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { + const self = of(ptr); + if (self.lanes.fetchRemove(s)) |kv| { + var l = kv.value; + l.seq.deinit(self.gpa); + } + } + }; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index e442ed8567..75339636e9 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -14,6 +14,7 @@ pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); +pub const native = @import("cuda_lanes.zig").Namespace(Tri); const p = aot.ptr; @@ -256,6 +257,7 @@ test { _ = ple; _ = mix; _ = prompt; + _ = native; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 8406b8b0fb..5dad8c369b 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -2,809 +2,58 @@ const std = @import("std"); const cuda = @import("cuda"); -const core = @import("core"); const embed = @import("cuda_embed.zig"); -const hc = @import("cuda_hc.zig"); const qmm = @import("cuda_qmm.zig"); const gdn = @import("cuda_gdn.zig"); const ple = @import("cuda_ple.zig"); - -const layers_n: usize = 48; -const linear_n: usize = 36; -const attn_n: usize = 12; -const dims: usize = 2560; -const streams: usize = 4; -const wide: usize = streams * dims; -const low: usize = 320; -const eps: f32 = 1e-6; -const vocab: usize = 248320; -const experts_n: usize = 512; -const slots: usize = 11; -const moe_w: usize = 640; -const block_bytes: usize = 160 * 4; -const tile: usize = 16; -const q_heads: usize = 24; -const kv_heads: usize = 2; -const head_dim: usize = 256; -const index_heads: usize = 4; -const index_dim: usize = 128; -const half: usize = 32; -const proj_n: usize = 13952; -const out_k: usize = 6144; -const nch: usize = 3; -const ple_layer: usize = 1; -const inj_stride: usize = 16; -const wts_stride: usize = 48; -const ple_tail_n: usize = 9; -const state_n: usize = gdn.nv * gdn.dv * gdn.dk; -const hc_sk: usize = 32; -const out_sk: usize = 8; - -const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; -const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; -const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; -const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; - -const Tensor = core.safetensors.Tensor; - -const Held = struct { rel: []u8, file: Shard }; - -const Shard = struct { - io: std.Io, - file: std.Io.File, - map: std.Io.File.MemoryMap, - at: usize, - arena: std.heap.ArenaAllocator, - names: std.StringHashMapUnmanaged(core.safetensors.Entry), - - fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Shard { - var file = try std.Io.Dir.cwd().openFile(io, path, .{}); - errdefer file.close(io); - const len: usize = @intCast(try file.length(io)); - if (len < 8) return error.BadSafetensors; - var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); - errdefer map.destroy(io); - const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); - if (header_len > len - 8) return error.BadSafetensors; - var arena = std.heap.ArenaAllocator.init(gpa); - errdefer arena.deinit(); - var names: std.StringHashMapUnmanaged(core.safetensors.Entry) = .empty; - const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), map.memory[8..][0..header_len], .{}); - defer parsed.deinit(); - var it = parsed.value.object.iterator(); - while (it.next()) |kv| { - if (std.mem.eql(u8, kv.key_ptr.*, "__metadata__")) continue; - const o = kv.value_ptr.object; - const dtype = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; - const shape = o.get("shape").?.array.items; - if (shape.len > core.safetensors.max_rank) continue; - var e: core.safetensors.Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = 0 }; - var n: usize = dtype.size(); - for (shape, 0..) |d, i| { - e.shape[i] = @intCast(d.integer); - n *= e.shape[i]; - } - const offs = o.get("data_offsets").?.array.items; - e.begin = @intCast(offs[0].integer); - e.end = @intCast(offs[1].integer); - if (e.end < e.begin or e.end - e.begin != n or e.end > len - 8 - header_len) return error.BadSafetensors; - try names.put(arena.allocator(), try arena.allocator().dupe(u8, kv.key_ptr.*), e); - } - return .{ .io = io, .file = file, .map = map, .at = 8 + header_len, .arena = arena, .names = names }; - } - - fn close(self: *Shard) void { - self.map.destroy(self.io); - self.file.close(self.io); - self.arena.deinit(); - self.* = undefined; - } - - fn get(self: *const Shard, name: []const u8) ?Tensor { - const e = self.names.get(name) orelse return null; - return .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .bytes = self.map.memory[self.at + e.begin .. self.at + e.end] }; - } -}; - -const Store = struct { - gpa: std.mem.Allocator, - io: std.Io, - dir: []const u8, - parsed: std.json.Parsed(std.json.Value), - path: []u8 = &.{}, - rel: []u8 = &.{}, - file: ?Shard = null, - held: [12]Held = undefined, - nheld: usize = 0, - - fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { - const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); - defer gpa.free(index); - const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); - defer gpa.free(text); - return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}) }; - } - - fn deinit(self: *Store) void { - self.release(); - self.closeFile(); - self.parsed.deinit(); - } - - fn release(self: *Store) void { - for (self.held[0..self.nheld]) |*h| { - h.file.close(); - self.gpa.free(h.rel); - } - self.nheld = 0; - } - - fn hold(self: *Store, key: []const u8) !Tensor { - const file = try self.shard(key); - for (self.held[0..self.nheld]) |*h| { - if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; - } - if (self.nheld == self.held.len) return error.UnexpectedTensor; - const rel = try self.gpa.dupe(u8, file); - errdefer self.gpa.free(rel); - const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); - defer self.gpa.free(path); - self.held[self.nheld] = .{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }; - self.nheld += 1; - return self.held[self.nheld - 1].file.get(key) orelse { - std.debug.print("absent {s}\n", .{key}); - return error.MissingTensor; - }; - } - - fn closeFile(self: *Store) void { - if (self.file) |*f| f.close(); - self.file = null; - if (self.path.len != 0) self.gpa.free(self.path); - self.path = &.{}; - if (self.rel.len != 0) self.gpa.free(self.rel); - self.rel = &.{}; - } - - fn shard(self: *Store, key: []const u8) ![]const u8 { - const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; - const v = wm.object.get(key) orelse { - std.debug.print("missing {s}\n", .{key}); - return error.MissingTensor; - }; - return v.string; - } - - fn use(self: *Store, key: []const u8) !void { - const file = try self.shard(key); - if (self.file != null and std.mem.eql(u8, file, self.rel)) return; - self.closeFile(); - self.rel = try self.gpa.dupe(u8, file); - self.path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); - self.file = try Shard.load(self.gpa, self.io, self.path); - } - - fn tensor(self: *Store, key: []const u8) !Tensor { - try self.use(key); - return self.file.?.get(key) orelse { - std.debug.print("absent {s}\n", .{key}); - return error.MissingTensor; - }; - } - - fn copyOf(self: *Store, key: []const u8) ![]u8 { - const t = try self.tensor(key); - const out = try self.gpa.alloc(u8, t.bytes.len); - @memcpy(out, t.bytes); - return out; - } -}; - -const Face = struct { - w: []const u8, - s: []const u8, - b: []const u8, - n: usize, - k: usize, - routed: bool, -}; - -const Q4 = struct { - w: cuda.DeviceBuffer, - s: cuda.DeviceBuffer, - b: cuda.DeviceBuffer, - n: usize, - k: usize, - - fn free(self: *Q4) void { - self.w.free(); - self.s.free(); - self.b.free(); - } - - fn empty(d: *cuda.Driver) Q4 { - const z: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; - return .{ .w = z, .s = z, .b = z, .n = 0, .k = 0 }; - } -}; - -const HcW = struct { - scale: cuda.DeviceBuffer, - down: Q4, - up: Q4, - n_down: usize, - has_inj: usize, - - fn free(self: *HcW) void { - self.scale.free(); - self.down.free(); - self.up.free(); - } -}; - -fn linearLayer(layer: usize) bool { - return layer % 4 != 3; -} - -fn linearIndex(layer: usize) usize { - return layer - layer / 4; -} - -fn attnIndex(layer: usize) usize { - return layer / 4; -} - -fn toBf16(v: f32) u16 { - const bits: u32 = @bitCast(v); - if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); - return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); -} - -fn promote(bits: u16) f32 { - return @bitCast(@as(u32, bits) << 16); -} - -fn cint(v: usize) c_int { - return @intCast(v); -} - -fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { - var n: usize = 0; - for (parts) |p| n += p.len; - const out = try gpa.alloc(u8, n); - var at: usize = 0; - for (parts) |p| { - @memcpy(out[at..][0..p.len], p); - at += p.len; - } - return out; -} - -fn asF32(gpa: std.mem.Allocator, t: Tensor) ![]f32 { - if (t.rank != 1) return error.UnexpectedTensor; - const out = try gpa.alloc(f32, t.dim(0)); - if (t.dtype == .f32 and t.bytes.len == out.len * 4) { - for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); - } else if (t.dtype == .bf16 and t.bytes.len == out.len * 2) { - for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)); - } else return error.UnexpectedTensor; - return out; -} - -fn uploadQ4(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { - var lane = try qmm.pack(gpa, words, scales, biases, n, k8); - defer lane.deinit(gpa); - var w = try cuda.DeviceBuffer.fromHost(driver, lane.weight); - errdefer w.free(); - var s = try cuda.DeviceBuffer.fromHost(driver, lane.scales); - errdefer s.free(); - const b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); - return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; -} - -fn uploadTiled(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { - const kg = k8 / 4; - const tiled = try hc.tileWords(gpa, words, n, k8); - defer gpa.free(tiled); - const ts = try hc.transposeBf16(gpa, scales, n, kg); - defer gpa.free(ts); - const tb = try hc.transposeBf16(gpa, biases, n, kg); - defer gpa.free(tb); - var w = try cuda.DeviceBuffer.fromHost(driver, tiled); - errdefer w.free(); - var s = try cuda.DeviceBuffer.fromHost(driver, ts); - errdefer s.free(); - const b = try cuda.DeviceBuffer.fromHost(driver, tb); - return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; -} - -fn join3(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8) !struct { w: []u8, s: []u8, b: []u8 } { - const w = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); - errdefer gpa.free(w); - const s = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); - errdefer gpa.free(s); - const b = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); - return .{ .w = w, .s = s, .b = b }; -} - -fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { - var buf: [180]u8 = undefined; - const down = try join3(gpa, store, &buf, prefix, ".input_mix_weight_down"); - defer gpa.free(down.w); - defer gpa.free(down.s); - defer gpa.free(down.b); - const up = try join3(gpa, store, &buf, prefix, ".input_mix_weight_up"); - defer gpa.free(up.w); - defer gpa.free(up.s); - defer gpa.free(up.b); - const gamma = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix})); - const scale = try asF32(gpa, gamma); - defer gpa.free(scale); - if (scale.len != wide) return error.UnexpectedTensor; - var n: usize = low; - var words: []u8 = undefined; - var scales: []u8 = undefined; - var biases: []u8 = undefined; - var inj_w: []u8 = &.{}; - var inj_s: []u8 = &.{}; - var inj_b: []u8 = &.{}; - if (inject) { - const inj = try join3(gpa, store, &buf, prefix, ".block_inject_weight"); - inj_w = inj.w; - inj_s = inj.s; - inj_b = inj.b; - n = low + 4; - words = try cat(gpa, &.{ down.w, inj.w }); - scales = try cat(gpa, &.{ down.s, inj.s }); - biases = try cat(gpa, &.{ down.b, inj.b }); - } else { - words = try gpa.dupe(u8, down.w); - scales = try gpa.dupe(u8, down.s); - biases = try gpa.dupe(u8, down.b); - } - defer gpa.free(words); - defer gpa.free(scales); - defer gpa.free(biases); - defer if (inject) gpa.free(inj_w); - defer if (inject) gpa.free(inj_s); - defer if (inject) gpa.free(inj_b); - const k8 = wide / 8; - if (words.len != n * k8 * 4) return error.UnexpectedTensor; - var down_q = try uploadTiled(gpa, driver, words, scales, biases, n, k8); - errdefer down_q.free(); - var up_q = try uploadTiled(gpa, driver, up.w, up.s, up.b, wide, low / 8); - errdefer up_q.free(); - const scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); - return .{ .scale = scale_b, .down = down_q, .up = up_q, .n_down = n, .has_inj = @intFromBool(inject) }; -} - -fn stitch(gpa: std.mem.Allocator, store: *Store, prefix: []const u8, names: []const []const u8) !struct { w: []u8, s: []u8, b: []u8, n: usize } { - var buf: [180]u8 = undefined; - var ws: [6][]u8 = undefined; - var ss: [6][]u8 = undefined; - var bs: [6][]u8 = undefined; - var n: usize = 0; - for (names, 0..) |part, i| { - const got = try join3(gpa, store, &buf, prefix, part); - ws[i] = got.w; - ss[i] = got.s; - bs[i] = got.b; - n += 1; - } - defer for (0..n) |i| { - gpa.free(ws[i]); - gpa.free(ss[i]); - gpa.free(bs[i]); - }; - return .{ .w = try cat(gpa, ws[0..n]), .s = try cat(gpa, ss[0..n]), .b = try cat(gpa, bs[0..n]), .n = n }; -} - -const Scratch = struct { - h: cuda.DeviceBuffer, - y: cuda.DeviceBuffer, - wts: cuda.DeviceBuffer, - inj: cuda.DeviceBuffer, - pss: cuda.DeviceBuffer, - norm: cuda.DeviceBuffer, - dn: cuda.DeviceBuffer, - hc_part: cuda.DeviceBuffer, - act: cuda.DeviceBuffer, - xs_act: cuda.DeviceBuffer, - mixed: cuda.DeviceBuffer, - xsm: cuda.DeviceBuffer, - proj: cuda.DeviceBuffer, - part: cuda.DeviceBuffer, - gout: cuda.DeviceBuffer, - gxs: cuda.DeviceBuffer, - st: cuda.DeviceBuffer, - ple_emb: cuda.DeviceBuffer, - ple_xs: cuda.DeviceBuffer, - ple_keys: cuda.DeviceBuffer, - ple_vals: cuda.DeviceBuffer, - ple_gated: cuda.DeviceBuffer, - ple_pss: cuda.DeviceBuffer, - ple_nrow: cuda.DeviceBuffer, - ple_part: cuda.DeviceBuffer, - pa: cuda.DeviceBuffer, - q: cuda.DeviceBuffer, - iq: cuda.DeviceBuffer, - po: cuda.DeviceBuffer, - pm: cuda.DeviceBuffer, - pl: cuda.DeviceBuffer, - ao: cuda.DeviceBuffer, - gated: cuda.DeviceBuffer, - gxs_attn: cuda.DeviceBuffer, - dummy: cuda.DeviceBuffer, - inv: cuda.DeviceBuffer, - pos: cuda.DeviceBuffer, - ids: cuda.DeviceBuffer, - logits: cuda.DeviceBuffer, - router_l: cuda.DeviceBuffer, - picks: cuda.DeviceBuffer, - - fn alloc(driver: *cuda.Driver, rows: usize) !Scratch { - var s: Scratch = undefined; - s.h = try cuda.DeviceBuffer.alloc(driver, rows * wide * 2); - s.y = try cuda.DeviceBuffer.alloc(driver, rows * slots * dims * 4); - s.wts = try cuda.DeviceBuffer.alloc(driver, rows * wts_stride); - s.inj = try cuda.DeviceBuffer.alloc(driver, rows * inj_stride); - s.pss = try cuda.DeviceBuffer.alloc(driver, (dims / 256) * streams * 4); - s.norm = try cuda.DeviceBuffer.alloc(driver, wide * 2); - s.dn = try cuda.DeviceBuffer.alloc(driver, (low + 4) * 2); - s.hc_part = try cuda.DeviceBuffer.alloc(driver, hc_sk * (low + 4) * 4); - s.act = try cuda.DeviceBuffer.alloc(driver, low * 2); - s.xs_act = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); - s.mixed = try cuda.DeviceBuffer.alloc(driver, dims * 2); - s.xsm = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); - s.proj = try cuda.DeviceBuffer.alloc(driver, gdn.proj_width * 2); - s.part = try cuda.DeviceBuffer.alloc(driver, out_sk * dims * 4); - s.gout = try cuda.DeviceBuffer.alloc(driver, gdn.value_dim * 2); - s.gxs = try cuda.DeviceBuffer.alloc(driver, (gdn.value_dim / 32) * 4); - s.st = try cuda.DeviceBuffer.alloc(driver, state_n * 4); - s.ple_emb = try cuda.DeviceBuffer.alloc(driver, dims * 2); - s.ple_xs = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); - s.ple_keys = try cuda.DeviceBuffer.alloc(driver, wide * 2); - s.ple_vals = try cuda.DeviceBuffer.alloc(driver, dims * 2); - s.ple_gated = try cuda.DeviceBuffer.alloc(driver, wide * 2); - s.ple_pss = try cuda.DeviceBuffer.alloc(driver, streams * 4); - s.ple_nrow = try cuda.DeviceBuffer.alloc(driver, wide * 2); - s.ple_part = try cuda.DeviceBuffer.alloc(driver, 4 * dims * 4); - s.pa = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); - s.q = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); - s.iq = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); - s.po = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); - s.pm = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); - s.pl = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); - s.ao = try cuda.DeviceBuffer.alloc(driver, out_k * 2); - s.gated = try cuda.DeviceBuffer.alloc(driver, out_k * 2); - s.gxs_attn = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); - s.dummy = try cuda.DeviceBuffer.alloc(driver, 4096); - var inv: [half]f32 = undefined; - for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); - s.inv = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); - s.pos = try cuda.DeviceBuffer.alloc(driver, 4); - s.ids = try cuda.DeviceBuffer.alloc(driver, 4); - s.logits = try cuda.DeviceBuffer.alloc(driver, vocab * 2); - s.router_l = try cuda.DeviceBuffer.alloc(driver, (experts_n + 1) * 4); - s.picks = try cuda.DeviceBuffer.alloc(driver, slots * 4); - return s; - } - - fn free(self: *Scratch) void { - self.h.free(); - self.y.free(); - self.wts.free(); - self.inj.free(); - self.pss.free(); - self.norm.free(); - self.dn.free(); - self.hc_part.free(); - self.act.free(); - self.xs_act.free(); - self.mixed.free(); - self.xsm.free(); - self.proj.free(); - self.part.free(); - self.gout.free(); - self.gxs.free(); - self.st.free(); - self.ple_emb.free(); - self.ple_xs.free(); - self.ple_keys.free(); - self.ple_vals.free(); - self.ple_gated.free(); - self.ple_pss.free(); - self.ple_nrow.free(); - self.ple_part.free(); - self.pa.free(); - self.q.free(); - self.iq.free(); - self.po.free(); - self.pm.free(); - self.pl.free(); - self.ao.free(); - self.gated.free(); - self.gxs_attn.free(); - self.dummy.free(); - self.inv.free(); - self.pos.free(); - self.ids.free(); - self.logits.free(); - self.router_l.free(); - self.picks.free(); - } -}; - -fn writeback(tri: anytype, s: *Scratch, h: u64, mode: usize, inj: u64, y: u64, wts: u64) !void { - const dummy = s.dummy.ptr; - if (mode == 0) { - try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 0, 1, 1, 1); - } else if (mode == 2) { - try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", dims, 1, dims, streams, 2, 10, 11, 1); - } else if (mode == 4) { - try tri.hcWriteback(h, h, s.pss.ptr, s.part.ptr, "*fp32", inj, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 4, 1, 1, out_sk); - } else return error.UnexpectedTensor; -} - -fn readout(tri: anytype, s: *Scratch, hc_w: HcW, h: u64, inj: u64) !void { - try tri.hcDown(h, s.pss.ptr, hc_w.scale.ptr, s.norm.ptr, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn.ptr, s.hc_part.ptr, eps, 1, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); - const inj_ptr = if (hc_w.has_inj == 1) inj else s.act.ptr; - try tri.hcReduceAct(s.hc_part.ptr, s.act.ptr, s.xs_act.ptr, inj_ptr, hc_sk, 1, streams, low, hc_w.n_down, hc_w.has_inj); - try tri.hcUpmix(s.act.ptr, s.xs_act.ptr, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.norm.ptr, s.mixed.ptr, s.xsm.ptr, 1, wide, low, dims, streams); -} - -fn shiftTail(gpa: std.mem.Allocator, stream: *cuda.Stream, window: cuda.DeviceBuffer, taps_n: usize, row_bytes: usize, fresh: []const u8) !void { - const bytes = taps_n * row_bytes; - const raw = try gpa.alloc(u8, bytes); - defer gpa.free(raw); - try stream.synchronize(); - try window.download(0, raw); - const next = try gpa.alloc(u8, bytes); - defer gpa.free(next); - const keep = taps_n - 1; - @memcpy(next[0 .. keep * row_bytes], raw[row_bytes..]); - @memcpy(next[keep * row_bytes ..], fresh); - try window.upload(0, next); -} - -fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s: *Scratch, proj_q: Q4, out_q: Q4, conv: cuda.DeviceBuffer, cw: u64, a_log: u64, dt: u64, norm: u64, state: *cuda.DeviceBuffer) !void { - try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.proj.ptr, 1, proj_q.n, proj_q.k); - try gdn.launch(driver, stream.*, s.proj.ptr, conv.ptr, cw, state.ptr, a_log, dt, norm, eps, s.gout.ptr, s.gxs.ptr, s.st.ptr); - try state.copyFrom(0, s.st.ptr, state_n * 4, stream.handle); - const row = try gpa.alloc(u8, gdn.conv_dim * 2); - defer gpa.free(row); - try stream.synchronize(); - try s.proj.download(0, row); - try shiftTail(gpa, stream, conv, gdn.taps - 1, gdn.conv_dim * 2, row); - try qmm.partials(driver, stream.*, s.gout.ptr, s.gxs.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, gdn.value_dim, out_sk); -} - -fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { - try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); - try s.pos.upload(0, std.mem.asBytes(&pos)); - try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); - try tri.attnChunks(s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, 1); - try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads); - try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); - try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); -} - -fn copyExpert(dst_w: []u8, dst_s: []u8, dst_b: []u8, routed: Face, shared: Face, id: i32, slot: usize) void { - const p = if (id == experts_n) shared else routed; - const k8 = p.k / 8; - const kg = p.k / 32; - const wb = p.n * k8 * 4; - const sb = p.n * kg * 2; - const at: usize = if (id == experts_n) 0 else @intCast(id); - @memcpy(dst_w[slot * wb ..][0..wb], p.w[at * wb ..][0..wb]); - @memcpy(dst_s[slot * sb ..][0..sb], p.s[at * sb ..][0..sb]); - @memcpy(dst_b[slot * sb ..][0..sb], p.b[at * sb ..][0..sb]); -} - -fn packDevice(driver: *cuda.Driver, stream: cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) !cuda.DeviceBuffer { - const kg = k / 32; - const nb = n / 32; - var w_b = try cuda.DeviceBuffer.fromHost(driver, words); - defer w_b.free(); - var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); - defer s_b.free(); - var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); - defer b_b.free(); - var out_b = try cuda.DeviceBuffer.alloc(driver, e * nb * kg * block_bytes); - errdefer out_b.free(); - var a: cuda.Args = .{}; - a.add(w_b.ptr); - a.add(s_b.ptr); - a.add(b_b.ptr); - a.add(out_b.ptr); - a.add(cint(n)); - a.add(cint(k / 8)); - a.add(cint(kg)); - a.add(cint(nb)); - try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream, &a); - try stream.synchronize(); - return out_b; -} - -fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { - if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; - const nblk = gate.len / block_bytes; - const out = try gpa.alloc(u8, gate.len + up.len); - for (0..nblk) |b| { - @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); - @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); - } - return out; -} - -fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { - var a: cuda.Args = .{}; - a.add(x); - a.add(cint(x_stride)); - a.add(cint(slot_n)); - a.add(w); - a.add(cint(kg)); - a.add(cint(nb)); - a.add(items); - a.add(counts); - a.add(members); - a.add(out); - a.add(cint(n)); - a.add(@as(f32, 0)); - try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((units + 3) / 4) }, .block = .{ .x = 128 } }, stream, &a); -} - -fn moeStep(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, router_w: u64, faces: [3]Face, shared: [3]Face, y: u64, wts: u64) !void { - try tri.router(s.mixed.ptr, router_w, s.router_l.ptr, dims); - try tri.topkRows(s.router_l.ptr, s.picks.ptr, wts); - var pick_raw: [slots * 4]u8 = undefined; - try stream.synchronize(); - try s.picks.download(0, &pick_raw); - var packed_ids: [slots]i32 = undefined; - var npack: usize = 0; - var local: [slots]i32 = undefined; - for (0..slots) |slot| { - const id = std.mem.readInt(i32, pick_raw[4 * slot ..][0..4], .little); - if (id < 0 or id > experts_n) return error.UnexpectedTensor; - var found: ?usize = null; - for (packed_ids[0..npack], 0..) |prev, i| if (prev == id) { - found = i; - }; - if (found) |i| { - local[slot] = @intCast(i); - } else { - packed_ids[npack] = id; - local[slot] = @intCast(npack); - npack += 1; - } - } - const gw = npack * moe_w * (dims / 8) * 4; - const gs = npack * moe_w * (dims / 32) * 2; - const dw = npack * dims * (moe_w / 8) * 4; - const ds = npack * dims * (moe_w / 32) * 2; - const gathered_w = try gpa.alloc(u8, gw); - defer gpa.free(gathered_w); - const gathered_s = try gpa.alloc(u8, gs); - defer gpa.free(gathered_s); - const gathered_b = try gpa.alloc(u8, gs); - defer gpa.free(gathered_b); - const up_w = try gpa.alloc(u8, gw); - defer gpa.free(up_w); - const up_s = try gpa.alloc(u8, gs); - defer gpa.free(up_s); - const up_b = try gpa.alloc(u8, gs); - defer gpa.free(up_b); - const down_w = try gpa.alloc(u8, dw); - defer gpa.free(down_w); - const down_s = try gpa.alloc(u8, ds); - defer gpa.free(down_s); - const down_b = try gpa.alloc(u8, ds); - defer gpa.free(down_b); - for (packed_ids[0..npack], 0..) |id, slot| { - copyExpert(gathered_w, gathered_s, gathered_b, faces[0], shared[0], id, slot); - copyExpert(up_w, up_s, up_b, faces[1], shared[1], id, slot); - copyExpert(down_w, down_s, down_b, faces[2], shared[2], id, slot); - } - var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); - defer pack_mod.unload(); - var exp_mod = try cuda.Module.load(driver, cuda.kernels.experts); - defer exp_mod.unload(); - const pack_fn = try pack_mod.function(pack_symbol); - var packed_g = try packDevice(driver, stream.*, pack_fn, gathered_w, gathered_s, gathered_b, npack, moe_w, dims); - defer packed_g.free(); - var packed_u = try packDevice(driver, stream.*, pack_fn, up_w, up_s, up_b, npack, moe_w, dims); - defer packed_u.free(); - var packed_d = try packDevice(driver, stream.*, pack_fn, down_w, down_s, down_b, npack, dims, moe_w); - defer packed_d.free(); - const g_raw = try gpa.alloc(u8, packed_g.len); - defer gpa.free(g_raw); - const u_raw = try gpa.alloc(u8, packed_u.len); - defer gpa.free(u_raw); - try packed_g.download(0, g_raw); - try packed_u.download(0, u_raw); - const stacked = try stackHalves(gpa, g_raw, u_raw); - defer gpa.free(stacked); - var up_wb = try cuda.DeviceBuffer.fromHost(driver, stacked); - defer up_wb.free(); - var local_pick: [slots]i32 = undefined; - for (0..slots) |slot| local_pick[slot] = local[slot]; - var picks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&local_pick)); - defer picks_b.free(); - var members_b = try cuda.DeviceBuffer.alloc(driver, slots * 4); - defer members_b.free(); - var items_b = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); - defer items_b.free(); - var counts_b = try cuda.DeviceBuffer.alloc(driver, 8); - defer counts_b.free(); - var act_b = try cuda.DeviceBuffer.alloc(driver, slots * moe_w * 2); - defer act_b.free(); - const plan_fn = try exp_mod.function(plan_symbol); - var plan_a: cuda.Args = .{}; - plan_a.add(picks_b.ptr); - plan_a.add(cint(slots)); - plan_a.add(cint(npack)); - plan_a.add(cint(tile)); - plan_a.add(members_b.ptr); - plan_a.add(items_b.ptr); - plan_a.add(counts_b.ptr); - try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); - const up_fn = try exp_mod.function(up_symbol); - const down_fn = try exp_mod.function(down_symbol); - try launchExpert(up_fn, stream.*, s.mixed.ptr, dims, slots, up_wb.ptr, dims / 32, moe_w / 32, items_b.ptr, counts_b.ptr, members_b.ptr, act_b.ptr, moe_w, slots * (moe_w / 32)); - try launchExpert(down_fn, stream.*, act_b.ptr, moe_w, 0, packed_d.ptr, moe_w / 32, dims / 32, items_b.ptr, counts_b.ptr, members_b.ptr, y, dims, slots * (dims / 32)); - try stream.synchronize(); -} - -fn faceOf(store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !Face { - const w = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); - const sc = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); - const b = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); - const k8 = k / 8; - const kg = k / 32; - if (routed) { - if (!w.is(.u32, &.{ experts_n, n, k8 }) or !sc.is(.bf16, &.{ experts_n, n, kg }) or !b.is(.bf16, &.{ experts_n, n, kg })) return error.UnexpectedTensor; - } else if (!w.is(.u32, &.{ n, k8 }) or !sc.is(.bf16, &.{ n, kg }) or !b.is(.bf16, &.{ n, kg })) return error.UnexpectedTensor; - return .{ .w = w.bytes, .s = sc.bytes, .b = b.bytes, .n = n, .k = k, .routed = routed }; -} - -fn reduceBf16(gpa: std.mem.Allocator, stream: *cuda.Stream, part: cuda.DeviceBuffer, sk: usize, n: usize, dst: cuda.DeviceBuffer) !void { - const raw = try gpa.alloc(u8, sk * n * 4); - defer gpa.free(raw); - try stream.synchronize(); - try part.download(0, raw); - const row = try gpa.alloc(u16, n); - defer gpa.free(row); - for (0..n) |i| { - var acc: f32 = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); - for (1..sk) |s| acc += @as(f32, @bitCast(std.mem.readInt(u32, raw[4 * (s * n + i) ..][0..4], .little))); - row[i] = toBf16(acc); - } - try dst.upload(0, std.mem.sliceAsBytes(row)); -} - -fn pleStep(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, model_dir: []const u8, history: *[2]i64, token: u32, h: u64, key_q: Q4, val_q: Q4, nk: u64, nq: u64, nc: u64, cw: u64, tail: *cuda.DeviceBuffer) !void { - const emb = try gpa.alloc(u16, dims); - defer gpa.free(emb); - const xs = try gpa.alloc(f32, dims / 32); - defer gpa.free(xs); - try ple.embedding(gpa, io, model_dir, history, token, emb, xs); - try s.ple_emb.upload(0, std.mem.sliceAsBytes(emb)); - try s.ple_xs.upload(0, std.mem.sliceAsBytes(xs)); - try qmm.matmul(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, key_q.w.ptr, key_q.s.ptr, key_q.b.ptr, s.ple_keys.ptr, 1, wide, dims); - try qmm.partials(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, val_q.w.ptr, val_q.s.ptr, val_q.b.ptr, s.ple_vals.ptr, s.ple_part.ptr, 1, dims, dims, 4); - try reduceBf16(gpa, stream, s.ple_part, 4, dims, s.ple_vals); - try tri.pleGate(s.ple_keys.ptr, s.ple_vals.ptr, h, nk, nq, s.ple_gated.ptr, s.ple_pss.ptr, eps, 1, dims, streams); - try tri.pleConv(s.ple_gated.ptr, s.ple_pss.ptr, nc, tail.ptr, cw, h, h, s.ple_nrow.ptr, eps, 1, dims, streams, 4, 3); - const nrow = try gpa.alloc(u8, wide * 2); - defer gpa.free(nrow); - try stream.synchronize(); - try s.ple_nrow.download(0, nrow); - try shiftTail(gpa, stream, tail.*, ple_tail_n, wide * 2, nrow); - history[0] = history[1]; - history[1] = token; -} - -fn argmax(raw: []const u8) u32 { +const weights = @import("cuda_prompt_wt.zig"); +const step = @import("cuda_prompt_step.zig"); +const load = @import("cuda_prompt_load.zig"); + +const layers_n = weights.layers_n; +const linear_n = weights.linear_n; +const attn_n = weights.attn_n; +const dims = weights.dims; +const streams = weights.streams; +const wide = weights.wide; +const vocab = weights.vocab; +const experts_n = weights.experts_n; +const slots = weights.slots; +const moe_w = weights.moe_w; +const head_dim = weights.head_dim; +const index_dim = weights.index_dim; +const proj_n = weights.proj_n; +const out_k = weights.out_k; +const ple_layer = weights.ple_layer; +const inj_stride = weights.inj_stride; +const wts_stride = weights.wts_stride; +const ple_tail_n = weights.ple_tail_n; +const state_n = weights.state_n; +const kv_heads = weights.kv_heads; +const Store = weights.Store; +const Face = weights.Face; +const Q4 = weights.Q4; +const HcW = weights.HcW; +const linearLayer = weights.linearLayer; +const linearIndex = weights.linearIndex; +const attnIndex = weights.attnIndex; +const promote = weights.promote; +const asF32 = weights.asF32; +const uploadQ4 = weights.uploadQ4; +const join3 = weights.join3; +const loadHc = weights.loadHc; +const stitch = weights.stitch; +const Scratch = step.Scratch; +const writeback = step.writeback; +const readout = step.readout; +const gdnStep = step.gdnStep; +const attnStep = step.attnStep; +const moeStep = step.moeStep; +const pleStep = step.pleStep; +const Layer = load.Layer; +const loadLayer = load.loadLayer; + +pub fn argmax(raw: []const u8) u32 { var best = promote(std.mem.readInt(u16, raw[0..2], .little)); var at: usize = 0; var i: usize = 1; @@ -838,6 +87,10 @@ const Pass = struct { tail: cuda.DeviceBuffer, mixer: HcW, head_q: Q4, + logits: []u8, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, }; fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { @@ -846,170 +99,7 @@ fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); } -fn oneLayer(comptime Tri: type, p: *Pass, tri: Tri, layer: usize, ids: []const u32, pos0: usize) !void { - var buf: [200]u8 = undefined; - const base = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{layer}); - var base_buf: [80]u8 = undefined; - @memcpy(base_buf[0..base.len], base); - const root = base_buf[0..base.len]; - var attn_hc = try loadHc(p.gpa, p.driver, p.store, try std.fmt.bufPrint(&buf, "{s}.attn_hyper_connection", .{root}), true); - defer attn_hc.free(); - var mlp_hc = try loadHc(p.gpa, p.driver, p.store, try std.fmt.bufPrint(&buf, "{s}.mlp_hyper_connection", .{root}), true); - defer mlp_hc.free(); - const mlp = try std.fmt.bufPrint(&buf, "{s}.mlp", .{root}); - var mlp_buf: [96]u8 = undefined; - @memcpy(mlp_buf[0..mlp.len], mlp); - const mlp_root = mlp_buf[0..mlp.len]; - const gate = try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.gate.weight", .{mlp_root})); - if (!gate.is(.bf16, &.{ experts_n, dims })) return error.UnexpectedTensor; - const gate_copy = try p.gpa.alloc(u8, gate.bytes.len); - defer p.gpa.free(gate_copy); - @memcpy(gate_copy, gate.bytes); - const sg = try join3(p.gpa, p.store, &buf, mlp_root, ".shared_expert_gate"); - defer p.gpa.free(sg.w); - defer p.gpa.free(sg.s); - defer p.gpa.free(sg.b); - const router = try p.gpa.alloc(u16, (experts_n + 1) * dims); - defer p.gpa.free(router); - @memcpy(std.mem.sliceAsBytes(router[0 .. experts_n * dims]), gate_copy); - try embed.dequant(sg.w, sg.s, sg.b, dims, 0, router[experts_n * dims ..][0..dims]); - var router_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(router)); - defer router_b.free(); - - const linear = linearLayer(layer); - var proj_q: Q4 = undefined; - var out_q: Q4 = undefined; - var cw_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var a_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var dt_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var nw_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var qs_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var ks_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var is_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var key_q = Q4.empty(p.driver); - var val_q = Q4.empty(p.driver); - var nk_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var nq_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var nc_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - var pcw_b: cuda.DeviceBuffer = .{ .d = p.driver, .ptr = 0, .len = 0 }; - defer if (cw_b.ptr != 0) cw_b.free(); - defer if (a_b.ptr != 0) a_b.free(); - defer if (dt_b.ptr != 0) dt_b.free(); - defer if (nw_b.ptr != 0) nw_b.free(); - defer if (qs_b.ptr != 0) qs_b.free(); - defer if (ks_b.ptr != 0) ks_b.free(); - defer if (is_b.ptr != 0) is_b.free(); - defer if (nk_b.ptr != 0) nk_b.free(); - defer if (nq_b.ptr != 0) nq_b.free(); - defer if (nc_b.ptr != 0) nc_b.free(); - defer if (pcw_b.ptr != 0) pcw_b.free(); - if (linear) { - const pre = try std.fmt.bufPrint(&buf, "{s}.linear_attn", .{root}); - var pre_buf: [120]u8 = undefined; - @memcpy(pre_buf[0..pre.len], pre); - const lp = pre_buf[0..pre.len]; - const parts = [_][]const u8{ ".in_proj_qkv", ".in_proj_z", ".in_proj_b", ".in_proj_a" }; - const stitched = try stitch(p.gpa, p.store, lp, &parts); - defer p.gpa.free(stitched.w); - defer p.gpa.free(stitched.s); - defer p.gpa.free(stitched.b); - if (stitched.w.len != gdn.proj_width * (dims / 8) * 4) return error.UnexpectedTensor; - proj_q = try uploadQ4(p.gpa, p.driver, stitched.w, stitched.s, stitched.b, gdn.proj_width, dims / 8); - const out = try join3(p.gpa, p.store, &buf, lp, ".out_proj"); - defer p.gpa.free(out.w); - defer p.gpa.free(out.s); - defer p.gpa.free(out.b); - out_q = try uploadQ4(p.gpa, p.driver, out.w, out.s, out.b, dims, gdn.value_dim / 8); - const conv = try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); - if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; - cw_b = try cuda.DeviceBuffer.fromHost(p.driver, conv.bytes); - const a_log = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); - defer p.gpa.free(a_log); - const dt = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.dt_bias", .{lp}))); - defer p.gpa.free(dt); - if (a_log.len != gdn.nv or dt.len != gdn.nv) return error.UnexpectedTensor; - a_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(a_log)); - dt_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(dt)); - const norm = try p.store.copyOf(try std.fmt.bufPrint(&buf, "{s}.norm.weight", .{lp})); - defer p.gpa.free(norm); - nw_b = try cuda.DeviceBuffer.fromHost(p.driver, norm); - } else { - const pre = try std.fmt.bufPrint(&buf, "{s}.self_attn", .{root}); - var pre_buf: [120]u8 = undefined; - @memcpy(pre_buf[0..pre.len], pre); - const ap = pre_buf[0..pre.len]; - const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; - const stitched = try stitch(p.gpa, p.store, ap, &parts); - defer p.gpa.free(stitched.w); - defer p.gpa.free(stitched.s); - defer p.gpa.free(stitched.b); - if (stitched.w.len != proj_n * (dims / 8) * 4) return error.UnexpectedTensor; - proj_q = try uploadQ4(p.gpa, p.driver, stitched.w, stitched.s, stitched.b, proj_n, dims / 8); - const out = try join3(p.gpa, p.store, &buf, ap, ".o_proj"); - defer p.gpa.free(out.w); - defer p.gpa.free(out.s); - defer p.gpa.free(out.b); - out_q = try uploadQ4(p.gpa, p.driver, out.w, out.s, out.b, dims, out_k / 8); - const qn = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); - defer p.gpa.free(qn); - const kn = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); - defer p.gpa.free(kn); - const iqn = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); - defer p.gpa.free(iqn); - if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim) return error.UnexpectedTensor; - qs_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(qn)); - ks_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(kn)); - is_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(iqn)); - } - defer proj_q.free(); - defer out_q.free(); - if (layer == ple_layer) { - const pre = try std.fmt.bufPrint(&buf, "{s}.ple", .{root}); - var pre_buf: [96]u8 = undefined; - @memcpy(pre_buf[0..pre.len], pre); - const pp = pre_buf[0..pre.len]; - const key = try join3(p.gpa, p.store, &buf, pp, ".key_proj"); - defer p.gpa.free(key.w); - defer p.gpa.free(key.s); - defer p.gpa.free(key.b); - const val = try join3(p.gpa, p.store, &buf, pp, ".value_proj"); - defer p.gpa.free(val.w); - defer p.gpa.free(val.s); - defer p.gpa.free(val.b); - key_q = try uploadQ4(p.gpa, p.driver, key.w, key.s, key.b, wide, dims / 8); - val_q = try uploadQ4(p.gpa, p.driver, val.w, val.s, val.b, dims, dims / 8); - const nk = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); - defer p.gpa.free(nk); - const nq = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); - defer p.gpa.free(nq); - const nc = try asF32(p.gpa, try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_conv.weight", .{pp}))); - defer p.gpa.free(nc); - if (nk.len != wide or nq.len != wide or nc.len != wide) return error.UnexpectedTensor; - nk_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(nk)); - nq_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(nq)); - nc_b = try cuda.DeviceBuffer.fromHost(p.driver, std.mem.sliceAsBytes(nc)); - const conv = try p.store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); - if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; - pcw_b = try cuda.DeviceBuffer.fromHost(p.driver, conv.bytes); - } - defer key_q.free(); - defer val_q.free(); - - // Shared experts are copied. Routed rows stay mapped, so this is the last shard switch. - var sh_g = try ownedFace(p.gpa, p.store, &buf, mlp_root, ".shared_expert.gate_proj", moe_w, dims); - defer sh_g.free(p.gpa); - var sh_u = try ownedFace(p.gpa, p.store, &buf, mlp_root, ".shared_expert.up_proj", moe_w, dims); - defer sh_u.free(p.gpa); - var sh_d = try ownedFace(p.gpa, p.store, &buf, mlp_root, ".shared_expert.down_proj", dims, moe_w); - defer sh_d.free(p.gpa); - p.store.release(); - var faces: [3]Face = undefined; - faces[0] = try faceOf(p.store, &buf, mlp_root, ".switch_mlp.gate_proj", true, moe_w, dims); - faces[1] = try faceOf(p.store, &buf, mlp_root, ".switch_mlp.up_proj", true, moe_w, dims); - faces[2] = try faceOf(p.store, &buf, mlp_root, ".switch_mlp.down_proj", true, dims, moe_w); - const sh = [3]Face{ sh_g.face, sh_u.face, sh_d.face }; - - std.debug.print("layer {d}\n", .{layer}); +fn runLayer(comptime Tri: type, p: *Pass, tri: Tri, lw: *Layer, layer: usize, ids: []const u32, pos0: usize) !void { for (ids, 0..) |token, t| { const h = p.s.h.ptr + @as(u64, t) * wide * 2; const inj = p.s.inj.ptr + @as(u64, t) * inj_stride; @@ -1019,51 +109,24 @@ fn oneLayer(comptime Tri: type, p: *Pass, tri: Tri, layer: usize, ids: []const u try writeback(tri, p.s, h, 2, inj, y, wt); p.mode[t] = 0; } - if (layer == ple_layer) try pleStep(Tri, p.gpa, p.io, p.driver, p.stream, tri, p.s, p.model_dir, p.history, token, h, key_q, val_q, nk_b.ptr, nq_b.ptr, nc_b.ptr, pcw_b.ptr, &p.tail); + if (layer == ple_layer) try pleStep(Tri, p.gpa, p.io, p.driver, p.stream, tri, p.s, p.model_dir, p.history, token, h, lw.key, lw.val, lw.nk.ptr, lw.nq.ptr, lw.nc.ptr, lw.ple_cw.ptr, &p.tail); if (p.mode[t] == 2) { try writeback(tri, p.s, h, 2, inj, y, wt); p.mode[t] = 0; } else try writeback(tri, p.s, h, 0, inj, y, wt); - try readout(tri, p.s, attn_hc, h, inj); - if (linear) { + try readout(tri, p.s, lw.attn_hc, h, inj); + if (lw.linear) { const li = linearIndex(layer); - try gdnStep(p.gpa, p.driver, p.stream, p.s, proj_q, out_q, p.conv[li], cw_b.ptr, a_b.ptr, dt_b.ptr, nw_b.ptr, &p.state[li]); + try gdnStep(p.gpa, p.driver, p.stream, p.s, lw.proj, lw.out, p.conv[li], lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, lw.norm.ptr, &p.state[li]); } else { const ai = attnIndex(layer); - try attnStep(Tri, p.driver, p.stream, tri, p.s, proj_q, out_q, qs_b.ptr, ks_b.ptr, is_b.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); + try attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); } try writeback(tri, p.s, h, 4, inj, y, wt); - try readout(tri, p.s, mlp_hc, h, inj); - try moeStep(Tri, p.gpa, p.driver, p.stream, tri, p.s, router_b.ptr, faces, sh, y, wt); + try readout(tri, p.s, lw.mlp_hc, h, inj); + try moeStep(Tri, tri, p.s, p.stream, p.plan_fn, p.up_fn, p.down_fn, lw.router.ptr, lw.experts_up.ptr, lw.experts_down.ptr, y, wt); p.mode[t] = 2; - if (t % 16 == 0) std.debug.print("layer {d} row {d}\n", .{ layer, t }); } - p.store.release(); -} - -const OwnedFace = struct { - face: Face, - w: []u8, - s: []u8, - b: []u8, - - fn free(self: *OwnedFace, gpa: std.mem.Allocator) void { - gpa.free(self.w); - gpa.free(self.s); - gpa.free(self.b); - } -}; - -fn ownedFace(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8, n: usize, k: usize) !OwnedFace { - const got = try join3(gpa, store, buf, prefix, name); - const k8 = k / 8; - if (got.w.len != n * k8 * 4 or got.s.len != n * (k / 32) * 2) { - gpa.free(got.w); - gpa.free(got.s); - gpa.free(got.b); - return error.UnexpectedTensor; - } - return .{ .face = .{ .w = got.w, .s = got.s, .b = got.b, .n = n, .k = k, .routed = false }, .w = got.w, .s = got.s, .b = got.b }; } fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![]cuda.DeviceBuffer { @@ -1080,7 +143,7 @@ fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![] return out; } -fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !u32 { +fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !void { if (p.mode[row] != 2) return error.UnexpectedTensor; const h = p.s.h.ptr + @as(u64, row) * wide * 2; const inj = p.s.inj.ptr + @as(u64, row) * inj_stride; @@ -1090,108 +153,240 @@ fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !u32 { p.mode[row] = 0; try readout(tri, p.s, p.mixer, h, inj); try qmm.matmul(p.driver, p.stream.*, p.s.mixed.ptr, p.s.xsm.ptr, p.head_q.w.ptr, p.head_q.s.ptr, p.head_q.b.ptr, p.s.logits.ptr, 1, vocab, dims); - const raw = try p.gpa.alloc(u8, vocab * 2); - defer p.gpa.free(raw); try p.stream.synchronize(); - try p.s.logits.download(0, raw); - return argmax(raw); + try p.s.logits.download(0, p.logits); } -fn forward(comptime Tri: type, p: *Pass, tri: Tri, ids: []const u32, pos0: usize) !u32 { +fn forward(comptime Tri: type, p: *Pass, tri: Tri, layers: []Layer, ids: []const u32, pos0: usize) !void { @memset(p.mode, 0); for (ids, 0..) |token, t| try embedRow(Tri, p, tri, token, p.s.h.ptr + @as(u64, t) * wide * 2); - for (0..layers_n) |layer| try oneLayer(Tri, p, tri, layer, ids, pos0); - return head(Tri, p, tri, ids.len - 1); + for (layers, 0..) |*lw, layer| try runLayer(Tri, p, tri, lw, layer, ids, pos0); + try head(Tri, p, tri, ids.len - 1); } -/// `out` receives `new_tokens` greedy ids. The first is the prompt's argmax. The rest are fed back one at a time. -pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { - if (!cuda.kernels.available) return error.BuiltWithoutKernels; - if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > 256) return error.UnexpectedTensor; - var store = try Store.open(gpa, io, model_dir); - defer store.deinit(); - const weight = try store.tensor(embed.weight_name); - const scales = try store.tensor(embed.scales_name); - const biases = try store.tensor(embed.biases_name); - var table = try embed.Table.upload(driver, weight, scales, biases); - defer table.deinit(); - const eos = try ple.eosOf(gpa, io, model_dir); - var history = [2]i64{ eos, eos }; - var scratch = try Scratch.alloc(driver, tokens.len); - defer scratch.free(); - try scratch.dummy.fill8(0, null); - const conv = try zeros(gpa, driver, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); - defer { - for (conv) |*b| b.free(); - gpa.free(conv); - } - const state = try zeros(gpa, driver, linear_n, state_n * 4); - defer { - for (state) |*b| b.free(); - gpa.free(state); - } - const cache = tokens.len + new_tokens; - const kc = try zeros(gpa, driver, attn_n, cache * kv_heads * head_dim * 2); - defer { - for (kc) |*b| b.free(); - gpa.free(kc); +pub const vocabulary: usize = vocab; + +/// One loaded checkpoint: the embedding table, the mixer, the head, and a scratch row the forwards share. +pub const Engine = struct { + gpa: std.mem.Allocator, + io: std.Io, + driver: *cuda.Driver, + stream: *cuda.Stream, + store: Store, + model_dir: []u8, + table: embed.Table, + mixer: HcW, + head_q: Q4, + scratch: Scratch, + scratch_rows: usize, + mode: []u8, + eos: i64, + experts: cuda.Module, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, + layers: [layers_n]Layer, + n_layers: usize, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.gpa = gpa; + e.io = io; + e.driver = driver; + e.stream = stream; + e.n_layers = 0; + errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); + e.model_dir = try gpa.dupe(u8, model_dir); + errdefer gpa.free(e.model_dir); + e.store = try Store.open(gpa, io, e.model_dir); + errdefer e.store.deinit(); + const weight = try e.store.tensor(embed.weight_name); + const scales = try e.store.tensor(embed.scales_name); + const biases = try e.store.tensor(embed.biases_name); + e.table = try embed.Table.upload(driver, weight, scales, biases); + errdefer e.table.deinit(); + e.eos = try ple.eosOf(gpa, io, e.model_dir); + e.scratch = try Scratch.alloc(driver, 1); + errdefer e.scratch.free(); + try e.scratch.dummy.fill8(0, null); + e.scratch_rows = 1; + e.mode = try gpa.alloc(u8, 1); + errdefer gpa.free(e.mode); + e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); + errdefer e.mixer.free(); + var buf: [80]u8 = undefined; + const head_w = try join3(gpa, &e.store, &buf, "language_model.lm_head", ""); + defer gpa.free(head_w.w); + defer gpa.free(head_w.s); + defer gpa.free(head_w.b); + if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; + e.head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); + errdefer e.head_q.free(); + e.experts = try cuda.Module.load(driver, cuda.kernels.experts); + errdefer e.experts.unload(); + e.plan_fn = try e.experts.function(weights.plan_symbol); + e.up_fn = try e.experts.function(weights.up_symbol); + e.down_fn = try e.experts.function(weights.down_symbol); + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + const pack_fn = try pack_mod.function(weights.pack_symbol); + for (0..layers_n) |i| { + e.layers[i] = try loadLayer(gpa, driver, stream, &e.store, pack_fn, i); + e.n_layers += 1; + } + e.store.release(); + e.store.closeFile(); + return e; + } + + pub fn deinit(e: *Engine) void { + e.experts.unload(); + for (e.layers[0..e.n_layers]) |*l| l.free(e.gpa); + e.head_q.free(); + e.mixer.free(); + e.gpa.free(e.mode); + e.scratch.free(); + e.table.deinit(); + e.store.deinit(); + e.gpa.free(e.model_dir); + e.gpa.destroy(e); + } + + fn grow(e: *Engine, rows: usize) !void { + if (rows <= e.scratch_rows) return; + var next = try Scratch.alloc(e.driver, rows); + errdefer next.free(); + try next.dummy.fill8(0, null); + const mode = try e.gpa.alloc(u8, rows); + e.scratch.free(); + e.gpa.free(e.mode); + e.scratch = next; + e.mode = mode; + e.scratch_rows = rows; + } + + /// A fresh cache. `cap` is how many tokens it can hold, prompt plus generated. + pub fn newSeq(e: *Engine, cap: usize) !Seq { + if (cap == 0) return error.UnexpectedTensor; + const logits = try e.gpa.alloc(u8, vocab * 2); + errdefer e.gpa.free(logits); + const conv = try zeros(e.gpa, e.driver, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); + errdefer freeOwned(e.gpa, conv); + const state = try zeros(e.gpa, e.driver, linear_n, state_n * 4); + errdefer freeOwned(e.gpa, state); + const kc = try zeros(e.gpa, e.driver, attn_n, cap * kv_heads * head_dim * 2); + errdefer freeOwned(e.gpa, kc); + const vc = try zeros(e.gpa, e.driver, attn_n, cap * kv_heads * head_dim * 2); + errdefer freeOwned(e.gpa, vc); + const ikc = try zeros(e.gpa, e.driver, attn_n, cap * index_dim * 2); + errdefer freeOwned(e.gpa, ikc); + var tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); + errdefer tail.free(); + try tail.fill8(0, null); + return .{ + .conv = conv, + .state = state, + .kc = kc, + .vc = vc, + .ikc = ikc, + .tail = tail, + .history = .{ e.eos, e.eos }, + .fed = 0, + .logits = logits, + .cap = cap, + }; } - const vc = try zeros(gpa, driver, attn_n, cache * kv_heads * head_dim * 2); - defer { - for (vc) |*b| b.free(); - gpa.free(vc); + + /// Run `ids` on from `seq.fed` and leave the next token's scores in `seq.logits`. + pub fn predict(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32) !void { + if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + try e.grow(ids.len); + var pass = Pass{ + .gpa = e.gpa, + .io = e.io, + .driver = e.driver, + .stream = e.stream, + .store = &e.store, + .model_dir = e.model_dir, + .s = &e.scratch, + .table = e.table, + .rows = ids.len, + .mode = e.mode, + .history = &seq.history, + .conv = seq.conv, + .state = seq.state, + .kc = seq.kc, + .vc = seq.vc, + .ikc = seq.ikc, + .tail = seq.tail, + .mixer = e.mixer, + .head_q = e.head_q, + .logits = seq.logits, + .plan_fn = e.plan_fn, + .up_fn = e.up_fn, + .down_fn = e.down_fn, + }; + try forward(Tri, &pass, tri, &e.layers, ids, seq.fed); + seq.fed += ids.len; } - const ikc = try zeros(gpa, driver, attn_n, cache * index_dim * 2); - defer { - for (ikc) |*b| b.free(); - gpa.free(ikc); +}; + +/// One stream's caches. `fed` tokens have already gone through the model. +pub const Seq = struct { + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + history: [2]i64, + fed: usize, + logits: []u8, + cap: usize, + + pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { + freeBufs(s.conv); + freeBufs(s.state); + freeBufs(s.kc); + freeBufs(s.vc); + freeBufs(s.ikc); + gpa.free(s.conv); + gpa.free(s.state); + gpa.free(s.kc); + gpa.free(s.vc); + gpa.free(s.ikc); + if (s.tail.ptr != 0) s.tail.free(); + gpa.free(s.logits); } - var tail = try cuda.DeviceBuffer.alloc(driver, ple_tail_n * wide * 2); - defer tail.free(); - try tail.fill8(0, null); - var mixer = try loadHc(gpa, driver, &store, "language_model.model.hyper_connection_mixer", false); - defer mixer.free(); - var buf: [80]u8 = undefined; - const head_w = try join3(gpa, &store, &buf, "language_model.lm_head", ""); - defer gpa.free(head_w.w); - defer gpa.free(head_w.s); - defer gpa.free(head_w.b); - if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; - var head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); - defer head_q.free(); - const mode = try gpa.alloc(u8, tokens.len); - defer gpa.free(mode); - var pass = Pass{ - .gpa = gpa, - .io = io, - .driver = driver, - .stream = stream, - .store = &store, - .model_dir = model_dir, - .s = &scratch, - .table = table, - .rows = tokens.len, - .mode = mode, - .history = &history, - .conv = conv, - .state = state, - .kc = kc, - .vc = vc, - .ikc = ikc, - .tail = tail, - .mixer = mixer, - .head_q = head_q, - }; - out[0] = try forward(Tri, &pass, tri, tokens, 0); +}; + +fn freeBufs(bs: []cuda.DeviceBuffer) void { + for (bs) |*b| b.free(); +} + +fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { + freeBufs(bs); + gpa.free(bs); +} + +/// `out` receives `new_tokens` greedy ids. The first is the prompt's argmax. The rest are fed back one at a time. +pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { + if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > 256) return error.UnexpectedTensor; + var e = try Engine.init(gpa, io, driver, stream, model_dir); + defer e.deinit(); + var seq = try e.newSeq(tokens.len + new_tokens); + defer seq.deinit(gpa); + try e.predict(Tri, tri, &seq, tokens); + out[0] = argmax(seq.logits); std.debug.print("token {d}\n", .{out[0]}); var tok = out[0]; - var pos: usize = tokens.len; for (1..new_tokens) |i| { var one = [_]u32{tok}; - // The new token reuses row 0. Clear every row so a stale prompt pending is not written back. - tok = try forward(Tri, &pass, tri, &one, pos); + try e.predict(Tri, tri, &seq, &one); + tok = argmax(seq.logits); out[i] = tok; - pos += 1; std.debug.print("token {d}\n", .{tok}); } } diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig new file mode 100644 index 0000000000..f7192ce768 --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -0,0 +1,389 @@ +//! One resident copy of every decoder layer. The prompt forward reads these, it does not reload them. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const gdn = @import("cuda_gdn.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +const layers_n = wt.layers_n; +const dims = wt.dims; +const wide = wt.wide; +const experts_n = wt.experts_n; +const moe_w = wt.moe_w; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const ple_layer = wt.ple_layer; +const block_bytes = wt.block_bytes; +const cint = wt.cint; +const Store = wt.Store; +const Face = wt.Face; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const linearLayer = wt.linearLayer; +const asF32 = wt.asF32; +const uploadQ4 = wt.uploadQ4; +const join3 = wt.join3; +const loadHc = wt.loadHc; +const stitch = wt.stitch; + +const OwnedFace = struct { + face: Face, + w: []u8, + s: []u8, + b: []u8, + + fn free(self: *OwnedFace, gpa: std.mem.Allocator) void { + gpa.free(self.w); + gpa.free(self.s); + gpa.free(self.b); + } +}; + +pub fn faceOf(store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !Face { + const w = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + const sc = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + const b = try store.hold(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + const k8 = k / 8; + const kg = k / 32; + if (routed) { + if (!w.is(.u32, &.{ experts_n, n, k8 }) or !sc.is(.bf16, &.{ experts_n, n, kg }) or !b.is(.bf16, &.{ experts_n, n, kg })) return error.UnexpectedTensor; + } else if (!w.is(.u32, &.{ n, k8 }) or !sc.is(.bf16, &.{ n, kg }) or !b.is(.bf16, &.{ n, kg })) return error.UnexpectedTensor; + return .{ .w = w.bytes, .s = sc.bytes, .b = b.bytes, .n = n, .k = k, .routed = routed }; +} + +pub const Layer = struct { + attn_hc: HcW, + mlp_hc: HcW, + proj: Q4, + out: Q4, + router: cuda.DeviceBuffer, + linear: bool, + cw: cuda.DeviceBuffer, + a_log: cuda.DeviceBuffer, + dt: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + q_scale: cuda.DeviceBuffer, + k_scale: cuda.DeviceBuffer, + i_scale: cuda.DeviceBuffer, + key: Q4, + val: Q4, + nk: cuda.DeviceBuffer, + nq: cuda.DeviceBuffer, + nc: cuda.DeviceBuffer, + ple_cw: cuda.DeviceBuffer, + /// Packed gate and up for every routed expert plus the shared expert at index 512. + experts_up: cuda.DeviceBuffer, + experts_down: cuda.DeviceBuffer, + + pub fn free(self: *Layer, gpa: std.mem.Allocator) void { + _ = gpa; + self.attn_hc.free(); + self.mlp_hc.free(); + self.proj.free(); + self.out.free(); + self.router.free(); + self.cw.free(); + self.a_log.free(); + self.dt.free(); + self.norm.free(); + self.q_scale.free(); + self.k_scale.free(); + self.i_scale.free(); + self.key.free(); + self.val.free(); + self.nk.free(); + self.nq.free(); + self.nc.free(); + self.ple_cw.free(); + self.experts_up.free(); + self.experts_down.free(); + } +}; + +fn deadBuf(d: *cuda.Driver) cuda.DeviceBuffer { + return .{ .d = d, .ptr = 0, .len = 0 }; +} + +fn blankFace() OwnedFace { + return .{ .face = .{ .w = &.{}, .s = &.{}, .b = &.{}, .n = 0, .k = 0, .routed = false }, .w = &.{}, .s = &.{}, .b = &.{} }; +} + +/// Copy one expert tensor out of the checkpoint so the shard can be closed. +fn keepFace(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8, routed: bool, n: usize, k: usize) !OwnedFace { + const view = try faceOf(store, buf, prefix, name, routed, n, k); + const w = try gpa.dupe(u8, view.w); + errdefer gpa.free(w); + const s = try gpa.dupe(u8, view.s); + errdefer gpa.free(s); + const b = try gpa.dupe(u8, view.b); + return .{ .face = .{ .w = w, .s = s, .b = b, .n = n, .k = k, .routed = routed }, .w = w, .s = s, .b = b }; +} + +fn packDevice(driver: *cuda.Driver, stream: *cuda.Stream, pack_fn: cuda.Function, words: []const u8, scales: []const u8, biases: []const u8, e: usize, n: usize, k: usize) !cuda.DeviceBuffer { + const kg = k / 32; + const nb = n / 32; + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, e * nb * kg * block_bytes); + errdefer out_b.free(); + var a: cuda.Args = .{}; + a.add(w_b.ptr); + a.add(s_b.ptr); + a.add(b_b.ptr); + a.add(out_b.ptr); + a.add(cint(n)); + a.add(cint(k / 8)); + a.add(cint(kg)); + a.add(cint(nb)); + try cuda.launch.launch(pack_fn, .{ .grid = .{ .x = @intCast(kg), .y = @intCast(nb), .z = @intCast(e) }, .block = .{ .x = 160 } }, stream.*, &a); + try stream.synchronize(); + return out_b; +} + +/// Routed experts, then the shared expert as index `experts_n`, in the kernel's packed layout. +fn packAll(driver: *cuda.Driver, stream: *cuda.Stream, pack_fn: cuda.Function, routed: OwnedFace, shared: OwnedFace) !cuda.DeviceBuffer { + const n = routed.face.n; + const k = routed.face.k; + var bulk = try packDevice(driver, stream, pack_fn, routed.w, routed.s, routed.b, experts_n, n, k); + defer bulk.free(); + var one = try packDevice(driver, stream, pack_fn, shared.w, shared.s, shared.b, 1, n, k); + defer one.free(); + if (bulk.len != experts_n * one.len) return error.UnexpectedTensor; + var out = try cuda.DeviceBuffer.alloc(driver, bulk.len + one.len); + errdefer out.free(); + try out.copyFrom(0, bulk.ptr, bulk.len, stream.handle); + try out.copyFrom(bulk.len, one.ptr, one.len, stream.handle); + try stream.synchronize(); + return out; +} + +fn stackHalves(gpa: std.mem.Allocator, gate: []const u8, up: []const u8) ![]u8 { + if (gate.len != up.len or gate.len % block_bytes != 0) return error.UnexpectedTensor; + const nblk = gate.len / block_bytes; + const out = try gpa.alloc(u8, gate.len + up.len); + for (0..nblk) |b| { + @memcpy(out[b * 2 * block_bytes ..][0..block_bytes], gate[b * block_bytes ..][0..block_bytes]); + @memcpy(out[b * 2 * block_bytes + block_bytes ..][0..block_bytes], up[b * block_bytes ..][0..block_bytes]); + } + return out; +} + +/// Gate and up blocks side by side, which is the SwiGLU weight the expert kernel reads. +fn stackUp(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, gate: cuda.DeviceBuffer, up: cuda.DeviceBuffer) !cuda.DeviceBuffer { + const graw = try gpa.alloc(u8, gate.len); + defer gpa.free(graw); + const uraw = try gpa.alloc(u8, up.len); + defer gpa.free(uraw); + try stream.synchronize(); + try gate.download(0, graw); + try up.download(0, uraw); + const stacked = try stackHalves(gpa, graw, uraw); + defer gpa.free(stacked); + return cuda.DeviceBuffer.fromHost(driver, stacked); +} + +pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, layer: usize) !Layer { + var buf: [200]u8 = undefined; + const base = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{layer}); + var base_buf: [80]u8 = undefined; + @memcpy(base_buf[0..base.len], base); + const root = base_buf[0..base.len]; + var attn_hc = try loadHc(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.attn_hyper_connection", .{root}), true); + errdefer attn_hc.free(); + var mlp_hc = try loadHc(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.mlp_hyper_connection", .{root}), true); + errdefer mlp_hc.free(); + const mlp = try std.fmt.bufPrint(&buf, "{s}.mlp", .{root}); + var mlp_buf: [96]u8 = undefined; + @memcpy(mlp_buf[0..mlp.len], mlp); + const mlp_root = mlp_buf[0..mlp.len]; + const gate = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.gate.weight", .{mlp_root})); + if (!gate.is(.bf16, &.{ experts_n, dims })) return error.UnexpectedTensor; + const gate_copy = try gpa.alloc(u8, gate.bytes.len); + defer gpa.free(gate_copy); + @memcpy(gate_copy, gate.bytes); + const sg = try join3(gpa, store, &buf, mlp_root, ".shared_expert_gate"); + defer gpa.free(sg.w); + defer gpa.free(sg.s); + defer gpa.free(sg.b); + const router = try gpa.alloc(u16, (experts_n + 1) * dims); + defer gpa.free(router); + @memcpy(std.mem.sliceAsBytes(router[0 .. experts_n * dims]), gate_copy); + try embed.dequant(sg.w, sg.s, sg.b, dims, 0, router[experts_n * dims ..][0..dims]); + var router_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); + errdefer router_b.free(); + + const linear = linearLayer(layer); + var proj_q = Q4.empty(driver); + var out_q = Q4.empty(driver); + errdefer proj_q.free(); + errdefer out_q.free(); + var cw_b = deadBuf(driver); + var a_b = deadBuf(driver); + var dt_b = deadBuf(driver); + var nw_b = deadBuf(driver); + var qs_b = deadBuf(driver); + var ks_b = deadBuf(driver); + var is_b = deadBuf(driver); + var key_q = Q4.empty(driver); + var val_q = Q4.empty(driver); + var nk_b = deadBuf(driver); + var nq_b = deadBuf(driver); + var nc_b = deadBuf(driver); + var pcw_b = deadBuf(driver); + errdefer cw_b.free(); + errdefer a_b.free(); + errdefer dt_b.free(); + errdefer nw_b.free(); + errdefer qs_b.free(); + errdefer ks_b.free(); + errdefer is_b.free(); + errdefer nk_b.free(); + errdefer nq_b.free(); + errdefer nc_b.free(); + errdefer pcw_b.free(); + if (linear) { + const pre = try std.fmt.bufPrint(&buf, "{s}.linear_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const lp = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".in_proj_qkv", ".in_proj_z", ".in_proj_b", ".in_proj_a" }; + const stitched = try stitch(gpa, store, lp, &parts); + defer gpa.free(stitched.w); + defer gpa.free(stitched.s); + defer gpa.free(stitched.b); + if (stitched.w.len != gdn.proj_width * (dims / 8) * 4) return error.UnexpectedTensor; + proj_q = try uploadQ4(gpa, driver, stitched.w, stitched.s, stitched.b, gdn.proj_width, dims / 8); + const out = try join3(gpa, store, &buf, lp, ".out_proj"); + defer gpa.free(out.w); + defer gpa.free(out.s); + defer gpa.free(out.b); + out_q = try uploadQ4(gpa, driver, out.w, out.s, out.b, dims, gdn.value_dim / 8); + const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); + if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; + cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + const a_log = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); + defer gpa.free(a_log); + const dt = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.dt_bias", .{lp}))); + defer gpa.free(dt); + if (a_log.len != gdn.nv or dt.len != gdn.nv) return error.UnexpectedTensor; + a_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(a_log)); + dt_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(dt)); + const norm = try store.copyOf(try std.fmt.bufPrint(&buf, "{s}.norm.weight", .{lp})); + defer gpa.free(norm); + nw_b = try cuda.DeviceBuffer.fromHost(driver, norm); + } else { + const pre = try std.fmt.bufPrint(&buf, "{s}.self_attn", .{root}); + var pre_buf: [120]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const ap = pre_buf[0..pre.len]; + const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; + const stitched = try stitch(gpa, store, ap, &parts); + defer gpa.free(stitched.w); + defer gpa.free(stitched.s); + defer gpa.free(stitched.b); + if (stitched.w.len != proj_n * (dims / 8) * 4) return error.UnexpectedTensor; + proj_q = try uploadQ4(gpa, driver, stitched.w, stitched.s, stitched.b, proj_n, dims / 8); + const out = try join3(gpa, store, &buf, ap, ".o_proj"); + defer gpa.free(out.w); + defer gpa.free(out.s); + defer gpa.free(out.b); + out_q = try uploadQ4(gpa, driver, out.w, out.s, out.b, dims, out_k / 8); + const qn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); + defer gpa.free(qn); + const kn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); + defer gpa.free(kn); + const iqn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); + defer gpa.free(iqn); + if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim) return error.UnexpectedTensor; + qs_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(qn)); + ks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(kn)); + is_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(iqn)); + } + if (layer == ple_layer) { + const pre = try std.fmt.bufPrint(&buf, "{s}.ple", .{root}); + var pre_buf: [96]u8 = undefined; + @memcpy(pre_buf[0..pre.len], pre); + const pp = pre_buf[0..pre.len]; + const key = try join3(gpa, store, &buf, pp, ".key_proj"); + defer gpa.free(key.w); + defer gpa.free(key.s); + defer gpa.free(key.b); + const val = try join3(gpa, store, &buf, pp, ".value_proj"); + defer gpa.free(val.w); + defer gpa.free(val.s); + defer gpa.free(val.b); + key_q = try uploadQ4(gpa, driver, key.w, key.s, key.b, wide, dims / 8); + val_q = try uploadQ4(gpa, driver, val.w, val.s, val.b, dims, dims / 8); + const nk = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); + defer gpa.free(nk); + const nq = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); + defer gpa.free(nq); + const nc = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_conv.weight", .{pp}))); + defer gpa.free(nc); + if (nk.len != wide or nq.len != wide or nc.len != wide) return error.UnexpectedTensor; + nk_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nk)); + nq_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nq)); + nc_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nc)); + const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); + if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; + pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); + } + errdefer key_q.free(); + errdefer val_q.free(); + + var shared: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&shared) |*f| f.free(gpa); + shared[0] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.gate_proj", false, moe_w, dims); + shared[1] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.up_proj", false, moe_w, dims); + shared[2] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.down_proj", false, dims, moe_w); + var routed: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&routed) |*f| f.free(gpa); + routed[0] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.gate_proj", true, moe_w, dims); + routed[1] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.up_proj", true, moe_w, dims); + routed[2] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.down_proj", true, dims, moe_w); + var gate_p = try packAll(driver, stream, pack_fn, routed[0], shared[0]); + errdefer gate_p.free(); + var up_p = try packAll(driver, stream, pack_fn, routed[1], shared[1]); + errdefer up_p.free(); + var down_p = try packAll(driver, stream, pack_fn, routed[2], shared[2]); + errdefer down_p.free(); + for (&shared) |*f| f.free(gpa); + for (&routed) |*f| f.free(gpa); + shared = .{ blankFace(), blankFace(), blankFace() }; + routed = .{ blankFace(), blankFace(), blankFace() }; + var experts_up = try stackUp(gpa, driver, stream, gate_p, up_p); + errdefer experts_up.free(); + gate_p.free(); + up_p.free(); + store.release(); + return .{ + .attn_hc = attn_hc, + .mlp_hc = mlp_hc, + .proj = proj_q, + .out = out_q, + .router = router_b, + .linear = linear, + .cw = cw_b, + .a_log = a_b, + .dt = dt_b, + .norm = nw_b, + .q_scale = qs_b, + .k_scale = ks_b, + .i_scale = is_b, + .key = key_q, + .val = val_q, + .nk = nk_b, + .nq = nq_b, + .nc = nc_b, + .ple_cw = pcw_b, + .experts_up = experts_up, + .experts_down = down_p, + }; +} diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig new file mode 100644 index 0000000000..de02eb2eef --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -0,0 +1,316 @@ +//! One fused decode row: hyper-connection, linear or full attention, experts, and the n-gram step. + +const std = @import("std"); +const cuda = @import("cuda"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const experts_n = wt.experts_n; +const slots = wt.slots; +const moe_w = wt.moe_w; +const block_bytes = wt.block_bytes; +const tile = wt.tile; +const q_heads = wt.q_heads; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_heads = wt.index_heads; +const index_dim = wt.index_dim; +const half = wt.half; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const nch = wt.nch; +const ple_tail_n = wt.ple_tail_n; +const state_n = wt.state_n; +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; +const inj_stride = wt.inj_stride; +const wts_stride = wt.wts_stride; +const vocab = wt.vocab; +const cint = wt.cint; +const toBf16 = wt.toBf16; +const pack_symbol = wt.pack_symbol; +const plan_symbol = wt.plan_symbol; +const up_symbol = wt.up_symbol; +const down_symbol = wt.down_symbol; +const Face = wt.Face; +const Q4 = wt.Q4; +const HcW = wt.HcW; + +pub const Scratch = struct { + h: cuda.DeviceBuffer, + y: cuda.DeviceBuffer, + wts: cuda.DeviceBuffer, + inj: cuda.DeviceBuffer, + pss: cuda.DeviceBuffer, + norm: cuda.DeviceBuffer, + dn: cuda.DeviceBuffer, + hc_part: cuda.DeviceBuffer, + act: cuda.DeviceBuffer, + xs_act: cuda.DeviceBuffer, + mixed: cuda.DeviceBuffer, + xsm: cuda.DeviceBuffer, + proj: cuda.DeviceBuffer, + part: cuda.DeviceBuffer, + gout: cuda.DeviceBuffer, + gxs: cuda.DeviceBuffer, + st: cuda.DeviceBuffer, + ple_emb: cuda.DeviceBuffer, + ple_xs: cuda.DeviceBuffer, + ple_keys: cuda.DeviceBuffer, + ple_vals: cuda.DeviceBuffer, + ple_gated: cuda.DeviceBuffer, + ple_pss: cuda.DeviceBuffer, + ple_nrow: cuda.DeviceBuffer, + ple_part: cuda.DeviceBuffer, + pa: cuda.DeviceBuffer, + q: cuda.DeviceBuffer, + iq: cuda.DeviceBuffer, + po: cuda.DeviceBuffer, + pm: cuda.DeviceBuffer, + pl: cuda.DeviceBuffer, + ao: cuda.DeviceBuffer, + gated: cuda.DeviceBuffer, + gxs_attn: cuda.DeviceBuffer, + dummy: cuda.DeviceBuffer, + inv: cuda.DeviceBuffer, + pos: cuda.DeviceBuffer, + ids: cuda.DeviceBuffer, + logits: cuda.DeviceBuffer, + router_l: cuda.DeviceBuffer, + picks: cuda.DeviceBuffer, + moe_members: cuda.DeviceBuffer, + moe_items: cuda.DeviceBuffer, + moe_counts: cuda.DeviceBuffer, + moe_act: cuda.DeviceBuffer, + + pub fn alloc(driver: *cuda.Driver, rows: usize) !Scratch { + var s: Scratch = undefined; + s.h = try cuda.DeviceBuffer.alloc(driver, rows * wide * 2); + s.y = try cuda.DeviceBuffer.alloc(driver, rows * slots * dims * 4); + s.wts = try cuda.DeviceBuffer.alloc(driver, rows * wts_stride); + s.inj = try cuda.DeviceBuffer.alloc(driver, rows * inj_stride); + s.pss = try cuda.DeviceBuffer.alloc(driver, (dims / 256) * streams * 4); + s.norm = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.dn = try cuda.DeviceBuffer.alloc(driver, (low + 4) * 2); + s.hc_part = try cuda.DeviceBuffer.alloc(driver, hc_sk * (low + 4) * 4); + s.act = try cuda.DeviceBuffer.alloc(driver, low * 2); + s.xs_act = try cuda.DeviceBuffer.alloc(driver, (low / 32) * 4); + s.mixed = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.xsm = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.proj = try cuda.DeviceBuffer.alloc(driver, gdn.proj_width * 2); + s.part = try cuda.DeviceBuffer.alloc(driver, out_sk * dims * 4); + s.gout = try cuda.DeviceBuffer.alloc(driver, gdn.value_dim * 2); + s.gxs = try cuda.DeviceBuffer.alloc(driver, (gdn.value_dim / 32) * 4); + s.st = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + s.ple_emb = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_xs = try cuda.DeviceBuffer.alloc(driver, (dims / 32) * 4); + s.ple_keys = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_vals = try cuda.DeviceBuffer.alloc(driver, dims * 2); + s.ple_gated = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_pss = try cuda.DeviceBuffer.alloc(driver, streams * 4); + s.ple_nrow = try cuda.DeviceBuffer.alloc(driver, wide * 2); + s.ple_part = try cuda.DeviceBuffer.alloc(driver, 4 * dims * 4); + s.pa = try cuda.DeviceBuffer.alloc(driver, proj_n * 2); + s.q = try cuda.DeviceBuffer.alloc(driver, q_heads * head_dim * 2); + s.iq = try cuda.DeviceBuffer.alloc(driver, index_heads * index_dim * 2); + s.po = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); + s.pm = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.pl = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * 4); + s.ao = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gated = try cuda.DeviceBuffer.alloc(driver, out_k * 2); + s.gxs_attn = try cuda.DeviceBuffer.alloc(driver, (out_k / 32) * 4); + s.dummy = try cuda.DeviceBuffer.alloc(driver, 4096); + var inv: [half]f32 = undefined; + for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); + s.inv = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(&inv)); + s.pos = try cuda.DeviceBuffer.alloc(driver, 4); + s.ids = try cuda.DeviceBuffer.alloc(driver, 4); + s.logits = try cuda.DeviceBuffer.alloc(driver, vocab * 2); + s.router_l = try cuda.DeviceBuffer.alloc(driver, (experts_n + 1) * 4); + s.picks = try cuda.DeviceBuffer.alloc(driver, slots * 4); + s.moe_members = try cuda.DeviceBuffer.alloc(driver, slots * 4); + s.moe_items = try cuda.DeviceBuffer.alloc(driver, slots * 3 * 4); + s.moe_counts = try cuda.DeviceBuffer.alloc(driver, 8); + s.moe_act = try cuda.DeviceBuffer.alloc(driver, slots * moe_w * 2); + return s; + } + + pub fn free(self: *Scratch) void { + self.h.free(); + self.y.free(); + self.wts.free(); + self.inj.free(); + self.pss.free(); + self.norm.free(); + self.dn.free(); + self.hc_part.free(); + self.act.free(); + self.xs_act.free(); + self.mixed.free(); + self.xsm.free(); + self.proj.free(); + self.part.free(); + self.gout.free(); + self.gxs.free(); + self.st.free(); + self.ple_emb.free(); + self.ple_xs.free(); + self.ple_keys.free(); + self.ple_vals.free(); + self.ple_gated.free(); + self.ple_pss.free(); + self.ple_nrow.free(); + self.ple_part.free(); + self.pa.free(); + self.q.free(); + self.iq.free(); + self.po.free(); + self.pm.free(); + self.pl.free(); + self.ao.free(); + self.gated.free(); + self.gxs_attn.free(); + self.dummy.free(); + self.inv.free(); + self.pos.free(); + self.ids.free(); + self.logits.free(); + self.router_l.free(); + self.picks.free(); + self.moe_members.free(); + self.moe_items.free(); + self.moe_counts.free(); + self.moe_act.free(); + } +}; + +pub fn writeback(tri: anytype, s: *Scratch, h: u64, mode: usize, inj: u64, y: u64, wts: u64) !void { + const dummy = s.dummy.ptr; + if (mode == 0) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 0, 1, 1, 1); + } else if (mode == 2) { + try tri.hcWriteback(h, h, s.pss.ptr, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", dims, 1, dims, streams, 2, 10, 11, 1); + } else if (mode == 4) { + try tri.hcWriteback(h, h, s.pss.ptr, s.part.ptr, "*fp32", inj, dummy, "*bf16", dummy, "*bf16", dims, 1, dims, streams, 4, 1, 1, out_sk); + } else return error.UnexpectedTensor; +} + +pub fn readout(tri: anytype, s: *Scratch, hc_w: HcW, h: u64, inj: u64) !void { + try tri.hcDown(h, s.pss.ptr, hc_w.scale.ptr, s.norm.ptr, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn.ptr, s.hc_part.ptr, eps, 1, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); + const inj_ptr = if (hc_w.has_inj == 1) inj else s.act.ptr; + try tri.hcReduceAct(s.hc_part.ptr, s.act.ptr, s.xs_act.ptr, inj_ptr, hc_sk, 1, streams, low, hc_w.n_down, hc_w.has_inj); + try tri.hcUpmix(s.act.ptr, s.xs_act.ptr, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, s.norm.ptr, s.mixed.ptr, s.xsm.ptr, 1, wide, low, dims, streams); +} + +fn shiftTail(gpa: std.mem.Allocator, stream: *cuda.Stream, window: cuda.DeviceBuffer, taps_n: usize, row_bytes: usize, fresh: []const u8) !void { + const bytes = taps_n * row_bytes; + const raw = try gpa.alloc(u8, bytes); + defer gpa.free(raw); + try stream.synchronize(); + try window.download(0, raw); + const next = try gpa.alloc(u8, bytes); + defer gpa.free(next); + const keep = taps_n - 1; + @memcpy(next[0 .. keep * row_bytes], raw[row_bytes..]); + @memcpy(next[keep * row_bytes ..], fresh); + try window.upload(0, next); +} + +pub fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, s: *Scratch, proj_q: Q4, out_q: Q4, conv: cuda.DeviceBuffer, cw: u64, a_log: u64, dt: u64, norm: u64, state: *cuda.DeviceBuffer) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.proj.ptr, 1, proj_q.n, proj_q.k); + try gdn.launch(driver, stream.*, s.proj.ptr, conv.ptr, cw, state.ptr, a_log, dt, norm, eps, s.gout.ptr, s.gxs.ptr, s.st.ptr); + try state.copyFrom(0, s.st.ptr, state_n * 4, stream.handle); + const row = try gpa.alloc(u8, gdn.conv_dim * 2); + defer gpa.free(row); + try stream.synchronize(); + try s.proj.download(0, row); + try shiftTail(gpa, stream, conv, gdn.taps - 1, gdn.conv_dim * 2, row); + try qmm.partials(driver, stream.*, s.gout.ptr, s.gxs.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, gdn.value_dim, out_sk); +} + +pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { + try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); + try s.pos.upload(0, std.mem.asBytes(&pos)); + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); + try tri.attnChunks(s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, 1); + try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads); + try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); + try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); +} + +fn launchExpert(f: cuda.Function, stream: cuda.Stream, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, items: u64, counts: u64, members: u64, out: u64, n: usize, units: usize) !void { + var a: cuda.Args = .{}; + a.add(x); + a.add(cint(x_stride)); + a.add(cint(slot_n)); + a.add(w); + a.add(cint(kg)); + a.add(cint(nb)); + a.add(items); + a.add(counts); + a.add(members); + a.add(out); + a.add(cint(n)); + a.add(@as(f32, 0)); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((units + 3) / 4) }, .block = .{ .x = 128 } }, stream, &a); +} + +pub fn moeStep(comptime Tri: type, tri: Tri, s: *Scratch, stream: *cuda.Stream, plan_fn: cuda.Function, up_fn: cuda.Function, down_fn: cuda.Function, router_w: u64, up: u64, down: u64, y: u64, wts: u64) !void { + try tri.router(s.mixed.ptr, router_w, s.router_l.ptr, dims); + try tri.topkRows(s.router_l.ptr, s.picks.ptr, wts); + var plan_a: cuda.Args = .{}; + plan_a.add(s.picks.ptr); + plan_a.add(cint(slots)); + plan_a.add(cint(experts_n + 1)); + plan_a.add(cint(tile)); + plan_a.add(s.moe_members.ptr); + plan_a.add(s.moe_items.ptr); + plan_a.add(s.moe_counts.ptr); + try cuda.launch.launch(plan_fn, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, stream.*, &plan_a); + try launchExpert(up_fn, stream.*, s.mixed.ptr, dims, slots, up, dims / 32, moe_w / 32, s.moe_items.ptr, s.moe_counts.ptr, s.moe_members.ptr, s.moe_act.ptr, moe_w, slots * (moe_w / 32)); + try launchExpert(down_fn, stream.*, s.moe_act.ptr, moe_w, 0, down, moe_w / 32, dims / 32, s.moe_items.ptr, s.moe_counts.ptr, s.moe_members.ptr, y, dims, slots * (dims / 32)); +} + +fn reduceBf16(gpa: std.mem.Allocator, stream: *cuda.Stream, part: cuda.DeviceBuffer, sk: usize, n: usize, dst: cuda.DeviceBuffer) !void { + const raw = try gpa.alloc(u8, sk * n * 4); + defer gpa.free(raw); + try stream.synchronize(); + try part.download(0, raw); + const row = try gpa.alloc(u16, n); + defer gpa.free(row); + for (0..n) |i| { + var acc: f32 = @bitCast(std.mem.readInt(u32, raw[4 * i ..][0..4], .little)); + for (1..sk) |s| acc += @as(f32, @bitCast(std.mem.readInt(u32, raw[4 * (s * n + i) ..][0..4], .little))); + row[i] = toBf16(acc); + } + try dst.upload(0, std.mem.sliceAsBytes(row)); +} + +pub fn pleStep(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, model_dir: []const u8, history: *[2]i64, token: u32, h: u64, key_q: Q4, val_q: Q4, nk: u64, nq: u64, nc: u64, cw: u64, tail: *cuda.DeviceBuffer) !void { + const emb = try gpa.alloc(u16, dims); + defer gpa.free(emb); + const xs = try gpa.alloc(f32, dims / 32); + defer gpa.free(xs); + try ple.embedding(gpa, io, model_dir, history, token, emb, xs); + try s.ple_emb.upload(0, std.mem.sliceAsBytes(emb)); + try s.ple_xs.upload(0, std.mem.sliceAsBytes(xs)); + try qmm.matmul(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, key_q.w.ptr, key_q.s.ptr, key_q.b.ptr, s.ple_keys.ptr, 1, wide, dims); + try qmm.partials(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, val_q.w.ptr, val_q.s.ptr, val_q.b.ptr, s.ple_vals.ptr, s.ple_part.ptr, 1, dims, dims, 4); + try reduceBf16(gpa, stream, s.ple_part, 4, dims, s.ple_vals); + try tri.pleGate(s.ple_keys.ptr, s.ple_vals.ptr, h, nk, nq, s.ple_gated.ptr, s.ple_pss.ptr, eps, 1, dims, streams); + try tri.pleConv(s.ple_gated.ptr, s.ple_pss.ptr, nc, tail.ptr, cw, h, h, s.ple_nrow.ptr, eps, 1, dims, streams, 4, 3); + const nrow = try gpa.alloc(u8, wide * 2); + defer gpa.free(nrow); + try stream.synchronize(); + try s.ple_nrow.download(0, nrow); + try shiftTail(gpa, stream, tail.*, ple_tail_n, wide * 2, nrow); + history[0] = history[1]; + history[1] = token; +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig new file mode 100644 index 0000000000..a2598efd7d --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -0,0 +1,398 @@ +//! Checkpoint bytes for one drafts-off prompt: the safetensors reader and the quantized linears. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const hc = @import("cuda_hc.zig"); +const qmm = @import("cuda_qmm.zig"); +const gdn = @import("cuda_gdn.zig"); + +pub const layers_n: usize = 48; +pub const linear_n: usize = 36; +pub const attn_n: usize = 12; +pub const dims: usize = 2560; +pub const streams: usize = 4; +pub const wide: usize = streams * dims; +pub const low: usize = 320; +pub const eps: f32 = 1e-6; +pub const vocab: usize = 248320; +pub const experts_n: usize = 512; +pub const slots: usize = 11; +pub const moe_w: usize = 640; +pub const block_bytes: usize = 160 * 4; +pub const tile: usize = 16; +pub const q_heads: usize = 24; +pub const kv_heads: usize = 2; +pub const head_dim: usize = 256; +pub const index_heads: usize = 4; +pub const index_dim: usize = 128; +pub const half: usize = 32; +pub const proj_n: usize = 13952; +pub const out_k: usize = 6144; +pub const nch: usize = 3; +pub const ple_layer: usize = 1; +pub const inj_stride: usize = 16; +pub const wts_stride: usize = 48; +pub const ple_tail_n: usize = 9; +pub const state_n: usize = gdn.nv * gdn.dv * gdn.dk; +pub const hc_sk: usize = 32; +pub const out_sk: usize = 8; + +pub const pack_symbol: [:0]const u8 = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +pub const plan_symbol: [:0]const u8 = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; +pub const up_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; +pub const down_symbol: [:0]const u8 = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + +const Tensor = core.safetensors.Tensor; + +const Held = struct { rel: []u8, file: Shard }; + +const Shard = struct { + io: std.Io, + file: std.Io.File, + map: std.Io.File.MemoryMap, + at: usize, + arena: std.heap.ArenaAllocator, + names: std.StringHashMapUnmanaged(core.safetensors.Entry), + + fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Shard { + var file = try std.Io.Dir.cwd().openFile(io, path, .{}); + errdefer file.close(io); + const len: usize = @intCast(try file.length(io)); + if (len < 8) return error.BadSafetensors; + var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); + errdefer map.destroy(io); + const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); + if (header_len > len - 8) return error.BadSafetensors; + var arena = std.heap.ArenaAllocator.init(gpa); + errdefer arena.deinit(); + var names: std.StringHashMapUnmanaged(core.safetensors.Entry) = .empty; + const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), map.memory[8..][0..header_len], .{}); + defer parsed.deinit(); + var it = parsed.value.object.iterator(); + while (it.next()) |kv| { + if (std.mem.eql(u8, kv.key_ptr.*, "__metadata__")) continue; + const o = kv.value_ptr.object; + const dtype = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + const shape = o.get("shape").?.array.items; + if (shape.len > core.safetensors.max_rank) continue; + var e: core.safetensors.Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = 0 }; + var n: usize = dtype.size(); + for (shape, 0..) |d, i| { + e.shape[i] = @intCast(d.integer); + n *= e.shape[i]; + } + const offs = o.get("data_offsets").?.array.items; + e.begin = @intCast(offs[0].integer); + e.end = @intCast(offs[1].integer); + if (e.end < e.begin or e.end - e.begin != n or e.end > len - 8 - header_len) return error.BadSafetensors; + try names.put(arena.allocator(), try arena.allocator().dupe(u8, kv.key_ptr.*), e); + } + return .{ .io = io, .file = file, .map = map, .at = 8 + header_len, .arena = arena, .names = names }; + } + + fn close(self: *Shard) void { + self.map.destroy(self.io); + self.file.close(self.io); + self.arena.deinit(); + self.* = undefined; + } + + fn get(self: *const Shard, name: []const u8) ?Tensor { + const e = self.names.get(name) orelse return null; + return .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .bytes = self.map.memory[self.at + e.begin .. self.at + e.end] }; + } +}; + +pub const Store = struct { + gpa: std.mem.Allocator, + io: std.Io, + dir: []const u8, + parsed: std.json.Parsed(std.json.Value), + path: []u8 = &.{}, + rel: []u8 = &.{}, + file: ?Shard = null, + held: [12]Held = undefined, + nheld: usize = 0, + + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { + const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); + defer gpa.free(text); + return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}) }; + } + + pub fn deinit(self: *Store) void { + self.release(); + self.closeFile(); + self.parsed.deinit(); + } + + pub fn release(self: *Store) void { + for (self.held[0..self.nheld]) |*h| { + h.file.close(); + self.gpa.free(h.rel); + } + self.nheld = 0; + } + + pub fn hold(self: *Store, key: []const u8) !Tensor { + const file = try self.shard(key); + for (self.held[0..self.nheld]) |*h| { + if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; + } + if (self.nheld == self.held.len) return error.UnexpectedTensor; + const rel = try self.gpa.dupe(u8, file); + errdefer self.gpa.free(rel); + const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + defer self.gpa.free(path); + self.held[self.nheld] = .{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }; + self.nheld += 1; + return self.held[self.nheld - 1].file.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + pub fn closeFile(self: *Store) void { + if (self.file) |*f| f.close(); + self.file = null; + if (self.path.len != 0) self.gpa.free(self.path); + self.path = &.{}; + if (self.rel.len != 0) self.gpa.free(self.rel); + self.rel = &.{}; + } + + fn shard(self: *Store, key: []const u8) ![]const u8 { + const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + const v = wm.object.get(key) orelse { + std.debug.print("missing {s}\n", .{key}); + return error.MissingTensor; + }; + return v.string; + } + + fn use(self: *Store, key: []const u8) !void { + const file = try self.shard(key); + if (self.file != null and std.mem.eql(u8, file, self.rel)) return; + self.closeFile(); + self.rel = try self.gpa.dupe(u8, file); + self.path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); + self.file = try Shard.load(self.gpa, self.io, self.path); + } + + pub fn tensor(self: *Store, key: []const u8) !Tensor { + try self.use(key); + return self.file.?.get(key) orelse { + std.debug.print("absent {s}\n", .{key}); + return error.MissingTensor; + }; + } + + pub fn copyOf(self: *Store, key: []const u8) ![]u8 { + const t = try self.tensor(key); + const out = try self.gpa.alloc(u8, t.bytes.len); + @memcpy(out, t.bytes); + return out; + } +}; + +pub const Face = struct { + w: []const u8, + s: []const u8, + b: []const u8, + n: usize, + k: usize, + routed: bool, +}; + +pub const Q4 = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + n: usize, + k: usize, + + pub fn free(self: *Q4) void { + self.w.free(); + self.s.free(); + self.b.free(); + } + + pub fn empty(d: *cuda.Driver) Q4 { + const z: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + return .{ .w = z, .s = z, .b = z, .n = 0, .k = 0 }; + } +}; + +pub const HcW = struct { + scale: cuda.DeviceBuffer, + down: Q4, + up: Q4, + n_down: usize, + has_inj: usize, + + pub fn free(self: *HcW) void { + self.scale.free(); + self.down.free(); + self.up.free(); + } +}; + +pub fn linearLayer(layer: usize) bool { + return layer % 4 != 3; +} + +pub fn linearIndex(layer: usize) usize { + return layer - layer / 4; +} + +pub fn attnIndex(layer: usize) usize { + return layer / 4; +} + +pub fn toBf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +pub fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +pub fn cint(v: usize) c_int { + return @intCast(v); +} + +pub fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { + var n: usize = 0; + for (parts) |p| n += p.len; + const out = try gpa.alloc(u8, n); + var at: usize = 0; + for (parts) |p| { + @memcpy(out[at..][0..p.len], p); + at += p.len; + } + return out; +} + +pub fn asF32(gpa: std.mem.Allocator, t: Tensor) ![]f32 { + if (t.rank != 1) return error.UnexpectedTensor; + const out = try gpa.alloc(f32, t.dim(0)); + if (t.dtype == .f32 and t.bytes.len == out.len * 4) { + for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); + } else if (t.dtype == .bf16 and t.bytes.len == out.len * 2) { + for (out, 0..) |*o, i| o.* = promote(std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)); + } else return error.UnexpectedTensor; + return out; +} + +pub fn uploadQ4(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + var lane = try qmm.pack(gpa, words, scales, biases, n, k8); + defer lane.deinit(gpa); + var w = try cuda.DeviceBuffer.fromHost(driver, lane.weight); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, lane.scales); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, lane.biases); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +pub fn uploadTiled(gpa: std.mem.Allocator, driver: *cuda.Driver, words: []const u8, scales: []const u8, biases: []const u8, n: usize, k8: usize) !Q4 { + const kg = k8 / 4; + const tiled = try hc.tileWords(gpa, words, n, k8); + defer gpa.free(tiled); + const ts = try hc.transposeBf16(gpa, scales, n, kg); + defer gpa.free(ts); + const tb = try hc.transposeBf16(gpa, biases, n, kg); + defer gpa.free(tb); + var w = try cuda.DeviceBuffer.fromHost(driver, tiled); + errdefer w.free(); + var s = try cuda.DeviceBuffer.fromHost(driver, ts); + errdefer s.free(); + const b = try cuda.DeviceBuffer.fromHost(driver, tb); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k8 * 8 }; +} + +pub fn join3(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u8, name: []const u8) !struct { w: []u8, s: []u8, b: []u8 } { + const w = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.weight", .{ prefix, name })); + errdefer gpa.free(w); + const s = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.scales", .{ prefix, name })); + errdefer gpa.free(s); + const b = try store.copyOf(try std.fmt.bufPrint(buf, "{s}{s}.biases", .{ prefix, name })); + return .{ .w = w, .s = s, .b = b }; +} + +pub fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { + var buf: [180]u8 = undefined; + const down = try join3(gpa, store, &buf, prefix, ".input_mix_weight_down"); + defer gpa.free(down.w); + defer gpa.free(down.s); + defer gpa.free(down.b); + const up = try join3(gpa, store, &buf, prefix, ".input_mix_weight_up"); + defer gpa.free(up.w); + defer gpa.free(up.s); + defer gpa.free(up.b); + const gamma = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix})); + const scale = try asF32(gpa, gamma); + defer gpa.free(scale); + if (scale.len != wide) return error.UnexpectedTensor; + var n: usize = low; + var words: []u8 = undefined; + var scales: []u8 = undefined; + var biases: []u8 = undefined; + var inj_w: []u8 = &.{}; + var inj_s: []u8 = &.{}; + var inj_b: []u8 = &.{}; + if (inject) { + const inj = try join3(gpa, store, &buf, prefix, ".block_inject_weight"); + inj_w = inj.w; + inj_s = inj.s; + inj_b = inj.b; + n = low + 4; + words = try cat(gpa, &.{ down.w, inj.w }); + scales = try cat(gpa, &.{ down.s, inj.s }); + biases = try cat(gpa, &.{ down.b, inj.b }); + } else { + words = try gpa.dupe(u8, down.w); + scales = try gpa.dupe(u8, down.s); + biases = try gpa.dupe(u8, down.b); + } + defer gpa.free(words); + defer gpa.free(scales); + defer gpa.free(biases); + defer if (inject) gpa.free(inj_w); + defer if (inject) gpa.free(inj_s); + defer if (inject) gpa.free(inj_b); + const k8 = wide / 8; + if (words.len != n * k8 * 4) return error.UnexpectedTensor; + var down_q = try uploadTiled(gpa, driver, words, scales, biases, n, k8); + errdefer down_q.free(); + var up_q = try uploadTiled(gpa, driver, up.w, up.s, up.b, wide, low / 8); + errdefer up_q.free(); + const scale_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(scale)); + return .{ .scale = scale_b, .down = down_q, .up = up_q, .n_down = n, .has_inj = @intFromBool(inject) }; +} + +pub fn stitch(gpa: std.mem.Allocator, store: *Store, prefix: []const u8, names: []const []const u8) !struct { w: []u8, s: []u8, b: []u8, n: usize } { + var buf: [180]u8 = undefined; + var ws: [6][]u8 = undefined; + var ss: [6][]u8 = undefined; + var bs: [6][]u8 = undefined; + var n: usize = 0; + for (names, 0..) |part, i| { + const got = try join3(gpa, store, &buf, prefix, part); + ws[i] = got.w; + ss[i] = got.s; + bs[i] = got.b; + n += 1; + } + defer for (0..n) |i| { + gpa.free(ws[i]); + gpa.free(ss[i]); + gpa.free(bs[i]); + }; + return .{ .w = try cat(gpa, ws[0..n]), .s = try cat(gpa, ss[0..n]), .b = try cat(gpa, bs[0..n]), .n = n }; +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..c599ba1582 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -4,13 +4,14 @@ const cuda = @import("cuda"); const api = @import("engine_api"); const lanes = @import("lanes"); const nemotron = @import("nemotron"); +const flashnext = @import("flashnext"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{nemotron.native}; +const registry = .{ nemotron.native, flashnext.native }; /// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ @@ -417,6 +418,8 @@ test "chip classes name the compute capability as gate entries do" { test "every registered family is listed for capabilities" { try std.testing.expectEqual(@as(usize, registry.len), families.len); try std.testing.expectEqualStrings("nemotron_h", families[0].model_type); + try std.testing.expectEqualStrings("qwen4_exp", families[1].model_type); + try std.testing.expectEqualStrings("mlx-q4g32", families[1].formats[0]); } test "variables: unset or empty means the default, anything but a whole number is refused" { From 785896e1981ecc17dac520a369fe72d2d54edc08 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:08:15 -0500 Subject: [PATCH 029/115] Flash Next CUDA: prompt ids and positions upload on the compute stream; prompts past 512 keys; LR7 capture. The ids and position uploads were synchronous copies, which do not wait for the non-blocking compute stream, so a later row overwrote a queued row's token id or position: wrong tokens on most prompts. Both now go on the stream. attnStep launches one 512-key chunk per started chunk of the row's keys (as attention.attention does) instead of always one, so prompts run to the captured NCH * 512 keys; the lanes engine defaults to that window and refuses a larger one. The prompt check prints each run's top five scores and its timing. zig/tests/cuda/flashnext: capture.py and capture.sh (Python engine 0.6.1, drafts off, greedy, 1536 cache slots) and the Unbound-LR7 prompts with their references; flashnext-prompt matches the teacher paths on all three. --- zig/src/families/flashnext/cuda_lanes.zig | 4 +- zig/src/families/flashnext/cuda_prompt.zig | 39 +- .../families/flashnext/cuda_prompt_step.zig | 8 +- zig/tests/cuda/flashnext/capture.py | 178 ++++ zig/tests/cuda/flashnext/capture.sh | 25 + zig/tests/cuda/flashnext/lr7/long.ids | 1 + zig/tests/cuda/flashnext/lr7/mid.ids | 1 + zig/tests/cuda/flashnext/lr7/prompts.json | 1 + zig/tests/cuda/flashnext/lr7/results.json | 920 ++++++++++++++++++ zig/tests/cuda/flashnext/lr7/short.ids | 1 + 10 files changed, 1173 insertions(+), 5 deletions(-) create mode 100644 zig/tests/cuda/flashnext/capture.py create mode 100755 zig/tests/cuda/flashnext/capture.sh create mode 100644 zig/tests/cuda/flashnext/lr7/long.ids create mode 100644 zig/tests/cuda/flashnext/lr7/mid.ids create mode 100644 zig/tests/cuda/flashnext/lr7/prompts.json create mode 100644 zig/tests/cuda/flashnext/lr7/results.json create mode 100644 zig/tests/cuda/flashnext/lr7/short.ids diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 7502108299..cfacce15b3 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -13,7 +13,8 @@ pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; pub const formats: []const []const u8 = &.{"mlx-q4g32"}; - pub const default_context: i64 = 262144; + // dense attention over the captured scratch (NCH chunks of 512 keys, QSA off): no sparse indexer yet + pub const default_context: i64 = @intCast(prompt.max_tokens); pub const prefill_step: u32 = 0; pub const Options = struct { context: usize, drafts: bool }; @@ -98,6 +99,7 @@ pub fn Namespace(comptime Tri: type) type { pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: []const u8, o: Options) !Loaded { _ = o.drafts; // this model is served one token at a time; there is no draft head + if (o.context > prompt.max_tokens) return error.ContextOverDenseAttention; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 5dad8c369b..e955b7d5d1 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -67,6 +67,26 @@ pub fn argmax(raw: []const u8) u32 { return @intCast(at); } +/// The five best scores of one row, best first, as "id:score" pairs on stderr (the capture's top5 field). +fn printTop(raw: []const u8) void { + var ids: [5]usize = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (0..vocab) |i| { + const v = promote(std.mem.readInt(u16, raw[2 * i ..][0..2], .little)); + if (v <= vals[4]) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = i; + } + std.debug.print("top5", .{}); + for (ids, vals) |id, v| std.debug.print(" {d}:{d:.4}", .{ id, v }); + std.debug.print("\n", .{}); +} + const Pass = struct { gpa: std.mem.Allocator, io: std.Io, @@ -95,7 +115,9 @@ const Pass = struct { fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { const id: i32 = @intCast(token); - try p.s.ids.upload(0, std.mem.asBytes(&id)); + // stream-ordered: a synchronous copy does not wait for the non-blocking stream, so a later row's id + // would overwrite this one before its queued _embed reads it + try p.s.ids.uploadAsync(0, std.mem.asBytes(&id), p.stream.handle); try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); } @@ -165,6 +187,8 @@ fn forward(comptime Tri: type, p: *Pass, tri: Tri, layers: []Layer, ids: []const } pub const vocabulary: usize = vocab; +/// Prompt plus reply tokens a sequence can hold: the attention kernels' NCH chunks of 512 keys. +pub const max_tokens: usize = weights.nch * 512; /// One loaded checkpoint: the embedding table, the mixer, the head, and a scratch row the forwards share. pub const Engine = struct { @@ -373,22 +397,33 @@ fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { /// `out` receives `new_tokens` greedy ids. The first is the prompt's argmax. The rest are fed back one at a time. pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { - if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > 256) return error.UnexpectedTensor; + if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > max_tokens) return error.UnexpectedTensor; var e = try Engine.init(gpa, io, driver, stream, model_dir); defer e.deinit(); var seq = try e.newSeq(tokens.len + new_tokens); defer seq.deinit(gpa); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); try e.predict(Tri, tri, &seq, tokens); + const t1 = std.Io.Clock.awake.now(io).toNanoseconds(); out[0] = argmax(seq.logits); + printTop(seq.logits); std.debug.print("token {d}\n", .{out[0]}); var tok = out[0]; for (1..new_tokens) |i| { var one = [_]u32{tok}; try e.predict(Tri, tri, &seq, &one); tok = argmax(seq.logits); + if (i < 4) printTop(seq.logits); out[i] = tok; std.debug.print("token {d}\n", .{tok}); } + const t2 = std.Io.Clock.awake.now(io).toNanoseconds(); + const ms = struct { + fn of(ns: i96) f64 { + return @as(f64, @floatFromInt(ns)) / 1e6; + } + }.of; + std.debug.print("timing prompt {d} rows {d:.1} ms, decode {d} tokens {d:.2} ms/token\n", .{ tokens.len, ms(t1 - t0), new_tokens - 1, if (new_tokens > 1) ms(t2 - t1) / @as(f64, @floatFromInt(new_tokens - 1)) else 0 }); } test "linear layers skip every fourth index and attention starts at 3" { diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig index de02eb2eef..44838b6077 100644 --- a/zig/src/families/flashnext/cuda_prompt_step.zig +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -237,9 +237,13 @@ pub fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Strea pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); - try s.pos.upload(0, std.mem.asBytes(&pos)); + // stream-ordered, like the ids: the previous row's queued attention kernels still read s.pos + try s.pos.uploadAsync(0, std.mem.asBytes(&pos), stream.handle); try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); - try tri.attnChunks(s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, 1); + // attention.attention: chunks of 512 keys up to the row's position, at most the scratch's nch + const keys: usize = @as(usize, @intCast(pos)) + 1; + if (keys > nch * 512) return error.UnexpectedTensor; + try tri.attnChunks(s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512); try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads); try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); diff --git a/zig/tests/cuda/flashnext/capture.py b/zig/tests/cuda/flashnext/capture.py new file mode 100644 index 0000000000..a92c58adc7 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +"""Flash Next's Python CUDA engine (qwen4_exp, drafts off, greedy) as the Zig port's oracle: every Triton launch, +prompt ids, and two greedy references per prompt (the served path: chunked prefill then serial decode; the teacher +path: every prompt token as its own decode row, the way zig's flashnext-prompt runs it).""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +PROMPTS = { + "short": "What is the capital of France?", + "mid": ("Explain in plain words why the sky looks blue during the day but turns red and orange at sunset. " + "Mention the role of the atmosphere and the wavelength of light, and keep it under five sentences."), +} +TOP_STEPS = 4 # sampled steps whose five best logits the teacher paths keep +LONG_TOKENS = 600 # the long prompt: module source cut to make the whole prompt about this many tokens + + +def chat(user: str) -> str: + return f"<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" + + +def prompts(model: Path) -> dict[str, list[int]]: + from tokenizers import Tokenizer + import textwrap + + tok = Tokenizer.from_file(str(model / "tokenizer.json")) + enc = lambda s: tok.encode(s, add_special_tokens=False).ids # noqa: E731 + out = {name: enc(chat(text)) for name, text in PROMPTS.items()} + head = enc("<|im_start|>user\nSummarize what this Python module does, then list its public functions:\n\n") + tail = enc("<|im_end|>\n<|im_start|>assistant\n\n\n\n\n") + body = enc(Path(textwrap.__file__).read_text())[:LONG_TOKENS - len(head) - len(tail)] + out["long"] = head + body + tail + return out + + +def specialization() -> dict: + """Each JIT function's parameter names and which ones Triton never specializes (aot_pack.py's --jit).""" + + from triton.runtime.jit import JITFunction + import gc + + out = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + out[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + return out + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", required=True, help="the checkpoint's snapshot dir") + ap.add_argument("--out", required=True) + ap.add_argument("--tools", default="", help="folder of triton_aot_manifest.py: record every Triton launch") + ap.add_argument("--tokens", type=int, default=32, help="greedy tokens per prompt") + ap.add_argument("--context", type=int, default=1536, + help="the cache window; 1025..1536 gives the port's NCH=3, QSA=0 attention specialization") + ap.add_argument("--only", default="", help="comma list of prompt names (default all)") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + rec = None + if a.tools: + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + + import torch + from contextlib import nullcontext + + def phase(name: str, detail: bool = False): + return rec.scope(name, detail) if rec is not None else nullcontext() + + model = Path(a.model) + ids = prompts(model) + if a.only: + ids = {k: v for k, v in ids.items() if k in a.only.split(",")} + (out / "prompts.json").write_text(json.dumps(ids) + "\n") + for name, p in ids.items(): + (out / f"{name}.ids").write_text(" ".join(map(str, p)) + "\n") + + from tensorfold.families.qwen4_exp.cuda.engine import FlashNextEngine + from tensorfold.families.qwen4_exp.cuda.forward import commit + + t0 = time.perf_counter() + with phase("startup"): + eng = FlashNextEngine(model, depth=0, max_len=a.context, context_explicit=True, streams=1, + ple_on_ssd=True, kv_dtype="bf16", graphs=False) + e = eng.e + info = {"startup_s": round(time.perf_counter() - t0, 2), "max_len": eng.max_len, "prefill_rows": eng.prefill_rows, + "nch": e.buf.attn.nch, "qsa": bool(e.buf.attn.qsa), "idw": e.buf.attn.idw, "eos": list(eng.eos), + "torch": torch.__version__, "gpu": torch.cuda.get_device_name(0), + "prompt_lengths": {k: len(v) for k, v in ids.items()}} + print(json.dumps(info), flush=True) + + results: dict[str, dict] = {} + from functools import partial + from tensorfold.families.qwen4_exp.cuda.state import Buffers + + teach = partial(_teach, phase, results, a, commit, nullcontext) + fp32 = None + for name, prompt in ids.items(): + # the served path: chunked prefill, then one token a round (draft=False is FlashNextEngine._serial) + toks: list[int] = [] + with phase(f"serial-{name}"): + torch.cuda.synchronize() + start = time.perf_counter() + stats = eng.generate(prompt, a.tokens, None, lambda new: toks.extend(map(int, new)) and False, + draft=False, stop_eos=False) + torch.cuda.synchronize() + wall = time.perf_counter() - start + results[f"serial/{name}"] = {"tokens": toks, "wall_s": round(wall, 4), **stats} + print(f"serial {name} ({len(prompt)} prompt): {toks} {stats}", flush=True) + if fp32 is None: + # teacher32: the decode-form experts (fp32 slot outputs, Buffers' default and the MTP head's), the + # arithmetic zig's flashnext-prompt runs; the serial engine's buffers take the prompt-form experts + fp32 = eng.serial.twin() + fp32.buf = Buffers(eng.e.w, eng.e.rows, eng.e.capacity) + + for label, t in (("teacher", eng.serial), ("teacher32", fp32)): + teach(label, name, prompt, t, toks) + + (out / "results.json").write_text(json.dumps({"info": info, "results": results}, indent=1) + "\n") + if rec is not None: + rec.dump(out / "launches.json") + (out / "jit.json").write_text(json.dumps(specialization(), indent=1, sort_keys=True) + "\n") + eng.close() + print("done", flush=True) + return 0 + + +def _teach(phase, results, a, commit, nullcontext, label, name, prompt, t, served): + """Every prompt token one decode row, as zig's flashnext-prompt feeds them, then greedy rows.""" + + import torch + + t.reset() + got: list[int] = [] + top: list = [] + steps = 0 + with phase(f"{label}-{name}"): + torch.cuda.synchronize() + start = time.perf_counter() + feed = list(prompt) + dec_start = None + while len(got) < a.tokens: + tok = feed.pop(0) if feed else got[-1] + with (phase(f"{label}-{name}-detail", detail=True) if name == "short" and steps < 2 + else nullcontext()): + logits = t.forward([tok]) + if not feed: + got.append(t.sample(logits[:1], [t.st.pos + 1], None)[0]) + if len(got) <= TOP_STEPS: + v, i = logits[0].float().topk(5) + top.append([[int(x), round(float(y), 4)] for x, y in zip(i.tolist(), v.tolist())]) + if dec_start is None: + torch.cuda.synchronize() + dec_start = time.perf_counter() + commit(t.w, t.st, t.buf, 1, 1) + steps += 1 + torch.cuda.synchronize() + end = time.perf_counter() + ms = (end - dec_start) / max(1, a.tokens - 1) * 1e3 + results[f"{label}/{name}"] = {"tokens": got, "wall_s": round(end - start, 4), + "decode_ms_per_token": round(ms, 3), "top5": top} + print(f"{label} {name}: {got} ({ms:.2f} ms/token) same as serial: {got == served}", flush=True) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture.sh b/zig/tests/cuda/flashnext/capture.sh new file mode 100755 index 0000000000..5bb769d5b8 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Flash Next (qwen4_exp) capture -> manifest -> kernel set for zig's flashnext-prompt and the native lanes engine. +# Usage: capture.sh WORK SNAPSHOT_DIR [IMAGE] +# WORK a host dir this script owns: WORK/cache (Triton + torch extension caches), WORK/out/cap, WORK/kernels +# SNAPSHOT_DIR the checkpoint's HF snapshot dir (its blobs must sit at ../../blobs, as in an HF cache) +# IMAGE TensorFold 0.6.1 image, default tensorfold-qwen38:v0.6.1 +# The Python engine runs drafts off, greedy, --context 1535 (1536 cache slots: _chunks/_merge NCH=3, QSA=0, the dense +# attention the port launches), --ple-on-ssd, bf16 KV. Needs the GPU free (~78 GiB). +set -euo pipefail +WORK=$(readlink -f "$1"); SNAP=$(readlink -f "$2"); IMAGE=${3:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +HF=$(dirname "$(dirname "$(dirname "$SNAP")")") # .../hub/models--X/snapshots/REV -> .../hub (holds the blobs) +REL=${SNAP#"$HF"/} +mkdir -p "$WORK/cache" "$WORK/out/cap" +docker run --rm --name fn-capture --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \ + --user "$(id -u):$(id -g)" -e HOME=/tmp -e HF_HUB_OFFLINE=1 -e TENSORFOLD_NO_UPDATE_CHECK=1 \ + -e TENSORFOLD_MEMORY_RESERVE_GIB=2 \ + -v "$HF":/hf:ro -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture.py --model "/hf/$REL" --out /out/cap \ + --tools /src/tools/zig --tokens 32 --context 1535 +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/cap/launches.json" --mount /cache/triton --out "$WORK/out/cap/manifest.json" +python3 -B "$REPO/zig/tests/cuda/nemotron/aot_pack.py" --manifest "$WORK/out/cap/manifest.json" \ + --cache "$WORK/cache/triton" --jit "$WORK/out/cap/jit.json" --out "$WORK/kernels" +echo "kernel set: $WORK/kernels; prompt ids and references: $WORK/out/cap/{short,mid,long}.ids, results.json" diff --git a/zig/tests/cuda/flashnext/lr7/long.ids b/zig/tests/cuda/flashnext/lr7/long.ids new file mode 100644 index 0000000000..81ee38d549 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/long.ids @@ -0,0 +1 @@ +248045 846 198 8917 5437 537 1092 411 12654 4536 1503 11 1179 1103 1141 569 5568 25 271 12237 1138 39834 321 20621 13 198 12237 271 2 2861 318 34 8 220 16 24 24 24 12 17 15 15 16 41550 387 13 25788 13 198 2 2861 318 34 8 220 17 15 15 17 11 220 17 15 15 18 12654 4236 4846 13 198 2 32770 539 15573 25788 361 70 1550 31 12305 4909 29 271 464 311 271 548 528 548 283 2429 1138 11210 504 359 9797 504 359 7320 504 359 9497 305 504 359 31752 504 359 8412 268 642 271 2 10906 1787 279 14541 35169 5587 310 279 2201 12 54556 198 2 35169 5587 13 220 561 1822 2781 364 3604 411 369 421 198 2 1010 34707 12258 318 4650 1088 84 15 15 64 15 8 513 2397 52483 21201 26453 13 198 83201 283 5031 83 1639 3351 15 65 3351 15 66 11674 359 271 1005 2824 11210 25 198 262 4071 198 262 2907 364 39834 6454 9121 1414 13 220 561 569 3622 16661 314 198 262 279 14623 363 321 4990 363 5245 26 279 975 5245 513 1066 1017 364 198 262 65878 310 2723 303 1906 310 49324 279 1581 16501 13 198 262 1368 488 1318 310 6379 8032 279 1822 39834 11767 11 198 262 488 3172 4507 599 310 2723 696 9797 63261 973 271 262 24876 2774 7944 2444 5090 13173 314 39834 25 198 406 2297 318 2186 25 220 22 15 8 198 285 279 6966 2297 314 18880 4965 318 37488 1387 17268 18408 198 285 369 867 8 198 406 2763 66708 318 2186 25 11363 198 285 886 421 668 381 20671 2789 310 279 1118 1500 314 18880 198 285 2468 13 220 82699 6757 279 1500 579 2297 13 198 406 16905 66708 318 2186 25 11363 198 285 886 421 668 381 20671 2789 310 660 4965 3462 279 1118 198 285 314 18880 2468 26 1048 14159 6757 1754 1500 579 2297 13 198 406 8952 56008 318 2186 25 804 8 198 285 48455 21700 303 1879 1414 310 12258 1518 4473 8427 13 198 285 8618 5474 668 3512 220 15 5074 359 5999 2073 6 12258 11 11313 383 1141 2234 198 285 303 1141 1500 13 220 1368 867 11 1754 5474 369 11421 430 264 3074 3542 13 198 406 5474 2073 318 2186 25 220 23 8 198 285 48455 21700 303 1879 1414 310 220 15 5074 359 5999 2073 6 12258 11 7014 198 285 359 31251 56008 6 369 867 13 198 406 8032 83201 318 2186 25 804 8 198 285 28592 660 35169 5587 303 279 1879 1414 539 12258 198 285 1238 5474 14046 13 220 6817 421 413 8952 56008 369 867 321 198 285 8032 83201 369 804 11 1396 5474 668 381 15625 310 264 198 285 3074 3433 0 198 406 4884 46720 6018 793 318 2186 25 867 8 198 285 28320 421 11316 64541 59429 369 2590 7854 198 285 539 1330 12258 13 220 3978 539 1581 1521 279 11767 369 198 285 318 354 46261 2752 8 65669 13 198 406 1387 17268 18408 318 2186 25 804 8 198 285 15164 4105 4860 1056 359 2998 4282 220 1368 867 11 1782 4105 668 524 198 285 381 10548 11 321 1010 4965 2493 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/mid.ids b/zig/tests/cuda/flashnext/lr7/mid.ids new file mode 100644 index 0000000000..e2a6571190 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/mid.ids @@ -0,0 +1 @@ +248045 846 198 814 20139 303 13983 4105 3069 279 12515 5686 6105 2261 279 1834 694 10263 2438 321 18016 506 41564 13 82210 279 3364 314 279 16078 321 279 43796 314 3000 11 321 2426 424 1172 4097 22157 13 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/prompts.json b/zig/tests/cuda/flashnext/lr7/prompts.json new file mode 100644 index 0000000000..c067510088 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/prompts.json @@ -0,0 +1 @@ +{"short": [248045, 846, 198, 3710, 369, 279, 6511, 314, 9338, 30, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], "mid": [248045, 846, 198, 814, 20139, 303, 13983, 4105, 3069, 279, 12515, 5686, 6105, 2261, 279, 1834, 694, 10263, 2438, 321, 18016, 506, 41564, 13, 82210, 279, 3364, 314, 279, 16078, 321, 279, 43796, 314, 3000, 11, 321, 2426, 424, 1172, 4097, 22157, 13, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271], "long": [248045, 846, 198, 8917, 5437, 537, 1092, 411, 12654, 4536, 1503, 11, 1179, 1103, 1141, 569, 5568, 25, 271, 12237, 1138, 39834, 321, 20621, 13, 198, 12237, 271, 2, 2861, 318, 34, 8, 220, 16, 24, 24, 24, 12, 17, 15, 15, 16, 41550, 387, 13, 25788, 13, 198, 2, 2861, 318, 34, 8, 220, 17, 15, 15, 17, 11, 220, 17, 15, 15, 18, 12654, 4236, 4846, 13, 198, 2, 32770, 539, 15573, 25788, 361, 70, 1550, 31, 12305, 4909, 29, 271, 464, 311, 271, 548, 528, 548, 283, 2429, 1138, 11210, 504, 359, 9797, 504, 359, 7320, 504, 359, 9497, 305, 504, 359, 31752, 504, 359, 8412, 268, 642, 271, 2, 10906, 1787, 279, 14541, 35169, 5587, 310, 279, 2201, 12, 54556, 198, 2, 35169, 5587, 13, 220, 561, 1822, 2781, 364, 3604, 411, 369, 421, 198, 2, 1010, 34707, 12258, 318, 4650, 1088, 84, 15, 15, 64, 15, 8, 513, 2397, 52483, 21201, 26453, 13, 198, 83201, 283, 5031, 83, 1639, 3351, 15, 65, 3351, 15, 66, 11674, 359, 271, 1005, 2824, 11210, 25, 198, 262, 4071, 198, 262, 2907, 364, 39834, 6454, 9121, 1414, 13, 220, 561, 569, 3622, 16661, 314, 198, 262, 279, 14623, 363, 321, 4990, 363, 5245, 26, 279, 975, 5245, 513, 1066, 1017, 364, 198, 262, 65878, 310, 2723, 303, 1906, 310, 49324, 279, 1581, 16501, 13, 198, 262, 1368, 488, 1318, 310, 6379, 8032, 279, 1822, 39834, 11767, 11, 198, 262, 488, 3172, 4507, 599, 310, 2723, 696, 9797, 63261, 973, 271, 262, 24876, 2774, 7944, 2444, 5090, 13173, 314, 39834, 25, 198, 406, 2297, 318, 2186, 25, 220, 22, 15, 8, 198, 285, 279, 6966, 2297, 314, 18880, 4965, 318, 37488, 1387, 17268, 18408, 198, 285, 369, 867, 8, 198, 406, 2763, 66708, 318, 2186, 25, 11363, 198, 285, 886, 421, 668, 381, 20671, 2789, 310, 279, 1118, 1500, 314, 18880, 198, 285, 2468, 13, 220, 82699, 6757, 279, 1500, 579, 2297, 13, 198, 406, 16905, 66708, 318, 2186, 25, 11363, 198, 285, 886, 421, 668, 381, 20671, 2789, 310, 660, 4965, 3462, 279, 1118, 198, 285, 314, 18880, 2468, 26, 1048, 14159, 6757, 1754, 1500, 579, 2297, 13, 198, 406, 8952, 56008, 318, 2186, 25, 804, 8, 198, 285, 48455, 21700, 303, 1879, 1414, 310, 12258, 1518, 4473, 8427, 13, 198, 285, 8618, 5474, 668, 3512, 220, 15, 5074, 359, 5999, 2073, 6, 12258, 11, 11313, 383, 1141, 2234, 198, 285, 303, 1141, 1500, 13, 220, 1368, 867, 11, 1754, 5474, 369, 11421, 430, 264, 3074, 3542, 13, 198, 406, 5474, 2073, 318, 2186, 25, 220, 23, 8, 198, 285, 48455, 21700, 303, 1879, 1414, 310, 220, 15, 5074, 359, 5999, 2073, 6, 12258, 11, 7014, 198, 285, 359, 31251, 56008, 6, 369, 867, 13, 198, 406, 8032, 83201, 318, 2186, 25, 804, 8, 198, 285, 28592, 660, 35169, 5587, 303, 279, 1879, 1414, 539, 12258, 198, 285, 1238, 5474, 14046, 13, 220, 6817, 421, 413, 8952, 56008, 369, 867, 321, 198, 285, 8032, 83201, 369, 804, 11, 1396, 5474, 668, 381, 15625, 310, 264, 198, 285, 3074, 3433, 0, 198, 406, 4884, 46720, 6018, 793, 318, 2186, 25, 867, 8, 198, 285, 28320, 421, 11316, 64541, 59429, 369, 2590, 7854, 198, 285, 539, 1330, 12258, 13, 220, 3978, 539, 1581, 1521, 279, 11767, 369, 198, 285, 318, 354, 46261, 2752, 8, 65669, 13, 198, 406, 1387, 17268, 18408, 318, 2186, 25, 804, 8, 198, 285, 15164, 4105, 4860, 1056, 359, 2998, 4282, 220, 1368, 867, 11, 1782, 4105, 668, 524, 198, 285, 381, 10548, 11, 321, 1010, 4965, 2493, 248046, 198, 248045, 74455, 198, 248068, 271, 248069, 271]} diff --git a/zig/tests/cuda/flashnext/lr7/results.json b/zig/tests/cuda/flashnext/lr7/results.json new file mode 100644 index 0000000000..50ff059a2d --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/results.json @@ -0,0 +1,920 @@ +{ + "info": { + "startup_s": 93.63, + "max_len": 1536, + "prefill_rows": 4096, + "nch": 3, + "qsa": false, + "idw": 2052, + "eos": [ + 248046, + 248044 + ], + "torch": "2.13.0a0+9186a08b2c.nv26.07", + "gpu": "NVIDIA GB10", + "prompt_lengths": { + "short": 19, + "mid": 52, + "long": 600 + } + }, + "results": { + "serial/short": { + "tokens": [ + 760, + 6511, + 314, + 9338, + 369, + 2972, + 57590, + 159034, + 248046, + 198, + 248044, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198 + ], + "wall_s": 14.6758, + "prefill_s": 0.3383, + "cached": 0, + "drafts": false, + "decode_s": 14.3201, + "rounds": 31, + "decode_tps": 2.16 + }, + "teacher/short": { + "tokens": [ + 760, + 6511, + 314, + 9338, + 369, + 2972, + 57590, + 159034, + 248046, + 198, + 248044, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198 + ], + "wall_s": 1.494, + "decode_ms_per_token": 28.396, + "top5": [ + [ + [ + 760, + 24.0 + ], + [ + 3710, + 15.875 + ], + [ + 3733, + 14.3125 + ], + [ + 561, + 14.3125 + ], + [ + 1206, + 13.625 + ] + ], + [ + [ + 6511, + 25.0 + ], + [ + 3177, + 14.1875 + ], + [ + 1994, + 13.875 + ], + [ + 5839, + 13.25 + ], + [ + 1822, + 13.125 + ] + ], + [ + [ + 314, + 25.5 + ], + [ + 321, + 19.125 + ], + [ + 3177, + 16.0 + ], + [ + 11, + 14.3125 + ], + [ + 318, + 14.1875 + ] + ], + [ + [ + 9338, + 28.5 + ], + [ + 279, + 17.75 + ], + [ + 2972, + 17.25 + ], + [ + 47358, + 15.1875 + ], + [ + 11751, + 14.75 + ] + ] + ] + }, + "teacher32/short": { + "tokens": [ + 760, + 6511, + 314, + 9338, + 369, + 2972, + 57590, + 159034, + 248046, + 198, + 248044, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198, + 248045, + 248046, + 198 + ], + "wall_s": 1.4951, + "decode_ms_per_token": 27.503, + "top5": [ + [ + [ + 760, + 23.875 + ], + [ + 3710, + 15.625 + ], + [ + 561, + 14.375 + ], + [ + 3733, + 14.1875 + ], + [ + 1206, + 13.5625 + ] + ], + [ + [ + 6511, + 24.75 + ], + [ + 3177, + 14.0 + ], + [ + 1994, + 13.625 + ], + [ + 1822, + 13.3125 + ], + [ + 5839, + 13.25 + ] + ], + [ + [ + 314, + 25.125 + ], + [ + 321, + 19.375 + ], + [ + 3177, + 16.0 + ], + [ + 318, + 14.4375 + ], + [ + 11, + 14.375 + ] + ], + [ + [ + 9338, + 28.125 + ], + [ + 2972, + 17.625 + ], + [ + 279, + 17.625 + ], + [ + 47358, + 14.6875 + ], + [ + 11751, + 14.625 + ] + ] + ] + }, + "serial/mid": { + "tokens": [ + 29108, + 4009, + 5435, + 660, + 7736, + 11, + 694, + 22602, + 6105, + 89661, + 43000, + 777, + 6497, + 948, + 781, + 4062, + 6605, + 33470, + 303, + 279, + 16078, + 11, + 3154, + 279, + 12515, + 4833, + 6105, + 2261, + 279, + 1834, + 13, + 2332 + ], + "wall_s": 1.3186, + "prefill_s": 0.3477, + "cached": 0, + "drafts": false, + "decode_s": 0.9708, + "rounds": 31, + "decode_tps": 31.93 + }, + "teacher/mid": { + "tokens": [ + 29108, + 4009, + 5435, + 660, + 7736, + 11, + 694, + 22602, + 6105, + 89661, + 43000, + 777, + 6497, + 948, + 781, + 4062, + 6605, + 33470, + 303, + 279, + 16078, + 11, + 3154, + 279, + 12515, + 4833, + 6105, + 2261, + 279, + 1834, + 13, + 2332 + ], + "wall_s": 2.2782, + "decode_ms_per_token": 28.083, + "top5": [ + [ + [ + 29108, + 23.25 + ], + [ + 13672, + 19.0 + ], + [ + 15542, + 18.25 + ], + [ + 760, + 16.25 + ], + [ + 10025, + 14.75 + ] + ], + [ + [ + 4009, + 25.625 + ], + [ + 5603, + 15.625 + ], + [ + 3000, + 14.375 + ], + [ + 37728, + 13.9375 + ], + [ + 4022, + 13.5625 + ] + ], + [ + [ + 5435, + 22.75 + ], + [ + 369, + 20.0 + ], + [ + 7701, + 19.875 + ], + [ + 16661, + 19.0 + ], + [ + 27891, + 18.75 + ] + ], + [ + [ + 660, + 26.75 + ], + [ + 264, + 21.0 + ], + [ + 1599, + 19.0 + ], + [ + 1396, + 17.75 + ], + [ + 7736, + 16.75 + ] + ] + ] + }, + "teacher32/mid": { + "tokens": [ + 29108, + 4009, + 5435, + 660, + 7736, + 11, + 694, + 22602, + 6105, + 89661, + 43000, + 777, + 6497, + 948, + 781, + 4062, + 6605, + 33470, + 303, + 279, + 16078, + 11, + 3154, + 279, + 12515, + 4833, + 6105, + 2261, + 279, + 1834, + 13, + 2332 + ], + "wall_s": 2.2279, + "decode_ms_per_token": 27.005, + "top5": [ + [ + [ + 29108, + 23.125 + ], + [ + 13672, + 18.875 + ], + [ + 15542, + 18.125 + ], + [ + 760, + 16.25 + ], + [ + 37728, + 14.75 + ] + ], + [ + [ + 4009, + 25.5 + ], + [ + 5603, + 15.625 + ], + [ + 3000, + 14.3125 + ], + [ + 37728, + 13.9375 + ], + [ + 4022, + 13.5625 + ] + ], + [ + [ + 5435, + 22.75 + ], + [ + 369, + 19.875 + ], + [ + 7701, + 19.75 + ], + [ + 16661, + 19.125 + ], + [ + 27891, + 18.75 + ] + ], + [ + [ + 660, + 26.75 + ], + [ + 264, + 21.0 + ], + [ + 1599, + 19.0 + ], + [ + 1396, + 17.625 + ], + [ + 7736, + 16.875 + ] + ] + ] + }, + "serial/long": { + "tokens": [ + 13962, + 20858, + 314, + 279, + 13317, + 271, + 1919, + 12654, + 4536, + 5529, + 28459, + 364, + 2972, + 1272, + 39834, + 321, + 20621, + 159034, + 1049, + 369, + 5995, + 310, + 3443, + 1414, + 539, + 14293, + 1248, + 4965, + 1083, + 22602, + 5981, + 421 + ], + "wall_s": 1.8294, + "prefill_s": 0.852, + "cached": 0, + "drafts": false, + "decode_s": 0.9773, + "rounds": 31, + "decode_tps": 31.72 + }, + "teacher/long": { + "tokens": [ + 13962, + 20858, + 314, + 279, + 13317, + 271, + 1919, + 12654, + 4536, + 5529, + 28459, + 364, + 2972, + 1272, + 39834, + 321, + 20621, + 159034, + 11116, + 5839, + 7193, + 369, + 310, + 3443, + 1414, + 539, + 14293, + 1248, + 4965, + 1083, + 22602, + 5981 + ], + "wall_s": 17.2408, + "decode_ms_per_token": 29.059, + "top5": [ + [ + [ + 13962, + 22.375 + ], + [ + 1919, + 20.0 + ], + [ + 760, + 19.125 + ], + [ + 27775, + 18.75 + ], + [ + 550, + 18.375 + ] + ], + [ + [ + 20858, + 22.5 + ], + [ + 13317, + 19.875 + ], + [ + 3437, + 16.125 + ], + [ + 33655, + 14.5 + ], + [ + 2972, + 14.4375 + ] + ], + [ + [ + 314, + 22.625 + ], + [ + 198, + 21.75 + ], + [ + 271, + 21.0 + ], + [ + 25, + 15.25 + ], + [ + 13317, + 13.0 + ] + ], + [ + [ + 279, + 22.75 + ], + [ + 13317, + 20.125 + ], + [ + 5534, + 18.875 + ], + [ + 12654, + 17.5 + ], + [ + 1510, + 17.5 + ] + ] + ] + }, + "teacher32/long": { + "tokens": [ + 13962, + 20858, + 314, + 279, + 13317, + 271, + 1919, + 12654, + 4536, + 5529, + 28459, + 364, + 2972, + 1272, + 39834, + 321, + 20621, + 159034, + 11116, + 5839, + 7193, + 369, + 310, + 3443, + 1414, + 539, + 14293, + 1248, + 4965, + 1083, + 22602, + 5981 + ], + "wall_s": 17.1927, + "decode_ms_per_token": 28.848, + "top5": [ + [ + [ + 13962, + 22.375 + ], + [ + 1919, + 20.0 + ], + [ + 760, + 19.125 + ], + [ + 27775, + 18.625 + ], + [ + 550, + 18.375 + ] + ], + [ + [ + 20858, + 22.5 + ], + [ + 13317, + 20.0 + ], + [ + 3437, + 16.125 + ], + [ + 33655, + 14.5625 + ], + [ + 2972, + 14.4375 + ] + ], + [ + [ + 314, + 22.625 + ], + [ + 198, + 21.875 + ], + [ + 271, + 21.0 + ], + [ + 25, + 15.25 + ], + [ + 13317, + 13.0 + ] + ], + [ + [ + 279, + 22.75 + ], + [ + 13317, + 20.125 + ], + [ + 5534, + 18.625 + ], + [ + 1510, + 17.625 + ], + [ + 12654, + 17.5 + ] + ] + ] + } + } +} diff --git a/zig/tests/cuda/flashnext/lr7/short.ids b/zig/tests/cuda/flashnext/lr7/short.ids new file mode 100644 index 0000000000..23c34f9c55 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/short.ids @@ -0,0 +1 @@ +248045 846 198 3710 369 279 6511 314 9338 30 248046 198 248045 74455 198 248068 271 248069 271 From 8391cac7eff1f9634f4fc638489a7f4405256f25 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:17:39 -0500 Subject: [PATCH 030/115] native/cuda.zig takes each CUDA family's prompt-cache copies, cache rules and device weight bytes from what it declares; Flash Next registers The registry gains flashnext.native. What was Nemotron's by name is now read from the family's declarations: its prompt-cache copies (snap_bytes, snap_save, snap_restore, snap_drop, and snap_ready's spare buffers), the store's rules (cache_rules; without them a state keeps at any row with the MTP head's one token on, #670), the quiet wait before a warm pass and the lone driver's hand-over budget. Nemotron declares the copies, the wait and the budget, so its paths are unchanged; a family without copies gets no prompt cache, and the startup line says so. A family that keeps part of its checkpoint off the device says what it uploads (deviceWeightBytes) and admission counts that, not every shard: Flash Next leaves its n-gram tables mapped and loads no vision tower, so the MLX 4-bit checkpoint uploads 73.3 of its 105.4 GiB, where counting every shard made the default reserve on a GB10 refuse it by 0.1 GiB. --slide is refused before anything loads for a family that declares no learner. The Flash Next family names what the host reads: one segment, its prompt rows, each stream's bytes, its refusals in words, and the native host's lone-driver signature. --- zig/src/families/flashnext/cuda_lanes.zig | 37 ++++++++++- zig/src/families/flashnext/cuda_prompt.zig | 14 +++++ zig/src/families/flashnext/cuda_prompt_wt.zig | 41 +++++++++++++ zig/src/native/cuda.zig | 61 ++++++++++++------- 4 files changed, 129 insertions(+), 24 deletions(-) diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index cfacce15b3..9e7eddf98d 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -4,6 +4,7 @@ const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); const prompt = @import("cuda_prompt.zig"); +const weights = @import("cuda_prompt_wt.zig"); const be = lanes.backend; const ring = 1024; @@ -15,18 +16,44 @@ pub fn Namespace(comptime Tri: type) type { pub const formats: []const []const u8 = &.{"mlx-q4g32"}; // dense attention over the captured scratch (NCH chunks of 512 keys, QSA off): no sparse indexer yet pub const default_context: i64 = @intCast(prompt.max_tokens); - pub const prefill_step: u32 = 0; + // a prompt runs as one pass of fused decode rows, the whole prompt at once: no staggered segments + pub const max_segments: u32 = 1; + pub const prompt_rows: u32 = @intCast(prompt.max_tokens); - pub const Options = struct { context: usize, drafts: bool }; + pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + + /// One stream's prompt-and-reply driver, as native/cuda.zig types it; this family has none. + pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; pub const Loaded = struct { backend: be.Backend, facts: lanes.Model, rows: u32, + /// Device bytes each admitted stream allocates: its own caches at the server's window. + stream_bytes: usize, + /// Every stream allocates its own sequence; none rides on the engine's. + free_streams: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + learn_bytes: usize = 0, }; + /// What the device holds of the checkpoint, for admission: the n-gram tables stay in the mapped shards. + /// 0 when the headers do not read, and the load then names what is wrong. + pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + return weights.deviceBytes(gpa, io, dir) catch 0; + } + + /// A request or load this engine refuses, in words; null: none of its own. + pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", + error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + else => null, + }; + } + const Lane = struct { seq: prompt.Seq }; const Cuda = struct { @@ -97,9 +124,12 @@ pub fn Namespace(comptime Tri: type) type { } }; - pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: []const u8, o: Options) !Loaded { + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: Options) !Loaded { _ = o.drafts; // this model is served one token at a time; there is no draft head if (o.context > prompt.max_tokens) return error.ContextOverDenseAttention; + if (o.segments != 1 or o.slide) return error.UnsupportedOption; + // every launch is a captured Triton cubin: TENSORFOLD_CUDA_KERNELS names the set + const kernels = captured orelse return error.CapturedKernelsRequired; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); @@ -116,6 +146,7 @@ pub fn Namespace(comptime Tri: type) type { .backend = own.backend(), .facts = Cuda.facts(), .rows = 1, + .stream_bytes = prompt.seqBytes(o.context), .ctx = own, .deinit = release, }; diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index e955b7d5d1..d248885d00 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -190,6 +190,14 @@ pub const vocabulary: usize = vocab; /// Prompt plus reply tokens a sequence can hold: the attention kernels' NCH chunks of 512 keys. pub const max_tokens: usize = weights.nch * 512; +/// Device bytes `Engine.newSeq(cap)` allocates: the convolution tails, the recurrent states, the key and value +/// caches with the indexer keys, and the n-gram tail. +pub fn seqBytes(cap: usize) usize { + const linear = linear_n * ((gdn.taps - 1) * gdn.conv_dim * 2 + state_n * 4); + const attn = attn_n * cap * (2 * kv_heads * head_dim * 2 + index_dim * 2); + return linear + attn + ple_tail_n * wide * 2; +} + /// One loaded checkpoint: the embedding table, the mixer, the head, and a scratch row the forwards share. pub const Engine = struct { gpa: std.mem.Allocator, @@ -436,3 +444,9 @@ test "linear layers skip every fourth index and attention starts at 3" { try std.testing.expectEqual(@as(usize, 0), attnIndex(3)); try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); } + +test "a sequence's device bytes grow by the twelve attention layers' keys, values and indexer keys per token" { + const per_token = attn_n * (2 * kv_heads * head_dim * 2 + index_dim * 2); + try std.testing.expectEqual(per_token, seqBytes(2) - seqBytes(1)); + try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index a2598efd7d..d91216d5c7 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -278,6 +278,38 @@ pub fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { return out; } +/// A tensor the engine leaves off the device: the n-gram tables (rows gathered from the mapped shards), the draft +/// head (served without drafts) and the vision tower (text only). +pub fn hostOnly(name: []const u8) bool { + return std.mem.indexOf(u8, name, ".ple_embedding.ngram_embedding.") != null or + std.mem.startsWith(u8, name, "language_model.mtp.") or + std.mem.startsWith(u8, name, "vision_tower."); +} + +/// The checkpoint's tensor bytes the engine uploads, every shard the index names read for its header alone. +pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { + var store = try Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = 0; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var sh = try Shard.load(gpa, io, path); + defer sh.close(); + var ni = sh.names.iterator(); + while (ni.next()) |e| { + if (!hostOnly(e.key_ptr.*)) total += e.value_ptr.end - e.value_ptr.begin; + } + } + return total; +} + pub fn asF32(gpa: std.mem.Allocator, t: Tensor) ![]f32 { if (t.rank != 1) return error.UnexpectedTensor; const out = try gpa.alloc(f32, t.dim(0)); @@ -396,3 +428,12 @@ pub fn stitch(gpa: std.mem.Allocator, store: *Store, prefix: []const u8, names: }; return .{ .w = try cat(gpa, ws[0..n]), .s = try cat(gpa, ss[0..n]), .b = try cat(gpa, bs[0..n]), .n = n }; } + +test "the n-gram tables, the draft head and the vision tower stay off the device" { + try std.testing.expect(hostOnly("language_model.model.layers.2.ple.ple_embedding.ngram_embedding.shards.7.weight")); + try std.testing.expect(hostOnly("language_model.mtp.layers.0.mlp.switch_mlp.up_proj.weight")); + try std.testing.expect(hostOnly("vision_tower.blocks.0.attn.qkv.weight")); + try std.testing.expect(!hostOnly("language_model.model.layers.2.ple.ple_embedding.ngram_heads_offsets")); + try std.testing.expect(!hostOnly("language_model.model.layers.0.mlp.switch_mlp.up_proj.weight")); + try std.testing.expect(!hostOnly("language_model.lm_head.weight")); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index c599ba1582..564fd63ee1 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -13,13 +13,30 @@ const nemotron_slide = @import("nemotron_slide.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{ nemotron.native, flashnext.native }; -/// Nemotron's prompt-cache copies. A file-level value so the store's pointer stays valid for the process. -const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ - .bytes = nemotron.native.snap_bytes, - .save = nemotron.native.snap_save, - .restore = nemotron.native.snap_restore, - .drop = nemotron.native.snap_drop, -}; +/// A family's prompt-cache copies, a container-level value so the store's pointer lives on; null: it keeps none. +fn snapshots(comptime F: type) ?*const api.prompt_cache.Snapshots.VTable { + if (!@hasDecl(F, "snap_bytes")) return null; + return &struct { + const v: api.prompt_cache.Snapshots.VTable = .{ .bytes = F.snap_bytes, .save = F.snap_save, .restore = F.snap_restore, .drop = F.snap_drop }; + }.v; +} + +/// Where a family's kept states sit: its own `cache_rules`, else any row with the head's one token on (#670). +fn cacheRules(comptime F: type) api.prompt_cache.Rules { + if (!@hasDecl(F, "cache_rules")) return .{ .lookahead = 1, .min_prompt = 0, .warm = if (@hasDecl(F, "warm_quiet")) .waits else .off }; + const r = F.cache_rules; + return .{ .lookahead = r.lookahead, .planned = r.planned, .grid = if (r.planned) F.prompt_rows else 0, .min_gap = r.min_gap, .min_prompt = r.min_prompt }; +} + +/// What the lane host takes from the family: its quiet wait before a warm pass, its lone driver's hand-over budget. +const PerFamily = struct { warm_quiet: u64 = 0, lone_min_budget: u32 = 0 }; + +fn perFamily(comptime F: type) PerFamily { + return .{ + .warm_quiet = if (@hasDecl(F, "warm_quiet")) F.warm_quiet else 0, + .lone_min_budget = if (@hasDecl(F, "lone_min_budget")) F.lone_min_budget else 0, + }; +} pub const backends: []const []const u8 = &.{"cuda"}; pub const families: []const api.Family = blk: { @@ -186,7 +203,7 @@ const Host = struct { } /// The round loop and the lane host over the family's backend, its thread started. - fn serve(h: *Host, io: std.Io, facts: lanes.Model, rows: u32, info: api.Info, explain: ?api.Explain) !void { + fn serve(h: *Host, io: std.Io, facts: lanes.Model, rows: u32, info: api.Info, explain: ?api.Explain, per: PerFamily) !void { h.cfg = try lanes.Config.init(h.gpa, facts, rows, rows - 1); errdefer h.cfg.deinit(h.gpa); h.clock = .{ .io = io }; @@ -196,8 +213,8 @@ const Host = struct { h.host.memory = if (h.gpu != null) .{ .read = readMemory } else null; h.host.explain = explain; h.host.cache = h.store; - h.host.warm_quiet = if (info.warm_turns) nemotron.native.warm_quiet else 0; - if (h.lone != null) h.host.lone = .{ .ctx = h, .run = loneRun, .sampled = true, .logprobs = h.inner.vtable.first_row != null, .structures = h.inner.vtable.first_masked != null, .min_budget = nemotron.native.lone_min_budget, .resumes = true }; + h.host.warm_quiet = if (info.warm_turns) per.warm_quiet else 0; + if (h.lone != null) h.host.lone = .{ .ctx = h, .run = loneRun, .sampled = true, .logprobs = h.inner.vtable.first_row != null, .structures = h.inner.vtable.first_masked != null, .min_budget = per.lone_min_budget, .resumes = true }; if (h.learn != null) h.host.learner = h.learner(); try h.host.start(); } @@ -273,8 +290,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O if (!@hasDecl(F, "slide")) { problem.* = try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on CUDA yet", .{o.model_type}); return null; - } - if (try nemotron_slide.refusal(F, a, io, o.dir, "`cp -RL` (a copy that follows links)")) |why| { + } else if (try nemotron_slide.refusal(F, a, io, o.dir, "`cp -RL` (a copy that follows links)")) |why| { problem.* = why; return null; } @@ -323,7 +339,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }, }; const before = try pool(a, io, &g.ctx, problem) orelse return null; - const weights = weightBytes(io, o.dir); + // a family that keeps part of its checkpoint off the device says what it uploads; else every shard's bytes + const weights = if (@hasDecl(F, "deviceWeightBytes")) F.deviceWeightBytes(a, io, o.dir) else weightBytes(io, o.dir); const held0 = cuda.usage(false).device; if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); @@ -348,24 +365,26 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + const snaps = comptime snapshots(F); + const cache = if (snaps != null) budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)) else 0; h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (snaps) |vtable| if (cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); - // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet - store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = &nemotron_snaps }, .{ .lookahead = 1, .min_prompt = 0, .warm = .waits }, cache - try nemotron.native.snap_ready(loaded.target, cache)); + // states kept where the family's rules allow; its spare state buffers come out of the budget + const spare = if (@hasDecl(F, "snap_ready")) try F.snap_ready(loaded.target, cache) else 0; + store.* = pc.Store.init(gpa, .{ .ptr = loaded.target, .vtable = vtable }, comptime cacheRules(F), cache - spare); h.store = store; - } + }; errdefer if (h.store) |store| { store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { + if (@hasDecl(F, "slide") and loaded.learns != null) { const Slide = nemotron_slide.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); @@ -385,7 +404,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null and @hasDecl(F, "warm_quiet"), .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }, comptime perFamily(F)); opened = true; return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } @@ -450,7 +469,7 @@ test "a cancel during a long prompt stops it, and the next request's reply is un const h = try gpa.create(Host); defer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = null, .family = &target, .release = Nothing.release, .inner = target.backend(), .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined }; - try h.serve(std.testing.io, .{ .exact_width = 8, .mtp = true, .speculate = true, .drafts = 7, .hidden_rows = true }, 8, .{ .lanes = 2 }, null); + try h.serve(std.testing.io, .{ .exact_width = 8, .mtp = true, .speculate = true, .drafts = 7, .hidden_rows = true }, 8, .{ .lanes = 2 }, null, .{}); defer { h.host.stop(); h.core.deinit(); From bd359401580cea7e9f941f5c4dbe078ce91af093 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:10:31 -0500 Subject: [PATCH 031/115] CUDA experts over a window's or a prompt chunk's rows, each row with the decode kernel's sums: experts_rows.cu experts_rows.cu is experts.cu's decode expert kernel with two 16-row tiles a warp, so a warp reads each group of an expert's weights once for up to 32 pair rows. Every row keeps the decode kernel's arithmetic: the same group sums, the same two m16n8k16 steps from zero a group, the same fma epilogue order and the same SwiGLU, built with experts.cu's flags. A row's bits therefore do not depend on how many rows share its window or chunk, which is what the lane core's exactness contract needs of a lane kernel. It is instantiated at groups of 32 (gate and up stacked with SwiGLU, and the down projection to fp32), the MLX 4-bit group-32 experts Flash Next reads. cuda.experts.RowsExperts resolves the two instances and launches them on the shared plan (cuda.grouped.Router at its 32-pair tile). `tf-cuda-test experts-rows` packs random group-32 experts with experts_pack.cu, routes 1, 2, 5, 16, 33, 100 and 300 rows (three routed experts and a shared one a row) and checks every pair row of the rows kernel against experts.cu's one-row launches of the same row, byte for byte, for the SwiGLU rows and the fp32 down rows. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/experts_rows.cu | 144 +++++++++++++++++++++ zig/src/cuda/experts.zig | 38 +++++- zig/src/cuda/kernels.zig | 1 + zig/tests/cuda/experts_rows_tests.zig | 177 ++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 6 files changed, 361 insertions(+), 3 deletions(-) create mode 100644 zig/kernels/cuda/experts_rows.cu create mode 100644 zig/tests/cuda/experts_rows_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index dc2d682294..7bb23f2285 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -23,6 +23,7 @@ const kernels = [_]Kernel{ .{ .name = "experts", .flags = &.{"-O3"} }, // cuda/experts.py, tensorfold_experts_v7 .{ .name = "experts_prefill", .flags = &.{"-O3"} }, .{ .name = "experts_pack", .flags = &.{"-O3"} }, + .{ .name = "experts_rows", .flags = &.{"-O3"} }, // experts.cu's decode sums over a window or prompt chunk's rows .{ .name = "prefill_attention", .flags = &.{ "-O3", "--fmad=false" } }, // tensorfold_prefill_attention_v1 .{ .name = "scan_rows", .flags = &.{ "-O3", "--fmad=false" } }, // nemotron_h/cuda/mamba.py .{ .name = "nemotron_ops", .flags = &.{"-O3"} }, // ours: Nemotron's serial feed, routed plan and rests diff --git a/zig/kernels/cuda/experts_rows.cu b/zig/kernels/cuda/experts_rows.cu new file mode 100644 index 0000000000..8ec6d0d162 --- /dev/null +++ b/zig/kernels/cuda/experts_rows.cu @@ -0,0 +1,144 @@ +// experts.cu's decode expert kernel over up to 16 * RT pair rows a warp; a row's bits do not depend on its window. + +#include +#include + +#include "experts.cuh" + +namespace tf_experts_rows { + +__device__ __forceinline__ uint4 ld_w(const uint4* p) { + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +__device__ __forceinline__ uint4 ld_x(const __nv_bfloat16* p) { return __ldg(reinterpret_cast(p)); } + +// experts.cu's group_sum: a lane's 8 inputs a row in order, then the quad's partners. +template +__device__ __forceinline__ float group_sum(const uint4 (&v)[XV]) { + float s = 0.f; +#pragma unroll + for (int c = 0; c < XV; ++c) +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t u = comp(v[c], i); + s += lo_f(u); + s += hi_f(u); + } + s += __shfl_xor_sync(0xffffffffu, s, 1); + s += __shfl_xor_sync(0xffffffffu, s, 2); + return s; +} + +template +__global__ void __launch_bounds__(WARPS * 32) + expert_rows(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, int KG, + int NB, const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, float limit) { + using G = Geo; + const int lane = threadIdx.x & 31, gq = lane >> 2, t = lane & 3; + const int units = __ldg(counts) * NB; + for (int unit = blockIdx.x * WARPS + (threadIdx.x >> 5); unit < units; unit += gridDim.x * WARPS) { + const int it = unit / NB, cb = unit - it * NB; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + bool v0[RT], v1[RT]; + int pr0[RT], pr1[RT]; + const __nv_bfloat16* x0[RT]; + const __nv_bfloat16* x1[RT]; + int nt = 0; // row tiles of 16 this item fills +#pragma unroll + for (int r = 0; r < RT; ++r) { + v0[r] = 16 * r + gq < cnt; + v1[r] = 16 * r + gq + 8 < cnt; + pr0[r] = v0[r] ? __ldg(members + first + 16 * r + gq) : 0; + pr1[r] = v1[r] ? __ldg(members + first + 16 * r + gq + 8) : 0; + const int r0 = slots ? pr0[r] / slots : pr0[r], r1 = slots ? pr1[r] / slots : pr1[r]; + x0[r] = X + (size_t)r0 * x_stride + t * (GS / 4); + x1[r] = X + (size_t)r1 * x_stride + t * (GS / 4); + if (16 * r < cnt) nt = r + 1; + } + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (M * G::BLOCK); + float acc[M][RT][NTW][4]; +#pragma unroll + for (int m = 0; m < M; ++m) +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[m][r][j][0] = acc[m][r][j][1] = acc[m][r][j][2] = acc[m][r][j][3] = 0.f; + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + for (int g = 0; g < KG; ++g) { + const uint4* b = blk + (size_t)g * (M * G::BLOCK); + uint4 w[M][G::WV], sb[M][G::SBV]; +#pragma unroll + for (int m = 0; m < M; ++m) { +#pragma unroll + for (int c = 0; c < G::WV; ++c) w[m][c] = ld_w(b + m * G::BLOCK + c * 32 + lane); +#pragma unroll + for (int c = 0; c < G::SBV; ++c) sb[m][c] = ld_w(b + m * G::BLOCK + 32 * G::WV + t * G::SBV + c); + } + const int k0 = g * GS; +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + uint4 xa[G::XV], xb[G::XV]; +#pragma unroll + for (int c = 0; c < G::XV; ++c) { + xa[c] = v0[r] ? ld_x(x0[r] + k0 + 8 * c) : zero; + xb[c] = v1[r] ? ld_x(x1[r] + k0 + 8 * c) : zero; + } + const float xsa = group_sum(xa); + const float xsb = group_sum(xb); +#pragma unroll + for (int m = 0; m < M; ++m) { + float p[NTW][4]; +#pragma unroll + for (int j = 0; j < NTW; ++j) p[j][0] = p[j][1] = p[j][2] = p[j][3] = 0.f; +#pragma unroll + for (int ks = 0; ks < G::KS; ++ks) { + const uint32_t a0 = comp(xa[ks >> 1], 2 * (ks & 1)); + const uint32_t a2 = comp(xa[ks >> 1], 2 * (ks & 1) + 1); + const uint32_t a1 = comp(xb[ks >> 1], 2 * (ks & 1)); + const uint32_t a3 = comp(xb[ks >> 1], 2 * (ks & 1) + 1); + const int sh = (ks & 1) ? 8 : 0; +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const int wi = j * (GS / 32) + (ks >> 1); + const uint32_t word = comp(w[m][wi >> 2], wi & 3); + mma(p[j], a0, a1, a2, a3, nib2(word, sh), nib2(word, sh + 4)); + } + } +#pragma unroll + for (int j = 0; j < NTW; ++j) { + const uint32_t sp = comp(sb[m][j >> 2], j & 3); + const uint32_t bp = comp(sb[m][(NTW + j) >> 2], (NTW + j) & 3); + const float s0 = lo_f(sp), s1 = hi_f(sp), b0 = lo_f(bp), b1 = hi_f(bp); + float(&a)[4] = acc[m][r][j]; + a[0] = fmaf(xsa, b0, fmaf(p[j][0], s0, a[0])); + a[1] = fmaf(xsa, b1, fmaf(p[j][1], s1, a[1])); + if (v1[r]) { + a[2] = fmaf(xsb, b0, fmaf(p[j][2], s0, a[2])); + a[3] = fmaf(xsb, b1, fmaf(p[j][3], s1, a[3])); + } + } + } + } + } +#pragma unroll + for (int r = 0; r < RT; ++r) { + if (r >= nt) break; + epilogue(acc, r, out, N, cb * COLS + 2 * t, pr0[r], pr1[r], v0[r], v1[r], limit); + } + } +} + +} // namespace tf_experts_rows + +// Groups of 32, two row tiles (32 rows) a warp: SwiGLU gate and up together (epilogue 2), down to fp32 (epilogue 0). +template __global__ void tf_experts_rows::expert_rows<32, 2, 2, 2, 4>(const __nv_bfloat16*, int, int, const uint4*, + int, int, const int*, const int*, const int*, void*, int, float); +template __global__ void tf_experts_rows::expert_rows<32, 1, 0, 2, 4>(const __nv_bfloat16*, int, int, const uint4*, + int, int, const int*, const int*, const int*, void*, int, float); diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..16b641ef3a 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -126,7 +126,7 @@ pub const Affine4Experts = struct { return if (up) .{ .ptr = l.up, .kg = l.dims / 64, .n = l.width } else .{ .ptr = l.down, .kg = l.width / 64, .n = l.dims }; } - fn args(x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize) launch_.Args { + pub fn argsOf(x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize) launch_.Args { var a: launch_.Args = .{}; a.add(x); a.add(int(x_stride)); @@ -144,7 +144,7 @@ pub const Affine4Experts = struct { pub fn run(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_units: usize) !void { const grid = @min((max_units + 3) / 4, e.resident[if (up) 0 else 1]); if (grid < 1) return; - var a = args(x, x_stride, slots, w, kg, nb, p, out, n); + var a = argsOf(x, x_stride, slots, w, kg, nb, p, out, n); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 }, .shared = 0 }; try launch_.launch(if (up) e.up_fn else e.down_fn, cfg, s, &a); } @@ -153,7 +153,7 @@ pub const Affine4Experts = struct { pub fn runPrompt(e: *const Affine4Experts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, max_items: usize) !void { const grid = max_items * ((nb + 3) / 4); if (grid < 1) return; - var a = args(x, x_stride, slots, w, kg, nb, p, out, n); + var a = argsOf(x, x_stride, slots, w, kg, nb, p, out, n); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 256 }, .shared = prompt_smem }; try launch_.launch(if (up) e.pre_up_fn else e.pre_down_fn, cfg, s, &a); } @@ -174,6 +174,38 @@ pub const Affine4Experts = struct { } }; +/// Affine-4 SwiGLU experts at groups of 32 over pair rows (experts_rows.cu): each row's bits are those it gets alone. +pub const RowsExperts = struct { + up_fn: Function, + down_fn: Function, + + /// Pairs a plan item holds: the kernel's two row tiles of 16. + pub const tile: usize = 32; + + /// The instances' names (cuobjdump -symbols of the experts_rows fatbin): gate and up with SwiGLU, then down. + pub const symbols = struct { + pub const up = "_ZN15tf_experts_rows11expert_rowsILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const down = "_ZN15tf_experts_rows11expert_rowsILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + /// experts.cu's one-row forms at groups of 32, whose bits these give every row (tf-cuda-test experts-rows). + pub const decode_up = "_ZN10tf_experts13expert_kernelILi32ELi2ELi2ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + pub const decode_down = "_ZN10tf_experts13expert_kernelILi32ELi1ELi0ELi2ELi4EEEvPK13__nv_bfloat16iiPK5uint4iiPKiS8_S8_Pvif"; + }; + + /// From a loaded experts_rows module. + pub fn resolve(mod: Module) !RowsExperts { + return .{ .up_fn = try mod.function(symbols.up), .down_fn = try mod.function(symbols.down) }; + } + + /// `items` plan items: `up` maps token rows to SwiGLU's bf16, else pair rows to fp32; K in `kg` groups of 32. + pub fn run(e: *const RowsExperts, s: Stream, up: bool, x: u64, x_stride: usize, slots: usize, w: u64, kg: usize, nb: usize, p: Plan, out: u64, n: usize, items: usize) !void { + const grid = @min((items * nb + 3) / 4, 4096); + if (grid < 1) return; + var a = Affine4Experts.argsOf(x, x_stride, slots, w, kg, nb, p, out, n); + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 }, .shared = 0 }; + try launch_.launch(if (up) e.up_fn else e.down_fn, cfg, s, &a); + } +}; + /// Block-FP8 experts (fp8/experts.cu): the Python packing, kernels and grid choice, decode and prompt forms. pub const Fp8Experts = struct { mod: Module, diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 4f3a78a50c..78260422ba 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -21,6 +21,7 @@ pub const qmm_prefill: []const u8 = if (available) &Blob("fatbin_qmm_prefill").b pub const experts: []const u8 = if (available) &Blob("fatbin_experts").bytes else &.{}; pub const experts_prefill: []const u8 = if (available) &Blob("fatbin_experts_prefill").bytes else &.{}; pub const experts_pack: []const u8 = if (available) &Blob("fatbin_experts_pack").bytes else &.{}; +pub const experts_rows: []const u8 = if (available) &Blob("fatbin_experts_rows").bytes else &.{}; pub const prefill_attention: []const u8 = if (available) &Blob("fatbin_prefill_attention").bytes else &.{}; pub const scan_rows: []const u8 = if (available) &Blob("fatbin_scan_rows").bytes else &.{}; pub const nemotron_ops: []const u8 = if (available) &Blob("fatbin_nemotron_ops").bytes else &.{}; diff --git a/zig/tests/cuda/experts_rows_tests.zig b/zig/tests/cuda/experts_rows_tests.zig new file mode 100644 index 0000000000..cb5ecd86ac --- /dev/null +++ b/zig/tests/cuda/experts_rows_tests.zig @@ -0,0 +1,177 @@ +//! experts_rows.cu against experts.cu's decode kernel: every pair row of a window or chunk, byte for byte. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const grouped = cuda.grouped; +const Rows = cuda.experts.RowsExperts; + +/// experts_pack.cu's groups of 32 (one block: 160 uint32 words of codes, scales and biases). +const pack_symbol = "_ZN15tf_experts_pack11pack_kernelILi1EEEvPKjPKtS4_Pjiiii"; +const block_bytes = 160 * 4; +const dims = 256; +const width = 128; +const routed = 16; +const experts = routed + 1; +const slots = 4; + +const Rig = struct { + gpu: Gpu, + s: cuda.Stream, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn zeros(r: *Rig, n: usize) !u64 { + var b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(n, 16)); + try b.fill8(0, null); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn get(r: *Rig, ptr: u64, n: usize) ![]u8 { + try r.s.synchronize(); + const out = try r.gpu.gpa.alloc(u8, n); + try r.gpu.d.check(r.gpu.d.api.cuMemcpyDtoH_v2(out.ptr, ptr, n), "cuMemcpyDtoH"); + return out; + } + + fn bf16(v: f32) u16 { + return @truncate(@as(u32, @bitCast(v)) >> 16); + } + + /// Random MLX 4-bit words with small scales and biases, packed for `count` experts of n outputs from k inputs. + fn table(r: *Rig, pack: cuda.Function, count: usize, n: usize, k: usize) !u64 { + const a = r.gpu.gpa; + const words = try a.alloc(u32, count * n * k / 8); + defer a.free(words); + for (words) |*w| w.* = r.rng.int(u32); + const groups = count * n * k / 32; + const scales = try a.alloc(u16, groups); + defer a.free(scales); + const biases = try a.alloc(u16, groups); + defer a.free(biases); + for (scales, biases) |*s, *b| { + s.* = bf16(0.002 + 0.004 * r.rng.float(f32)); + b.* = bf16(-0.015 + 0.01 * r.rng.float(f32)); + } + const out = try r.zeros(count * (n / 32) * (k / 32) * block_bytes); + var args: cuda.Args = .{}; + for ([_]u64{ try r.dev(std.mem.sliceAsBytes(words)), try r.dev(std.mem.sliceAsBytes(scales)), try r.dev(std.mem.sliceAsBytes(biases)), out }) |v| args.add(v); + for ([_]usize{ n, k / 8, k / 32, n / 32 }) |v| args.add(@as(c_int, @intCast(v))); + try cuda.launch.launch(pack, .{ .grid = .{ .x = @intCast(k / 32), .y = @intCast(n / 32), .z = @intCast(count) }, .block = .{ .x = 160 } }, r.s, &args); + return out; + } + + /// Gate and up blocks side by side, the SwiGLU table the kernels read. + fn stack(r: *Rig, gate: u64, up: u64, blocks: usize) !u64 { + const g = try r.get(gate, blocks * block_bytes); + defer r.gpu.gpa.free(g); + const u = try r.get(up, blocks * block_bytes); + defer r.gpu.gpa.free(u); + const both = try r.gpu.gpa.alloc(u8, 2 * blocks * block_bytes); + defer r.gpu.gpa.free(both); + for (0..blocks) |b| { + @memcpy(both[2 * b * block_bytes ..][0..block_bytes], g[b * block_bytes ..][0..block_bytes]); + @memcpy(both[(2 * b + 1) * block_bytes ..][0..block_bytes], u[b * block_bytes ..][0..block_bytes]); + } + return r.dev(both); + } +}; + +/// The plan scratch for `pairs` pairs at `tile` pairs an item. +fn planFor(r: *Rig, pairs: usize, tile: usize) !grouped.Plan { + const wide = pairs > grouped.small; + return .{ + .members = try r.zeros(pairs * 4), + .items = try r.zeros(grouped.maxItems(pairs, experts, tile) * 12), + .counts = try r.zeros(8), + .rank = try r.zeros((if (wide) pairs else 1) * 4), + .hist = try r.zeros((if (wide) (pairs + 1023) / 1024 * experts else 1) * 4), + }; +} + +fn launch(r: *Rig, f: cuda.Function, x: u64, x_stride: usize, slot_n: usize, w: u64, kg: usize, nb: usize, p: grouped.Plan, out: u64, n: usize, units: usize) !void { + var a = cuda.experts.Affine4Experts.argsOf(x, x_stride, slot_n, w, kg, nb, p, out, n); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast((units + 3) / 4) }, .block = .{ .x = 128 } }, r.s, &a); +} + +pub fn run(gpu: Gpu) !void { + const a = gpu.gpa; + var experts_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer experts_mod.unload(); + var rows_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts_rows); + defer rows_mod.unload(); + var pack_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts_pack); + defer pack_mod.unload(); + const router = try grouped.Router.resolve(experts_mod); + const kernels = try Rows.resolve(rows_mod); + const one_up = try experts_mod.function(Rows.symbols.decode_up); + const one_down = try experts_mod.function(Rows.symbols.decode_down); + var prng = std.Random.DefaultPrng.init(7); + var r: Rig = .{ .gpu = gpu, .s = try cuda.Stream.init(gpu.d, true), .rng = prng.random() }; + defer r.s.deinit(); + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(a); + } + const pack = try pack_mod.function(pack_symbol); + const up_w = try r.stack(try r.table(pack, experts, width, dims), try r.table(pack, experts, width, dims), experts * (width / 32) * (dims / 32)); + const down_w = try r.table(pack, experts, dims, width); + for ([_]usize{ 1, 2, 5, 16, 33, 100, 300 }) |n_rows| { + const pairs = n_rows * slots; + const x = try a.alloc(u16, n_rows * dims); + defer a.free(x); + for (x) |*v| v.* = Rig.bf16(r.rng.float(f32) * 2 - 1); + // each row's routed picks distinct, then the shared expert in the last slot, as Flash Next routes + const picks = try a.alloc(i32, pairs); + defer a.free(picks); + for (0..n_rows) |row| { + var used: [routed]bool = @splat(false); + for (0..slots - 1) |j| { + var e = r.rng.uintLessThan(usize, routed); + while (used[e]) e = (e + 1) % routed; + used[e] = true; + picks[row * slots + j] = @intCast(e); + } + picks[row * slots + slots - 1] = routed; + } + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dpicks = try r.dev(std.mem.sliceAsBytes(picks)); + // the rows kernel: every row in one plan of 32-pair items + const p = try planFor(&r, pairs, Rows.tile); + try router.route(r.s, dpicks, pairs, experts, Rows.tile, p); + const items = pairs / Rows.tile + experts; + const act = try r.zeros(pairs * width * 2); + const y = try r.zeros(pairs * dims * 4); + try kernels.run(r.s, true, dx, dims, slots, up_w, dims / 32, width / 32, p, act, width, items); + try kernels.run(r.s, false, act, width, 0, down_w, width / 32, dims / 32, p, y, dims, items); + // the decode kernel: each row alone, its pairs written where the rows kernel writes them + const act1 = try r.zeros(pairs * width * 2); + const y1 = try r.zeros(pairs * dims * 4); + const p1 = try planFor(&r, slots, 16); + for (0..n_rows) |row| { + try router.route(r.s, dpicks + row * slots * 4, slots, experts, 16, p1); + try launch(&r, one_up, dx + row * dims * 2, dims, slots, up_w, dims / 32, width / 32, p1, act1 + row * slots * width * 2, width, slots * (width / 32)); + try launch(&r, one_down, act1 + row * slots * width * 2, width, 0, down_w, width / 32, dims / 32, p1, y1 + row * slots * dims * 4, dims, slots * (dims / 32)); + } + const got_act = try r.get(act, pairs * width * 2); + defer a.free(got_act); + const want_act = try r.get(act1, pairs * width * 2); + defer a.free(want_act); + try check.sameBytes("experts_rows SwiGLU rows", got_act, want_act); + const got_y = try r.get(y, pairs * dims * 4); + defer a.free(got_y); + const want_y = try r.get(y1, pairs * dims * 4); + defer a.free(want_y); + try check.sameBytes("experts_rows down rows", got_y, want_y); + check.pass("experts-rows: {d} rows ({d} pairs) equal the decode kernel's one-row launches byte for byte", .{ n_rows, pairs }); + } +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..a95f29cef7 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -29,6 +29,7 @@ const norms_tests = @import("norms_tests.zig"); const norm_rows_tests = @import("norm_rows_tests.zig"); const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); +const experts_rows_tests = @import("experts_rows_tests.zig"); const usage = \\usage: tf-cuda-test @@ -50,6 +51,7 @@ const usage = \\ sample sample.cu's keyed draws against the Metal rule's host references (synthetic rows) \\ glue Nemotron's glue kernels against their host references at its shapes (synthetic rows) \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too + \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ window-profile MODEL WIDTHS IDS_FILE a decode window's GPU ms by kernel class (WIDTHS like 1,4,16), beside its graph \\ tree-accept MODEL IDS_FILE COUNT OUT a greedy reply's MTP chain and level-1 runner-up branches at every position (JSON) \\ chunk-costs MODEL IDS_FILE WIDTHS r rows as the decode window graph (r <= 16) and as a prompt chunk (any r) @@ -131,6 +133,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "norms")) return norms_tests.run(gpu); if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); + if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "pocket")) return pocket_tests.run(gpu, if (rest.len > 0) try std.fmt.parseInt(usize, rest[0], 10) else 0); if (std.mem.eql(u8, cmd, "window-profile")) return window_profile.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); From 2484c939c18d1e8282b15ac97ecdc810c79d4a3b Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:10:31 -0500 Subject: [PATCH 032/115] triton_aot_pack.py: a constexpr Triton folded from None is left out, and --jit is repeatable A Triton kernel called with None for a pointer it does not read gets that argument folded into the kernel as a constexpr None; Flash Next's _pool kernel has two (ROPE and DELTA). The packer stopped on them ("constexpr None has no Zig form"); it now leaves them out of the variant's constexprs, since there is nothing for a launch to match. --jit is repeatable like --manifest, and the captures' jit.json files are merged, so one call packs several captures whose kernels come from different JIT functions. A single --jit, or none, works as before. test_aot_manifest.py's fake kernel now has a None constexpr (the packer before this change stops on it) and packs once more with a second capture's jit.json beside the first, which must give the same set. --- tools/zig/triton_aot_pack.py | 8 +++++--- zig/tests/cuda/nemotron/test_aot_manifest.py | 21 +++++++++++++++----- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/tools/zig/triton_aot_pack.py b/tools/zig/triton_aot_pack.py index 61d617a1bb..c9510e9f4c 100644 --- a/tools/zig/triton_aot_pack.py +++ b/tools/zig/triton_aot_pack.py @@ -60,7 +60,8 @@ def entry(h: str, k: dict, nospec: set[str], cache: Path | None, out: Path) -> d "global_scratch": md.get("global_scratch_size", 0), "global_align": md.get("global_scratch_align", 1), "profile_scratch": md.get("profile_scratch_size", 0), "pdl": bool(md.get("launch_pdl", False)), "params": [{"name": n, "type": sig[n], "div16": n in div16, "nospec": n in nospec} for n in runtime], - "consts": {n: const(v) for n, v in k["constexprs"].items() if n not in runtime}, + # a None argument Triton folded into the kernel (an unused pointer) leaves nothing to match + "consts": {n: const(v) for n, v in k["constexprs"].items() if n not in runtime and v is not None}, } @@ -69,14 +70,15 @@ def main() -> int: ap.add_argument("--manifest", action="append", required=True, help="a manifest.json or a Recorder's launches.json (repeatable)") ap.add_argument("--cache", action="append", help="each capture's Triton cache as seen here, in --manifest order") - ap.add_argument("--jit", help="the capture's jit.json; without it no parameter is marked do_not_specialize") + ap.add_argument("--jit", action="append", + help="a capture's jit.json (repeatable, merged); without one no parameter is do_not_specialize") ap.add_argument("--out", required=True) a = ap.parse_args() docs = [json.loads(Path(m).read_text()) for m in a.manifest] caches = a.cache or [d.get("cache") for d in docs] if len(caches) != len(docs): raise SystemExit("--cache is given once per --manifest, or not at all") - jit = json.loads(Path(a.jit).read_text()) if a.jit else None + jit = {f: v for j in a.jit for f, v in json.loads(Path(j).read_text()).items()} if a.jit else None out = Path(a.out) (out / "cubins").mkdir(parents=True, exist_ok=True) kernels, seen = [], set() diff --git a/zig/tests/cuda/nemotron/test_aot_manifest.py b/zig/tests/cuda/nemotron/test_aot_manifest.py index e87ff8551c..7a82c87028 100755 --- a/zig/tests/cuda/nemotron/test_aot_manifest.py +++ b/zig/tests/cuda/nemotron/test_aot_manifest.py @@ -37,10 +37,10 @@ def fake_cache(cache: Path) -> dict: for leaf in ("__grp__fake_kernel.json", "fake_kernel.cubin", "fake_kernel.ptx")} info = {"name": "fake_kernel", "function": "fakemod.FakeKernel", "source": {"file": "tensorfold/families/nemotron_h/cuda/kernels.py", "line": 10}, - "params": ["p0", "n", "BLOCK", "SCALE"], - "signature": {"p0": "*i8", "n": "i32", "BLOCK": "constexpr", "SCALE": "constexpr"}, + "params": ["p0", "n", "BLOCK", "SCALE", "ROPE"], + "signature": {"p0": "*i8", "n": "i32", "BLOCK": "constexpr", "SCALE": "constexpr", "ROPE": "constexpr"}, "constexprs": {"BLOCK": 64, "SCALE": {"float": 0.125, "fp32_bits": "0x3e000000", - "fp64_bits": "0x3fc0000000000000"}, "n": 3}, + "fp64_bits": "0x3fc0000000000000"}, "n": 3, "ROPE": None}, "attrs": {"p0": [["tt.divisibility", 16]], "n": [["tt.contiguity", 1]]}, "metadata": {"num_warps": 4}, "files": files, "n_regs": 24, "n_spills": 0} return {"generator": "test", "written_unix": 0, "kernels": {HASH: info}, @@ -73,7 +73,7 @@ def main() -> int: assert k["launches"] == {"startup": 3} and k["grids"] == [[128, 1, 1]] jit = out / "jit.json" - jit.write_text(json.dumps({"fakemod.FakeKernel": {"params": ["p0", "n", "BLOCK", "SCALE"], + jit.write_text(json.dumps({"fakemod.FakeKernel": {"params": ["p0", "n", "BLOCK", "SCALE", "ROPE"], "do_not_specialize": [], "no_align": []}}) + "\n") # the manifest shape, with a cache and a jit file named on the command line packed = work / "packed" @@ -112,7 +112,18 @@ def main() -> int: print(r.stdout + r.stderr) return 1 assert (default / "aot.json").read_text() == (packed / "aot.json").read_text() - print("test_aot_manifest: one kernel through the manifest, the dump and the defaults; entries checked") + + # a second capture's jit.json beside the first: merged, the same set again + jit2 = out / "jit2.json" + jit2.write_text(json.dumps({"othermod.Other": {"params": ["x"], "do_not_specialize": ["x"], "no_align": []}})) + merged = work / "merged" + r = subprocess.run([sys.executable, "-B", str(PACK), "--manifest", str(manifest), "--cache", str(cache), + "--jit", str(jit), "--jit", str(jit2), "--out", str(merged)], capture_output=True, text=True) + if r.returncode != 0: + print(r.stdout + r.stderr) + return 1 + assert (merged / "aot.json").read_text() == (packed / "aot.json").read_text() + print("test_aot_manifest: one kernel through the manifest, the dump, the defaults and two jit files; entries checked") return 0 From 5cf42918ebd0411c0d006acace97bca191037115 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:15:58 -0500 Subject: [PATCH 033/115] Flash Next CUDA: prompts and decode rounds on one batched pass, with the sparse indexer and the 262,144-token window cuda_batch.zig runs a chunk of rows through each layer at once, on kernels that keep a row's bits whatever the row count: the group-32 lane tile (qmm.cu), the core's grouped expert plan (cuda.grouped.Router) and rows kernel (cuda.experts.RowsExperts, experts_rows.cu), the DeltaNet chain over the chunk (flashnext_batch.cu: fn_gdn_front, fn_gdn_rows without block barriers, fn_gdn_norm, each value by the one-row chain's expression), attention over each row's causal prefix, and the captured Triton forms at more than one row (capture_rows.sh). A chunk therefore leaves the logits, KV cache, DeltaNet conv windows and states and the n-gram tail exactly as feeding its rows one at a time. The n-gram table stays mapped for the engine's life and a chunk's rows are gathered on host threads while the GPU runs layer 0. Engine.round runs several sequences' next tokens as the rows of one pass: the stateless kernels over every row, each row's n-gram conv, DeltaNet step, rope, cache writes and attention on its own sequence at its own position. Past 2,048 keys each attention layer reads 512 blocks of 4 keys a row, as the model does: the pooled indexer keys (_pool), the block scores (_scores), a native selection (fn_qsa_select: the lists Triton's _select and _select_tiles make, checked by flashnext-select on capture_sparse.sh's fixtures with many ties) and the sparse _chunks/_merge (QSA 1, NCH 5). The family serves the model's 262,144-token window. The family's own pieces only: config.json's quantization and shapes are checked at load against the ones the kernels and buffers are built for (cuda_config.zig: a 6-bit checkpoint or another width is refused by name); the shards are read with core.safetensors.File; a folder whose shared experts carry Living Weights neurons is refused in words. The row-by-row path Adrian's prompt check was built on moves out of the engine into cuda_rows.zig, where flashnext-prompt (TF_FLASHNEXT_PROMPT=rows) still uses it as the batched pass's reference. Captures: capture_rows.sh and capture_sparse.sh record the extra Triton forms, and pack.sh packs any set of captures with tools/zig/triton_aot_pack.py. flashnext-prompt (32 greedy tokens) matches the Python engine's teacher path on short, mid and long, and on long3k (3,001 tokens) and long12k (12,001 tokens) past the indexer's budget (results-sparse.json). --- zig/build/cuda.zig | 7 + zig/kernels/cuda/flashnext_batch.cu | 316 +++++++++++ zig/kernels/cuda/flashnext_gdn.cu | 3 +- zig/kernels/cuda/qmm.cu | 2 +- zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_attn.zig | 6 +- zig/src/families/flashnext/cuda_batch.zig | 521 ++++++++++++++++++ .../families/flashnext/cuda_batch_state.zig | 345 ++++++++++++ zig/src/families/flashnext/cuda_config.zig | 85 +++ .../families/flashnext/cuda_embed_main.zig | 2 +- zig/src/families/flashnext/cuda_launch.zig | 102 +++- zig/src/families/flashnext/cuda_moe.zig | 2 +- zig/src/families/flashnext/cuda_ple.zig | 413 ++++++-------- zig/src/families/flashnext/cuda_ple_check.zig | 235 ++++++++ zig/src/families/flashnext/cuda_prompt.zig | 348 +++++------- .../families/flashnext/cuda_prompt_load.zig | 25 +- .../families/flashnext/cuda_prompt_main.zig | 4 +- .../families/flashnext/cuda_prompt_step.zig | 12 +- zig/src/families/flashnext/cuda_prompt_wt.zig | 67 +-- zig/src/families/flashnext/cuda_rows.zig | 222 ++++++++ .../families/flashnext/cuda_select_main.zig | 113 ++++ zig/tests/cuda/flashnext/capture.py | 28 +- zig/tests/cuda/flashnext/capture.sh | 27 +- zig/tests/cuda/flashnext/capture_rows.py | 86 +++ zig/tests/cuda/flashnext/capture_rows.sh | 11 + zig/tests/cuda/flashnext/capture_sparse.py | 108 ++++ zig/tests/cuda/flashnext/capture_sparse.sh | 11 + zig/tests/cuda/flashnext/lr7/long12k.ids | 1 + zig/tests/cuda/flashnext/lr7/long3k.ids | 1 + .../cuda/flashnext/lr7/results-sparse.json | 364 ++++++++++++ zig/tests/cuda/flashnext/pack.sh | 8 + 31 files changed, 2896 insertions(+), 580 deletions(-) create mode 100644 zig/kernels/cuda/flashnext_batch.cu create mode 100644 zig/src/families/flashnext/cuda_batch.zig create mode 100644 zig/src/families/flashnext/cuda_batch_state.zig create mode 100644 zig/src/families/flashnext/cuda_config.zig create mode 100644 zig/src/families/flashnext/cuda_ple_check.zig create mode 100644 zig/src/families/flashnext/cuda_rows.zig create mode 100644 zig/src/families/flashnext/cuda_select_main.zig create mode 100644 zig/tests/cuda/flashnext/capture_rows.py create mode 100755 zig/tests/cuda/flashnext/capture_rows.sh create mode 100644 zig/tests/cuda/flashnext/capture_sparse.py create mode 100755 zig/tests/cuda/flashnext/capture_sparse.sh create mode 100644 zig/tests/cuda/flashnext/lr7/long12k.ids create mode 100644 zig/tests/cuda/flashnext/lr7/long3k.ids create mode 100644 zig/tests/cuda/flashnext/lr7/results-sparse.json create mode 100755 zig/tests/cuda/flashnext/pack.sh diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 7bb23f2285..7d759a2715 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -16,6 +16,7 @@ const glue = &[_][]const u8{ "-O3", "--fmad=false", "--ftz=false", "--prec-div=t const kernels = [_]Kernel{ .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2 .{ .name = "flashnext_gdn", .flags = &.{ "-O3", "--fmad=false" } }, // qwen4_exp/cuda/gdn.cu chain, one row + .{ .name = "flashnext_batch", .flags = &.{ "-O3", "--fmad=false" } }, // ours: Flash Next's batched prompt pass glue .{ .name = "probe", .flags = &.{"-O3"} }, .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 .{ .name = "qmm", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py lane matmul, groups of 32 @@ -166,6 +167,12 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const prompt_count = b.option([]const u8, "flashnext_count", "greedy tokens to emit") orelse "8"; if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); + // not installed: the native block selection against Triton's lists (capture_sparse.sh's fixtures) + const select_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_select_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + select_check.addImport("cuda", cuda); + const select_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-select", .root_module = select_check })); + if (b.option([]const u8, "flashnext_fixtures", "capture_sparse.sh's selection fixtures folder")) |dir| select_run.addArgs(&.{dir}); + b.step("flashnext-select", "fn_qsa_select against Triton's _select / _select_tiles lists on captured score rows").dependOn(&select_run.step); nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, mods.tokenizer, build_options, true).root_module.strip = strip; } diff --git a/zig/kernels/cuda/flashnext_batch.cu b/zig/kernels/cuda/flashnext_batch.cu new file mode 100644 index 0000000000..c4ad0e2691 --- /dev/null +++ b/zig/kernels/cuda/flashnext_batch.cu @@ -0,0 +1,316 @@ +// The small device steps of Flash Next's batched prompt pass (cuda_batch.zig), built without FMA contraction. + +#include +#include + +// out[i] = the sum of the sk slices' part[s][i] in order, rounded once (qmm.cu's reduce_kernel): bf16 if out_bf16. +extern "C" __global__ void fn_sum_slices(const float* __restrict__ part, void* __restrict__ out, int out_bf16, + long long n, int sk) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + float acc = part[i]; + for (int s = 1; s < sk; ++s) acc = acc + part[static_cast(s) * n + i]; + if (out_bf16) reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); + else reinterpret_cast(out)[i] = acc; +} + +// Shifts the last-`taps`-rows window [taps, c] in place past `rows`, n rows `stride` apart (conv history, n-gram tail). +extern "C" __global__ void fn_shift_window(__nv_bfloat16* __restrict__ win, const __nv_bfloat16* __restrict__ rows, + long long stride, int n, int taps, int c) { + const int ch = blockIdx.x * blockDim.x + threadIdx.x; + if (ch >= c) return; + __nv_bfloat16 next[16]; + for (int j = 0; j < taps; ++j) { + const int src = j + n; // row src of [window (taps rows); new rows] + next[j] = src < taps ? win[static_cast(src) * c + ch] + : rows[static_cast(src - taps) * stride + ch]; + } + for (int j = 0; j < taps; ++j) win[static_cast(j) * c + ch] = next[j]; +} + +// ---- the DeltaNet chain over a chunk: flashnext_gdn.cu's chain cut into front, rows and norm, with its bits ---- + +namespace fn_gdn { +constexpr int NK = 16, NV = 48, DK = 128, DV = 128, TAPS = 4; +constexpr int C = 2 * NK * DK + NV * DV; +constexpr int PW = C + NV * DV + 2 * NV; + +static __device__ __forceinline__ float bf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } + +static __device__ __forceinline__ float warp_sum(float x) { + for (int o = 16; o; o >>= 1) x += __shfl_xor_sync(0xffffffffu, x, o); + return x; +} + +static __device__ __forceinline__ float sigmoidf_(float x) { return 1.0f / (1.0f + expf(-x)); } + +static __device__ __forceinline__ float softplusf_(float x) { return x > 20.0f ? x : log1pf(expf(x)); } + +static __device__ __forceinline__ float conv_act(const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, + const __nv_bfloat16* __restrict__ cw, int r, int c) { + float acc = 0.0f; +#pragma unroll + for (int tap = 0; tap < TAPS; ++tap) { + const int at = r + tap; + const float x = at < TAPS - 1 ? __bfloat162float(cs[at * C + c]) + : __bfloat162float(P[(size_t)(at - (TAPS - 1)) * PW + c]); + acc = acc + __bfloat162float(cw[c * TAPS + tap]) * x; + } + return bf(acc / (1.0f + expf(-acc))); +} +} // namespace fn_gdn + +// Program (row, unit), 256 threads: key head u's normed q and k (u < NK), else a value head's v, decay g and beta. +extern "C" __global__ void __launch_bounds__(256) fn_gdn_front( + const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ cs, const __nv_bfloat16* __restrict__ cw, + const float* __restrict__ a_log, const float* __restrict__ dt_bias, float* __restrict__ qn, + float* __restrict__ kn, float* __restrict__ v, float* __restrict__ gb) { + using namespace fn_gdn; + const int r = blockIdx.x, u = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31; + if (u < NK) { + __shared__ float x[2][DK]; + const int hk = u; + x[t / DK][t % DK] = conv_act(P, cs, cw, r, t < DK ? hk * DK + t : NK * DK + hk * DK + (t - DK)); + __syncthreads(); + if (warp < 2) { + float v4[4], ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + v4[i] = x[warp][lane * 4 + i]; + ss = ss + v4[i] * v4[i]; + } + ss = warp_sum(ss); + float inv = 1.0f / sqrtf(ss + 1e-6f); + if (warp == 0) inv = inv * (1.0f / sqrtf((float)DK)); + float* dst = (warp == 0 ? qn : kn) + ((size_t)r * NK + hk) * DK; +#pragma unroll + for (int i = 0; i < 4; ++i) dst[lane * 4 + i] = v4[i] * inv; + } + } else { + const int hv = u - NK; + if (t < DV) v[(size_t)r * NV * DV + hv * DV + t] = conv_act(P, cs, cw, r, 2 * NK * DK + hv * DV + t); + else if (t == DV) { + const float b = __bfloat162float(P[(size_t)r * PW + C + NV * DV + hv]); + const float a = __bfloat162float(P[(size_t)r * PW + C + NV * DV + NV + hv]); + gb[((size_t)r * NV + hv) * 2] = expf(-expf(a_log[hv]) * softplusf_(a + dt_bias[hv])); + gb[((size_t)r * NV + hv) * 2 + 1] = bf(sigmoidf_(b)); + } + } +} + +// The chain through the chunk from state_in to state_out (in place, or apart for a rollback): a warp 4 state rows. +static __device__ __forceinline__ void gdn_rows( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { + using namespace fn_gdn; + const int hv = blockIdx.x, hk = hv / (NV / NK), warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int row0 = blockIdx.y * 16 + warp * 4; + const size_t sbase = (size_t)hv * DV * DK; + float s[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = state_in[sbase + (size_t)(row0 + j) * DK + lane * 4 + i]; + float kk[4], qq[4], vv[4], g = 0.0f, beta = 0.0f; + auto load = [&](int r, float (&k_)[4], float (&q_)[4], float (&v_)[4], float& g_, float& b_) { + const float4 k4 = *reinterpret_cast(kn + ((size_t)r * NK + hk) * DK + lane * 4); + const float4 q4 = *reinterpret_cast(qn + ((size_t)r * NK + hk) * DK + lane * 4); + const float4 v4 = *reinterpret_cast(v + (size_t)r * NV * DV + hv * DV + row0); + k_[0] = k4.x; k_[1] = k4.y; k_[2] = k4.z; k_[3] = k4.w; + q_[0] = q4.x; q_[1] = q4.y; q_[2] = q4.z; q_[3] = q4.w; + v_[0] = v4.x; v_[1] = v4.y; v_[2] = v4.z; v_[3] = v4.w; + g_ = gb[((size_t)r * NV + hv) * 2]; + b_ = gb[((size_t)r * NV + hv) * 2 + 1]; + }; + if (rows > 0) load(0, kk, qq, vv, g, beta); + for (int r = 0; r < rows; ++r) { + float nk[4], nq[4], nv[4], ng = 0.0f, nb = 0.0f; + if (r + 1 < rows) load(r + 1, nk, nq, nv, ng, nb); +#pragma unroll + for (int j = 0; j < 4; ++j) { + float kv = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + s[j][i] = s[j][i] * g; + kv = kv + s[j][i] * kk[i]; + } + kv = warp_sum(kv); + const float delta = (vv[j] - kv) * beta; +#pragma unroll + for (int i = 0; i < 4; ++i) s[j][i] = s[j][i] + kk[i] * delta; + } + float o4[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float o = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) o = o + s[j][i] * qq[i]; + o4[j] = bf(warp_sum(o)); + } + if (lane == 0) *reinterpret_cast(ys_out + ((size_t)r * NV + hv) * DV + row0) = make_float4(o4[0], o4[1], o4[2], o4[3]); + if (r + 1 < rows) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + kk[i] = nk[i]; + qq[i] = nq[i]; + vv[i] = nv[i]; + } + g = ng; + beta = nb; + } + } +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) state_out[sbase + (size_t)(row0 + j) * DK + lane * 4 + i] = s[j][i]; +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, float* state, int rows, float* __restrict__ ys_out) { + gdn_rows(qn, kn, v, gb, state, state, rows, ys_out); +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_rows_to( + const float* __restrict__ qn, const float* __restrict__ kn, const float* __restrict__ v, + const float* __restrict__ gb, const float* state_in, float* state_out, int rows, float* __restrict__ ys_out) { + gdn_rows(qn, kn, v, gb, state_in, state_out, rows, ys_out); +} + +extern "C" __global__ void __launch_bounds__(128) fn_gdn_norm( + const float* __restrict__ ys_in, const __nv_bfloat16* __restrict__ P, const __nv_bfloat16* __restrict__ norm_w, + float eps, __nv_bfloat16* __restrict__ out, float* __restrict__ xs) { + using namespace fn_gdn; + const int r = blockIdx.x, hv = blockIdx.y, t = threadIdx.x, warp = t >> 5, lane = t & 31; + __shared__ float ys[DV]; + __shared__ float rinv; + ys[t] = ys_in[((size_t)r * NV + hv) * DV + t]; + __syncthreads(); + if (warp == 0) { + float ss = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float y = ys[lane * 4 + i]; + ss = ss + y * y; + } + ss = warp_sum(ss); + if (lane == 0) rinv = 1.0f / sqrtf(ss / (float)DV + eps); + } + __syncthreads(); + const float yn = bf(bf(ys[t] * rinv) * __bfloat162float(norm_w[t])); + const float z = __bfloat162float(P[(size_t)r * PW + C + hv * DV + t]); + const float o = bf(yn * sigmoidf_(z)); + out[(size_t)r * NV * DV + hv * DV + t] = __float2bfloat16_rn(o); + const float gs = warp_sum(o); + if (lane == 0) xs[(size_t)r * (NV * DV / 32) + hv * (DV / 32) + warp] = gs; +} + +// ---- sparse attention (QSA): each row's blocks ---------------------------------------------------------------------- + +// attention._select's order-preserving key of a block score: negative floats flipped, positive ones topped. +__device__ __forceinline__ uint32_t fn_qsa_key(float v) { + const uint32_t bits = __float_as_uint(v); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +// An inclusive block-wide sum of one int a thread (blockDim.x a multiple of 32, at most 1024); *total gets the sum. +__device__ __forceinline__ int fn_block_scan(int v, int* warp_sums, int* total) { + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int warps = blockDim.x >> 5; + for (int o = 1; o < 32; o <<= 1) { + const int n = __shfl_up_sync(0xffffffffu, v, o); + if (lane >= o) v += n; + } + if (lane == 31) warp_sums[warp] = v; + __syncthreads(); + if (warp == 0) { + int w = lane < warps ? warp_sums[lane] : 0; + for (int o = 1; o < 32; o <<= 1) { + const int n = __shfl_up_sync(0xffffffffu, w, o); + if (lane >= o) w += n; + } + if (lane < warps) warp_sums[lane] = w; + } + __syncthreads(); + if (warp > 0) v += warp_sums[warp - 1]; + *total = warp_sums[warps - 1]; + __syncthreads(); + return v; +} + +// attention._select and _select_tiles, a row a block: a row past `top` complete blocks lists the best, then its tail. +extern "C" __global__ void __launch_bounds__(1024) fn_qsa_select(const float* __restrict__ sc, + const int* __restrict__ pos0, int* __restrict__ ids, + int* __restrict__ nkr, int* __restrict__ spr, int nb, + int idw, int top, int ratio) { + __shared__ int hist[256]; + __shared__ int warp_sums[32]; + __shared__ uint32_t cut_s; + __shared__ int need_s; + const int r = blockIdx.x; + const int end = *pos0 + r + 1; + const int complete = end / ratio; + if (complete <= top) { + if (threadIdx.x == 0) { + nkr[r] = end; + spr[r] = 0; + } + return; + } + const float* row = sc + static_cast(r) * nb; + int* out = ids + static_cast(r) * idw; + uint32_t cut = 0; + int need = top; + for (int p = 0; p < 4; ++p) { + const int shift = 24 - 8 * p; + for (int i = threadIdx.x; i < 256; i += blockDim.x) hist[i] = 0; + __syncthreads(); + for (int b = threadIdx.x; b < complete; b += blockDim.x) { + const uint32_t k = fn_qsa_key(row[b]); + if (p == 0 || (k >> (shift + 8)) == (cut >> (shift + 8))) atomicAdd(&hist[(k >> shift) & 0xFF], 1); + } + __syncthreads(); + if (threadIdx.x == 0) { + int above = 0; // keys under the prefix with a larger byte + int d = 255; + for (; d > 0; --d) { + if (above + hist[d] >= need) break; + above += hist[d]; + } + need_s = need - above; + cut_s = cut | (static_cast(d) << shift); + } + __syncthreads(); + cut = cut_s; + need = need_s; + __syncthreads(); + } + // every key above the cut, then the first `need` keys equal to it, placed in block order + int seen = 0; + int placed = 0; + for (int t0 = 0; t0 < complete && placed < top; t0 += blockDim.x) { + const int b = t0 + threadIdx.x; + const bool ok = b < complete; + const uint32_t k = ok ? fn_qsa_key(row[b]) : 0u; + const bool above = ok && k > cut; + const bool equal = ok && k == cut; + int eq_total = 0; + const int rank = seen + fn_block_scan(equal ? 1 : 0, warp_sums, &eq_total); + const bool chosen = above || (equal && rank <= need); + int ch_total = 0; + const int place = placed + fn_block_scan(chosen ? 1 : 0, warp_sums, &ch_total) - 1; + if (chosen) + for (int j = 0; j < ratio; ++j) out[place * ratio + j] = b * ratio + j; + seen += eq_total; + placed += ch_total; + } + if (threadIdx.x < ratio) { + const int tail = ratio * complete + threadIdx.x; + if (tail < end) out[top * ratio + threadIdx.x] = tail; + } + if (threadIdx.x == 0) { + nkr[r] = top * ratio + end - ratio * complete; + spr[r] = 1; + } +} diff --git a/zig/kernels/cuda/flashnext_gdn.cu b/zig/kernels/cuda/flashnext_gdn.cu index 5ce476b7df..681528a89a 100644 --- a/zig/kernels/cuda/flashnext_gdn.cu +++ b/zig/kernels/cuda/flashnext_gdn.cu @@ -1,5 +1,4 @@ -// One value head per block: depthwise conv, the delta-rule step, then the gated RMSNorm. 16 key heads, 48 value heads. -// Same operations as the one-row chain, compiled without FMA contraction. +// One value head per block: depthwise conv, the delta-rule step, then the gated RMSNorm, without FMA contraction. #include #include diff --git a/zig/kernels/cuda/qmm.cu b/zig/kernels/cuda/qmm.cu index 773d1d5926..4901978a3a 100644 --- a/zig/kernels/cuda/qmm.cu +++ b/zig/kernels/cuda/qmm.cu @@ -1,4 +1,4 @@ -// Device code of src/tensorfold/cuda/kernels/qmm.cu (lines 1-230, comments and ATen includes dropped), checked by zig/tests/cuda/copies.py. +// Device code of src/tensorfold/cuda/kernels/qmm.cu (lines 1-230, comments and ATen includes dropped). #include #include diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 78260422ba..791d394ce0 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -14,6 +14,7 @@ pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; +pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes else &.{}; pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_attn.zig b/zig/src/families/flashnext/cuda_attn.zig index 47d823fc1a..e93e1a2e30 100644 --- a/zig/src/families/flashnext/cuda_attn.zig +++ b/zig/src/families/flashnext/cuda_attn.zig @@ -258,7 +258,7 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, defer ik_b.free(); var dummy = try cuda.DeviceBuffer.alloc(driver, 256); defer dummy.free(); - try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, dummy.ptr, dummy.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0); + try tri.attnPrep(pa_b.ptr, pos_b.ptr, qw_b.ptr, kw_b.ptr, iw_b.ptr, inv_b.ptr, q_b.ptr, kc_b.ptr, vc_b.ptr, iq_b.ptr, ik_b.ptr, dummy.ptr, dummy.ptr, eps, 1, 0); var po_b = try cuda.DeviceBuffer.alloc(driver, nch * q_heads * head_dim * 4); defer po_b.free(); @@ -268,8 +268,8 @@ pub fn decode(comptime Tri: type, gpa: std.mem.Allocator, driver: *cuda.Driver, defer pl_b.free(); var o_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); defer o_b.free(); - try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, dummy.ptr, dummy.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1); - try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, kv_heads); + try tri.attnChunks(q_b.ptr, kc_b.ptr, vc_b.ptr, pos_b.ptr, po_b.ptr, pm_b.ptr, pl_b.ptr, dummy.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, 1, false); + try tri.attnMerge(po_b.ptr, pm_b.ptr, pl_b.ptr, pos_b.ptr, o_b.ptr, dummy.ptr, dummy.ptr, 1, kv_heads, false); var gated_b = try cuda.DeviceBuffer.alloc(driver, out_k * 2); defer gated_b.free(); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig new file mode 100644 index 0000000000..1b3c0aa7a5 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -0,0 +1,521 @@ +//! The batched pass: a chunk of rows through each layer at once, each row's bits those of a one-row pass. + +const std = @import("std"); +const cuda = @import("cuda"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_prompt_load.zig"); +const parts = @import("cuda_batch_state.zig"); + +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const vocab = wt.vocab; +const experts_n = wt.experts_n; +const slots = wt.slots; +const moe_w = wt.moe_w; +const tile = wt.tile; +const q_heads = wt.q_heads; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_heads = wt.index_heads; +const index_dim = wt.index_dim; +const half = wt.half; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const nch = wt.nch; +const ple_layer = wt.ple_layer; +const ple_tail_n = wt.ple_tail_n; +const state_n = wt.state_n; +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const Layer = load.Layer; + +/// Normed-stream partial sums a row: D / 256 blocks of each stream. +pub const pss_n = (dims / 256) * streams; +/// The widest projection row: the DeltaNet's q|k|v|z|b|a (attention's q|k|v|indexer is 13952). +pub const proj_row = @max(gdn.proj_width, proj_n); +/// Attention rows a _chunks launch takes (its scratch: rows x nch x heads x dim fp32). +pub const att_rows = parts.att_rows; +/// Rows a pass holds at most unless asked for more. +pub const default_chunk = parts.default_chunk; +const ratio = wt.index_ratio; +const top = wt.index_top; +const idw = wt.index_idw; +const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +const qmm_smem: u32 = 9472; +/// Pairs an expert item takes: the core rows kernel's tile. +const rows_tile: usize = cuda.experts.RowsExperts.tile; +const plan_small: usize = cuda.grouped.small; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} +pub const Logits = parts.Logits; +pub const LayerKeep = parts.LayerKeep; +pub const Keep = parts.Keep; +pub const Hooks = parts.Hooks; +pub const Kernels = parts.Kernels; +pub const Scratch = parts.Scratch; +pub const SeqState = parts.SeqState; +pub const scoreStride = parts.scoreStride; +pub const scoreBytes = parts.scoreBytes; +pub const Lane = parts.Lane; +pub const Model = parts.Model; + +/// TF_FLASHNEXT_PROFILE=1: each chunk's wall time (synchronized) and its host n-gram gather on stderr. +fn profile() bool { + const v = std.c.getenv("TF_FLASHNEXT_PROFILE") orelse return false; + return std.mem.eql(u8, std.mem.span(v), "1"); +} + +pub fn Pass(comptime Tri: type) type { + return struct { + const Self = @This(); + + gpa: std.mem.Allocator, + io: std.Io, + stream: cuda.Stream, + tri: Tri, + k: *const Kernels, + s: *const Scratch, + m: Model, + seq: SeqState, + hooks: Hooks, + /// A decode round: row r is lanes[r]'s token at its own position and state; empty, the rows continue `seq`. + lanes: []const Lane = &.{}, + rows: usize = 0, + pos0: usize = 0, + /// The MoE write-back the next layer (or the head) applies first, and the layer it came from. + pending: ?usize = null, + + // ---- launches ---------------------------------------------------------------------------------------- + + /// qmm.cu's lane tile over m rows: out bf16 [m, n] (sk 1) or the unreduced fp32 slices [sk, m, n] in part. + pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !void { + const bm: usize = 16; + const bn: usize = 64; + const k = w.k; + const rows_t = cdiv(m, bm); + const cols_t = cdiv(w.n, bn); + const sweep = @max(@as(usize, 1), @min(rows_t, (12 << 20) / (bm * k * 2))); + const npad = (w.n + 127) / 128 * 128; + var a: cuda.Args = .{}; + a.add(x); + a.add(xs); + a.add(w.w.ptr); + a.add(w.s.ptr); + a.add(w.b.ptr); + a.add(out); + a.add(part); + for ([_]usize{ m, w.n, k, sk, npad, k, sweep }) |v| a.add(cint(v)); + try cuda.launch.launch(self.k.lane, .{ .grid = .{ .x = @intCast(rows_t * cols_t), .z = @intCast(sk) }, .block = .{ .x = 128 }, .shared = qmm_smem }, self.stream, &a); + } + + pub fn sumSlices(self: *Self, part: u64, out: u64, bf16: bool, n: usize, sk: usize) !void { + var a: cuda.Args = .{}; + a.add(part); + a.add(out); + a.add(@as(c_int, @intFromBool(bf16))); + a.add(@as(c_longlong, @intCast(n))); + a.add(cint(sk)); + try cuda.launch.launch(self.k.sum, .{ .grid = .{ .x = @intCast(cdiv(n, 256)) }, .block = .{ .x = 256 } }, self.stream, &a); + } + + pub fn shiftWindow(self: *Self, win: u64, rows_ptr: u64, stride: usize, n: usize, taps: usize, c: usize) !void { + var a: cuda.Args = .{}; + a.add(win); + a.add(rows_ptr); + a.add(@as(c_longlong, @intCast(stride))); + a.add(cint(n)); + a.add(cint(taps)); + a.add(cint(c)); + try cuda.launch.launch(self.k.shift, .{ .grid = .{ .x = @intCast(cdiv(c, 256)) }, .block = .{ .x = 256 } }, self.stream, &a); + } + + pub fn copy(self: *Self, dst: u64, src: u64, n: usize) !void { + if (dst == 0 or n == 0) return; + const d = self.stream.d; + try d.check(d.api.cuMemcpyDtoDAsync_v2(dst, src, n, self.stream.handle), "cuMemcpyDtoDAsync"); + } + + pub fn upload(self: *Self, dst: u64, bytes: []const u8) !void { + const d = self.stream.d; + try d.check(d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, self.stream.handle), "cuMemcpyHtoDAsync"); + } + + // ---- the blocks ---------------------------------------------------------------------------------------- + + /// The pending MoE output (or none: mode 0, the streams' sums alone) written back into rows [r0, r0 + n). + pub fn writeback(self: *Self, r0: usize, n: usize) !void { + const s = self.s; + const dummy = s.dummy; + const h = s.h + r0 * wide * 2; + const pss = s.pss + r0 * pss_n * 4; + const inj = s.inj + r0 * streams * 2; + const rs = self.rows * dims; + if (self.pending == null) { + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", dummy, dummy, "*bf16", dummy, "*bf16", rs, n, dims, streams, 0, 1, 1, 1); + return; + } + self.pending = null; + const y = s.y + r0 * slots * dims * 4; + const wts = s.wts + r0 * slots * 4; + try self.tri.hcWriteback(h, h, pss, dummy, "*bf16", inj, y, "*fp32", wts, "*fp32", rs, n, dims, streams, 2, 10, 11, 1); + } + + /// A read-out over rows [r0, r0 + n) of the streams: mixed and its group sums, and the inject gates. + pub fn readout(self: *Self, hc_w: HcW, r0: usize, n: usize) !void { + const s = self.s; + const h = s.h + r0 * wide * 2; + const pss = s.pss + r0 * pss_n * 4; + const inj = s.inj + r0 * streams * 2; + const norm = s.norm + r0 * wide * 2; + const act = s.act + r0 * low * 2; + const xs_act = s.xs_act + r0 * (low / 32) * 4; + try self.tri.hcDown(h, pss, hc_w.scale.ptr, norm, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn, s.hc_part, eps, n, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); + const inj_ptr = if (hc_w.has_inj == 1) inj else act; + try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, hc_sk, n, streams, low, hc_w.n_down, hc_w.has_inj); + try self.tri.hcUpmix(act, xs_act, hc_w.up.w.ptr, hc_w.up.s.ptr, hc_w.up.b.ptr, norm, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, wide, low, dims, streams); + } + + fn pleBlock(self: *Self, lw: *const Layer, ids: []const u32) !void { + const s = self.s; + const n = self.rows; + // the rows' n-gram embeddings (host gathers from the mapped table) are already uploaded by `stage` + _ = ids; + try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); + try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); + try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, 4); + try self.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, n, dims, streams); + if (self.lanes.len == 0) { + try self.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, self.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, n, dims, streams, 4, 3); + try self.shiftWindow(self.seq.tail.ptr, s.ple_nrow, wide, n, ple_tail_n, wide); + return; + } + // a round: each row's convolution over its own stream's tail, the one-row launch a lone decode makes + for (self.lanes, 0..) |l, r| { + const g = r * wide * 2; + try self.tri.pleConv(s.ple_gated + g, s.ple_pss + r * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, 1, dims, streams, 4, 3); + try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, 1, ple_tail_n, wide); + } + } + + fn gdnBlock(self: *Self, lw: *const Layer, li: usize, keep: LayerKeep) !void { + const s = self.s; + const n = self.rows; + try self.lane(s.mixed, s.xsm, lw.proj, s.proj, s.dummy, n, 1); + try self.copy(keep.proj, s.proj, n * gdn.proj_width * 2); + try self.copy(keep.hist, self.seq.conv[li].ptr, (gdn.taps - 1) * gdn.conv_dim * 2); + // every row's conv, norms and gates, the recurrence, then the gated norm: the one-block chain's bits + if (self.lanes.len == 0) { + try self.gdnFront(lw, s.proj, self.seq.conv[li].ptr, n, s.gq, s.gk, s.gv, s.ggb); + try self.gdnChain(s.gq, s.gk, s.gv, s.ggb, self.seq.state[li].ptr, keep.state_out, n, s.ys); + } else for (self.lanes, 0..) |l, r| { + // a round: each row's front and recurrence on its own stream's conv window and state + const qk = r * gdn.nk * gdn.dk * 4; + const v = r * gdn.value_dim * 4; + const gb = r * gdn.nv * 2 * 4; + try self.gdnFront(lw, s.proj + r * gdn.proj_width * 2, l.seq.conv[li].ptr, 1, s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb); + try self.gdnChain(s.gq + qk, s.gk + qk, s.gv + v, s.ggb + gb, l.seq.state[li].ptr, 0, 1, s.ys + v); + } + var b: cuda.Args = .{}; + b.add(s.ys); + b.add(s.proj); + b.add(lw.norm.ptr); + b.add(eps); + b.add(s.gout); + b.add(s.gxs); + try cuda.launch.launch(self.k.gnorm, .{ .grid = .{ .x = @intCast(n), .y = gdn.nv }, .block = .{ .x = gdn.dv } }, self.stream, &b); + if (self.lanes.len == 0) { + try self.shiftWindow(self.seq.conv[li].ptr, s.proj, gdn.proj_width, n, gdn.taps - 1, gdn.conv_dim); + } else for (self.lanes, 0..) |l, r| { + try self.shiftWindow(l.seq.conv[li].ptr, s.proj + r * gdn.proj_width * 2, gdn.proj_width, 1, gdn.taps - 1, gdn.conv_dim); + } + try self.lane(s.gout, s.gxs, lw.out, s.dummy, s.part, n, out_sk); + } + + /// fn_gdn_front over `rows` projection rows after the conv window `conv`. + fn gdnFront(self: *Self, lw: *const Layer, proj: u64, conv: u64, rows: usize, gq: u64, gk: u64, gv: u64, ggb: u64) !void { + var f: cuda.Args = .{}; + f.add(proj); + f.add(conv); + f.add(lw.cw.ptr); + f.add(lw.a_log.ptr); + f.add(lw.dt.ptr); + f.add(gq); + f.add(gk); + f.add(gv); + f.add(ggb); + try cuda.launch.launch(self.k.front, .{ .grid = .{ .x = @intCast(rows), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, self.stream, &f); + } + + /// fn_gdn_rows over `rows` rows from `state`, in place or (state_out) to fn_gdn_rows_to's; readouts into ys. + fn gdnChain(self: *Self, gq: u64, gk: u64, gv: u64, ggb: u64, state: u64, state_out: u64, rows: usize, ys: u64) !void { + var a: cuda.Args = .{}; + a.add(gq); + a.add(gk); + a.add(gv); + a.add(ggb); + a.add(state); + if (state_out != 0) a.add(state_out); + a.add(cint(rows)); + a.add(ys); + try cuda.launch.launch(if (state_out != 0) self.k.chain_to else self.k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, self.stream, &a); + } + + pub fn attnBlock(self: *Self, lw: *const Layer, ai: usize) !void { + const s = self.s; + const n = self.rows; + const pa = s.proj; + try self.lane(s.mixed, s.xsm, lw.proj, pa, s.dummy, n, 1); + if (self.lanes.len > 0) { + // a round: each row's rope, cache writes, selection and attention on its own stream's caches + const own_seq = self.seq; + const own_pos0 = self.pos0; + defer { + self.seq = own_seq; + self.pos0 = own_pos0; + } + for (self.lanes, 0..) |l, r| { + self.seq = l.seq; + self.pos0 = l.pos; + try self.attnRows(lw, ai, pa + r * proj_n * 2, 1, .{ .pos = s.pos + r * 16, .q = s.q + r * q_heads * head_dim * 2, .iq = s.iq + r * index_heads * index_dim * 2, .ao = s.ao + r * out_k * 2 }); + } + } else try self.attnRows(lw, ai, pa, n, .{ .pos = s.pos, .q = s.q, .iq = s.iq, .ao = s.ao }); + try self.tri.attnGate(s.ao, pa, s.gated, s.gxs_attn, n, proj_n, q_heads, head_dim); + try self.lane(s.gated, s.gxs_attn, lw.out, s.dummy, s.part, n, out_sk); + } + + /// Where attnRows reads its rows' positions (16 bytes apart) and writes queries, indexer queries and outputs. + const RowsAt = struct { pos: u64, q: u64, iq: u64, ao: u64 }; + + /// `n` rows of self.seq from self.pos0: rope, cache writes and pooling, then scores, selection and attention. + fn attnRows(self: *Self, lw: *const Layer, ai: usize, pa: u64, n: usize, at: RowsAt) !void { + const s = self.s; + const kc = self.seq.kc[ai].ptr; + const vc = self.seq.vc[ai].ptr; + const ikc = self.seq.ikc[ai].ptr; + const pooled = self.seq.pooled[ai].ptr; + try self.tri.attnPrep(pa, at.pos, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, s.inv, at.q, kc, vc, at.iq, ikc, s.dummy, s.dummy, eps, n, 0); + // attention.qsa_pool for the blocks these rows complete, then each block of rows' scores and attention + try self.tri.qsaPool(ikc, pooled, at.pos, lw.ik_scale.ptr, s.inv, eps, n); + const nb = scoreStride(self.seq.cap); + var r0: usize = 0; + var b: usize = 0; + while (r0 < n) : ({ + r0 += att_rows; + b += 1; + }) { + const m = @min(att_rows, n - r0); + const ends = self.pos0 + r0 + m; + if (ends > self.seq.cap) return error.UnexpectedTensor; + const pos = at.pos + b * 16; + const iq = at.iq + r0 * index_heads * index_dim * 2; + // rows past `top` complete blocks score theirs; _scores leaves the rest alone, so skip it when none is + if (ends / ratio > top) try self.tri.qsaScores(iq, pooled, pos, self.seq.scores, nb, m, @min(nb, cdiv(ends, ratio))); + var a: cuda.Args = .{}; + a.add(self.seq.scores); + a.add(pos); + a.add(s.qids); + a.add(s.qnk); + a.add(s.qspr); + a.add(cint(nb)); + a.add(cint(idw)); + a.add(cint(top)); + a.add(cint(ratio)); + try cuda.launch.launch(self.k.select, .{ .grid = .{ .x = @intCast(m) }, .block = .{ .x = 1024 } }, self.stream, &a); + const chunks = cdiv(@min(ends, wt.dense_keys), 512); + try self.tri.attnChunks(at.q + r0 * q_heads * head_dim * 2, kc, vc, pos, s.po, s.pm, s.pl, s.qids, s.qnk, s.qspr, m, kv_heads, chunks, true); + try self.tri.attnMerge(s.po, s.pm, s.pl, pos, at.ao + r0 * out_k * 2, s.qnk, s.qspr, m, kv_heads, true); + } + } + + /// Router, each row's ten experts and the shared one, the grouped expert launches: y [rows, 11, dims] fp32. + pub fn moeBlock(self: *Self, lw: *const Layer) !void { + const s = self.s; + const n = self.rows; + const pairs = n * slots; + try self.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); + try self.tri.topkMany(s.router_l, s.picks, s.wts, n); + const e = experts_n + 1; + const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; + try self.k.router.route(self.stream, s.picks, pairs, e, rows_tile, plan); + const items = cdiv(pairs, rows_tile) + e; + try self.k.moe.run(self.stream, true, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items); + try self.k.moe.run(self.stream, false, s.moe_act, moe_w, 0, lw.experts_down.ptr, moe_w / 32, dims / 32, plan, s.y, dims, items); + } + + // ---- the pass ------------------------------------------------------------------------------------------ + + /// Token ids, positions and (when the n-gram layer is there) each row's n-gram embedding onto the device. + pub fn stage(self: *Self, ids: []const u32) !void { + const s = self.s; + const n = ids.len; + const id32 = try self.gpa.alloc(i32, n); + defer self.gpa.free(id32); + for (id32, ids) |*o, v| o.* = @intCast(v); + try self.upload(s.ids, std.mem.sliceAsBytes(id32)); + const blocks = if (self.lanes.len > 0) n else cdiv(n, att_rows); + // one int a block (a round: a row), 16 bytes apart: _chunks and _merge take POS0 16-byte aligned + const pos = try self.gpa.alloc(i32, blocks * 4); + defer self.gpa.free(pos); + @memset(pos, 0); + for (0..blocks) |b| pos[4 * b] = @intCast(if (self.lanes.len > 0) self.lanes[b].pos else self.pos0 + b * att_rows); + try self.upload(s.pos, std.mem.sliceAsBytes(pos)); + } + + /// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. + const Gather = struct { + ngram: *const ple.Ngram, + history: []i64, + ids: []const u32, + emb: []u16, + xs: []f32, + failed: ?anyerror = null, + ns: i96 = 0, + io: std.Io, + /// A round's rows, each gathered after its own stream's history (empty: every row after `history`). + lanes: []const Lane = &.{}, + + fn main(g: *Gather) void { + const t0 = std.Io.Clock.awake.now(g.io).toNanoseconds(); + defer g.ns = std.Io.Clock.awake.now(g.io).toNanoseconds() - t0; + if (g.lanes.len == 0) { + g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { + g.failed = err; + }; + return; + } + for (g.lanes, 0..) |l, r| { + g.ngram.rows(l.seq.history[0..], g.ids[r..][0..1], g.emb[r * dims ..][0..dims], g.xs[r * (dims / 32) ..][0 .. dims / 32]) catch |err| { + g.failed = err; + return; + }; + } + } + }; + + /// One chunk of rows at positions pos0.. through every layer and the head. + pub fn chunk(self: *Self, ids: []const u32, pos0: usize, logits_host: ?[]u8) !void { + const s = self.s; + const n = ids.len; + const t0 = std.Io.Clock.awake.now(self.io).toNanoseconds(); + defer if (profile()) { + self.stream.synchronize() catch {}; + std.debug.print("batch: {d} rows at {d} in {d:.1} ms\n", .{ n, pos0, @as(f64, @floatFromInt(std.Io.Clock.awake.now(self.io).toNanoseconds() - t0)) / 1e6 }); + }; + if (n == 0 or n > s.cap) return error.UnexpectedTensor; + if (self.lanes.len > 0) { + // a round reads every row's scores and keeps nothing for a rollback + if (n != self.lanes.len or self.hooks.logits != .all or self.hooks.keep != null or self.hooks.streams != 0) return error.UnexpectedTensor; + for (self.lanes) |l| if (l.pos + 1 > l.seq.cap) return error.UnexpectedTensor; + } else if (pos0 + n > self.seq.cap) return error.UnexpectedTensor; + self.rows = n; + self.pos0 = pos0; + self.pending = null; + const emb = try self.gpa.alloc(u16, n * dims); + defer self.gpa.free(emb); + const xs = try self.gpa.alloc(f32, n * (dims / 32)); + defer self.gpa.free(xs); + var gather: Gather = .{ .ngram = self.m.ngram, .history = self.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = self.io, .lanes = self.lanes }; + var thread: ?std.Thread = try std.Thread.spawn(.{}, Gather.main, .{&gather}); + defer if (thread) |th| th.join(); + try self.stage(ids); + try self.tri.embed(s.ids, self.m.table_w, self.m.table_s, self.m.table_b, s.h, dims, streams, n); + for (self.m.layers, 0..) |*lw, layer| { + const keep: LayerKeep = if (self.hooks.keep) |kp| kp.layers[layer] else .{}; + if (layer == ple_layer) { + thread.?.join(); + thread = null; + if (gather.failed) |err| return err; + if (profile()) std.debug.print("batch: {d} rows' n-gram rows gathered on the host in {d:.1} ms\n", .{ n, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); + try self.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); + try self.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); + if (self.pending != null) try self.writeback(0, n); + try self.pleBlock(lw, ids); + } + try self.writeback(0, n); + try self.readout(lw.attn_hc, 0, n); + if (lw.linear) { + try self.gdnBlock(lw, wt.linearIndex(layer), keep); + } else { + try self.attnBlock(lw, wt.attnIndex(layer)); + } + try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, out_sk); + try self.readout(lw.mlp_hc, 0, n); + try self.moeBlock(lw); + self.pending = layer; + } + try self.head(logits_host); + } + + fn head(self: *Self, logits_host: ?[]u8) !void { + const s = self.s; + const n = self.rows; + switch (self.hooks.logits) { + .none => { + if (self.hooks.streams == 0) return; + try self.writeback(0, n); + try self.copy(self.hooks.streams, s.h, n * wide * 2); + return; + }, + .all => { + try self.writeback(0, n); + try self.copy(self.hooks.streams, s.h, n * wide * 2); + try self.readout(self.m.mixer, 0, n); + const out = if (self.hooks.logits_out != 0) self.hooks.logits_out else return error.UnexpectedTensor; + try self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); + if (logits_host) |lh| { + try self.stream.synchronize(); + const d = self.stream.d; + try d.check(d.api.cuMemcpyDtoH_v2(lh.ptr, out + (n - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH"); + } + }, + .last => { + // the last row's streams, sums, gates and MoE slots moved to row 0, which keeps 16-byte alignment + const r = n - 1; + if (self.hooks.streams != 0) { + // every row written back and kept, then the last row's streams and sums at row 0 + try self.writeback(0, n); + try self.copy(self.hooks.streams, s.h, n * wide * 2); + if (r > 0) { + try self.copy(s.h, s.h + r * wide * 2, wide * 2); + try self.copy(s.pss, s.pss + r * pss_n * 4, pss_n * 4); + } + } else if (r > 0) { + try self.copy(s.h, s.h + r * wide * 2, wide * 2); + try self.copy(s.pss, s.pss + r * pss_n * 4, pss_n * 4); + try self.copy(s.inj, s.inj + r * streams * 2, streams * 2); + try self.copy(s.y, s.y + r * slots * dims * 4, slots * dims * 4); + try self.copy(s.wts, s.wts + r * slots * 4, slots * 4); + } + if (self.hooks.streams == 0) try self.writeback(0, 1); + try self.readout(self.m.mixer, 0, 1); + try self.lane(s.mixed, s.xsm, self.m.head, s.logits, s.dummy, 1, 1); + if (logits_host) |lh| { + try self.stream.synchronize(); + const d = self.stream.d; + try d.check(d.api.cuMemcpyDtoH_v2(lh.ptr, s.logits, vocab * 2), "cuMemcpyDtoH"); + } + }, + } + } + }; +} + +test "a pass's scratch at the dense cap stays near a gigabyte and a half" { + try std.testing.expect(Scratch.bytes(default_chunk) < 2 << 30); + try std.testing.expect(Scratch.bytes(1) < Scratch.bytes(2)); +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig new file mode 100644 index 0000000000..8cd1fb2df8 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -0,0 +1,345 @@ +//! The batched pass's kernels, scratch, a sequence's view, and the hooks a verify window or the MTP head passes. + +const std = @import("std"); +const cuda = @import("cuda"); +const gdn = @import("cuda_gdn.zig"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_prompt_load.zig"); + +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const low = wt.low; +const eps = wt.eps; +const vocab = wt.vocab; +const experts_n = wt.experts_n; +const slots = wt.slots; +const moe_w = wt.moe_w; +const tile = wt.tile; +const q_heads = wt.q_heads; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_heads = wt.index_heads; +const index_dim = wt.index_dim; +const half = wt.half; +const proj_n = wt.proj_n; +const out_k = wt.out_k; +const nch = wt.nch; +const ple_layer = wt.ple_layer; +const ple_tail_n = wt.ple_tail_n; +const state_n = wt.state_n; +const hc_sk = wt.hc_sk; +const out_sk = wt.out_sk; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const Layer = load.Layer; + +/// Normed-stream partial sums a row: D / 256 blocks of each stream. +pub const pss_n = (dims / 256) * streams; +/// The widest projection row: the DeltaNet's q|k|v|z|b|a (attention's q|k|v|indexer is 13952). +pub const proj_row = @max(gdn.proj_width, proj_n); +/// Attention rows a _chunks launch takes (its scratch: rows x nch x heads x dim fp32). +pub const att_rows: usize = 256; +/// Rows a pass holds at most unless asked for more. +pub const default_chunk: usize = 1536; +const ratio = wt.index_ratio; +const top = wt.index_top; +const idw = wt.index_idw; +const cubin_qmm: [:0]const u8 = "_ZN6tf_qmm10qmm_kernelILi32ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKfPKjS3_S3_PvPfiiiiiii"; +const qmm_smem: u32 = 9472; +/// Pairs an expert item takes: the core rows kernel's tile. +const rows_tile: usize = cuda.experts.RowsExperts.tile; +const plan_small: usize = cuda.grouped.small; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// Which rows get the head: the last (a prompt), every row (a verify window or a round), or none. +pub const Logits = enum { last, all, none }; + +/// A DeltaNet layer's window keep: projection rows, conv history before the chunk, and where its end state goes. +pub const LayerKeep = struct { + proj: u64 = 0, + hist: u64 = 0, + /// The chain's end state goes here, the sequence's own left as before the chunk; 0: in place. + state_out: u64 = 0, +}; + +/// What a drafted window keeps for its rollback (cuda_mtp.zig): each DeltaNet layer's rows and states. +pub const Keep = struct { + layers: [layers_n]LayerKeep = @splat(.{}), +}; + +/// A caller's view into the pass. Without any, the pass is the server's prompt pass. +pub const Hooks = struct { + logits: Logits = .last, + /// Device [rows, vocab] bf16 for `.all` (else the pass's own buffer for the last row). + logits_out: u64 = 0, + keep: ?*Keep = null, + /// Device [rows, wide] bf16: every row's streams after the last MoE write-back (the MTP head's input); 0: none. + streams: u64 = 0, +}; + +/// The kernels the pass launches besides the captured Triton ones, loaded once. +pub const Kernels = struct { + qmm: cuda.Module, + experts: cuda.Module, + rows: cuda.Module, + batch: cuda.Module, + /// The core's expert plan and its rows kernels. + router: cuda.grouped.Router, + moe: cuda.experts.RowsExperts, + lane: cuda.Function, + front: cuda.Function, + chain: cuda.Function, + chain_to: cuda.Function, + gnorm: cuda.Function, + sum: cuda.Function, + shift: cuda.Function, + /// Each sparse row's top blocks (attention._select / _select_tiles), native. + select: cuda.Function, + + pub fn load(d: *const cuda.Driver) !Kernels { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var k: Kernels = undefined; + k.qmm = try cuda.Module.load(d, cuda.kernels.qmm); + errdefer k.qmm.unload(); + k.experts = try cuda.Module.load(d, cuda.kernels.experts); + errdefer k.experts.unload(); + k.rows = try cuda.Module.load(d, cuda.kernels.experts_rows); + errdefer k.rows.unload(); + k.batch = try cuda.Module.load(d, cuda.kernels.flashnext_batch); + errdefer k.batch.unload(); + k.lane = try k.qmm.function(cubin_qmm); + try k.lane.allowDynamicShared(qmm_smem); + k.router = try cuda.grouped.Router.resolve(k.experts); + k.moe = try cuda.experts.RowsExperts.resolve(k.rows); + k.sum = try k.batch.function("fn_sum_slices"); + k.shift = try k.batch.function("fn_shift_window"); + k.select = try k.batch.function("fn_qsa_select"); + k.front = try k.batch.function("fn_gdn_front"); + k.chain = try k.batch.function("fn_gdn_rows"); + k.chain_to = try k.batch.function("fn_gdn_rows_to"); + k.gnorm = try k.batch.function("fn_gdn_norm"); + return k; + } + + pub fn deinit(k: *Kernels) void { + k.batch.unload(); + k.rows.unload(); + k.experts.unload(); + k.qmm.unload(); + } +}; + +/// Device scratch for chunks of up to `cap` rows. +pub const Scratch = struct { + cap: usize, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + ids: u64 = 0, + pos: u64 = 0, + h: u64 = 0, + pss: u64 = 0, + norm: u64 = 0, + dn: u64 = 0, + hc_part: u64 = 0, + act: u64 = 0, + xs_act: u64 = 0, + inj: u64 = 0, + mixed: u64 = 0, + xsm: u64 = 0, + proj: u64 = 0, + part: u64 = 0, + gout: u64 = 0, + gxs: u64 = 0, + ys: u64 = 0, + gq: u64 = 0, + gk: u64 = 0, + gv: u64 = 0, + ggb: u64 = 0, + q: u64 = 0, + iq: u64 = 0, + po: u64 = 0, + pm: u64 = 0, + pl: u64 = 0, + /// An attention block's rows' key lists, key counts and sparse flags (attention.AttnScratch ids/nk/sparse). + qids: u64 = 0, + qnk: u64 = 0, + qspr: u64 = 0, + ao: u64 = 0, + gated: u64 = 0, + gxs_attn: u64 = 0, + router_l: u64 = 0, + picks: u64 = 0, + wts: u64 = 0, + members: u64 = 0, + items: u64 = 0, + counts: u64 = 0, + rank: u64 = 0, + hist: u64 = 0, + moe_act: u64 = 0, + y: u64 = 0, + branch: u64 = 0, + ple_emb: u64 = 0, + ple_xs: u64 = 0, + ple_keys: u64 = 0, + ple_vals: u64 = 0, + ple_part: u64 = 0, + ple_gated: u64 = 0, + ple_pss: u64 = 0, + ple_nrow: u64 = 0, + logits: u64 = 0, + inv: u64 = 0, + dummy: u64 = 0, + + fn add(s: *Scratch, gpa: std.mem.Allocator, d: *const cuda.Driver, len: usize) !u64 { + var b = try cuda.DeviceBuffer.alloc(d, @max(len, 256)); + errdefer b.free(); + try s.bufs.append(gpa, b); + return b.ptr; + } + + /// Bytes `alloc(cap)` takes on the device. + pub fn bytes(cap: usize) usize { + const pairs = cap * slots; + const nblk = cdiv(pairs, plan_small); + const items = cdiv(pairs, tile) + experts_n + 1; + var n: usize = 0; + n += cap * 4 + (cap + 1) * 16; + n += cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; + n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; + n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; + n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; + n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; + n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; + n += att_rows * (idw + 2) * 4; + n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; + n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; + n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; + n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + half * 4 + 4096; + return n; + } + + pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, cap: usize) !Scratch { + var s: Scratch = .{ .cap = cap }; + errdefer s.free(gpa); + const pairs = cap * slots; + const nblk = cdiv(pairs, plan_small); + const items = cdiv(pairs, tile) + experts_n + 1; + s.ids = try s.add(gpa, d, cap * 4); + // a position every 16 bytes: one a block of att_rows rows, or one a row of a decode round + s.pos = try s.add(gpa, d, (cap + 1) * 16); + s.h = try s.add(gpa, d, cap * wide * 2); + s.pss = try s.add(gpa, d, cap * pss_n * 4); + s.norm = try s.add(gpa, d, cap * wide * 2); + s.dn = try s.add(gpa, d, cap * (low + 4) * 2); + s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); + s.act = try s.add(gpa, d, cap * low * 2); + s.xs_act = try s.add(gpa, d, cap * (low / 32) * 4); + s.inj = try s.add(gpa, d, cap * streams * 2); + s.mixed = try s.add(gpa, d, cap * dims * 2); + s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); + s.proj = try s.add(gpa, d, cap * proj_row * 2); + s.part = try s.add(gpa, d, out_sk * cap * dims * 4); + s.gout = try s.add(gpa, d, cap * gdn.value_dim * 2); + s.gxs = try s.add(gpa, d, cap * (gdn.value_dim / 32) * 4); + s.ys = try s.add(gpa, d, cap * gdn.value_dim * 4); + s.gq = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); + s.gk = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); + s.gv = try s.add(gpa, d, cap * gdn.value_dim * 4); + s.ggb = try s.add(gpa, d, cap * gdn.nv * 2 * 4); + s.q = try s.add(gpa, d, cap * q_heads * head_dim * 2); + s.iq = try s.add(gpa, d, cap * index_heads * index_dim * 2); + s.po = try s.add(gpa, d, att_rows * nch * q_heads * head_dim * 4); + s.pm = try s.add(gpa, d, att_rows * nch * q_heads * 4); + s.pl = try s.add(gpa, d, att_rows * nch * q_heads * 4); + s.qids = try s.add(gpa, d, att_rows * idw * 4); + s.qnk = try s.add(gpa, d, att_rows * 4); + s.qspr = try s.add(gpa, d, att_rows * 4); + s.ao = try s.add(gpa, d, cap * out_k * 2); + s.gated = try s.add(gpa, d, cap * out_k * 2); + s.gxs_attn = try s.add(gpa, d, cap * (out_k / 32) * 4); + s.router_l = try s.add(gpa, d, cap * (experts_n + 1) * 4); + s.picks = try s.add(gpa, d, pairs * 4); + s.wts = try s.add(gpa, d, pairs * 4); + s.members = try s.add(gpa, d, pairs * 4); + s.items = try s.add(gpa, d, items * 3 * 4); + s.counts = try s.add(gpa, d, 8); + s.rank = try s.add(gpa, d, pairs * 4); + s.hist = try s.add(gpa, d, nblk * (experts_n + 1) * 4); + s.moe_act = try s.add(gpa, d, pairs * moe_w * 2); + s.y = try s.add(gpa, d, pairs * dims * 4); + s.branch = try s.add(gpa, d, cap * dims * 2); + s.ple_emb = try s.add(gpa, d, cap * dims * 2); + s.ple_xs = try s.add(gpa, d, cap * (dims / 32) * 4); + s.ple_keys = try s.add(gpa, d, cap * wide * 2); + s.ple_vals = try s.add(gpa, d, cap * dims * 2); + s.ple_part = try s.add(gpa, d, 4 * cap * dims * 4); + s.ple_gated = try s.add(gpa, d, cap * wide * 2); + s.ple_pss = try s.add(gpa, d, cap * streams * 4); + s.ple_nrow = try s.add(gpa, d, cap * wide * 2); + s.logits = try s.add(gpa, d, vocab * 2); + var inv: [half]f32 = undefined; + for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); + s.inv = try s.add(gpa, d, half * 4); + try s.bufs.items[s.bufs.items.len - 1].upload(0, std.mem.sliceAsBytes(&inv)); + s.dummy = try s.add(gpa, d, 4096); + try s.bufs.items[s.bufs.items.len - 1].fill8(0, null); + return s; + } + + pub fn free(s: *Scratch, gpa: std.mem.Allocator) void { + for (s.bufs.items) |*b| b.free(); + s.bufs.deinit(gpa); + s.* = undefined; + } +}; + +/// A sequence's state the pass reads and leaves as the row-by-row path would. +pub const SeqState = struct { + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + /// Each attention layer's pooled indexer keys [cap / 4, 128] bf16. + pooled: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + history: *[2]i64, + /// Tokens the caches hold. + cap: usize, + /// Block scores for an attention block's rows [min(rows, att_rows), scoreStride(cap)] fp32 (the engine's). + scores: u64, +}; + +/// The block scores' row stride for a cache of `cap` tokens: its blocks rounded up to 16 (the captured _scores'). +pub fn scoreStride(cap: usize) usize { + return (cdiv(cap, ratio) + 15) / 16 * 16; +} + +/// Bytes of block scores a pass over a cache of `cap` tokens needs for chunks of up to `rows` rows. +pub fn scoreBytes(cap: usize, rows: usize) usize { + return @min(rows, att_rows) * scoreStride(cap) * 4; +} + +/// One row of a decode round: its stream's next token, on that stream's own caches at its own position. +pub const Lane = struct { seq: SeqState, pos: usize }; + +/// The engine's resident weights the pass reads. +pub const Model = struct { + table_w: u64, + table_s: u64, + table_b: u64, + layers: []const Layer, + mixer: HcW, + head: Q4, + ngram: *const ple.Ngram, +}; diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig new file mode 100644 index 0000000000..5160d3eaf6 --- /dev/null +++ b/zig/src/families/flashnext/cuda_config.zig @@ -0,0 +1,85 @@ +//! The checkpoint's config.json against the shapes the CUDA path is built for: a different one is refused at load. + +const std = @import("std"); +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); + +/// One config field and the value the kernels and buffers assume. +const Want = struct { name: []const u8, value: i64 }; + +/// The text config's shapes the CUDA path's constants hold. +const shapes = [_]Want{ + .{ .name = "hidden_size", .value = wt.dims }, + .{ .name = "num_hidden_layers", .value = wt.layers_n }, + .{ .name = "full_attention_interval", .value = 4 }, + .{ .name = "vocab_size", .value = wt.vocab }, + .{ .name = "num_attention_heads", .value = wt.q_heads }, + .{ .name = "num_key_value_heads", .value = wt.kv_heads }, + .{ .name = "head_dim", .value = wt.head_dim }, + .{ .name = "num_experts", .value = wt.experts_n }, + .{ .name = "num_experts_per_tok", .value = wt.slots - 1 }, + .{ .name = "moe_intermediate_size", .value = wt.moe_w }, + .{ .name = "shared_expert_intermediate_size", .value = wt.moe_w }, + .{ .name = "hc_count", .value = wt.streams }, + .{ .name = "hc_lowrank", .value = wt.low }, + .{ .name = "linear_num_key_heads", .value = gdn.nk }, + .{ .name = "linear_num_value_heads", .value = gdn.nv }, + .{ .name = "linear_key_head_dim", .value = gdn.dk }, + .{ .name = "linear_value_head_dim", .value = gdn.dv }, + .{ .name = "linear_conv_kernel_dim", .value = gdn.taps }, + .{ .name = "indexer_n_heads", .value = wt.index_heads }, + .{ .name = "indexer_head_dim", .value = wt.index_dim }, + .{ .name = "indexer_compress_ratio", .value = wt.index_ratio }, + .{ .name = "mtp_num_hidden_layers", .value = 1 }, +}; + +fn int(o: std.json.ObjectMap, name: []const u8) ?i64 { + const v = o.get(name) orelse return null; + return if (v == .integer) v.integer else null; +} + +/// The first config value that differs from what the CUDA path serves, or null when every one matches. +pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u8 { + const q = quant orelse return "quantization"; + if (int(q, "bits") != 4) return "quantization.bits"; + if (int(q, "group_size") != 32) return "quantization.group_size"; + if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return "quantization.mode"; + for (shapes) |w| if (int(text, w.name) != w.value) return w.name; + return null; +} + +/// Reads `dir`/config.json; error.UnsupportedCheckpoint names the field in a log line when one differs. +pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(bytes); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + if (parsed.value != .object) return error.UnsupportedCheckpoint; + const root = parsed.value.object; + const text = if (root.get("text_config")) |v| (if (v == .object) v.object else root) else root; + const quant = if (root.get("quantization")) |v| (if (v == .object) v.object else null) else null; + if (mismatch(text, quant)) |field| { + std.log.err("flashnext: config.json's {s} differs from the MLX 4-bit group-32 Flash Next the CUDA path serves", .{field}); + return error.UnsupportedCheckpoint; + } +} + +test "the served shapes pass and a 6-bit checkpoint or another width is named" { + const gpa = std.testing.allocator; + var text: std.json.ObjectMap = .empty; + defer text.deinit(gpa); + for (shapes) |w| try text.put(gpa, w.name, .{ .integer = w.value }); + var quant: std.json.ObjectMap = .empty; + defer quant.deinit(gpa); + try quant.put(gpa, "bits", .{ .integer = 4 }); + try quant.put(gpa, "group_size", .{ .integer = 32 }); + try std.testing.expectEqual(@as(?[]const u8, null), mismatch(text, quant)); + try quant.put(gpa, "bits", .{ .integer = 6 }); + try std.testing.expectEqualStrings("quantization.bits", mismatch(text, quant).?); + try quant.put(gpa, "bits", .{ .integer = 4 }); + try text.put(gpa, "hidden_size", .{ .integer = 4096 }); + try std.testing.expectEqualStrings("hidden_size", mismatch(text, quant).?); + try std.testing.expectEqualStrings("quantization", mismatch(text, null).?); +} diff --git a/zig/src/families/flashnext/cuda_embed_main.zig b/zig/src/families/flashnext/cuda_embed_main.zig index 294e0e6b69..d91d104d9a 100644 --- a/zig/src/families/flashnext/cuda_embed_main.zig +++ b/zig/src/families/flashnext/cuda_embed_main.zig @@ -118,7 +118,7 @@ pub fn main(init: std.process.Init) !u8 { const norm_mismatch = try runDown(gpa, io, args[1], &driver, &stream, &set, &mapped, wide.ptr, pss.ptr, wide_u16, pss_bytes, table.dims, streams, nc); if (norm_mismatch != 0) return 1; - const ngram = try flash.ple.table(flash.Tri, gpa, io, args[1], &driver, &stream, tri, token); + const ngram = try flash.ple_check.table(flash.Tri, gpa, io, args[1], &driver, &stream, tri, token); const scored = try flash.mix.scores(flash.Tri, gpa, io, args[1], &driver, &stream, tri, wide_u16); return if (ngram == 0 and scored == 0) 0 else 1; } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 75339636e9..3f4bfcfe46 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -12,8 +12,11 @@ pub const mlp = @import("cuda_mlp.zig"); pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); +pub const ple_check = @import("cuda_ple_check.zig"); pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); +/// The row-by-row reference path, for flashnext-prompt +pub const row_path = @import("cuda_rows.zig"); pub const native = @import("cuda_lanes.zig").Namespace(Tri); const p = aot.ptr; @@ -63,6 +66,27 @@ pub const Tri = struct { }, &.{ ci("D", dims), ci("S_COPIES", copies) }); } + /// glue.rmsnorm per row and `group` features: bf16 out [rows, dims] and 32-group sums; x rows `x_stride` apart. + pub fn rmsnorm(t: Tri, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { + try t.run("_rmsnorm", .{ u(rows), u(dims / group), 1 }, &.{ + p("X", "*bf16", x), + p("W", "*fp32", w), + p("OUT", "*bf16", out), + p("XS", "*fp32", xs), + aot.float("eps", eps), + int("x_stride", x_stride), + }, &.{ ci("D", dims), ci("G", group), ci("BLOCK", rmsBlock(group)) }); + } + + /// glue.add_streams: out[r, s] = bf16(e[r] + hs[r, s]), the MTP head's input streams. Program (row, 256 columns). + pub fn addStreams(t: Tri, e: u64, hs: u64, out: u64, rows: usize, dims: usize, streams: usize) !void { + try t.run("_add_streams", .{ u(rows), u(dims / 256), 1 }, &.{ + p("E", "*bf16", e), + p("HS", "*bf16", hs), + p("OUT", "*bf16", out), + }, &.{ ci("D", dims), ci("S", streams), ci("BLOCK", 256) }); + } + /// glue.attn_gate: one program per row and per query head. pub fn attnGate(t: Tri, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { try t.run("_attn_gate", .{ u(rows), u(q_heads), 1 }, &.{ @@ -73,8 +97,8 @@ pub const Tri = struct { }, &.{ ci("PW", pair_width), ci("NQ", q_heads), ci("HD", head_dim) }); } - /// glue.attn_prep for text positions. The grid includes the indexer key. `length` is unused when MODE is 0. - pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, ks: u64, vs: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { + /// glue.attn_prep for text positions, bf16 KV (BITS 0, the scale pointers unread); `length` unused at MODE 0. + pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ p("P", "*bf16", proj), p("POS0", "*i32", pos), @@ -85,8 +109,8 @@ pub const Tri = struct { p("Q", "*bf16", q), p("KC", "*bf16", kc), p("VC", "*bf16", vc), - p("KS", "*fp16", ks), - p("VS", "*fp16", vs), + p("KS", "*fp16", rope), + p("VS", "*fp16", rope), p("IQ", "*bf16", iq), p("IKC", "*bf16", ikc), p("ROPE", "*i32", rope), @@ -96,14 +120,14 @@ pub const Tri = struct { }, &.{ ci("PW", 13952), ci("NQ", 24), ci("NKV", 2), ci("HD", 256), ci("NI", 4), ci("IHD", 128), ci("HALF", 32), ci("BITS", 0), ci("MODE", 0), ci("S1", 11), ci("S2", 10) }); } - /// attention._chunks: one decode row, dense, the captured three-chunk scratch, one chunk launched. - pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, ks: u64, vs: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize) !void { + /// attention._chunks over 512-key chunks, bf16 KV; `qsa` the sparse form (NCH 5, QSA 1), else dense (NCH 3, QSA 0). + pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ p("Q", "*bf16", q), p("KC", "*bf16", kc), p("VC", "*bf16", vc), - p("KSC", "*fp16", ks), - p("VSC", "*fp16", vs), + p("KSC", "*fp16", sparse), + p("VSC", "*fp16", sparse), p("POS0", "*i32", pos), p("PO", "*fp32", po), p("PM", "*fp32", pm), @@ -111,11 +135,11 @@ pub const Tri = struct { p("IDS", "*i32", ids), p("NKR", "*i32", nk), p("SPR", "*i32", sparse), - }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", 3), aot.cf("SCALE", 0.0625), ci("IDW", 2052), ci("QSA", 0), ci("BITS", 0) }); + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", if (qsa) 5 else 3), aot.cf("SCALE", 0.0625), ci("IDW", 2052), ci("QSA", @intFromBool(qsa)), ci("BITS", 0) }); } - /// attention._merge: one program per row and per KV head. - pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize) !void { + /// attention._merge: one program per row and per KV head, in attnChunks' form. + pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { try t.run("_merge", .{ u(rows), u(kv), 1 }, &.{ p("PO", "*fp32", po), p("PM", "*fp32", pm), @@ -124,7 +148,32 @@ pub const Tri = struct { p("OUT", "*bf16", out), p("NKR", "*i32", nk), p("SPR", "*i32", sparse), - }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", 3), ci("QSA", 0), ci("BITS", 0) }); + }, &.{ ci("H", 24), ci("HK", 2), ci("D", 256), ci("G", 12), ci("CH", 512), ci("NCH", if (qsa) 5 else 3), ci("QSA", @intFromBool(qsa)), ci("BITS", 0) }); + } + + /// attention.qsa_pool for text rows (MODE 0): every complete block of 4 indexer keys pooled, normed and roped. + pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { + try t.run("_pool", .{ u(rows / 4 + 2), 1, 1 }, &.{ + p("IKC", "*bf16", ikc), + p("POOLED", "*bf16", pooled), + p("POS0", "*i32", pos), + p("W", "*fp32", k_scale), + p("INV", "*fp32", inv), + aot.float("eps", eps), + int("R", rows), + int("length", 0), + }, &.{ ci("DI", 128), ci("HALF", 32), ci("RATIO", 4), ci("MODE", 0), ci("S1", 11), ci("S2", 10) }); + } + + /// attention._scores: rows' complete blocks scored; `nb` is the row stride (a multiple of 16), `blocks` the grid. + pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { + try t.run("_scores", .{ u(rows), cdiv(blocks, 64), 1 }, &.{ + p("IQ", "*bf16", iq), + p("POOLED", "*bf16", pooled), + p("POS0", "*i32", pos), + p("SC", "*fp32", sc), + int("NB", nb), + }, &.{ ci("HI", 4), ci("DI", 128), ci("RATIO", 4), ci("TOP", 512), ci("BB", 64) }); } /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. @@ -227,23 +276,33 @@ pub const Tri = struct { /// moe.router for one decode row. M folds to a constexpr. 513 rows: the experts, then the shared gate. pub fn router(t: Tri, x: u64, w: u64, out: u64, x_stride: usize) !void { - try t.run("_router", .{ 1, cdiv(513, 32), 1 }, &.{ - p("X", "*bf16", x), - p("W", "*bf16", w), - p("OUT", "*fp32", out), - int("M", 1), - int("x_stride", x_stride), - }, &.{ ci("D", 2560), ci("NE", 513), ci("BM", 16), ci("BLOCK_E", 32), ci("BK", 256) }); + try t.routerMany(x, w, out, x_stride, 1); } /// moe._topk_rows: one program, the row's 10 experts and the shared expert. pub fn topkRows(t: Tri, logits: u64, pick: u64, wts: u64) !void { - try t.run("_topk_rows", .{ 1, 1, 1 }, &.{ + try t.topkMany(logits, pick, wts, 1); + } + + /// moe._topk_rows over `rows` rows of [rows, 513] logits: picks and weights [rows, 11]. + pub fn topkMany(t: Tri, logits: u64, pick: u64, wts: u64, rows: usize) !void { + try t.run("_topk_rows", .{ u(rows), 1, 1 }, &.{ p("L", "*fp32", logits), p("PICK", "*i32", pick), p("WTS", "*fp32", wts), }, &.{ ci("NE", 512), ci("NL", 513), ci("TOPK", 10), ci("SLOTS", 11), ci("BLOCK", 1024), ci("SLOTP", 16) }); } + + /// moe.router over `rows` rows at the one-row tiles (BM 16, BLOCK_E 32, BK 256), each row keeping its own bits. + pub fn routerMany(t: Tri, x: u64, w: u64, out: u64, x_stride: usize, rows: usize) !void { + try t.run("_router", .{ cdiv(rows, 16), cdiv(513, 32), 1 }, &.{ + p("X", "*bf16", x), + p("W", "*bf16", w), + p("OUT", "*fp32", out), + int("M", rows), + int("x_stride", x_stride), + }, &.{ ci("D", 2560), ci("NE", 513), ci("BM", 16), ci("BLOCK_E", 32), ci("BK", 256) }); + } }; test { @@ -257,7 +316,10 @@ test { _ = ple; _ = mix; _ = prompt; + _ = prompt.batch; _ = native; + _ = @import("cuda_config.zig"); + _ = row_path; } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_moe.zig b/zig/src/families/flashnext/cuda_moe.zig index 5a0030dfc0..a59407876f 100644 --- a/zig/src/families/flashnext/cuda_moe.zig +++ b/zig/src/families/flashnext/cuda_moe.zig @@ -108,7 +108,7 @@ fn routerDots(x: []const u16, w: []const u16, out: []f32) void { } } -/// Largest `k` of the first `n` logits (lower id on a tie), weights exp(l - l0) / sum rounded to bf16, then the shared gate. +/// Largest `k` of the first `n` logits (lower id on a tie), softmax weights rounded to bf16, then the shared gate. fn choose(logits: []const f32, n: usize, k: usize, picks: []i32, wts: []f32) void { var scratch: [count]f32 = undefined; for (0..n) |i| scratch[i] = logits[i]; diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index 751f8a8c75..96348db971 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -1,4 +1,4 @@ -//! Layer 2's n-gram table for one decode token: hash the row ids, gather those rows, dequant like glue._ple_embed. +//! Layer 2's n-gram table: the hash of each token's row ids, and the mapped table a pass gathers its rows from. const std = @import("std"); const cuda = @import("cuda"); @@ -10,14 +10,14 @@ const golden: u64 = 0x9E3779B97F4A7C15; const mix1: u64 = 0xBF58476D1CE4E5B9; const mix2: u64 = 0x94D049BB133111EB; const prime_step: u64 = 10007; -const max_n: usize = 8; -const max_heads: usize = 64; +pub const max_n: usize = 8; +pub const max_heads: usize = 64; const max_shards: usize = 128; -const layer: usize = 1; +pub const layer: usize = 1; const Tensor = core.safetensors.Tensor; -const Hash = struct { +pub const Hash = struct { n: usize, per: usize, heads: usize, @@ -29,7 +29,7 @@ const Hash = struct { mult: [max_n]i64 = undefined, starts: [max_shards + 1]i64 = undefined, - fn init(args: struct { + pub fn init(args: struct { vocab: u64, ngram: usize, per: usize, @@ -73,7 +73,7 @@ const Hash = struct { return h; } - fn locate(h: Hash, row: i64) !struct { shard: usize, local: usize } { + pub fn locate(h: Hash, row: i64) !struct { shard: usize, local: usize } { if (row < 0 or row >= h.starts[h.shards]) return error.UnexpectedTensor; var lo: usize = 0; var hi = h.shards; @@ -114,7 +114,7 @@ fn divCeil(a: i64, b: i64) i64 { } /// Row ids of one new token after `history` (n - 1 ids). A fresh decode window is n - 1 copies of eos. -fn lastIds(h: Hash, history: []const i64, token: i64, out: []i64) !void { +pub fn lastIds(h: Hash, history: []const i64, token: i64, out: []i64) !void { if (history.len != h.n - 1 or out.len != h.heads) return error.UnexpectedTensor; var seq: [max_n]i64 = undefined; for (history, 0..) |v, i| seq[i] = v; @@ -149,13 +149,13 @@ fn lastIds(h: Hash, history: []const i64, token: i64, out: []i64) !void { } } -fn toBf16(v: f32) u16 { +pub fn toBf16(v: f32) u16 { const bits: u32 = @bitCast(v); if (std.math.isNan(v)) return @intCast((bits >> 16) | 0x40); return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); } -fn sumGroups(row: []const u16, out: []f32) void { +pub fn sumGroups(row: []const u16, out: []f32) void { const groups = row.len / 32; for (0..groups) |g| { var buf: [32]f32 = undefined; @@ -169,79 +169,52 @@ fn sumGroups(row: []const u16, out: []f32) void { } } -const Open = struct { +/// One mapped shard (core.safetensors.File) with a typed lookup. +pub const Open = struct { + file: core.safetensors.File, io: std.Io, - path: []u8, - file: std.Io.File, - map: std.Io.File.MemoryMap, - parsed: std.json.Parsed(std.json.Value), - at: usize, - - fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Open { - const owned = try gpa.dupe(u8, path); - errdefer gpa.free(owned); - var file = try std.Io.Dir.cwd().openFile(io, owned, .{}); - errdefer file.close(io); - const len: usize = @intCast(try file.length(io)); - if (len < 8) return error.BadSafetensors; - var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); - errdefer map.destroy(io); - const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); - if (header_len > len - 8) return error.BadSafetensors; - const parsed = try std.json.parseFromSlice(std.json.Value, gpa, map.memory[8..][0..header_len], .{}); - return .{ .io = io, .path = owned, .file = file, .map = map, .parsed = parsed, .at = 8 + header_len }; + + pub fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Open { + return .{ .file = try core.safetensors.File.open(gpa, io, path), .io = io }; } - fn close(self: *Open, gpa: std.mem.Allocator) void { - self.parsed.deinit(); - self.map.destroy(self.io); + pub fn close(self: *Open, gpa: std.mem.Allocator) void { + _ = gpa; self.file.close(self.io); - gpa.free(self.path); - self.* = undefined; } - fn tensor(self: *const Open, name: []const u8, dtype: core.safetensors.DType) !Tensor { - const o = (self.parsed.value.object.get(name) orelse return error.MissingTensor).object; - const got = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; - if (got != dtype) return error.UnexpectedTensor; - const shape = o.get("shape").?.array.items; - if (shape.len == 0 or shape.len > 4) return error.UnexpectedTensor; - const offs = o.get("data_offsets").?.array.items; - const begin: usize = @intCast(offs[0].integer); - const end: usize = @intCast(offs[1].integer); - const data = self.map.memory[self.at..]; - if (end < begin or end > data.len) return error.BadSafetensors; - var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = data[begin..end] }; - for (shape, 0..) |d, i| t.shape[i] = @intCast(d.integer); + pub fn tensor(self: *const Open, name: []const u8, dtype: core.safetensors.DType) !Tensor { + const t = self.file.get(name) orelse return error.MissingTensor; + if (t.dtype != dtype) return error.UnexpectedTensor; return t; } }; -fn readI64(t: Tensor, out: []i64) !void { +pub fn readI64(t: Tensor, out: []i64) !void { if (t.dtype != .i64 or t.rank != 1 or t.dim(0) != out.len or t.bytes.len != out.len * 8) return error.UnexpectedTensor; for (out, 0..) |*o, i| o.* = std.mem.readInt(i64, t.bytes[8 * i ..][0..8], .little); } -fn weightName(buf: []u8, shard: usize, nested: bool) ![]const u8 { +pub fn weightName(buf: []u8, shard: usize, nested: bool) ![]const u8 { return if (nested) std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.weight", .{ layer, shard }) else std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.weight", .{ layer, shard }); } -fn partName(buf: []u8, shard: usize, nested: bool, part: []const u8) ![]const u8 { +pub fn partName(buf: []u8, shard: usize, nested: bool, part: []const u8) ![]const u8 { return if (nested) std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.{s}", .{ layer, shard, part }) else std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.{s}", .{ layer, shard, part }); } -fn findOpen(opened: []Open, path: []const u8) ?*Open { - for (opened) |*o| if (std.mem.eql(u8, o.path, path)) return o; +pub fn findOpen(opened: []Open, path: []const u8) ?*Open { + for (opened) |*o| if (std.mem.eql(u8, o.file.path, path)) return o; return null; } -fn copyRow(file: *const Open, shard: usize, nested: bool, local: usize, dh: usize, words: []u8, scales: []u8, biases: []u8) !void { +pub fn copyRow(file: *const Open, shard: usize, nested: bool, local: usize, dh: usize, words: []u8, scales: []u8, biases: []u8) !void { var name: [180]u8 = undefined; const w = try file.tensor(try partName(&name, shard, nested, "weight"), .u32); const s = try file.tensor(try partName(&name, shard, nested, "scales"), .bf16); @@ -256,17 +229,14 @@ fn copyRow(file: *const Open, shard: usize, nested: bool, local: usize, dh: usiz @memcpy(biases, b.bytes[local * gb ..][0..gb]); } -fn tableScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap) !f32 { +pub fn tableScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap) !f32 { const name = "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale"; const file = wm.get(name) orelse return 1; const path = try std.fs.path.join(gpa, &.{ dir, file.string }); defer gpa.free(path); var opened = try Open.load(gpa, io, path); defer opened.close(gpa); - const t = try opened.tensor(name, blk: { - const o = opened.parsed.value.object.get(name).?.object; - break :blk core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; - }); + const t = opened.file.get(name) orelse return error.MissingTensor; if (t.numel() != 1) return error.UnexpectedTensor; return switch (t.dtype) { .f32 => @bitCast(std.mem.readInt(u32, t.bytes[0..4], .little)), @@ -301,7 +271,7 @@ const NgramSpec = struct { }; /// `ple_layer_ids` value 2 is decoder layer 1. The eos id is the text config's, which the hash mixes in. -fn readSpec(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !NgramSpec { +pub fn readSpec(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !NgramSpec { const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); defer gpa.free(path); const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); @@ -354,215 +324,142 @@ pub fn eosOf(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !i64 { return cfg.eos; } -/// Host dequant of one token's sixteen n-gram rows. `history` is the previous n - 1 ids. -pub fn embedding(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, history: []const i64, token: u32, out: []u16, xs: []f32) !void { - const cfg = try readSpec(gpa, io, model_dir); - const hash = try Hash.init(.{ - .vocab = cfg.vocab, - .ngram = cfg.ngram, - .per = cfg.per, - .base = cfg.base, - .divisor = cfg.divisor, - .shards = cfg.shards, - .seed = cfg.seed, - .eos = cfg.eos, - .ple_dim = cfg.ple_dim, - .ple_index = cfg.ple_index, - }); - if (history.len != hash.n - 1 or out.len != hash.heads * hash.dh or xs.len != out.len / 32) return error.UnexpectedTensor; - var ids: [max_heads]i64 = undefined; - try lastIds(hash, history, @intCast(token), ids[0..hash.heads]); - const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); - defer gpa.free(index_path); - const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); - defer gpa.free(text); - const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); - defer parsed.deinit(); - const wm = parsed.value.object.get("weight_map").?.object; - const scale = try tableScale(gpa, io, model_dir, wm); - const heads = hash.heads; - const dh = hash.dh; - const words = try gpa.alloc(u8, heads * (dh / 8) * 4); - defer gpa.free(words); - const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); - defer gpa.free(scales); - const biases = try gpa.alloc(u8, scales.len); - defer gpa.free(biases); - var opened: [max_heads]Open = undefined; - var nopen: usize = 0; - defer for (opened[0..nopen]) |*o| o.close(gpa); - var key: [180]u8 = undefined; - for (0..heads) |h| { - const loc = try hash.locate(ids[h]); - var nested = false; - const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| - v.string - else blk: { - nested = true; - break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; - }; - const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); - defer gpa.free(path); - const file = findOpen(opened[0..nopen], path) orelse file: { - opened[nopen] = try Open.load(gpa, io, path); - nopen += 1; - break :file &opened[nopen - 1]; - }; - const wb = (dh / 8) * 4; - const gb = (dh / 32) * 2; - try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); - } - for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, out[h * dh ..][0..dh]); - if (scale != 1) { - for (out) |*v| { - const f: f32 = @bitCast(@as(u32, v.*) << 16); - v.* = toBf16(f * scale); +/// The n-gram table held open for an engine's life: the hash, the scale, and each shard's slices of its mapped file. +pub const Ngram = struct { + gpa: std.mem.Allocator, + hash: Hash, + scale: f32, + files: std.ArrayList(Open) = .empty, + w: [max_shards][]const u8 = undefined, + s: [max_shards][]const u8 = undefined, + b: [max_shards][]const u8 = undefined, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { + const cfg = try readSpec(gpa, io, model_dir); + var g: Ngram = .{ .gpa = gpa, .hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }), .scale = 1 }; + errdefer g.deinit(); + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + g.scale = try tableScale(gpa, io, model_dir, wm); + const dh = g.hash.dh; + var key: [180]u8 = undefined; + for (0..g.hash.shards) |shard| { + var nested = false; + const mapped = if (wm.get(try weightName(&key, shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(g.files.items, path) orelse file: { + var o = try Open.load(gpa, io, path); + errdefer o.close(gpa); + std.posix.madvise(o.file.map.memory.ptr, o.file.map.memory.len, std.posix.MADV.RANDOM) catch {}; + try g.files.append(gpa, o); + break :file &g.files.items[g.files.items.len - 1]; + }; + const w = try file.tensor(try partName(&key, shard, nested, "weight"), .u32); + const s = try file.tensor(try partName(&key, shard, nested, "scales"), .bf16); + const b = try file.tensor(try partName(&key, shard, nested, "biases"), .bf16); + if (w.dim(1) != dh / 8 or s.dim(1) != dh / 32 or b.dim(1) != dh / 32) return error.UnexpectedTensor; + g.w[shard] = w.bytes; + g.s[shard] = s.bytes; + g.b[shard] = b.bytes; } + return g; } - sumGroups(out, xs); -} -/// One decode token. `ple_layer_ids` `[2]` is this layer. Returns 0 when the gathered rows match the host dequant. -pub fn table(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, token: u32) !u32 { - const cfg = try readSpec(gpa, io, model_dir); - const hash = try Hash.init(.{ - .vocab = cfg.vocab, - .ngram = cfg.ngram, - .per = cfg.per, - .base = cfg.base, - .divisor = cfg.divisor, - .shards = cfg.shards, - .seed = cfg.seed, - .eos = cfg.eos, - .ple_dim = cfg.ple_dim, - .ple_index = cfg.ple_index, - }); - var history: [max_n]i64 = undefined; - for (0..hash.n - 1) |i| history[i] = hash.eos; - var ids: [max_heads]i64 = undefined; - try lastIds(hash, history[0 .. hash.n - 1], token, ids[0..hash.heads]); - - const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); - defer gpa.free(index_path); - const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); - defer gpa.free(text); - const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); - defer parsed.deinit(); - const wm = parsed.value.object.get("weight_map").?.object; - - var cname: [160]u8 = undefined; - const mult_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.layer_multipliers", .{layer}); - const const_file = wm.get(mult_name) orelse return error.MissingTensor; - const const_path = try std.fs.path.join(gpa, &.{ model_dir, const_file.string }); - defer gpa.free(const_path); - var constants = try Open.load(gpa, io, const_path); - defer constants.close(gpa); - var got_m: [max_n]i64 = undefined; - var got_o: [max_heads]i64 = undefined; - var got_s: [max_heads]i64 = undefined; - try readI64(try constants.tensor(mult_name, .i64), got_m[0..hash.n]); - const off_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_offsets", .{layer}); - try readI64(try constants.tensor(off_name, .i64), got_o[0..hash.heads]); - const size_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_vocab_sizes", .{layer}); - try readI64(try constants.tensor(size_name, .i64), got_s[0..hash.heads]); - if (!std.mem.eql(i64, got_m[0..hash.n], hash.mult[0..hash.n]) or !std.mem.eql(i64, got_o[0..hash.heads], hash.offsets[0..hash.heads]) or !std.mem.eql(i64, got_s[0..hash.heads], hash.sizes[0..hash.heads])) { - std.debug.print("ngram constants differ\n", .{}); - return 1; + pub fn deinit(g: *Ngram) void { + for (g.files.items) |*o| o.close(g.gpa); + g.files.deinit(g.gpa); } - const scale = try tableScale(gpa, io, model_dir, wm); - - const heads = hash.heads; - const dh = hash.dh; - const words = try gpa.alloc(u8, heads * (dh / 8) * 4); - defer gpa.free(words); - const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); - defer gpa.free(scales); - const biases = try gpa.alloc(u8, scales.len); - defer gpa.free(biases); - var opened: [max_heads]Open = undefined; - var nopen: usize = 0; - defer for (opened[0..nopen]) |*o| o.close(gpa); - var key: [180]u8 = undefined; - for (0..heads) |h| { - const loc = try hash.locate(ids[h]); - var nested = false; - const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| - v.string - else blk: { - nested = true; - break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; - }; - const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); - defer gpa.free(path); - const file = findOpen(opened[0..nopen], path) orelse file: { - opened[nopen] = try Open.load(gpa, io, path); - nopen += 1; - break :file &opened[nopen - 1]; - }; - const wb = (dh / 8) * 4; - const gb = (dh / 32) * 2; - try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + + /// Values a row's embedding holds (heads x head width). + pub fn width(g: *const Ngram) usize { + return g.hash.heads * g.hash.dh; } - const host = try gpa.alloc(u16, heads * dh); - defer gpa.free(host); - for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, host[h * dh ..][0..dh]); - if (scale != 1) { - for (host) |*v| { - const f: f32 = @bitCast(@as(u32, v.*) << 16); - v.* = toBf16(f * scale); + /// Rows [lo, hi) of a chunk's token-major row ids `ids`, each dequantized into its head's slice of `out`. + fn fill(g: *const Ngram, ids: []const i64, out: []u16, lo: usize, hi: usize) !void { + const dh = g.hash.dh; + for (lo..hi) |i| { + const loc = try g.hash.locate(ids[i]); + try embed.dequant(g.w[loc.shard], g.s[loc.shard], g.b[loc.shard], dh, loc.local, out[i * dh ..][0..dh]); } } - const host_xs = try gpa.alloc(f32, heads * dh / 32); - defer gpa.free(host_xs); - sumGroups(host, host_xs); - - var w_b = try cuda.DeviceBuffer.fromHost(driver, words); - defer w_b.free(); - var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); - defer s_b.free(); - var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); - defer b_b.free(); - var out_b = try cuda.DeviceBuffer.alloc(driver, host.len * 2); - defer out_b.free(); - var xs_b = try cuda.DeviceBuffer.alloc(driver, host_xs.len * 4); - defer xs_b.free(); - try tri.pleEmbed(w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, xs_b.ptr, 1, heads, dh, scale); - const got_bytes = try gpa.alloc(u8, host.len * 2); - defer gpa.free(got_bytes); - const xs_bytes = try gpa.alloc(u8, host_xs.len * 4); - defer gpa.free(xs_bytes); - try stream.synchronize(); - try out_b.download(0, got_bytes); - try xs_b.download(0, xs_bytes); - const gpu = try gpa.alloc(u16, host.len); - defer gpa.free(gpu); - for (gpu, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); - var off: usize = 0; - var steps: u32 = 0; - var at: usize = 0; - for (gpu, host, 0..) |a, b, i| { - const d = hc.mixedSteps(a, b); - if (d != 0) off += 1; - if (d > steps) { - steps = d; - at = i; - } + + fn worker(g: *const Ngram, ids: []const i64, out: []u16, lo: usize, hi: usize, failed: *?anyerror) void { + g.fill(ids, out, lo, hi) catch |err| { + failed.* = err; + }; } - var xs_ulp: u32 = 0; - var xs_at: usize = 0; - for (host_xs, 0..) |want, i| { - const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); - const dist = hc.ulps(got, want); - if (dist > xs_ulp) { - xs_ulp = dist; - xs_at = i; + + /// `embedding` for each token after `history` (moved past the chunk) into out and xs, read by up to 128 threads. + pub fn rows(g: *const Ngram, history: []i64, tokens: []const u32, out: []u16, xs: []f32) !void { + const heads = g.hash.heads; + const w = g.width(); + if (history.len != g.hash.n - 1 or out.len != tokens.len * w or xs.len != out.len / 32) return error.UnexpectedTensor; + const items = tokens.len * heads; + const ids = try g.gpa.alloc(i64, items); + defer g.gpa.free(ids); + var hist: [max_n]i64 = undefined; + @memcpy(hist[0..history.len], history); + for (tokens, 0..) |token, t| { + try lastIds(g.hash, hist[0..history.len], @intCast(token), ids[t * heads ..][0..heads]); + for (0..history.len - 1) |i| hist[i] = hist[i + 1]; + hist[history.len - 1] = token; + } + const max_threads = 128; + const n_threads = @min(max_threads, items / 4); + if (n_threads <= 1) { + try g.fill(ids, out, 0, items); + } else { + var threads: [max_threads]std.Thread = undefined; + var failed: [max_threads]?anyerror = @splat(null); + var spawned: usize = 0; + defer for (threads[0..spawned]) |th| th.join(); + const per = (items + n_threads - 1) / n_threads; + for (0..n_threads) |i| { + const lo = @min(items, i * per); + const hi = @min(items, lo + per); + threads[i] = try std.Thread.spawn(.{}, worker, .{ g, ids, out, lo, hi, &failed[i] }); + spawned += 1; + } + for (threads[0..spawned]) |th| th.join(); + spawned = 0; + for (failed[0..n_threads]) |f| if (f) |err| return err; } + for (0..tokens.len) |t| { + const row = out[t * w ..][0..w]; + if (g.scale != 1) { + for (row) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * g.scale); + } + } + sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); + } + @memcpy(history, hist[0..history.len]); } - std.debug.print("ngram token {d} heads {d} values {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} xs_ulp {d} xs_at {d} head {x:0>4}\n", .{ token, heads, host.len, off, steps, at, gpu[at], host[at], xs_ulp, xs_at, gpu[0] }); - if (off != 0 or steps != 0 or xs_ulp != 0) return 1; - return 0; -} +}; test "a fresh window hashes token 7 into sixteen rows" { const h = try Hash.init(.{ diff --git a/zig/src/families/flashnext/cuda_ple_check.zig b/zig/src/families/flashnext/cuda_ple_check.zig new file mode 100644 index 0000000000..b31417259a --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple_check.zig @@ -0,0 +1,235 @@ +//! The n-gram table's host checks: one token's rows dequantized on the host (the row path) and against the GPU gather. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const embed = @import("cuda_embed.zig"); +const hc = @import("cuda_hc.zig"); +const ple = @import("cuda_ple.zig"); + +const max_n = ple.max_n; +const max_heads = ple.max_heads; +const layer = ple.layer; +const Hash = ple.Hash; +const Open = ple.Open; +const Tensor = core.safetensors.Tensor; +const lastIds = ple.lastIds; +const toBf16 = ple.toBf16; +const sumGroups = ple.sumGroups; +const readI64 = ple.readI64; +const weightName = ple.weightName; +const partName = ple.partName; +const findOpen = ple.findOpen; +const copyRow = ple.copyRow; +const tableScale = ple.tableScale; +const readSpec = ple.readSpec; + +/// Host dequant of one token's sixteen n-gram rows. `history` is the previous n - 1 ids. +pub fn embedding(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, history: []const i64, token: u32, out: []u16, xs: []f32) !void { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + if (history.len != hash.n - 1 or out.len != hash.heads * hash.dh or xs.len != out.len / 32) return error.UnexpectedTensor; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history, @intCast(token), ids[0..hash.heads]); + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + const scale = try tableScale(gpa, io, model_dir, wm); + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, out[h * dh ..][0..dh]); + if (scale != 1) { + for (out) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + sumGroups(out, xs); +} + +/// One decode token. `ple_layer_ids` `[2]` is this layer. Returns 0 when the gathered rows match the host dequant. +pub fn table(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, token: u32) !u32 { + const cfg = try readSpec(gpa, io, model_dir); + const hash = try Hash.init(.{ + .vocab = cfg.vocab, + .ngram = cfg.ngram, + .per = cfg.per, + .base = cfg.base, + .divisor = cfg.divisor, + .shards = cfg.shards, + .seed = cfg.seed, + .eos = cfg.eos, + .ple_dim = cfg.ple_dim, + .ple_index = cfg.ple_index, + }); + var history: [max_n]i64 = undefined; + for (0..hash.n - 1) |i| history[i] = hash.eos; + var ids: [max_heads]i64 = undefined; + try lastIds(hash, history[0 .. hash.n - 1], token, ids[0..hash.heads]); + + const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); + defer gpa.free(index_path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + defer parsed.deinit(); + const wm = parsed.value.object.get("weight_map").?.object; + + var cname: [160]u8 = undefined; + const mult_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.layer_multipliers", .{layer}); + const const_file = wm.get(mult_name) orelse return error.MissingTensor; + const const_path = try std.fs.path.join(gpa, &.{ model_dir, const_file.string }); + defer gpa.free(const_path); + var constants = try Open.load(gpa, io, const_path); + defer constants.close(gpa); + var got_m: [max_n]i64 = undefined; + var got_o: [max_heads]i64 = undefined; + var got_s: [max_heads]i64 = undefined; + try readI64(try constants.tensor(mult_name, .i64), got_m[0..hash.n]); + const off_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_offsets", .{layer}); + try readI64(try constants.tensor(off_name, .i64), got_o[0..hash.heads]); + const size_name = try std.fmt.bufPrint(&cname, "language_model.model.layers.{d}.ple.ple_embedding.ngram_heads_vocab_sizes", .{layer}); + try readI64(try constants.tensor(size_name, .i64), got_s[0..hash.heads]); + if (!std.mem.eql(i64, got_m[0..hash.n], hash.mult[0..hash.n]) or !std.mem.eql(i64, got_o[0..hash.heads], hash.offsets[0..hash.heads]) or !std.mem.eql(i64, got_s[0..hash.heads], hash.sizes[0..hash.heads])) { + std.debug.print("ngram constants differ\n", .{}); + return 1; + } + const scale = try tableScale(gpa, io, model_dir, wm); + + const heads = hash.heads; + const dh = hash.dh; + const words = try gpa.alloc(u8, heads * (dh / 8) * 4); + defer gpa.free(words); + const scales = try gpa.alloc(u8, heads * (dh / 32) * 2); + defer gpa.free(scales); + const biases = try gpa.alloc(u8, scales.len); + defer gpa.free(biases); + var opened: [max_heads]Open = undefined; + var nopen: usize = 0; + defer for (opened[0..nopen]) |*o| o.close(gpa); + var key: [180]u8 = undefined; + for (0..heads) |h| { + const loc = try hash.locate(ids[h]); + var nested = false; + const mapped = if (wm.get(try weightName(&key, loc.shard, false))) |v| + v.string + else blk: { + nested = true; + break :blk (wm.get(try weightName(&key, loc.shard, true)) orelse return error.MissingTensor).string; + }; + const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); + defer gpa.free(path); + const file = findOpen(opened[0..nopen], path) orelse file: { + opened[nopen] = try Open.load(gpa, io, path); + nopen += 1; + break :file &opened[nopen - 1]; + }; + const wb = (dh / 8) * 4; + const gb = (dh / 32) * 2; + try copyRow(file, loc.shard, nested, loc.local, dh, words[h * wb ..][0..wb], scales[h * gb ..][0..gb], biases[h * gb ..][0..gb]); + } + + const host = try gpa.alloc(u16, heads * dh); + defer gpa.free(host); + for (0..heads) |h| try embed.dequant(words, scales, biases, dh, h, host[h * dh ..][0..dh]); + if (scale != 1) { + for (host) |*v| { + const f: f32 = @bitCast(@as(u32, v.*) << 16); + v.* = toBf16(f * scale); + } + } + const host_xs = try gpa.alloc(f32, heads * dh / 32); + defer gpa.free(host_xs); + sumGroups(host, host_xs); + + var w_b = try cuda.DeviceBuffer.fromHost(driver, words); + defer w_b.free(); + var s_b = try cuda.DeviceBuffer.fromHost(driver, scales); + defer s_b.free(); + var b_b = try cuda.DeviceBuffer.fromHost(driver, biases); + defer b_b.free(); + var out_b = try cuda.DeviceBuffer.alloc(driver, host.len * 2); + defer out_b.free(); + var xs_b = try cuda.DeviceBuffer.alloc(driver, host_xs.len * 4); + defer xs_b.free(); + try tri.pleEmbed(w_b.ptr, s_b.ptr, b_b.ptr, out_b.ptr, xs_b.ptr, 1, heads, dh, scale); + const got_bytes = try gpa.alloc(u8, host.len * 2); + defer gpa.free(got_bytes); + const xs_bytes = try gpa.alloc(u8, host_xs.len * 4); + defer gpa.free(xs_bytes); + try stream.synchronize(); + try out_b.download(0, got_bytes); + try xs_b.download(0, xs_bytes); + const gpu = try gpa.alloc(u16, host.len); + defer gpa.free(gpu); + for (gpu, 0..) |*o, i| o.* = std.mem.readInt(u16, got_bytes[2 * i ..][0..2], .little); + var off: usize = 0; + var steps: u32 = 0; + var at: usize = 0; + for (gpu, host, 0..) |a, b, i| { + const d = hc.mixedSteps(a, b); + if (d != 0) off += 1; + if (d > steps) { + steps = d; + at = i; + } + } + var xs_ulp: u32 = 0; + var xs_at: usize = 0; + for (host_xs, 0..) |want, i| { + const got: f32 = @bitCast(std.mem.readInt(u32, xs_bytes[4 * i ..][0..4], .little)); + const dist = hc.ulps(got, want); + if (dist > xs_ulp) { + xs_ulp = dist; + xs_at = i; + } + } + std.debug.print("ngram token {d} heads {d} values {d} off {d} max_steps {d} at {d} got {x:0>4} host {x:0>4} xs_ulp {d} xs_at {d} head {x:0>4}\n", .{ token, heads, host.len, off, steps, at, gpu[at], host[at], xs_ulp, xs_at, gpu[0] }); + if (off != 0 or steps != 0 or xs_ulp != 0) return 1; + return 0; +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index d248885d00..d0b96fafdd 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -1,14 +1,13 @@ -//! One prompt, drafts off: each token is one fused decode row, layers in order, then the argmax. +//! The loaded engine and one stream's caches: prompts and windows on the batched pass, decode rounds over streams. const std = @import("std"); const cuda = @import("cuda"); const embed = @import("cuda_embed.zig"); -const qmm = @import("cuda_qmm.zig"); const gdn = @import("cuda_gdn.zig"); const ple = @import("cuda_ple.zig"); const weights = @import("cuda_prompt_wt.zig"); -const step = @import("cuda_prompt_step.zig"); const load = @import("cuda_prompt_load.zig"); +pub const batch = @import("cuda_batch.zig"); const layers_n = weights.layers_n; const linear_n = weights.linear_n; @@ -43,13 +42,6 @@ const uploadQ4 = weights.uploadQ4; const join3 = weights.join3; const loadHc = weights.loadHc; const stitch = weights.stitch; -const Scratch = step.Scratch; -const writeback = step.writeback; -const readout = step.readout; -const gdnStep = step.gdnStep; -const attnStep = step.attnStep; -const moeStep = step.moeStep; -const pleStep = step.pleStep; const Layer = load.Layer; const loadLayer = load.loadLayer; @@ -67,90 +59,6 @@ pub fn argmax(raw: []const u8) u32 { return @intCast(at); } -/// The five best scores of one row, best first, as "id:score" pairs on stderr (the capture's top5 field). -fn printTop(raw: []const u8) void { - var ids: [5]usize = @splat(0); - var vals: [5]f32 = @splat(-std.math.inf(f32)); - for (0..vocab) |i| { - const v = promote(std.mem.readInt(u16, raw[2 * i ..][0..2], .little)); - if (v <= vals[4]) continue; - var at: usize = 4; - while (at > 0 and v > vals[at - 1]) : (at -= 1) { - vals[at] = vals[at - 1]; - ids[at] = ids[at - 1]; - } - vals[at] = v; - ids[at] = i; - } - std.debug.print("top5", .{}); - for (ids, vals) |id, v| std.debug.print(" {d}:{d:.4}", .{ id, v }); - std.debug.print("\n", .{}); -} - -const Pass = struct { - gpa: std.mem.Allocator, - io: std.Io, - driver: *cuda.Driver, - stream: *cuda.Stream, - store: *Store, - model_dir: []const u8, - s: *Scratch, - table: embed.Table, - rows: usize, - mode: []u8, - history: *[2]i64, - conv: []cuda.DeviceBuffer, - state: []cuda.DeviceBuffer, - kc: []cuda.DeviceBuffer, - vc: []cuda.DeviceBuffer, - ikc: []cuda.DeviceBuffer, - tail: cuda.DeviceBuffer, - mixer: HcW, - head_q: Q4, - logits: []u8, - plan_fn: cuda.Function, - up_fn: cuda.Function, - down_fn: cuda.Function, -}; - -fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { - const id: i32 = @intCast(token); - // stream-ordered: a synchronous copy does not wait for the non-blocking stream, so a later row's id - // would overwrite this one before its queued _embed reads it - try p.s.ids.uploadAsync(0, std.mem.asBytes(&id), p.stream.handle); - try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); -} - -fn runLayer(comptime Tri: type, p: *Pass, tri: Tri, lw: *Layer, layer: usize, ids: []const u32, pos0: usize) !void { - for (ids, 0..) |token, t| { - const h = p.s.h.ptr + @as(u64, t) * wide * 2; - const inj = p.s.inj.ptr + @as(u64, t) * inj_stride; - const y = p.s.y.ptr + @as(u64, t) * slots * dims * 4; - const wt = p.s.wts.ptr + @as(u64, t) * wts_stride; - if (layer == ple_layer and p.mode[t] == 2) { - try writeback(tri, p.s, h, 2, inj, y, wt); - p.mode[t] = 0; - } - if (layer == ple_layer) try pleStep(Tri, p.gpa, p.io, p.driver, p.stream, tri, p.s, p.model_dir, p.history, token, h, lw.key, lw.val, lw.nk.ptr, lw.nq.ptr, lw.nc.ptr, lw.ple_cw.ptr, &p.tail); - if (p.mode[t] == 2) { - try writeback(tri, p.s, h, 2, inj, y, wt); - p.mode[t] = 0; - } else try writeback(tri, p.s, h, 0, inj, y, wt); - try readout(tri, p.s, lw.attn_hc, h, inj); - if (lw.linear) { - const li = linearIndex(layer); - try gdnStep(p.gpa, p.driver, p.stream, p.s, lw.proj, lw.out, p.conv[li], lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, lw.norm.ptr, &p.state[li]); - } else { - const ai = attnIndex(layer); - try attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); - } - try writeback(tri, p.s, h, 4, inj, y, wt); - try readout(tri, p.s, lw.mlp_hc, h, inj); - try moeStep(Tri, tri, p.s, p.stream, p.plan_fn, p.up_fn, p.down_fn, lw.router.ptr, lw.experts_up.ptr, lw.experts_down.ptr, y, wt); - p.mode[t] = 2; - } -} - fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![]cuda.DeviceBuffer { const out = try gpa.alloc(cuda.DeviceBuffer, n); var i: usize = 0; @@ -165,40 +73,26 @@ fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![] return out; } -fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !void { - if (p.mode[row] != 2) return error.UnexpectedTensor; - const h = p.s.h.ptr + @as(u64, row) * wide * 2; - const inj = p.s.inj.ptr + @as(u64, row) * inj_stride; - const y = p.s.y.ptr + @as(u64, row) * slots * dims * 4; - const wt = p.s.wts.ptr + @as(u64, row) * wts_stride; - try writeback(tri, p.s, h, 2, inj, y, wt); - p.mode[row] = 0; - try readout(tri, p.s, p.mixer, h, inj); - try qmm.matmul(p.driver, p.stream.*, p.s.mixed.ptr, p.s.xsm.ptr, p.head_q.w.ptr, p.head_q.s.ptr, p.head_q.b.ptr, p.s.logits.ptr, 1, vocab, dims); - try p.stream.synchronize(); - try p.s.logits.download(0, p.logits); -} - -fn forward(comptime Tri: type, p: *Pass, tri: Tri, layers: []Layer, ids: []const u32, pos0: usize) !void { - @memset(p.mode, 0); - for (ids, 0..) |token, t| try embedRow(Tri, p, tri, token, p.s.h.ptr + @as(u64, t) * wide * 2); - for (layers, 0..) |*lw, layer| try runLayer(Tri, p, tri, lw, layer, ids, pos0); - try head(Tri, p, tri, ids.len - 1); -} - pub const vocabulary: usize = vocab; -/// Prompt plus reply tokens a sequence can hold: the attention kernels' NCH chunks of 512 keys. -pub const max_tokens: usize = weights.nch * 512; +/// Prompt plus reply tokens a sequence can hold: the model's window. +pub const max_tokens: usize = weights.window; -/// Device bytes `Engine.newSeq(cap)` allocates: the convolution tails, the recurrent states, the key and value -/// caches with the indexer keys, and the n-gram tail. +/// One position's keys and values in one attention layer, bf16. +pub const kv_row_bytes: usize = kv_heads * head_dim * 2; + +/// Device bytes `Engine.newSeq(cap)` allocates: conv tails, recurrent states, KV with indexer keys, the n-gram tail. pub fn seqBytes(cap: usize) usize { const linear = linear_n * ((gdn.taps - 1) * gdn.conv_dim * 2 + state_n * 4); - const attn = attn_n * cap * (2 * kv_heads * head_dim * 2 + index_dim * 2); + const attn = attn_n * (cap * (2 * kv_row_bytes + index_dim * 2) + pooledBytes(cap)); return linear + attn + ple_tail_n * wide * 2; } -/// One loaded checkpoint: the embedding table, the mixer, the head, and a scratch row the forwards share. +/// One attention layer's pooled indexer keys for a cache of `cap` tokens: a bf16 row of 128 a block of 4. +pub fn pooledBytes(cap: usize) usize { + return (cap + weights.index_ratio - 1) / weights.index_ratio * index_dim * 2; +} + +/// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. pub const Engine = struct { gpa: std.mem.Allocator, io: std.Io, @@ -209,16 +103,18 @@ pub const Engine = struct { table: embed.Table, mixer: HcW, head_q: Q4, - scratch: Scratch, - scratch_rows: usize, - mode: []u8, eos: i64, - experts: cuda.Module, - plan_fn: cuda.Function, - up_fn: cuda.Function, - down_fn: cuda.Function, layers: [layers_n]Layer, n_layers: usize, + batch_k: batch.Kernels, + /// The n-gram table, open for the engine's life (the batched pass gathers its rows from it). + ngram: ple.Ngram, + /// The pass's scratch, grown to the widest chunk run. + batch_s: ?batch.Scratch = null, + /// Block scores for the sparse rows of an attention block, grown to the longest sequence's need. + scores: ?cuda.DeviceBuffer = null, + /// Rows a batched chunk takes at most (TF_FLASHNEXT_CHUNK; batch.default_chunk by default). + chunk: usize, pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; @@ -240,12 +136,6 @@ pub const Engine = struct { e.table = try embed.Table.upload(driver, weight, scales, biases); errdefer e.table.deinit(); e.eos = try ple.eosOf(gpa, io, e.model_dir); - e.scratch = try Scratch.alloc(driver, 1); - errdefer e.scratch.free(); - try e.scratch.dummy.fill8(0, null); - e.scratch_rows = 1; - e.mode = try gpa.alloc(u8, 1); - errdefer gpa.free(e.mode); e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); errdefer e.mixer.free(); var buf: [80]u8 = undefined; @@ -256,11 +146,6 @@ pub const Engine = struct { if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; e.head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); errdefer e.head_q.free(); - e.experts = try cuda.Module.load(driver, cuda.kernels.experts); - errdefer e.experts.unload(); - e.plan_fn = try e.experts.function(weights.plan_symbol); - e.up_fn = try e.experts.function(weights.up_symbol); - e.down_fn = try e.experts.function(weights.down_symbol); var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); defer pack_mod.unload(); const pack_fn = try pack_mod.function(weights.pack_symbol); @@ -270,35 +155,31 @@ pub const Engine = struct { } e.store.release(); e.store.closeFile(); + e.ngram = try ple.Ngram.init(gpa, io, e.model_dir); + errdefer e.ngram.deinit(); + e.batch_k = try batch.Kernels.load(driver); + e.batch_s = null; + e.scores = null; + e.chunk = batch.default_chunk; + if (std.c.getenv("TF_FLASHNEXT_CHUNK")) |v| e.chunk = std.fmt.parseInt(usize, std.mem.span(v), 10) catch batch.default_chunk; + if (e.chunk == 0) e.chunk = batch.default_chunk; return e; } pub fn deinit(e: *Engine) void { - e.experts.unload(); + if (e.batch_s) |*bs| bs.free(e.gpa); + if (e.scores) |*b| b.free(); + e.batch_k.deinit(); + e.ngram.deinit(); for (e.layers[0..e.n_layers]) |*l| l.free(e.gpa); e.head_q.free(); e.mixer.free(); - e.gpa.free(e.mode); - e.scratch.free(); e.table.deinit(); e.store.deinit(); e.gpa.free(e.model_dir); e.gpa.destroy(e); } - fn grow(e: *Engine, rows: usize) !void { - if (rows <= e.scratch_rows) return; - var next = try Scratch.alloc(e.driver, rows); - errdefer next.free(); - try next.dummy.fill8(0, null); - const mode = try e.gpa.alloc(u8, rows); - e.scratch.free(); - e.gpa.free(e.mode); - e.scratch = next; - e.mode = mode; - e.scratch_rows = rows; - } - /// A fresh cache. `cap` is how many tokens it can hold, prompt plus generated. pub fn newSeq(e: *Engine, cap: usize) !Seq { if (cap == 0) return error.UnexpectedTensor; @@ -314,6 +195,8 @@ pub const Engine = struct { errdefer freeOwned(e.gpa, vc); const ikc = try zeros(e.gpa, e.driver, attn_n, cap * index_dim * 2); errdefer freeOwned(e.gpa, ikc); + const pooled = try zeros(e.gpa, e.driver, attn_n, pooledBytes(cap)); + errdefer freeOwned(e.gpa, pooled); var tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); errdefer tail.free(); try tail.fill8(0, null); @@ -323,6 +206,7 @@ pub const Engine = struct { .kc = kc, .vc = vc, .ikc = ikc, + .pooled = pooled, .tail = tail, .history = .{ e.eos, e.eos }, .fed = 0, @@ -333,36 +217,98 @@ pub const Engine = struct { /// Run `ids` on from `seq.fed` and leave the next token's scores in `seq.logits`. pub fn predict(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32) !void { + return e.forward(Tri, tri, seq, ids, .{}); + } + + /// Block scores for a `cap`-token sequence's chunks of up to `rows` rows, reallocated only to grow. + pub fn reserveScores(e: *Engine, cap: usize, rows: usize) !u64 { + const need = batch.scoreBytes(cap, rows); + if (e.scores) |*b| { + if (b.len >= need) return b.ptr; + try e.stream.synchronize(); + b.free(); + e.scores = null; + } + e.scores = try cuda.DeviceBuffer.alloc(e.driver, need); + return e.scores.?.ptr; + } + + /// Scratch for chunks of `rows` rows, reallocated only to grow. + pub fn reserve(e: *Engine, rows: usize) !*batch.Scratch { + if (e.batch_s) |*bs| { + if (bs.cap >= rows) return bs; + try e.stream.synchronize(); + bs.free(e.gpa); + e.batch_s = null; + } + e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, rows); + return &e.batch_s.?; + } + + /// The batched pass over `ids` from `seq.fed` in chunks of at most `chunk` rows; a window's hooks take one chunk. + pub fn forward(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; - try e.grow(ids.len); - var pass = Pass{ + if ((hooks.keep != null or hooks.logits == .all) and ids.len > e.chunk) return error.UnexpectedTensor; + const s = try e.reserve(@min(ids.len, e.chunk)); + const scores = try e.reserveScores(seq.cap, s.cap); + var p: batch.Pass(Tri) = .{ .gpa = e.gpa, .io = e.io, - .driver = e.driver, - .stream = e.stream, - .store = &e.store, - .model_dir = e.model_dir, - .s = &e.scratch, - .table = e.table, - .rows = ids.len, - .mode = e.mode, - .history = &seq.history, - .conv = seq.conv, - .state = seq.state, - .kc = seq.kc, - .vc = seq.vc, - .ikc = seq.ikc, - .tail = seq.tail, - .mixer = e.mixer, - .head_q = e.head_q, - .logits = seq.logits, - .plan_fn = e.plan_fn, - .up_fn = e.up_fn, - .down_fn = e.down_fn, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, + .seq = seq.view(scores), + .hooks = hooks, }; - try forward(Tri, &pass, tri, &e.layers, ids, seq.fed); + var at: usize = 0; + while (at < ids.len) { + const n = @min(s.cap, ids.len - at); + const last = at + n == ids.len; + p.hooks.logits = if (last) hooks.logits else .none; + p.hooks.streams = if (hooks.streams != 0) hooks.streams + at * wide * 2 else 0; + try p.chunk(ids[at..][0..n], seq.fed + at, if (last and hooks.logits != .none) seq.logits else null); + at += n; + } + if (hooks.logits == .none) try e.stream.synchronize(); seq.fed += ids.len; } + + /// One decode round: each of `seqs` takes its token of `ids` in one batched pass; scores as if alone. + pub fn round(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { + if (seqs.len == 0 or seqs.len != ids.len) return error.UnexpectedTensor; + var cap: usize = 0; + for (seqs) |q| { + if (q.fed + 1 > q.cap) return error.UnexpectedTensor; + cap = @max(cap, q.cap); + } + const rows = try e.gpa.alloc(batch.Lane, seqs.len); + defer e.gpa.free(rows); + const s = try e.reserve(seqs.len); + // each row's attention runs alone on the block scores + const scores = try e.reserveScores(cap, 1); + for (rows, seqs) |*l, q| l.* = .{ .seq = q.view(scores), .pos = q.fed }; + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, + .seq = rows[0].seq, + .lanes = rows, + .hooks = .{ .logits = .all, .logits_out = out }, + }; + try p.chunk(ids, 0, null); + try e.stream.synchronize(); + const d = e.stream.d; + for (seqs, 0..) |q, r| { + try d.check(d.api.cuMemcpyDtoH_v2(q.logits.ptr, out + r * vocab * 2, vocab * 2), "cuMemcpyDtoH"); + q.fed += 1; + } + } }; /// One stream's caches. `fed` tokens have already gone through the model. @@ -372,23 +318,31 @@ pub const Seq = struct { kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, + pooled: []cuda.DeviceBuffer, tail: cuda.DeviceBuffer, history: [2]i64, fed: usize, logits: []u8, cap: usize, + /// The caches a batched pass reads and leaves; `scores` the engine's block scores (Engine.reserveScores). + pub fn view(s: *Seq, scores: u64) batch.SeqState { + return .{ .conv = s.conv, .state = s.state, .kc = s.kc, .vc = s.vc, .ikc = s.ikc, .pooled = s.pooled, .tail = s.tail, .history = &s.history, .cap = s.cap, .scores = scores }; + } + pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { freeBufs(s.conv); freeBufs(s.state); freeBufs(s.kc); freeBufs(s.vc); freeBufs(s.ikc); + freeBufs(s.pooled); gpa.free(s.conv); gpa.free(s.state); gpa.free(s.kc); gpa.free(s.vc); gpa.free(s.ikc); + gpa.free(s.pooled); if (s.tail.ptr != 0) s.tail.free(); gpa.free(s.logits); } @@ -403,37 +357,6 @@ fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { gpa.free(bs); } -/// `out` receives `new_tokens` greedy ids. The first is the prompt's argmax. The rest are fed back one at a time. -pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { - if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > max_tokens) return error.UnexpectedTensor; - var e = try Engine.init(gpa, io, driver, stream, model_dir); - defer e.deinit(); - var seq = try e.newSeq(tokens.len + new_tokens); - defer seq.deinit(gpa); - const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); - try e.predict(Tri, tri, &seq, tokens); - const t1 = std.Io.Clock.awake.now(io).toNanoseconds(); - out[0] = argmax(seq.logits); - printTop(seq.logits); - std.debug.print("token {d}\n", .{out[0]}); - var tok = out[0]; - for (1..new_tokens) |i| { - var one = [_]u32{tok}; - try e.predict(Tri, tri, &seq, &one); - tok = argmax(seq.logits); - if (i < 4) printTop(seq.logits); - out[i] = tok; - std.debug.print("token {d}\n", .{tok}); - } - const t2 = std.Io.Clock.awake.now(io).toNanoseconds(); - const ms = struct { - fn of(ns: i96) f64 { - return @as(f64, @floatFromInt(ns)) / 1e6; - } - }.of; - std.debug.print("timing prompt {d} rows {d:.1} ms, decode {d} tokens {d:.2} ms/token\n", .{ tokens.len, ms(t1 - t0), new_tokens - 1, if (new_tokens > 1) ms(t2 - t1) / @as(f64, @floatFromInt(new_tokens - 1)) else 0 }); -} - test "linear layers skip every fourth index and attention starts at 3" { try std.testing.expect(linearLayer(0)); try std.testing.expect(!linearLayer(3)); @@ -445,8 +368,9 @@ test "linear layers skip every fourth index and attention starts at 3" { try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); } -test "a sequence's device bytes grow by the twelve attention layers' keys, values and indexer keys per token" { +test "a sequence's device bytes grow by the twelve attention layers' keys, values, indexer keys and pooled blocks" { const per_token = attn_n * (2 * kv_heads * head_dim * 2 + index_dim * 2); - try std.testing.expectEqual(per_token, seqBytes(2) - seqBytes(1)); + try std.testing.expectEqual(per_token, seqBytes(6) - seqBytes(5)); + try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5) - seqBytes(4)); try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); } diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index f7192ce768..7626925352 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -68,6 +68,8 @@ pub const Layer = struct { q_scale: cuda.DeviceBuffer, k_scale: cuda.DeviceBuffer, i_scale: cuda.DeviceBuffer, + /// The indexer's k_layernorm (fp32), the pooled keys' RMSNorm scale. + ik_scale: cuda.DeviceBuffer, key: Q4, val: Q4, nk: cuda.DeviceBuffer, @@ -92,6 +94,7 @@ pub const Layer = struct { self.q_scale.free(); self.k_scale.free(); self.i_scale.free(); + self.ik_scale.free(); self.key.free(); self.val.free(); self.nk.free(); @@ -190,8 +193,14 @@ fn stackUp(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, g } pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, layer: usize) !Layer { - var buf: [200]u8 = undefined; + var buf: [80]u8 = undefined; const base = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{layer}); + return loadLayerAt(gpa, driver, stream, store, pack_fn, base, linearLayer(layer), layer == ple_layer); +} + +/// The decoder layer under `base` (the MTP head's "language_model.mtp.layers.0": full attention, no n-gram branch). +pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, base: []const u8, linear: bool, with_ple: bool) !Layer { + var buf: [200]u8 = undefined; var base_buf: [80]u8 = undefined; @memcpy(base_buf[0..base.len], base); const root = base_buf[0..base.len]; @@ -219,7 +228,6 @@ pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Str var router_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); errdefer router_b.free(); - const linear = linearLayer(layer); var proj_q = Q4.empty(driver); var out_q = Q4.empty(driver); errdefer proj_q.free(); @@ -231,6 +239,7 @@ pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Str var qs_b = deadBuf(driver); var ks_b = deadBuf(driver); var is_b = deadBuf(driver); + var iks_b = deadBuf(driver); var key_q = Q4.empty(driver); var val_q = Q4.empty(driver); var nk_b = deadBuf(driver); @@ -244,6 +253,7 @@ pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Str errdefer qs_b.free(); errdefer ks_b.free(); errdefer is_b.free(); + errdefer iks_b.free(); errdefer nk_b.free(); errdefer nq_b.free(); errdefer nc_b.free(); @@ -301,12 +311,15 @@ pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Str defer gpa.free(kn); const iqn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); defer gpa.free(iqn); - if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim) return error.UnexpectedTensor; + const ikn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.k_layernorm.weight", .{ap}))); + defer gpa.free(ikn); + if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim or ikn.len != index_dim) return error.UnexpectedTensor; qs_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(qn)); ks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(kn)); is_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(iqn)); + iks_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(ikn)); } - if (layer == ple_layer) { + if (with_ple) { const pre = try std.fmt.bufPrint(&buf, "{s}.ple", .{root}); var pre_buf: [96]u8 = undefined; @memcpy(pre_buf[0..pre.len], pre); @@ -340,6 +353,9 @@ pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Str var shared: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; errdefer for (&shared) |*f| f.free(gpa); + // a folder Living Weights taught holds a wider shared expert (its 640 neurons, then the pocket's) + const width = (try store.hold(try std.fmt.bufPrint(&buf, "{s}.shared_expert.gate_proj.weight", .{mlp_root}))).dim(0); + if (width != moe_w) return error.LearnedPocket; shared[0] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.gate_proj", false, moe_w, dims); shared[1] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.up_proj", false, moe_w, dims); shared[2] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.down_proj", false, dims, moe_w); @@ -377,6 +393,7 @@ pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Str .q_scale = qs_b, .k_scale = ks_b, .i_scale = is_b, + .ik_scale = iks_b, .key = key_q, .val = val_q, .nk = nk_b, diff --git a/zig/src/families/flashnext/cuda_prompt_main.zig b/zig/src/families/flashnext/cuda_prompt_main.zig index bb84cafefc..0815b17d20 100644 --- a/zig/src/families/flashnext/cuda_prompt_main.zig +++ b/zig/src/families/flashnext/cuda_prompt_main.zig @@ -15,7 +15,7 @@ pub fn main(init: std.process.Init) !u8 { return 2; } const count = try std.fmt.parseInt(usize, args[4], 10); - const text = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .limited(1 << 20)); + const text = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .limited(1 << 24)); defer gpa.free(text); var ids: std.ArrayList(u32) = .empty; defer ids.deinit(gpa); @@ -34,7 +34,7 @@ pub fn main(init: std.process.Init) !u8 { const tri: flash.Tri = .{ .set = &set, .s = stream }; const got = try gpa.alloc(u32, count); defer gpa.free(got); - try flash.prompt.run(flash.Tri, gpa, io, &driver, &stream, tri, args[1], ids.items, count, got); + try flash.row_path.run(flash.Tri, gpa, io, &driver, &stream, tri, args[1], ids.items, count, got); std.debug.print("tokens", .{}); for (got) |id| std.debug.print(" {d}", .{id}); std.debug.print("\n", .{}); diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig index 44838b6077..f80b0c2f01 100644 --- a/zig/src/families/flashnext/cuda_prompt_step.zig +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -239,12 +239,12 @@ pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); // stream-ordered, like the ids: the previous row's queued attention kernels still read s.pos try s.pos.uploadAsync(0, std.mem.asBytes(&pos), stream.handle); - try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); - // attention.attention: chunks of 512 keys up to the row's position, at most the scratch's nch + try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); + // attention.attention, dense: chunks of 512 keys up to the row's position (the row-by-row path stays dense) const keys: usize = @as(usize, @intCast(pos)) + 1; - if (keys > nch * 512) return error.UnexpectedTensor; - try tri.attnChunks(s.q.ptr, kc, vc, s.dummy.ptr, s.dummy.ptr, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512); - try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads); + if (keys > wt.dense_nch * 512) return error.UnexpectedTensor; + try tri.attnChunks(s.q.ptr, kc, vc, s.pos.ptr, s.po.ptr, s.pm.ptr, s.pl.ptr, s.dummy.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, (keys + 511) / 512, false); + try tri.attnMerge(s.po.ptr, s.pm.ptr, s.pl.ptr, s.pos.ptr, s.ao.ptr, s.dummy.ptr, s.dummy.ptr, 1, kv_heads, false); try tri.attnGate(s.ao.ptr, s.pa.ptr, s.gated.ptr, s.gxs_attn.ptr, 1, proj_n, q_heads, head_dim); try qmm.partials(driver, stream.*, s.gated.ptr, s.gxs_attn.ptr, out_q.w.ptr, out_q.s.ptr, out_q.b.ptr, s.dummy.ptr, s.part.ptr, 1, dims, out_k, out_sk); } @@ -302,7 +302,7 @@ pub fn pleStep(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: * defer gpa.free(emb); const xs = try gpa.alloc(f32, dims / 32); defer gpa.free(xs); - try ple.embedding(gpa, io, model_dir, history, token, emb, xs); + try @import("cuda_ple_check.zig").embedding(gpa, io, model_dir, history, token, emb, xs); try s.ple_emb.upload(0, std.mem.sliceAsBytes(emb)); try s.ple_xs.upload(0, std.mem.sliceAsBytes(xs)); try qmm.matmul(driver, stream.*, s.ple_emb.ptr, s.ple_xs.ptr, key_q.w.ptr, key_q.s.ptr, key_q.b.ptr, s.ple_keys.ptr, 1, wide, dims); diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index d91216d5c7..127afe2128 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -29,7 +29,18 @@ pub const index_dim: usize = 128; pub const half: usize = 32; pub const proj_n: usize = 13952; pub const out_k: usize = 6144; -pub const nch: usize = 3; +/// Chunks of 512 keys a row's attention reads at most: 2,048 selected keys plus a tail, or 2,051 dense keys. +pub const nch: usize = 5; +pub const dense_nch: usize = 3; +/// The block indexer (QSA): keys pooled in blocks of 4, a row past 512 complete blocks reads its 512 best. +pub const index_ratio: usize = 4; +pub const index_top: usize = 512; +/// A row's key list: the budget's 2,048 keys plus a tail of up to ratio - 1 (attention.AttnScratch.idw). +pub const index_idw: usize = index_top * index_ratio + index_ratio; +/// Keys a dense row reads at most before selection starts: (top + 1) x ratio - 1. +pub const dense_keys: usize = (index_top + 1) * index_ratio - 1; +/// The model's window (config max_position_embeddings). +pub const window: usize = 262144; pub const ple_layer: usize = 1; pub const inj_stride: usize = 16; pub const wts_stride: usize = 48; @@ -47,60 +58,21 @@ const Tensor = core.safetensors.Tensor; const Held = struct { rel: []u8, file: Shard }; +/// One mapped shard: core.safetensors.File. const Shard = struct { + f: core.safetensors.File, io: std.Io, - file: std.Io.File, - map: std.Io.File.MemoryMap, - at: usize, - arena: std.heap.ArenaAllocator, - names: std.StringHashMapUnmanaged(core.safetensors.Entry), fn load(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Shard { - var file = try std.Io.Dir.cwd().openFile(io, path, .{}); - errdefer file.close(io); - const len: usize = @intCast(try file.length(io)); - if (len < 8) return error.BadSafetensors; - var map = try std.Io.File.MemoryMap.create(io, file, .{ .len = len, .protection = .{ .read = true, .write = false }, .populate = false }); - errdefer map.destroy(io); - const header_len: usize = @intCast(std.mem.readInt(u64, map.memory[0..8], .little)); - if (header_len > len - 8) return error.BadSafetensors; - var arena = std.heap.ArenaAllocator.init(gpa); - errdefer arena.deinit(); - var names: std.StringHashMapUnmanaged(core.safetensors.Entry) = .empty; - const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), map.memory[8..][0..header_len], .{}); - defer parsed.deinit(); - var it = parsed.value.object.iterator(); - while (it.next()) |kv| { - if (std.mem.eql(u8, kv.key_ptr.*, "__metadata__")) continue; - const o = kv.value_ptr.object; - const dtype = core.safetensors.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; - const shape = o.get("shape").?.array.items; - if (shape.len > core.safetensors.max_rank) continue; - var e: core.safetensors.Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = 0 }; - var n: usize = dtype.size(); - for (shape, 0..) |d, i| { - e.shape[i] = @intCast(d.integer); - n *= e.shape[i]; - } - const offs = o.get("data_offsets").?.array.items; - e.begin = @intCast(offs[0].integer); - e.end = @intCast(offs[1].integer); - if (e.end < e.begin or e.end - e.begin != n or e.end > len - 8 - header_len) return error.BadSafetensors; - try names.put(arena.allocator(), try arena.allocator().dupe(u8, kv.key_ptr.*), e); - } - return .{ .io = io, .file = file, .map = map, .at = 8 + header_len, .arena = arena, .names = names }; + return .{ .f = try core.safetensors.File.open(gpa, io, path), .io = io }; } fn close(self: *Shard) void { - self.map.destroy(self.io); - self.file.close(self.io); - self.arena.deinit(); - self.* = undefined; + self.f.close(self.io); } fn get(self: *const Shard, name: []const u8) ?Tensor { - const e = self.names.get(name) orelse return null; - return .{ .dtype = e.dtype, .rank = e.rank, .shape = e.shape, .bytes = self.map.memory[self.at + e.begin .. self.at + e.end] }; + return self.f.get(name); } }; @@ -278,8 +250,7 @@ pub fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { return out; } -/// A tensor the engine leaves off the device: the n-gram tables (rows gathered from the mapped shards), the draft -/// head (served without drafts) and the vision tower (text only). +/// A tensor the engine leaves off the device: the n-gram tables, the draft head and the vision tower. pub fn hostOnly(name: []const u8) bool { return std.mem.indexOf(u8, name, ".ple_embedding.ngram_embedding.") != null or std.mem.startsWith(u8, name, "language_model.mtp.") or @@ -302,7 +273,7 @@ pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { defer gpa.free(path); var sh = try Shard.load(gpa, io, path); defer sh.close(); - var ni = sh.names.iterator(); + var ni = sh.f.names.iterator(); while (ni.next()) |e| { if (!hostOnly(e.key_ptr.*)) total += e.value_ptr.end - e.value_ptr.begin; } diff --git a/zig/src/families/flashnext/cuda_rows.zig b/zig/src/families/flashnext/cuda_rows.zig new file mode 100644 index 0000000000..0204205632 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rows.zig @@ -0,0 +1,222 @@ +//! The row-by-row reference: each prompt row one fused decode step, the batched pass's check in flashnext-prompt. + +const std = @import("std"); +const cuda = @import("cuda"); +const embed = @import("cuda_embed.zig"); +const qmm = @import("cuda_qmm.zig"); +const weights = @import("cuda_prompt_wt.zig"); +const step = @import("cuda_prompt_step.zig"); +const load = @import("cuda_prompt_load.zig"); +const prompt = @import("cuda_prompt.zig"); + +const dims = weights.dims; +const streams = weights.streams; +const wide = weights.wide; +const vocab = weights.vocab; +const slots = weights.slots; +const ple_layer = weights.ple_layer; +const inj_stride = weights.inj_stride; +const wts_stride = weights.wts_stride; +const Q4 = weights.Q4; +const HcW = weights.HcW; +const Scratch = step.Scratch; +const Layer = load.Layer; +const promote = weights.promote; + +/// The row path's own scratch and expert kernels, made for a check and dropped after it. +pub const Rows = struct { + gpa: std.mem.Allocator, + driver: *cuda.Driver, + scratch: Scratch, + scratch_rows: usize, + mode: []u8, + experts: cuda.Module, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, + + pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver) !Rows { + var r: Rows = .{ .gpa = gpa, .driver = driver, .scratch = try Scratch.alloc(driver, 1), .scratch_rows = 1, .mode = undefined, .experts = undefined, .plan_fn = undefined, .up_fn = undefined, .down_fn = undefined }; + errdefer r.scratch.free(); + try r.scratch.dummy.fill8(0, null); + r.mode = try gpa.alloc(u8, 1); + errdefer gpa.free(r.mode); + r.experts = try cuda.Module.load(driver, cuda.kernels.experts); + errdefer r.experts.unload(); + r.plan_fn = try r.experts.function(weights.plan_symbol); + r.up_fn = try r.experts.function(weights.up_symbol); + r.down_fn = try r.experts.function(weights.down_symbol); + return r; + } + + pub fn deinit(r: *Rows) void { + r.experts.unload(); + r.gpa.free(r.mode); + r.scratch.free(); + } + + fn grow(r: *Rows, rows: usize) !void { + if (rows <= r.scratch_rows) return; + var next = try Scratch.alloc(r.driver, rows); + errdefer next.free(); + try next.dummy.fill8(0, null); + const mode = try r.gpa.alloc(u8, rows); + r.scratch.free(); + r.gpa.free(r.mode); + r.scratch = next; + r.mode = mode; + r.scratch_rows = rows; + } + + /// Every row of `ids` from seq.fed one fused decode step; the last row's scores land in seq.logits. + pub fn predict(r: *Rows, e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq, ids: []const u32) !void { + if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + // dense attention only: the reference stops where the indexer would start choosing + if (seq.fed + ids.len > weights.dense_nch * 512) return error.UnexpectedTensor; + try r.grow(ids.len); + var pass = Pass{ .gpa = e.gpa, .io = e.io, .driver = e.driver, .stream = e.stream, .model_dir = e.model_dir, .s = &r.scratch, .table = e.table, .mode = r.mode, .history = &seq.history, .conv = seq.conv, .state = seq.state, .kc = seq.kc, .vc = seq.vc, .ikc = seq.ikc, .tail = seq.tail, .mixer = e.mixer, .head_q = e.head_q, .logits = seq.logits, .plan_fn = r.plan_fn, .up_fn = r.up_fn, .down_fn = r.down_fn }; + @memset(pass.mode, 0); + for (ids, 0..) |token, t| try embedRow(Tri, &pass, tri, token, pass.s.h.ptr + @as(u64, t) * wide * 2); + for (e.layers[0..e.n_layers], 0..) |*lw, layer| try runLayer(Tri, &pass, tri, lw, layer, ids, seq.fed); + try head(Tri, &pass, tri, ids.len - 1); + seq.fed += ids.len; + } +}; + +const Pass = struct { + gpa: std.mem.Allocator, + io: std.Io, + driver: *cuda.Driver, + stream: *cuda.Stream, + model_dir: []const u8, + s: *Scratch, + table: embed.Table, + mode: []u8, + history: *[2]i64, + conv: []cuda.DeviceBuffer, + state: []cuda.DeviceBuffer, + kc: []cuda.DeviceBuffer, + vc: []cuda.DeviceBuffer, + ikc: []cuda.DeviceBuffer, + tail: cuda.DeviceBuffer, + mixer: HcW, + head_q: Q4, + logits: []u8, + plan_fn: cuda.Function, + up_fn: cuda.Function, + down_fn: cuda.Function, +}; + +fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { + const id: i32 = @intCast(token); + // on the compute stream: a synchronous copy can land before an earlier row's queued _embed reads its id + try p.s.ids.uploadAsync(0, std.mem.asBytes(&id), p.stream.handle); + try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); +} + +fn runLayer(comptime Tri: type, p: *Pass, tri: Tri, lw: *Layer, layer: usize, ids: []const u32, pos0: usize) !void { + for (ids, 0..) |token, t| { + const h = p.s.h.ptr + @as(u64, t) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, t) * inj_stride; + const y = p.s.y.ptr + @as(u64, t) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, t) * wts_stride; + if (layer == ple_layer and p.mode[t] == 2) { + try step.writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } + if (layer == ple_layer) try step.pleStep(Tri, p.gpa, p.io, p.driver, p.stream, tri, p.s, p.model_dir, p.history, token, h, lw.key, lw.val, lw.nk.ptr, lw.nq.ptr, lw.nc.ptr, lw.ple_cw.ptr, &p.tail); + if (p.mode[t] == 2) { + try step.writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[t] = 0; + } else try step.writeback(tri, p.s, h, 0, inj, y, wt); + try step.readout(tri, p.s, lw.attn_hc, h, inj); + if (lw.linear) { + const li = weights.linearIndex(layer); + try step.gdnStep(p.gpa, p.driver, p.stream, p.s, lw.proj, lw.out, p.conv[li], lw.cw.ptr, lw.a_log.ptr, lw.dt.ptr, lw.norm.ptr, &p.state[li]); + } else { + const ai = weights.attnIndex(layer); + try step.attnStep(Tri, p.driver, p.stream, tri, p.s, lw.proj, lw.out, lw.q_scale.ptr, lw.k_scale.ptr, lw.i_scale.ptr, p.kc[ai].ptr, p.vc[ai].ptr, p.ikc[ai].ptr, @intCast(pos0 + t)); + } + try step.writeback(tri, p.s, h, 4, inj, y, wt); + try step.readout(tri, p.s, lw.mlp_hc, h, inj); + try step.moeStep(Tri, tri, p.s, p.stream, p.plan_fn, p.up_fn, p.down_fn, lw.router.ptr, lw.experts_up.ptr, lw.experts_down.ptr, y, wt); + p.mode[t] = 2; + } +} + +fn head(comptime Tri: type, p: *Pass, tri: Tri, row: usize) !void { + if (p.mode[row] != 2) return error.UnexpectedTensor; + const h = p.s.h.ptr + @as(u64, row) * wide * 2; + const inj = p.s.inj.ptr + @as(u64, row) * inj_stride; + const y = p.s.y.ptr + @as(u64, row) * slots * dims * 4; + const wt = p.s.wts.ptr + @as(u64, row) * wts_stride; + try step.writeback(tri, p.s, h, 2, inj, y, wt); + p.mode[row] = 0; + try step.readout(tri, p.s, p.mixer, h, inj); + try qmm.matmul(p.driver, p.stream.*, p.s.mixed.ptr, p.s.xsm.ptr, p.head_q.w.ptr, p.head_q.s.ptr, p.head_q.b.ptr, p.s.logits.ptr, 1, vocab, dims); + try p.stream.synchronize(); + try p.s.logits.download(0, p.logits); +} + +/// The five best scores of one row, best first, as "id:score" pairs on stderr (the capture's top5 field). +fn printTop(raw: []const u8) void { + var ids: [5]usize = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (0..vocab) |i| { + const v = promote(std.mem.readInt(u16, raw[2 * i ..][0..2], .little)); + if (v <= vals[4]) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = i; + } + std.debug.print("top5", .{}); + for (ids, vals) |id, v| std.debug.print(" {d}:{d:.4}", .{ id, v }); + std.debug.print("\n", .{}); +} + +fn ms(ns: i96) f64 { + return @as(f64, @floatFromInt(ns)) / 1e6; +} + +fn gib(b: u64) f64 { + return @as(f64, @floatFromInt(b)) / (1 << 30); +} + +/// `out` gets `new_tokens` greedy ids: the prompt's argmax, then one fed back at a time (TF_FLASHNEXT_PROMPT=rows). +pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, tokens: []const u32, new_tokens: usize, out: []u32) !void { + if (tokens.len == 0 or new_tokens == 0 or out.len < new_tokens or tokens.len + new_tokens > prompt.max_tokens) return error.UnexpectedTensor; + var e = try prompt.Engine.init(gpa, io, driver, stream, model_dir); + defer e.deinit(); + const by_rows = if (std.c.getenv("TF_FLASHNEXT_PROMPT")) |v| std.mem.eql(u8, std.mem.span(v), "rows") else false; + var rows: ?Rows = if (by_rows) try Rows.init(gpa, driver) else null; + defer if (rows) |*r| r.deinit(); + var seq = try e.newSeq(tokens.len + new_tokens); + defer seq.deinit(gpa); + const loaded = cuda.usage(true); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + if (rows) |*r| try r.predict(e, Tri, tri, &seq, tokens) else try e.predict(Tri, tri, &seq, tokens); + const t1 = std.Io.Clock.awake.now(io).toNanoseconds(); + out[0] = prompt.argmax(seq.logits); + printTop(seq.logits); + // TF_FLASHNEXT_LOGITS=path: the prompt's last row of scores (bf16), to compare the two paths byte for byte + if (std.c.getenv("TF_FLASHNEXT_LOGITS")) |path| try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = std.mem.span(path), .data = seq.logits }); + std.debug.print("token {d}\n", .{out[0]}); + var tok = out[0]; + for (1..new_tokens) |i| { + var one = [_]u32{tok}; + if (rows) |*r| try r.predict(e, Tri, tri, &seq, &one) else try e.predict(Tri, tri, &seq, &one); + tok = prompt.argmax(seq.logits); + if (i < 4) printTop(seq.logits); + out[i] = tok; + std.debug.print("token {d}\n", .{tok}); + } + const t2 = std.Io.Clock.awake.now(io).toNanoseconds(); + const per = if (new_tokens > 1) ms(t2 - t1) / @as(f64, @floatFromInt(new_tokens - 1)) else 0; + std.debug.print("timing prompt {d} rows {d:.1} ms, decode {d} tokens {d:.2} ms/token\n", .{ tokens.len, ms(t1 - t0), new_tokens - 1, per }); + const used = cuda.usage(false); + std.debug.print("memory: {d:.2} GiB on the device after the load and the sequence, peak {d:.2} GiB during the prompt and decode (+{d:.2} GiB)\n", .{ gib(loaded.device), gib(used.peak), gib(used.peak - loaded.device) }); +} diff --git a/zig/src/families/flashnext/cuda_select_main.zig b/zig/src/families/flashnext/cuda_select_main.zig new file mode 100644 index 0000000000..80b69a4d6e --- /dev/null +++ b/zig/src/families/flashnext/cuda_select_main.zig @@ -0,0 +1,113 @@ +//! fn_qsa_select (flashnext_batch.cu) against the lists Triton's _select and _select_tiles made (capture_sparse.sh). + +const std = @import("std"); +const cuda = @import("cuda"); + +const Case = struct { name: []const u8, rows: usize, pos0: i32, nb: usize, idw: usize }; + +fn readAll(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8, ext: []const u8, len: usize) ![]u8 { + const path = try std.fmt.allocPrint(gpa, "{s}/{s}.{s}", .{ dir, name, ext }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 30)); + if (bytes.len != len) return error.BadFixture; + return bytes; +} + +fn int(bytes: []const u8, i: usize) i32 { + return std.mem.readInt(i32, bytes[4 * i ..][0..4], .little); +} + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("usage: flashnext-select \n", .{}); + return 2; + } + const dir = args[1]; + const index = try std.fmt.allocPrint(gpa, "{s}/cases.json", .{dir}); + defer gpa.free(index); + const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 20)); + defer gpa.free(text); + const parsed = try std.json.parseFromSlice([]Case, gpa, text, .{ .ignore_unknown_fields = true }); + defer parsed.deinit(); + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + var module = try cuda.Module.load(&driver, cuda.kernels.flashnext_batch); + defer module.unload(); + const select = try module.function("fn_qsa_select"); + + var failed = false; + for (parsed.value) |c| { + const scores = try readAll(gpa, io, dir, c.name, "scores", c.rows * c.nb * 4); + defer gpa.free(scores); + const want_ids = try readAll(gpa, io, dir, c.name, "ids", c.rows * c.idw * 4); + defer gpa.free(want_ids); + const want_nk = try readAll(gpa, io, dir, c.name, "nk", c.rows * 4); + defer gpa.free(want_nk); + const want_sp = try readAll(gpa, io, dir, c.name, "sparse", c.rows * 4); + defer gpa.free(want_sp); + + var sc_b = try cuda.DeviceBuffer.fromHost(&driver, scores); + defer sc_b.free(); + var pos_b = try cuda.DeviceBuffer.fromHost(&driver, std.mem.asBytes(&c.pos0)); + defer pos_b.free(); + var ids_b = try cuda.DeviceBuffer.alloc(&driver, c.rows * c.idw * 4); + defer ids_b.free(); + try ids_b.fill8(0xFF, null); + var nk_b = try cuda.DeviceBuffer.alloc(&driver, c.rows * 4); + defer nk_b.free(); + var sp_b = try cuda.DeviceBuffer.alloc(&driver, c.rows * 4); + defer sp_b.free(); + var a: cuda.Args = .{}; + a.add(sc_b.ptr); + a.add(pos_b.ptr); + a.add(ids_b.ptr); + a.add(nk_b.ptr); + a.add(sp_b.ptr); + a.add(@as(c_int, @intCast(c.nb))); + a.add(@as(c_int, @intCast(c.idw))); + a.add(@as(c_int, 512)); + a.add(@as(c_int, 4)); + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + try cuda.launch.launch(select, .{ .grid = .{ .x = @intCast(c.rows) }, .block = .{ .x = 1024 } }, stream, &a); + try stream.synchronize(); + const us = @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - t0)) / 1e3; + const got_ids = try gpa.alloc(u8, c.rows * c.idw * 4); + defer gpa.free(got_ids); + const got_nk = try gpa.alloc(u8, c.rows * 4); + defer gpa.free(got_nk); + const got_sp = try gpa.alloc(u8, c.rows * 4); + defer gpa.free(got_sp); + try ids_b.download(0, got_ids); + try nk_b.download(0, got_nk); + try sp_b.download(0, got_sp); + var bad: usize = 0; + var sparse_rows: usize = 0; + for (0..c.rows) |r| { + const nk = int(want_nk, r); + if (int(got_nk, r) != nk or int(got_sp, r) != int(want_sp, r)) { + bad += 1; + continue; + } + if (int(want_sp, r) == 0) continue; + sparse_rows += 1; + for (0..@intCast(nk)) |i| { + if (int(got_ids, r * c.idw + i) != int(want_ids, r * c.idw + i)) { + bad += 1; + break; + } + } + } + std.debug.print("{s}: {d} rows at {d}, {d} blocks a row, {d} sparse: {d} differ ({d:.0} us)\n", .{ c.name, c.rows, c.pos0, c.nb, sparse_rows, bad, us }); + if (bad != 0) failed = true; + } + std.debug.print("{s}\n", .{if (failed) "differ" else "match"}); + return if (failed) 1 else 0; +} diff --git a/zig/tests/cuda/flashnext/capture.py b/zig/tests/cuda/flashnext/capture.py index a92c58adc7..ad8fa443c4 100644 --- a/zig/tests/cuda/flashnext/capture.py +++ b/zig/tests/cuda/flashnext/capture.py @@ -1,7 +1,5 @@ #!/usr/bin/env python3 -"""Flash Next's Python CUDA engine (qwen4_exp, drafts off, greedy) as the Zig port's oracle: every Triton launch, -prompt ids, and two greedy references per prompt (the served path: chunked prefill then serial decode; the teacher -path: every prompt token as its own decode row, the way zig's flashnext-prompt runs it).""" +"""Flash Next's Python CUDA engine (drafts off, greedy) as the Zig oracle: Triton launches, prompt ids, references.""" from __future__ import annotations @@ -18,6 +16,10 @@ } TOP_STEPS = 4 # sampled steps whose five best logits the teacher paths keep LONG_TOKENS = 600 # the long prompt: module source cut to make the whole prompt about this many tokens +# prompts past the block indexer's 2,048-key budget, cut from pydoc_data.topics to exactly this many tokens +SPARSE_TOKENS = {"long3k": 3001, "long12k": 12001} +SPARSE_ASK = ("\n\nQuestion: according to the excerpt above, what does an assert statement do when Python runs with " + "optimization on, and what does an augmented assignment evaluate only once? Answer briefly.") def chat(user: str) -> str: @@ -35,6 +37,14 @@ def prompts(model: Path) -> dict[str, list[int]]: tail = enc("<|im_end|>\n<|im_start|>assistant\n\n\n\n\n") body = enc(Path(textwrap.__file__).read_text())[:LONG_TOKENS - len(head) - len(tail)] out["long"] = head + body + tail + from pydoc_data.topics import topics + + reference = "\n\n".join(topics[k] for k in sorted(topics)) + head = enc("<|im_start|>user\nBelow is an excerpt of the Python language reference.\n\n") + tail = enc(SPARSE_ASK + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n") + text = enc(reference) + for name, total in SPARSE_TOKENS.items(): + out[name] = head + text[:total - len(head) - len(tail)] + tail return out @@ -62,7 +72,9 @@ def main() -> int: ap.add_argument("--tokens", type=int, default=32, help="greedy tokens per prompt") ap.add_argument("--context", type=int, default=1536, help="the cache window; 1025..1536 gives the port's NCH=3, QSA=0 attention specialization") - ap.add_argument("--only", default="", help="comma list of prompt names (default all)") + ap.add_argument("--only", default="", help="comma list of prompt names (default: every prompt the context fits)") + ap.add_argument("--teachers", default="teacher,teacher32", + help="teacher paths to run after the served one (comma list of teacher, teacher32; empty: none)") a = ap.parse_args() out = Path(a.out) out.mkdir(parents=True, exist_ok=True) @@ -83,6 +95,8 @@ def phase(name: str, detail: bool = False): ids = prompts(model) if a.only: ids = {k: v for k, v in ids.items() if k in a.only.split(",")} + else: + ids = {k: v for k, v in ids.items() if len(v) + a.tokens < a.context} (out / "prompts.json").write_text(json.dumps(ids) + "\n") for name, p in ids.items(): (out / f"{name}.ids").write_text(" ".join(map(str, p)) + "\n") @@ -120,13 +134,13 @@ def phase(name: str, detail: bool = False): results[f"serial/{name}"] = {"tokens": toks, "wall_s": round(wall, 4), **stats} print(f"serial {name} ({len(prompt)} prompt): {toks} {stats}", flush=True) if fp32 is None: - # teacher32: the decode-form experts (fp32 slot outputs, Buffers' default and the MTP head's), the - # arithmetic zig's flashnext-prompt runs; the serial engine's buffers take the prompt-form experts + # teacher32: the decode-form experts (fp32 slot outputs), the arithmetic zig's flashnext-prompt runs fp32 = eng.serial.twin() fp32.buf = Buffers(eng.e.w, eng.e.rows, eng.e.capacity) for label, t in (("teacher", eng.serial), ("teacher32", fp32)): - teach(label, name, prompt, t, toks) + if label in a.teachers.split(","): + teach(label, name, prompt, t, toks) (out / "results.json").write_text(json.dumps({"info": info, "results": results}, indent=1) + "\n") if rec is not None: diff --git a/zig/tests/cuda/flashnext/capture.sh b/zig/tests/cuda/flashnext/capture.sh index 5bb769d5b8..8aaf8b9ccb 100755 --- a/zig/tests/cuda/flashnext/capture.sh +++ b/zig/tests/cuda/flashnext/capture.sh @@ -1,25 +1,22 @@ #!/usr/bin/env bash -# Flash Next (qwen4_exp) capture -> manifest -> kernel set for zig's flashnext-prompt and the native lanes engine. -# Usage: capture.sh WORK SNAPSHOT_DIR [IMAGE] -# WORK a host dir this script owns: WORK/cache (Triton + torch extension caches), WORK/out/cap, WORK/kernels -# SNAPSHOT_DIR the checkpoint's HF snapshot dir (its blobs must sit at ../../blobs, as in an HF cache) -# IMAGE TensorFold 0.6.1 image, default tensorfold-qwen38:v0.6.1 -# The Python engine runs drafts off, greedy, --context 1535 (1536 cache slots: _chunks/_merge NCH=3, QSA=0, the dense -# attention the port launches), --ple-on-ssd, bf16 KV. Needs the GPU free (~78 GiB). +# Usage: capture.sh WORK SNAPSHOT_DIR [IMAGE]; captures Flash Next's launches, references and a kernel set under WORK. set -euo pipefail WORK=$(readlink -f "$1"); SNAP=$(readlink -f "$2"); IMAGE=${3:-tensorfold-qwen38:v0.6.1} REPO=$(cd "$(dirname "$0")/../../../.." && pwd) -HF=$(dirname "$(dirname "$(dirname "$SNAP")")") # .../hub/models--X/snapshots/REV -> .../hub (holds the blobs) +HF=$(dirname "$(dirname "$(dirname "$SNAP")")") REL=${SNAP#"$HF"/} -mkdir -p "$WORK/cache" "$WORK/out/cap" +# CONTEXT=16383 ONLY=long3k,long12k TEACHERS=teacher32 OUT=sparse: the prompts past the indexer's 2,048-key budget +CONTEXT=${CONTEXT:-1535}; ONLY=${ONLY:-}; TEACHERS=${TEACHERS-teacher,teacher32}; OUT=${OUT:-cap} +mkdir -p "$WORK/cache" "$WORK/out/$OUT" +KSET=$WORK/kernels; [ "$OUT" = cap ] || KSET=$WORK/kernels-$OUT +# SNAPSHOT_DIR is the HF snapshot (blobs at ../../blobs); the default context 1535 keeps attention dense (QSA 0) docker run --rm --name fn-capture --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \ --user "$(id -u):$(id -g)" -e HOME=/tmp -e HF_HUB_OFFLINE=1 -e TENSORFOLD_NO_UPDATE_CHECK=1 \ -e TENSORFOLD_MEMORY_RESERVE_GIB=2 \ -v "$HF":/hf:ro -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -w /out \ - "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture.py --model "/hf/$REL" --out /out/cap \ - --tools /src/tools/zig --tokens 32 --context 1535 + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture.py --model "/hf/$REL" --out "/out/$OUT" \ + --tools /src/tools/zig --tokens 32 --context "$CONTEXT" --only "$ONLY" --teachers "$TEACHERS" python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ - --launches "$WORK/out/cap/launches.json" --mount /cache/triton --out "$WORK/out/cap/manifest.json" -python3 -B "$REPO/zig/tests/cuda/nemotron/aot_pack.py" --manifest "$WORK/out/cap/manifest.json" \ - --cache "$WORK/cache/triton" --jit "$WORK/out/cap/jit.json" --out "$WORK/kernels" -echo "kernel set: $WORK/kernels; prompt ids and references: $WORK/out/cap/{short,mid,long}.ids, results.json" + --launches "$WORK/out/$OUT/launches.json" --mount /cache/triton --out "$WORK/out/$OUT/manifest.json" +"$(dirname "$0")/pack.sh" "$KSET" "$WORK/out/$OUT" +echo "prompt ids (*.ids) and references (results.json) in $WORK/out/$OUT" diff --git a/zig/tests/cuda/flashnext/capture_rows.py b/zig/tests/cuda/flashnext/capture_rows.py new file mode 100644 index 0000000000..85930056c3 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_rows.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""The decode kernels Flash Next's batched prompt pass runs at M > 1 (M 17 and 32), no checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", required=True, help="folder of triton_aot_manifest.py") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + from tensorfold.cuda import moe + from tensorfold.families.qwen4_exp.cuda import glue, qmm + + dev = "cuda" + g = torch.Generator(device=dev).manual_seed(0) + streams, dims, low = 4, 2560, 320 + wide = streams * dims + + def q4(n: int, k: int): + w = torch.randint(-2**31, 2**31 - 1, (n, k // 8), dtype=torch.int32, device=dev, generator=g) + s = (torch.rand((n, k // 32), device=dev, generator=g) * 0.01).to(torch.bfloat16) + b = (torch.rand((n, k // 32), device=dev, generator=g) * 0.01 - 0.005).to(torch.bfloat16) + return qmm.make_q4(w, s, b, layout="tiled") + + downs = {n: q4(n, wide) for n in (low + streams, low)} + up = q4(wide, low) + router_rows = (torch.randn((513, dims), device=dev, generator=g) * 0.02).to(torch.bfloat16) + with rec.scope("rows"): + for m in (17, 32): + h = torch.randn((m, wide), device=dev, generator=g).to(torch.bfloat16) + pss = torch.rand((m, dims // 256, streams), device=dev, generator=g) * 100 + scale = torch.rand((wide,), device=dev, generator=g) + normed = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + act = torch.empty((m, low), dtype=torch.bfloat16, device=dev) + xs_act = torch.empty((m, low // 32), dtype=torch.float32, device=dev) + inj = torch.empty((m, streams), dtype=torch.bfloat16, device=dev) + for n, q in downs.items(): + dn = torch.empty((m, n), dtype=torch.bfloat16, device=dev) + part = torch.empty((32 * m * n,), dtype=torch.float32, device=dev) + got = qmm.hc_down(h, pss, scale, normed, q, 1e-6, streams, out=dn, part=part) + glue.hc_reduce_act(got, act, xs_act, inj if n == low + streams else None, streams, low) + mixed = torch.empty((m, dims), dtype=torch.bfloat16, device=dev) + xs_mixed = torch.empty((m, dims // 32), dtype=torch.float32, device=dev) + qmm.hc_upmix(act, xs_act, up, normed, mixed, xs_mixed, streams) + # the up projection alone (a training pass keeps it): qmm.matmul's tiled kernel at SHAPES16 + up_rows = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + qmm._qmm[(-(-m // 16), wide // 64, 1)](act, xs_act, up.weight, up.scales, up.biases, up_rows, up_rows, m, + act.stride(0), N=wide, K=low, SK=1, BM=16, BLOCK_N=64, GPI=2, + F32=False, SBN=64, num_warps=4, num_stages=3) + x = torch.randn((m, dims), device=dev, generator=g).to(torch.bfloat16) + logits = torch.empty((m, 513), dtype=torch.float32, device=dev) + grid = (-(-m // 16), -(-513 // 32)) + moe._router[grid](x, router_rows, logits, m, x.stride(0), D=dims, NE=513, BM=16, BLOCK_E=32, BK=256, + num_warps=4, num_stages=4) + torch.cuda.synchronize() + rec.dump(out / "launches.json") + from triton.runtime.jit import JITFunction + import gc + + spec = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + spec[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + (out / "jit.json").write_text(json.dumps(spec, indent=1, sort_keys=True) + "\n") + print("done", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture_rows.sh b/zig/tests/cuda/flashnext/capture_rows.sh new file mode 100755 index 0000000000..fcbd158b11 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_rows.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Usage: capture_rows.sh WORK [IMAGE]: the batched pass's M > 1 kernel forms (capture_rows.py) into WORK/out. +set -euo pipefail +WORK=$(readlink -f "$1"); IMAGE=${2:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +mkdir -p "$WORK/cache" "$WORK/out" +docker run --rm --name fn-capture-rows --gpus all --ipc=host --user "$(id -u):$(id -g)" -e HOME=/tmp \ + -e TENSORFOLD_NO_UPDATE_CHECK=1 -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture_rows.py --out /out --tools /src/tools/zig +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/launches.json" --mount /cache/triton --out "$WORK/out/manifest.json" diff --git a/zig/tests/cuda/flashnext/capture_sparse.py b/zig/tests/cuda/flashnext/capture_sparse.py new file mode 100644 index 0000000000..dc5540df82 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_sparse.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Flash Next's sparse attention kernels and _select fixtures on synthetic tensors, no checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", required=True, help="folder of triton_aot_manifest.py") + ap.add_argument("--fixtures", default="", help="folder for the selection fixtures (empty: none)") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + import triton + from tensorfold.families.qwen4_exp.cuda import attention as A + from tensorfold.families.qwen4_exp.cuda import glue, kvcache + + dev = "cuda" + g = torch.Generator(device=dev).manual_seed(0) + heads, kv, hd, ih, idim, ratio = 24, 2, 256, 4, 128, 4 + cap = 8192 # nb = 2048: NB % 16 == 0, as the Zig scratch keeps it + inv = (10_000_000.0 ** (-torch.arange(32, dtype=torch.float64) / 32)).to(torch.float32).to(dev) + with rec.scope("sparse"): + for rows in (1, 17, 32): + sc = A.AttnScratch(rows, heads, hd, cap, dev) + assert sc.qsa and sc.nch == 5 and sc.nb % 16 == 0 + cache = kvcache.KVCache(cap, kv, hd, dev, "bf16") + cache.k.copy_(torch.randn(cache.k.shape, device=dev, generator=g).to(torch.bfloat16)) + cache.v.copy_(torch.randn(cache.v.shape, device=dev, generator=g).to(torch.bfloat16)) + ikc = torch.randn((cap, idim), device=dev, generator=g).to(torch.bfloat16) + pooled = torch.zeros((sc.nb, idim), dtype=torch.bfloat16, device=dev) + w = torch.rand((idim,), device=dev, generator=g) + 0.5 + # every block before the window pooled first (one row at a time would do it the same way) + pos0 = torch.zeros((1,), dtype=torch.int32, device=dev) + p0 = 3001 + A.qsa_pool(ikc, pooled, pos0, w, inv, 1e-6, sc, p0 - p0 % 32) if rows != 1 else None + pos0.fill_(p0) + A.qsa_pool(ikc, pooled, pos0, w, inv, 1e-6, sc, rows) + iq = torch.randn((rows, ih, idim), device=dev, generator=g).to(torch.bfloat16) + q = torch.randn((rows, heads, hd), device=dev, generator=g).to(torch.bfloat16) + ends = p0 + rows + blocks = min(sc.nb, max(1, triton.cdiv(ends, ratio))) + A._scores[(rows, triton.cdiv(blocks, 64))](iq, pooled, pos0, sc.scores, sc.nb, HI=ih, DI=idim, + RATIO=ratio, TOP=sc.budget // ratio, BB=64, num_warps=4) + A._launch_select(sc, pos0, rows, blocks) + A.attention(q, cache.k, cache.v, pos0, sc, rows, hd ** -0.5, context=ends, ks=cache.ks, vs=cache.vs) + torch.cuda.synchronize() + rec.dump(out / "launches.json") + from triton.runtime.jit import JITFunction + import gc + + spec = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + spec[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + (out / "jit.json").write_text(json.dumps(spec, indent=1, sort_keys=True) + "\n") + if a.fixtures: + fixtures(Path(a.fixtures), A, torch, triton, g) + print("done", flush=True) + return 0 + + +def fixtures(fx: Path, A, torch, triton, g) -> None: + """Score rows (few distinct values: ties everywhere) and the lists _select / _select_tiles make of them.""" + + fx.mkdir(parents=True, exist_ok=True) + dev = "cuda" + cases = [] + # (capacity, pos0, rows, distinct score values): register select, then rows past SELECT_REGS blocks (tiles) + for i, (cap, p0, rows, levels) in enumerate([(8192, 2040, 24, 7), (8192, 5000, 9, 0), (65536, 40000, 5, 3), + (262144, 150000, 3, 11), (262144, 262000, 2, 0)]): + sc = A.AttnScratch(rows, 24, 256, cap, dev) + if levels: + s = torch.randint(0, levels, (rows, sc.nb), device=dev, generator=g).float() / 7.0 + else: + s = torch.rand((rows, sc.nb), device=dev, generator=g) * 3.0 + s[:, ::97] = 0.0 # zeros among them, as relu sums give + sc.scores.copy_(s) + sc.ids.fill_(-1) + pos0 = torch.full((1,), p0, dtype=torch.int32, device=dev) + blocks = min(sc.nb, max(1, triton.cdiv(p0 + rows, 4))) + A._launch_select(sc, pos0, rows, blocks) + torch.cuda.synchronize() + name = f"case{i}" + sc.scores.cpu().numpy().tofile(fx / f"{name}.scores") + sc.ids.cpu().numpy().tofile(fx / f"{name}.ids") + sc.nk.cpu().numpy().tofile(fx / f"{name}.nk") + sc.sparse.cpu().numpy().tofile(fx / f"{name}.sparse") + cases.append({"name": name, "rows": rows, "pos0": p0, "nb": sc.nb, "idw": sc.idw}) + (fx / "cases.json").write_text(json.dumps(cases, indent=1) + "\n") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture_sparse.sh b/zig/tests/cuda/flashnext/capture_sparse.sh new file mode 100755 index 0000000000..1598002374 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_sparse.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Usage: capture_sparse.sh WORK [IMAGE]: the sparse attention's kernels and selection fixtures (capture_sparse.py). +set -euo pipefail +WORK=$(readlink -f "$1"); IMAGE=${2:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +mkdir -p "$WORK/cache" "$WORK/out" "$WORK/fixtures" +docker run --rm --name fn-capture-sparse --gpus all --ipc=host --user "$(id -u):$(id -g)" -e HOME=/tmp \ + -e TENSORFOLD_NO_UPDATE_CHECK=1 -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -v "$WORK/fixtures":/fixtures -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture_sparse.py --out /out --tools /src/tools/zig --fixtures /fixtures +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/launches.json" --mount /cache/triton --out "$WORK/out/manifest.json" diff --git a/zig/tests/cuda/flashnext/lr7/long12k.ids b/zig/tests/cuda/flashnext/lr7/long12k.ids new file mode 100644 index 0000000000..e922e33cbc --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/long12k.ids @@ -0,0 +1 @@ +248045 846 198 36948 369 449 47801 314 279 12654 3992 5604 13 271 760 328 2138 1 4952 198 865 6610 271 8272 11889 513 264 16219 1560 310 5479 26805 52992 198 17583 264 1957 25 271 256 1992 21152 61987 328 2138 1 7258 498 2129 7258 60 271 760 4145 1304 11 328 2138 7258 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 25 4673 63990 271 760 11244 1304 11 328 2138 7258 16 11 7258 17 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 16 25 4673 63990 46718 17 8 271 9205 42018 2303 9370 421 12880 8090 548 1 321 328 66317 1401 1 8042 198 951 279 5617 3309 7101 440 1782 4874 13 220 733 279 1428 198 13957 11 279 5617 3309 3759 12880 8090 548 1 369 328 2434 1 1172 198 7994 12779 11 328 3913 1 948 24460 369 10897 318 5454 198 1021 2904 6320 46 1764 220 561 1428 1970 13427 70307 874 1970 364 449 198 2138 4952 948 24460 369 10897 506 18612 854 13 220 6817 198 8761 424 369 24366 310 2830 279 2450 1970 364 279 7258 198 8761 4490 303 279 1412 1876 26 424 668 381 12234 430 919 314 279 198 7450 11319 13 271 27991 1320 310 12880 8090 548 1 513 11477 13 220 561 869 364 279 5617 3309 198 9521 369 10521 948 279 37999 8211 13 1358 40238 11889 198 865 17731 271 40238 11889 513 1429 310 318 265 8 7423 4874 310 2663 321 310 198 32925 7944 466 3470 314 24983 5978 25 271 256 15840 21152 61987 318 5480 1951 35036 7030 318 11531 1114 27160 735 7302 27160 8 198 256 2100 1951 257 61987 2100 318 2129 2100 4653 498 487 1293 198 256 2100 670 61987 12448 198 1025 735 32467 498 5480 1951 60 56439 198 1025 735 48975 498 5480 1951 60 84576 198 1025 735 16819 330 42514 198 1025 735 14701 198 1025 735 81987 198 1025 735 63669 2100 271 7 9538 3644 34019 5262 364 279 18890 16954 364 348 42648 330 42514 11961 198 9 33261 11961 321 348 82 86047 9 5946 271 2014 15840 4952 64013 279 7258 1103 318 28321 421 198 559 628 381 264 3074 7258 466 264 30644 70221 1103 11 279 14827 198 28726 286 264 13992 8 321 48233 279 3074 12571 1576 310 1754 314 198 1719 2100 11140 11 494 2047 310 1245 13 271 40238 369 4364 51070 11313 383 279 1304 314 279 2100 198 7080 553 3095 264 2100 369 919 314 264 24983 1576 318 276 6816 198 15812 11 14701 466 81987 681 279 24983 1576 1902 198 484 6887 2649 279 15840 321 9973 883 1141 30794 11 321 198 17812 4673 449 4531 413 279 15840 369 41565 13 220 561 5426 198 5308 2683 539 5090 4346 321 279 18987 8953 513 2574 440 279 198 17829 314 279 1576 4346 318 3924 3644 561 5129 913 198 71 32759 553 271 40238 314 449 1576 310 264 2100 1103 11 44007 42360 303 198 39708 37052 466 9057 37641 11 369 51070 4364 430 10695 13 271 9 1368 279 2100 1103 369 264 3074 2100 440 874 26849 30644 11 198 220 44007 303 71458 11 279 1576 369 12244 310 421 2100 13 271 9 18236 25 271 220 348 1368 279 2100 1103 5435 799 2100 78685 440 449 32815 3086 11 198 262 2512 264 1001 11531 1114 828 2100 25 561 1576 1902 381 449 49124 440 506 198 262 3140 430 1599 3470 430 1017 513 11218 303 279 2100 1103 11 26401 198 262 799 13 220 561 1118 3470 314 279 49124 513 12244 11 494 2047 310 198 262 1245 11 310 279 11218 1518 279 56276 2100 13 220 561 1534 3470 198 262 314 279 49124 513 12244 310 279 11218 1238 279 56276 198 262 2100 13 220 357 1103 314 279 9376 3470 303 279 49124 369 1179 198 262 12244 310 279 56276 2100 318 1719 1103 628 381 4147 553 271 220 348 18236 25 561 1576 1902 381 449 49124 440 279 1788 1324 314 3470 198 262 430 1017 513 11218 303 279 2100 1103 11 321 279 3470 513 198 262 12244 11 494 2047 310 1245 11 310 279 11811 11218 13 271 40238 314 449 1576 310 264 3074 2100 369 51070 4364 430 198 18000 82 13 271 9 1368 279 2100 369 449 12448 318 591 1590 271 220 348 1368 279 803 1503 524 11822 303 264 328 9461 1 466 328 6082 2359 1 4952 198 262 303 279 1428 1970 2424 25 279 803 369 6608 310 279 1576 303 279 198 262 1428 2136 4325 13 271 220 348 17673 25 279 803 369 6608 310 279 1576 303 279 3521 4325 198 262 466 279 15484 4325 10521 539 328 6082 2359 487 15119 13 271 220 561 803 369 40158 413 424 557 2582 6608 13 220 1061 1189 5073 279 198 220 5604 1698 364 279 1576 8334 6608 310 279 803 310 5372 198 220 6942 11 13969 279 1576 310 381 3372 38157 321 1141 44462 318 331 424 198 220 682 799 8 310 381 2512 13 271 9 1368 279 2100 369 449 6816 5604 25 561 5839 7258 303 198 220 279 5604 369 24276 13 220 1049 1220 7302 449 1576 440 198 220 86370 7944 26 413 411 369 524 279 1105 11 328 77318 1 369 198 220 8953 13 220 2844 1576 369 1179 4439 310 9501 279 12244 1576 310 198 220 279 2574 6816 26 413 424 4021 2649 279 15840 11 424 24096 198 220 449 4531 318 41547 694 524 13902 328 3776 1401 1764 271 220 6817 25 1368 279 1576 369 264 523 2774 321 279 6816 5604 198 220 13273 383 2107 10744 314 279 15840 5497 11 279 1245 23641 3008 198 220 7258 11 328 64 1926 1 628 2528 2892 449 2774 6816 466 318 331 874 198 220 2774 6816 6513 8 264 523 6816 13 220 561 2047 23641 3008 198 220 2100 328 64 1926 1 369 2590 716 430 449 2774 6816 11 6611 424 413 198 220 5689 13 220 13891 11 279 1330 54133 314 328 64 1926 1 635 524 13902 198 220 8042 310 279 1788 6816 25 413 279 1245 23641 3008 7258 198 220 18675 310 264 523 6816 11 279 2047 23641 3008 11121 264 491 198 220 2774 6816 430 279 2100 314 279 15840 25 271 257 523 2356 82 25 198 260 830 283 220 18 1734 653 523 3759 198 257 1699 283 2356 82 363 198 257 1699 1926 283 1699 1926 478 220 16 256 653 13514 1699 1926 430 220 19 9101 2356 82 1926 430 220 18 271 220 1061 3874 1503 524 13902 3668 310 15933 198 220 7944 11 1680 430 5706 3354 440 328 3573 363 3158 271 9 1368 279 2100 369 264 14701 25 561 5839 7258 303 279 198 220 5604 369 24276 13 220 1049 1220 7302 2892 264 24983 8240 198 220 1576 318 20171 430 264 1103 8 466 264 12366 1576 318 20171 430 264 10675 553 198 220 9019 11 279 94860 7258 369 24276 13 271 220 1368 279 5839 369 264 24983 8240 1576 318 20171 430 264 1103 681 279 198 220 94860 1902 7302 449 7308 13 220 1368 424 369 7968 11 279 8240 725 198 220 2984 369 3568 310 424 13 220 561 12571 869 1902 381 264 2397 40823 198 220 7308 2599 1056 279 8240 725 2984 11 321 279 8240 369 4439 198 220 310 9501 279 12244 1576 310 1141 1455 440 421 1857 13 220 1368 279 198 220 1857 369 680 314 2020 11 328 1498 1401 1 369 8953 318 28973 310 264 198 220 94860 290 8240 4021 884 491 3470 310 264 1103 553 271 220 1368 279 5839 369 264 12366 1576 318 20171 430 264 10675 681 279 198 220 94860 1902 599 264 913 17606 440 279 12366 725 1328 913 11 198 220 321 279 12366 369 1179 4439 310 1791 264 1328 55616 6505 864 13640 198 220 279 94860 310 279 12244 1576 13 220 1061 628 2892 8032 449 198 220 6149 1328 55616 6505 440 279 1788 1328 869 11 466 5479 264 491 198 220 1328 55616 6505 318 331 874 1328 440 279 1788 869 23533 553 271 220 1690 1156 37435 5978 11 279 12880 723 1163 548 33905 1654 369 2512 440 198 220 8053 5789 13 271 9 1368 279 2100 369 264 81987 25 561 5839 7258 303 279 5604 198 220 369 24276 13 220 1049 1220 7302 264 24983 8240 1576 318 20171 430 264 198 220 1103 553 220 561 12244 1576 1220 381 264 8240 1576 314 279 1788 198 220 913 13 220 9019 11 279 4570 321 8156 6608 22666 513 24276 11 198 220 303 90349 781 513 2945 26 16184 513 6942 321 279 8240 725 198 220 2984 13 220 561 13854 1220 14989 310 24959 13 1368 2892 6608 369 198 220 7968 11 279 8240 725 2984 369 3568 310 424 13 220 561 12571 198 220 13854 513 70575 310 9940 1881 6942 321 279 8240 725 2984 11 198 220 27389 13 220 16864 11 279 8240 1576 369 4439 310 8032 279 198 220 15511 440 279 3470 314 279 12244 8240 13 220 561 2984 314 279 198 220 15511 1189 381 2086 494 279 2984 314 279 12244 8240 11 198 220 8190 9722 279 2984 314 279 2100 8240 11 413 279 2100 198 220 8240 5954 424 13 271 332 6901 26033 7873 7472 64700 733 279 1428 7873 11 279 198 54304 364 11218 369 4284 310 381 279 1788 430 364 22666 11 321 198 11469 18890 369 17030 2261 279 1970 9191 9956 11 13969 198 1657 11346 1412 6397 13 271 15337 279 7044 314 15840 23197 421 73406 1881 198 1719 2047 23641 3008 321 279 1245 23641 3008 513 3260 14352 484 17140 515 318 1891 198 8422 328 64 11 292 283 292 11 264 1 70664 1330 7101 681 73406 348 53977 9 279 198 13034 314 12244 4534 7101 11822 2047 4534 6491 11 6807 198 1334 286 303 20684 13 220 1690 2774 11 279 2614 1957 22741 198 35643 15 11 220 17 60 763 271 256 830 283 498 15 11 220 16 60 198 256 585 283 220 15 198 256 585 11 830 957 60 283 220 16 11 220 17 260 653 585 369 5860 11 1179 830 957 60 369 5860 198 256 1133 2007 8 271 9538 1048 25 271 220 2972 1679 47 220 18 16 18 17 332 471 39226 36760 1188 91841 198 257 561 24332 364 279 15171 5480 1 4416 13 1358 21199 26115 15840 11889 198 1464 43414 271 21199 26115 15840 369 279 10286 11 303 264 3074 4952 11 314 264 198 25064 5488 321 449 15840 4952 25 271 256 54109 50145 21152 61987 18720 5480 18720 445 318 27188 1951 735 7302 27160 8 198 256 18720 5480 336 61987 12448 735 16819 330 42514 735 14701 735 81987 198 256 18720 445 3696 61987 6421 420 735 6320 420 735 15171 420 735 8168 420 735 3405 420 735 34397 420 735 5775 420 735 65667 420 198 1734 735 328 2372 420 735 14291 420 735 13012 420 735 38591 420 735 24403 420 271 7 9538 3644 34019 5262 364 279 18890 16954 314 279 1483 2250 198 65151 5946 271 2014 54109 15840 64013 279 2100 318 7949 11 19231 4472 198 28973 11889 11 4021 381 449 30047 286 8 321 279 7258 198 1551 11 25822 279 7620 5488 3050 310 279 913 314 15840 198 263 279 1330 52680 11 321 48233 279 1067 310 279 3889 2100 13 198 760 2100 369 1132 24276 2957 13 271 2014 54109 15840 7258 1040 328 87 1373 220 16 1 628 381 56546 430 198 63230 283 830 478 220 16 1 310 10752 264 4283 11 694 524 6681 5951 2376 13 733 279 198 18599 26115 2243 11 328 87 1 369 1132 24276 2957 13 7054 11 948 3102 11 198 1719 4880 5488 369 10345 348 258 40284 11961 7063 421 4598 1056 198 43408 264 491 1576 321 58706 421 310 279 2100 11 279 2235 1576 198 284 10490 4370 13 271 41657 4472 30152 11 54109 30152 14989 279 2047 12 198 10346 3008 348 14372 9 36319 279 1245 23641 3008 13 220 1690 3010 11 328 64 957 60 198 8236 281 2007 9646 1118 5686 5072 328 64 957 18503 1179 424 64013 328 69 2007 9646 321 25822 198 1719 5089 11 321 1483 391 11 424 13514 279 1067 1142 310 328 64 957 60 3158 271 2278 279 4531 314 58706 310 43716 321 5081 11218 303 264 198 14892 4952 11 279 15840 2725 539 54109 15840 198 91318 369 17083 279 1788 1560 430 4472 30152 13 33105 11 198 4058 279 4531 314 279 3102 348 258 40284 9 7465 11 279 7620 198 8987 10345 539 54109 15840 369 279 1788 430 279 4472 198 25064 7289 13 271 2381 11218 864 513 6816 14619 11 279 1788 83081 883 198 1005 321 2774 7944 16298 430 364 5611 30152 13 1358 2014 84027 15840 11889 198 1464 43414 271 9 19109 9 15840 369 279 10286 11 303 264 3074 4952 11 314 198 64 3759 466 6816 20572 321 449 9801 15840 198 23425 25 271 256 60736 50145 21152 61987 18720 5480 22698 7258 198 761 4241 420 318 11531 1114 27160 735 7302 27160 7025 271 760 6463 494 4472 33284 11889 369 421 1132 264 3074 198 5480 369 5252 13 271 2381 4145 4874 430 15840 11218 11 413 303 523 466 4536 6676 11 198 1719 31144 513 24276 321 9476 303 264 3175 523 466 4536 198 8843 12880 38314 548 1 421 369 264 10675 12366 494 3759 198 11072 318 76 37652 413 843 8 310 24276 31144 13 1061 6816 369 198 86 21833 321 369 9184 3354 506 279 1151 314 523 466 4536 198 2513 10993 11 413 31144 513 1669 88034 13 271 2381 22666 430 15840 11218 11 279 31144 513 24276 198 331 303 523 466 4536 6676 11 694 524 9476 13 271 2592 264 803 369 60736 303 264 709 6676 11 1179 411 803 369 2136 198 1891 421 6676 13 90117 513 2496 24276 321 9476 303 709 198 78134 13 271 2592 279 1245 1375 3008 369 2945 11 449 60736 15840 25822 198 1719 4880 15840 1518 36319 31144 318 2777 198 658 44575 553 1368 279 1245 1375 3008 369 524 2945 364 449 7258 198 5480 11 1179 279 37999 64013 279 2100 3525 364 279 1483 198 1 548 723 1163 548 33905 466 12880 84163 548 33905 1562 13 271 9538 1048 25 271 220 2972 1679 47 220 20 17 21 332 471 31058 364 12051 90117 198 257 561 13340 421 3568 18890 364 34848 1060 279 4346 314 198 257 7101 318 15693 523 7101 321 2774 7101 681 198 257 4370 314 35493 1070 1472 5853 13 271 220 2972 1679 47 220 19 23 19 332 471 3856 29642 198 257 561 13340 421 3568 279 328 86711 1 4536 310 3300 264 5129 198 257 18890 364 913 31144 421 628 381 1429 303 1064 6157 198 257 7141 321 27069 82 13 271 5221 303 2243 220 18 13 23 25 4543 60736 30152 2070 279 1788 198 13392 880 303 279 1245 1375 3008 430 5611 30152 13 56703 11 198 14265 22666 318 4650 632 2195 8878 80019 13992 22666 8 8613 264 198 54304 1412 13 1358 10266 27693 198 584 332 271 3446 303 2243 220 18 13 20 13 1358 31099 709 7044 198 1464 34195 271 256 3203 9310 727 61987 498 91430 60 328 7449 1 328 727 1 2821 591 32467 498 15704 1951 60 56439 198 3696 4241 397 1 7258 60 22698 15705 271 19678 314 12654 1766 27693 628 381 20950 321 55483 506 1599 198 7462 318 3924 348 5816 13754 187525 328 11093 1 22666 11 328 7449 364 1 321 328 7449 198 4058 1 628 1132 381 1429 303 279 2407 314 264 75140 709 13 271 24407 4364 440 328 7449 687 1 18890 513 2590 75140 198 20752 11 1442 413 781 635 524 6435 328 11093 1 466 328 7449 1 20209 13 271 2064 369 264 328 32769 1401 1 310 958 264 328 28726 494 1 7258 4613 279 2407 198 1020 264 75140 709 13 271 2014 3010 314 264 75140 709 25 271 256 3203 687 2821 9453 16 11 1626 17 1590 198 950 635 92099 363 198 950 2454 1010 14308 13754 363 271 5221 303 2243 220 18 13 22 25 328 11093 1 321 328 7449 1 513 1381 20209 26 198 1666 17737 781 998 1132 11421 430 1680 4613 279 2407 314 264 198 5816 13754 709 13 1358 760 328 7449 364 1 4952 198 1464 41615 271 256 3203 5305 21152 61987 328 7449 1 364 21152 271 2014 348 299 26930 49124 9 5529 449 12880 1269 261 548 1 1654 421 198 19194 391 4523 449 348 299 26930 14650 11961 864 628 1562 198 299 26930 1970 303 1141 12880 276 419 548 1 1654 13 271 760 328 7449 364 1 4952 5954 16219 18971 888 198 299 26930 5199 4635 13 271 760 2614 1970 25 271 256 3203 364 39678 303 84788 25 198 950 15035 5803 198 256 745 25 198 950 15035 5803 17 271 3742 5067 78157 13185 310 25 271 256 5199 283 318 23954 8 198 256 5199 283 913 26202 19768 1269 261 3676 1947 8 198 256 4162 283 2912 271 256 1345 4162 25 198 950 1380 25 198 1779 39678 283 2454 913 26202 19768 276 419 3676 1947 8 198 950 3525 13809 6322 51317 25 198 1779 4162 283 3439 198 950 745 25 198 1779 15035 5803 198 256 745 25 198 950 15035 5803 17 271 9538 1048 12880 1269 261 548 33905 321 12880 276 419 548 33905 364 3447 13 271 2064 369 264 328 32769 1401 1 310 958 449 328 7449 364 1 4952 4732 279 2407 198 1020 264 75140 709 13 1358 760 328 7449 440 1 4952 198 1464 41830 271 256 3203 6406 21152 61987 328 7449 1 440 21152 271 2014 348 299 26930 2193 6436 9 369 264 348 2078 6436 9 421 369 2858 198 951 25410 10993 303 1141 348 1883 9 321 348 13259 9 5245 13 271 760 2614 1970 25 271 256 3203 440 13073 78847 430 39678 25 198 950 15035 5803 271 284 5067 78157 13185 310 25 271 256 6436 283 318 3152 61670 8 198 256 264 1883 283 913 53314 19768 64 1883 548 198 256 264 13259 283 913 53314 19768 64 13259 548 198 256 869 283 2454 264 1883 53314 8 198 256 4062 52797 283 3439 271 256 1380 25 198 950 39678 283 869 198 950 15035 5803 198 256 3525 25 198 950 4062 52797 283 2912 198 950 413 524 2454 264 13259 53314 11 348 7545 62538 3009 75458 198 1779 4673 198 256 5326 25 198 950 413 524 4062 52797 25 198 1779 2454 264 13259 53314 11 2168 11 2168 11 2168 8 271 9538 1048 12880 64 1883 548 33905 321 12880 64 13259 548 33905 364 3447 13 271 2064 369 264 328 32769 1401 1 310 958 449 328 7449 440 1 4952 4732 279 198 2513 314 264 75140 709 13 271 9538 1048 25 271 220 2972 1679 47 220 19 24 17 332 471 4315 27693 440 3203 321 2454 18890 198 257 561 13340 421 1801 1766 27693 264 5976 41954 7059 303 198 257 12654 11 321 3568 12250 18890 13 1358 27382 11497 318 7729 8 198 865 11862 271 2014 12448 29853 430 449 18546 369 264 803 13 220 3382 3644 21808 11497 198 429 20209 364 74931 7044 321 3644 93141 321 10649 364 198 73693 314 33791 321 10649 13 271 4350 279 803 369 6608 310 449 1576 11 15977 314 279 18546 34237 198 8761 1576 13 3095 264 803 369 524 6608 11 449 4621 310 14989 424 198 31207 264 328 657 1401 1 4531 13 271 332 16295 803 48509 2631 64700 3095 449 12448 421 1414 1769 13273 303 198 64 523 7044 11690 440 1330 466 777 50359 5587 321 198 26189 524 809 303 1330 466 777 69351 11 424 369 6306 264 348 1929 198 591 9 314 421 523 13 9391 4874 513 22773 310 264 4860 1304 198 14372 1970 369 7658 364 1070 13 220 561 17454 47640 279 198 1005 803 11 440 6187 69351 6747 321 264 3074 50359 198 4069 290 11 303 3929 314 279 803 13 220 1690 3010 11 279 12448 12880 72967 1 198 499 44216 303 264 523 6725 328 38617 1 668 381 22773 310 8728 38617 548 72967 3158 198 1919 17454 369 9209 314 279 80316 34826 2193 303 864 198 1719 12448 369 1429 13 220 1368 279 22773 803 369 8744 1248 198 12265 261 1056 220 17 20 20 5587 681 7873 4364 60735 362 1189 198 71 61608 13 1368 279 523 803 16661 1132 314 69351 11 874 198 1405 1570 369 2725 13 1358 43 36552 198 584 271 29292 11222 886 321 5639 72710 321 5090 23311 198 75 36552 25 271 256 22419 61987 886 34600 735 5639 34600 198 2215 735 7308 735 2153 3946 735 712 3141 862 271 79520 314 264 22419 34237 449 1576 314 279 2574 913 318 889 11 198 9363 11 7308 11 18484 1406 1324 11 6150 1324 8 440 279 2574 198 927 13 220 561 869 1189 381 9556 6921 303 279 1105 314 18484 1406 198 429 48271 318 22523 8 72710 13 220 3382 3644 24661 1091 364 3447 13 271 2324 72710 7618 310 49835 795 4346 11 321 15612 279 198 1640 725 9284 369 2599 2894 1056 1141 869 13 220 27869 198 13764 35840 314 72710 440 279 1788 869 318 48219 279 1788 198 13581 19934 303 279 1957 1414 466 264 2086 30520 8 1189 6637 198 1719 1788 1576 466 264 2086 1576 440 279 1788 869 13 1358 9962 4694 6816 2528 198 34581 271 760 2614 5245 628 381 4364 310 30069 279 7063 314 198 8843 2528 318 784 314 11 15840 310 11 466 34872 314 328 87 2557 871 364 198 1005 12740 13 271 1640 4692 78913 3676 701 11 803 8 271 256 20505 948 279 1581 6816 2528 14107 440 449 198 256 328 3776 1401 1 318 48219 12880 447 8843 548 33905 24096 449 198 256 328 3776 1401 1 1521 348 591 9 369 524 449 2774 6816 466 449 198 256 6816 303 279 523 4757 364 328 701 4961 466 12880 447 548 33905 314 264 348 591 9 198 256 3234 24096 328 3776 1401 1764 220 1061 1654 1220 2892 198 256 460 279 318 68600 8 6816 869 466 4673 449 328 3776 1401 1 198 256 4531 13 271 256 6817 421 413 279 6816 369 1669 1472 279 4472 16456 11 198 256 12880 78913 548 33905 369 524 2512 13 220 318 1919 369 449 45299 51505 31179 198 256 1881 12880 78913 548 33905 321 12880 84163 91514 3158 8 1061 369 2725 2107 364 198 256 14588 7715 321 1521 5752 12880 78913 548 33905 1000 599 198 256 874 1560 310 2528 975 7944 314 279 2774 13 220 6817 421 506 198 256 3140 364 2774 7101 11 488 628 12062 2702 2444 539 524 198 256 37029 866 2663 303 279 2774 6816 10675 318 7834 198 256 4370 37029 1070 303 2361 1576 553 220 3382 279 198 256 12880 447 8843 548 33905 1654 3559 364 264 1560 310 3404 615 2702 198 256 2444 888 6816 2528 13 271 1640 4692 447 8843 3676 701 11 803 8 271 256 20505 632 92649 310 4072 6816 74992 364 198 256 12740 314 279 523 13 1368 279 523 1048 18091 12880 78913 548 48812 198 256 279 14827 668 524 381 2512 7014 12880 447 8843 548 33905 2892 198 256 6526 424 20335 466 24096 449 328 3776 1401 3158 1061 1654 198 256 1220 460 279 318 68600 8 6816 869 466 4673 449 198 256 328 3776 1401 1 4531 13 733 1906 310 5471 23064 48575 303 198 256 411 1654 11 1141 7873 1220 2590 1562 279 2255 523 198 256 1654 440 279 1788 803 310 2528 866 7944 424 3749 11 364 198 256 3010 11 328 1640 4692 447 8843 3676 701 11 803 90765 271 256 6817 25 271 257 1061 1654 1189 1990 381 29715 290 948 3221 685 3175 5245 198 257 430 279 1067 314 17457 27757 4417 3992 18890 466 198 257 5617 3309 5568 13 3382 9493 1654 18054 13 271 256 1690 3529 15739 6816 74992 11 24096 449 89304 1484 198 256 328 1640 4692 78913 548 1 440 5789 328 2220 1 321 328 591 3158 271 1640 4692 84163 3676 701 11 803 11 869 8 271 256 20505 948 449 6816 15840 369 16693 13 220 1061 369 2512 198 256 4370 314 279 4472 16456 318 72 1673 13 3436 279 869 303 279 198 256 2774 10675 553 348 591 9 369 279 6816 803 11 348 927 9 369 279 198 256 869 310 381 12244 310 424 13 271 256 1368 12880 84163 548 33905 6587 310 9501 310 449 2774 6816 11 424 198 256 1220 1562 279 2255 523 1654 440 279 1788 803 11 364 3010 11 198 256 328 1640 4692 84163 3676 701 11 803 11 869 90765 271 256 1690 3529 15739 6816 30152 11 24096 449 89304 198 256 1484 328 1640 4692 84163 548 1 440 5789 328 2220 487 328 591 487 328 927 3158 271 1640 4692 9302 2896 3676 701 11 803 8 271 256 8639 12880 84163 548 33905 694 364 6816 34872 4370 314 198 256 15840 13 220 1061 1220 1132 381 11205 413 328 9302 2746 2557 1 369 198 256 21716 364 279 1576 13 271 256 1690 3529 15739 6816 18686 880 11 24096 449 89304 1484 198 256 328 1640 4692 9302 2896 548 1 440 5789 328 2220 1 321 328 591 3158 271 1640 4692 3614 3676 701 8 271 256 20505 948 328 3614 33905 369 2512 383 279 1576 13 1473 49124 1902 381 198 256 5743 13 328 3614 33905 31637 279 5743 49124 310 264 1103 321 198 256 20218 424 13 1358 9962 4694 4536 6816 2528 198 2960 8442 271 19747 4874 12880 78913 548 1 321 12880 3614 548 1 628 381 1048 1429 310 198 91125 2528 310 4536 7944 13 561 12880 78913 548 1 709 506 198 1719 4536 2119 1220 4054 799 5515 864 369 279 803 314 449 198 8843 321 460 279 23424 869 466 4673 449 328 3776 1401 3158 198 2592 449 6816 369 524 1669 383 264 4536 1576 1472 279 4472 198 20377 11 585 1673 13 328 1640 4692 447 8843 548 48812 1179 12880 78913 548 1 369 198 1773 290 303 279 4536 12880 8217 548 1 1518 17987 449 328 3776 1401 3158 198 2592 1669 11 424 369 2512 440 279 6816 803 321 279 1067 369 198 74936 13 271 760 12880 3614 548 1 709 1220 4054 874 5789 11 321 460 449 198 1947 470 314 8796 421 10551 279 4874 15156 383 4536 13 1368 198 27803 11 411 709 34282 279 5129 328 3614 33905 2624 383 264 198 4211 13 271 2381 264 777 6699 1046 2540 46439 314 279 4536 7465 318 14874 198 11987 11 5706 11 4831 23618 799 628 716 279 12880 1005 548 1 6816 198 1020 264 4536 1576 310 264 35959 314 328 8969 26094 900 3158 1690 3010 25 271 256 1120 5530 198 256 494 4346 1120 13317 900 271 256 523 6054 8039 3223 72729 900 1590 198 950 687 1259 29828 3676 701 1590 198 1779 460 281 6 41197 81 8039 313 701 4692 591 548 11325 271 950 687 1259 84163 3676 701 11 6175 11 869 1590 198 1779 1133 925 12887 49369 313 2896 92 1076 837 198 1779 2184 25658 84163 3676 2896 11 869 8 271 256 5530 25949 88970 591 548 906 548 1005 548 283 6054 8039 3223 271 8839 25 271 220 3761 5562 4536 12880 78913 548 1 321 6047 4536 12880 1005 548 1 1132 198 220 7556 1353 8339 1801 1608 279 6816 2528 18890 1317 5774 198 220 30745 279 4536 36536 318 47132 539 1970 2785 279 4536 11 466 198 220 4417 264 5604 310 279 4536 725 36536 10675 8 369 74872 13 271 5221 303 2243 220 18 13 20 25 12880 1005 548 1 4536 6816 369 1381 44375 13 271 3446 303 2243 220 18 13 22 25 12880 78913 548 1 321 12880 3614 548 1 4536 7944 13 271 9538 1048 25 271 220 2972 1679 47 220 20 21 17 332 471 13317 1259 78913 548 321 1259 3614 548 198 257 3744 53283 279 12880 78913 548 1 321 12880 3614 548 1 5568 383 13067 13 1358 60856 286 3744 24887 198 1464 857 271 760 2614 5245 1132 3668 948 449 2774 314 279 523 198 747 1988 279 1654 318 64 748 18330 348 51347 9 523 8 7701 303 449 198 9 7862 9 523 318 1719 15933 1902 381 303 2892 279 6171 725 523 198 34491 466 303 279 523 10675 364 799 314 1141 6357 553 220 733 279 198 49927 3559 11 1001 1719 6816 828 18675 310 279 6816 6483 803 369 198 1719 1328 314 279 3234 303 279 6171 523 515 12880 8217 548 3158 271 1640 4692 447 3676 701 11 2774 11 6171 5675 8 271 256 20505 310 615 279 6816 314 279 6171 523 318 1005 6816 198 256 2528 8 466 314 449 2774 314 421 523 318 4696 6816 198 256 2528 553 561 9801 348 7862 9 5515 369 279 6171 523 11 1345 198 256 348 4696 9 369 279 2774 421 279 6816 557 23975 1472 11 198 256 466 328 3928 1 948 279 6816 369 23975 1472 279 348 7862 19315 271 256 1061 1654 1220 460 279 23424 6816 869 466 4673 449 198 256 328 3776 1401 1 4531 13 271 256 2972 1679 47 220 17 20 17 332 28151 421 12880 447 548 33905 369 39571 440 799 466 1330 198 256 5789 13 220 12654 725 1765 5617 3309 45933 1761 411 198 256 24332 26 4611 11 424 369 4222 421 1010 4688 23280 7141 198 256 599 45933 421 1325 2107 5789 13 220 12654 725 1765 198 256 12880 447 8843 548 33905 7873 2590 15734 303 2107 5789 198 256 3315 781 513 2483 466 524 13 271 1640 4692 723 3676 701 11 2774 11 869 8 271 256 20505 310 716 279 6816 383 449 2774 348 4696 9 314 279 6171 198 256 523 310 264 491 869 11 348 927 19315 271 256 6817 11 7594 12880 723 548 33905 466 12880 4394 548 33905 4203 279 2993 314 198 256 15933 310 264 1001 673 15933 10878 220 3382 18239 10430 3744 24887 364 198 256 777 3447 13 271 1640 4692 4394 3676 701 11 2774 8 271 256 20505 310 3572 279 6816 383 449 2774 348 4696 9 314 279 198 256 6171 523 13 271 41310 314 45933 1189 1048 599 279 12880 2220 1005 548 1 6816 198 27803 25 271 1640 4692 2220 1005 548 271 256 561 6816 12880 2220 1005 548 1 369 31233 539 279 328 79221 1 4536 198 256 430 36588 279 523 1332 411 1576 557 4364 318 14874 411 198 256 33744 628 7543 303 15135 24715 15603 314 8476 523 198 256 7944 553 1690 1562 4635 11 424 1189 12832 421 449 2774 314 279 198 256 2574 913 318 269 264 35959 8 369 3481 466 2483 430 279 1118 198 256 65269 5515 318 1891 3010 11 351 29292 7062 411 6816 364 198 256 632 10573 5245 421 513 11205 303 351 553 1358 14732 10430 3744 24887 198 1464 590 271 623 4437 11 264 15933 369 449 1576 6816 440 1001 30018 198 53992 9059 799 6483 6816 2528 682 978 49724 539 5245 198 258 279 15933 11176 25 220 12880 447 548 48812 12880 723 548 48812 321 198 1 548 4394 91514 3158 1368 866 314 1782 5245 513 4364 364 449 1576 11 424 198 284 1018 310 381 264 15933 13 271 760 1581 7465 364 6816 2528 369 310 615 11 716 11 466 3572 198 1719 6816 494 449 1576 725 10675 13 1690 2774 11 328 64 1926 1 682 264 198 20377 8516 5732 440 328 64 4692 8217 548 659 87 72203 1179 198 42993 2784 19768 8217 548 659 87 72203 321 13942 1472 279 2255 6632 314 198 42993 2784 9646 42328 2197 565 32748 13 271 10883 11 413 279 6749 5072 869 369 449 1576 25471 799 314 279 198 51347 5245 11 1179 12654 1189 2723 279 1581 7465 321 198 21731 279 15933 1654 4370 13 220 10646 411 13273 303 279 198 25317 738 8516 13402 383 864 15933 5245 998 4364 321 198 4993 781 998 2512 13 271 760 5732 1406 364 15933 27757 369 264 10649 11 328 64 1926 3158 2500 198 1719 5789 513 32878 13402 383 328 64 763 271 15555 6920 198 256 561 43070 321 3140 4047 1562 369 948 1156 1970 5774 198 256 63821 264 15933 1654 25 262 328 87 4692 447 3676 64 90765 271 2443 25073 198 256 1368 10649 310 449 1576 2774 11 328 64 1926 1 369 22773 1083 279 198 256 1562 25 328 1267 2784 19768 8217 548 659 87 642 4692 447 3676 64 11 913 2784 578 3158 271 1890 25073 198 256 1368 10649 310 264 523 11 328 32 1926 1 369 22773 1083 279 1562 25 198 256 328 32 4692 8217 548 659 87 642 4692 447 3676 3928 11 357 90765 271 18701 25073 198 256 357 56150 18054 1680 430 328 9241 4205 11 264 553 87 1 25628 198 256 328 64 4692 1005 47719 76 298 548 1 364 264 2255 523 328 33 1 2614 328 32 1 321 1179 198 256 4523 328 33 4692 8217 548 659 87 642 4692 447 3676 64 11 357 90765 220 1368 524 264 15933 11 328 87 1 198 256 369 5743 33702 13 271 2381 2774 34519 11 279 51274 314 15933 27757 13402 198 263 864 15933 5245 513 4364 13 220 357 15933 628 6762 866 198 72474 314 12880 447 548 48812 12880 723 548 33905 321 12880 4394 91514 3158 220 1368 424 198 26189 524 6762 12880 447 548 48812 1179 30745 279 6816 668 460 198 1719 15933 1576 4924 7014 1017 369 264 869 303 279 1576 725 198 4696 10675 13 220 1368 279 15933 18091 12880 723 548 33905 321 4979 198 1 548 4394 548 48812 424 369 264 795 15933 26 413 424 18091 13468 11 424 369 198 64 2397 13545 15933 13 220 49941 11 795 45933 6762 2107 198 1 548 447 548 33905 321 12880 723 548 48812 1345 2397 13545 45933 599 1066 279 198 1 548 447 548 33905 1654 13 220 2791 45933 440 12880 447 548 33905 321 12880 723 548 33905 198 57889 4979 12880 4394 91514 871 4364 2590 2723 264 311 17829 303 449 198 4696 10675 13 220 733 12503 11 2397 13545 45933 628 381 198 1908 40787 539 12740 13 271 29292 5245 318 15693 1782 34816 440 8168 21527 1 321 198 92943 25901 871 513 11205 430 2397 13545 45933 13 220 60694 11 198 46230 628 80403 321 2723 5245 13 220 1061 5954 3714 198 46230 310 20604 26228 421 1684 494 975 12740 314 279 198 23310 523 13 271 760 328 3573 33905 709 369 11205 430 264 795 15933 13 198 10865 391 11 12740 4021 2723 279 7465 314 264 3234 13 1358 548 48093 548 198 41615 271 9 548 48093 548 9 2070 586 310 20335 15444 795 3492 318 4650 198 12811 8 321 22380 279 9400 314 12880 8217 548 1 321 35832 12768 1062 548 9 198 17818 1657 20335 13866 303 35832 48093 548 9 466 2420 303 264 2594 5946 271 760 3433 6568 888 1608 12880 8217 548 1 628 381 4927 13 16260 198 20377 4478 628 381 11602 12573 430 1575 13 271 1640 4692 48093 548 271 256 1061 523 3759 628 381 12244 264 886 11 49124 11 466 8240 198 256 314 8796 440 3759 4874 1429 539 12740 13 220 35832 48093 548 9 198 256 28538 3433 364 279 13866 7101 321 27028 279 198 256 16465 9400 314 12880 8217 548 1 321 35832 12768 1062 548 9 364 1754 198 256 2774 13 271 21003 383 1608 35832 48093 548 77340 271 9 3095 14722 5672 494 264 523 1973 35832 48093 548 11961 279 12880 8217 548 1 321 198 220 35832 12768 1062 548 9 6816 314 279 12740 668 2590 381 15156 13 271 9 16643 264 12880 8217 548 1 3759 11 12740 4021 381 12244 491 198 220 7101 524 9711 303 279 35832 48093 548 9 7044 13 220 80427 310 198 220 9501 310 449 632 30214 3759 803 24096 328 3776 1401 3158 1368 198 220 8476 15840 314 491 7101 369 12321 11 1179 884 198 220 6952 548 8217 548 14482 310 279 8240 314 8796 303 279 35832 48093 548 9 198 220 17466 13 271 9 16643 264 35832 12768 1062 548 9 3759 364 1754 2774 11 6632 25471 198 220 35832 48093 548 9 635 524 1761 328 12768 14619 1 310 1141 12740 13 1368 198 220 7234 5604 1761 369 4221 11 1179 884 6952 548 12768 1062 548 14482 310 279 198 220 8240 314 8796 303 279 35832 48093 548 9 17466 13 271 9 35832 48093 548 9 513 11205 506 279 523 2119 539 6611 198 220 45933 364 1754 3759 803 13 220 1577 264 1067 11 523 7944 198 220 4021 381 1429 310 716 1581 2663 364 2774 7101 4364 198 220 539 35832 48093 548 243343 5752 11 279 523 6816 1000 31153 279 198 220 15933 15840 13 271 9 561 1852 314 264 35832 48093 548 9 17466 369 524 6973 310 279 523 198 220 1332 424 369 4364 13 220 35832 48093 548 9 13866 303 6357 513 2420 198 220 303 1623 6632 13 4213 11 1623 65878 668 615 264 12880 8217 548 1 198 220 321 35832 12768 1062 548 9 7014 781 1048 6762 35832 48093 548 9 318 7949 1220 198 220 1132 6435 4874 314 866 348 34325 9 15250 553 271 9 1368 264 523 18091 264 9167 1048 4364 303 264 2255 523 11 279 2774 198 220 3759 4364 539 279 2255 523 9167 369 78909 318 11347 539 198 220 46616 1141 15933 5774 494 279 2255 523 553 1061 198 220 42812 279 7063 314 279 1957 5438 13 220 733 279 3725 11 264 198 220 1716 1189 381 3568 310 5190 411 13 271 9 328 77318 1 668 381 8953 413 2397 3092 35832 48093 548 9 513 4364 364 264 198 220 523 13849 494 264 1535 9521 28365 1 5617 3309 913 1 1680 430 198 220 328 390 487 328 9363 487 321 328 23814 3158 271 9 5586 2397 29868 348 1947 470 9 1189 381 12244 310 35832 48093 548 19315 271 9 1368 264 328 34491 1 369 1429 310 9501 35832 48093 548 11961 279 10675 6679 198 220 668 381 1429 430 279 9167 4874 13 561 2663 314 279 10675 628 381 198 220 1429 310 3300 791 12 8843 4479 18034 421 668 381 37382 539 198 220 328 79221 652 5069 33905 321 12234 303 279 2468 314 328 8389 363 3158 271 9 12880 1005 548 1 15840 4138 1132 413 2107 6632 599 279 1788 198 220 35832 48093 548 19315 271 9 27869 40443 440 5081 1678 15065 2594 6632 628 381 198 220 1429 11 694 1132 799 2594 369 5252 310 599 7944 3354 539 198 220 15250 318 1719 975 22372 1902 599 4147 9167 47410 8 471 25709 198 220 4673 328 77318 3158 271 9 1368 449 348 6640 9 369 1429 364 35832 48093 548 9 1179 264 348 51347 9 369 198 220 3354 364 1754 314 279 14650 725 2663 13 4213 11 279 35832 48093 548 9 198 220 6816 668 381 449 4147 14650 13 1358 3776 14619 198 30896 271 2014 6816 5604 369 264 5839 7854 539 264 4031 321 264 803 25 271 256 16819 330 42514 61987 5839 22052 12448 271 760 5839 1902 14989 310 449 1576 314 264 913 421 11222 198 8843 14619 11 864 1379 5978 635 13 220 1061 1576 369 1179 198 1038 290 310 7936 279 6816 6483 803 369 279 12448 13 561 913 198 429 869 8677 369 10521 539 279 1576 13 220 27869 53231 198 1020 279 1788 6816 5604 1189 7302 2086 5978 13 271 1919 5492 628 381 30650 539 65581 279 198 1 548 447 8843 548 33905 1654 466 279 12880 78913 548 33905 1654 13 220 561 198 1 548 447 8843 548 33905 1654 369 2512 1118 321 2892 4523 264 869 198 269 24096 328 3776 1401 1 413 279 6816 369 524 2420 13 271 2592 449 328 3776 1401 1 369 8953 321 279 1576 682 264 12880 78913 548 33905 198 4249 11 421 1654 369 2512 430 264 31686 13 1358 21199 26115 15840 11889 198 13052 37745 271 21199 26115 15840 369 279 10286 11 303 264 3074 4952 11 314 264 198 25064 5488 321 449 15840 4952 25 271 256 54109 50145 21152 61987 18720 5480 18720 445 318 27188 1951 735 7302 27160 8 198 256 18720 5480 336 61987 12448 735 16819 330 42514 735 14701 735 81987 198 256 18720 445 3696 61987 6421 420 735 6320 420 735 15171 420 735 8168 420 735 3405 420 735 34397 420 735 5775 420 735 65667 420 198 1734 735 328 2372 420 735 14291 420 735 13012 420 735 38591 420 735 24403 420 271 7 9538 3644 34019 5262 364 279 18890 16954 314 279 1483 2250 198 65151 5946 271 2014 54109 15840 64013 279 2100 318 7949 11 19231 4472 198 28973 11889 11 4021 381 449 30047 286 8 321 279 7258 198 1551 11 25822 279 7620 5488 3050 310 279 913 314 15840 198 263 279 1330 52680 11 321 48233 279 1067 310 279 3889 2100 13 198 760 2100 369 1132 24276 2957 13 271 2014 54109 15840 7258 1040 328 87 1373 220 16 1 628 381 56546 430 198 63230 283 830 478 220 16 1 310 10752 264 4283 11 694 524 6681 5951 2376 13 733 279 198 18599 26115 2243 11 328 87 1 369 1132 24276 2957 13 7054 11 948 3102 11 198 1719 4880 5488 369 10345 348 258 40284 11961 7063 421 4598 1056 198 43408 264 491 1576 321 58706 421 310 279 2100 11 279 2235 1576 198 284 10490 4370 13 271 41657 4472 30152 11 54109 30152 14989 279 2047 12 198 10346 3008 348 14372 9 36319 279 1245 23641 3008 13 220 1690 3010 11 328 64 957 60 198 8236 281 2007 9646 1118 5686 5072 328 64 957 18503 1179 424 64013 328 69 2007 9646 321 25822 198 1719 5089 11 321 1483 391 11 424 13514 279 1067 1142 310 328 64 957 60 3158 271 2278 279 4531 314 58706 310 43716 321 5081 11218 303 264 198 14892 4952 11 279 15840 2725 539 54109 15840 198 91318 369 17083 279 1788 1560 430 4472 30152 13 33105 11 198 4058 279 4531 314 279 3102 348 258 40284 9 7465 11 279 7620 198 8987 10345 539 54109 15840 369 279 1788 430 279 4472 198 25064 7289 13 271 2381 11218 864 513 6816 14619 11 279 1788 83081 883 198 1005 321 2774 7944 16298 430 364 5611 30152 13 1358 36117 7258 198 865 271 50 12393 279 10993 314 348 5816 13754 9 383 449 348 11093 470 9 1576 13 2885 198 3138 381 1429 4613 264 348 5816 13754 709 19315 271 256 2454 21239 61987 328 11093 1 5839 271 3446 303 2243 220 18 13 20 13 1358 20682 33633 7289 198 34581 271 760 7620 33633 7289 599 279 20066 10304 198 41150 13 220 6817 421 1010 314 1439 7289 1048 3668 310 3529 2397 12 198 19000 4346 13 220 33554 494 279 2279 5497 11 1017 513 1132 1330 198 41150 11 799 364 12197 63062 19024 321 799 364 61103 198 31750 25 271 256 295 21239 61987 560 21239 735 295 21239 63669 560 21239 735 295 21239 93504 295 21239 735 198 1025 295 21239 34397 1 560 21239 735 295 21239 16397 560 21239 735 198 1025 295 21239 76104 560 21239 198 256 264 21239 61987 295 21239 735 264 21239 57488 295 21239 735 264 21239 25675 295 21239 271 760 63669 318 75305 1633 8 5497 34237 279 1918 314 1141 5789 13 198 760 5789 1902 2892 2107 381 4947 11 466 799 5515 1902 381 449 198 11326 321 279 975 1902 381 264 8240 13 733 279 4588 1105 11 279 198 36142 513 15625 310 264 4047 913 321 1179 53070 3658 13 198 623 279 14827 1105 11 8240 51623 369 10345 26 264 7968 198 265 53897 7911 34237 449 4147 8240 13 271 1919 5488 628 381 30650 1608 279 3175 12880 23402 548 33905 321 198 1 548 8454 355 548 33905 5245 13 271 760 93504 318 266 8 5497 369 10287 310 381 1429 364 5979 198 75305 1633 13 220 2233 57610 12654 4346 4072 411 5497 13 271 3446 303 2243 220 18 13 20 13 271 760 16397 318 27822 8 321 34397 1 318 29457 12436 8 19024 7302 279 198 13104 1128 314 836 5789 13 220 561 23311 5789 513 1118 198 73596 310 264 4047 913 13 14074 314 24959 34237 264 2153 11 1345 198 29457 12436 314 24959 2961 303 449 7308 26 279 1067 369 421 198 1020 34780 12436 440 279 3260 29457 515 709 8977 310 279 198 1334 13 220 14074 539 6942 24096 279 328 17461 49513 1401 1 4531 13 271 1919 5488 628 381 30650 1608 279 3175 12880 371 3155 340 548 33905 321 198 1 548 86715 526 340 548 33905 5245 13 271 760 76104 318 84934 8 5497 34237 279 25472 494 279 12436 314 198 1719 1118 5515 539 279 2018 13 220 561 23311 5789 513 1118 198 73596 310 264 4047 913 13 220 357 6942 1245 5515 24096 279 198 1 17461 49513 1401 1 4531 13 220 561 5789 1189 381 18484 1406 198 36142 11 378 1257 2487 328 18 13 16 19 4 15 13 22 1 16327 328 15 13 18 19 1 318 10957 328 18 13 16 19 1 16327 328 19 9 15 13 22 478 198 15 13 18 19 3158 8 220 561 57844 5497 2590 34237 264 1067 440 279 1788 1777 198 299 1141 2018 26336 318 269 6942 1174 279 10424 869 314 279 1067 369 198 6418 391 8882 1056 279 10424 869 314 279 2018 26336 498 16 906 271 760 6220 12436 321 57844 19024 513 8173 539 279 2614 198 16417 25 328 87 606 318 87 320 88 4653 88 478 318 87 4 88 90765 220 16093 12436 321 57844 513 1048 198 14844 440 279 5617 3309 709 328 596 2508 363 763 328 596 2508 2007 11 374 8 606 198 2007 320 88 11 830 4 88 90765 498 17 906 271 623 5089 310 15899 279 57844 5488 383 4947 11 279 76104 198 7636 369 1048 72052 539 886 5978 310 2649 2235 10970 198 889 35358 318 18092 3750 430 35282 553 220 561 18890 364 198 889 35358 369 7246 303 279 12654 11300 16699 11 198 2720 3964 10970 894 86508 13 271 760 348 84934 9 5488 628 381 30650 1608 279 3175 12880 2508 548 33905 198 4249 13 271 760 6220 12436 5497 11 279 57844 5497 11 321 279 328 596 2508 33905 198 1628 513 524 4364 364 6150 4947 13 220 11746 11 5335 310 264 198 79665 1406 1324 1608 279 328 3325 33905 709 413 8053 13 271 760 57488 318 698 662 8 5497 34237 279 2542 314 1141 5789 13 220 561 198 15889 1902 2892 2107 381 4947 466 2107 381 22958 314 279 1788 198 1267 13 220 733 279 4588 1105 11 279 4947 513 15625 310 264 4047 913 198 429 1179 3568 3658 13 733 279 14827 1105 11 279 22958 513 198 19551 268 639 13 271 1919 5488 628 381 30650 1608 279 3175 12880 698 548 33905 321 198 1 548 81 698 548 33905 5245 13 271 760 25675 318 1899 25580 8 5497 34237 279 6463 314 1141 5789 13 198 760 23311 5789 513 1118 15625 310 264 4047 913 13 271 1919 5488 628 381 30650 1608 279 3175 12880 1899 548 33905 1654 13 1358 20682 94589 7289 198 13052 9 271 4699 314 279 2250 94589 7289 682 264 2086 10304 2119 25 271 256 321 21239 61987 6206 21239 735 321 21239 13012 1 6206 21239 198 256 52131 21239 61987 321 21239 735 52131 21239 38591 1 321 21239 198 256 466 21239 220 61987 52131 21239 735 466 21239 58383 52131 21239 271 760 13012 1 5497 34237 279 94589 3449 314 1141 5789 11 864 1902 198 1323 24959 466 799 314 1070 1902 381 264 2446 1576 65581 198 1 548 429 548 33905 466 12880 11008 548 33905 3175 5245 13 271 760 38591 1 5497 34237 279 94589 67521 318 86055 2639 8 314 1141 198 15889 11 864 1902 381 24959 466 799 314 1070 1902 381 264 2446 198 1640 65581 12880 68890 548 33905 466 12880 12288 269 548 33905 3175 5245 13 271 760 58383 5497 34237 279 94589 318 94964 8 2639 314 1141 5789 11 198 7949 1902 381 24959 466 799 314 1070 1902 381 264 2446 1576 198 1908 48008 12880 269 548 33905 466 12880 7389 548 33905 3175 5245 13 1358 2010 23359 198 584 422 271 2010 5978 513 1429 539 279 7873 310 3875 1001 66561 12 198 49134 828 31097 12654 1970 1680 430 264 709 2407 13 2302 1684 198 1445 709 5978 1521 781 1459 1357 6435 264 5604 310 836 198 9461 10993 4424 13 220 5927 5978 513 5743 539 279 5617 12 198 258 328 19810 33905 709 321 628 381 26543 494 709 5978 198 19139 836 12880 1787 548 1 6816 13 3382 1048 279 328 1787 1 4536 13 271 5865 286 12880 1787 548 1 24096 449 89304 1484 328 1640 4692 78913 548 1 198 4058 5789 328 2220 1 321 1535 548 1787 548 67008 271 32 1970 1576 628 381 15234 466 24276 539 11946 424 318 62643 314 264 198 2348 886 8 310 279 328 11411 33905 466 328 13764 33905 220 5617 3309 5568 13 271 9538 561 5129 913 27980 364 777 1928 13 1358 760 13150 45823 2907 198 865 11862 271 1919 1576 369 16138 1429 539 81987 318 3924 326 407 793 553 220 1049 11222 198 2083 3175 7289 13 220 2532 369 6681 799 24792 45823 1576 11 6725 198 1 41295 45823 1 318 64 5617 3309 803 553 220 328 1267 10406 636 45823 8 33905 18082 279 198 1 41295 45823 1 35680 13 271 2064 369 5158 430 328 41295 45823 1 466 37431 3158 1358 760 17545 2907 198 584 37745 271 1919 1576 369 5743 539 5568 421 1459 1357 20335 460 264 198 927 13 220 1049 11222 874 3175 7289 13 220 2532 369 6681 799 819 198 1640 11 6725 328 3928 1 318 64 5617 3309 803 553 220 328 1267 25864 8 33905 18082 279 198 23310 35680 13 271 2064 369 5158 430 328 3928 3158 1358 900 23359 198 584 422 271 900 5978 3875 279 5090 1576 4346 13 220 1473 1576 725 913 369 198 5045 290 539 279 5617 3309 709 328 1267 363 3158 220 2532 513 874 3175 198 36901 383 4346 13 220 561 5129 4536 328 8969 1 18091 4874 364 198 528 5129 5617 3309 4346 13 271 4036 513 5158 1040 411 25 3919 1005 359 390 6073 3158 1358 6675 7289 198 865 332 271 256 466 4306 220 61987 321 4306 735 466 4306 328 269 1 321 4306 198 256 321 4306 61987 524 4306 735 321 4306 328 429 1 524 4306 198 256 524 4306 61987 11960 735 328 1856 1 524 4306 271 623 279 2193 314 6774 7289 11 321 1048 948 22666 513 198 2506 539 2444 6195 11889 11 279 2614 2663 513 31233 198 299 867 25 328 3913 487 328 3928 487 23311 6942 314 660 4346 11 321 4147 198 18034 321 23107 318 15693 8796 11 43716 11 11140 11 198 67 42723 11 7062 321 12431 17991 1367 553 220 1942 975 2663 513 31233 198 299 804 13 220 2570 37435 5978 628 30069 836 7793 869 539 198 751 1256 286 264 12880 2554 548 33905 1654 13 271 760 5497 328 1856 1 34237 328 2434 1 413 1141 5515 369 867 11 328 3913 1 198 57913 13 271 760 7258 328 87 321 374 1 1118 64013 348 87 243343 413 348 87 9 369 867 11 1141 198 927 369 5743 26 5752 11 348 88 9 369 24276 321 279 12571 869 198 284 5743 13 271 760 7258 328 87 466 374 1 1118 64013 348 87 243343 413 348 87 9 369 804 11 1141 869 198 284 5743 26 5752 11 348 88 9 369 24276 321 279 12571 869 369 198 74936 13 271 8839 421 13468 328 429 1 6130 328 269 1 8623 279 869 321 913 781 198 671 310 328 3913 1 321 328 2434 487 694 4598 460 279 1483 24276 198 14064 13 220 1061 369 6807 5222 11 378 1257 2487 413 328 82 1 369 264 886 421 198 5274 381 12215 539 264 1581 869 413 424 369 4147 11 279 7258 198 39441 466 359 7724 14482 34237 279 12321 869 13 220 8938 328 1856 1 682 310 1791 264 198 902 869 11 424 4523 264 2623 869 15029 314 279 913 314 1141 198 14064 318 1891 3010 11 328 1856 359 7724 14482 18082 328 3913 1 4598 1056 328 4456 3158 8 1358 760 328 8689 1 4952 198 865 17731 271 256 1387 21152 61987 328 8689 1 271 1 8689 1 1189 1132 11822 80316 519 2586 23283 303 264 328 1891 1 466 328 3361 1 198 10185 11 694 524 23283 303 264 709 466 523 7044 2785 421 198 10185 13 271 2064 84591 279 22746 82000 6136 11 41244 279 9801 328 1449 1 198 79248 413 279 6136 682 799 13 271 2592 264 328 1891 1 6136 369 30245 539 328 8689 487 279 6136 2444 2100 198 432 6902 1141 1428 869 13 271 4350 328 8689 1 15734 2444 680 314 264 328 1485 1 4952 440 264 328 37879 1 198 79248 11 421 328 37879 1 20984 369 15234 1518 2098 9101 279 198 10185 13 1358 2196 15031 39571 5978 198 13052 332 271 1640 4692 6450 3676 701 36038 2736 1076 2387 271 256 20505 948 279 2774 369 1001 42431 828 430 264 709 26 413 411 1654 198 256 369 4364 11 328 87 9125 16 11 1344 17 11 2423 9646 16757 45351 310 198 256 328 1267 2007 19768 6450 3676 87 11 1344 16 11 45347 3158 1358 53435 198 17731 271 32 1562 6526 264 39571 1576 318 68 1257 2487 264 348 1628 3636 440 264 10451 198 3092 3878 314 348 15889 77340 271 256 1562 336 61987 5839 32467 498 14064 1951 498 487 1293 735 59036 60 56439 198 256 5515 1951 285 61987 65269 41997 498 2129 56276 8119 49627 60 198 5015 498 2129 20209 41997 60 198 3696 735 56276 8119 49627 498 2129 20209 41997 60 198 3696 735 20209 41997 198 256 65269 41997 61987 65269 5457 318 2129 65269 5457 4653 198 256 65269 5457 406 61987 15840 27160 735 63669 7258 198 256 56276 8119 49627 61987 3376 61946 7258 735 15698 5457 8 198 967 318 2129 63669 7258 735 15096 15698 5457 4653 198 256 20209 41997 256 61987 318 19258 5457 735 65667 1 7258 8 198 3388 318 2129 15698 5457 735 15096 65667 1 7258 4653 198 256 15698 5457 260 61987 12448 35036 7258 271 2014 9801 26849 30644 1189 381 2945 1238 279 65269 321 198 19258 5789 694 1503 524 7556 279 50921 13 271 760 5839 1902 14989 310 264 39571 1576 318 846 37435 198 20752 11 5617 3309 5568 11 5245 314 5617 3309 5978 11 523 198 18630 11 5245 314 523 12740 11 321 660 5978 3322 264 198 1 548 6450 548 33905 1654 513 39571 553 220 1942 5515 22666 513 198 13764 12484 1518 279 1562 369 16693 13 220 5044 8042 310 3644 198 4987 16954 364 279 18890 314 15440 348 15704 9 11140 13 271 2592 15698 5789 513 2945 11 781 513 1118 15625 310 198 936 2910 5789 11 430 10695 13 220 5339 11 264 1103 314 9194 4233 15250 369 198 6898 364 279 15440 4868 13 220 1368 1017 513 443 65269 198 15889 11 781 513 8826 303 279 1118 443 15250 13 220 9019 11 364 1754 198 19258 5515 11 279 12448 369 1429 310 7995 279 198 5816 6604 286 9167 318 331 279 12448 369 279 1788 430 279 1118 15440 198 15704 803 11 279 1118 9167 369 1429 11 321 748 383 553 220 1368 279 9167 369 198 43967 9896 11 264 328 77318 1 4531 369 8953 13 17673 11 279 198 14064 369 8826 303 279 9167 11 20621 424 318 16296 413 279 7258 369 198 1 3928 487 424 39248 279 9167 553 220 3095 660 5789 599 978 14789 11 198 1719 15250 421 513 1990 9194 4233 513 9896 440 279 11811 198 2186 869 494 279 709 7044 13 220 318 3549 2663 513 198 5248 48129 11 2957 11 948 279 709 369 4364 26 8190 11 264 24983 1576 198 20171 430 264 1103 466 10675 1429 430 1581 869 668 381 5904 539 198 528 6526 421 1459 1357 13441 449 5515 869 364 279 11811 198 20877 26 411 1220 5802 381 29784 5946 220 1368 1017 513 866 9194 4233 198 48093 364 864 874 1581 869 369 5024 11 264 328 77318 1 4531 198 284 8953 13 220 17673 11 279 1103 314 9896 15250 369 1429 430 279 198 14064 1103 364 279 1562 13 271 332 6901 26033 7873 7472 64700 1473 7873 1189 3300 198 44697 3309 5568 6483 65269 4868 635 524 599 4874 11 1442 198 331 781 513 3260 29259 515 364 279 7193 314 9417 11 321 864 198 17977 1278 4021 381 16713 539 15698 13 220 733 351 29292 11 411 369 279 1105 198 1891 5568 11205 303 351 421 958 328 13432 2648 74691 27742 33905 310 4563 198 48055 5789 13 271 2592 1017 513 777 65269 5789 1056 1017 513 15440 5555 198 48093 11 264 328 77318 1 4531 369 8953 11 7014 264 15440 5555 198 940 279 18890 15171 15441 1 369 2945 26 303 411 1105 11 421 15440 198 15704 20539 264 13992 8222 279 13230 65269 5789 198 77167 449 4147 13992 413 1017 998 874 13230 65269 5789 553 271 2592 866 15698 5515 1503 524 7618 310 264 15440 5555 198 591 11 264 328 77318 1 4531 369 8953 11 7014 264 15440 5555 198 940 279 18890 65667 15441 1 369 2945 26 303 411 1105 11 421 15440 198 15704 20539 264 10675 8222 279 13230 15698 198 15889 318 940 279 20209 430 6679 321 279 5515 2663 430 198 5816 6604 286 2663 681 466 264 318 902 8 4147 10675 413 1017 998 874 198 325 1084 15698 5789 13 271 2592 279 18890 15171 27188 1 7701 303 279 709 1562 11 328 27188 1 198 24028 14989 310 449 348 1947 470 19315 220 33027 494 1439 5199 4635 513 198 83 2757 430 413 781 998 4945 65269 5789 13 220 1690 279 1562 198 92845 2007 16 11 830 17 11 348 88 11 830 18 11 830 19 11250 413 348 88 9 64013 310 264 8240 348 88 16 11961 4444 11 348 88 44 11961 198 559 369 13185 310 264 1562 440 380 10 19 65269 5789 348 87 16 11961 348 87 17 11961 198 32144 16 11961 4444 11 348 88 44 11961 348 87 18 11961 348 87 19 19315 271 32 27735 314 411 369 421 7643 279 15171 27188 1 18890 1189 198 39420 348 10378 9 11135 15698 5789 11 424 369 14789 348 14372 9 198 1719 15698 5789 318 429 866 65667 27188 1 5789 1317 1436 3559 553 198 4272 25 271 256 11765 687 281 2784 11 292 1590 198 256 2423 257 1133 2784 11 292 8 198 256 2423 198 256 11765 281 1818 28 16 11 12492 17 11 578 198 256 220 17 220 16 198 256 11765 281 2784 28 16 11 12492 17 11 578 198 256 26289 1371 318 3444 3110 1562 1483 1590 198 257 2793 3919 49863 20507 1500 220 16 11 303 361 4211 29 198 256 24237 25 281 363 2597 5081 2663 364 15698 5515 359 64 6 198 256 11765 281 7 16 11 12492 17 11 578 198 256 220 16 220 17 271 2064 369 17957 364 2107 15698 5789 321 279 15171 27188 1 18890 198 951 381 1429 303 279 1788 1562 11 748 303 6380 411 20684 1503 524 198 58768 29783 13 271 2592 279 18890 65667 27188 1 7701 303 279 709 1562 11 198 1 27188 1 1902 14989 310 264 348 39446 11961 279 8529 314 864 513 198 83 2757 430 4945 15698 5789 13 1368 264 5555 12219 264 1328 198 4497 2582 978 2574 264 869 318 1634 449 11135 15698 5515 11 466 198 1445 2361 30047 286 681 264 328 77318 1 4531 369 8953 13 271 4350 65667 27188 1 369 1429 11 1754 1328 303 411 12366 1902 381 264 198 889 13 8618 869 494 279 12366 369 12244 310 279 1118 15440 198 15704 16515 364 15698 15840 6483 803 369 5951 310 279 198 767 13 357 1328 1144 524 381 264 12654 12448 318 68 1257 13 1535 2721 65884 35087 37 2918 369 198 35788 11 7643 424 668 524 2353 866 15440 5555 421 1362 198 1323 13866 553 1368 1017 369 874 2353 310 264 15440 5555 279 1328 18509 198 12306 369 14129 539 279 65667 1 5555 11 413 1017 369 799 11 466 413 1017 198 284 524 11 264 328 77318 1 4531 369 8953 13 271 1774 278 4868 1608 279 18890 15171 15441 1 466 65667 15441 1 198 32157 381 1429 430 65269 5515 15250 466 430 15698 5515 198 11072 13 271 5221 303 2243 220 18 13 20 25 5534 6526 4054 866 1324 314 63669 321 198 175267 1 30047 793 11 65269 5789 1189 1732 49124 30047 793 198 28625 3842 321 15698 5789 1189 1732 10675 30047 793 3376 332 1764 198 36030 10660 539 2972 1679 47 220 19 19 23 159034 271 32 1562 2590 4523 1010 869 11 10451 328 3928 487 7014 424 24096 449 198 7557 13 220 2500 411 869 369 23424 13402 383 279 913 314 279 198 91907 1576 13 271 2592 424 369 2218 271 64 1156 37435 709 25 198 256 561 1970 2424 364 279 709 369 15234 11 11946 424 279 198 256 5515 1103 13 220 561 1118 3065 279 1970 2424 668 635 369 10403 279 198 256 15440 4868 310 279 5789 26 411 369 7246 303 3644 198 256 5534 16954 13 220 3095 279 1970 2424 49703 264 328 671 1 198 256 4952 11 411 28151 279 460 869 314 279 709 1562 13 271 64 5617 3309 709 466 1654 25 198 256 561 1067 369 685 310 279 37999 26 1436 31959 3309 22815 364 279 198 256 26884 314 5617 3309 5568 321 5245 13 271 64 523 1576 25 198 256 357 491 2774 314 421 523 369 5743 13 271 64 523 2774 1654 25 198 256 561 11811 1156 37435 709 369 2512 11 440 449 5515 198 256 1103 421 369 799 4860 1056 279 5515 1103 314 279 1562 25 279 198 256 2774 8772 279 1118 5515 13 271 64 523 2774 25 198 256 561 523 1902 6762 264 12880 6450 548 33905 1654 26 279 2376 369 1179 279 198 256 1788 430 413 421 1654 557 2512 13 1358 1890 16954 198 865 9 271 32 523 7044 18091 264 523 1576 318 3924 3644 561 5129 198 1267 27980 1590 271 256 523 727 262 61987 498 91430 60 328 1005 1 68259 498 1267 6533 60 498 12457 663 60 22698 15705 198 256 40443 61987 32467 498 14064 1951 60 56439 198 256 68259 256 61987 12448 271 32 523 7044 369 449 31097 4952 13 220 561 40443 1103 198 41547 6486 264 1103 314 2255 6632 318 3924 6017 565 32748 364 777 198 40400 5533 681 748 1754 1455 303 279 1103 1220 14989 310 264 523 198 1640 864 5954 35959 286 13 220 32749 1973 449 40443 1103 198 12457 11 539 1581 11 494 279 2255 523 328 1640 4961 15612 11 271 256 523 32320 25 198 950 1440 271 284 13185 310 271 256 523 32320 4253 1590 198 950 1440 271 760 523 725 15705 369 1179 15234 303 264 491 10993 3899 318 3924 198 82459 321 10649 681 1608 264 13238 3354 2136 4325 321 279 198 9596 3521 4325 13 318 68253 11 279 15705 5435 9712 198 1628 16954 5946 220 3095 279 523 725 15705 31941 10993 11 1141 198 50947 3899 369 41784 694 1141 2136 4325 369 6568 13 498 20 60 357 198 1005 1576 369 1179 3354 1608 279 40443 1103 364 279 2255 198 8588 321 279 6568 2136 4325 364 279 6816 10675 13 198 760 523 803 369 6608 310 411 523 1576 303 279 3889 2136 198 2159 13 271 760 1906 303 864 7944 513 4364 303 279 523 2407 369 198 22885 2683 303 279 491 523 725 12880 8217 548 3158 220 6817 421 411 369 14294 198 3138 1245 1238 279 523 369 3354 321 1132 364 6632 421 998 198 9423 1608 279 7044 18890 13 271 1890 9400 628 381 30650 16421 1608 2197 565 32748 13 271 19465 628 1048 381 34816 25 1066 1040 948 46963 5568 11 271 256 554 69 16 9125 8 198 256 554 69 17 198 256 523 32320 25 1440 271 284 16757 13185 310 271 256 523 32320 25 1440 198 256 32320 283 281 16 9125 2299 69 17 7584 2537 578 271 760 15977 5426 364 279 48973 22666 513 279 1788 430 364 198 1628 60565 13 220 561 1067 369 1179 6608 310 279 523 803 13 271 5221 303 2243 220 18 13 24 25 32749 1189 381 34816 440 866 2610 198 1 28973 27160 3158 56703 11 279 30394 557 1691 777 198 26993 520 26 1436 2972 1679 47 220 21 16 19 332 364 3447 13 271 32 1103 314 913 4868 1189 381 2574 303 9057 37641 6849 198 10378 279 523 725 803 13 1061 14378 310 1064 913 1716 382 421 198 1719 523 369 13553 13 2332 15135 11 279 913 4868 628 381 29413 198 1445 279 523 725 12880 1267 6533 548 1 6816 13 3382 20627 6632 364 198 6183 13 271 5221 303 2243 271 14162 25 3956 310 279 47801 3294 11 1092 1503 449 1992 4952 635 948 12654 8213 440 24460 383 11 321 1092 1503 449 54109 15840 14989 1132 2957 30 21134 25899 13 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/long3k.ids b/zig/tests/cuda/flashnext/lr7/long3k.ids new file mode 100644 index 0000000000..04e5bd5759 --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/long3k.ids @@ -0,0 +1 @@ +248045 846 198 36948 369 449 47801 314 279 12654 3992 5604 13 271 760 328 2138 1 4952 198 865 6610 271 8272 11889 513 264 16219 1560 310 5479 26805 52992 198 17583 264 1957 25 271 256 1992 21152 61987 328 2138 1 7258 498 2129 7258 60 271 760 4145 1304 11 328 2138 7258 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 25 4673 63990 271 760 11244 1304 11 328 2138 7258 16 11 7258 17 487 369 13185 310 271 256 413 1259 8090 77573 198 950 413 524 7258 16 25 4673 63990 46718 17 8 271 9205 42018 2303 9370 421 12880 8090 548 1 321 328 66317 1401 1 8042 198 951 279 5617 3309 7101 440 1782 4874 13 220 733 279 1428 198 13957 11 279 5617 3309 3759 12880 8090 548 1 369 328 2434 1 1172 198 7994 12779 11 328 3913 1 948 24460 369 10897 318 5454 198 1021 2904 6320 46 1764 220 561 1428 1970 13427 70307 874 1970 364 449 198 2138 4952 948 24460 369 10897 506 18612 854 13 220 6817 198 8761 424 369 24366 310 2830 279 2450 1970 364 279 7258 198 8761 4490 303 279 1412 1876 26 424 668 381 12234 430 919 314 279 198 7450 11319 13 271 27991 1320 310 12880 8090 548 1 513 11477 13 220 561 869 364 279 5617 3309 198 9521 369 10521 948 279 37999 8211 13 1358 40238 11889 198 865 17731 271 40238 11889 513 1429 310 318 265 8 7423 4874 310 2663 321 310 198 32925 7944 466 3470 314 24983 5978 25 271 256 15840 21152 61987 318 5480 1951 35036 7030 318 11531 1114 27160 735 7302 27160 8 198 256 2100 1951 257 61987 2100 318 2129 2100 4653 498 487 1293 198 256 2100 670 61987 12448 198 1025 735 32467 498 5480 1951 60 56439 198 1025 735 48975 498 5480 1951 60 84576 198 1025 735 16819 330 42514 198 1025 735 14701 198 1025 735 81987 198 1025 735 63669 2100 271 7 9538 3644 34019 5262 364 279 18890 16954 364 348 42648 330 42514 11961 198 9 33261 11961 321 348 82 86047 9 5946 271 2014 15840 4952 64013 279 7258 1103 318 28321 421 198 559 628 381 264 3074 7258 466 264 30644 70221 1103 11 279 14827 198 28726 286 264 13992 8 321 48233 279 3074 12571 1576 310 1754 314 198 1719 2100 11140 11 494 2047 310 1245 13 271 40238 369 4364 51070 11313 383 279 1304 314 279 2100 198 7080 553 3095 264 2100 369 919 314 264 24983 1576 318 276 6816 198 15812 11 14701 466 81987 681 279 24983 1576 1902 198 484 6887 2649 279 15840 321 9973 883 1141 30794 11 321 198 17812 4673 449 4531 413 279 15840 369 41565 13 220 561 5426 198 5308 2683 539 5090 4346 321 279 18987 8953 513 2574 440 279 198 17829 314 279 1576 4346 318 3924 3644 561 5129 913 198 71 32759 553 271 40238 314 449 1576 310 264 2100 1103 11 44007 42360 303 198 39708 37052 466 9057 37641 11 369 51070 4364 430 10695 13 271 9 1368 279 2100 1103 369 264 3074 2100 440 874 26849 30644 11 198 220 44007 303 71458 11 279 1576 369 12244 310 421 2100 13 271 9 18236 25 271 220 348 1368 279 2100 1103 5435 799 2100 78685 440 449 32815 3086 11 198 262 2512 264 1001 11531 1114 828 2100 25 561 1576 1902 381 449 49124 440 506 198 262 3140 430 1599 3470 430 1017 513 11218 303 279 2100 1103 11 26401 198 262 799 13 220 561 1118 3470 314 279 49124 513 12244 11 494 2047 310 198 262 1245 11 310 279 11218 1518 279 56276 2100 13 220 561 1534 3470 198 262 314 279 49124 513 12244 310 279 11218 1238 279 56276 198 262 2100 13 220 357 1103 314 279 9376 3470 303 279 49124 369 1179 198 262 12244 310 279 56276 2100 318 1719 1103 628 381 4147 553 271 220 348 18236 25 561 1576 1902 381 449 49124 440 279 1788 1324 314 3470 198 262 430 1017 513 11218 303 279 2100 1103 11 321 279 3470 513 198 262 12244 11 494 2047 310 1245 11 310 279 11811 11218 13 271 40238 314 449 1576 310 264 3074 2100 369 51070 4364 430 198 18000 82 13 271 9 1368 279 2100 369 449 12448 318 591 1590 271 220 348 1368 279 803 1503 524 11822 303 264 328 9461 1 466 328 6082 2359 1 4952 198 262 303 279 1428 1970 2424 25 279 803 369 6608 310 279 1576 303 279 198 262 1428 2136 4325 13 271 220 348 17673 25 279 803 369 6608 310 279 1576 303 279 3521 4325 198 262 466 279 15484 4325 10521 539 328 6082 2359 487 15119 13 271 220 561 803 369 40158 413 424 557 2582 6608 13 220 1061 1189 5073 279 198 220 5604 1698 364 279 1576 8334 6608 310 279 803 310 5372 198 220 6942 11 13969 279 1576 310 381 3372 38157 321 1141 44462 318 331 424 198 220 682 799 8 310 381 2512 13 271 9 1368 279 2100 369 449 6816 5604 25 561 5839 7258 303 198 220 279 5604 369 24276 13 220 1049 1220 7302 449 1576 440 198 220 86370 7944 26 413 411 369 524 279 1105 11 328 77318 1 369 198 220 8953 13 220 2844 1576 369 1179 4439 310 9501 279 12244 1576 310 198 220 279 2574 6816 26 413 424 4021 2649 279 15840 11 424 24096 198 220 449 4531 318 41547 694 524 13902 328 3776 1401 1764 271 220 6817 25 1368 279 1576 369 264 523 2774 321 279 6816 5604 198 220 13273 383 2107 10744 314 279 15840 5497 11 279 1245 23641 3008 198 220 7258 11 328 64 1926 1 628 2528 2892 449 2774 6816 466 318 331 874 198 220 2774 6816 6513 8 264 523 6816 13 220 561 2047 23641 3008 198 220 2100 328 64 1926 1 369 2590 716 430 449 2774 6816 11 6611 424 413 198 220 5689 13 220 13891 11 279 1330 54133 314 328 64 1926 1 635 524 13902 198 220 8042 310 279 1788 6816 25 413 279 1245 23641 3008 7258 198 220 18675 310 264 523 6816 11 279 2047 23641 3008 11121 264 491 198 220 2774 6816 430 279 2100 314 279 15840 25 271 257 523 2356 82 25 198 260 830 283 220 18 1734 653 523 3759 198 257 1699 283 2356 82 363 198 257 1699 1926 283 1699 1926 478 220 16 256 653 13514 1699 1926 430 220 19 9101 2356 82 1926 430 220 18 271 220 1061 3874 1503 524 13902 3668 310 15933 198 220 7944 11 1680 430 5706 3354 440 328 3573 363 3158 271 9 1368 279 2100 369 264 14701 25 561 5839 7258 303 279 198 220 5604 369 24276 13 220 1049 1220 7302 2892 264 24983 8240 198 220 1576 318 20171 430 264 1103 8 466 264 12366 1576 318 20171 430 264 10675 553 198 220 9019 11 279 94860 7258 369 24276 13 271 220 1368 279 5839 369 264 24983 8240 1576 318 20171 430 264 1103 681 279 198 220 94860 1902 7302 449 7308 13 220 1368 424 369 7968 11 279 8240 725 198 220 2984 369 3568 310 424 13 220 561 12571 869 1902 381 264 2397 40823 198 220 7308 2599 1056 279 8240 725 2984 11 321 279 8240 369 4439 198 220 310 9501 279 12244 1576 310 1141 1455 440 421 1857 13 220 1368 279 198 220 1857 369 680 314 2020 11 328 1498 1401 1 369 8953 318 28973 310 264 198 220 94860 290 8240 4021 884 491 3470 310 264 1103 553 271 220 1368 279 5839 369 264 12366 1576 318 20171 430 264 10675 681 279 198 220 94860 1902 599 264 913 17606 440 279 12366 725 1328 913 11 198 220 321 279 12366 369 1179 4439 310 1791 264 1328 55616 6505 864 13640 198 220 279 94860 310 279 12244 1576 13 220 1061 628 2892 8032 449 198 220 6149 1328 55616 6505 440 279 1788 1328 869 11 466 5479 264 491 198 220 1328 55616 6505 318 331 874 1328 440 279 1788 869 23533 553 271 220 1690 1156 37435 5978 11 279 12880 723 1163 548 33905 1654 369 2512 440 198 220 8053 5789 13 271 9 1368 279 2100 369 264 81987 25 561 5839 7258 303 279 5604 198 220 369 24276 13 220 1049 1220 7302 264 24983 8240 1576 318 20171 430 264 198 220 1103 553 220 561 12244 1576 1220 381 264 8240 1576 314 279 1788 198 220 913 13 220 9019 11 279 4570 321 8156 6608 22666 513 24276 11 198 220 303 90349 781 513 2945 26 16184 513 6942 321 279 8240 725 198 220 2984 13 220 561 13854 1220 14989 310 24959 13 1368 2892 6608 369 198 220 7968 11 279 8240 725 2984 369 3568 310 424 13 220 561 12571 198 220 13854 513 70575 310 9940 1881 6942 321 279 8240 725 2984 11 198 220 27389 13 220 16864 11 279 8240 1576 369 4439 310 8032 279 198 220 15511 440 279 3470 314 279 12244 8240 13 220 561 2984 314 279 198 220 15511 1189 381 2086 494 279 2984 314 279 12244 8240 11 198 220 8190 9722 279 2984 314 279 2100 8240 11 413 279 2100 198 220 8240 5954 424 13 271 332 6901 26033 7873 7472 64700 733 279 1428 7873 11 279 198 54304 364 11218 369 4284 310 381 279 1788 430 364 22666 11 321 198 11469 18890 369 17030 2261 279 1970 9191 9956 11 13969 198 1657 11346 1412 6397 13 271 15337 279 7044 314 15840 23197 421 73406 1881 198 1719 2047 23641 3008 321 279 1245 23641 3008 513 3260 14352 484 17140 515 318 1891 198 8422 328 64 11 292 283 292 11 264 1 70664 1330 7101 681 73406 348 53977 9 279 198 13034 314 12244 4534 7101 11822 2047 4534 6491 11 6807 198 1334 286 303 20684 13 220 1690 2774 11 279 2614 1957 22741 198 35643 15 11 220 17 60 763 271 256 830 283 498 15 11 220 16 60 198 256 585 283 220 15 198 256 585 11 830 957 60 283 220 16 11 220 17 260 653 585 369 5860 11 1179 830 957 60 369 5860 198 256 1133 2007 8 271 9538 1048 25 271 220 2972 1679 47 220 18 16 18 17 332 471 39226 36760 1188 91841 198 257 561 24332 364 279 15171 5480 1 4416 13 1358 21199 26115 15840 11889 198 1464 43414 271 21199 26115 15840 369 279 10286 11 303 264 3074 4952 11 314 264 198 25064 5488 321 449 15840 4952 25 271 256 54109 50145 21152 61987 18720 5480 18720 445 318 27188 1951 735 7302 27160 8 198 256 18720 5480 336 61987 12448 735 16819 330 42514 735 14701 735 81987 198 256 18720 445 3696 61987 6421 420 735 6320 420 735 15171 420 735 8168 420 735 3405 420 735 34397 420 735 5775 420 735 65667 420 198 1734 735 328 2372 420 735 14291 420 735 13012 420 735 38591 420 735 24403 420 271 7 9538 3644 34019 5262 364 279 18890 16954 314 279 1483 2250 198 65151 5946 271 2014 54109 15840 64013 279 2100 318 7949 11 19231 4472 198 28973 11889 11 4021 381 449 30047 286 8 321 279 7258 198 1551 11 25822 279 7620 5488 3050 310 279 913 314 15840 198 263 279 1330 52680 11 321 48233 279 1067 310 279 3889 2100 13 198 760 2100 369 1132 24276 2957 13 271 2014 54109 15840 7258 1040 328 87 1373 220 16 1 628 381 56546 430 198 63230 283 830 478 220 16 1 310 10752 264 4283 11 694 524 6681 5951 2376 13 733 279 198 18599 26115 2243 11 328 87 1 369 1132 24276 2957 13 7054 11 948 3102 11 198 1719 4880 5488 369 10345 348 258 40284 11961 7063 421 4598 1056 198 43408 264 491 1576 321 58706 421 310 279 2100 11 279 2235 1576 198 284 10490 4370 13 271 41657 4472 30152 11 54109 30152 14989 279 2047 12 198 10346 3008 348 14372 9 36319 279 1245 23641 3008 13 220 1690 3010 11 328 64 957 60 198 8236 281 2007 9646 1118 5686 5072 328 64 957 18503 1179 424 64013 328 69 2007 9646 321 25822 198 1719 5089 11 321 1483 391 11 424 13514 279 1067 1142 310 328 64 957 60 3158 271 2278 279 4531 314 58706 310 43716 321 5081 11218 303 264 198 14892 4952 11 279 15840 2725 539 54109 15840 198 91318 369 17083 279 1788 1560 430 4472 30152 13 33105 11 198 4058 279 4531 314 279 3102 348 258 40284 9 7465 11 279 7620 198 8987 10345 539 54109 15840 369 279 1788 430 279 4472 198 25064 7289 13 271 2381 11218 864 513 6816 14619 11 279 1788 83081 883 198 1005 321 2774 7944 16298 430 364 5611 30152 13 1358 2014 84027 15840 11889 198 1464 43414 271 9 19109 9 15840 369 279 10286 11 303 264 3074 4952 11 314 198 64 3759 466 6816 20572 321 449 9801 15840 198 23425 25 271 256 60736 50145 21152 61987 18720 5480 22698 7258 198 761 4241 420 318 11531 1114 27160 735 7302 27160 7025 271 760 6463 494 4472 33284 11889 369 421 1132 264 3074 198 5480 369 5252 13 271 2381 4145 4874 430 15840 11218 11 413 303 523 466 4536 6676 11 198 1719 31144 513 24276 321 9476 303 264 3175 523 466 4536 198 8843 12880 38314 548 1 421 369 264 10675 12366 494 3759 198 11072 318 76 37652 413 843 8 310 24276 31144 13 1061 6816 369 198 86 21833 321 369 9184 3354 506 279 1151 314 523 466 4536 198 2513 10993 11 413 31144 513 1669 88034 13 271 2381 22666 430 15840 11218 11 279 31144 513 24276 198 331 303 523 466 4536 6676 11 694 524 9476 13 271 2592 264 803 369 60736 303 264 709 6676 11 1179 411 803 369 2136 198 1891 421 6676 13 90117 513 2496 24276 321 9476 303 709 198 78134 13 271 2592 279 1245 1375 3008 369 2945 11 449 60736 15840 25822 198 1719 4880 15840 1518 36319 31144 318 2777 198 658 44575 553 1368 279 1245 1375 3008 369 524 2945 364 449 7258 198 5480 11 1179 279 37999 64013 279 2100 3525 364 279 1483 198 1 548 723 1163 548 33905 466 12880 84163 548 33905 1562 13 271 9538 1048 25 271 220 2972 1679 47 220 20 17 21 332 471 31058 364 12051 90117 198 257 561 13340 421 3568 18890 364 34848 1060 279 4346 314 198 257 7101 318 15693 523 7101 321 2774 7101 681 198 257 4370 314 35493 1070 1472 5853 13 271 220 2972 1679 47 220 19 23 19 332 471 3856 29642 198 257 561 13340 421 3568 279 328 86711 1 4536 310 3300 264 5129 198 257 18890 364 913 31144 421 628 381 1429 303 1064 6157 198 257 7141 321 27069 82 13 271 5221 303 2243 220 18 13 23 25 4543 60736 30152 2070 279 1788 198 13392 880 303 279 1245 1375 3008 430 5611 30152 13 56703 11 198 14265 22666 318 4650 632 2195 8878 80019 13992 22666 8 8613 264 198 54304 1412 13 1358 10266 27693 198 584 332 271 3446 303 2243 220 18 13 20 13 1358 31099 709 7044 198 1464 34195 271 256 3203 9310 727 61987 498 91430 60 328 7449 1 328 727 1 2821 591 32467 498 15704 1951 60 56439 198 3696 4241 397 1 7258 60 22698 15705 271 19678 314 12654 1766 27693 628 381 20950 321 55483 506 1599 198 7462 318 3924 348 5816 13754 187525 328 11093 1 22666 11 328 7449 364 1 321 328 7449 198 4058 1 628 1132 381 1429 303 279 2407 314 264 75140 709 13 271 24407 4364 440 328 7449 687 1 18890 513 2590 75140 198 20752 11 1442 413 781 635 524 6435 328 11093 1 466 328 7449 1 20209 13 271 2064 369 264 328 32769 1401 1 310 958 264 328 28726 494 1 7258 4613 279 2407 198 1020 264 75140 709 13 271 2014 3010 314 264 75140 709 25 271 256 3203 687 2821 9453 16 11 1626 17 1590 198 950 635 92099 363 198 950 2454 1010 14308 13754 363 271 5221 303 2243 220 271 14162 25 3956 310 279 47801 3294 11 1092 1503 449 1992 4952 635 948 12654 8213 440 24460 383 11 321 1092 1503 449 54109 15840 14989 1132 2957 30 21134 25899 13 248046 198 248045 74455 198 248068 271 248069 271 diff --git a/zig/tests/cuda/flashnext/lr7/results-sparse.json b/zig/tests/cuda/flashnext/lr7/results-sparse.json new file mode 100644 index 0000000000..1a4198200a --- /dev/null +++ b/zig/tests/cuda/flashnext/lr7/results-sparse.json @@ -0,0 +1,364 @@ +{ + "info": { + "startup_s": 95.44, + "max_len": 16384, + "prefill_rows": 4096, + "nch": 5, + "qsa": true, + "idw": 2052, + "eos": [ + 248046, + 248044 + ], + "torch": "2.13.0a0+9186a08b2c.nv26.07", + "gpu": "NVIDIA GB10", + "prompt_lengths": { + "long3k": 3001, + "long12k": 12001 + } + }, + "results": { + "serial/long3k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1510, + 2138, + 63, + 4952, + 70307, + 874, + 1970, + 318, + 275, + 369, + 13057, + 11693, + 553, + 733, + 449, + 54109, + 15840 + ], + "wall_s": 16.7331, + "prefill_s": 2.4387, + "cached": 0, + "drafts": false, + "decode_s": 14.2676, + "rounds": 31, + "decode_tps": 2.17 + }, + "teacher32/long3k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1970, + 13427, + 70307, + 2972, + 2083, + 1970, + 332, + 364, + 449, + 1510, + 2138, + 63, + 4952, + 13, + 271, + 623, + 449 + ], + "wall_s": 86.062, + "decode_ms_per_token": 30.764, + "top5": [ + [ + [ + 4350, + 21.25 + ], + [ + 2278, + 18.25 + ], + [ + 16, + 18.0 + ], + [ + 27775, + 17.625 + ], + [ + 2014, + 17.375 + ] + ], + [ + [ + 12654, + 22.375 + ], + [ + 24460, + 20.125 + ], + [ + 4162, + 19.125 + ], + [ + 1542, + 16.125 + ], + [ + 32237, + 15.75 + ] + ], + [ + [ + 8213, + 24.75 + ], + [ + 369, + 19.5 + ], + [ + 1542, + 13.875 + ], + [ + 24460, + 13.25 + ], + [ + 4162, + 13.1875 + ] + ], + [ + [ + 440, + 23.875 + ], + [ + 303, + 13.9375 + ], + [ + 383, + 12.75 + ], + [ + 4058, + 12.5625 + ], + [ + 24460, + 12.375 + ] + ] + ] + }, + "serial/long12k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1510, + 2138, + 63, + 4952, + 70307, + 2972, + 2083, + 1970, + 332, + 318, + 275, + 369, + 13057, + 33993, + 14, + 56497, + 553 + ], + "wall_s": 6.012, + "prefill_s": 5.117, + "cached": 0, + "drafts": false, + "decode_s": 0.8949, + "rounds": 31, + "decode_tps": 34.64 + }, + "teacher32/long12k": { + "tokens": [ + 4350, + 12654, + 8213, + 440, + 24460, + 383, + 318, + 940, + 279, + 90563, + 46, + 63, + 5016, + 681, + 279, + 1970, + 13427, + 70307, + 2972, + 2083, + 1970, + 332, + 364, + 449, + 1510, + 2138, + 63, + 4952, + 13, + 271, + 623, + 449 + ], + "wall_s": 338.3455, + "decode_ms_per_token": 30.42, + "top5": [ + [ + [ + 4350, + 21.375 + ], + [ + 27775, + 19.375 + ], + [ + 10865, + 18.125 + ], + [ + 2278, + 18.125 + ], + [ + 16, + 18.0 + ] + ], + [ + [ + 12654, + 22.5 + ], + [ + 24460, + 19.75 + ], + [ + 4162, + 19.25 + ], + [ + 1542, + 15.9375 + ], + [ + 32237, + 15.5 + ] + ], + [ + [ + 8213, + 24.375 + ], + [ + 369, + 19.25 + ], + [ + 1542, + 13.5625 + ], + [ + 4162, + 12.9375 + ], + [ + 725, + 12.9375 + ] + ], + [ + [ + 440, + 23.125 + ], + [ + 303, + 13.3125 + ], + [ + 383, + 12.5625 + ], + [ + 24460, + 12.5 + ], + [ + 4058, + 12.4375 + ] + ] + ] + } + } +} diff --git a/zig/tests/cuda/flashnext/pack.sh b/zig/tests/cuda/flashnext/pack.sh new file mode 100755 index 0000000000..8e3c27d874 --- /dev/null +++ b/zig/tests/cuda/flashnext/pack.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +# Usage: pack.sh OUT_KERNELS CAPTURE_OUT...: one kernel set from the captures' manifest.json and jit.json. +set -euo pipefail +OUT=$1; shift +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +ARGS=() +for d in "$@"; do ARGS+=(--manifest "$d/manifest.json" --jit "$d/jit.json"); done +python3 -B "$REPO/tools/zig/triton_aot_pack.py" "${ARGS[@]}" --out "$OUT" From 035b7f96efe826562e58012562c4e9c2b240be77 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:15:58 -0500 Subject: [PATCH 034/115] Flash Next CUDA: the MTP head, its draft chains, verify windows and the DeltaNet rollback The head (cuda_mtp.zig, as the Python engine's mtp.py and decode.draft) reads the next token's embedding and the main model's four streams after the last MoE write-back, runs its full-attention layer with its own hyper-connections, experts and KV cache, then its mixer and the lm_head rows of the 79,591-id draft vocabulary. A chain keeps its first draft and stops before a later one whose temperature-1 probability falls under the confidence (6 drafts and 0.60 by default; TF_FLASHNEXT_DRAFTS and TF_FLASHNEXT_CONFIDENCE). The head runs over a prompt a chunk at a time beside the main pass (Ops.prefillChunk), so it keeps only a chunk's streams, not the whole prompt's. A verify window (the pending token and the drafts) is one batched pass whose rows keep the bits they have alone, each row drawn at its own keyed position with the draw serial decoding makes. The window runs the DeltaNet chain out of place: kept whole, its end states swap in; cut, the kept rows replay from the window's start states and the projection rows the pass kept, the conv windows and the n-gram tail shift by the kept rows alone. capture_mtp.sh records the head's own Triton forms. flashnext-prompt with TF_FLASHNEXT_DRAFTS decodes each prompt serially and drafted and compares them token for token (TF_FLASHNEXT_SEED for keyed sampled draws). The server does not use the head yet. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/flashnext_mtp.cu | 70 ++++ zig/src/cuda/kernels.zig | 1 + zig/src/families/flashnext/cuda_launch.zig | 4 + zig/src/families/flashnext/cuda_mtp.zig | 293 +++++++++++++++++ zig/src/families/flashnext/cuda_mtp_check.zig | 171 ++++++++++ zig/src/families/flashnext/cuda_mtp_draft.zig | 178 +++++++++++ zig/src/families/flashnext/cuda_mtp_ops.zig | 300 ++++++++++++++++++ .../families/flashnext/cuda_prompt_main.zig | 33 ++ zig/tests/cuda/flashnext/capture_mtp.py | 64 ++++ zig/tests/cuda/flashnext/capture_mtp.sh | 11 + 11 files changed, 1126 insertions(+) create mode 100644 zig/kernels/cuda/flashnext_mtp.cu create mode 100644 zig/src/families/flashnext/cuda_mtp.zig create mode 100644 zig/src/families/flashnext/cuda_mtp_check.zig create mode 100644 zig/src/families/flashnext/cuda_mtp_draft.zig create mode 100644 zig/src/families/flashnext/cuda_mtp_ops.zig create mode 100644 zig/tests/cuda/flashnext/capture_mtp.py create mode 100755 zig/tests/cuda/flashnext/capture_mtp.sh diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 7d759a2715..89f4d77503 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -17,6 +17,7 @@ const kernels = [_]Kernel{ .{ .name = "gdn", .flags = &.{ "-O3", "--fmad=false" } }, // cuda/kernels/gdn.py, tensorfold_gdn_v2 .{ .name = "flashnext_gdn", .flags = &.{ "-O3", "--fmad=false" } }, // qwen4_exp/cuda/gdn.cu chain, one row .{ .name = "flashnext_batch", .flags = &.{ "-O3", "--fmad=false" } }, // ours: Flash Next's batched prompt pass glue + .{ .name = "flashnext_mtp", .flags = &.{"-O3"} }, // ours: Flash Next's MTP drafts, a row's best score and log-sum-exp .{ .name = "probe", .flags = &.{"-O3"} }, .{ .name = "qmm_group", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py, tensorfold_qmm_v5 .{ .name = "qmm", .flags = &.{"-O3"} }, // cuda/kernels/qmm.py lane matmul, groups of 32 diff --git a/zig/kernels/cuda/flashnext_mtp.cu b/zig/kernels/cuda/flashnext_mtp.cu new file mode 100644 index 0000000000..1ac2e6fa4b --- /dev/null +++ b/zig/kernels/cuda/flashnext_mtp.cu @@ -0,0 +1,70 @@ +// Flash Next's MTP drafts (cuda_mtp.zig): a score row's best value, its first id and the row's log-sum-exp. + +#include +#include +#include + +namespace fn_mtp { + +static __device__ __forceinline__ void better(float& v, int& i, float ov, int oi) { + if (ov > v || (ov == v && oi < i)) { + v = ov; + i = oi; + } +} + +} // namespace fn_mtp + +// Program (row), 1024 threads: out [rows, 4] fp32 holds the best of a row's first V values, its log-sum-exp, its id, 0. +extern "C" __global__ void __launch_bounds__(1024) fn_row_best(const __nv_bfloat16* __restrict__ L, int V, long long ld, + float* __restrict__ out) { + using namespace fn_mtp; + const int r = blockIdx.x, t = threadIdx.x, lane = t & 31, warp = t >> 5; + const __nv_bfloat16* row = L + static_cast(r) * ld; + float v = -CUDART_INF_F; + int at = 0x7fffffff; + for (int i = t; i < V; i += blockDim.x) { + const float x = __bfloat162float(row[i]); + if (x > v) { + v = x; + at = i; + } + } + for (int o = 16; o; o >>= 1) better(v, at, __shfl_xor_sync(0xffffffffu, v, o), __shfl_xor_sync(0xffffffffu, at, o)); + __shared__ float sv[32]; + __shared__ int si[32]; + __shared__ float top; + __shared__ int top_at; + if (lane == 0) { + sv[warp] = v; + si[warp] = at; + } + __syncthreads(); + if (warp == 0) { + v = lane < (int)(blockDim.x >> 5) ? sv[lane] : -CUDART_INF_F; + at = lane < (int)(blockDim.x >> 5) ? si[lane] : 0x7fffffff; + for (int o = 16; o; o >>= 1) better(v, at, __shfl_xor_sync(0xffffffffu, v, o), __shfl_xor_sync(0xffffffffu, at, o)); + if (lane == 0) { + top = v; + top_at = at; + } + } + __syncthreads(); + const float m = top; + float sum = 0.0f; + for (int i = t; i < V; i += blockDim.x) sum += expf(__bfloat162float(row[i]) - m); + for (int o = 16; o; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o); + __shared__ float ss[32]; + if (lane == 0) ss[warp] = sum; + __syncthreads(); + if (warp == 0) { + sum = lane < (int)(blockDim.x >> 5) ? ss[lane] : 0.0f; + for (int o = 16; o; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o); + if (lane == 0) { + out[r * 4 + 0] = m; + out[r * 4 + 1] = m + logf(sum); + out[r * 4 + 2] = __int_as_float(top_at); + out[r * 4 + 3] = 0.0f; + } + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 791d394ce0..72868085a0 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -15,6 +15,7 @@ pub const available = options.with_kernels; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; +pub const flashnext_mtp: []const u8 = if (available) &Blob("fatbin_flashnext_mtp").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; pub const qmm_group: []const u8 = if (available) &Blob("fatbin_qmm_group").bytes else &.{}; pub const qmm: []const u8 = if (available) &Blob("fatbin_qmm").bytes else &.{}; diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 3f4bfcfe46..f040561f86 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -17,6 +17,8 @@ pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); /// The row-by-row reference path, for flashnext-prompt pub const row_path = @import("cuda_rows.zig"); +/// The MTP head's drafts: its forward, the chain, the verify window's rows and their rollback +pub const mtp = @import("cuda_mtp.zig"); pub const native = @import("cuda_lanes.zig").Namespace(Tri); const p = aot.ptr; @@ -317,9 +319,11 @@ test { _ = mix; _ = prompt; _ = prompt.batch; + _ = mtp; _ = native; _ = @import("cuda_config.zig"); _ = row_path; + _ = @import("cuda_mtp_check.zig"); } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_mtp.zig b/zig/src/families/flashnext/cuda_mtp.zig new file mode 100644 index 0000000000..73cd614672 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -0,0 +1,293 @@ +//! Flash Next's MTP drafts on CUDA (Python qwen4_exp/cuda/mtp.py): the head, its draft chains and verify windows. + +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const wt = @import("cuda_prompt_wt.zig"); +const load = @import("cuda_prompt_load.zig"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const gdn = @import("cuda_gdn.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const vocab = wt.vocab; +const eps = wt.eps; +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const state_n = wt.state_n; +const out_sk = wt.out_sk; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const ple_tail_n = wt.ple_tail_n; +const Q4 = wt.Q4; +const HcW = wt.HcW; + +pub const Sampling = lanes.sampling.Sampling; + +/// Drafts a round at most (Python DEPTH). +pub const max_chain: usize = 6; +/// The live recipe's --mtp-confidence (Python's CONFIDENCE default is 0.7). +pub const default_confidence: f64 = 0.6; +/// A verify window's rows at most: the pending token and a chain, or a copied continuation. +pub const max_rows: usize = 16; +/// qmm.split_for(2560, 2560): fc_embedding and fc_hidden sum K in four slices, in order. +pub const fc_sk: usize = 4; +/// Head rows one launch takes at most (a prompt's rows go through in pieces). +pub const piece_rows: usize = 256; +/// Normed-stream partial sums a row (cuda_batch.zig's pss_n). +pub const pss_n = (dims / 256) * streams; +pub const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; +pub const proj_bytes = gdn.proj_width * 2; + +const vocab_text = @embedFile("draft_vocab_default.txt"); + +fn cint(v: usize) c_int { + return @intCast(v); +} + +/// The draft vocabulary's ids, sorted (Python draft_token_ids("default")): column j of the draft head is ids[j]. +pub fn draftIds(gpa: std.mem.Allocator) ![]u32 { + var out: std.ArrayList(u32) = .empty; + errdefer out.deinit(gpa); + var it = std.mem.tokenizeAny(u8, vocab_text, " \t\r\n,"); + while (it.next()) |part| { + const id = try std.fmt.parseInt(u32, part, 10); + if (id >= vocab) continue; + try out.append(gpa, id); + } + std.mem.sort(u32, out.items, {}, std.sort.asc(u32)); + var n: usize = 0; + for (out.items, 0..) |id, i| { + if (i > 0 and id == out.items[n - 1]) continue; + out.items[n] = id; + n += 1; + } + out.shrinkRetainingCapacity(n); + return out.toOwnedSlice(gpa); +} + +// ---- sampling ------------------------------------------------------------------------------------------------------ + +pub fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +/// A row's first maximum (argmax's choice): its column. +pub fn argmaxRow(row: []const u16) usize { + var best = promote(row[0]); + var at: usize = 0; + for (row[1..], 1..) |b, i| { + const v = promote(b); + if (v > best) { + best = v; + at = i; + } + } + return at; +} + +/// The keyed draw over one bf16 score row, column j being token ids[j] (or j), as lanes.sampling.choose draws. +pub fn drawRow(gpa: std.mem.Allocator, row: []const u16, ids: ?[]const u32, position: u64, sm: ?Sampling) !u32 { + const id = struct { + fn of(map: ?[]const u32, col: usize) u64 { + return if (map) |m| m[col] else col; + } + }.of; + const s = sm orelse return @intCast(id(ids, argmaxRow(row))); + if (s.temperature <= 0) return @intCast(id(ids, argmaxRow(row))); + const n = row.len; + const k = if (s.top_k == 0 or s.top_k >= n) n else @as(usize, s.top_k); + const vals = try gpa.alloc(f64, k); + defer gpa.free(vals); + const keys = try gpa.alloc(u64, k); + defer gpa.free(keys); + if (k == n) { + for (vals, keys, row, 0..) |*v, *key, b, col| { + v.* = promote(b); + key.* = id(ids, col); + } + } else { + // the k best columns, best first; a later column with an equal value ranks below (the id order) + const top_v = try gpa.alloc(f32, k); + defer gpa.free(top_v); + const top_c = try gpa.alloc(usize, k); + defer gpa.free(top_c); + var held: usize = 0; + for (row, 0..) |b, col| { + const v = promote(b); + if (held == k and !(v > top_v[k - 1])) continue; + var at: usize = if (held < k) held else k - 1; + if (held < k) held += 1; + while (at > 0 and v > top_v[at - 1]) : (at -= 1) { + top_v[at] = top_v[at - 1]; + top_c[at] = top_c[at - 1]; + } + top_v[at] = v; + top_c[at] = col; + } + for (vals, keys, top_v, top_c) |*v, *key, tv, c| { + v.* = tv; + key.* = id(ids, c); + } + } + return @intCast(try lanes.sampling.choose(gpa, vals, keys, position, s)); +} + +// ---- the head's weights -------------------------------------------------------------------------------------------- + +pub const Head = struct { + gpa: std.mem.Allocator, + norm_e: cuda.DeviceBuffer, // [D] fp32 + norm_h: cuda.DeviceBuffer, // [S D] fp32 + fc_e: Q4, + fc_h: Q4, + layer: load.Layer, + mixer: HcW, + /// lm_head's rows at `ids`, in their order. + draft: Q4, + ids: []u32, + module: cuda.Module, + best_fn: cuda.Function, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, dir: []const u8) !*Head { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + const h = try gpa.create(Head); + errdefer gpa.destroy(h); + h.gpa = gpa; + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + var buf: [120]u8 = undefined; + const ne = try wt.asF32(gpa, try store.tensor("language_model.mtp.pre_fc_norm_embedding.weight")); + defer gpa.free(ne); + const nh = try wt.asF32(gpa, try store.tensor("language_model.mtp.pre_fc_norm_hidden.weight")); + defer gpa.free(nh); + if (ne.len != dims or nh.len != wide) return error.UnexpectedTensor; + h.norm_e = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(ne)); + errdefer h.norm_e.free(); + h.norm_h = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nh)); + errdefer h.norm_h.free(); + h.fc_e = try fc(gpa, driver, &store, &buf, "language_model.mtp.fc_embedding"); + errdefer h.fc_e.free(); + h.fc_h = try fc(gpa, driver, &store, &buf, "language_model.mtp.fc_hidden"); + errdefer h.fc_h.free(); + h.mixer = try wt.loadHc(gpa, driver, &store, "language_model.mtp.hyper_connection_mixer", false); + errdefer h.mixer.free(); + { + var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); + defer pack_mod.unload(); + const pack_fn = try pack_mod.function(wt.pack_symbol); + h.layer = try load.loadLayerAt(gpa, driver, stream, &store, pack_fn, "language_model.mtp.layers.0", false, false); + } + errdefer h.layer.free(gpa); + h.ids = try draftIds(gpa); + errdefer gpa.free(h.ids); + { + const head_w = try wt.join3(gpa, &store, &buf, "language_model.lm_head", ""); + defer gpa.free(head_w.w); + defer gpa.free(head_w.s); + defer gpa.free(head_w.b); + const wrow = (dims / 8) * 4; + const grow = (dims / 32) * 2; + if (head_w.w.len != vocab * wrow or head_w.s.len != vocab * grow) return error.UnexpectedTensor; + const n = h.ids.len; + const w = try gpa.alloc(u8, n * wrow); + defer gpa.free(w); + const s = try gpa.alloc(u8, n * grow); + defer gpa.free(s); + const b = try gpa.alloc(u8, n * grow); + defer gpa.free(b); + for (h.ids, 0..) |id, j| { + @memcpy(w[j * wrow ..][0..wrow], head_w.w[id * wrow ..][0..wrow]); + @memcpy(s[j * grow ..][0..grow], head_w.s[id * grow ..][0..grow]); + @memcpy(b[j * grow ..][0..grow], head_w.b[id * grow ..][0..grow]); + } + h.draft = try wt.uploadQ4(gpa, driver, w, s, b, n, dims / 8); + } + errdefer h.draft.free(); + h.module = try cuda.Module.load(driver, cuda.kernels.flashnext_mtp); + errdefer h.module.unload(); + h.best_fn = try h.module.function("fn_row_best"); + return h; + } + + fn fc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *wt.Store, buf: []u8, name: []const u8) !Q4 { + const t = try wt.join3(gpa, store, buf, name, ""); + defer gpa.free(t.w); + defer gpa.free(t.s); + defer gpa.free(t.b); + if (t.w.len != dims * (dims / 8) * 4) return error.UnexpectedTensor; + return wt.uploadQ4(gpa, driver, t.w, t.s, t.b, dims, dims / 8); + } + + pub fn deinit(h: *Head) void { + h.module.unload(); + h.draft.free(); + h.gpa.free(h.ids); + h.layer.free(h.gpa); + h.mixer.free(); + h.fc_h.free(); + h.fc_e.free(); + h.norm_h.free(); + h.norm_e.free(); + h.gpa.destroy(h); + } + + /// fn_row_best over `rows` rows of `v` scores `ld` apart: [rows, 4] fp32 at `out` (best, log-sum-exp, id bits). + pub fn best(h: *const Head, stream: cuda.Stream, scores: u64, v: usize, ld: usize, rows: usize, out: u64) !void { + var a: cuda.Args = .{}; + a.add(scores); + a.add(cint(v)); + a.add(@as(c_longlong, @intCast(ld))); + a.add(out); + try cuda.launch.launch(h.best_fn, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 1024 } }, stream, &a); + } +}; + +/// Device bytes the head's weights take beyond the main model's (for the server's admission estimate). +pub fn headBytes() usize { + const experts = (wt.experts_n + 1) * 3 * wt.moe_w * dims / 2 * 5 / 4; + const attn = (wt.proj_n + dims) * dims / 2 * 5 / 4; + const draft = 79_591 * dims / 2 * 5 / 4; + return experts + attn + draft + 2 * dims * dims / 2 * 5 / 4; +} + +pub const Pick = @import("cuda_mtp_draft.zig").Pick; +pub const Chain = @import("cuda_mtp_draft.zig").Chain; +pub const Draft = @import("cuda_mtp_draft.zig").Draft; +pub const Ops = @import("cuda_mtp_ops.zig").Ops; +pub const check = @import("cuda_mtp_check.zig").check; + +test "the draft vocabulary is 79,591 sorted ids under the vocabulary" { + const ids = try draftIds(std.testing.allocator); + defer std.testing.allocator.free(ids); + try std.testing.expectEqual(@as(usize, 79_591), ids.len); + for (ids[1..], ids[0 .. ids.len - 1]) |b, a| try std.testing.expect(b > a); +} + +test "a keyed draw over the top-k columns is the draw over every column" { + const gpa = std.testing.allocator; + var row: [300]u16 = undefined; + var prng = std.Random.DefaultPrng.init(3); + for (&row) |*v| v.* = @as(u16, 0x3f00) + prng.random().int(u8); // values near 0.5..1, many ties + const all = struct { + fn draw(a: std.mem.Allocator, r: []const u16, position: u64, s: Sampling) !u64 { + const vals = try a.alloc(f64, r.len); + defer a.free(vals); + const ids = try a.alloc(u64, r.len); + defer a.free(ids); + for (vals, ids, r, 0..) |*v, *id, b, i| { + v.* = promote(b); + id.* = i; + } + return lanes.sampling.choose(a, vals, ids, position, s); + } + }.draw; + for (0..64) |position| { + const s: Sampling = .{ .seed = 11, .temperature = 1.0, .top_k = 20, .top_p = 0.95 }; + try std.testing.expectEqual(try all(gpa, &row, position, s), try drawRow(gpa, &row, null, position, s)); + } + try std.testing.expectEqual(@as(u32, @intCast(argmaxRow(&row))), try drawRow(gpa, &row, null, 0, null)); +} diff --git a/zig/src/families/flashnext/cuda_mtp_check.zig b/zig/src/families/flashnext/cuda_mtp_check.zig new file mode 100644 index 0000000000..5ba7762f0a --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_check.zig @@ -0,0 +1,171 @@ +//! The drafted check (flashnext-prompt with TF_FLASHNEXT_DRAFTS): serial and drafted decoding, token for token. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const prompt = @import("cuda_prompt.zig"); +const mtp = @import("cuda_mtp.zig"); + +const vocab = wt.vocab; +const max_chain = mtp.max_chain; +const max_rows = mtp.max_rows; +const default_confidence = mtp.default_confidence; +const Head = mtp.Head; +const Draft = mtp.Draft; +const Chain = mtp.Chain; +const Ops = mtp.Ops; +const Sampling = mtp.Sampling; +const drawRow = mtp.drawRow; + +/// What a drafted decode did: rounds, drafts offered and kept, and its wall time. +pub const Stats = struct { + rounds: usize = 0, + drafted: usize = 0, + accepted: usize = 0, + ms: f64 = 0, + /// Wall time in the windows (forward and draws) and in the head (rollback, absorb, chain). + verify_ms: f64 = 0, + draft_ms: f64 = 0, +}; + +/// Python mtp_decode over a prefilled sequence: `out` holds the first token; decode until `count` tokens. +pub fn decode(comptime Tri: type, o: Ops(Tri), seq: *prompt.Seq, d: *Draft, out: *std.ArrayList(u32), count: usize, depth: usize, confidence: f64, sm: ?Sampling, stop: []const u32) !Stats { + const gpa = o.e.gpa; + var st: Stats = .{}; + var ch: Chain = .{}; + var sampled: [max_rows]u32 = undefined; + var ids: [max_rows]u32 = undefined; + const t0 = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + const isStop = struct { + fn f(list: []const u32, t: u32) bool { + return std.mem.indexOfScalar(u32, list, t) != null; + } + }.f; + try o.draft(d, out.items[out.items.len - 1 ..], d.prompt_last, seq.fed + 1, @min(depth, count - out.items.len), sm, confidence, &ch); + while (out.items.len < count and !isStop(stop, out.items[out.items.len - 1])) { + ids[0] = out.items[out.items.len - 1]; + @memcpy(ids[1..][0..ch.n], ch.slice()); + const rows = 1 + ch.n; + const tv = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + try o.verify(seq, d, ids[0..rows]); + try o.draws(d, sm, sampled[0..rows]); + const td = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + st.verify_ms += @as(f64, @floatFromInt(td - tv)) / 1e6; + defer st.draft_ms += @as(f64, @floatFromInt(std.Io.Clock.awake.now(o.e.io).toNanoseconds() - td)) / 1e6; + var keep: usize = 1; + for (ch.slice(), 0..) |t, i| { + if (sampled[i] != t or isStop(stop, sampled[i])) break; + keep += 1; + } + try o.rollback(seq, d, keep); + st.rounds += 1; + st.drafted += ch.n; + st.accepted += keep - 1; + const take = @min(keep, count - out.items.len); + try out.appendSlice(gpa, sampled[0..take]); + if (out.items.len >= count or isStop(stop, out.items[out.items.len - 1])) break; + try o.draft(d, sampled[0..keep], d.streams.ptr, seq.fed + 1, @min(depth, count - out.items.len), sm, confidence, &ch); + } + st.ms = @as(f64, @floatFromInt(std.Io.Clock.awake.now(o.e.io).toNanoseconds() - t0)) / 1e6; + return st; +} + +/// Serial decoding through the same pass and draws: one row a forward (the reference drafted decoding must equal). +pub fn serial(comptime Tri: type, o: Ops(Tri), seq: *prompt.Seq, out: *std.ArrayList(u32), count: usize, sm: ?Sampling, stop: []const u32) !f64 { + const gpa = o.e.gpa; + const t0 = std.Io.Clock.awake.now(o.e.io).toNanoseconds(); + const row = try gpa.alloc(u16, vocab); + defer gpa.free(row); + while (out.items.len < count and std.mem.indexOfScalar(u32, stop, out.items[out.items.len - 1]) == null) { + var one = [_]u32{out.items[out.items.len - 1]}; + try o.e.predict(Tri, o.tri, seq, &one); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + try out.append(gpa, try drawRow(gpa, row, null, seq.fed, sm)); + } + return @as(f64, @floatFromInt(std.Io.Clock.awake.now(o.e.io).toNanoseconds() - t0)) / 1e6; +} + +fn envUsize(name: [:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; +} + +fn envF64(name: [:0]const u8, default: f64) f64 { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseFloat(f64, std.mem.span(v)) catch default; +} + +/// The drafted check (TF_FLASHNEXT_DRAFTS): serial and drafted decoding of `count` tokens a prompt, token for token. +pub fn check(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, model_dir: []const u8, prompts: []const []const u32, count: usize) !bool { + var e = try prompt.Engine.init(gpa, io, driver, stream, model_dir); + defer e.deinit(); + const h = try Head.init(gpa, io, driver, stream, model_dir); + defer h.deinit(); + const o: Ops(Tri) = .{ .e = e, .tri = tri, .h = h }; + const depth = envUsize("TF_FLASHNEXT_DRAFTS", max_chain); + const confidence = envF64("TF_FLASHNEXT_CONFIDENCE", default_confidence); + const sm: ?Sampling = if (std.c.getenv("TF_FLASHNEXT_SEED")) |seed| .{ + .seed = std.fmt.parseInt(u64, std.mem.span(seed), 10) catch 0, + .temperature = envF64("TF_FLASHNEXT_TEMPERATURE", 1.0), + .top_p = envF64("TF_FLASHNEXT_TOP_P", 0.95), + .top_k = @intCast(envUsize("TF_FLASHNEXT_TOP_K", 20)), + } else null; + const repeat = envUsize("TF_FLASHNEXT_REPEAT", 1); + const stop: []const u32 = &.{}; + var same = true; + for (prompts, 0..) |ids, pi| { + const cap = ids.len + count + max_rows; + var seq = try e.newSeq(cap); + defer seq.deinit(gpa); + var d = try Draft.init(gpa, driver, cap, h.ids.len, e.chunk); + defer d.deinit(); + for (0..repeat) |rep| { + // serial: the plain prompt pass, then one row a forward + seq.fed = 0; + seq.history = .{ e.eos, e.eos }; + try resetSeq(&seq); + var ref: std.ArrayList(u32) = .empty; + defer ref.deinit(gpa); + try e.predict(Tri, tri, &seq, ids); + const row = try gpa.alloc(u16, vocab); + defer gpa.free(row); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + try ref.append(gpa, try drawRow(gpa, row, null, seq.fed, sm)); + const serial_ms = try serial(Tri, o, &seq, &ref, count, sm, stop); + // drafted: the prompt pass with its streams and the head's cache, then windows + seq.fed = 0; + seq.history = .{ e.eos, e.eos }; + try resetSeq(&seq); + var got: std.ArrayList(u32) = .empty; + defer got.deinit(gpa); + try o.prefill(&seq, &d, ids); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + try got.append(gpa, try drawRow(gpa, row, null, seq.fed, sm)); + const st = try decode(Tri, o, &seq, &d, &got, count, depth, confidence, sm, stop); + const equal = std.mem.eql(u32, ref.items, got.items); + same = same and equal; + const n_ref: f64 = @floatFromInt(ref.items.len - 1); + const n_got: f64 = @floatFromInt(got.items.len - 1); + std.debug.print("prompt {d} ({d} ids) run {d}: serial {d:.2} tok/s, drafted {d:.2} tok/s ({d:.2}x), {d} rounds, {d} drafts, {d} kept ({d:.1}%), {d:.2} tokens a round: {s}\n", .{ + pi, ids.len, rep, + n_ref / (serial_ms / 1e3), n_got / (st.ms / 1e3), (n_got / (st.ms / 1e3)) / (n_ref / (serial_ms / 1e3)), + st.rounds, st.drafted, st.accepted, + if (st.drafted > 0) 100.0 * @as(f64, @floatFromInt(st.accepted)) / @as(f64, @floatFromInt(st.drafted)) else 0.0, if (st.rounds > 0) n_got / @as(f64, @floatFromInt(st.rounds)) else 0.0, if (equal) "same tokens" else "DIFFER", + }); + std.debug.print(" drafted rounds: {d:.2} ms a window ({d:.2} rows), {d:.2} ms in the head a round\n", .{ st.verify_ms / @as(f64, @floatFromInt(@max(st.rounds, 1))), @as(f64, @floatFromInt(st.rounds + st.drafted)) / @as(f64, @floatFromInt(@max(st.rounds, 1))), st.draft_ms / @as(f64, @floatFromInt(@max(st.rounds, 1))) }); + std.debug.print("serial ", .{}); + for (ref.items) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\ndrafted", .{}); + for (got.items) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + } + } + return same; +} + +/// A sequence's caches back to empty (the check reruns a prompt on one allocation). +fn resetSeq(seq: *prompt.Seq) !void { + for (seq.conv) |*b| try b.fill8(0, null); + for (seq.state) |*b| try b.fill8(0, null); + try seq.tail.fill8(0, null); +} diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig new file mode 100644 index 0000000000..cab7607f3e --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -0,0 +1,178 @@ +//! A drafting stream's state: the head's cache, a window's buffers and what a partly kept window rolls back to. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const gdn = @import("cuda_gdn.zig"); +const mtp = @import("cuda_mtp.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const vocab = wt.vocab; +const eps = wt.eps; +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const state_n = wt.state_n; +const out_sk = wt.out_sk; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const ple_tail_n = wt.ple_tail_n; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const max_chain = mtp.max_chain; +const max_rows = mtp.max_rows; +const conv_bytes = mtp.conv_bytes; +const proj_bytes = mtp.proj_bytes; + +/// One draft and its temperature-1 probability under the head. +pub const Pick = struct { token: u32, prob: f64 }; + +/// A chain of drafts the head holds for the next window. +pub const Chain = struct { + tokens: [max_chain]u32 = undefined, + probs: [max_chain]f64 = undefined, + n: usize = 0, + + pub fn slice(c: *const Chain) []const u32 { + return c.tokens[0..c.n]; + } +}; + +/// A sequence's head cache and the buffers a window needs: its rows' scores and streams, and what it rolls back to. +pub const Draft = struct { + gpa: std.mem.Allocator, + kc: cuda.DeviceBuffer, + vc: cuda.DeviceBuffer, + ikc: cuda.DeviceBuffer, + /// Its pooled indexer keys [cap / 4, 128] bf16, for sparse attention past 2,048 keys (Python State.mtp_pooled). + pooled: cuda.DeviceBuffer, + /// Entries the head's cache holds (Python mtp_len), the last `drafted` of them a chain's. + len: usize = 0, + drafted: usize = 0, + cap: usize, + /// The last pass's rows of main-model streams [rows_cap, wide]: a prompt chunk's, then each window's. + streams: cuda.DeviceBuffer, + rows_cap: usize, + /// The prompt's last row of streams in `streams`, which the first draft reads with the first token. + prompt_last: u64 = 0, + /// The head's last output row (its streams after the MoE write-back): the next chained draft's input. + out: cuda.DeviceBuffer, + /// A window's scores [max_rows, vocab], the draft head's [ids], best values [max_rows, 4]. + logits: cuda.DeviceBuffer, + dlogits: cuda.DeviceBuffer, + best: cuda.DeviceBuffer, + /// Before a window: each DeltaNet layer's conv window and (from the pass) projection rows; the n-gram tail. + saved: cuda.DeviceBuffer, + /// A window's DeltaNet end states: swapped in when every row is kept, else the kept rows replay on the sequence's. + spare: [linear_n]cuda.DeviceBuffer = undefined, + /// A window ran with the states written to `spare` and is not yet kept or rolled back. + open: bool = false, + keep: batch.Keep = .{}, + history: [2]i64 = .{ 0, 0 }, + unused_history: [2]i64 = .{ 0, 0 }, + fed0: usize = 0, + rows: usize = 0, + ids: [max_rows]u32 = undefined, + /// The scores a sampled draw reads on the host (a window's rows, or the draft head's row). + host: []u16, + best_host: [max_rows * 4]f32 = undefined, + + const hist_at: usize = 0; + const proj_at: usize = hist_at + linear_n * conv_bytes; + pub const tail_at: usize = proj_at + linear_n * max_rows * proj_bytes; + /// The window's n-gram rows (the pass's ple_nrow), saved for a rollback to shift the kept ones into the tail. + pub const nrow_at: usize = tail_at + ple_tail_n * wide * 2; + const saved_bytes: usize = nrow_at + max_rows * wide * 2; + + /// Device bytes `init(cap, draft_n, rows)` allocates. + pub fn bytes(cap: usize, draft_n: usize, rows: usize) usize { + return cap * (2 * prompt.kv_row_bytes + index_dim * 2) + prompt.pooledBytes(cap) + @max(rows, max_rows) * wide * 2 + wide * 2 + max_rows * vocab * 2 + draft_n * 2 + max_rows * 16 + saved_bytes + linear_n * state_n * 4; + } + + /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps. + pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, cap: usize, draft_n: usize, rows: usize) !Draft { + var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; + d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + errdefer d.kc.free(); + d.vc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); + errdefer d.vc.free(); + d.ikc = try cuda.DeviceBuffer.alloc(driver, cap * index_dim * 2); + errdefer d.ikc.free(); + d.pooled = try cuda.DeviceBuffer.alloc(driver, prompt.pooledBytes(cap)); + errdefer d.pooled.free(); + d.streams = try cuda.DeviceBuffer.alloc(driver, d.rows_cap * wide * 2); + errdefer d.streams.free(); + d.out = try cuda.DeviceBuffer.alloc(driver, wide * 2); + errdefer d.out.free(); + d.logits = try cuda.DeviceBuffer.alloc(driver, max_rows * vocab * 2); + errdefer d.logits.free(); + d.dlogits = try cuda.DeviceBuffer.alloc(driver, draft_n * 2); + errdefer d.dlogits.free(); + d.best = try cuda.DeviceBuffer.alloc(driver, max_rows * 16); + errdefer d.best.free(); + d.saved = try cuda.DeviceBuffer.alloc(driver, saved_bytes); + errdefer d.saved.free(); + var made: usize = 0; + errdefer for (d.spare[0..made]) |*b| b.free(); + while (made < linear_n) : (made += 1) d.spare[made] = try cuda.DeviceBuffer.alloc(driver, state_n * 4); + d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); + errdefer gpa.free(d.host); + try d.kc.fill8(0, null); + try d.vc.fill8(0, null); + try d.ikc.fill8(0, null); + try d.pooled.fill8(0, null); + var li: usize = 0; + for (&d.keep.layers, 0..) |*lk, layer| { + if (!wt.linearLayer(layer)) continue; + lk.hist = d.saved.ptr + hist_at + li * conv_bytes; + lk.proj = d.saved.ptr + proj_at + li * max_rows * proj_bytes; + li += 1; + } + d.aim(); + return d; + } + + /// Each DeltaNet layer's window end state into its spare buffer. + fn aim(d: *Draft) void { + var li: usize = 0; + for (&d.keep.layers, 0..) |*lk, layer| { + if (!wt.linearLayer(layer)) continue; + lk.state_out = d.spare[li].ptr; + li += 1; + } + } + + /// An open window kept whole: its end states become the sequence's (buffers swapped, nothing copied). + pub fn settle(d: *Draft, seq: *prompt.Seq) void { + if (!d.open) return; + for (seq.state, &d.spare) |*a, *b| std.mem.swap(cuda.DeviceBuffer, a, b); + d.aim(); + d.open = false; + } + + pub fn deinit(d: *Draft) void { + d.gpa.free(d.host); + for (&d.spare) |*b| b.free(); + d.saved.free(); + d.best.free(); + d.dlogits.free(); + d.logits.free(); + d.out.free(); + d.streams.free(); + d.pooled.free(); + d.ikc.free(); + d.vc.free(); + d.kc.free(); + } + + /// Drop the head's cache (a new prompt). + pub fn reset(d: *Draft) void { + d.len = 0; + d.drafted = 0; + d.rows = 0; + } +}; diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig new file mode 100644 index 0000000000..80e08aeb06 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -0,0 +1,300 @@ +//! The head's work over a sequence: chunked prompt passes, its forward and chains, verify windows and rollbacks. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const batch = @import("cuda_batch.zig"); +const prompt = @import("cuda_prompt.zig"); +const gdn = @import("cuda_gdn.zig"); +const mtp = @import("cuda_mtp.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const vocab = wt.vocab; +const eps = wt.eps; +const layers_n = wt.layers_n; +const linear_n = wt.linear_n; +const state_n = wt.state_n; +const out_sk = wt.out_sk; +const kv_heads = wt.kv_heads; +const head_dim = wt.head_dim; +const index_dim = wt.index_dim; +const ple_tail_n = wt.ple_tail_n; +const Q4 = wt.Q4; +const HcW = wt.HcW; +const max_chain = mtp.max_chain; +const max_rows = mtp.max_rows; +const conv_bytes = mtp.conv_bytes; +const proj_bytes = mtp.proj_bytes; +const fc_sk = mtp.fc_sk; +const piece_rows = mtp.piece_rows; +const pss_n = mtp.pss_n; +const Head = mtp.Head; +const Draft = mtp.Draft; +const Chain = mtp.Chain; +const Pick = mtp.Pick; +const Sampling = mtp.Sampling; +const drawRow = mtp.drawRow; +const promote = mtp.promote; + +fn cint(v: usize) c_int { + return @intCast(v); +} + +fn copyAsync(stream: cuda.Stream, dst: u64, src: u64, n: usize) !void { + if (n == 0) return; + const d = stream.d; + try d.check(d.api.cuMemcpyDtoDAsync_v2(dst, src, n, stream.handle), "cuMemcpyDtoDAsync"); +} + +fn download(stream: cuda.Stream, dst: []u8, src: u64) !void { + const d = stream.d; + try stream.synchronize(); + try d.check(d.api.cuMemcpyDtoH_v2(dst.ptr, src, dst.len), "cuMemcpyDtoH"); +} + +/// The head's work over a sequence: prompt pass, verify windows, rollback, forward and chains; `Tri` launches kernels. +pub fn Ops(comptime Tri: type) type { + return struct { + e: *prompt.Engine, + tri: Tri, + h: *Head, + + const Self = @This(); + + fn stream(o: Self) cuda.Stream { + return o.e.stream.*; + } + + /// Prompt rows [at, stop): the pass with its streams kept, then the head over each row with a known next token. + pub fn prefillChunk(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32, at: usize, stop: usize) !void { + if (stop > d.cap or stop <= at or stop - at > d.rows_cap) return error.PromptTooLong; + if (seq.fed != at or d.len != at or d.drafted != 0) return error.PositionMismatch; + const last = stop == ids.len; + try o.e.forward(Tri, o.tri, seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr }); + const head_to = @min(stop, ids.len - 1); + if (head_to > at) try o.forwardHead(d, ids[at + 1 .. head_to + 1], d.streams.ptr, false); + if (last) d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; + } + + /// The whole prompt a chunk of the engine's rows at a time, from an empty head cache. + pub fn prefill(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void { + d.reset(); + var at: usize = 0; + while (at < ids.len) { + const stop = @min(ids.len, at + @min(o.e.chunk, d.rows_cap)); + try o.prefillChunk(seq, d, ids, at, stop); + at = stop; + } + } + + /// The head over `next.len` rows of streams `rows`: cache at d.len, last row in d.out, scores in d.dlogits. + pub fn forwardHead(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + const e = o.e; + const tri = o.tri; + const h = o.h; + if (next.len == 0) return; + if (d.len + next.len > d.cap) return error.PromptTooLong; + const want = @max(@as(usize, 32), 2 * @min(next.len, piece_rows)); + const s = try e.reserve(want); + const piece = @max(@as(usize, 1), @min(piece_rows, s.cap / 2)); + // the engine's block scores (the passes share one stream), for rows past 2,048 keys of the head's cache + const score_rows = try e.reserveScores(d.cap, piece); + var p: batch.Pass(Tri) = .{ + .gpa = e.gpa, + .io = e.io, + .stream = e.stream.*, + .tri = tri, + .k = &e.batch_k, + .s = s, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, + .seq = .{ .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, + .hooks = .{ .logits = .none }, + }; + var at: usize = 0; + while (at < next.len) { + const m = @min(piece, next.len - at); + p.rows = m; + p.pos0 = d.len; + p.pending = null; + try p.stage(next[at..][0..m]); + // the embedding branch: the token's row (stream 0 of the four copies), normed, fc_embedding + try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, m); + try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, m, dims, dims); + try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, m, fc_sk); + try p.sumSlices(s.part, s.branch, true, m * dims, fc_sk); + // the hidden branch: the four streams normed as one row, fc_hidden on each stream + try tri.rmsnorm(rows + at * wide * 2, h.norm_h.ptr, s.norm, s.gq, eps, wide, m, wide, wide); + try p.lane(s.norm, s.gq, h.fc_h, s.dummy, s.part, streams * m, fc_sk); + try p.sumSlices(s.part, s.ple_gated, true, streams * m * dims, fc_sk); + try tri.addStreams(s.branch, s.ple_gated, s.h, m, dims, streams); + // the layer: attention over the head's own cache, its experts + try p.writeback(0, m); + try p.readout(h.layer.attn_hc, 0, m); + try p.attnBlock(&h.layer, 0); + try tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", m * dims, m, dims, streams, 4, 1, 1, out_sk); + try p.readout(h.layer.mlp_hc, 0, m); + try p.moeBlock(&h.layer); + p.pending = 0; + try p.writeback(0, m); + try p.copy(d.out.ptr, s.h + (m - 1) * wide * 2, wide * 2); + d.len += m; + at += m; + if (at == next.len and scores) { + // the last row's streams and sums at row 0, where the captured kernels' pointers stay aligned + if (m > 1) { + try p.copy(s.h, s.h + (m - 1) * wide * 2, wide * 2); + try p.copy(s.pss, s.pss + (m - 1) * pss_n * 4, pss_n * 4); + } + try p.readout(h.mixer, 0, 1); + try p.lane(s.mixed, s.xsm, h.draft, d.dlogits.ptr, s.dummy, 1, 1); + } + } + } + + /// Python draft.absorb: the chain's entries trimmed, then the kept rows (main streams `rows`, next tokens). + pub fn absorb(o: Self, d: *Draft, next: []const u32, rows: u64, scores: bool) !void { + if (d.drafted > 0) { + d.len -= d.drafted; + d.drafted = 0; + } + try o.forwardHead(d, next, rows, scores); + } + + /// The draft at `position` and its temperature-1 probability (Python Engine.sample_draft). + pub fn pick(o: Self, d: *Draft, position: u64, sm: ?Sampling) !Pick { + const h = o.h; + const n = h.ids.len; + try h.best(o.stream(), d.dlogits.ptr, n, n, 1, d.best.ptr); + try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0..4]), d.best.ptr); + const top = d.best_host[0]; + const lse = d.best_host[1]; + const greedy = if (sm) |s| s.temperature <= 0 else true; + if (greedy) { + const col: usize = @intCast(@as(i32, @bitCast(d.best_host[2]))); + return .{ .token = h.ids[col], .prob = @exp(@as(f64, top) - @as(f64, lse)) }; + } + const row = d.host[0..n]; + try download(o.stream(), std.mem.sliceAsBytes(row), d.dlogits.ptr); + const tok = try drawRow(o.e.gpa, row, h.ids, position, sm); + const col = std.sort.binarySearch(u32, h.ids, tok, struct { + fn order(t: u32, x: u32) std.math.Order { + return std.math.order(t, x); + } + }.order) orelse return .{ .token = tok, .prob = 0 }; + return .{ .token = tok, .prob = @exp(@as(f64, promote(row[col])) - @as(f64, lse)) }; + } + + /// Python decode.draft: absorb the kept rows, then chain up to `count` drafts, the first always kept. + pub fn draft(o: Self, d: *Draft, next: []const u32, rows: u64, position: u64, count: usize, sm: ?Sampling, confidence: f64, out: *Chain) !void { + out.n = 0; + try o.absorb(d, next, rows, count > 0); + var j: usize = 0; + while (j < @min(count, max_chain)) : (j += 1) { + const got = try o.pick(d, position + j, sm); + const low = confidence > 0 and got.prob < confidence; + if (low and j > 0) break; + out.tokens[out.n] = got.token; + out.probs[out.n] = got.prob; + out.n += 1; + if (low) break; + if (j + 1 < count) { + try o.forwardHead(d, &.{got.token}, d.out.ptr, true); + d.drafted += 1; + } + } + } + + /// One forward over the window `ids` from seq.fed: each row's scores in d.logits, its streams in d.streams. + pub fn verify(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void { + const n = ids.len; + if (n == 0 or n > max_rows) return error.WindowTooWide; + if (seq.fed + n > seq.cap) return error.PromptTooLong; + d.settle(seq); + d.fed0 = seq.fed; + d.rows = n; + d.history = seq.history; + @memcpy(d.ids[0..n], ids); + const keeps = n > 1; + if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.tail_at, seq.tail.ptr, ple_tail_n * wide * 2); + try o.e.forward(Tri, o.tri, seq, ids, .{ .logits = .all, .logits_out = d.logits.ptr, .keep = if (keeps) &d.keep else null, .streams = d.streams.ptr }); + if (keeps) try copyAsync(o.stream(), d.saved.ptr + Draft.nrow_at, o.e.batch_s.?.ple_nrow, n * wide * 2); + d.open = keeps; + } + + /// The window's rows drawn at their keyed positions (seq position fed0 + 1 + r), as serial decoding draws. + pub fn draws(o: Self, d: *Draft, sm: ?Sampling, out: []u32) !void { + const n = d.rows; + const greedy = if (sm) |s| s.temperature <= 0 else true; + if (greedy) { + try o.h.best(o.stream(), d.logits.ptr, vocab, vocab, n, d.best.ptr); + try download(o.stream(), std.mem.sliceAsBytes(d.best_host[0 .. 4 * n]), d.best.ptr); + for (out[0..n], 0..) |*t, r| t.* = @intCast(@as(i32, @bitCast(d.best_host[4 * r + 2]))); + return; + } + const rows = d.host[0 .. n * vocab]; + try download(o.stream(), std.mem.sliceAsBytes(rows), d.logits.ptr); + for (out[0..n], 0..) |*t, r| t.* = try drawRow(o.e.gpa, rows[r * vocab ..][0..vocab], null, d.fed0 + 1 + r, sm); + } + + /// Keep the window's first `keep` rows (Python forward.commit) from its saved rows; a repeat is a no-op. + pub fn rollback(o: Self, seq: *prompt.Seq, d: *Draft, keep: usize) !void { + if (keep == 0 or keep > d.rows) return error.CannotRewind; + if (keep == d.rows) return d.settle(seq); + if (!d.open) return error.CannotRewind; + const e = o.e; + const st = o.stream(); + const s = if (e.batch_s) |*bs| bs else return error.CannotRewind; + const k = &e.batch_k; + seq.fed = d.fed0 + keep; + var hist = d.history; + for (d.ids[0..keep]) |t| hist = .{ hist[1], t }; + seq.history = hist; + try copyAsync(st, seq.tail.ptr, d.saved.ptr + Draft.tail_at, ple_tail_n * wide * 2); + try shift(k, st, seq.tail.ptr, d.saved.ptr + Draft.nrow_at, wide, keep, ple_tail_n, wide); + var li: usize = 0; + for (e.layers[0..e.n_layers], 0..) |*lw, layer| { + if (!lw.linear) continue; + defer li += 1; + // the sequence's state is still the window's start (the pass wrote its end state to the spare) + const lk = d.keep.layers[layer]; + try copyAsync(st, seq.conv[li].ptr, lk.hist, conv_bytes); + var f: cuda.Args = .{}; + f.add(lk.proj); + f.add(seq.conv[li].ptr); + f.add(lw.cw.ptr); + f.add(lw.a_log.ptr); + f.add(lw.dt.ptr); + f.add(s.gq); + f.add(s.gk); + f.add(s.gv); + f.add(s.ggb); + try cuda.launch.launch(k.front, .{ .grid = .{ .x = @intCast(keep), .y = gdn.nk + gdn.nv }, .block = .{ .x = 256 } }, st, &f); + var a: cuda.Args = .{}; + a.add(s.gq); + a.add(s.gk); + a.add(s.gv); + a.add(s.ggb); + a.add(seq.state[li].ptr); + a.add(cint(keep)); + a.add(s.ys); + try cuda.launch.launch(k.chain, .{ .grid = .{ .x = gdn.nv, .y = gdn.dv / 16 }, .block = .{ .x = 128 } }, st, &a); + try shift(k, st, seq.conv[li].ptr, lk.proj, gdn.proj_width, keep, gdn.taps - 1, gdn.conv_dim); + } + d.rows = keep; + d.open = false; + } + + fn shift(k: *const batch.Kernels, st: cuda.Stream, win: u64, rows_ptr: u64, stride: usize, n: usize, taps: usize, c: usize) !void { + var a: cuda.Args = .{}; + a.add(win); + a.add(rows_ptr); + a.add(@as(c_longlong, @intCast(stride))); + a.add(cint(n)); + a.add(cint(taps)); + a.add(cint(c)); + try cuda.launch.launch(k.shift, .{ .grid = .{ .x = @intCast((c + 255) / 256) }, .block = .{ .x = 256 } }, st, &a); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_prompt_main.zig b/zig/src/families/flashnext/cuda_prompt_main.zig index 0815b17d20..c5fe7157bb 100644 --- a/zig/src/families/flashnext/cuda_prompt_main.zig +++ b/zig/src/families/flashnext/cuda_prompt_main.zig @@ -15,6 +15,8 @@ pub fn main(init: std.process.Init) !u8 { return 2; } const count = try std.fmt.parseInt(usize, args[4], 10); + // TF_FLASHNEXT_DRAFTS: the MTP check (cuda_mtp.zig) over every comma-separated ids file, one load + if (std.c.getenv("TF_FLASHNEXT_DRAFTS") != null) return drafted(gpa, io, args[1], args[2], args[3], count); const text = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .limited(1 << 24)); defer gpa.free(text); var ids: std.ArrayList(u32) = .empty; @@ -48,3 +50,34 @@ pub fn main(init: std.process.Init) !u8 { } return 0; } + +/// Serial and drafted decoding of each prompt compared token for token (flash.mtp.check), with their speeds. +fn drafted(gpa: std.mem.Allocator, io: std.Io, model: []const u8, kernels: []const u8, files: []const u8, count: usize) !u8 { + var prompts: std.ArrayList([]u32) = .empty; + defer { + for (prompts.items) |p| gpa.free(p); + prompts.deinit(gpa); + } + var names = std.mem.tokenizeScalar(u8, files, ','); + while (names.next()) |name| { + const text = try std.Io.Dir.cwd().readFileAlloc(io, name, gpa, .limited(1 << 20)); + defer gpa.free(text); + var ids: std.ArrayList(u32) = .empty; + errdefer ids.deinit(gpa); + var it = std.mem.tokenizeAny(u8, text, ", \t\r\n"); + while (it.next()) |part| try ids.append(gpa, try std.fmt.parseInt(u32, part, 10)); + try prompts.append(gpa, try ids.toOwnedSlice(gpa)); + } + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, kernels); + defer set.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + const tri: flash.Tri = .{ .set = &set, .s = stream }; + const same = try flash.mtp.check(flash.Tri, gpa, io, &driver, &stream, tri, model, prompts.items, count); + std.debug.print("{s}\n", .{if (same) "drafted == serial on every prompt" else "drafted DIFFERS from serial"}); + return if (same) 0 else 1; +} diff --git a/zig/tests/cuda/flashnext/capture_mtp.py b/zig/tests/cuda/flashnext/capture_mtp.py new file mode 100644 index 0000000000..6cf9323f39 --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_mtp.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +"""The Triton kernels Flash Next's MTP head adds to the main model's (its RMSNorms and _add_streams), no checkpoint.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", required=True) + ap.add_argument("--tools", required=True, help="folder of triton_aot_manifest.py") + a = ap.parse_args() + out = Path(a.out) + out.mkdir(parents=True, exist_ok=True) + sys.path.insert(0, a.tools) + import triton_aot_manifest as aot + + rec = aot.Recorder().install() + import torch + from tensorfold.families.qwen4_exp.cuda import glue + + dev = "cuda" + g = torch.Generator(device=dev).manual_seed(0) + streams, dims = 4, 2560 + wide = streams * dims + with rec.scope("mtp"): + for m in (1, 7): + emb4 = torch.randn((m, wide), device=dev, generator=g).to(torch.bfloat16) + norm_e = torch.rand((dims,), device=dev, generator=g) + mixed = torch.empty((m, dims), dtype=torch.bfloat16, device=dev) + xs = torch.empty((m, dims // 32), dtype=torch.float32, device=dev) + glue.rmsnorm(emb4[:, :dims], norm_e, 1e-6, out=mixed, xs=xs) + h = torch.randn((m, wide), device=dev, generator=g).to(torch.bfloat16) + norm_h = torch.rand((wide,), device=dev, generator=g) + hn = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + xh = torch.empty((m, wide // 32), dtype=torch.float32, device=dev) + glue.rmsnorm(h, norm_h, 1e-6, out=hn, xs=xh) + eo = torch.randn((m, dims), device=dev, generator=g).to(torch.bfloat16) + hs = torch.randn((m * streams, dims), device=dev, generator=g).to(torch.bfloat16) + hout = torch.empty((m, wide), dtype=torch.bfloat16, device=dev) + glue.add_streams(eo, hs, hout, streams) + torch.cuda.synchronize() + rec.dump(out / "launches.json") + from triton.runtime.jit import JITFunction + import gc + + spec = {} + for fn in gc.get_objects(): + if isinstance(fn, JITFunction): + name = f"{fn.fn.__module__}.{fn.fn.__qualname__}" + spec[name] = {"params": [p.name for p in fn.params], + "do_not_specialize": [p.name for p in fn.params if p.do_not_specialize], + "no_align": [p.name for p in fn.params if p.do_not_specialize_on_alignment]} + (out / "jit.json").write_text(json.dumps(spec, indent=1, sort_keys=True) + "\n") + print("done", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/zig/tests/cuda/flashnext/capture_mtp.sh b/zig/tests/cuda/flashnext/capture_mtp.sh new file mode 100755 index 0000000000..15ea4d36aa --- /dev/null +++ b/zig/tests/cuda/flashnext/capture_mtp.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Usage: capture_mtp.sh WORK [IMAGE]: the MTP head's own Triton kernels (capture_mtp.py) into WORK/out, no checkpoint. +set -euo pipefail +WORK=$(readlink -f "$1"); IMAGE=${2:-tensorfold-qwen38:v0.6.1} +REPO=$(cd "$(dirname "$0")/../../../.." && pwd) +mkdir -p "$WORK/cache" "$WORK/out" +docker run --rm --name fn-capture-mtp --gpus all --ipc=host --user "$(id -u):$(id -g)" -e HOME=/tmp \ + -e TENSORFOLD_NO_UPDATE_CHECK=1 -v "$WORK/cache":/cache -v "$REPO":/src:ro -v "$WORK/out":/out -w /out \ + "$IMAGE" python -u /src/zig/tests/cuda/flashnext/capture_mtp.py --out /out --tools /src/tools/zig +python3 -B "$REPO/tools/zig/triton_aot_manifest.py" --cache "$WORK/cache/triton" \ + --launches "$WORK/out/launches.json" --mount /cache/triton --out "$WORK/out/manifest.json" From 1ac9bd013fcdd7d02b1dfb47f79ba303076169de Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:16:09 -0500 Subject: [PATCH 035/115] Flash Next CUDA serves several streams through the lane core: drafted windows, shared rounds, kept prompt states and prompt steps The lane backend (cuda_lanes.zig) plugs into the core's round loop, prompt cache and lane host; the family adds no scheduler. Each stream owns its caches (the 36 DeltaNet conv windows and states, the 12 attention layers' KV, indexer keys and pooled blocks, the n-gram tail and history), keyed by the core's stream. - With the MTP head, a drafting stream's verify window is one batched pass; several streams' windows come in one verify call, and each window keeps its own n-gram rows so another stream's pass can reuse the scratch before it is kept. - Streams without drafts (--no-drafts, or "draft": false) take their next token together in one decode round (Engine.round) of up to 32 rows: each stream's reply is the one it gets alone. - prefill_step runs a prompt a chunk at a time (the engine's 1,536 rows, cut at the cache's marks), so the lane host runs the other streams' rounds between a long prompt's chunks instead of after the whole prompt. - The prompt cache keeps a stream's state at the marks the store asks for (cuda_snap.zig: the DeltaNet states, the attention rows and complete pooled blocks, the n-gram tail, and for a drafting stream the MTP head's cache rows too), and a later prompt restores it and runs the rest; the batched pass gives the same bits however a prompt is cut. The head reads one token past a row, so states keep with the core's lookahead of 1, as Nemotron's do. A drafting stream resumes only from a state that holds the head's rows; one without them serves a stream without drafts. - A stream's bytes are its caches and, with the head, its head cache and a chunk's streams; the pass's scratch is reserved once at load. --- zig/src/families/flashnext/cuda_lanes.zig | 357 +++++++++++++++++---- zig/src/families/flashnext/cuda_launch.zig | 1 + zig/src/families/flashnext/cuda_snap.zig | 111 +++++++ 3 files changed, 402 insertions(+), 67 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_snap.zig diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 9e7eddf98d..3eaee65d50 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -1,24 +1,41 @@ -//! Flash Next on the native CUDA server: one stream's prompt, then one token at a time. No draft head. +//! Flash Next's lane backend on the native CUDA server: chunked prompts, MTP draft windows, and shared rounds. const std = @import("std"); const cuda = @import("cuda"); const lanes = @import("lanes"); const prompt = @import("cuda_prompt.zig"); const weights = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); +const snap = @import("cuda_snap.zig"); const be = lanes.backend; const ring = 1024; +/// Streams a round takes at most (a round with more live streams takes them in turns). +const round_streams = 32; /// `Tri` is the captured-kernel launcher. `open` loads the checkpoint and serves greedy or keyed draws. pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; pub const formats: []const []const u8 = &.{"mlx-q4g32"}; - // dense attention over the captured scratch (NCH chunks of 512 keys, QSA off): no sparse indexer yet + // the model's window: past 2,048 keys the block indexer picks each row's 512 best blocks of 4 (cuda_batch.zig) pub const default_context: i64 = @intCast(prompt.max_tokens); - // a prompt runs as one pass of fused decode rows, the whole prompt at once: no staggered segments + // a prompt runs on the batched pass in chunks of this many rows, one a prefill_step: no staggered segments pub const max_segments: u32 = 1; - pub const prompt_rows: u32 = @intCast(prompt.max_tokens); + pub const prompt_rows: u32 = @intCast(prompt.batch.default_chunk); + + /// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`. + pub const snap_bytes = snapBytes; + pub const snap_save = snapSave; + pub const snap_restore = snapRestore; + pub const snap_drop = snapDrop; + /// The pass ignores its cuts, so states keep at any token; the MTP head reads one token on (lookahead 1). + pub const cache_rules = struct { + pub const lookahead: u32 = 1; + pub const planned = false; + pub const min_gap: u32 = 256; + pub const min_prompt: u32 = 512; + }; pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; @@ -29,20 +46,23 @@ pub fn Namespace(comptime Tri: type) type { backend: be.Backend, facts: lanes.Model, rows: u32, - /// Device bytes each admitted stream allocates: its own caches at the server's window. + /// Device bytes each admitted stream allocates: its caches (and with the head, its cache and window). stream_bytes: usize, /// Every stream allocates its own sequence; none rides on the engine's. free_streams: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, lone: ?LoneRun = null, + learns: ?*anyopaque = null, learn_bytes: usize = 0, + /// What the prompt cache's copies are taken from: the backend itself (snap_* look up the owner stream). + target: *anyopaque, }; - /// What the device holds of the checkpoint, for admission: the n-gram tables stay in the mapped shards. - /// 0 when the headers do not read, and the load then names what is wrong. + /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { - return weights.deviceBytes(gpa, io, dir) catch 0; + const main = weights.deviceBytes(gpa, io, dir) catch return 0; + return main + mtp.headBytes(); } /// A request or load this engine refuses, in words; null: none of its own. @@ -50,19 +70,36 @@ pub fn Namespace(comptime Tri: type) type { return switch (err) { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + error.UnsupportedCheckpoint => "Flash Next on CUDA serves MLX affine 4-bit group-32 Qwen3.8-Flash-Next checkpoints; the log names the config field that differs", + error.LearnedPocket => "this folder holds Living Weights pockets, which Flash Next on CUDA does not serve yet", + error.CapturedKernelsRequired => "Flash Next on CUDA replays captured Triton kernels: set TENSORFOLD_CUDA_KERNELS to the set's folder", else => null, }; } - const Lane = struct { seq: prompt.Seq }; + /// A stream's caches; with the head, its cache, a window's buffers and the drafts it holds for the next round. + const Lane = struct { + seq: prompt.Seq, + draft: ?mtp.Draft = null, + chain: mtp.Chain = .{}, + /// The prompt's next row while its pass runs a chunk a call; null once the prompt is in. + filled: ?usize = 0, + }; + + const Ops = mtp.Ops(Tri); const Cuda = struct { gpa: std.mem.Allocator, context: usize, engine: *prompt.Engine, + head: ?*mtp.Head = null, + depth: usize = mtp.max_chain, + confidence: f64 = mtp.default_confidence, set: cuda.aot.Set, stream: cuda.Stream, lanes: std.AutoHashMapUnmanaged(*lanes.Stream, Lane) = .empty, + /// A round's scores, every row [round_streams, vocab] bf16. + logits: cuda.DeviceBuffer, drawn: [ring]u32 = undefined, next: u64 = 0, @@ -70,9 +107,14 @@ pub fn Namespace(comptime Tri: type) type { return .{ .set = &self.set, .s = self.stream }; } + fn ops(self: *Cuda) Ops { + return .{ .e = self.engine, .tri = self.tri(), .h = self.head.? }; + } + fn backend(self: *Cuda) be.Backend { return .{ .ptr = self, .vtable = &.{ .prefill = prefillFn, + .prefill_step = prefillStepFn, .first = firstFn, .queue = queueFn, .read = readFn, @@ -80,18 +122,26 @@ pub fn Namespace(comptime Tri: type) type { .keep = keepFn, .draft = draftFn, .release = releaseFn, + .tree = treeFn, + .probabilities = probabilitiesFn, } }; } - fn facts() lanes.Model { + /// With the head: windows of up to 16 rows cut by its confidence; else shared one-row rounds. + fn facts(self: *const Cuda) lanes.Model { + const drafting = self.head != null; return .{ - .exact_width = 1, + .exact_width = if (drafting) mtp.max_rows else 2, .gpu_tokens = false, - .mtp = false, - .speculate = false, - .drafts = 0, - .batch_rows = 1, - .max_streams = 1, + .mtp = drafting, + .speculate = drafting, + .speculate_early = false, + .drafts = if (drafting) @intCast(self.depth) else 0, + .draft_probabilities = drafting, + .streams_exact = true, + .hidden_rows = true, + .batch_rows = if (drafting) 4 * mtp.max_rows else round_streams, + .max_streams = round_streams, }; } @@ -102,92 +152,176 @@ pub fn Namespace(comptime Tri: type) type { return h; } + fn value(self: *const Cuda, feed: be.Feed) u32 { + return switch (feed) { + .handle => |h| self.drawn[h % ring], + .value => |v| v, + }; + } + + /// The token after the sequence's last row (seq.logits) at `position`: the draws serial decoding makes. fn tokenAt(self: *Cuda, seq: *prompt.Seq, s: *lanes.Stream, position: u64) !u32 { - if (s.sampling) |sm| { - const n = prompt.vocabulary; - const vals = try self.gpa.alloc(f64, n); - defer self.gpa.free(vals); - const ids = try self.gpa.alloc(u64, n); - defer self.gpa.free(ids); - for (vals, ids, 0..) |*v, *id, i| { - const bits = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); - v.* = @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); - id.* = @intCast(i); - } - return @intCast(try lanes.sampling.choose(self.gpa, vals, ids, position, sm)); - } - return prompt.argmax(seq.logits); + if (s.sampling == null) return prompt.argmax(seq.logits); + const row = try self.gpa.alloc(u16, prompt.vocabulary); + defer self.gpa.free(row); + for (row, 0..) |*v, i| v.* = std.mem.readInt(u16, seq.logits[2 * i ..][0..2], .little); + return mtp.drawRow(self.gpa, row, null, position, s.sampling); } fn lane(self: *Cuda, s: *lanes.Stream) !*Lane { return self.lanes.getPtr(s) orelse return error.UnknownStream; } + + /// The stream's caches, from the kept state the cache names on a restore; its prompt pass starts there. + fn open(self: *Cuda, s: *lanes.Stream) !*Lane { + const ids = s.prompt(); + if (ids.len == 0 or ids.len + s.max_new > self.context) return error.PromptTooLong; + const gop = try self.lanes.getOrPut(self.gpa, s); + if (gop.found_existing) dropLane(self, gop.value_ptr); + const drafting = self.head != null and s.drafts; + // without the head's drafts, one row a stream a round: only a drafted window can rewind its states + if (!drafting) { + s.drafts = false; + s.proposer = null; + } + // a window's rows past the reply's last token (its drafts) still need cache slots + const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); + gop.value_ptr.* = .{ .seq = self.engine.newSeq(cap) catch |err| { + self.lanes.removeByPtr(gop.key_ptr); + return err; + } }; + const l = gop.value_ptr; + errdefer { + dropLane(self, l); + self.lanes.removeByPtr(gop.key_ptr); + } + if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, cap, self.head.?.ids.len, self.engine.chunk); + s.cached = 0; + const saved = s.reuse.saved orelse return l; + const kept: *snap.Snap = @ptrCast(@alignCast(saved)); + if (kept.at != s.reuse.at or kept.at == 0 or kept.at >= ids.len) { + s.reuse_failed = true; + } else if (snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { + l.filled = kept.at; + s.cached = @intCast(kept.at); + } else |err| switch (err) { + error.SequenceTooShort, error.NoHeadState => s.reuse_failed = true, + else => return err, + } + return l; + } + + /// One prompt chunk, cut at the next cache mark; true once the last row's scores are in seq.logits. + fn step(self: *Cuda, s: *lanes.Stream, l: *Lane) !bool { + const ids = s.prompt(); + const at = l.filled orelse return true; + if (s.isCancelled()) return error.Cancelled; + var stop = @min(ids.len, at + self.engine.chunk); + for (s.reuse.marks) |m| if (m > at and m < stop) { + stop = m; + break; + }; + const last = stop == ids.len; + if (l.draft) |*d| { + try self.ops().prefillChunk(&l.seq, d, ids, at, stop); + } else try self.engine.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none }); + l.filled = if (last) null else stop; + if (!last and std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(stop)) != null) { + if (s.reuse.hook) |k| k.at(k.ptr, s, @intCast(stop)); + } + return last; + } }; + fn envUsize(name: [:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; + } + pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: Options) !Loaded { - _ = o.drafts; // this model is served one token at a time; there is no draft head - if (o.context > prompt.max_tokens) return error.ContextOverDenseAttention; + if (o.context > prompt.max_tokens) return error.ContextOverWindow; if (o.segments != 1 or o.slide) return error.UnsupportedOption; - // every launch is a captured Triton cubin: TENSORFOLD_CUDA_KERNELS names the set + // every launch past our own kernels is a captured Triton cubin: TENSORFOLD_CUDA_KERNELS names the set const kernels = captured orelse return error.CapturedKernelsRequired; + try @import("cuda_config.zig").check(gpa, io, dir); const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); - own.gpa = gpa; - own.context = o.context; + own.* = .{ .gpa = gpa, .context = o.context, .engine = undefined, .set = undefined, .stream = undefined, .logits = undefined }; own.stream = try cuda.Stream.init(driver, true); errdefer own.stream.deinit(); own.set = try cuda.aot.Set.load(gpa, io, driver, ctx.device, kernels); errdefer own.set.deinit(); own.engine = try prompt.Engine.init(gpa, io, driver, &own.stream, dir); + errdefer own.engine.deinit(); + // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) + _ = try own.engine.reserve(@min(o.context, own.engine.chunk)); + // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) + if (o.drafts) { + own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); + own.depth = @max(1, @min(mtp.max_chain, envUsize("TF_FLASHNEXT_DRAFTS", mtp.max_chain))); + if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; + } + errdefer if (own.head) |h| h.deinit(); + const window = @min(o.context + mtp.max_rows, prompt.max_tokens); + // and an attention block's block scores at the window (64 MiB at 262,144 tokens) + _ = try own.engine.reserveScores(window, @min(o.context, own.engine.chunk)); + own.logits = try cuda.DeviceBuffer.alloc(driver, round_streams * prompt.vocabulary * 2); + errdefer own.logits.free(); own.lanes = .empty; own.next = 0; + var stream_bytes = prompt.seqBytes(window); + if (own.head) |h| stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); return .{ .backend = own.backend(), - .facts = Cuda.facts(), - .rows = 1, - .stream_bytes = prompt.seqBytes(o.context), + .facts = own.facts(), + .rows = if (own.head != null) mtp.max_rows else 2, + .stream_bytes = stream_bytes, .ctx = own, .deinit = release, + .target = own, }; } fn release(p: *anyopaque) void { const own: *Cuda = @ptrCast(@alignCast(p)); var it = own.lanes.valueIterator(); - while (it.next()) |l| l.seq.deinit(own.gpa); + while (it.next()) |l| dropLane(own, l); own.lanes.deinit(own.gpa); + own.logits.free(); + if (own.head) |h| h.deinit(); own.engine.deinit(); own.set.deinit(); own.stream.deinit(); own.gpa.destroy(own); } + fn dropLane(self: *Cuda, l: *Lane) void { + if (l.draft) |*d| d.deinit(); + l.seq.deinit(self.gpa); + } + fn of(ptr: *anyopaque) *Cuda { return @ptrCast(@alignCast(ptr)); } fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { const self = of(ptr); - const ids = s.prompt(); - if (ids.len == 0 or ids.len + s.max_new > self.context) return error.PromptTooLong; - const gop = try self.lanes.getOrPut(self.gpa, s); - if (gop.found_existing) gop.value_ptr.seq.deinit(self.gpa); - gop.value_ptr.* = .{ .seq = self.engine.newSeq(ids.len + s.max_new) catch |err| { - self.lanes.removeByPtr(gop.key_ptr); - return err; - } }; - self.engine.predict(Tri, self.tri(), &gop.value_ptr.seq, ids) catch |err| { - gop.value_ptr.seq.deinit(self.gpa); - self.lanes.removeByPtr(gop.key_ptr); - return err; - }; + const l = try self.open(s); + while (!try self.step(s, l)) {} + } + + /// One chunk of the stream's prompt between the other streams' rounds; the first call opens its caches. + fn prefillStepFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!bool { + const self = of(ptr); + const l = if (self.lanes.getPtr(s)) |x| (if (x.filled != null) x else try self.open(s)) else try self.open(s); + return self.step(s, l); } fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const self = of(ptr); const l = try self.lane(s); - if (position != l.seq.fed) return error.PositionMismatch; + if (l.filled != null or position != l.seq.fed) return error.PositionMismatch; return self.take(try self.tokenAt(&l.seq, s, position)); } @@ -205,44 +339,133 @@ pub fn Namespace(comptime Tri: type) type { return self.drawn[handle % ring]; } + /// A drafting stream's window in one batched pass; streams without drafts share rounds, each row its own bits. fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { const self = of(ptr); - for (windows, out) |w, *o| { + const seqs = try self.gpa.alloc(*prompt.Seq, windows.len); + defer self.gpa.free(seqs); + const ids = try self.gpa.alloc(u32, windows.len); + defer self.gpa.free(ids); + const at_window = try self.gpa.alloc(usize, windows.len); + defer self.gpa.free(at_window); + var plain: usize = 0; + for (windows, out, 0..) |w, *o, wi| { if (w.held != 0 or w.parents != null) return error.HeldDraftsNotBuilt; const l = try self.lane(w.stream); const rows = w.rows(); - var input = w.pending; - for (0..rows) |r| { - if (w.positions[r] != l.seq.fed + 1) return error.PositionMismatch; - var one = [_]u32{input}; - try self.engine.predict(Tri, self.tri(), &l.seq, &one); - o.sampled[r] = try self.tokenAt(&l.seq, w.stream, w.positions[r]); - if (r + 1 < rows) input = w.tokens[r]; + if (l.draft) |*d| { + if (rows > mtp.max_rows) return error.WindowTooWide; + for (w.positions, 0..) |p, r| if (p != l.seq.fed + 1 + r) return error.PositionMismatch; + var row_ids: [mtp.max_rows]u32 = undefined; + row_ids[0] = w.pending; + @memcpy(row_ids[1..rows], w.tokens); + try self.ops().verify(&l.seq, d, row_ids[0..rows]); + try self.ops().draws(d, w.stream.sampling, o.sampled[0..rows]); + } else { + if (rows != 1) return error.CannotRewind; + if (w.positions[0] != l.seq.fed + 1) return error.PositionMismatch; + seqs[plain] = &l.seq; + ids[plain] = w.pending; + at_window[plain] = wi; + plain += 1; } if (w.tokens.len > 0) @memcpy(o.drafts[0..w.tokens.len], w.tokens); } + var at: usize = 0; + while (at < plain) : (at += round_streams) { + const n = @min(round_streams, plain - at); + try self.engine.round(Tri, self.tri(), seqs[at..][0..n], ids[at..][0..n], self.logits.ptr); + } + for (seqs[0..plain], at_window[0..plain]) |q, wi| out[wi].sampled[0] = try self.tokenAt(q, windows[wi].stream, windows[wi].positions[0]); } + /// Each window keeps its first rows (a chain's prefix): the DeltaNet states replayed over them alone. fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { - _ = ptr; + const self = of(ptr); for (windows, paths) |w, path| { - if (path.len != w.rows()) return error.CannotRewind; + if (path.len == 0 or path.len > w.rows()) return error.CannotRewind; for (path, 0..) |r, i| if (r != i) return error.CannotRewind; + const l = try self.lane(w.stream); + if (l.draft) |*d| try self.ops().rollback(&l.seq, d, path.len) else if (path.len != w.rows()) return error.CannotRewind; } } + /// The head absorbs the kept rows and chains the next window's drafts (Python decode.draft). fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { - _ = ptr; - _ = requests; - return error.NoDraftHead; + const self = of(ptr); + if (self.head == null) return error.NoDraftHead; + for (requests) |r| { + if (r.lanes != null) return error.TreesNotBuilt; + const l = try self.lane(r.stream); + const d = if (l.draft) |*x| x else return error.NoDraftHead; + // the window's kept rows: rollback is idempotent, so a later keep finds it done + if (r.rows) |path| try self.ops().rollback(&l.seq, d, path.len) else d.settle(&l.seq); + if (r.position != l.seq.fed + 1) return error.PositionMismatch; + // the window's rows past the cache's room cannot run: fewer drafts there + const room = l.seq.cap - l.seq.fed - 1; + const count = @min(@min(r.depth, self.depth), room); + if (r.rows) |path| { + for (path, 0..) |row, i| if (row != i) return error.TreesNotBuilt; + if (r.follow.len < path.len) return error.PositionMismatch; + try self.ops().draft(d, r.follow[0..path.len], d.streams.ptr, r.position, count, r.stream.sampling, self.confidence, &l.chain); + } else { + const first = self.value(r.first orelse return error.NoFirstToken); + try self.ops().draft(d, &.{first}, d.prompt_last, r.position, count, r.stream.sampling, self.confidence, &l.chain); + } + } + } + + /// The held chain as host tokens: the round loop verifies them as it would copied ones. + fn treeFn(ptr: *anyopaque, s: *lanes.Stream, gpa: std.mem.Allocator) anyerror!?lanes.stream.Held { + const self = of(ptr); + const l = try self.lane(s); + if (l.draft == null) return null; + return .{ .count = @intCast(l.chain.n), .tokens = try gpa.dupe(u32, l.chain.slice()) }; + } + + /// Each held draft's temperature-1 probability under the head. + fn probabilitiesFn(ptr: *anyopaque, s: *lanes.Stream, out: []f64) anyerror!bool { + const self = of(ptr); + const l = try self.lane(s); + if (l.draft == null or out.len > l.chain.n) return false; + @memcpy(out, l.chain.probs[0..out.len]); + return true; } fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { const self = of(ptr); if (self.lanes.fetchRemove(s)) |kv| { var l = kv.value; - l.seq.deinit(self.gpa); + self.stream.synchronize() catch {}; + dropLane(self, &l); } } + + // -- the prompt cache's copies ------------------------------------------------------------------------ + + fn snapBytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(at, of(ptr).head != null); + } + + /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). + fn snapSave(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + if (l.seq.fed != at) return error.NotAtMark; + return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); + } + + fn snapRestore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); + } + + fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { + const self = of(ptr); + snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); + } }; } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index f040561f86..b779560d7f 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -321,6 +321,7 @@ test { _ = prompt.batch; _ = mtp; _ = native; + _ = @import("cuda_snap.zig"); _ = @import("cuda_config.zig"); _ = row_path; _ = @import("cuda_mtp_check.zig"); diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig new file mode 100644 index 0000000000..f3f4a72cdc --- /dev/null +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -0,0 +1,111 @@ +//! Flash Next prompt states for the shared prompt cache: one stream's caches after `at` tokens, copied on the device. + +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("cuda_prompt.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const gdn = @import("cuda_gdn.zig"); +const Draft = @import("cuda_mtp_draft.zig").Draft; + +const conv_bytes = (gdn.taps - 1) * gdn.conv_dim * 2; +const state_bytes = wt.state_n * 4; +const index_row = wt.index_dim * 2; +const tail_bytes = wt.ple_tail_n * wt.wide * 2; + +/// One attention cache's first `at` key, value and indexer-key rows, and the pooled blocks they complete. +fn attnBytes(at: usize) usize { + return at * (2 * prompt.kv_row_bytes + index_row) + at / wt.index_ratio * index_row; +} + +/// Device bytes a state after `at` tokens takes: the DeltaNet's are fixed, each attention cache grows by the token. +pub fn bytes(at: usize, head: bool) usize { + return wt.linear_n * (conv_bytes + state_bytes) + (wt.attn_n + @intFromBool(head)) * attnBytes(at) + tail_bytes; +} + +/// One kept state; `head`: it holds the MTP head's cache over the same rows, which a drafting stream needs. +pub const Snap = struct { + buf: cuda.DeviceBuffer, + at: usize, + history: [2]i64, + head: bool, +}; + +const Way = enum { save, restore }; + +const Piece = struct { bufs: []cuda.DeviceBuffer, n: usize }; + +/// An attention cache's pieces for `at` rows: complete pooled blocks only (the rest pool again as later rows land). +fn attnPieces(kc: []cuda.DeviceBuffer, vc: []cuda.DeviceBuffer, ikc: []cuda.DeviceBuffer, pooled: []cuda.DeviceBuffer, at: usize) [4]Piece { + return .{ + .{ .bufs = kc, .n = at * prompt.kv_row_bytes }, + .{ .bufs = vc, .n = at * prompt.kv_row_bytes }, + .{ .bufs = ikc, .n = at * index_row }, + .{ .bufs = pooled, .n = at / wt.index_ratio * index_row }, + }; +} + +/// Each piece of the sequence (and the head's cache, with `d`) to or from its place in the buffer at `base`. +fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, way: Way) !void { + const dr = stream.d; + var off: u64 = 0; + var pieces: [11]Piece = undefined; + pieces[0] = .{ .bufs = seq.conv, .n = conv_bytes }; + pieces[1] = .{ .bufs = seq.state, .n = state_bytes }; + pieces[2..6].* = attnPieces(seq.kc, seq.vc, seq.ikc, seq.pooled, at); + pieces[6] = .{ .bufs = (&seq.tail)[0..1], .n = tail_bytes }; + var count: usize = 7; + if (d) |h| { + pieces[7..11].* = attnPieces((&h.kc)[0..1], (&h.vc)[0..1], (&h.ikc)[0..1], (&h.pooled)[0..1], at); + count = 11; + } + for (pieces[0..count]) |p| for (p.bufs) |b| { + if (p.n == 0) continue; + const dst, const src = if (way == .save) .{ base + off, b.ptr } else .{ b.ptr, base + off }; + try dr.check(dr.api.cuMemcpyDtoDAsync_v2(dst, src, p.n, stream.handle), "cuMemcpyDtoDAsync"); + off += p.n; + }; + std.debug.assert(off == bytes(at, d != null)); +} + +/// A copy of `seq` (and of the head's cache, with `d`) as it stands after `seq.fed` tokens. +pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft) !*Snap { + const at = seq.fed; + if (d) |h| if (h.len != at or h.drafted != 0) return error.NotAtMark; + const snap = try gpa.create(Snap); + errdefer gpa.destroy(snap); + snap.* = .{ .buf = try cuda.DeviceBuffer.alloc(driver, bytes(at, d != null)), .at = at, .history = seq.history, .head = d != null }; + errdefer snap.buf.free(); + try move(stream, seq, d, snap.buf.ptr, at, .save); + try stream.synchronize(); + return snap; +} + +/// `seq` (and the head's cache, with `d`) becomes `snap`'s state; refused before any copy when it cannot hold it. +pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { + if (snap.at > seq.cap) return error.SequenceTooShort; + if (d) |h| if (!snap.head or snap.at > h.cap) return error.NoHeadState; + // a state with the head's rows serves a stream without drafts too: those rows stay in the buffer + try move(stream, seq, d, snap.buf.ptr, snap.at, .restore); + try stream.synchronize(); + seq.history = snap.history; + seq.fed = snap.at; + if (d) |h| { + h.reset(); + h.len = snap.at; + } +} + +pub fn drop(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { + stream.synchronize() catch {}; + snap.buf.free(); + gpa.destroy(snap); +} + +test "a state's bytes: about 110 MiB of DeltaNet states and the n-gram tail, then 27 KiB a token" { + try std.testing.expectEqual(@as(usize, 12 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, false) - bytes(0, false)); + // a complete block of 4 adds its pooled row + try std.testing.expectEqual(@as(usize, 12 * (4 * (2 * 2 * 256 * 2 + 128 * 2) + 128 * 2)), bytes(4, false) - bytes(0, false)); + // the head's cache is one more attention layer's rows + try std.testing.expectEqual(@as(usize, 13 * (2 * 2 * 256 * 2 + 128 * 2)), bytes(1, true) - bytes(0, true)); + try std.testing.expect(bytes(0, false) > 100 << 20 and bytes(0, false) < 120 << 20); +} From 62865ab4f431e1b425f9ae925d68dac6883dfcc1 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:50:35 -0500 Subject: [PATCH 036/115] Flash Next CUDA: buffers are zeroed on the compute stream, so a pass cannot read them before the fill lands A fill with no stream (DeviceBuffer.fill8(..., null), cuMemsetD8) runs on the legacy stream and, for device memory, returns before the memset is done; the engine's passes run on a non-blocking stream, which the legacy stream does not order against. A new sequence's DeltaNet conv windows and states, attention caches and n-gram tail, a drafting stream's head cache, and the pass scratch's zero row were filled that way, so the first chunk of a prompt could read a buffer before it was zeroed, for example the previous sequence's DeltaNet state. Every fill now goes on the compute stream (and the scratch's inverse frequencies upload there too), ahead of the passes that read them, as the token-id upload does. The drafted check showed it: it zeroes the DeltaNet states between a prompt's serial and drafted runs, and in one run of six prompts the prose prompt's drafted decode left the serial tokens at token 9 (113 drafts where its other runs make 111). A variant that skips the reset, so the drafted run starts from a state that is not zero, gives a mismatch of the same kind; with the fills ordered, the check's repeats are equal, under CPU load too. The same pattern was in the row path the prompt check uses, which is fixed alike. --- .../families/flashnext/cuda_batch_state.zig | 7 ++++--- zig/src/families/flashnext/cuda_lanes.zig | 2 +- zig/src/families/flashnext/cuda_mtp_check.zig | 14 ++++++------- zig/src/families/flashnext/cuda_mtp_draft.zig | 8 +++---- zig/src/families/flashnext/cuda_prompt.zig | 21 ++++++++++--------- zig/src/families/flashnext/cuda_rows.zig | 12 +++++------ 6 files changed, 32 insertions(+), 32 deletions(-) diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 8cd1fb2df8..b91273d58b 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -228,7 +228,7 @@ pub const Scratch = struct { return n; } - pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, cap: usize) !Scratch { + pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize) !Scratch { var s: Scratch = .{ .cap = cap }; errdefer s.free(gpa); const pairs = cap * slots; @@ -290,9 +290,10 @@ pub const Scratch = struct { var inv: [half]f32 = undefined; for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); s.inv = try s.add(gpa, d, half * 4); - try s.bufs.items[s.bufs.items.len - 1].upload(0, std.mem.sliceAsBytes(&inv)); + try s.bufs.items[s.bufs.items.len - 1].uploadAsync(0, std.mem.sliceAsBytes(&inv), stream.handle); s.dummy = try s.add(gpa, d, 4096); - try s.bufs.items[s.bufs.items.len - 1].fill8(0, null); + try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); + try stream.synchronize(); return s; } diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 3eaee65d50..9b97a94081 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -195,7 +195,7 @@ pub fn Namespace(comptime Tri: type) type { dropLane(self, l); self.lanes.removeByPtr(gop.key_ptr); } - if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, cap, self.head.?.ids.len, self.engine.chunk); + if (drafting) l.draft = try mtp.Draft.init(self.gpa, self.engine.driver, self.stream, cap, self.head.?.ids.len, self.engine.chunk); s.cached = 0; const saved = s.reuse.saved orelse return l; const kept: *snap.Snap = @ptrCast(@alignCast(saved)); diff --git a/zig/src/families/flashnext/cuda_mtp_check.zig b/zig/src/families/flashnext/cuda_mtp_check.zig index 5ba7762f0a..aea0f37daf 100644 --- a/zig/src/families/flashnext/cuda_mtp_check.zig +++ b/zig/src/families/flashnext/cuda_mtp_check.zig @@ -117,13 +117,13 @@ pub fn check(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cu const cap = ids.len + count + max_rows; var seq = try e.newSeq(cap); defer seq.deinit(gpa); - var d = try Draft.init(gpa, driver, cap, h.ids.len, e.chunk); + var d = try Draft.init(gpa, driver, stream.*, cap, h.ids.len, e.chunk); defer d.deinit(); for (0..repeat) |rep| { // serial: the plain prompt pass, then one row a forward seq.fed = 0; seq.history = .{ e.eos, e.eos }; - try resetSeq(&seq); + try resetSeq(&seq, stream.*); var ref: std.ArrayList(u32) = .empty; defer ref.deinit(gpa); try e.predict(Tri, tri, &seq, ids); @@ -135,7 +135,7 @@ pub fn check(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cu // drafted: the prompt pass with its streams and the head's cache, then windows seq.fed = 0; seq.history = .{ e.eos, e.eos }; - try resetSeq(&seq); + try resetSeq(&seq, stream.*); var got: std.ArrayList(u32) = .empty; defer got.deinit(gpa); try o.prefill(&seq, &d, ids); @@ -164,8 +164,8 @@ pub fn check(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cu } /// A sequence's caches back to empty (the check reruns a prompt on one allocation). -fn resetSeq(seq: *prompt.Seq) !void { - for (seq.conv) |*b| try b.fill8(0, null); - for (seq.state) |*b| try b.fill8(0, null); - try seq.tail.fill8(0, null); +fn resetSeq(seq: *prompt.Seq, stream: cuda.Stream) !void { + for (seq.conv) |*b| try b.fill8(0, stream.handle); + for (seq.state) |*b| try b.fill8(0, stream.handle); + try seq.tail.fill8(0, stream.handle); } diff --git a/zig/src/families/flashnext/cuda_mtp_draft.zig b/zig/src/families/flashnext/cuda_mtp_draft.zig index cab7607f3e..7f354f48de 100644 --- a/zig/src/families/flashnext/cuda_mtp_draft.zig +++ b/zig/src/families/flashnext/cuda_mtp_draft.zig @@ -94,7 +94,7 @@ pub const Draft = struct { } /// A head cache of `cap` entries; `rows` the widest prompt chunk whose streams it keeps. - pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, cap: usize, draft_n: usize, rows: usize) !Draft { + pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, cap: usize, draft_n: usize, rows: usize) !Draft { var d: Draft = .{ .gpa = gpa, .cap = cap, .rows_cap = @max(rows, max_rows), .kc = undefined, .vc = undefined, .ikc = undefined, .pooled = undefined, .streams = undefined, .out = undefined, .logits = undefined, .dlogits = undefined, .best = undefined, .saved = undefined, .host = undefined }; d.kc = try cuda.DeviceBuffer.alloc(driver, cap * kv_heads * head_dim * 2); errdefer d.kc.free(); @@ -121,10 +121,8 @@ pub const Draft = struct { while (made < linear_n) : (made += 1) d.spare[made] = try cuda.DeviceBuffer.alloc(driver, state_n * 4); d.host = try gpa.alloc(u16, @max(max_rows * vocab, draft_n)); errdefer gpa.free(d.host); - try d.kc.fill8(0, null); - try d.vc.fill8(0, null); - try d.ikc.fill8(0, null); - try d.pooled.fill8(0, null); + // zeroed on the compute stream, ahead of the passes that read them + for ([_]cuda.DeviceBuffer{ d.kc, d.vc, d.ikc, d.pooled }) |b| try b.fill8(0, stream.handle); var li: usize = 0; for (&d.keep.layers, 0..) |*lk, layer| { if (!wt.linearLayer(layer)) continue; diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index d0b96fafdd..7e022fa2cf 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -59,7 +59,8 @@ pub fn argmax(raw: []const u8) u32 { return @intCast(at); } -fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![]cuda.DeviceBuffer { +/// `n` zeroed buffers, the fill on the compute stream: a legacy-stream memset returns before it lands. +fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, n: usize, len: usize) ![]cuda.DeviceBuffer { const out = try gpa.alloc(cuda.DeviceBuffer, n); var i: usize = 0; errdefer { @@ -68,7 +69,7 @@ fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, n: usize, len: usize) ![] } while (i < n) : (i += 1) { out[i] = try cuda.DeviceBuffer.alloc(driver, len); - try out[i].fill8(0, null); + try out[i].fill8(0, stream.handle); } return out; } @@ -185,21 +186,21 @@ pub const Engine = struct { if (cap == 0) return error.UnexpectedTensor; const logits = try e.gpa.alloc(u8, vocab * 2); errdefer e.gpa.free(logits); - const conv = try zeros(e.gpa, e.driver, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); + const conv = try zeros(e.gpa, e.driver, e.stream.*, linear_n, (gdn.taps - 1) * gdn.conv_dim * 2); errdefer freeOwned(e.gpa, conv); - const state = try zeros(e.gpa, e.driver, linear_n, state_n * 4); + const state = try zeros(e.gpa, e.driver, e.stream.*, linear_n, state_n * 4); errdefer freeOwned(e.gpa, state); - const kc = try zeros(e.gpa, e.driver, attn_n, cap * kv_heads * head_dim * 2); + const kc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, kc); - const vc = try zeros(e.gpa, e.driver, attn_n, cap * kv_heads * head_dim * 2); + const vc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * kv_heads * head_dim * 2); errdefer freeOwned(e.gpa, vc); - const ikc = try zeros(e.gpa, e.driver, attn_n, cap * index_dim * 2); + const ikc = try zeros(e.gpa, e.driver, e.stream.*, attn_n, cap * index_dim * 2); errdefer freeOwned(e.gpa, ikc); - const pooled = try zeros(e.gpa, e.driver, attn_n, pooledBytes(cap)); + const pooled = try zeros(e.gpa, e.driver, e.stream.*, attn_n, pooledBytes(cap)); errdefer freeOwned(e.gpa, pooled); var tail = try cuda.DeviceBuffer.alloc(e.driver, ple_tail_n * wide * 2); errdefer tail.free(); - try tail.fill8(0, null); + try tail.fill8(0, e.stream.handle); return .{ .conv = conv, .state = state, @@ -241,7 +242,7 @@ pub const Engine = struct { bs.free(e.gpa); e.batch_s = null; } - e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, rows); + e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows); return &e.batch_s.?; } diff --git a/zig/src/families/flashnext/cuda_rows.zig b/zig/src/families/flashnext/cuda_rows.zig index 0204205632..a3081891e8 100644 --- a/zig/src/families/flashnext/cuda_rows.zig +++ b/zig/src/families/flashnext/cuda_rows.zig @@ -35,10 +35,10 @@ pub const Rows = struct { up_fn: cuda.Function, down_fn: cuda.Function, - pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver) !Rows { + pub fn init(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream) !Rows { var r: Rows = .{ .gpa = gpa, .driver = driver, .scratch = try Scratch.alloc(driver, 1), .scratch_rows = 1, .mode = undefined, .experts = undefined, .plan_fn = undefined, .up_fn = undefined, .down_fn = undefined }; errdefer r.scratch.free(); - try r.scratch.dummy.fill8(0, null); + try r.scratch.dummy.fill8(0, stream.handle); r.mode = try gpa.alloc(u8, 1); errdefer gpa.free(r.mode); r.experts = try cuda.Module.load(driver, cuda.kernels.experts); @@ -55,11 +55,11 @@ pub const Rows = struct { r.scratch.free(); } - fn grow(r: *Rows, rows: usize) !void { + fn grow(r: *Rows, stream: cuda.Stream, rows: usize) !void { if (rows <= r.scratch_rows) return; var next = try Scratch.alloc(r.driver, rows); errdefer next.free(); - try next.dummy.fill8(0, null); + try next.dummy.fill8(0, stream.handle); const mode = try r.gpa.alloc(u8, rows); r.scratch.free(); r.gpa.free(r.mode); @@ -73,7 +73,7 @@ pub const Rows = struct { if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; // dense attention only: the reference stops where the indexer would start choosing if (seq.fed + ids.len > weights.dense_nch * 512) return error.UnexpectedTensor; - try r.grow(ids.len); + try r.grow(e.stream.*, ids.len); var pass = Pass{ .gpa = e.gpa, .io = e.io, .driver = e.driver, .stream = e.stream, .model_dir = e.model_dir, .s = &r.scratch, .table = e.table, .mode = r.mode, .history = &seq.history, .conv = seq.conv, .state = seq.state, .kc = seq.kc, .vc = seq.vc, .ikc = seq.ikc, .tail = seq.tail, .mixer = e.mixer, .head_q = e.head_q, .logits = seq.logits, .plan_fn = r.plan_fn, .up_fn = r.up_fn, .down_fn = r.down_fn }; @memset(pass.mode, 0); for (ids, 0..) |token, t| try embedRow(Tri, &pass, tri, token, pass.s.h.ptr + @as(u64, t) * wide * 2); @@ -192,7 +192,7 @@ pub fn run(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cuda var e = try prompt.Engine.init(gpa, io, driver, stream, model_dir); defer e.deinit(); const by_rows = if (std.c.getenv("TF_FLASHNEXT_PROMPT")) |v| std.mem.eql(u8, std.mem.span(v), "rows") else false; - var rows: ?Rows = if (by_rows) try Rows.init(gpa, driver) else null; + var rows: ?Rows = if (by_rows) try Rows.init(gpa, driver, stream) else null; defer if (rows) |*r| r.deinit(); var seq = try e.newSeq(tokens.len + new_tokens); defer seq.deinit(gpa); From 8d32cd33a043705b579350276d3cda839b64a6d4 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:37:07 -0500 Subject: [PATCH 037/115] Docs: Flash Next on CUDA, its recipe, the CUDA and Flash Next pages, the recipe index, RUNBOOK and CHANGELOG docs/recipes/flash-next-cuda.md says what the CUDA engine serves and how: the checkpoint it takes and the ones it refuses, the captured kernel set and how to capture and pack one, the checks (the prompt and drafted checks, the block selection, the rows kernel, served exactness) and measurements on a GB10, with what it does not do yet. The Flash Next, CUDA and zig-family pages point to it, and CHANGELOG has the Unreleased entries. --- CHANGELOG.md | 7 ++ RUNBOOK.md | 2 + docs/recipes/README.md | 2 +- docs/recipes/adding-a-zig-family.md | 6 +- docs/recipes/cuda.md | 2 +- docs/recipes/flash-next-cuda.md | 104 ++++++++++++++++++++++++++++ docs/recipes/qwen3.8-flash-next.md | 7 +- 7 files changed, 122 insertions(+), 8 deletions(-) create mode 100644 docs/recipes/flash-next-cuda.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 1be875947f..1e856dfcb6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,13 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 release notes](RELEASE-NOTES-1.0.0.md) for the native binary's supported models and migration details. +## Unreleased + +- Qwen3.8 Flash Next serves on CUDA from its MLX 4-bit group-32 checkpoint with the MTP head, through the lane core: drafted windows, shared rounds of streams without drafts, kept prompt states (a drafting stream's too, with the head's cache rows and a lookahead of 1, as Nemotron's) and prompts a chunk at a time between other streams' rounds. Each prompt chunk, verify window and decode round is one batched pass whose kernels give every row the bits it gets alone, so drafted replies equal `"draft": false`, concurrent replies equal their solo runs and resumed prompts equal fresh ones; past 2,048 keys attention reads the 512 blocks the model's indexer picks, up to the 262,144-token window. It replays the Python 0.6.1 engine's Triton kernels from a captured set (`TENSORFOLD_CUDA_KERNELS`; the capture and pack scripts are in `zig/tests/cuda/flashnext`), and 32 greedy tokens equal that engine's teacher path on prompts of 19 to 12,001 tokens. On a GB10 one stream decodes 82.6 tok/s drafted on code and 61.8 on chat, 35.0 without drafts, and 4 sampled streams without drafts 77.3 in all; a cold 12,001-token prompt's first token takes 10.0 s and a 45,030-token one's 37.9 s, and another stream waits at most 1.3 s between tokens while it runs. A checkpoint of another width or shape is refused at load, by its config field. @AdrianBinDC started the port in #476, matching every layer of one decode token to host checks; see [Flash Next on CUDA](docs/recipes/flash-next-cuda.md) (#472). Thanks to @AdrianBinDC. +- A Flash Next sequence's buffers are zeroed on the engine's compute stream. A fill with no stream returns before the memset lands and is not ordered against the non-blocking stream the passes run on, so a new sequence's first prompt chunk could read its DeltaNet state before it was zeroed; the drafted check, which zeroes the states between a prompt's serial and drafted runs, once left the serial tokens that way. +- CUDA has a lane kernel for affine 4-bit experts over a window's or a prompt chunk's rows (`experts_rows.cu`, `cuda.experts.RowsExperts`): experts.cu's decode sums for each row, two row tiles of 16 a warp, so a warp reads an expert's weights once for 32 pairs. `tf-cuda-test experts-rows` checks it against the decode kernel's one-row launches byte for byte at 1 to 300 rows. +- `tools/zig/triton_aot_pack.py` leaves out a constexpr Triton folded from `None` (an unused pointer argument), where it stopped, and takes `--jit` once per capture, merged, so one call packs several captures; `test_aot_manifest.py` covers both. + ## 1.0.11 - `--cortheon`'s instruction filter keeps reading past a long word. A word longer than the filter's scan buffer ended the scan early, so an order written after one was never checked; each word now keeps at most 16 bytes, so a long word can no longer end the scan (#829). Thanks to @joeynyc. diff --git a/RUNBOOK.md b/RUNBOOK.md index 408f2d91e8..32a3638f75 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -123,6 +123,8 @@ The archives require glibc 2.28 or newer and an NVIDIA driver providing `libcuda CUDA serves Nemotron 3.5 Lightning, qualified on GB10 (`linux-aarch64`, greedy and sampled) and on Ampere cards with compute capability 8.6 (`linux-x86_64`, greedy: the RTX 30 series, RTX A6000, A10 and A40), and checked on an RTX PRO 6000 Blackwell (compute capability 12.0, `linux-x86_64`). +Flash Next also serves on CUDA from its MLX 4-bit checkpoint and a captured kernel set named by `TENSORFOLD_CUDA_KERNELS`, +run on a GB10: see [Flash Next on CUDA](docs/recipes/flash-next-cuda.md). On a GB10 arm64 host (an x86_64 host uses `linux-x86_64` in the same commands): diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 1e38e07605..c1bdf11971 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -8,7 +8,7 @@ the native engine, or the Python engine frozen at 0.6.6. | Nemotron 3.5 Lightning | [MLX](nemotron-3.5.md) | | Qwen3.8-27B | [MLX, quantization and CUDA](qwen3.8-27b.md) | | Qwen3.5-2B | [Native Metal preview, tied affine 4-bit](qwen3.5-2b-native.md) | -| Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [CUDA images](flash-next-vision.md) | +| Qwen3.8 Flash Next | [MLX prefill and CUDA](qwen3.8-flash-next.md), [native CUDA](flash-next-cuda.md), [CUDA images](flash-next-vision.md) | | Ternary Bonsai 2 27B | [MLX](ternary-bonsai-2.md) | | GLM-5.3-Flash | [MLX on a 256 GB Mac, two-rank CUDA](glm-5.3-flash.md) | | Gemma 4 26B-A4B | [MLX, fused one-row decode](gemma-4.md) | diff --git a/docs/recipes/adding-a-zig-family.md b/docs/recipes/adding-a-zig-family.md index 7252bd9cf7..7372013d35 100644 --- a/docs/recipes/adding-a-zig-family.md +++ b/docs/recipes/adding-a-zig-family.md @@ -38,7 +38,7 @@ Some of the host side already lives in core, but so far each piece serves one fa - staggered segments (`zig/src/core/segments.zig`, and `zig/src/cuda/segments.zig` on CUDA) serve Flash Next on Metal and Nemotron on CUDA; - every host returns the HTTP contract in `zig/src/core/engine_api.zig`, and `zig/src/core/lane_host.zig` serves Nemotron on Metal - and CUDA. Flash Next has its own host and serves one reply at a time. + and CUDA, and Flash Next on CUDA. Flash Next on Metal has its own host and serves one reply at a time. The GPU side is barely shared yet, and that costs time. In the GLM-5.3 port each extra row in a window cost about 4.4 ms of GPU time, against about 2.3 ms for the row's own experts. Row kernels read every replicated weight again for @@ -71,11 +71,11 @@ lane kernels go in, check against the engine's own plain output (step 4). | Piece | Shared now | Family copies or gaps | | --- | --- | --- | -| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels | +| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels. On CUDA, `zig/kernels/cuda/experts_rows.cu` runs affine 4-bit group-32 experts over a window's rows with the decode kernel's sums (Flash Next) | | GPU-side rounds | the host loop, depth rule and copy proposer in `zig/src/core/lanes`, used by Nemotron | Flash Next: its own host loop, plus `fz_accept`, the ring and event chaining (`zig/src/families/flashnext/replay.zig`). Nemotron on Metal: `zig/src/families/nemotron/gpu_round.zig` for a lone stream. Kimi K3: none | | Several Macs | the MCDMA fabric (`zig/src/fabric`) | Flash Next: `zig/src/families/flashnext/tp.zig` splits rows and holds the whole model on each Mac. Kimi K3: `zig/src/families/kimi_k3/parallel.zig` is a split plan with no link yet | | Prompt path | segments; the prompt cache, used by Flash Next, GLM, Nemotron and the Qwen3.8-27B | per-family prompt kernels and snapshots | -| Serving | `engine_api` for every host; `lane_host` for Nemotron | Flash Next: its own host, one reply at a time | +| Serving | `engine_api` for every host; `lane_host` for Nemotron, and for Flash Next on CUDA | Flash Next on Metal: its own host, one reply at a time | **The order.** 1. GLM-5.3 builds the core pieces first: diff --git a/docs/recipes/cuda.md b/docs/recipes/cuda.md index c777a1cfb7..72bcb3550d 100644 --- a/docs/recipes/cuda.md +++ b/docs/recipes/cuda.md @@ -8,7 +8,7 @@ They are not the native Zig release's support matrix. The kernel paths, the prom (`--prefill-fp8`, `--no-prefill-fp8`, `--precision`) and the request and memory notes below are that engine's too, and the native binary refuses those flags. -The native binary registers only Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40). +The native binary registers Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), and Flash Next from its MLX 4-bit checkpoint with a captured kernel set, run on a GB10 ([Flash Next on CUDA](flash-next-cuda.md)). A user's receipt for 1.0.4 on a bare-metal RTX 3090, at 225 W and 350 W, is in [#592](https://github.com/ashhart/TensorFold/issues/592). GLM-5.3-Flash on CUDA, EXL3 loading and DFlash2 drafting remain Python-engine features, not native upgrade paths. Inspect the installed binary with `tensorfold-native --version` and `tensorfold-native capabilities --json`. diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md new file mode 100644 index 0000000000..853e255a94 --- /dev/null +++ b/docs/recipes/flash-next-cuda.md @@ -0,0 +1,104 @@ +# Qwen3.8 Flash Next on CUDA + +The native CUDA build serves Flash Next (`qwen4_exp`) on one NVIDIA GPU from an MLX affine 4-bit, group-32 checkpoint +with its MTP head. It replays the Python 0.6.1 engine's Triton kernels from a captured set beside its own CUDA +kernels, so it gives that engine's teacher-path tokens. @AdrianBinDC started the port in +[#476](https://github.com/ashhart/TensorFold/pull/476); it is tracked in +[#472](https://github.com/ashhart/TensorFold/issues/472). It has been run on a GB10 (sm_121) only. + +## Run + +The engine needs a captured kernel set (see [The kernel set](#the-kernel-set)); without `TENSORFOLD_CUDA_KERNELS` it +refuses to load. + +```bash +TENSORFOLD_CUDA_KERNELS=/path/to/flashnext-kernels tensorfold-native serve /path/to/checkpoint --temperature 0 +``` + +- `--context N` takes up to the model's 262,144 tokens; the default is the whole window. Past 2,048 keys each attention + layer reads 512 blocks of 4 keys a row, chosen by the model's block indexer, as the Python engine does. +- `--parallel N` admits N streams. Each stream's caches are its own: about 110 MiB of DeltaNet state and the n-gram + tail, plus 28 KiB a token of keys, values and indexer keys, and with drafts its MTP head's cache. +- `--no-drafts` serves without the MTP head; `"draft": false` on a request does the same for that request. +- `TF_FLASHNEXT_DRAFTS` (1 to 6, default 6) and `TF_FLASHNEXT_CONFIDENCE` (default 0.60) set the head's chain. +- The engine uploads 73.3 GiB of the checkpoint: the n-gram tables stay mapped from the shards, and the vision tower is + not loaded. +- It refuses another quantization or other shapes (config.json is checked at load, and the log names the field), + `--segments` and `--slide`, and a folder whose shared experts carry Living Weights neurons. + +## What the core runs and what the family runs + +The family plugs into the core's pieces and adds no scheduler: + +- The core's round loop and lane host drive every token. Drafted windows, shared rounds of streams without drafts, and + prompts a chunk at a time between other streams' rounds (`prefill_step`) all go through them. +- The core's prompt cache keeps and resumes a stream's state at the marks it asks for. The family copies the state + (`cuda_snap.zig`), with the MTP head's cache rows for a drafting stream; the head reads one token past a row, so + states keep with a lookahead of 1, as Nemotron's do. +- The routed experts over a window's or a chunk's rows run on the core's plan (`cuda/grouped.zig`) and its rows kernel + (`experts_rows.cu`, `cuda.experts.RowsExperts`), which gives every row the decode kernel's sums. +- Captured Triton launches go through `cuda.aot.Set`, so a launch with no exact variant takes the general one. + +The family supplies its config check, weight loading and packing, the layer graph on one batched pass +(`cuda_batch.zig`), and kernels for what is its own: the Gated DeltaNet chain, the sparse indexer's block selection, +the n-gram gather, the hyper-connection and attention launches, and the MTP head. + +## The kernel set + +The set is built from the Python 0.6.1 engine in a container (`IMAGE`, default tag `tensorfold-qwen38:v0.6.1`) with +the scripts in `zig/tests/cuda/flashnext`, then packed with `tools/zig/triton_aot_pack.py`: + +```bash +zig/tests/cuda/flashnext/capture.sh WORK/cap /path/to/hf/snapshot # the model's own run, about 78 GiB of GPU +zig/tests/cuda/flashnext/capture_rows.sh WORK/rows # the batched pass's forms, no checkpoint +zig/tests/cuda/flashnext/capture_mtp.sh WORK/mtp # the MTP head's own kernels +zig/tests/cuda/flashnext/capture_sparse.sh WORK/sparse # the sparse attention's kernels +zig/tests/cuda/flashnext/pack.sh KERNELS WORK/cap/out/cap WORK/rows/out WORK/mtp/out WORK/sparse/out +``` + +`capture.sh` also writes the prompts and the Python engine's references the checks below read. The set has +76 kernels; `capture_sparse.sh` writes the selection fixtures `flashnext-select` reads. + +## Checks + +- `zig build flashnext-prompt -Dflashnext_model=DIR -Dflashnext_kernels=SET -Dflashnext_ids=FILE -Dflashnext_count=32`: + 32 greedy tokens through the batched pass against the Python engine's teacher path + (`zig/tests/cuda/flashnext/lr7/results.json`, `results-sparse.json`). `TF_FLASHNEXT_PROMPT=rows` runs the row-by-row + reference instead; `TF_FLASHNEXT_DRAFTS=6` decodes each prompt serially and drafted and compares them token for + token (`TF_FLASHNEXT_SEED` for keyed sampled draws). +- `zig build flashnext-select -Dflashnext_fixtures=DIR`: the native block selection against Triton's lists. +- `tf-cuda-test experts-rows`: the rows kernel against the decode kernel's one-row launches, byte for byte. +- Served: `tools/bench_concurrent.py --alone --serial` (concurrent equal to solo, drafted equal to `"draft": false`), + and a prompt sent twice (the second resumed from a kept state, equal to the first). + +## Measurements + +On a GB10 (DGX Spark, driver 580.178, CUDA 13.0), an MLX 4-bit group-32 Flash Next checkpoint with its MTP head, +the 76-kernel set, greedy unless said, `tools/bench_openai.py` (256 tokens, median of 3) and +`tools/bench_concurrent.py` (256 tokens, aggregate tok/s): + +| | tok/s | +| --- | --- | +| One stream, drafted, code / chat | 82.6 / 61.8 | +| One stream, `--no-drafts`, code / chat | 35.0 / 34.9 | +| `--no-drafts`, 1 / 2 / 4 streams, sampled code | 35.0 / 53.0 / 77.3 | +| drafted, 1 / 2 / 4 streams, greedy code | 81.6 / 80.6 / 80.4 | + +Every concurrent reply equalled its solo run, and every drafted one its `"draft": false` run. Drafted streams verify +their windows a pass each, so drafted throughput does not grow with streams yet. A cold prompt's first token takes 2.65 +s at 3,001 tokens, 10.0 s at 12,001 and 37.9 s at 45,030 (a 65,536-token window, no prompt cache). While a 12,001-token +prompt runs, another stream's decode waits at most 1.3 s between tokens (one prompt chunk). Each stream at a +16,384-token window takes 0.74 GiB with drafts, 0.54 GiB without. + +## Not yet + +- The set replays the Python engine's Triton kernels for norms, hyper-connections, attention, routing and the head. The + core's shared row kernels (`cuda/norms.cu`) sum in another order, so moving onto them is its own change. +- The dense projections run a copy of the Python engine's group-32 lane matmul (`qmm.cu`); the core's `qmm_group` and + `lane_gemv` are built for groups of 64. +- No CUDA graphs or GPU-side rounds; drafted streams verify their windows one pass each, and the MTP head drafts one + stream at a time. +- Prompts run slower than the Python engine's: a cold 12,001-token prompt takes 10.0 s against its 6.7 s. +- Draws are made on the host from each row's scores. +- Replies equal the Python engine's teacher path and part from its served path after 10 to 36 tokens, from the + routed-experts kernel's rounding form. diff --git a/docs/recipes/qwen3.8-flash-next.md b/docs/recipes/qwen3.8-flash-next.md index af65c4fccd..60ac2381eb 100644 --- a/docs/recipes/qwen3.8-flash-next.md +++ b/docs/recipes/qwen3.8-flash-next.md @@ -2,7 +2,8 @@ The `qwen4_exp` family has Gated DeltaNet, sparse attention, MoE, hyper-connections and hashed n-gram embeddings. The native engine's supported Metal checkpoint uses MLX affine 6-bit weights in groups of 32 -and includes an MTP head. The native engine refuses the 4-bit/group-32 checkpoint at load. +and includes an MTP head. On a Mac the native engine refuses the 4-bit/group-32 checkpoint at load; on CUDA it serves +that one ([Flash Next on CUDA](flash-next-cuda.md)). ```bash tensorfold pull TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP @@ -10,8 +11,8 @@ tensorfold serve TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP --name bench --tempe ``` For native serving, see the [Metal runbook](../../RUNBOOK.md#flash-next-and-paired-metal-serving). -Native Flash Next has no CUDA backend yet; its port is tracked in -[#472](https://github.com/ashhart/TensorFold/issues/472). +Native Flash Next on CUDA serves the MLX 4-bit group-32 checkpoint from a captured kernel set: see +[Flash Next on CUDA](flash-next-cuda.md) ([#472](https://github.com/ashhart/TensorFold/issues/472)). The sections below describe the Python 0.6.6 engine, available at the [`v0.6.6` tag](https://github.com/ashhart/TensorFold/tree/v0.6.6). Their flags, variables and paths are that engine's, From 52b82fc59954ec299ebd597ced7aaa6a32eea759 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 15 Oct 2026 06:15:42 -0500 Subject: [PATCH 038/115] prompt cache: an entry removed while pinned is dead to lookups and freed only at its last unpin; a job reports resumed before it unpins --- zig/src/core/lane_host.zig | 2 +- zig/src/core/prompt_cache.zig | 13 +++++---- zig/src/core/prompt_cache_test.zig | 47 ++++++++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 6 deletions(-) diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 9ddbf6a8fe..73937e925c 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -85,7 +85,7 @@ pub const LaneHost = struct { const e = job.entry orelse return; job.entry = null; const store = job.host.cache orelse return; - store.unpin(e); + defer store.unpin(e); // after resumed, which reads e if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index 6fe20982f7..81d5cf4ed4 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -109,6 +109,7 @@ pub const Entry = struct { shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded pins: u32 = 0, // admitted jobs that restore it and have not yet: eviction passes it over + dead: bool = false, // removed while pinned: no lookup finds it, and its last unpin frees it }; pub const Counts = struct { hits: u64 = 0, misses: u64 = 0, kept: u64 = 0, evicted: u64 = 0, refused: u64 = 0, failed: u64 = 0 }; @@ -149,6 +150,7 @@ pub const Store = struct { } fn free(s: *Store, e: *Entry) void { + s.held -= e.bytes; s.family.vtable.drop(s.family.ptr, e.saved); s.gpa.free(e.tokens); s.gpa.free(e.decode_spans); @@ -158,18 +160,19 @@ pub const Store = struct { fn remove(s: *Store, i: usize) void { const e = s.entries.orderedRemove(i); - s.held -= e.bytes; - s.free(e); + e.dead = true; + if (e.pins == 0) s.free(e); } - /// `e` waits for a restore: no lookup or keep evicts it until `unpin` (jobs set up after it run their own first). + /// `e` waits for a restore: nothing frees it until its last `unpin` (removed meanwhile, no lookup finds it). pub fn pin(_: *Store, e: *Entry) *Entry { e.pins += 1; return e; } - pub fn unpin(_: *Store, e: *Entry) void { + pub fn unpin(s: *Store, e: *Entry) void { e.pins -= 1; + if (e.pins == 0 and e.dead) s.free(e); } /// The family's free storage kept for later saves. @@ -484,11 +487,11 @@ pub const Store = struct { const charged = if (s.family.vtable.charged) |f| f(s.family.ptr, saved) else bytes; e.* = .{ .tokens = tokens, .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .last = last, .shared = shared }; owns_decoded = false; + s.held += charged; // free takes it back s.entries.append(s.gpa, e) catch { s.free(e); return s.fail(at, error.OutOfMemory); }; - s.held += charged; s.counts.kept += 1; if (shared) learned.learn(s, e, starts); if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held); // spare storage only inside what the budget leaves diff --git a/zig/src/core/prompt_cache_test.zig b/zig/src/core/prompt_cache_test.zig index b408b3c7ad..cde079b745 100644 --- a/zig/src/core/prompt_cache_test.zig +++ b/zig/src/core/prompt_cache_test.zig @@ -421,3 +421,50 @@ test "an entry pinned for a restore still to come outlives the lookups set up af try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); try std.testing.expectEqual(s.entries.items.len, f.live); } + +test "two jobs pin one entry and the first restore fails: lookups miss it, and it lives until the second job unpins" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const xb = [_]u32{ 5, 5, 5, 5, 8 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + try std.testing.expect(s.pin((try s.lookup(a, &xb, 3, &.{}, &.{}, &.{})).entry.?) == e); + const held = s.held; + s.resumed(e, &xa, false); // the first job's restore failed: it prefilled from 0 + s.unpin(e); + try std.testing.expect(s.find(&xb, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 0), s.entries.items.len); + try std.testing.expectEqual(@as(usize, 1), f.live); // the second job's restore still reads it + try std.testing.expectEqual(held, s.held); + try std.testing.expectEqual(@as(u32, 1), e.pins); + s.resumed(e, &xb, true); // no longer kept: its report changes nothing + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); + try std.testing.expectEqual(@as(u64, 0), s.held); +} + +test "forget leaves a pinned entry to its restore" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + s.forget(&xa, 3, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 1), f.live); + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); +} From b45e26ac6fbf3795af57c6180b950b7781faaf59 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:14:38 -0500 Subject: [PATCH 039/115] Flash Next CUDA refuses the config.json settings its kernels fix: RMS epsilon, RoPE, gates, activations, layer roles cuda_config.zig checked the shapes and the quantization but let other settings through while the kernels evaluate fixed constants. semantics() now refuses, by field name, a checkpoint whose text config changes what the CUDA path hard-codes: rms_norm_eps (taken as fp32 1e-6), RoPE (rope_theta 10,000,000 in rope_parameters or at the top level, partial_rotary_factor 0.25, a type other than default, a scaling factor or rope_scaling), output_gate_type (sigmoid), hidden_act (silu), attention_bias, attn_output_gate, norm_topk_prob, mamba_ssm_dtype (float32), tied embeddings, dedicated MTP embeddings, indexer_budget and indexer_kv_heads, the n-gram layer (ple_layer_ids [2], ple_embed_dim, ple_conv_kernel_size 4, ngram_size 3, which is the conv's dilation), the layer roles (three linear-attention layers, then one full) and the MTP layer's (one full-attention layer at the same RoPE base). An absent field means the Python engine's default, as its config reader takes it. The n-gram hash's own fields (heads_per_ngram, the vocabulary base and divisor, split parts, seed) are read and honoured by the n-gram reader; mrope_section moves image rows only. Tests change each setting in turn and expect its name back. --- zig/src/families/flashnext/cuda_config.zig | 201 ++++++++++++++++++++- 1 file changed, 192 insertions(+), 9 deletions(-) diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig index 5160d3eaf6..3d87549150 100644 --- a/zig/src/families/flashnext/cuda_config.zig +++ b/zig/src/families/flashnext/cuda_config.zig @@ -1,4 +1,4 @@ -//! The checkpoint's config.json against the shapes the CUDA path is built for: a different one is refused at load. +//! The checkpoint's config.json against what the CUDA path is built for: another shape or setting is refused at load. const std = @import("std"); const wt = @import("cuda_prompt_wt.zig"); @@ -33,11 +33,64 @@ const shapes = [_]Want{ .{ .name = "mtp_num_hidden_layers", .value = 1 }, }; +/// RoPE's base: Scratch.inv's 10,000,000 (the Python engine's default when the config names none). +pub const rope_theta: f64 = 10_000_000; +/// The rotated share of a head: 64 of 256 features, `wt.half` pairs. +pub const partial_rotary: f64 = @as(f64, @floatFromInt(2 * wt.half)) / @as(f64, @floatFromInt(wt.head_dim)); + +/// Settings the kernels fix, absent meaning the value they assume: a whole number, a float, a word or a flag. +const Fixed = struct { name: []const u8, value: union(enum) { int: i64, float: f64, word: []const u8, flag: bool } }; + +/// The text config's semantic settings the evaluation hard-codes (norms, gates, activations, layer roles, n-grams). +const fixed = [_]Fixed{ + .{ .name = "output_gate_type", .value = .{ .word = "sigmoid" } }, + .{ .name = "hidden_act", .value = .{ .word = "silu" } }, + .{ .name = "attention_bias", .value = .{ .flag = false } }, + .{ .name = "attn_output_gate", .value = .{ .flag = true } }, + .{ .name = "norm_topk_prob", .value = .{ .flag = true } }, + .{ .name = "mamba_ssm_dtype", .value = .{ .word = "float32" } }, + .{ .name = "tie_word_embeddings", .value = .{ .flag = false } }, + .{ .name = "mtp_use_dedicated_embeddings", .value = .{ .flag = false } }, + .{ .name = "indexer_budget", .value = .{ .int = wt.index_top * wt.index_ratio } }, + .{ .name = "indexer_kv_heads", .value = .{ .int = 1 } }, + .{ .name = "ple_embed_dim", .value = .{ .int = wt.dims } }, + .{ .name = "ple_conv_kernel_size", .value = .{ .int = 4 } }, + .{ .name = "ngram_size", .value = .{ .int = 3 } }, +}; + fn int(o: std.json.ObjectMap, name: []const u8) ?i64 { const v = o.get(name) orelse return null; return if (v == .integer) v.integer else null; } +fn float(v: std.json.Value) ?f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => null, + }; +} + +fn object(o: std.json.ObjectMap, name: []const u8) ?std.json.ObjectMap { + const v = o.get(name) orelse return null; + return if (v == .object) v.object else null; +} + +/// A value present and not null. +fn given(o: std.json.ObjectMap, name: []const u8) ?std.json.Value { + const v = o.get(name) orelse return null; + return if (v == .null) null else v; +} + +fn holds(v: std.json.Value, want: @FieldType(Fixed, "value")) bool { + return switch (want) { + .int => |i| v == .integer and v.integer == i, + .float => |f| float(v) == f, + .word => |w| v == .string and std.mem.eql(u8, v.string, w), + .flag => |b| v == .bool and v.bool == b, + }; +} + /// The first config value that differs from what the CUDA path serves, or null when every one matches. pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u8 { const q = quant orelse return "quantization"; @@ -45,6 +98,45 @@ pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u if (int(q, "group_size") != 32) return "quantization.group_size"; if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return "quantization.mode"; for (shapes) |w| if (int(text, w.name) != w.value) return w.name; + return semantics(text); +} + +/// The first semantic setting the kernels would ignore: eps, RoPE, gates, activations, layer roles, the head's layer. +pub fn semantics(text: std.json.ObjectMap) ?[]const u8 { + const eps = float(given(text, "rms_norm_eps") orelse return "rms_norm_eps") orelse return "rms_norm_eps"; + if (@as(f32, @floatCast(eps)) != wt.eps) return "rms_norm_eps"; // the kernels take eps as fp32 + for (fixed) |f| if (given(text, f.name)) |v| if (!holds(v, f.value)) return f.name; + if (rope(text)) |field| return field; + if (given(text, "layer_types")) |v| { + if (v != .array or v.array.items.len != wt.layers_n) return "layer_types"; + for (v.array.items, 0..) |kind, i| { + if (kind != .string) return "layer_types"; + if (std.mem.eql(u8, kind.string, "linear_attention") != wt.linearLayer(i)) return "layer_types"; + } + } + const ple = given(text, "ple_layer_ids") orelse return "ple_layer_ids"; + if (ple != .array or ple.array.items.len != 1 or !holds(ple.array.items[0], .{ .int = wt.ple_layer + 1 })) return "ple_layer_ids"; + if (object(text, "mtp")) |m| { + if (given(m, "rope_theta")) |v| if (float(v) != rope_theta) return "mtp.rope_theta"; + if (given(m, "num_hidden_layers")) |v| if (!holds(v, .{ .int = 1 })) return "mtp.num_hidden_layers"; + if (given(m, "layer_types")) |v| { + if (v != .array or v.array.items.len != 1 or !holds(v.array.items[0], .{ .word = "full_attention" })) return "mtp.layer_types"; + } + } + return null; +} + +/// RoPE as the kernels apply it: base 10,000,000 on 64 of 256 features, no scaling (text positions only). +fn rope(text: std.json.ObjectMap) ?[]const u8 { + if (given(text, "rope_scaling") != null) return "rope_scaling"; + const params = object(text, "rope_parameters"); + const theta = (if (params) |p| given(p, "rope_theta") else null) orelse given(text, "rope_theta"); + if (theta) |v| if (float(v) != rope_theta) return "rope_theta"; + const partial = (if (params) |p| given(p, "partial_rotary_factor") else null) orelse given(text, "partial_rotary_factor"); + if (partial) |v| if (float(v) != partial_rotary) return "partial_rotary_factor"; + const p = params orelse return null; + for ([_][]const u8{ "rope_type", "type" }) |name| if (given(p, name)) |v| if (!holds(v, .{ .word = "default" })) return "rope_parameters.rope_type"; + if (given(p, "factor")) |v| if (float(v) != 1) return "rope_parameters.factor"; return null; } @@ -60,19 +152,48 @@ pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { const root = parsed.value.object; const text = if (root.get("text_config")) |v| (if (v == .object) v.object else root) else root; const quant = if (root.get("quantization")) |v| (if (v == .object) v.object else null) else null; - if (mismatch(text, quant)) |field| { - std.log.err("flashnext: config.json's {s} differs from the MLX 4-bit group-32 Flash Next the CUDA path serves", .{field}); - return error.UnsupportedCheckpoint; - } + var field = mismatch(text, quant); + if (field == null) if (given(root, "tie_word_embeddings")) |v| if (!holds(v, .{ .flag = false })) { + field = "tie_word_embeddings"; + }; + std.log.err("flashnext: config.json's {s} differs from the MLX 4-bit group-32 Flash Next the CUDA path serves", .{field orelse return}); + return error.UnsupportedCheckpoint; +} + +/// The served config's text fields, as the checked-in checkpoint has them (a test's base to change one at a time). +fn served(gpa: std.mem.Allocator, text: *std.json.ObjectMap, rope_params: *std.json.ObjectMap, mtp: *std.json.ObjectMap, kinds: *std.json.Array, ple: *std.json.Array) !void { + for (shapes) |w| try text.put(gpa, w.name, .{ .integer = w.value }); + try text.put(gpa, "rms_norm_eps", .{ .float = 1e-06 }); + try text.put(gpa, "output_gate_type", .{ .string = "sigmoid" }); + try text.put(gpa, "hidden_act", .{ .string = "silu" }); + try text.put(gpa, "attention_bias", .{ .bool = false }); + try text.put(gpa, "partial_rotary_factor", .{ .float = 0.25 }); + try text.put(gpa, "indexer_budget", .{ .integer = 2048 }); + try text.put(gpa, "ple_conv_kernel_size", .{ .integer = 4 }); + try text.put(gpa, "ngram_size", .{ .integer = 3 }); + for (0..wt.layers_n) |i| try kinds.append(.{ .string = if (wt.linearLayer(i)) "linear_attention" else "full_attention" }); + try text.put(gpa, "layer_types", .{ .array = kinds.* }); + try ple.append(.{ .integer = 2 }); + try text.put(gpa, "ple_layer_ids", .{ .array = ple.* }); + try rope_params.put(gpa, "rope_theta", .{ .integer = 10_000_000 }); + try rope_params.put(gpa, "partial_rotary_factor", .{ .float = 0.25 }); + try rope_params.put(gpa, "type", .{ .string = "default" }); + try text.put(gpa, "rope_parameters", .{ .object = rope_params.* }); + try mtp.put(gpa, "rope_theta", .{ .integer = 10_000_000 }); + try text.put(gpa, "mtp", .{ .object = mtp.* }); } test "the served shapes pass and a 6-bit checkpoint or another width is named" { - const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); var text: std.json.ObjectMap = .empty; - defer text.deinit(gpa); - for (shapes) |w| try text.put(gpa, w.name, .{ .integer = w.value }); + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); var quant: std.json.ObjectMap = .empty; - defer quant.deinit(gpa); try quant.put(gpa, "bits", .{ .integer = 4 }); try quant.put(gpa, "group_size", .{ .integer = 32 }); try std.testing.expectEqual(@as(?[]const u8, null), mismatch(text, quant)); @@ -83,3 +204,65 @@ test "the served shapes pass and a 6-bit checkpoint or another width is named" { try std.testing.expectEqualStrings("hidden_size", mismatch(text, quant).?); try std.testing.expectEqualStrings("quantization", mismatch(text, null).?); } + +test "a changed epsilon, RoPE, gate, activation or layer role is refused by name" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); + const Change = struct { field: []const u8, value: std.json.Value, refused: []const u8 }; + const changes = [_]Change{ + .{ .field = "rms_norm_eps", .value = .{ .float = 1e-5 }, .refused = "rms_norm_eps" }, + .{ .field = "output_gate_type", .value = .{ .string = "silu" }, .refused = "output_gate_type" }, + .{ .field = "hidden_act", .value = .{ .string = "gelu" }, .refused = "hidden_act" }, + .{ .field = "attention_bias", .value = .{ .bool = true }, .refused = "attention_bias" }, + .{ .field = "attn_output_gate", .value = .{ .bool = false }, .refused = "attn_output_gate" }, + .{ .field = "norm_topk_prob", .value = .{ .bool = false }, .refused = "norm_topk_prob" }, + .{ .field = "rope_scaling", .value = .{ .string = "yarn" }, .refused = "rope_scaling" }, + .{ .field = "indexer_budget", .value = .{ .integer = 4096 }, .refused = "indexer_budget" }, + .{ .field = "ngram_size", .value = .{ .integer = 4 }, .refused = "ngram_size" }, + .{ .field = "ple_layer_ids", .value = .{ .array = .init(gpa) }, .refused = "ple_layer_ids" }, + .{ .field = "layer_types", .value = .{ .array = .init(gpa) }, .refused = "layer_types" }, + }; + for (changes) |c| { + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); + try std.testing.expectEqual(@as(?[]const u8, null), semantics(text)); + try text.put(gpa, c.field, c.value); + try std.testing.expectEqualStrings(c.refused, semantics(text) orelse "accepted"); + } +} + +test "RoPE's base and type inside rope_parameters, and the MTP layer's base, are checked too" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); + const Case = struct { in_mtp: bool, field: []const u8, value: std.json.Value, refused: []const u8 }; + const cases = [_]Case{ + .{ .in_mtp = false, .field = "rope_theta", .value = .{ .integer = 1_000_000 }, .refused = "rope_theta" }, + .{ .in_mtp = false, .field = "partial_rotary_factor", .value = .{ .float = 0.5 }, .refused = "partial_rotary_factor" }, + .{ .in_mtp = false, .field = "type", .value = .{ .string = "yarn" }, .refused = "rope_parameters.rope_type" }, + .{ .in_mtp = false, .field = "factor", .value = .{ .float = 4 }, .refused = "rope_parameters.factor" }, + .{ .in_mtp = true, .field = "rope_theta", .value = .{ .integer = 5_000_000 }, .refused = "mtp.rope_theta" }, + .{ .in_mtp = true, .field = "layer_types", .value = .{ .array = .init(gpa) }, .refused = "mtp.layer_types" }, + }; + for (cases) |c| { + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); + if (c.in_mtp) { + try mtp.put(gpa, c.field, c.value); + try text.put(gpa, "mtp", .{ .object = mtp }); + } else { + try rope_params.put(gpa, c.field, c.value); + try text.put(gpa, "rope_parameters", .{ .object = rope_params }); + } + try std.testing.expectEqualStrings(c.refused, semantics(text) orelse "accepted"); + } +} From 9fa628bc5d558ddbe48e461763b548b1ca7485f9 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:14:38 -0500 Subject: [PATCH 040/115] Flash Next CUDA: a failed zero fill frees the buffer it was filling zeros() counted a buffer only after its fill, so when the fill failed the errdefer freed the buffers before it and leaked the one just allocated. The loop is now made(): a buffer is counted as soon as it is allocated, before its fill, so a failure frees every buffer made, the last one too. A test with a fake buffer fails the fill at the first, second and fourth of four and checks none is left live. --- zig/src/families/flashnext/cuda_prompt.zig | 62 +++++++++++++++++++--- 1 file changed, 56 insertions(+), 6 deletions(-) diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 7e022fa2cf..17fd7463bb 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -61,15 +61,32 @@ pub fn argmax(raw: []const u8) u32 { /// `n` zeroed buffers, the fill on the compute stream: a legacy-stream memset returns before it lands. fn zeros(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: cuda.Stream, n: usize, len: usize) ![]cuda.DeviceBuffer { - const out = try gpa.alloc(cuda.DeviceBuffer, n); - var i: usize = 0; + const Zeroed = struct { + driver: *cuda.Driver, + stream: cuda.Stream, + len: usize, + fn make(z: @This()) !cuda.DeviceBuffer { + return cuda.DeviceBuffer.alloc(z.driver, z.len); + } + fn fill(z: @This(), b: *cuda.DeviceBuffer) !void { + try b.fill8(0, z.stream.handle); + } + }; + return made(cuda.DeviceBuffer, gpa, n, Zeroed{ .driver = driver, .stream = stream, .len = len }); +} + +/// `n` buffers from `src.make`, each through `src.fill`; on a failure every buffer made so far is freed, the last too. +fn made(comptime B: type, gpa: std.mem.Allocator, n: usize, src: anytype) ![]B { + const out = try gpa.alloc(B, n); + var held: usize = 0; errdefer { - for (out[0..i]) |*b| b.free(); + for (out[0..held]) |*b| b.free(); gpa.free(out); } - while (i < n) : (i += 1) { - out[i] = try cuda.DeviceBuffer.alloc(driver, len); - try out[i].fill8(0, stream.handle); + while (held < n) { + out[held] = try src.make(); + held += 1; // counted before its fill, so a failed fill frees it with the rest + try src.fill(&out[held - 1]); } return out; } @@ -375,3 +392,36 @@ test "a sequence's device bytes grow by the twelve attention layers' keys, value try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5) - seqBytes(4)); try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); } + +test "a failed fill frees the buffer it was filling and every one before it" { + const Fake = struct { + live: *usize, + pub fn free(b: *@This()) void { + b.live.* -= 1; + } + }; + const Src = struct { + live: *usize, + fills: *usize, + fail_at: usize, + fn make(s: @This()) !Fake { + s.live.* += 1; + return .{ .live = s.live }; + } + fn fill(s: @This(), _: *Fake) !void { + s.fills.* += 1; + if (s.fills.* == s.fail_at) return error.FillFailed; + } + }; + for ([_]usize{ 1, 2, 4 }) |fail_at| { + var live: usize = 0; + var fills: usize = 0; + try std.testing.expectError(error.FillFailed, made(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); + try std.testing.expectEqual(@as(usize, 0), live); + } + var live: usize = 0; + var fills: usize = 0; + const all = try made(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); + defer std.testing.allocator.free(all); + try std.testing.expectEqual(@as(usize, 3), live); +} From 475df9299098683ac2403ac6ba8831d086524bab Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:14:38 -0500 Subject: [PATCH 041/115] Flash Next CUDA: a pass's uploads read pinned staging the scratch owns until the copies land Pass.stage() copied a chunk's token ids and attention positions, and chunk() the rows' n-gram embeddings and their group sums, from allocator memory freed when the function returned, with cuMemcpyHtoDAsync on the compute stream. For pageable memory CUDA documents that call only as "might be synchronous": the driver stages the bytes before it returns on the drivers we run, but nothing promises it, and freed memory may be reused before a copy that was not staged lands. The scratch now owns one pinned buffer for all four (ids, positions, n-gram rows and sums, at 16-byte-aligned offsets) and an event recorded after each pass's last upload; the next pass waits on the event before it writes the buffer (the gather thread writes the n-gram rows into it directly). The row-by-row check path's two uploads from stack values now wait for their copy before the frame goes. --- zig/src/families/flashnext/cuda_batch.zig | 20 +++++----- .../families/flashnext/cuda_batch_state.zig | 40 +++++++++++++++++++ .../families/flashnext/cuda_prompt_step.zig | 3 +- zig/src/families/flashnext/cuda_rows.zig | 3 +- 4 files changed, 55 insertions(+), 11 deletions(-) diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 1b3c0aa7a5..8dcdf58e3c 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -150,6 +150,7 @@ pub fn Pass(comptime Tri: type) type { try d.check(d.api.cuMemcpyDtoDAsync_v2(dst, src, n, self.stream.handle), "cuMemcpyDtoDAsync"); } + /// A queued copy from pinned host bytes (the scratch's staging), which must stay unchanged until it lands. pub fn upload(self: *Self, dst: u64, bytes: []const u8) !void { const d = self.stream.d; try d.check(d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, self.stream.handle), "cuMemcpyHtoDAsync"); @@ -359,21 +360,21 @@ pub fn Pass(comptime Tri: type) type { // ---- the pass ------------------------------------------------------------------------------------------ - /// Token ids, positions and (when the n-gram layer is there) each row's n-gram embedding onto the device. + /// Token ids and positions onto the device from the scratch's pinned copies, which stay put until they land. pub fn stage(self: *Self, ids: []const u32) !void { const s = self.s; const n = ids.len; - const id32 = try self.gpa.alloc(i32, n); - defer self.gpa.free(id32); + try s.ready(); + const id32 = s.hostIds(n); for (id32, ids) |*o, v| o.* = @intCast(v); try self.upload(s.ids, std.mem.sliceAsBytes(id32)); const blocks = if (self.lanes.len > 0) n else cdiv(n, att_rows); // one int a block (a round: a row), 16 bytes apart: _chunks and _merge take POS0 16-byte aligned - const pos = try self.gpa.alloc(i32, blocks * 4); - defer self.gpa.free(pos); + const pos = s.hostPos(blocks * 4); @memset(pos, 0); for (0..blocks) |b| pos[4 * b] = @intCast(if (self.lanes.len > 0) self.lanes[b].pos else self.pos0 + b * att_rows); try self.upload(s.pos, std.mem.sliceAsBytes(pos)); + try s.uploaded(self.stream); } /// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. @@ -425,10 +426,10 @@ pub fn Pass(comptime Tri: type) type { self.rows = n; self.pos0 = pos0; self.pending = null; - const emb = try self.gpa.alloc(u16, n * dims); - defer self.gpa.free(emb); - const xs = try self.gpa.alloc(f32, n * (dims / 32)); - defer self.gpa.free(xs); + // the gather writes the scratch's pinned rows, so the last pass's uploads must have read them first + try s.ready(); + const emb = s.hostEmb(n); + const xs = s.hostXs(n); var gather: Gather = .{ .ngram = self.m.ngram, .history = self.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = self.io, .lanes = self.lanes }; var thread: ?std.Thread = try std.Thread.spawn(.{}, Gather.main, .{&gather}); defer if (thread) |th| th.join(); @@ -443,6 +444,7 @@ pub fn Pass(comptime Tri: type) type { if (profile()) std.debug.print("batch: {d} rows' n-gram rows gathered on the host in {d:.1} ms\n", .{ n, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); try self.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); try self.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); + try s.uploaded(self.stream); if (self.pending != null) try self.writeback(0, n); try self.pleBlock(lw, ids); } diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index b91273d58b..6e9fd2e437 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -199,6 +199,10 @@ pub const Scratch = struct { logits: u64 = 0, inv: u64 = 0, dummy: u64 = 0, + /// Pinned host copies of what a pass uploads (ids, positions, n-gram rows and sums), owned until the copies land. + host: ?cuda.HostBuffer = null, + /// Recorded on the compute stream after a pass's last upload: the next pass waits on it before writing `host`. + staged: ?cuda.Event = null, fn add(s: *Scratch, gpa: std.mem.Allocator, d: *const cuda.Driver, len: usize) !u64 { var b = try cuda.DeviceBuffer.alloc(d, @max(len, 256)); @@ -207,6 +211,38 @@ pub const Scratch = struct { return b.ptr; } + /// Byte offsets of the staged ids, positions, n-gram rows and their sums in `host`, and its length. + fn hostAt(cap: usize) [5]usize { + const pos = std.mem.alignForward(usize, cap * 4, 16); + const emb = pos + (cap + 1) * 16; + const xs = emb + cap * dims * 2; + return .{ 0, pos, emb, xs, xs + cap * (dims / 32) * 4 }; + } + + /// Waits until the last pass's uploads have read `host` (they are queued behind its kernels), so it may change. + pub fn ready(s: *const Scratch) !void { + if (s.staged) |e| try e.synchronize(); + } + + /// Marks `host` read once the stream reaches here: every upload queued before this point has copied its bytes. + pub fn uploaded(s: *const Scratch, stream: cuda.Stream) !void { + if (s.staged) |e| try e.record(stream); + } + + /// The staged ids [n], positions [n] (in 16-byte rows), n-gram rows [n, dims] and their 32-group sums. + pub fn hostIds(s: *const Scratch, n: usize) []i32 { + return @alignCast(std.mem.bytesAsSlice(i32, s.host.?.bytes[hostAt(s.cap)[0]..][0 .. n * 4])); + } + pub fn hostPos(s: *const Scratch, n: usize) []i32 { + return @alignCast(std.mem.bytesAsSlice(i32, s.host.?.bytes[hostAt(s.cap)[1]..][0 .. n * 4])); + } + pub fn hostEmb(s: *const Scratch, n: usize) []u16 { + return @alignCast(std.mem.bytesAsSlice(u16, s.host.?.bytes[hostAt(s.cap)[2]..][0 .. n * dims * 2])); + } + pub fn hostXs(s: *const Scratch, n: usize) []f32 { + return @alignCast(std.mem.bytesAsSlice(f32, s.host.?.bytes[hostAt(s.cap)[3]..][0 .. n * (dims / 32) * 4])); + } + /// Bytes `alloc(cap)` takes on the device. pub fn bytes(cap: usize) usize { const pairs = cap * slots; @@ -294,10 +330,14 @@ pub const Scratch = struct { s.dummy = try s.add(gpa, d, 4096); try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); try stream.synchronize(); + s.host = try cuda.HostBuffer.alloc(d, hostAt(cap)[4]); + s.staged = try cuda.Event.init(d, false); return s; } pub fn free(s: *Scratch, gpa: std.mem.Allocator) void { + if (s.staged) |*e| e.deinit(); + if (s.host) |*h| h.free(); for (s.bufs.items) |*b| b.free(); s.bufs.deinit(gpa); s.* = undefined; diff --git a/zig/src/families/flashnext/cuda_prompt_step.zig b/zig/src/families/flashnext/cuda_prompt_step.zig index f80b0c2f01..36d364152f 100644 --- a/zig/src/families/flashnext/cuda_prompt_step.zig +++ b/zig/src/families/flashnext/cuda_prompt_step.zig @@ -237,8 +237,9 @@ pub fn gdnStep(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Strea pub fn attnStep(comptime Tri: type, driver: *cuda.Driver, stream: *cuda.Stream, tri: Tri, s: *Scratch, proj_q: Q4, out_q: Q4, q_scale: u64, k_scale: u64, i_scale: u64, kc: u64, vc: u64, ikc: u64, pos: i32) !void { try qmm.matmul(driver, stream.*, s.mixed.ptr, s.xsm.ptr, proj_q.w.ptr, proj_q.s.ptr, proj_q.b.ptr, s.pa.ptr, 1, proj_n, dims); - // stream-ordered, like the ids: the previous row's queued attention kernels still read s.pos + // stream-ordered (queued kernels still read s.pos), then waited for: `pos` lives on this frame try s.pos.uploadAsync(0, std.mem.asBytes(&pos), stream.handle); + try stream.synchronize(); try tri.attnPrep(s.pa.ptr, s.pos.ptr, q_scale, k_scale, i_scale, s.inv.ptr, s.q.ptr, kc, vc, s.iq.ptr, ikc, s.dummy.ptr, s.dummy.ptr, eps, 1, 0); // attention.attention, dense: chunks of 512 keys up to the row's position (the row-by-row path stays dense) const keys: usize = @as(usize, @intCast(pos)) + 1; diff --git a/zig/src/families/flashnext/cuda_rows.zig b/zig/src/families/flashnext/cuda_rows.zig index a3081891e8..e005b1bde2 100644 --- a/zig/src/families/flashnext/cuda_rows.zig +++ b/zig/src/families/flashnext/cuda_rows.zig @@ -109,8 +109,9 @@ const Pass = struct { fn embedRow(comptime Tri: type, p: *Pass, tri: Tri, token: u32, h: u64) !void { const id: i32 = @intCast(token); - // on the compute stream: a synchronous copy can land before an earlier row's queued _embed reads its id + // on the compute stream (queued _embed launches still read ids), then waited for: `id` lives on this frame try p.s.ids.uploadAsync(0, std.mem.asBytes(&id), p.stream.handle); + try p.stream.synchronize(); try tri.embed(p.s.ids.ptr, p.table.w.ptr, p.table.s.ptr, p.table.b.ptr, h, dims, streams, 1); } From 1e86036438a5ca6c6acf02f2d333385de6575fac Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:14:38 -0500 Subject: [PATCH 042/115] Core: router logits and softmax top-k rows for every CUDA family (route_rows.cu) tf_route_logits gives OUT[m, e] = x[m] . w[e] in fp32 for bf16 rows and router rows on mma tiles: one chain over K in 16-wide steps from zero, the inputs staged with cp.async and read back with ldmatrix. A row's sums never depend on the row count or the tile, so the few-row form (up to 16 rows, three 256-wide stages in flight) and the prompt form (two 128-wide stages) give the same bits. tf_route_topk takes a row's top k of its first E logits (the lower id among equal ones) with weights bf16(e^(l_k - l_0) / their sum, summed in pick order), and optionally a shared expert as slot k at bf16(sigmoid(bf16(logit E))), the softmax routing Qwen-style MoE families use. core/route_ref.zig is the float64 reference, and tf-cuda-test route-rows checks both kernels against it at 1, 3, 16, 17 and 40 rows (logits within the fp32 dot's bound, picks equal to the reference rule's on the kernel's own logits, a row of ties too, weights and the shared gate within half a bf16 ulp plus 2^-18), and every row count's bits equal. --- zig/build/cuda.zig | 1 + zig/kernels/cuda/route_rows.cu | 167 ++++++++++++++++++++++++++++ zig/src/core/root.zig | 2 + zig/src/core/route_ref.zig | 42 +++++++ zig/src/cuda/kernels.zig | 1 + zig/src/cuda/root.zig | 2 + zig/src/cuda/route_rows.zig | 73 ++++++++++++ zig/tests/cuda/main.zig | 3 + zig/tests/cuda/route_rows_tests.zig | 145 ++++++++++++++++++++++++ 9 files changed, 436 insertions(+) create mode 100644 zig/kernels/cuda/route_rows.cu create mode 100644 zig/src/core/route_ref.zig create mode 100644 zig/src/cuda/route_rows.zig create mode 100644 zig/tests/cuda/route_rows_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 89f4d77503..6a200ad469 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -44,6 +44,7 @@ const kernels = [_]Kernel{ .{ .name = "nemotron_keyed", .flags = glue }, .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig) .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig) + .{ .name = "route_rows", .flags = glue }, // ours: router logits and softmax top-k rows (cuda/route_rows.zig) .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts diff --git a/zig/kernels/cuda/route_rows.cu b/zig/kernels/cuda/route_rows.cu new file mode 100644 index 0000000000..9407493f75 --- /dev/null +++ b/zig/kernels/cuda/route_rows.cu @@ -0,0 +1,167 @@ +// Softmax top-k routing rows for every CUDA family (float64 reference: core/route_ref.zig), a row's bits its own. + +// Router logits on mma tiles, K ascending from zero (a row's sums never depend on the row count), and top-k rows. + +#include +#include +#include + +#include "glue_math.cuh" +#include "qmm_frag.cuh" + +using namespace tf_glue; + +namespace tf_route { + +constexpr int NT = 128; +constexpr int MOST = 8; // logits a thread holds: up to NT * MOST experts + +// The block's largest (value, index), lower index on ties; slots alternate by `parity`: one barrier an iteration. +__device__ __forceinline__ void block_best(float& v, int& i, float (*sv)[NT / 32], int (*si)[NT / 32], int parity) { + for (int off = 16; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, v, off); + const int oi = __shfl_xor_sync(0xffffffffu, i, off); + if (ov > v || (ov == v && oi < i)) { + v = ov; + i = oi; + } + } + const int warp = threadIdx.x >> 5; + if ((threadIdx.x & 31) == 0) { + sv[parity][warp] = v; + si[parity][warp] = i; + } + __syncthreads(); + v = sv[parity][0]; + i = si[parity][0]; + for (int w = 1; w < NT / 32; ++w) + if (sv[parity][w] > v || (sv[parity][w] == v && si[parity][w] < i)) { + v = sv[parity][w]; + i = si[parity][w]; + } +} + +} // namespace tf_route + +using namespace tf_route; + +// OUT[m, e] = x[m] . w[e], a warp of 8 experts: one mma chain over K from zero, STAGES of KC staged (D % 16 == 0). +template +__device__ __forceinline__ void route_logits(const __nv_bfloat16* __restrict__ X, int x_ld, const __nv_bfloat16* __restrict__ W, + float* __restrict__ OUT, int M, int D, int E) { + constexpr int LD = KC + 8, EB = 8 * WARPS, ROWS = 16 + EB; + extern __shared__ __align__(16) unsigned char smem_raw[]; + auto tile = reinterpret_cast<__nv_bfloat16 (*)[ROWS][LD]>(smem_raw); // [STAGES] 16 input rows, then EB expert rows + const int m0 = blockIdx.x * 16, e0 = blockIdx.y * EB; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + const int stages = (D + KC - 1) / KC; + const int xr = min(16, M - m0); // input rows past M stay zero: written once, never staged + for (int v = threadIdx.x; v < STAGES * (16 - xr) * (KC / 8); v += blockDim.x) { + const int per = (16 - xr) * (KC / 8); + *reinterpret_cast(&tile[v / per][xr + (v % per) / (KC / 8)][8 * (v % (KC / 8))]) = make_uint4(0, 0, 0, 0); + } + auto stage = [&](int st) { + const int k0 = st * KC, buf = st % STAGES; + for (int v = threadIdx.x; v < (xr + EB) * (KC / 8); v += blockDim.x) { + const int i = v / (KC / 8), k = k0 + 8 * (v % (KC / 8)); + const bool x_row = i < xr; + const int row = x_row ? i : 16 + i - xr; + const int r = x_row ? m0 + i : e0 + i - xr; + const bool read = k < D && (x_row || r < E); + const __nv_bfloat16* src = x_row ? X + static_cast(r) * x_ld + (read ? k : 0) + : W + static_cast(read ? r : 0) * D + (read ? k : 0); + qmm_frag::cp16z(&tile[buf][row][8 * (v % (KC / 8))], src, read); + } + }; + float acc[4] = {}; + for (int st = 0; st < STAGES - 1; ++st) { + if (st < stages) stage(st); + qmm_frag::commit(); + } + for (int st = 0; st < stages; ++st) { + if (st + STAGES - 1 < stages) stage(st + STAGES - 1); + qmm_frag::commit(); + qmm_frag::wait(); + __syncthreads(); + const int buf = st % STAGES; + const int steps = min(KC, D - st * KC) / 16; + for (int kk = 0; kk < steps; ++kk) { + uint32_t a[4], b[2]; + qmm_frag::ldmatrix4(a, &tile[buf][lane % 16][kk * 16 + (lane / 16) * 8]); + qmm_frag::ldmatrix2(b, &tile[buf][16 + warp * 8 + lane % 8][kk * 16 + ((lane / 8) % 2) * 8]); + qmm_frag::mma(acc, a, b[0], b[1]); + } + __syncthreads(); + } + const int e = e0 + warp * 8 + 2 * t; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = m0 + g + 8 * h; + if (r >= M) continue; + if (e < E) OUT[static_cast(r) * E + e] = acc[2 * h]; + if (e + 1 < E) OUT[static_cast(r) * E + e + 1] = acc[2 * h + 1]; + } +} + +// route_logits over 32 experts a block (prompt chunks: the input rows are read once per 32 experts). +extern "C" __global__ void __launch_bounds__(128) tf_route_logits(const __nv_bfloat16* __restrict__ X, int x_ld, + const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, + int M, int D, int E) { + route_logits<4, 2, 128>(X, x_ld, W, OUT, M, D, E); +} + +// route_logits with three stages of 256 in flight (up to 16 rows, where the router rows' reads set its time). +extern "C" __global__ void __launch_bounds__(128) tf_route_logits_few(const __nv_bfloat16* __restrict__ X, int x_ld, + const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, + int M, int D, int E) { + route_logits<4, 3, 256>(X, x_ld, W, OUT, M, D, E); +} + +// A row's TOPK largest of NE logits (lower id on ties), bf16(e^(l_k - l_0) / sum); `shared`: slot TOPK gated by NE. +extern "C" __global__ void __launch_bounds__(NT) tf_route_topk(const float* __restrict__ L, int NL, int NE, int TOPK, + int* __restrict__ PICK, float* __restrict__ WTS, int SLOTS, + int shared) { + __shared__ float sv[2][NT / 32]; + __shared__ int si[2][NT / 32]; + const int r = blockIdx.x; + const float* row = L + static_cast(r) * NL; + float v[MOST]; +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int i = threadIdx.x + j * NT; + v[j] = i < NE ? row[i] : -INFINITY; + } + float top = 0.0f, total = 0.0f, mine = 0.0f; + int pick = 0; + for (int k = 0; k < TOPK; ++k) { + float bv = -INFINITY; + int bi = 0x7fffffff; +#pragma unroll + for (int j = 0; j < MOST; ++j) { + const int i = threadIdx.x + j * NT; + if (v[j] > bv || (v[j] == bv && i < bi)) { + bv = v[j]; + bi = i; + } + } + block_best(bv, bi, sv, si, k & 1); + if (k == 0) top = bv; +#pragma unroll + for (int j = 0; j < MOST; ++j) + if (threadIdx.x + j * NT == bi) v[j] = -INFINITY; + const float ex = tf_exp(__fsub_rn(bv, top)); + total = __fadd_rn(total, ex); + if (threadIdx.x == k) { + mine = ex; + pick = bi; + } + } + if (threadIdx.x < TOPK) { + PICK[static_cast(r) * SLOTS + threadIdx.x] = pick; + WTS[static_cast(r) * SLOTS + threadIdx.x] = rbf(__fdiv_rn(mine, total)); + } + if (shared && threadIdx.x == 0) { + PICK[static_cast(r) * SLOTS + TOPK] = NE; + WTS[static_cast(r) * SLOTS + TOPK] = rbf(tf_sigmoid(rbf(row[NE]))); + } +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 736248a93d..b6c0d6cf4c 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -18,6 +18,7 @@ pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); pub const act_ref = @import("act_ref.zig"); +pub const route_ref = @import("route_ref.zig"); test { _ = safetensors; @@ -29,6 +30,7 @@ test { _ = tuning; _ = norms_ref; _ = act_ref; + _ = route_ref; _ = exl3_format; _ = exlr8; _ = @import("exl3_rect_test.zig"); diff --git a/zig/src/core/route_ref.zig b/zig/src/core/route_ref.zig new file mode 100644 index 0000000000..21200dae9f --- /dev/null +++ b/zig/src/core/route_ref.zig @@ -0,0 +1,42 @@ +//! Softmax top-k routing in float64, the reference every backend's routing kernels are checked against. +const std = @import("std"); + +/// x . w for one router row. +pub fn dot(x: []const f64, w: []const f64) f64 { + var t: f64 = 0; + for (x, w) |a, b| t += a * b; + return t; +} + +/// An fp32 dot's tolerance: every product exact, at most one rounding a term of 2^-24 of the running |sum|. +pub fn dotTolerance(x: []const f64, w: []const f64) f64 { + var t: f64 = 0; + for (x, w) |a, b| t += @abs(a * b); + return t * @as(f64, @floatFromInt(x.len)) * std.math.pow(f64, 2, -24) + 1e-30; +} + +/// The `ids.len` largest logits (the lower id among equal ones), largest first, and weights e^(l_k - l_0) / their sum. +pub fn topk(logits: []const f64, ids: []usize, weights: []f64) void { + for (ids, 0..) |*id, k| { + var best: ?usize = null; + for (logits, 0..) |v, i| { + if (std.mem.indexOfScalar(usize, ids[0..k], i) != null) continue; + if (best == null or v > logits[best.?]) best = i; + } + id.* = best.?; + } + var total: f64 = 0; + for (ids) |id| total += @exp(logits[id] - logits[ids[0]]); + for (ids, weights) |id, *w| w.* = @exp(logits[id] - logits[ids[0]]) / total; +} + +test "the top picks are the largest, ties to the lower id, weights summing to one" { + const l = [_]f64{ 1, 3, 3, -1, 2 }; + var ids: [3]usize = undefined; + var w: [3]f64 = undefined; + topk(&l, &ids, &w); + try std.testing.expectEqualSlices(usize, &.{ 1, 2, 4 }, &ids); + try std.testing.expectApproxEqAbs(@as(f64, 1), w[0] + w[1] + w[2], 1e-12); + try std.testing.expectEqual(w[0], w[1]); + try std.testing.expect(dotTolerance(&.{ 1, 2 }, &.{ 3, -4 }) > 0); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 72868085a0..a351ea7e25 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -51,6 +51,7 @@ pub const nemotron_attention: []const u8 = if (available) &Blob("fatbin_nemotron pub const nemotron_keyed: []const u8 = if (available) &Blob("fatbin_nemotron_keyed").bytes else &.{}; pub const norms: []const u8 = if (available) &Blob("fatbin_norms").bytes else &.{}; pub const embed_rows: []const u8 = if (available) &Blob("fatbin_embed_rows").bytes else &.{}; +pub const route_rows: []const u8 = if (available) &Blob("fatbin_route_rows").bytes else &.{}; /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn). pub const gdn_symbols = struct { diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..dd9b45046c 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -33,6 +33,7 @@ pub const grouped = @import("grouped.zig"); pub const qlinear = @import("qlinear.zig"); pub const norms = @import("norms.zig"); pub const embed_rows = @import("embed_rows.zig"); +pub const route_rows = @import("route_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); @@ -43,6 +44,7 @@ test { _ = @import("qlinear.zig"); _ = @import("norms.zig"); _ = @import("embed_rows.zig"); + _ = @import("route_rows.zig"); _ = @import("experts.zig"); _ = @import("sample.zig"); _ = launch; diff --git a/zig/src/cuda/route_rows.zig b/zig/src/cuda/route_rows.zig new file mode 100644 index 0000000000..665463a1df --- /dev/null +++ b/zig/src/cuda/route_rows.zig @@ -0,0 +1,73 @@ +//! route_rows.cu for every CUDA family: router logits on mma tiles and softmax top-k rows, a row's bits its own. +const std = @import("std"); +const kernels = @import("kernels.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); + +const Driver = @import("driver.zig").Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// Logits a top-k row reads at most (the kernel's NT * MOST) and picks it keeps (its scratch). +pub const most_experts = 1024; +pub const most_picks = 128; + +/// tf_route_logits' staged rows (16 inputs and 32 experts, padded by 8): two stages of 128, or three of 256. +const many_shared: u32 = 2 * 48 * (128 + 8) * 2; +const few_shared: u32 = 3 * 48 * (256 + 8) * 2; + +pub const RouteRows = struct { + mod: Module, + logits_fn: Function, + few_fn: Function, + topk_fn: Function, + + pub fn load(d: *const Driver) !RouteRows { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.route_rows); + errdefer mod.unload(); + const r: RouteRows = .{ .mod = mod, .logits_fn = try mod.function("tf_route_logits"), .few_fn = try mod.function("tf_route_logits_few"), .topk_fn = try mod.function("tf_route_topk") }; + try r.few_fn.allowDynamicShared(few_shared); + return r; + } + + pub fn unload(g: *RouteRows) void { + g.mod.unload(); + } + + /// out [rows, e] fp32 = x [rows, d] (bf16, x_ld apart) times w [e, d] bf16 transposed; d % 16 == 0, 16-byte rows. + pub fn logits(g: *const RouteRows, s: Stream, x: u64, x_ld: usize, w: u64, out: u64, rows: usize, d: usize, e: usize) !void { + if (d % 16 != 0 or x_ld % 8 != 0 or x % 16 != 0 or w % 16 != 0) return error.RowShape; + if (rows == 0 or e == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i32, @intCast(x_ld))); + a.add(w); + a.add(out); + for ([_]usize{ rows, d, e }) |v| a.add(@as(i32, @intCast(v))); + // up to 16 rows, deeper staging (the same sums: the router rows' reads set the time) + const few = rows <= 16; + try launch_.launch(if (few) g.few_fn else g.logits_fn, .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((e + 31) / 32), .z = 1 }, .block = .{ .x = 128 }, .shared = if (few) few_shared else many_shared }, s, &a); + } + + /// Each row's top_k of its first `experts` logits into picks and weights [rows, slots]; `shared`: slot top_k too. + pub fn topk(g: *const RouteRows, s: Stream, l: u64, ld: usize, experts: usize, top_k: usize, picks: u64, wts: u64, slots: usize, rows: usize, shared: bool) !void { + const need = experts + @intFromBool(shared); + if (experts > most_experts or top_k == 0 or top_k > most_picks or top_k > experts or slots < top_k + @intFromBool(shared) or ld < need) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(l); + for ([_]usize{ ld, experts, top_k }) |v| a.add(@as(i32, @intCast(v))); + a.add(picks); + a.add(wts); + a.add(@as(i32, @intCast(slots))); + a.add(@as(i32, @intFromBool(shared))); + try launch_.launch(g.topk_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 128 } }, s, &a); + } +}; + +test "the top-k scratch holds a 512-expert router's ten picks" { + try std.testing.expect(most_experts >= 513 and most_picks >= 11); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index a95f29cef7..d836d45502 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -30,6 +30,7 @@ const norm_rows_tests = @import("norm_rows_tests.zig"); const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); +const route_rows_tests = @import("route_rows_tests.zig"); const usage = \\usage: tf-cuda-test @@ -52,6 +53,7 @@ const usage = \\ glue Nemotron's glue kernels against their host references at its shapes (synthetic rows) \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches + \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) \\ window-profile MODEL WIDTHS IDS_FILE a decode window's GPU ms by kernel class (WIDTHS like 1,4,16), beside its graph \\ tree-accept MODEL IDS_FILE COUNT OUT a greedy reply's MTP chain and level-1 runner-up branches at every position (JSON) \\ chunk-costs MODEL IDS_FILE WIDTHS r rows as the decode window graph (r <= 16) and as a prompt chunk (any r) @@ -134,6 +136,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "norm-rows")) return norm_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "pocket")) return pocket_tests.run(gpu, if (rest.len > 0) try std.fmt.parseInt(usize, rest[0], 10) else 0); if (std.mem.eql(u8, cmd, "window-profile")) return window_profile.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); diff --git a/zig/tests/cuda/route_rows_tests.zig b/zig/tests/cuda/route_rows_tests.zig new file mode 100644 index 0000000000..5cb1795c86 --- /dev/null +++ b/zig/tests/cuda/route_rows_tests.zig @@ -0,0 +1,145 @@ +//! route_rows.cu's router logits and softmax top-k rows against float64, bits equal at every row count. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +const rref = core.route_ref; +const aref = core.act_ref; +const counts = [_]usize{ 1, 3, 16, 17, 40 }; +const most = counts[counts.len - 1]; + +fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + worst: f64 = 0, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + fn out(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0xff, r.s.handle); + return b.ptr; + } + + fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try r.a.alloc(T, n); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(o)); + return o; + } + + fn bfs(r: *Rig, n: usize, scale: f32) []u16 { + const v = r.a.alloc(u16, n) catch unreachable; + for (v) |*x| x.* = bf16((r.rng.float(f32) * 2 - 1) * scale); + return v; + } + + fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { + const q = @abs(got - want) / tol; + if (!(q <= 1)) try check.expect(false, "{s}: within float64 tolerance (got {d}, want {d}, {d:.3} tolerances)", .{ what, got, want, q }); + r.worst = @max(r.worst, q); + } +}; + +/// Router logits at every row count (the few-row and the many-row launches): fp64 dots, then the same bits a row. +fn logits(r: *Rig, g: *const cuda.route_rows.RouteRows, d: usize, e: usize) !void { + const ld = d + 64; // rows with something after them + const x = r.bfs(most * ld, 1); + const w = r.bfs(e * d, 0.05); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dw = try r.dev(std.mem.sliceAsBytes(w)); + var outs: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + outs[c] = try r.out(m * e * 4); + try g.logits(r.s, dx, ld, dw, outs[c], m, d, e); + } + const got = try r.back(f32, outs[counts.len - 1], most * e); + const xv = try r.a.alloc(f64, d); + const wv = try r.a.alloc(f64, d); + for (0..most) |row| { + for (xv, 0..) |*o, i| o.* = bf(x[row * ld + i]); + for (0..e) |ex| { + for (wv, 0..) |*o, i| o.* = bf(w[ex * d + i]); + try r.within("route logits", got[row * e + ex], rref.dot(xv, wv), rref.dotTolerance(xv, wv)); + } + } + for (counts[0 .. counts.len - 1], 0..) |m, c| try check.sameBytes("route logits rows", std.mem.sliceAsBytes(try r.back(f32, outs[c], m * e)), std.mem.sliceAsBytes(got[0 .. m * e])); +} + +/// Top-k of each row (a row of ties too): the picks the fp64 rule makes, weights and the shared gate within bf16. +fn topk(r: *Rig, g: *const cuda.route_rows.RouteRows, experts: usize, k: usize, shared: bool) !void { + const ld = experts + 1; + const slots = k + @intFromBool(shared); + const l = try r.a.alloc(f32, most * ld); + for (l) |*v| v.* = (r.rng.float(f32) * 2 - 1) * 4; + for (0..experts) |i| l[i] = @floatFromInt(i % 3); // row 0: every pick a tie broken by the lower id + const dl = try r.dev(std.mem.sliceAsBytes(l)); + var picks: [counts.len]u64 = undefined; + var wts: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + picks[c] = try r.out(m * slots * 4); + wts[c] = try r.out(m * slots * 4); + try g.topk(r.s, dl, ld, experts, k, picks[c], wts[c], slots, m, shared); + } + const p = try r.back(i32, picks[counts.len - 1], most * slots); + const wt = try r.back(f32, wts[counts.len - 1], most * slots); + const row = try r.a.alloc(f64, experts); + const ids = try r.a.alloc(usize, k); + const want = try r.a.alloc(f64, k); + for (0..most) |i| { + for (row, 0..) |*o, j| o.* = l[i * ld + j]; + rref.topk(row, ids, want); + for (0..k) |j| { + try check.expect(p[i * slots + j] == @as(i32, @intCast(ids[j])), "route top-k row {d} pick {d}: expert {d}, the rule picks {d}", .{ i, j, p[i * slots + j], ids[j] }); + try r.within("route top-k weight", wt[i * slots + j], want[j], aref.tolerance(want[j])); + } + if (shared) { + try check.expect(p[i * slots + k] == @as(i32, @intCast(experts)), "route top-k row {d}: the shared slot names expert {d}", .{ i, p[i * slots + k] }); + const gate = aref.sigmoid(bf(bf16(l[i * ld + experts]))); + try r.within("route shared gate", wt[i * slots + k], gate, aref.tolerance(gate)); + } + } + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes("route top-k picks", std.mem.sliceAsBytes(try r.back(i32, picks[c], m * slots)), std.mem.sliceAsBytes(p[0 .. m * slots])); + try check.sameBytes("route top-k weights", std.mem.sliceAsBytes(try r.back(f32, wts[c], m * slots)), std.mem.sliceAsBytes(wt[0 .. m * slots])); + } +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(41); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = try cuda.Stream.init(gpu.d, false), .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + r.s.deinit(); + } + var g = try cuda.route_rows.RouteRows.load(gpu.d); + defer g.unload(); + try logits(&r, &g, 2560, 513); // Flash Next: 512 experts and the shared expert's gate row + try logits(&r, &g, 2688, 128); + try logits(&r, &g, 1040, 37); // a K tail past the staged 256s, and experts past a block + try topk(&r, &g, 512, 10, true); + try topk(&r, &g, 128, 8, false); + check.pass("route rows: logits and top-k within float64 tolerance (worst {d:.3}), the fp64 rule's picks, every row count's bits equal", .{r.worst}); +} From 794d273713f7baf4f25dc68a8e4980733cd12778 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:14:38 -0500 Subject: [PATCH 043/115] Core: token rows in groups of 32 with a copy a stream, and RMS rows with their group sums embed_rows.cu's affine 4-bit rows take the scale group (32 or 64) and a copy count (EmbedRows.affine4Groups; affine4 keeps groups of 64 and one copy, so Nemotron's launches are unchanged). norms.cu gains tf_rms_rows_sums (Norms.rmsSums): bf16 rows read x_ld apart, y = bf16(x * rsqrt(mean(x^2) + eps) * w) with fp32 weights and IEEE division and square root, then each group's sum of the stored y, for a group-quantized matmul's input. tf-cuda-test norm-rows checks both against float64 at every row count (Flash Next's 2,560-wide embedding in 4 copies, its MTP head's 2,560- and 10,240-wide norms) with the bits of every row count equal. --- zig/kernels/cuda/embed_rows.cu | 9 ++-- zig/kernels/cuda/norms.cu | 21 +++++++++ zig/src/cuda/embed_rows.zig | 11 +++-- zig/src/cuda/norms.zig | 16 +++++++ zig/tests/cuda/norm_rows_tests.zig | 71 +++++++++++++++++++++++++----- 5 files changed, 109 insertions(+), 19 deletions(-) diff --git a/zig/kernels/cuda/embed_rows.cu b/zig/kernels/cuda/embed_rows.cu index 2ed32c9665..ca5f2602d2 100644 --- a/zig/kernels/cuda/embed_rows.cu +++ b/zig/kernels/cuda/embed_rows.cu @@ -11,18 +11,19 @@ constexpr int NT = 256; using namespace tf_embed; -// out[r, i] = bf16(q * s + b) of token ids[r]: q 4 bits (8 a word, low first), s and b bf16 a group of 64; d % 64 == 0. +// out[r, k, i] = bf16(q * s + b) of token ids[r], k < copies: q 4 bits (low first), s and b a group's (32 or 64). extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_affine4(const int* __restrict__ ids, const uint32_t* __restrict__ w, const __nv_bfloat16* __restrict__ s, const __nv_bfloat16* __restrict__ b, - __nv_bfloat16* __restrict__ out, int d) { + __nv_bfloat16* __restrict__ out, int d, int group, + int copies) { const int64_t r = blockIdx.x; const int c = blockIdx.y * NT + threadIdx.x; // 8 values, one word if (8 * c >= d) return; const int64_t tok = ids[r]; const uint32_t word = w[tok * (d / 8) + c]; - const int64_t g = tok * (d / 64) + c / 8; + const int64_t g = tok * (d / group) + 8 * c / group; const float sc = __bfloat162float(s[g]), bi = __bfloat162float(b[g]); uint4 o; __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&o); @@ -30,7 +31,7 @@ extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_affine4(const int for (int i = 0; i < 4; ++i) h[i] = __floats2bfloat162_rn(__fmaf_rn(static_cast((word >> (8 * i)) & 0xFu), sc, bi), __fmaf_rn(static_cast((word >> (8 * i + 4)) & 0xFu), sc, bi)); - *reinterpret_cast(out + r * d + 8 * c) = o; + for (int k = 0; k < copies; ++k) *reinterpret_cast(out + (r * copies + k) * d + 8 * c) = o; } // out[r] = w[ids[r]] for a bf16 table, 16 bytes a thread (d % 8 == 0). diff --git a/zig/kernels/cuda/norms.cu b/zig/kernels/cuda/norms.cu index c0797a9ff3..c0089edf04 100644 --- a/zig/kernels/cuda/norms.cu +++ b/zig/kernels/cuda/norms.cu @@ -175,6 +175,27 @@ extern "C" __global__ void __launch_bounds__(NT) tf_rms_rows(const void* x, int #undef TF_RMS } +// bf16 rows `x_ld` apart: y = bf16(x * rsqrt(mean(x^2) + eps) * w), then xs = the stored y's sums over each `group`. +extern "C" __global__ void __launch_bounds__(NT) tf_rms_rows_sums(const __nv_bfloat16* __restrict__ x, long long x_ld, + const float* __restrict__ w, __nv_bfloat16* __restrict__ y, + float* __restrict__ xs, int width, int group, float eps) { + const int64_t at = static_cast(blockIdx.x) * x_ld, out = static_cast(blockIdx.x) * width; + float part = 0.f; + for (int i = threadIdx.x; i < width; i += NT) { + const float a = __bfloat162float(x[at + i]); + part = __fmaf_rn(a, a, part); + } + const float inv = inv_rms(part, width, eps); + for (int i = threadIdx.x; i < width; i += NT) + y[out + i] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(__bfloat162float(x[at + i]), inv), w[i])); + __syncthreads(); + for (int j = threadIdx.x; j < width / group; j += NT) { + float sum = 0.f; + for (int k = 0; k < group; ++k) sum = __fadd_rn(sum, __bfloat162float(y[out + j * group + k])); + xs[static_cast(blockIdx.x) * (width / group) + j] = sum; + } +} + // res (fp32, in place) += rms(x) * w1 (a sandwich norm; null: res += x), then y = rms(res) * w2 in y's kind. extern "C" __global__ void __launch_bounds__(NT) tf_add_rms_rows(const void* x, int x_kind, float* res, const float* w1, const float* w2, void* y, int y_kind, int width, diff --git a/zig/src/cuda/embed_rows.zig b/zig/src/cuda/embed_rows.zig index a8d4a67b4c..973b658f7c 100644 --- a/zig/src/cuda/embed_rows.zig +++ b/zig/src/cuda/embed_rows.zig @@ -1,4 +1,4 @@ -//! embed_rows.cu's token rows for every CUDA family: 4-bit affine (MLX layout, groups of 64) or bf16 tables. +//! embed_rows.cu's token rows for every CUDA family: 4-bit affine (MLX layout, groups of 32 or 64) or bf16 tables. const std = @import("std"); const kernels = @import("kernels.zig"); const module = @import("module.zig"); @@ -30,11 +30,16 @@ pub const EmbedRows = struct { /// out [rows, d] bf16 from ids (i32 [rows]): w u32 [vocab, d / 8], scales and biases bf16 [vocab, d / 64]. pub fn affine4(g: *const EmbedRows, s: Stream, ids: u64, w: u64, scales: u64, biases: u64, out: u64, rows: usize, d: usize) !void { - if (d % 64 != 0 or out % 16 != 0) return error.RowShape; + return g.affine4Groups(s, ids, w, scales, biases, out, rows, d, 64, 1); + } + + /// out [rows, copies, d] bf16, each row `copies` times: scales and biases a group of `group` (32 or 64) values. + pub fn affine4Groups(g: *const EmbedRows, s: Stream, ids: u64, w: u64, scales: u64, biases: u64, out: u64, rows: usize, d: usize, group: usize, copies: usize) !void { + if ((group != 32 and group != 64) or d % group != 0 or out % 16 != 0 or copies == 0) return error.RowShape; if (rows == 0) return; var a: launch_.Args = .{}; for ([_]u64{ ids, w, scales, biases, out }) |v| a.add(v); - a.add(@as(i32, @intCast(d))); + for ([_]usize{ d, group, copies }) |v| a.add(@as(i32, @intCast(v))); try launch_.launch(g.affine4_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); } diff --git a/zig/src/cuda/norms.zig b/zig/src/cuda/norms.zig index fd343b4635..6f774d50d3 100644 --- a/zig/src/cuda/norms.zig +++ b/zig/src/cuda/norms.zig @@ -42,6 +42,7 @@ pub const Norms = struct { add_rms_h_fn: Function, add_rms_h_f32w_fn: Function, gated_fns: [3]Function, // head dim 64, 128, 256 + rms_sums_fn: Function, pub fn load(d: *const Driver) !Norms { if (!kernels.available) return error.BuiltWithoutKernels; @@ -56,6 +57,7 @@ pub const Norms = struct { .add_rms_h_fn = try mod.function("tf_add_rms_rows_bf16h"), .add_rms_h_f32w_fn = try mod.function("tf_add_rms_rows_bf16h_f32w"), .gated_fns = .{ try mod.function("tf_gated_rms_rows_64"), try mod.function("tf_gated_rms_rows_128"), try mod.function("tf_gated_rms_rows_256") }, + .rms_sums_fn = try mod.function("tf_rms_rows_sums"), }; } @@ -87,6 +89,20 @@ pub const Norms = struct { try launch_.launch(n.rms_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); } + /// bf16 rows `x_ld` apart: y = bf16(rms(x) * w) [rows, width], fp32 w, and xs the stored y's sums of each `group`. + pub fn rmsSums(n: *const Norms, s: Stream, x: u64, x_ld: usize, w: u64, y: u64, xs: u64, rows: usize, width: usize, group: usize, eps: f32) !void { + if (group == 0 or width % group != 0 or x_ld < width) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(@as(i64, @intCast(x_ld))); + for ([_]u64{ w, y, xs }) |v| a.add(v); + a.add(@as(i32, @intCast(width))); + a.add(@as(i32, @intCast(group))); + a.add(eps); + try launch_.launch(n.rms_sums_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = threads } }, s, &a); + } + /// The residual step: res (fp32, in place) += rms(x) * w1 (w1 null: res += x), then y = rms(res) * w2. pub fn addRms(n: *const Norms, s: Stream, x: u64, x_kind: Kind, res: u64, w1: ?u64, w2: u64, y: u64, y_kind: Kind, rows: usize, width: usize, eps: f32) !void { if (rows == 0) return; diff --git a/zig/tests/cuda/norm_rows_tests.zig b/zig/tests/cuda/norm_rows_tests.zig index 8f1d295ec1..ec95d7fe36 100644 --- a/zig/tests/cuda/norm_rows_tests.zig +++ b/zig/tests/cuda/norm_rows_tests.zig @@ -171,35 +171,78 @@ fn gatedRms(r: *Rig, n: *const cuda.norms.Norms, heads: usize, dv: usize) !void } } -fn embed(r: *Rig, d: usize) !void { +fn embed(r: *Rig, d: usize, group: usize, copies: usize) !void { var e = try cuda.embed_rows.EmbedRows.load(r.gpu.d); defer e.unload(); const vocab = 300; const w = try r.a.alloc(u32, vocab * d / 8); for (w) |*x| x.* = r.rng.int(u32); - const sc = r.bfs(vocab * d / 64, 0.1); - const bi = r.bfs(vocab * d / 64, 0.5); + const sc = r.bfs(vocab * d / group, 0.1); + const bi = r.bfs(vocab * d / group, 0.5); const tab = r.bfs(vocab * d, 1); const ids = try r.a.alloc(i32, most); for (ids) |*x| x.* = @intCast(r.rng.uintLessThan(usize, vocab)); const did = try r.dev(std.mem.sliceAsBytes(ids)); - const o4 = try r.out(most * d * 2); + const o4 = try r.out(most * copies * d * 2); const o16 = try r.out(most * d * 2); - try e.affine4(r.s, did, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(sc)), try r.dev(std.mem.sliceAsBytes(bi)), o4, most, d); + const dw = try r.dev(std.mem.sliceAsBytes(w)); + const ds = try r.dev(std.mem.sliceAsBytes(sc)); + const db = try r.dev(std.mem.sliceAsBytes(bi)); + if (group == 64 and copies == 1) try e.affine4(r.s, did, dw, ds, db, o4, most, d) else try e.affine4Groups(r.s, did, dw, ds, db, o4, most, d, group, copies); try e.bf16(r.s, did, try r.dev(std.mem.sliceAsBytes(tab)), o16, most, d); - const g4 = try r.back(u16, o4, most * d); + const g4 = try r.back(u16, o4, most * copies * d); const g16 = try r.back(u16, o16, most * d); for (0..most) |row| { const t: usize = @intCast(ids[row]); for (0..d) |i| { const q: u4 = @truncate(w[t * d / 8 + i / 8] >> @intCast(4 * (i % 8))); - const want = aref.affine4(q, bf(sc[t * d / 64 + i / 64]), bf(bi[t * d / 64 + i / 64])); - try r.within("embed affine4", bf(g4[row * d + i]), want, nref.tolerance(want, .bf16)); + const want = aref.affine4(q, bf(sc[t * d / group + i / group]), bf(bi[t * d / group + i / group])); + try r.within("embed affine4", bf(g4[row * copies * d + i]), want, nref.tolerance(want, .bf16)); } + for (1..copies) |c| try check.sameBytes("embed affine4 copies", std.mem.sliceAsBytes(g4[(row * copies + c) * d ..][0..d]), std.mem.sliceAsBytes(g4[row * copies * d ..][0..d])); try check.sameBytes("embed bf16", std.mem.sliceAsBytes(g16[row * d ..][0..d]), std.mem.sliceAsBytes(tab[t * d ..][0..d])); } } +/// RMS rows `x_ld` apart with fp32 weights, and the stored outputs' sums of each `group`. +fn rmsSums(r: *Rig, n: *const cuda.norms.Norms, width: usize, x_ld: usize, group: usize) !void { + const x = r.bfs(most * x_ld, 1); + const wf = r.weights(width); + const wvals = try r.a.alloc(f64, width); + for (wvals, wf) |*o, v| o.* = v; + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dw = try r.dev(std.mem.sliceAsBytes(wf)); + var ys: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + ys[c] = try r.out(m * width * 2); + xss[c] = try r.out(m * width / group * 4); + try n.rmsSums(r.s, dx, x_ld, dw, ys[c], xss[c], m, width, group, eps); + } + const y = try r.back(u16, ys[counts.len - 1], most * width); + const xs = try r.back(f32, xss[counts.len - 1], most * width / group); + const xv = try r.a.alloc(f64, width); + const want = try r.a.alloc(f64, width); + for (0..most) |row| { + for (xv, 0..) |*v, i| v.* = bf(x[row * x_ld + i]); + nref.rms(xv, wvals, eps, want); + for (want, 0..) |w, i| try r.within("rms sums", bf(y[row * width + i]), w, nref.tolerance(w, .bf16)); + for (0..width / group) |j| { + var t: f64 = 0; + var abs: f64 = 0; + for (y[row * width + j * group ..][0..group]) |v| { + t += bf(v); + abs += @abs(bf(v)); + } + try r.within("rms sums group sums", xs[row * (width / group) + j], t, abs * @as(f64, @floatFromInt(group)) * std.math.pow(f64, 2, -24) + 1e-30); + } + } + for (counts[0 .. counts.len - 1], 0..) |m, c| { + try check.sameBytes("rms sums", std.mem.sliceAsBytes(try r.back(u16, ys[c], m * width)), std.mem.sliceAsBytes(y[0 .. m * width])); + try check.sameBytes("rms sums group sums", std.mem.sliceAsBytes(try r.back(f32, xss[c], m * width / group)), std.mem.sliceAsBytes(xs[0 .. m * width / group])); + } +} + pub fn run(gpu: check.Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); @@ -220,8 +263,12 @@ pub fn run(gpu: check.Gpu) !void { try gatedRms(&r, &n, 48, 128); // the 27B's GDN output try gatedRms(&r, &n, 5, 64); try gatedRms(&r, &n, 3, 256); - try embed(&r, 5120); // the 27B's hidden - try embed(&r, 2688); - try embed(&r, 192); - check.pass("norm rows: add_rms_h, gated rms, embed rows within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); + try embed(&r, 5120, 64, 1); // the 27B's hidden + try embed(&r, 2688, 64, 1); + try embed(&r, 192, 64, 1); + try embed(&r, 2560, 32, 4); // Flash Next: group 32, a copy a hyper-connection stream + try embed(&r, 96, 32, 1); + try rmsSums(&r, &n, 2560, 10240, 32); // Flash Next's MTP head: the embedding, its first stream of four + try rmsSums(&r, &n, 10240, 10240, 32); // and the four streams normed as one row + check.pass("norm rows: add_rms_h, gated rms, rms with group sums, embed rows within float64 tolerance (worst {d:.3}), group sums of the stored outputs, every row count's bits equal", .{r.worst}); } From ea7d098a130655e6e7d57548d6b627bd4f7c1680 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:14:38 -0500 Subject: [PATCH 044/115] Flash Next CUDA: its own glue kernels, flashnext_hc.cu and flashnext_attn.cu, against float64 Our kernels for every entry point the batched pass took from a captured Triton set, with the Python engine's formulas and bf16 rounding points, IEEE division and square root and the core's within-one-ulp exp (glue_math.cuh), built without contraction: - flashnext_hc.cu: the hyper-connection write-back (stream sums alone, the MoE slots, K slices), the down projection with its normed rows (the tiled 4-bit dots on mma tiles, groups in order, per K slice), the slices' sum with the bf16 SiLU and inject gates, the up projection with the stream mix, the n-gram gate and dilated conv, and the MTP head's stream adds; - flashnext_attn.cu: the q, k and indexer heads' norm and RoPE with the cache writes, the output gate, the indexer's pooled blocks and block scores, and attention over 512-key chunks (q . k and p . v on mma tiles, an online softmax with fp32 sums, dense or through a sparse row's list) with the chunks merged in order. Every kernel keeps a row's own sums in a fixed order, so a row's bits are the same at any row count; attention has a form that reads the next tile ahead for a decode round's few rows, with the same sums. cuda_glue.zig launches them (Glue), with the core's token rows, RMS rows and route rows for the embedding, the MTP head's norms and the router. tf-cuda-test flashnext-hc and flashnext-attn check each kernel at the model's shapes against a float64 chain that rounds where the kernel rounds: each bf16 output within one bf16 ulp (an intermediate rounding can tip either way), fp32 sums within their sum-of-|terms| bound, attention within 2^-8 of sum p |v| / l plus half an ulp (p enters p . v as bf16), copies and the stream adds exact, and the bits of 1, 3 and 17 rows (attention: 1, 3 and 20, both forms) equal. --- zig/build/cuda.zig | 3 + zig/kernels/cuda/flashnext_attn.cu | 370 +++++++++++++++++++ zig/kernels/cuda/flashnext_hc.cu | 394 +++++++++++++++++++++ zig/src/cuda/kernels.zig | 2 + zig/src/families/flashnext/cuda_glue.zig | 270 ++++++++++++++ zig/src/families/flashnext/cuda_launch.zig | 2 + zig/tests/cuda/flashnext_attn_tests.zig | 282 +++++++++++++++ zig/tests/cuda/flashnext_hc_tests.zig | 339 ++++++++++++++++++ zig/tests/cuda/flashnext_rig.zig | 119 +++++++ zig/tests/cuda/main.zig | 6 + 10 files changed, 1787 insertions(+) create mode 100644 zig/kernels/cuda/flashnext_attn.cu create mode 100644 zig/kernels/cuda/flashnext_hc.cu create mode 100644 zig/src/families/flashnext/cuda_glue.zig create mode 100644 zig/tests/cuda/flashnext_attn_tests.zig create mode 100644 zig/tests/cuda/flashnext_hc_tests.zig create mode 100644 zig/tests/cuda/flashnext_rig.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 6a200ad469..6aca30d30c 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -45,6 +45,8 @@ const kernels = [_]Kernel{ .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig) .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig) .{ .name = "route_rows", .flags = glue }, // ours: router logits and softmax top-k rows (cuda/route_rows.zig) + .{ .name = "flashnext_hc", .flags = glue }, // ours: Flash Next's hyper-connection and n-gram glue + .{ .name = "flashnext_attn", .flags = glue }, // ours: Flash Next's attention glue and its indexer's pooling and scores .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts @@ -141,6 +143,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("core", mods.core); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); + runner.addImport("flashnext", mods.flashnext); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); runner.addImport("exlr8", exlr8); const exlr8_tests = exlr8_build.testModules(b, target, optimize, exlr8); diff --git a/zig/kernels/cuda/flashnext_attn.cu b/zig/kernels/cuda/flashnext_attn.cu new file mode 100644 index 0000000000..71aaad7026 --- /dev/null +++ b/zig/kernels/cuda/flashnext_attn.cu @@ -0,0 +1,370 @@ +// Flash Next's attention glue and its sparse indexer's pooling and scores: each row's own fixed-order sums. + +// glue.py and attention.py's arithmetic and bf16 rounding points; the attention dots on mma tiles, K ascending. + +#include +#include +#include + +#include "glue_math.cuh" +#include "qmm_frag.cuh" + +using namespace tf_glue; + +namespace fn_attn { + +constexpr int TK = 32; // keys an attention tile holds +constexpr int HD = 256; // head width the attention tiles are built for +constexpr int QR = 16; // query heads a KV head's tile holds (G <= 16, the rest zero) + +// A block's fixed-order sum: each warp's butterfly, then the warps' sums in warp order. +__device__ __forceinline__ float block_sum(float v, float* scratch) { + v = warp_sum(v); + const int warps = blockDim.x >> 5; + if ((threadIdx.x & 31) == 0) scratch[threadIdx.x >> 5] = v; + __syncthreads(); + float total = 0.0f; + for (int w = 0; w < warps; ++w) total = __fadd_rn(total, scratch[w]); + __syncthreads(); + return total; +} + +__device__ __forceinline__ float inv_rms(float ss, int width, float eps) { + return __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(ss, static_cast(width)), eps))); +} + +// Rotate-half RoPE over the first 2 * half features at `pos`: cos and sin of the fp32 angle (full range reduction). +__device__ __forceinline__ float rope(float xn, float xpn, int d, int half, int pos, const float* __restrict__ inv) { + if (d >= 2 * half) return xn; + const int i = d < half ? d : d - half; + float s, c; + sincosf(__fmul_rn(static_cast(pos), inv[i]), &s, &c); + return d < half ? __fsub_rn(__fmul_rn(xn, c), __fmul_rn(xpn, s)) : __fadd_rn(__fmul_rn(xpn, s), __fmul_rn(xn, c)); +} + +// ldmatrix .x4 .trans: four 8x8 bf16 tiles of shared memory, each thread's pair down a column (the B layout of rows). +__device__ __forceinline__ void ldmatrix4t(uint32_t (&r)[4], const void* p) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(qmm_frag::smem(p))); +} + +__device__ __forceinline__ int partner(int d, int half) { + return d < half ? d + half : (d < 2 * half ? d - half : d); +} + +} // namespace fn_attn + +using namespace fn_attn; + +// glue._attn_prep (text positions, bf16 cache), program (row, head): heads normed, rounded, roped, rounded, stored. +extern "C" __global__ void __launch_bounds__(256) fn_attn_prep(const __nv_bfloat16* __restrict__ P, const int* __restrict__ POS0, + const float* __restrict__ QW, const float* __restrict__ KW, + const float* __restrict__ IW, const float* __restrict__ INV, + __nv_bfloat16* __restrict__ Q, __nv_bfloat16* __restrict__ KC, + __nv_bfloat16* __restrict__ VC, __nv_bfloat16* __restrict__ IQ, + __nv_bfloat16* __restrict__ IKC, float eps, int PW, int NQ, int NKV, + int NI, int IHD, int HALF) { + __shared__ float scratch[32]; + const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; + const int pos = POS0[0] + r; + const long long row = static_cast(r) * PW; + if (head >= NQ + NKV + NI) { + if (d < IHD) IKC[static_cast(pos) * IHD + d] = P[row + NQ * 2 * HD + 2 * NKV * HD + NI * IHD + d]; + return; + } + const bool index = head >= NQ + NKV; + const int width = index ? IHD : HD; + const bool live = d < width; + const long long src = row + (head < NQ ? head * 2 * HD : (head < NQ + NKV ? NQ * 2 * HD + (head - NQ) * HD : NQ * 2 * HD + 2 * NKV * HD + (head - NQ - NKV) * IHD)); + const float* w = head < NQ ? QW : (head < NQ + NKV ? KW : IW); + const float x = live ? bf(P[src + d]) : 0.0f; + const float rinv = inv_rms(block_sum(__fmul_rn(x, x), scratch), width, eps); + if (!live) return; + const int p = partner(d, HALF); + const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), w[d])); + const float xpn = rbf(__fmul_rn(__fmul_rn(bf(P[src + p]), rinv), w[p])); + const __nv_bfloat16 out = __float2bfloat16_rn(rope(xn, xpn, d, HALF, pos, INV)); + if (head < NQ) { + Q[(static_cast(r) * NQ + head) * HD + d] = out; + } else if (!index) { + const long long slot = static_cast(pos) * NKV + head - NQ; + KC[slot * HD + d] = out; + VC[slot * HD + d] = P[row + NQ * 2 * HD + NKV * HD + (head - NQ) * HD + d]; + } else { + IQ[(static_cast(r) * NI + head - NQ - NKV) * IHD + d] = out; + } +} + +// glue._attn_gate, program (row, head) of HD threads: bf16(o / (1 + e^-gate)) and its 32-group sums. +extern "C" __global__ void __launch_bounds__(256) fn_attn_gate(const __nv_bfloat16* __restrict__ O, const __nv_bfloat16* __restrict__ P, + __nv_bfloat16* __restrict__ OUT, float* __restrict__ XS, int PW, int NQ) { + const int r = blockIdx.x, h = blockIdx.y, d = threadIdx.x; + const float o = bf(O[(static_cast(r) * NQ + h) * HD + d]); + const float g = bf(P[static_cast(r) * PW + h * 2 * HD + HD + d]); + const float out = rbf(__fdiv_rn(o, __fadd_rn(1.0f, tf_exp(-g)))); + OUT[static_cast(r) * NQ * HD + h * HD + d] = __float2bfloat16_rn(out); + const float sum = warp_sum(out); + if ((d & 31) == 0) XS[static_cast(r) * (NQ * HD / 32) + h * (HD / 32) + d / 32] = sum; +} + +// attention._pool, a block a completed block of RATIO indexer keys: their mean, rounded, normed and roped at its start. +extern "C" __global__ void __launch_bounds__(128) fn_qsa_pool(const __nv_bfloat16* __restrict__ IKC, __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, const float* __restrict__ W, + const float* __restrict__ INV, float eps, int R, int DI, int HALF, + int RATIO) { + __shared__ float scratch[32]; + __shared__ float xs[256]; + const int p0 = POS0[0]; + const int b = p0 / RATIO + blockIdx.x, d = threadIdx.x; + if (RATIO * b + RATIO > p0 + R) return; + float acc = bf(IKC[static_cast(RATIO) * b * DI + d]); + for (int k = 1; k < RATIO; ++k) acc = __fadd_rn(acc, bf(IKC[(static_cast(RATIO) * b + k) * DI + d])); + const float x = rbf(__fdiv_rn(acc, static_cast(RATIO))); + xs[d] = x; + const float rinv = inv_rms(block_sum(__fmul_rn(x, x), scratch), DI, eps); + const int p = partner(d, HALF); + const float xn = rbf(__fmul_rn(__fmul_rn(x, rinv), W[d])); + const float xpn = rbf(__fmul_rn(__fmul_rn(xs[p], rinv), W[p])); + POOLED[static_cast(b) * DI + d] = __float2bfloat16_rn(rope(xn, xpn, d, HALF, RATIO * b, INV)); +} + +// attention._scores, program (row, 64 blocks): the heads' sum of relu(q . k) / sqrt(DI), rows past TOP (DI 128). +extern "C" __global__ void __launch_bounds__(256) fn_qsa_scores(const __nv_bfloat16* __restrict__ IQ, const __nv_bfloat16* __restrict__ POOLED, + const int* __restrict__ POS0, float* __restrict__ SC, int NB, int HI, + int RATIO, int TOP) { + constexpr int DI = 128; + __shared__ float qs[8][DI]; + const int r = blockIdx.x, j = blockIdx.y; + const int complete = (POS0[0] + r + 1) / RATIO; + if (complete <= TOP || j * 64 >= complete) return; + for (int e = threadIdx.x; e < HI * DI; e += blockDim.x) qs[e / DI][e % DI] = bf(IQ[static_cast(r) * HI * DI + e]); + __syncthreads(); + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const float scale = __fsqrt_rn(static_cast(DI)); + uint2 raw[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int b = j * 64 + warp * 8 + i; + raw[i] = b < complete ? *reinterpret_cast(POOLED + static_cast(b) * DI + 4 * lane) : make_uint2(0, 0); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int b = j * 64 + warp * 8 + i; + if (b >= complete) break; + const __nv_bfloat16* kb = reinterpret_cast(&raw[i]); + float total = 0.0f; + for (int h = 0; h < HI; ++h) { + float part = 0.0f; +#pragma unroll + for (int e = 0; e < 4; ++e) part = __fmaf_rn(bf(kb[e]), qs[h][4 * lane + e], part); + total = __fadd_rn(total, fmaxf(warp_sum(part), 0.0f)); + } + if (lane == 0) SC[static_cast(r) * NB + b] = __fdiv_rn(total, scale); + } +} + +// attention._chunks, program (row, KV head, chunk): the chunk's keys TK at a time, online softmax, partial o, m, l. +template +__device__ __forceinline__ void attn_chunks(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, const int* __restrict__ SPR, + int H, int HK, int G, int CH, int NCH, float SCALE, int IDW, int QSA) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + typedef __nv_bfloat16 Row[HD + 8]; + Row* qs = reinterpret_cast(smem_raw); // [QR] the row's query heads + Row* tiles_kv = qs + QR; // [NB][2][TK] each buffer's keys, then values + __nv_bfloat16 (*ps)[TK + 8] = reinterpret_cast<__nv_bfloat16 (*)[TK + 8]>(tiles_kv + NB * 2 * TK); // [QR] bf16 p + __shared__ float red_m[4][QR], red_l[4][QR]; + __shared__ int key_ids[NB][TK]; + const int r = blockIdx.x, hk = blockIdx.y, c = blockIdx.z; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + int n = POS0[0] + r + 1; + bool sparse = false; + if (QSA) { + sparse = SPR[r] != 0; + if (sparse) n = NKR[r]; + } + const int start = c * CH; + if (start >= n) return; + for (int e = threadIdx.x; e < QR * HD / 8; e += blockDim.x) { + const int i = e / (HD / 8), v = e % (HD / 8); + qmm_frag::cp16z(&qs[i][8 * v], Q + (static_cast(r) * H + hk * G + min(i, G - 1)) * HD + 8 * v, i < G); + } + // a tile's key ids, then its keys and values staged as two cp.async groups + auto load = [&](int t0, int buf) { + if (threadIdx.x < TK) { + const int ki = start + t0 + threadIdx.x; + key_ids[buf][threadIdx.x] = ki < n ? (sparse ? IDS[static_cast(r) * IDW + ki] : ki) : -1; + } + __syncthreads(); + for (int pass = 0; pass < 2; ++pass) { + const __nv_bfloat16* src = pass == 0 ? KC : VC; + Row* dst = tiles_kv + (buf * 2 + pass) * TK; + for (int e = threadIdx.x; e < TK * HD / 8; e += blockDim.x) { + const int k = e / (HD / 8), v = e % (HD / 8), key = key_ids[buf][k]; + qmm_frag::cp16z(&dst[k][8 * v], src + (static_cast(max(key, 0)) * HK + hk) * HD + 8 * v, key >= 0); + } + qmm_frag::commit(); + } + }; + float m[2] = {-INFINITY, -INFINITY}, l[2] = {0.0f, 0.0f}; + float o[8][4] = {}; + const int tiles = min(n - start, CH); + if (NB == 2) load(0, 0); + for (int t0 = 0, it = 0; t0 < tiles; t0 += TK, ++it) { + const int buf = NB == 2 ? it & 1 : 0; + const bool ahead = NB == 2 && t0 + TK < tiles; // the next tile's loads go out before this tile's math + if (NB == 1) load(t0, 0); + if (ahead) load(t0 + TK, buf ^ 1); + Row* kt = tiles_kv + buf * 2 * TK; + Row* vt = kt + TK; + const int* keys = key_ids[buf]; + if (ahead) qmm_frag::wait<3>(); else qmm_frag::wait<1>(); + __syncthreads(); + constexpr int JT = TK / 32; // a warp's 8-key tiles + float s[JT][4] = {}; +#pragma unroll 4 + for (int ks = 0; ks < HD / 16; ks += 2) { + uint32_t a0[4], a1[4]; + qmm_frag::ldmatrix4(a0, &qs[lane % 16][ks * 16 + (lane / 16) * 8]); + qmm_frag::ldmatrix4(a1, &qs[lane % 16][(ks + 1) * 16 + (lane / 16) * 8]); +#pragma unroll + for (int jt = 0; jt < JT; ++jt) { + uint32_t b[4]; // two k16 steps of the tile's 8 keys + qmm_frag::ldmatrix4(b, &kt[warp * 8 * JT + jt * 8 + lane % 8][ks * 16 + (lane / 8) * 8]); + qmm_frag::mma(s[jt], a0, b[0], b[1]); + qmm_frag::mma(s[jt], a1, b[2], b[3]); + } + } + float tile_m[2] = {-INFINITY, -INFINITY}; +#pragma unroll + for (int jt = 0; jt < JT; ++jt) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const bool valid = keys[warp * 8 * JT + jt * 8 + 2 * t + (e & 1)] >= 0; + s[jt][e] = valid ? __fmul_rn(s[jt][e], SCALE) : -INFINITY; + tile_m[e >> 1] = fmaxf(tile_m[e >> 1], s[jt][e]); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + tile_m[h] = fmaxf(tile_m[h], __shfl_xor_sync(0xffffffffu, tile_m[h], 1)); + tile_m[h] = fmaxf(tile_m[h], __shfl_xor_sync(0xffffffffu, tile_m[h], 2)); + if (t == 0) red_m[warp][g + 8 * h] = tile_m[h]; + } + __syncthreads(); + float next_m[2], alpha[2], rows[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = g + 8 * h; + const float tm = fmaxf(fmaxf(red_m[0][row], red_m[1][row]), fmaxf(red_m[2][row], red_m[3][row])); + const bool active = tm != -INFINITY; + next_m[h] = active ? fmaxf(m[h], tm) : m[h]; + alpha[h] = active ? (m[h] == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m[h], next_m[h]))) : 1.0f; + rows[h] = 0.0f; + } +#pragma unroll + for (int jt = 0; jt < JT; ++jt) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int h = e >> 1; + const float p = s[jt][e] == -INFINITY ? 0.0f : tf_exp(__fsub_rn(s[jt][e], next_m[h])); + rows[h] = __fadd_rn(rows[h], p); + ps[g + 8 * h][warp * 8 * JT + jt * 8 + 2 * t + (e & 1)] = __float2bfloat16_rn(p); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + rows[h] = __fadd_rn(rows[h], __shfl_xor_sync(0xffffffffu, rows[h], 1)); + rows[h] = __fadd_rn(rows[h], __shfl_xor_sync(0xffffffffu, rows[h], 2)); + if (t == 0) red_l[warp][g + 8 * h] = rows[h]; + } + if (ahead) qmm_frag::wait<2>(); else qmm_frag::wait<0>(); + __syncthreads(); +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int row = g + 8 * h; + const float sum = __fadd_rn(__fadd_rn(__fadd_rn(red_l[0][row], red_l[1][row]), red_l[2][row]), red_l[3][row]); + l[h] = __fmaf_rn(l[h], alpha[h], sum); + m[h] = next_m[h]; + } +#pragma unroll + for (int q = 0; q < 8; ++q) +#pragma unroll + for (int e = 0; e < 4; ++e) o[q][e] = __fmul_rn(o[q][e], alpha[e >> 1]); + for (int ks = 0; ks < TK / 16; ++ks) { + uint32_t a[4]; + qmm_frag::ldmatrix4(a, &ps[lane % 16][ks * 16 + (lane / 16) * 8]); +#pragma unroll + for (int q = 0; q < 8; q += 2) { + uint32_t b[4]; // two 8-column tiles of the values, transposed into the B layout + ldmatrix4t(b, &vt[ks * 16 + lane % 16][warp * 64 + q * 8 + (lane / 16) * 8]); + qmm_frag::mma(o[q], a, b[0], b[1]); + qmm_frag::mma(o[q + 1], a, b[2], b[3]); + } + } + __syncthreads(); + } +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = g + 8 * h; + if (i >= G) continue; + const long long base = (static_cast(r) * NCH + c) * H + hk * G + i; +#pragma unroll + for (int q = 0; q < 8; ++q) { + float2 v = make_float2(o[q][2 * h], o[q][2 * h + 1]); + *reinterpret_cast(PO + base * HD + warp * 64 + q * 8 + 2 * t) = v; + } + if (warp == 0 && t == 0) { + PM[base] = m[h]; + PL[base] = l[h]; + } + } +} + +// attn_chunks one tile at a time (prompt chunks: two blocks a multiprocessor). +extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + float SCALE, int IDW, int QSA) { + attn_chunks<1>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); +} + +// attn_chunks with the next tile's loads in flight (a decode round's few rows): the same sums. +extern "C" __global__ void __launch_bounds__(128) fn_attn_chunks_ahead(const __nv_bfloat16* __restrict__ Q, const __nv_bfloat16* __restrict__ KC, + const __nv_bfloat16* __restrict__ VC, const int* __restrict__ POS0, + float* __restrict__ PO, float* __restrict__ PM, float* __restrict__ PL, + const int* __restrict__ IDS, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + float SCALE, int IDW, int QSA) { + attn_chunks<2>(Q, KC, VC, POS0, PO, PM, PL, IDS, NKR, SPR, H, HK, G, CH, NCH, SCALE, IDW, QSA); +} + +// attention._merge, program (row, head) of HD threads: the head's chunk partials merged in chunk order, o / l. +extern "C" __global__ void __launch_bounds__(256) fn_attn_merge(const float* __restrict__ PO, const float* __restrict__ PM, + const float* __restrict__ PL, const int* __restrict__ POS0, + __nv_bfloat16* __restrict__ OUT, const int* __restrict__ NKR, + const int* __restrict__ SPR, int H, int HK, int G, int CH, int NCH, + int QSA) { + const int r = blockIdx.x, head = blockIdx.y, d = threadIdx.x; + int n = POS0[0] + r + 1; + if (QSA && SPR[r] != 0) n = NKR[r]; + const int chunks = (n + CH - 1) / CH; + { + float m = -INFINITY, l = 0.0f, o = 0.0f; + for (int c = 0; c < chunks; ++c) { + const long long base = (static_cast(r) * NCH + c) * H + head; + const float cm = PM[base], cl = PL[base], co = PO[base * HD + d]; + const bool active = cl > 0.0f; + const float next = active ? fmaxf(m, cm) : m; + const float a = active ? (m == -INFINITY ? 0.0f : tf_exp(__fsub_rn(m, next))) : 1.0f; + const float b = active ? tf_exp(__fsub_rn(cm, next)) : 0.0f; + o = __fmaf_rn(co, b, __fmul_rn(o, a)); + l = __fmaf_rn(cl, b, __fmul_rn(l, a)); + m = next; + } + OUT[(static_cast(r) * H + head) * HD + d] = __float2bfloat16_rn(__fdiv_rn(o, l)); + } +} diff --git a/zig/kernels/cuda/flashnext_hc.cu b/zig/kernels/cuda/flashnext_hc.cu new file mode 100644 index 0000000000..3c85e54762 --- /dev/null +++ b/zig/kernels/cuda/flashnext_hc.cu @@ -0,0 +1,394 @@ +// Flash Next's hyper-connection and n-gram glue: each row's own fixed-order fp32 sums, IEEE division and square root. + +// glue.py and qmm.py's arithmetic and bf16 rounding points, the 4-bit dots on mma tiles with K in ascending order. + +#include +#include + +#include "glue_math.cuh" +#include "qmm_frag.cuh" + +using namespace tf_glue; + +namespace fn_hc { + +constexpr int TILE_N = 64; // columns a stored weight tile holds ([N/64][K/32][64][4] words) + +__device__ __forceinline__ float bsig(float x) { return rbf(tf_sigmoid(x)); } +__device__ __forceinline__ float bsilu(float x) { return rbf(tf_silu(x)); } + +// bf16x2 of nibbles i (low half) and i + 1 (high half) of an MLX word: the integers 0..15, exact. +__device__ __forceinline__ uint32_t nibbles(uint32_t w, int i) { + const float lo = static_cast((w >> (4 * i)) & 0xFu), hi = static_cast((w >> (4 * i + 4)) & 0xFu); + return (__float_as_uint(lo) >> 16) | (__float_as_uint(hi) & 0xFFFF0000u); +} + +__device__ __forceinline__ uint32_t pack(float lo, float hi) { + return static_cast(__bfloat16_as_ushort(__float2bfloat16_rn(lo))) | + (static_cast(__bfloat16_as_ushort(__float2bfloat16_rn(hi))) << 16); +} + +// A [16, 32] bf16 tile at `x` (row stride `ld` elements) as the two k16 steps' A fragments. +__device__ __forceinline__ void a_frags(const __nv_bfloat16* x, int ld, uint32_t (&a)[2][4]) { + const int lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; +#pragma unroll + for (int s = 0; s < 2; ++s) { + const __nv_bfloat16* p = x + 16 * s + 2 * t; + a[s][0] = *reinterpret_cast(p + g * ld); + a[s][1] = *reinterpret_cast(p + (g + 8) * ld); + a[s][2] = *reinterpret_cast(p + g * ld + 8); + a[s][3] = *reinterpret_cast(p + (g + 8) * ld + 8); + } +} + +// p = x[16, 32] . q[32, 8]: one column's four words (32 nibbles), two chained k16 steps from zero. +__device__ __forceinline__ void group_dot(float (&p)[4], const uint32_t (&a)[2][4], uint4 words) { + const int t = (threadIdx.x & 31) & 3; + qmm_frag::mma0(p, a[0], nibbles(words.x, 2 * t), nibbles(words.y, 2 * t)); + qmm_frag::mma(p, a[1], nibbles(words.z, 2 * t), nibbles(words.w, 2 * t)); +} + +// A block's fixed-order sum: each warp's butterfly, then the warps' sums in warp order. +__device__ __forceinline__ float block_sum(float v, float* scratch) { + v = warp_sum(v); + const int warps = blockDim.x >> 5; + if ((threadIdx.x & 31) == 0) scratch[threadIdx.x >> 5] = v; + __syncthreads(); + float total = 0.0f; + for (int w = 0; w < warps; ++w) total = __fadd_rn(total, scratch[w]); + __syncthreads(); + return total; +} + +} // namespace fn_hc + +using namespace fn_hc; + +// glue._hc_writeback, program (row, 256 columns): MODE 0 sums only, 2 MoE slots (fp32 y and weights), 4 K slices. +template +__device__ __forceinline__ void hc_writeback(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* __restrict__ PSS, + const float* __restrict__ BR, const __nv_bfloat16* __restrict__ INJ, + const float* __restrict__ Y, const float* __restrict__ WTS, long long RS, int D, + int S, int TOPK, int SLOTS, int WORLD) { + constexpr int MAXS = 8; + __shared__ float red[8][MAXS]; + const int r = blockIdx.x, c = blockIdx.y, nc = D / 256; + const int d = c * 256 + threadIdx.x; + float branch = 0.0f; + if constexpr (MODE == 4) { + float acc = BR[static_cast(r) * D + d]; + for (int k = 1; k < WORLD; ++k) acc = __fadd_rn(acc, BR[k * RS + static_cast(r) * D + d]); + branch = rbf(acc); + } else if constexpr (MODE == 2) { + float acc = 0.0f; + for (int k = 0; k <= TOPK; ++k) { + const long long at = static_cast(r) * SLOTS + k; + acc = __fmaf_rn(Y[at * D + d], WTS[at], acc); + } + branch = rbf(acc); + } + const int warp = threadIdx.x >> 5; + for (int s = 0; s < S; ++s) { + const long long at = static_cast(r) * S * D + static_cast(s) * D + d; + float hv = bf(H[at]); + if constexpr (MODE != 0) { + hv = rbf(__fadd_rn(hv, rbf(__fmul_rn(branch, bf(INJ[r * S + s]))))); + HOUT[at] = __float2bfloat16_rn(hv); + } + const float ss = warp_sum(__fmul_rn(hv, hv)); + if ((threadIdx.x & 31) == 0) red[warp][s] = ss; + } + __syncthreads(); + if (threadIdx.x < S) { + float total = 0.0f; + for (int w = 0; w < 8; ++w) total = __fadd_rn(total, red[w][threadIdx.x]); + PSS[(static_cast(r) * nc + c) * S + threadIdx.x] = total; + } +} + +extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback0(const __nv_bfloat16* H, float* PSS, int D, int S) { + hc_writeback<0>(H, nullptr, PSS, nullptr, nullptr, nullptr, nullptr, 0, D, S, 0, 1, 1); +} + +extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_moe(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* PSS, + const __nv_bfloat16* INJ, const float* Y, + const float* WTS, int D, int S, int TOPK, int SLOTS) { + hc_writeback<2>(H, HOUT, PSS, nullptr, INJ, Y, WTS, 0, D, S, TOPK, SLOTS, 1); +} + +extern "C" __global__ void __launch_bounds__(256) fn_hc_writeback_slices(const __nv_bfloat16* H, __nv_bfloat16* HOUT, float* PSS, + const float* BR, const __nv_bfloat16* INJ, + long long RS, int D, int S, int WORLD) { + hc_writeback<4>(H, HOUT, PSS, BR, INJ, nullptr, nullptr, RS, D, S, 0, 1, WORLD); +} + +// qmm._qmm_hcdown, block (16 rows, 64 columns, K slice): the slice's normed rows times the down rows, fp32 to PART. +extern "C" __global__ void __launch_bounds__(128) fn_hc_down(const __nv_bfloat16* __restrict__ H, const float* __restrict__ PSS, + const float* __restrict__ SCALE, __nv_bfloat16* __restrict__ NORMED, + const uint4* __restrict__ W, const __nv_bfloat16* __restrict__ S, + const __nv_bfloat16* __restrict__ B, float* __restrict__ PART, + int M, float eps, int N, int K, int D, int NC, int SS, int SK) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + __shared__ float rinv[16]; + const int kg = K / 32, per = kg / SK, width = per * 32, ld = width + 8; + __nv_bfloat16* xt = reinterpret_cast<__nv_bfloat16*>(smem_raw); // [16][width + 8] the slice's normed rows + float* xsum = reinterpret_cast(smem_raw + 16 * ld * 2); // [16][per] their 32-group sums + const int m0 = blockIdx.x * 16, n0 = blockIdx.y * TILE_N, slice = blockIdx.z, k0 = slice * width; + const int st = k0 / D; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + if (threadIdx.x < 16) { + const int r = m0 + threadIdx.x; + float total = 0.0f; + if (r < M) + for (int c = 0; c < NC; ++c) total = __fadd_rn(total, PSS[(static_cast(r) * NC + c) * SS + st]); + rinv[threadIdx.x] = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(total, static_cast(D)), eps))); + } + __syncthreads(); + const int vecs = width / 8; + for (int e = threadIdx.x; e < 16 * vecs; e += blockDim.x) { + const int rr = e / vecs, v = e % vecs, r = m0 + rr; + uint4 out = make_uint4(0, 0, 0, 0); + if (r < M) { + const long long at = static_cast(r) * K + k0 + 8 * v; + const uint4 raw = *reinterpret_cast(H + at); + const __nv_bfloat16* hv = reinterpret_cast(&raw); + float x[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) x[i] = rbf(__fmul_rn(__fmul_rn(bf(hv[i]), rinv[rr]), SCALE[k0 + 8 * v + i])); + out = make_uint4(pack(x[0], x[1]), pack(x[2], x[3]), pack(x[4], x[5]), pack(x[6], x[7])); + if (blockIdx.y == 0) *reinterpret_cast(NORMED + at) = out; + } + *reinterpret_cast(xt + rr * ld + 8 * v) = out; + } + __syncthreads(); + for (int e = threadIdx.x; e < 16 * per; e += blockDim.x) { + const int rr = e / per, i = e % per; + float sum = 0.0f; + for (int k = 0; k < 32; ++k) sum = __fadd_rn(sum, bf(xt[rr * ld + 32 * i + k])); + xsum[rr * per + i] = sum; + } + __syncthreads(); + float acc[2][4] = {}; + const uint4* tile = W + static_cast(n0 / TILE_N) * kg * TILE_N; +#pragma unroll 2 + for (int i = 0; i < per; ++i) { + const int grp = slice * per + i; + uint32_t a[2][4]; + a_frags(xt + 32 * i, ld, a); + const float xs0 = xsum[g * per + i], xs1 = xsum[(g + 8) * per + i]; +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int local = warp * 16 + j * 8 + g; + float p[4]; + group_dot(p, a, tile[static_cast(grp) * TILE_N + local]); + const int n = n0 + warp * 16 + j * 8 + 2 * t; + const float s0 = n < N ? bf(S[static_cast(grp) * N + n]) : 0.0f; + const float s1 = n + 1 < N ? bf(S[static_cast(grp) * N + n + 1]) : 0.0f; + const float b0 = n < N ? bf(B[static_cast(grp) * N + n]) : 0.0f; + const float b1 = n + 1 < N ? bf(B[static_cast(grp) * N + n + 1]) : 0.0f; + acc[j][0] = __fmaf_rn(xs0, b0, __fmaf_rn(p[0], s0, acc[j][0])); + acc[j][1] = __fmaf_rn(xs0, b1, __fmaf_rn(p[1], s1, acc[j][1])); + acc[j][2] = __fmaf_rn(xs1, b0, __fmaf_rn(p[2], s0, acc[j][2])); + acc[j][3] = __fmaf_rn(xs1, b1, __fmaf_rn(p[3], s1, acc[j][3])); + } + } +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int n = n0 + warp * 16 + j * 8 + 2 * t; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int r = m0 + g + 8 * h; + if (r >= M) continue; + float* out = PART + (static_cast(slice) * M + r) * N; + if (n < N) out[n] = acc[j][2 * h]; + if (n + 1 < N) out[n + 1] = acc[j][2 * h + 1]; + } + } +} + +// glue._hc_reduce_act, a thread a column: SK (<= 32) slices in order, bf16 SiLU and its sums, the inject gates. +extern "C" __global__ void __launch_bounds__(1024) fn_hc_reduce_act(const float* __restrict__ PART, __nv_bfloat16* __restrict__ ACT, + float* __restrict__ XS, __nv_bfloat16* __restrict__ INJ, + int SK, int M, int S, int LOW, int HAS_INJ, int NDN) { + __shared__ float a_s[1024]; + const int r = blockIdx.x, i = threadIdx.x; + if (i < LOW + (HAS_INJ ? S : 0)) { + float part[32]; +#pragma unroll + for (int s = 0; s < 32; ++s) part[s] = s < SK ? PART[(static_cast(s) * M + r) * NDN + i] : 0.0f; + float v = part[0]; +#pragma unroll + for (int s = 1; s < 32; ++s) + if (s < SK) v = __fadd_rn(v, part[s]); + v = rbf(__fdiv_rn(rbf(v), static_cast(S))); + if (i < LOW) { + const float a = bsilu(v); + ACT[static_cast(r) * LOW + i] = __float2bfloat16_rn(a); + a_s[i] = a; + } else { + INJ[static_cast(r) * S + (i - LOW)] = __float2bfloat16_rn(__fmul_rn(2.0f, bsig(v))); + } + } + __syncthreads(); + if (i < LOW / 32) { + float sum = 0.0f; + for (int k = 0; k < 32; ++k) sum = __fadd_rn(sum, a_s[i * 32 + k]); + XS[static_cast(r) * (LOW / 32) + i] = sum; + } +} + +// qmm._qmm_upmix, block (16 rows, 32 dims), two warps a stream (SS 4): the up rows, then the streams mixed in order. +extern "C" __global__ void __launch_bounds__(256) fn_hc_upmix(const __nv_bfloat16* __restrict__ X, const float* __restrict__ XS, + const uint4* __restrict__ W, const __nv_bfloat16* __restrict__ S, + const __nv_bfloat16* __restrict__ B, const __nv_bfloat16* __restrict__ NORMED, + __nv_bfloat16* __restrict__ MIXED, float* __restrict__ XSM, + int M, int N, int K, int D, int SS) { + extern __shared__ __align__(16) unsigned char smem_raw[]; + const int kg = K / 32; + const int ld = K + 8; + __nv_bfloat16* xt = reinterpret_cast<__nv_bfloat16*>(smem_raw); // [16][K + 8] + float* part = reinterpret_cast(smem_raw + 16 * ld * 2); // [SS][16][32] each stream's terms + const int m0 = blockIdx.x * 16, j = blockIdx.y; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, g = lane >> 2, t = lane & 3; + for (int e = threadIdx.x; e < 16 * (K / 8); e += blockDim.x) { + const int rr = e / (K / 8), v = e % (K / 8), r = m0 + rr; + *reinterpret_cast(xt + rr * ld + 8 * v) = r < M ? *reinterpret_cast(X + static_cast(r) * K + 8 * v) : make_uint4(0, 0, 0, 0); + } + __syncthreads(); + const int st = warp >> 1; + const int row0 = st * D + j * 32; + const int q0 = 2 * (warp & 1); // this warp's two of the stream's four 8-column tiles + float acc[4][4] = {}; +#pragma unroll 5 + for (int grp = 0; grp < kg; ++grp) { + uint32_t a[2][4]; + a_frags(xt + grp * 32, ld, a); + const float xs0 = m0 + g < M ? XS[static_cast(m0 + g) * kg + grp] : 0.0f; + const float xs1 = m0 + g + 8 < M ? XS[static_cast(m0 + g + 8) * kg + grp] : 0.0f; +#pragma unroll + for (int q = q0; q < q0 + 2; ++q) { + const int col = row0 + q * 8 + g; + float p[4]; + group_dot(p, a, W[(static_cast(col / TILE_N) * kg + grp) * TILE_N + col % TILE_N]); + const int n = row0 + q * 8 + 2 * t; + const float s0 = bf(S[static_cast(grp) * N + n]), s1 = bf(S[static_cast(grp) * N + n + 1]); + const float b0 = bf(B[static_cast(grp) * N + n]), b1 = bf(B[static_cast(grp) * N + n + 1]); + acc[q][0] = __fmaf_rn(xs0, b0, __fmaf_rn(p[0], s0, acc[q][0])); + acc[q][1] = __fmaf_rn(xs0, b1, __fmaf_rn(p[1], s1, acc[q][1])); + acc[q][2] = __fmaf_rn(xs1, b0, __fmaf_rn(p[2], s0, acc[q][2])); + acc[q][3] = __fmaf_rn(xs1, b1, __fmaf_rn(p[3], s1, acc[q][3])); + } + } +#pragma unroll + for (int q = q0; q < q0 + 2; ++q) +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int rr = g + 8 * (e >> 1), dd = q * 8 + 2 * t + (e & 1), r = m0 + rr; + const float nv = r < M ? bf(NORMED[static_cast(r) * N + row0 + q * 8 + 2 * t + (e & 1)]) : 0.0f; + part[(st * 16 + rr) * 32 + dd] = rbf(__fmul_rn(bsig(rbf(acc[q][e])), nv)); + } + __syncthreads(); + __shared__ float mixed_s[16][32]; + for (int e = threadIdx.x; e < 16 * 32; e += blockDim.x) { + const int rr = e >> 5, dd = e & 31; + float total = 0.0f; + for (int s = 0; s < SS; ++s) total = __fadd_rn(total, part[(s * 16 + rr) * 32 + dd]); + const float m = rbf(__fdiv_rn(total, static_cast(SS))); + mixed_s[rr][dd] = m; + if (m0 + rr < M) MIXED[static_cast(m0 + rr) * D + j * 32 + dd] = __float2bfloat16_rn(m); + } + __syncthreads(); + if (threadIdx.x < 16 && m0 + threadIdx.x < M) { + float sum = 0.0f; + for (int k = 0; k < 32; ++k) sum = __fadd_rn(sum, mixed_s[threadIdx.x][k]); + XSM[static_cast(m0 + threadIdx.x) * (D / 32) + j] = sum; + } +} + +// glue._ple_gate, a block a row: per stream normed keys and queries, their dot's signed root gating the values. +extern "C" __global__ void __launch_bounds__(512) fn_ple_gate(const __nv_bfloat16* __restrict__ KEYS, const __nv_bfloat16* __restrict__ VALS, + const __nv_bfloat16* __restrict__ H, const float* __restrict__ NK, + const float* __restrict__ NQ, __nv_bfloat16* __restrict__ GATED, + float* __restrict__ PSS, float eps, int D, int S) { + __shared__ float scratch[32]; + const int r = blockIdx.x; + const float fd = static_cast(D); + for (int s = 0; s < S; ++s) { + const long long row = static_cast(r) * S * D + static_cast(s) * D; + float kss = 0.0f, qss = 0.0f; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float kv = bf(KEYS[row + d]), hv = bf(H[row + d]); + kss = __fmaf_rn(kv, kv, kss); + qss = __fmaf_rn(hv, hv, qss); + } + kss = block_sum(kss, scratch); + qss = block_sum(qss, scratch); + const float krinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(kss, fd), eps))); + const float qrinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(qss, fd), eps))); + float dot = 0.0f; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float kn = rbf(__fmul_rn(__fmul_rn(bf(KEYS[row + d]), krinv), NK[s * D + d])); + const float qn = rbf(__fmul_rn(__fmul_rn(bf(H[row + d]), qrinv), NQ[s * D + d])); + dot = __fadd_rn(dot, rbf(__fmul_rn(kn, qn))); + } + dot = block_sum(dot, scratch); + float gate = rbf(__fdiv_rn(dot, __fsqrt_rn(fd))); + const float mag = __fsqrt_rn(fmaxf(fabsf(gate), 1e-6f)); + gate = rbf(gate > 0.0f ? mag : (gate < 0.0f ? -mag : 0.0f)); + const float sg = bsig(gate); + float gss = 0.0f; + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const float gv = rbf(__fmul_rn(sg, bf(VALS[static_cast(r) * D + d]))); + GATED[row + d] = __float2bfloat16_rn(gv); + gss = __fmaf_rn(gv, gv, gss); + } + gss = block_sum(gss, scratch); + if (threadIdx.x == 0) PSS[static_cast(r) * S + s] = gss; + } +} + +// glue._ple_conv, program (row, 512 columns): the dilated conv over [TAIL; rows], bf16 SiLU, the adds, NROW kept. +extern "C" __global__ void __launch_bounds__(512) fn_ple_conv(const __nv_bfloat16* __restrict__ GATED, const float* __restrict__ PSS, + const float* __restrict__ NC, const __nv_bfloat16* __restrict__ TAIL, + const __nv_bfloat16* __restrict__ CW, const __nv_bfloat16* H, + __nv_bfloat16* HOUT, __nv_bfloat16* __restrict__ NROW, float eps, + int D, int S, int TAPS, int DIL) { + const int r = blockIdx.x; + const int w = S * D, nt = (TAPS - 1) * DIL; + const int e = blockIdx.y * blockDim.x + threadIdx.x; + if (e >= w) return; + const int s = e / D; + const float fd = static_cast(D); + float acc = 0.0f; + for (int t = 0; t < TAPS; ++t) { + const int at = r + t * DIL; + float xv; + if (at < nt) { + xv = bf(TAIL[static_cast(at) * w + e]); + } else { + const int rr = at - nt; + const float rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(PSS[rr * S + s], fd), eps))); + xv = rbf(__fmul_rn(__fmul_rn(bf(GATED[static_cast(rr) * w + e]), rinv), NC[e])); + } + acc = __fmaf_rn(bf(CW[static_cast(e) * TAPS + t]), xv, acc); + } + const float act = bsilu(rbf(acc)); + const float g = bf(GATED[static_cast(r) * w + e]); + const float ple = rbf(__fadd_rn(g, act)); + const long long at = static_cast(r) * w + e; + HOUT[at] = __float2bfloat16_rn(__fadd_rn(bf(H[at]), ple)); + const float rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(PSS[r * S + s], fd), eps))); + NROW[at] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(g, rinv), NC[e])); +} + +// glue._add_streams: out[r, s] = bf16(e[r] + hs[r, s]), the MTP head's input streams. +extern "C" __global__ void fn_add_streams(const __nv_bfloat16* __restrict__ E, const __nv_bfloat16* __restrict__ HS, + __nv_bfloat16* __restrict__ OUT, int D, int S) { + const int r = blockIdx.x, d = blockIdx.y * blockDim.x + threadIdx.x; + if (d >= D) return; + const float e = bf(E[static_cast(r) * D + d]); + for (int s = 0; s < S; ++s) { + const long long at = (static_cast(r) * S + s) * D + d; + OUT[at] = __float2bfloat16_rn(__fadd_rn(e, bf(HS[at]))); + } +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index a351ea7e25..30307716e0 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -52,6 +52,8 @@ pub const nemotron_keyed: []const u8 = if (available) &Blob("fatbin_nemotron_key pub const norms: []const u8 = if (available) &Blob("fatbin_norms").bytes else &.{}; pub const embed_rows: []const u8 = if (available) &Blob("fatbin_embed_rows").bytes else &.{}; pub const route_rows: []const u8 = if (available) &Blob("fatbin_route_rows").bytes else &.{}; +pub const flashnext_hc: []const u8 = if (available) &Blob("fatbin_flashnext_hc").bytes else &.{}; +pub const flashnext_attn: []const u8 = if (available) &Blob("fatbin_flashnext_attn").bytes else &.{}; /// Symbols in the gdn image as cuobjdump lists them for the built fatbin (named namespace tf_gdn). pub const gdn_symbols = struct { diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig new file mode 100644 index 0000000000..0bebfcfa84 --- /dev/null +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -0,0 +1,270 @@ +//! Flash Next's own glue kernels (flashnext_hc.cu, flashnext_attn.cu) with the core's shared rows: the default path. + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); + +/// The attention tiles' head width and query heads a KV head (flashnext_attn.cu's HD and QR). +const tile_head_dim = 256; +const tile_heads = 16; +/// fn_attn_chunks' staged rows: the query heads, `buffers` 32-key tiles of keys and values (rows padded by 8), bf16 p. +fn chunkSmem(buffers: u32) u32 { + return (tile_heads + buffers * 2 * 32) * (tile_head_dim + 8) * 2 + tile_heads * (32 + 8) * 2; +} +/// Rows up to which attention reads the next tile ahead (fewer blocks, so latency, not occupancy, sets the time). +const ahead_rows = 16; + +fn int(x: usize) c_int { + return @intCast(x); +} + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// The native kernels, by their extern "C" names. +const Fns = struct { + wb0: cuda.Function, + wb_moe: cuda.Function, + wb_slices: cuda.Function, + down: cuda.Function, + reduce_act: cuda.Function, + upmix: cuda.Function, + ple_gate: cuda.Function, + ple_conv: cuda.Function, + add_streams: cuda.Function, + prep: cuda.Function, + gate: cuda.Function, + pool: cuda.Function, + scores: cuda.Function, + chunks: cuda.Function, + chunks_ahead: cuda.Function, + merge: cuda.Function, +}; + +pub const Glue = struct { + hc: cuda.Module, + attn: cuda.Module, + emb: cuda.embed_rows.EmbedRows, + norms: cuda.norms.Norms, + route: cuda.route_rows.RouteRows, + f: Fns, + + pub fn load(d: *const cuda.Driver) !Glue { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var g: Glue = undefined; + g.hc = try cuda.Module.load(d, cuda.kernels.flashnext_hc); + errdefer g.hc.unload(); + g.attn = try cuda.Module.load(d, cuda.kernels.flashnext_attn); + errdefer g.attn.unload(); + g.emb = try cuda.embed_rows.EmbedRows.load(d); + errdefer g.emb.unload(); + g.norms = try cuda.norms.Norms.load(d); + errdefer g.norms.unload(); + g.route = try cuda.route_rows.RouteRows.load(d); + errdefer g.route.unload(); + g.f = .{ + .wb0 = try g.hc.function("fn_hc_writeback0"), + .wb_moe = try g.hc.function("fn_hc_writeback_moe"), + .wb_slices = try g.hc.function("fn_hc_writeback_slices"), + .down = try g.hc.function("fn_hc_down"), + .reduce_act = try g.hc.function("fn_hc_reduce_act"), + .upmix = try g.hc.function("fn_hc_upmix"), + .ple_gate = try g.hc.function("fn_ple_gate"), + .ple_conv = try g.hc.function("fn_ple_conv"), + .add_streams = try g.hc.function("fn_add_streams"), + .prep = try g.attn.function("fn_attn_prep"), + .gate = try g.attn.function("fn_attn_gate"), + .pool = try g.attn.function("fn_qsa_pool"), + .scores = try g.attn.function("fn_qsa_scores"), + .chunks = try g.attn.function("fn_attn_chunks"), + .chunks_ahead = try g.attn.function("fn_attn_chunks_ahead"), + .merge = try g.attn.function("fn_attn_merge"), + }; + try g.f.chunks.allowDynamicShared(chunkSmem(1)); + try g.f.chunks_ahead.allowDynamicShared(chunkSmem(2)); + return g; + } + + pub fn deinit(g: *Glue) void { + g.route.unload(); + g.norms.unload(); + g.emb.unload(); + g.attn.unload(); + g.hc.unload(); + } + + fn go(f: cuda.Function, s: cuda.Stream, grid: [3]usize, block: u32, shared: u32, a: *cuda.Args) !void { + if (grid[0] == 0 or grid[1] == 0 or grid[2] == 0) return; + const cfg: cuda.Config = .{ .grid = .{ .x = @intCast(grid[0]), .y = @intCast(grid[1]), .z = @intCast(grid[2]) }, .block = .{ .x = block }, .shared = shared }; + try cuda.launch.launch(f, cfg, s, a); + } + + fn ptrs(a: *cuda.Args, xs: []const u64) void { + for (xs) |x| a.add(x); + } + + fn ints(a: *cuda.Args, xs: []const usize) void { + for (xs) |x| a.add(int(x)); + } + + /// Token rows of the MLX 4-bit group-32 table into `copies` streams each (the core's embed rows). + pub fn embed(g: *const Glue, s: cuda.Stream, ids: u64, w: u64, sc: u64, b: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + try g.emb.affine4Groups(s, ids, w, sc, b, out, rows, dims, 32, copies); + } + + /// One RMS norm over each row's `dims` features (rows `x_stride` apart), bf16 out and 32-group sums (core norms). + pub fn rmsnorm(g: *const Glue, s: cuda.Stream, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { + if (group != dims) return error.RowShape; // the head norms each row whole + try g.norms.rmsSums(s, x, x_stride, w, out, xs, rows, dims, 32, eps); + } + + pub fn addStreams(g: *const Glue, s: cuda.Stream, e: u64, hs: u64, out: u64, rows: usize, dims: usize, streams: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ e, hs, out }); + ints(&a, &.{ dims, streams }); + try go(g.f.add_streams, s, .{ rows, cdiv(dims, 256), 1 }, 256, 0, &a); + } + + pub fn attnGate(g: *const Glue, s: cuda.Stream, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { + if (head_dim != tile_head_dim) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ o, proj, out, xs }); + ints(&a, &.{ pair_width, q_heads }); + try go(g.f.gate, s, .{ rows, q_heads, 1 }, tile_head_dim, 0, &a); + } + + /// Text positions, a bf16 cache: the model's heads (cuda_prompt_wt.zig) normed, roped and stored. + pub fn attnPrep(g: *const Glue, s: cuda.Stream, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, eps: f32, rows: usize) !void { + if (wt.head_dim != tile_head_dim) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc }); + a.add(eps); + ints(&a, &.{ wt.proj_n, wt.q_heads, wt.kv_heads, wt.index_heads, wt.index_dim, wt.half }); + try go(g.f.prep, s, .{ rows, wt.q_heads + wt.kv_heads + wt.index_heads + 1, 1 }, tile_head_dim, 0, &a); + } + + /// Attention over 512-key chunks; `qsa` the sparse form (NCH 5, rows' key lists), else dense (NCH 3). + pub fn attnChunks(g: *const Glue, s: cuda.Stream, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { + const group = wt.q_heads / kv; + if (group > tile_heads) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ q, kc, vc, pos, po, pm, pl, ids, nk, sparse }); + ints(&a, &.{ wt.q_heads, kv, group, 512, if (qsa) wt.nch else wt.dense_nch }); + a.add(@as(f32, 0.0625)); + ints(&a, &.{ wt.index_idw, @intFromBool(qsa) }); + const ahead = rows <= ahead_rows; + try go(if (ahead) g.f.chunks_ahead else g.f.chunks, s, .{ rows, kv, chunks }, 128, chunkSmem(if (ahead) 2 else 1), &a); + } + + pub fn attnMerge(g: *const Glue, s: cuda.Stream, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ po, pm, pl, pos, out, nk, sparse }); + ints(&a, &.{ wt.q_heads, kv, wt.q_heads / kv, 512, if (qsa) wt.nch else wt.dense_nch, @intFromBool(qsa) }); + try go(g.f.merge, s, .{ rows, wt.q_heads, 1 }, tile_head_dim, 0, &a); + } + + pub fn qsaPool(g: *const Glue, s: cuda.Stream, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ ikc, pooled, pos, k_scale, inv }); + a.add(eps); + ints(&a, &.{ rows, wt.index_dim, wt.half, wt.index_ratio }); + try go(g.f.pool, s, .{ rows / wt.index_ratio + 2, 1, 1 }, wt.index_dim, 0, &a); + } + + pub fn qsaScores(g: *const Glue, s: cuda.Stream, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { + if (wt.index_dim != 128 or wt.index_heads > 8) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ iq, pooled, pos, sc }); + ints(&a, &.{ nb, wt.index_heads, wt.index_ratio, wt.index_top }); + try go(g.f.scores, s, .{ rows, cdiv(blocks, 64), 1 }, 256, 0, &a); + } + + pub fn hcReduceAct(g: *const Glue, s: cuda.Stream, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { + const cols = low + has_inj * streams; + if (cols > 1024 or sk > 32) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ part, act, xs, inj }); + ints(&a, &.{ sk, rows, streams, low, has_inj, ndn }); + try go(g.f.reduce_act, s, .{ rows, 1, 1 }, @intCast(cdiv(@max(cols, low / 32), 32) * 32), 0, &a); + } + + pub fn hcUpmix(g: *const Glue, s: cuda.Stream, act: u64, xs: u64, w: u64, sc: u64, b: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { + if (streams != 4 or dims % 32 != 0 or k % 32 != 0) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ act, xs, w, sc, b, normed, mixed, xs_mixed }); + ints(&a, &.{ rows, n, k, dims, streams }); + const shared: u32 = @intCast(16 * (k + 8) * 2 + streams * 16 * 32 * 4); + try go(g.f.upmix, s, .{ cdiv(rows, 16), dims / 32, 1 }, 256, shared, &a); + } + + /// The normed streams times the down rows, each K slice's fp32 sums to `part` [sk, rows, n] (sk > 1). + pub fn hcDown(g: *const Glue, s: cuda.Stream, h: u64, pss: u64, scale: u64, normed: u64, w: u64, sc: u64, b: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { + _ = out; + const per = k / 32 / sk; + if (sk < 2 or k % (32 * sk) != 0 or dims % (per * 32) != 0) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ h, pss, scale, normed, w, sc, b, part }); + a.add(int(rows)); + a.add(eps); + ints(&a, &.{ n, k, dims, nc, streams, sk }); + const shared: u32 = @intCast(16 * (per * 32 + 8) * 2 + 16 * per * 4); + try go(g.f.down, s, .{ cdiv(rows, 16), cdiv(n, 64), sk }, 128, shared, &a); + } + + /// Mode 0: the streams' sums; 2: the MoE slots (fp32 y and weights) written back; 4: `world` fp32 K slices. + pub fn hcWriteback(g: *const Glue, s: cuda.Stream, h: u64, hout: u64, pss: u64, branch: u64, inj: u64, y: u64, wts: u64, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + if (dims % 256 != 0 or streams > 8) return error.RowShape; + var a: cuda.Args = .{}; + const grid: [3]usize = .{ rows, dims / 256, 1 }; + switch (mode) { + 0 => { + ptrs(&a, &.{ h, pss }); + ints(&a, &.{ dims, streams }); + try go(g.f.wb0, s, grid, 256, 0, &a); + }, + 2 => { + ptrs(&a, &.{ h, hout, pss, inj, y, wts }); + ints(&a, &.{ dims, streams, topk, slots }); + try go(g.f.wb_moe, s, grid, 256, 0, &a); + }, + 4 => { + ptrs(&a, &.{ h, hout, pss, branch, inj }); + a.add(@as(i64, @intCast(rs))); + ints(&a, &.{ dims, streams, world }); + try go(g.f.wb_slices, s, grid, 256, 0, &a); + }, + else => return error.RowShape, + } + } + + pub fn pleGate(g: *const Glue, s: cuda.Stream, keys: u64, vals: u64, h: u64, norm_key: u64, norm_query: u64, gated: u64, pss: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ keys, vals, h, norm_key, norm_query, gated, pss }); + a.add(eps); + ints(&a, &.{ dims, streams }); + try go(g.f.ple_gate, s, .{ rows, 1, 1 }, 512, 0, &a); + } + + pub fn pleConv(g: *const Glue, s: cuda.Stream, gated: u64, pss: u64, norm_conv: u64, tail: u64, conv_w: u64, h: u64, hout: u64, nrow: u64, eps: f32, rows: usize, dims: usize, streams: usize, taps: usize, dilation: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ gated, pss, norm_conv, tail, conv_w, h, hout, nrow }); + a.add(eps); + ints(&a, &.{ dims, streams, taps, dilation }); + try go(g.f.ple_conv, s, .{ rows, cdiv(streams * dims, 512), 1 }, 512, 0, &a); + } + + /// Router logits [rows, experts + 1] fp32: the experts' rows, then the shared expert's gate row (core route rows). + pub fn routerMany(g: *const Glue, s: cuda.Stream, x: u64, w: u64, out: u64, x_stride: usize, rows: usize) !void { + try g.route.logits(s, x, x_stride, w, out, rows, wt.dims, wt.experts_n + 1); + } + + /// Each row's ten experts and the shared one: picks and weights [rows, 11] (core route rows). + pub fn topkMany(g: *const Glue, s: cuda.Stream, logits: u64, pick: u64, wts: u64, rows: usize) !void { + try g.route.topk(s, logits, wt.experts_n + 1, wt.experts_n, wt.slots - 1, pick, wts, wt.slots, rows, true); + } +}; + +test "the attention tiles hold the model's heads" { + try std.testing.expect(wt.head_dim == tile_head_dim and wt.q_heads / wt.kv_heads <= tile_heads); +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index b779560d7f..e6de83ba1e 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -20,6 +20,8 @@ pub const row_path = @import("cuda_rows.zig"); /// The MTP head's drafts: its forward, the chain, the verify window's rows and their rollback pub const mtp = @import("cuda_mtp.zig"); pub const native = @import("cuda_lanes.zig").Namespace(Tri); +/// Our own glue kernels, which every launch takes when no captured set is loaded +pub const Glue = @import("cuda_glue.zig").Glue; const p = aot.ptr; diff --git a/zig/tests/cuda/flashnext_attn_tests.zig b/zig/tests/cuda/flashnext_attn_tests.zig new file mode 100644 index 0000000000..806a7207bf --- /dev/null +++ b/zig/tests/cuda/flashnext_attn_tests.zig @@ -0,0 +1,282 @@ +//! Flash Next's own attention glue and its indexer's pooling and scores against float64 references. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const Rig = rig.Rig; +const bf = rig.bf; +const rbf = rig.rbf; +const ulps = rig.ulps; +const sumTol = rig.sumTol; + +const q_heads = 24; +const kv_heads = 2; +const head_dim = 256; +const index_heads = 4; +const index_dim = 128; +const half = 32; +const proj_n = 13952; +const eps: f32 = 1e-6; +const cap = 3072; +const idw = 2052; + +/// RoPE's inverse frequencies as the pass uploads them (base 10,000,000 over 32 pairs). +fn inverse(r: *Rig) ![]f32 { + const inv = try r.a.alloc(f32, half); + for (inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / half)); + return inv; +} + +/// The kernels' rotate-half RoPE of normed values xn, xpn (feature d of 2 * half rotated) at `pos`. +fn rope(xn: f64, xpn: f64, d: usize, pos: usize, inv: []const f32) f64 { + if (d >= 2 * half) return xn; + const i = if (d < half) d else d - half; + const ang: f64 = @as(f32, @floatFromInt(pos)) * inv[i]; + return if (d < half) xn * @cos(ang) - xpn * @sin(ang) else xpn * @sin(ang) + xn * @cos(ang); +} + +fn partner(d: usize) usize { + return if (d < half) d + half else if (d < 2 * half) d - half else d; +} + +/// A head of `width` features at src normed by w, rounded, roped at pos and rounded: the float64 chain. +fn prepped(p: []const u16, src: usize, width: usize, w: []const f32, pos: usize, inv: []const f32, d: usize) f64 { + var ss: f64 = 0; + for (0..width) |i| ss += bf(p[src + i]) * bf(p[src + i]); + const rinv = 1 / @sqrt(ss / @as(f64, @floatFromInt(width)) + eps); + const xn = rbf(bf(p[src + d]) * rinv * w[d]); + const xpn = rbf(bf(p[src + partner(d)]) * rinv * w[partner(d)]); + return rbf(rope(xn, xpn, d, pos, inv)); +} + +fn prep(r: *Rig, g: *const flash.Glue) !void { + const counts = [_]usize{ 1, 3, 17 }; + const most = counts[counts.len - 1]; + const pos0 = 700; + const p = r.bfs(most * proj_n, 2); + const qw = r.f32s(head_dim, 0.5, 1.5); + const kw = r.f32s(head_dim, 0.5, 1.5); + const iw = r.f32s(index_dim, 0.5, 1.5); + const inv = try inverse(r); + const dp = try r.dev(std.mem.sliceAsBytes(p)); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const dqw = try r.dev(std.mem.sliceAsBytes(qw)); + const dkw = try r.dev(std.mem.sliceAsBytes(kw)); + const diw = try r.dev(std.mem.sliceAsBytes(iw)); + const dinv = try r.dev(std.mem.sliceAsBytes(inv)); + var qs: [counts.len]u64 = undefined; + var iqs: [counts.len]u64 = undefined; + var kcs: [counts.len]u64 = undefined; + var vcs: [counts.len]u64 = undefined; + var ikcs: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + qs[c] = try r.out(m * q_heads * head_dim * 2); + iqs[c] = try r.out(m * index_heads * index_dim * 2); + kcs[c] = try r.out(cap * kv_heads * head_dim * 2); + vcs[c] = try r.out(cap * kv_heads * head_dim * 2); + ikcs[c] = try r.out(cap * index_dim * 2); + try g.attnPrep(r.s, dp, dpos, dqw, dkw, diw, dinv, qs[c], kcs[c], vcs[c], iqs[c], ikcs[c], eps, m); + } + const last = counts.len - 1; + const q = try r.back(u16, qs[last], most * q_heads * head_dim); + const iq = try r.back(u16, iqs[last], most * index_heads * index_dim); + const kc = try r.back(u16, kcs[last], cap * kv_heads * head_dim); + const vc = try r.back(u16, vcs[last], cap * kv_heads * head_dim); + const ikc = try r.back(u16, ikcs[last], cap * index_dim); + for (0..most) |row| { + const pos = pos0 + row; + const base = row * proj_n; + for (0..q_heads) |h| for (0..head_dim) |d| { + const want = prepped(p, base + h * 2 * head_dim, head_dim, qw, pos, inv, d); + try r.within("prep queries", bf(q[(row * q_heads + h) * head_dim + d]), want, ulps(want, 1) + 1e-3); + }; + for (0..kv_heads) |h| for (0..head_dim) |d| { + const want = prepped(p, base + q_heads * 2 * head_dim + h * head_dim, head_dim, kw, pos, inv, d); + const at = (pos * kv_heads + h) * head_dim + d; + try r.within("prep keys", bf(kc[at]), want, ulps(want, 1) + 1e-3); + try check.expect(vc[at] == p[base + q_heads * 2 * head_dim + kv_heads * head_dim + h * head_dim + d], "prep values: row {d} head {d} copied", .{ row, h }); + }; + for (0..index_heads) |h| for (0..index_dim) |d| { + const want = prepped(p, base + q_heads * 2 * head_dim + 2 * kv_heads * head_dim + h * index_dim, index_dim, iw, pos, inv, d); + try r.within("prep indexer queries", bf(iq[(row * index_heads + h) * index_dim + d]), want, ulps(want, 1) + 1e-3); + }; + for (0..index_dim) |d| try check.expect(ikc[pos * index_dim + d] == p[base + q_heads * 2 * head_dim + 2 * kv_heads * head_dim + index_heads * index_dim + d], "prep indexer keys: row {d} copied", .{row}); + } + for (counts[0..last], 0..) |m, c| { + try r.sameRows("prep queries", u16, qs[c], q[0 .. m * q_heads * head_dim]); + try r.sameRows("prep keys", u16, kcs[c] + pos0 * kv_heads * head_dim * 2, kc[pos0 * kv_heads * head_dim ..][0 .. m * kv_heads * head_dim]); + } + // the output gate: bf16(o / (1 + e^-gate)) and its group sums + const o = r.bfs(most * q_heads * head_dim, 2); + const out = try r.out(most * q_heads * head_dim * 2); + const xs = try r.out(most * q_heads * head_dim / 32 * 4); + try g.attnGate(r.s, try r.dev(std.mem.sliceAsBytes(o)), dp, out, xs, most, proj_n, q_heads, head_dim); + const go = try r.back(u16, out, most * q_heads * head_dim); + const gx = try r.back(f32, xs, most * q_heads * head_dim / 32); + for (0..most) |row| for (0..q_heads) |h| { + for (0..head_dim / 32) |gi| { + var sum: f64 = 0; + var abs: f64 = 0; + for (0..32) |k| { + const d = gi * 32 + k; + const at = (row * q_heads + h) * head_dim + d; + const want = bf(o[at]) * rig.sigmoid(bf(p[row * proj_n + h * 2 * head_dim + head_dim + d])); + try r.within("output gate", bf(go[at]), want, ulps(want, 0.5) + @abs(want) * 0x1p-18); + sum += bf(go[at]); + abs += @abs(bf(go[at])); + } + try r.within("output gate sums", gx[(row * q_heads + h) * (head_dim / 32) + gi], sum, sumTol(abs, 32)); + } + }; +} + +/// The indexer's pooled blocks and each row's block scores. +fn indexer(r: *Rig, g: *const flash.Glue) !void { + const pos0 = 2196; + const rows = 5; + const ikc = r.bfs(cap * index_dim, 2); + const w = r.f32s(index_dim, 0.5, 1.5); + const inv = try inverse(r); + const dikc = try r.dev(std.mem.sliceAsBytes(ikc)); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ pos0, 0, 0, 0 })); + const pooled = try r.out((cap / 4) * index_dim * 2); + try g.qsaPool(r.s, dikc, pooled, dpos, try r.dev(std.mem.sliceAsBytes(w)), try r.dev(std.mem.sliceAsBytes(inv)), eps, rows); + const got = try r.back(u16, pooled, (cap / 4) * index_dim); + var b: usize = pos0 / 4; + while (4 * b + 4 <= pos0 + rows) : (b += 1) { + var x: [index_dim]f64 = undefined; + var ss: f64 = 0; + for (&x, 0..) |*v, d| { + var acc: f64 = 0; + for (0..4) |k| acc += bf(ikc[(4 * b + k) * index_dim + d]); + v.* = rbf(acc / 4); + ss += v.* * v.*; + } + const rinv = 1 / @sqrt(ss / index_dim + eps); + for (0..index_dim) |d| { + const xn = rbf(x[d] * rinv * w[d]); + const xpn = rbf(x[partner(d)] * rinv * w[partner(d)]); + const want = rbf(rope(xn, xpn, d, 4 * b, inv)); + try r.within("indexer pool", bf(got[b * index_dim + d]), want, ulps(want, 1) + 1e-3); + } + } + // scores of every complete block for rows past 512 of them + const keys = r.bfs((cap / 4) * index_dim, 1); + const iq = r.bfs(rows * index_heads * index_dim, 1); + const nb = 784; + const dkeys = try r.dev(std.mem.sliceAsBytes(keys)); + const diq = try r.dev(std.mem.sliceAsBytes(iq)); + var scs: [2]u64 = undefined; + for ([_]usize{ 1, rows }, 0..) |m, c| { + scs[c] = try r.out(m * nb * 4); + try g.qsaScores(r.s, diq, dkeys, dpos, scs[c], nb, m, (pos0 + rows) / 4); + } + const sc = try r.back(f32, scs[1], rows * nb); + for (0..rows) |row| { + const complete = (pos0 + row + 1) / 4; + for (0..complete) |blk| { + var total: f64 = 0; + var abs: f64 = 0; + for (0..index_heads) |h| { + var dot: f64 = 0; + for (0..index_dim) |d| { + const t = bf(iq[(row * index_heads + h) * index_dim + d]) * bf(keys[blk * index_dim + d]); + dot += t; + abs += @abs(t); + } + total += @max(dot, 0); + } + const want = total / @sqrt(@as(f64, index_dim)); + try r.within("indexer scores", sc[row * nb + blk], want, sumTol(abs, 2 * index_dim) / @sqrt(@as(f64, index_dim))); + } + } + try r.sameRows("indexer scores", f32, scs[0], sc[0 .. (pos0 + 1) / 4]); +} + +/// Each row's attention over its keys (dense from position 0, or a sparse row's list) in 512-key chunks, merged. +fn attention(r: *Rig, g: *const flash.Glue, sparse: bool) !void { + const counts = [_]usize{ 1, 3, 20 }; // up to 16 rows read the next tile ahead; 20 take one tile at a time + const rows = counts[counts.len - 1]; + const pos0: usize = if (sparse) 2600 else 1100; + const q = r.bfs(rows * q_heads * head_dim, 1); + const kc = r.bfs(cap * kv_heads * head_dim, 1); + const vc = r.bfs(cap * kv_heads * head_dim, 1); + const ids = try r.a.alloc(i32, rows * idw); + const nk = try r.a.alloc(i32, rows); + const spr = try r.a.alloc(i32, rows); + for (0..rows) |row| { + nk[row] = @intCast(if (sparse) idw - 1 - row % 3 else pos0 + row + 1); + spr[row] = @intFromBool(sparse); + // a sparse row's list: distinct keys before its position, in a scrambled order + for (0..idw) |i| ids[row * idw + i] = @intCast((i * 1237 + row * 31) % (pos0 + row + 1)); + } + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const dkc = try r.dev(std.mem.sliceAsBytes(kc)); + const dvc = try r.dev(std.mem.sliceAsBytes(vc)); + const dpos = try r.dev(std.mem.sliceAsBytes(&[_]i32{ @intCast(pos0), 0, 0, 0 })); + const dids = try r.dev(std.mem.sliceAsBytes(ids)); + const dnk = try r.dev(std.mem.sliceAsBytes(nk)); + const dspr = try r.dev(std.mem.sliceAsBytes(spr)); + const nch = 5; + const po = try r.out(rows * nch * q_heads * head_dim * 4); + const pm = try r.out(rows * nch * q_heads * 4); + const pl = try r.out(rows * nch * q_heads * 4); + var outs: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + outs[c] = try r.out(m * q_heads * head_dim * 2); + const chunks = if (sparse) (idw + 511) / 512 else (pos0 + m + 511) / 512; + try g.attnChunks(r.s, dq, dkc, dvc, dpos, po, pm, pl, dids, dnk, dspr, m, kv_heads, chunks, true); + try g.attnMerge(r.s, po, pm, pl, dpos, outs[c], dnk, dspr, m, kv_heads, true); + } + const got = try r.back(u16, outs[counts.len - 1], rows * q_heads * head_dim); + const name = if (sparse) "attention sparse" else "attention dense"; + const s = try r.a.alloc(f64, idw + pos0 + rows); + for (0..rows) |row| for (0..q_heads) |h| { + const kvh = h / (q_heads / kv_heads); + const n: usize = @intCast(nk[row]); + var top = -std.math.inf(f64); + for (0..n) |j| { + const key: usize = if (sparse) @intCast(ids[row * idw + j]) else j; + var dot: f64 = 0; + for (0..head_dim) |d| dot += bf(q[(row * q_heads + h) * head_dim + d]) * bf(kc[(key * kv_heads + kvh) * head_dim + d]); + s[j] = dot * 0.0625; + top = @max(top, s[j]); + } + var l: f64 = 0; + for (s[0..n]) |*v| { + v.* = @exp(v.* - top); + l += v.*; + } + for (0..head_dim) |d| { + var o: f64 = 0; + var abs: f64 = 0; + for (0..n) |j| { + const key: usize = if (sparse) @intCast(ids[row * idw + j]) else j; + const v = bf(vc[(key * kv_heads + kvh) * head_dim + d]); + o += s[j] * v; + abs += s[j] * @abs(v); + } + const want = o / l; + // p goes into p . v as bf16 (2^-9 of each term), its sum l stays fp32 + try r.within(name, bf(got[(row * q_heads + h) * head_dim + d]), want, abs / l * 0x1p-8 + ulps(want, 0.5)); + } + }; + for (counts[0 .. counts.len - 1], 0..) |m, c| try r.sameRows(name, u16, outs[c], got[0 .. m * q_heads * head_dim]); +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(59); + var r = try Rig.init(gpu, arena.allocator(), prng.random()); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + try prep(&r, &g); + try indexer(&r, &g); + try attention(&r, &g, false); + try attention(&r, &g, true); + r.report("flashnext-attn: q, k and indexer heads, the output gate, pooled blocks and scores, dense and sparse attention"); +} diff --git a/zig/tests/cuda/flashnext_hc_tests.zig b/zig/tests/cuda/flashnext_hc_tests.zig new file mode 100644 index 0000000000..bfca6aff72 --- /dev/null +++ b/zig/tests/cuda/flashnext_hc_tests.zig @@ -0,0 +1,339 @@ +//! Flash Next's own hyper-connection and n-gram glue against float64 chains that round where the kernels round. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); + +const Rig = rig.Rig; +const bf = rig.bf; +const rbf = rig.rbf; +const ulps = rig.ulps; +const sumTol = rig.sumTol; + +const dims = 2560; +const streams = 4; +const wide = streams * dims; +const low = 320; +const n_down = low + streams; +const slices = 32; +const nc = dims / 256; +const eps: f32 = 1e-6; +const counts = [_]usize{ 1, 3, 17 }; +const most = counts[counts.len - 1]; + +/// Tiled 4-bit rows on the device: words [N/64][K/32][64][4], scales and biases [K/32][N] (cuda_hc.zig's layout). +const Tiled = struct { w: []u32, s: []u16, b: []u16, n: usize, k: usize, dw: u64, ds: u64, db: u64 }; + +fn tiled(r: *Rig, n: usize, k: usize, scale: f32) !Tiled { + const w = r.words(n * k / 8); + const s = r.bfs(n * k / 32, scale); + for (s) |*v| v.* = rig.bf16(@floatCast(@abs(bf(v.*)))); + const b = r.bfs(n * k / 32, scale * 4); + const tw = try flash.hc.tileWords(r.a, std.mem.sliceAsBytes(w), n, k / 8); + const ts = try flash.hc.transposeBf16(r.a, std.mem.sliceAsBytes(s), n, k / 32); + const tb = try flash.hc.transposeBf16(r.a, std.mem.sliceAsBytes(b), n, k / 32); + return .{ .w = w, .s = s, .b = b, .n = n, .k = k, .dw = try r.dev(tw), .ds = try r.dev(ts), .db = try r.dev(tb) }; +} + +/// Row `col` of a tiled matrix over inputs [k0, k0 + 32 g): sum over groups of (x . q) s + (sum x) b, and its |terms|. +fn rowDot(t: Tiled, col: usize, x: []const f64, k0: usize, groups: usize, xs: ?[]const f64) [2]f64 { + var total: f64 = 0; + var abs: f64 = 0; + for (0..groups) |gi| { + const g = k0 / 32 + gi; + var p: f64 = 0; + var sum: f64 = 0; + var pa: f64 = 0; + for (0..32) |k| { + const q = rig.nibble(t.w, t.k / 8, col, g * 32 + k); + p += x[gi * 32 + k] * q; + pa += @abs(x[gi * 32 + k]) * q; + sum += x[gi * 32 + k]; + } + if (xs) |v| sum = v[gi]; + const s = bf(t.s[col * (t.k / 32) + g]); + const b = bf(t.b[col * (t.k / 32) + g]); + total += p * s + sum * b; + abs += pa * @abs(s) + @abs(sum * b); + } + return .{ total, abs }; +} + +fn writeback(r: *Rig, g: *const flash.Glue, mode: usize) !void { + const h = r.bfs(most * wide, 1); + const inj = r.bfs(most * streams, 2); + const y = r.f32s(most * 11 * dims, -1, 1); + const wts = r.f32s(most * 11, 0, 0.2); + const br = r.f32s(8 * most * dims, -0.5, 0.5); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dinj = try r.dev(std.mem.sliceAsBytes(inj)); + const dy = try r.dev(std.mem.sliceAsBytes(y)); + const dwts = try r.dev(std.mem.sliceAsBytes(wts)); + const dbr = try r.dev(std.mem.sliceAsBytes(br)); + var outs: [counts.len]u64 = undefined; + var psss: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + outs[c] = try r.out(m * wide * 2); + psss[c] = try r.out(m * nc * streams * 4); + try g.hcWriteback(r.s, dh, outs[c], psss[c], dbr, dinj, dy, dwts, most * dims, m, dims, streams, mode, 10, 11, 8); + } + const name = switch (mode) { + 0 => "writeback sums", + 2 => "writeback moe", + else => "writeback slices", + }; + const hout = if (mode == 0) h else try r.back(u16, outs[counts.len - 1], most * wide); + const pss = try r.back(f32, psss[counts.len - 1], most * nc * streams); + for (0..most) |row| { + for (0..dims) |d| { + if (mode == 0) break; + var acc: f64 = 0; + if (mode == 2) { + for (0..11) |k| acc += y[(row * 11 + k) * dims + d] * wts[row * 11 + k]; + } else for (0..8) |k| { + acc += br[k * most * dims + row * dims + d]; + } + const branch = rbf(acc); + for (0..streams) |s| { + const at = row * wide + s * dims + d; + const want = rbf(bf(h[at]) + rbf(branch * bf(inj[row * streams + s]))); + try r.within(name, bf(hout[at]), want, ulps(want, 1)); + } + } + for (0..nc) |c| for (0..streams) |s| { + var t: f64 = 0; + for (0..256) |i| t += bf(hout[row * wide + s * dims + c * 256 + i]) * bf(hout[row * wide + s * dims + c * 256 + i]); + try r.within(name, pss[(row * nc + c) * streams + s], t, sumTol(t, 256)); + }; + } + for (counts[0 .. counts.len - 1], 0..) |m, c| { + if (mode != 0) try r.sameRows(name, u16, outs[c], hout[0 .. m * wide]); + try r.sameRows(name, f32, psss[c], pss[0 .. m * nc * streams]); + } +} + +/// The read-out chain as a pass runs it: down (normed rows, K slices), the slices' sum and activation, the up mix. +fn readout(r: *Rig, g: *const flash.Glue) !void { + const h = r.bfs(most * wide, 1); + const pss = try r.a.alloc(f32, most * nc * streams); + for (0..most) |row| for (0..nc) |c| for (0..streams) |s| { + var t: f32 = 0; + for (0..256) |i| t += @floatCast(bf(h[row * wide + s * dims + c * 256 + i]) * bf(h[row * wide + s * dims + c * 256 + i])); + pss[(row * nc + c) * streams + s] = t; + }; + const scale = r.f32s(wide, 0.5, 1.5); + const down = try tiled(r, n_down, wide, 0.004); + const up = try tiled(r, wide, low, 0.05); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dpss = try r.dev(std.mem.sliceAsBytes(pss)); + const dscale = try r.dev(std.mem.sliceAsBytes(scale)); + var norms: [counts.len]u64 = undefined; + var parts: [counts.len]u64 = undefined; + var acts: [counts.len]u64 = undefined; + var xss: [counts.len]u64 = undefined; + var injs: [counts.len]u64 = undefined; + var mixes: [counts.len]u64 = undefined; + var xsms: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + norms[c] = try r.out(m * wide * 2); + parts[c] = try r.out(slices * m * n_down * 4); + acts[c] = try r.out(m * low * 2); + xss[c] = try r.out(m * low / 32 * 4); + injs[c] = try r.out(m * streams * 2); + mixes[c] = try r.out(m * dims * 2); + xsms[c] = try r.out(m * dims / 32 * 4); + try g.hcDown(r.s, dh, dpss, dscale, norms[c], down.dw, down.ds, down.db, 0, parts[c], eps, m, n_down, wide, dims, nc, streams, slices); + try g.hcReduceAct(r.s, parts[c], acts[c], xss[c], injs[c], slices, m, streams, low, n_down, 1); + try g.hcUpmix(r.s, acts[c], xss[c], up.dw, up.ds, up.db, norms[c], mixes[c], xsms[c], m, wide, low, dims, streams); + } + const last = counts.len - 1; + const normed = try r.back(u16, norms[last], most * wide); + const part = try r.back(f32, parts[last], slices * most * n_down); + const act = try r.back(u16, acts[last], most * low); + const xs = try r.back(f32, xss[last], most * low / 32); + const inj = try r.back(u16, injs[last], most * streams); + const mixed = try r.back(u16, mixes[last], most * dims); + const xsm = try r.back(f32, xsms[last], most * dims / 32); + const xv = try r.a.alloc(f64, wide); + for (0..most) |row| { + for (0..streams) |s| { + var t: f64 = 0; + for (0..nc) |c| t += pss[(row * nc + c) * streams + s]; + const rinv = 1 / @sqrt(t / dims + eps); + for (0..dims) |d| { + const at = s * dims + d; + const want = bf(h[row * wide + at]) * rinv * scale[at]; + try r.within("down normed", bf(normed[row * wide + at]), want, ulps(want, 0.5) + @abs(want) * 0x1p-20); + xv[at] = bf(normed[row * wide + at]); + } + } + // each K slice's sums from the stored normed values + const per = wide / slices; + for (0..slices) |z| for (0..n_down) |col| { + const d = rowDot(down, col, xv[z * per ..][0..per], z * per, per / 32, null); + try r.within("down slices", part[(z * most + row) * n_down + col], d[0], d[1] * 0x1p-17 + 1e-9); + }; + var asum: [low / 32]f64 = @splat(0); + var aabs: [low / 32]f64 = @splat(0); + for (0..n_down) |col| { + var v: f64 = 0; + for (0..slices) |z| v += part[(z * most + row) * n_down + col]; + const vv = rbf(rbf(v) / streams); + if (col < low) { + const want = rbf(vv * rig.sigmoid(vv)); + try r.within("reduce act", bf(act[row * low + col]), want, ulps(want, 1)); + asum[col / 32] += bf(act[row * low + col]); + aabs[col / 32] += @abs(bf(act[row * low + col])); + } else { + const want = rbf(2 * rbf(rig.sigmoid(vv))); + try r.within("reduce inject", bf(inj[row * streams + col - low]), want, ulps(want, 1)); + } + } + for (0..low / 32) |gi| try r.within("reduce act sums", xs[row * (low / 32) + gi], asum[gi], sumTol(aabs[gi], 32)); + // the up mix from the stored activations, their sums and the normed streams + const av = try r.a.alloc(f64, low); + for (av, 0..) |*o, i| o.* = bf(act[row * low + i]); + const xsv = try r.a.alloc(f64, low / 32); + for (xsv, 0..) |*o, i| o.* = xs[row * (low / 32) + i]; + for (0..dims / 32) |j| { + var gsum: f64 = 0; + var gabs: f64 = 0; + for (0..32) |dd| { + const d = j * 32 + dd; + var total: f64 = 0; + var tabs: f64 = 0; + for (0..streams) |s| { + const u = rowDot(up, s * dims + d, av, 0, low / 32, xsv)[0]; + const term = rbf(rbf(rig.sigmoid(rbf(u))) * bf(normed[row * wide + s * dims + d])); + total += term; + tabs += @abs(term); + } + const want = rbf(total / streams); + try r.within("up mix", bf(mixed[row * dims + d]), want, ulps(tabs / streams, 1) + ulps(want, 0.5)); + gsum += bf(mixed[row * dims + d]); + gabs += @abs(bf(mixed[row * dims + d])); + } + try r.within("up mix sums", xsm[row * (dims / 32) + j], gsum, sumTol(gabs, 32)); + } + } + for (counts[0..last], 0..) |m, c| { + try r.sameRows("down normed", u16, norms[c], normed[0 .. m * wide]); + for (0..slices) |z| try r.sameRows("down slices", f32, parts[c] + z * m * n_down * 4, part[z * most * n_down ..][0 .. m * n_down]); + try r.sameRows("reduce act", u16, acts[c], act[0 .. m * low]); + try r.sameRows("reduce act sums", f32, xss[c], xs[0 .. m * low / 32]); + try r.sameRows("reduce inject", u16, injs[c], inj[0 .. m * streams]); + try r.sameRows("up mix", u16, mixes[c], mixed[0 .. m * dims]); + try r.sameRows("up mix sums", f32, xsms[c], xsm[0 .. m * dims / 32]); + } +} + +/// The n-gram gate then its conv over [tail; rows], and the MTP head's stream adds. +fn ngram(r: *Rig, g: *const flash.Glue) !void { + const keys = r.bfs(most * wide, 1); + const vals = r.bfs(most * dims, 1); + const h = r.bfs(most * wide, 1); + const nk = r.f32s(wide, 0.5, 1.5); + const nq = r.f32s(wide, 0.5, 1.5); + const ncw = r.f32s(wide, 0.5, 1.5); + const tail = r.bfs(9 * wide, 1); + const cw = r.bfs(wide * 4, 0.5); + const dk = try r.dev(std.mem.sliceAsBytes(keys)); + const dv = try r.dev(std.mem.sliceAsBytes(vals)); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dnk = try r.dev(std.mem.sliceAsBytes(nk)); + const dnq = try r.dev(std.mem.sliceAsBytes(nq)); + const dnc = try r.dev(std.mem.sliceAsBytes(ncw)); + const dtail = try r.dev(std.mem.sliceAsBytes(tail)); + const dcw = try r.dev(std.mem.sliceAsBytes(cw)); + var gated: [counts.len]u64 = undefined; + var gpss: [counts.len]u64 = undefined; + var houts: [counts.len]u64 = undefined; + var nrows: [counts.len]u64 = undefined; + for (counts, 0..) |m, c| { + gated[c] = try r.out(m * wide * 2); + gpss[c] = try r.out(m * streams * 4); + houts[c] = try r.out(m * wide * 2); + nrows[c] = try r.out(m * wide * 2); + try g.pleGate(r.s, dk, dv, dh, dnk, dnq, gated[c], gpss[c], eps, m, dims, streams); + try g.pleConv(r.s, gated[c], gpss[c], dnc, dtail, dcw, dh, houts[c], nrows[c], eps, m, dims, streams, 4, 3); + } + const last = counts.len - 1; + const gt = try r.back(u16, gated[last], most * wide); + const ps = try r.back(f32, gpss[last], most * streams); + const hout = try r.back(u16, houts[last], most * wide); + const nrow = try r.back(u16, nrows[last], most * wide); + for (0..most) |row| for (0..streams) |s| { + const at = row * wide + s * dims; + var kss: f64 = 0; + var qss: f64 = 0; + for (0..dims) |d| { + kss += bf(keys[at + d]) * bf(keys[at + d]); + qss += bf(h[at + d]) * bf(h[at + d]); + } + const kr = 1 / @sqrt(kss / dims + eps); + const qr = 1 / @sqrt(qss / dims + eps); + var dot: f64 = 0; + for (0..dims) |d| dot += rbf(rbf(bf(keys[at + d]) * kr * nk[s * dims + d]) * rbf(bf(h[at + d]) * qr * nq[s * dims + d])); + var gate = rbf(dot / @sqrt(@as(f64, dims))); + const mag = @sqrt(@max(@abs(gate), 1e-6)); + gate = rbf(if (gate > 0) mag else if (gate < 0) -mag else 0); + const sg = rbf(rig.sigmoid(gate)); + var gss: f64 = 0; + for (0..dims) |d| { + const want = rbf(sg * bf(vals[row * dims + d])); + try r.within("ngram gate", bf(gt[at + d]), want, ulps(want, 1)); + gss += bf(gt[at + d]) * bf(gt[at + d]); + } + try r.within("ngram gate sums", ps[row * streams + s], gss, sumTol(gss, dims)); + }; + // the conv of row r reads rows r - 9 + 3t, taken from the tail before the chunk + for (0..most) |row| for (0..wide) |e| { + const s = e / dims; + var acc: f64 = 0; + for (0..4) |t| { + const at = row + 3 * t; + const xv = if (at < 9) bf(tail[at * wide + e]) else blk: { + const rr = at - 9; + const rinv = 1 / @sqrt(@as(f64, ps[rr * streams + s]) / dims + eps); + break :blk rbf(bf(gt[rr * wide + e]) * rinv * ncw[e]); + }; + acc += bf(cw[e * 4 + t]) * xv; + } + const conv = rbf(acc); + const ple = rbf(bf(gt[row * wide + e]) + rbf(conv * rig.sigmoid(conv))); + const want = rbf(bf(h[row * wide + e]) + ple); + try r.within("ngram conv", bf(hout[row * wide + e]), want, ulps(ple, 1) + ulps(want, 0.5)); + const rinv = 1 / @sqrt(@as(f64, ps[row * streams + s]) / dims + eps); + const own = rbf(bf(gt[row * wide + e]) * rinv * ncw[e]); + try r.within("ngram conv rows", bf(nrow[row * wide + e]), own, ulps(own, 1)); + }; + for (counts[0..last], 0..) |m, c| { + try r.sameRows("ngram gate", u16, gated[c], gt[0 .. m * wide]); + try r.sameRows("ngram gate sums", f32, gpss[c], ps[0 .. m * streams]); + try r.sameRows("ngram conv rows", u16, nrows[c], nrow[0 .. m * wide]); + } + // the MTP head's input streams: bf16(e + h), each sum exact before its one rounding + const e = r.bfs(most * dims, 1); + const out = try r.out(most * wide * 2); + try g.addStreams(r.s, try r.dev(std.mem.sliceAsBytes(e)), dh, out, most, dims, streams); + const got = try r.back(u16, out, most * wide); + for (0..most) |row| for (0..wide) |i| { + const want = rbf(bf(e[row * dims + i % dims]) + bf(h[row * wide + i])); + try check.expect(bf(got[row * wide + i]) == want, "add streams: row {d} value {d} is the rounded exact sum", .{ row, i }); + }; +} + +pub fn run(gpu: check.Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(53); + var r = try Rig.init(gpu, arena.allocator(), prng.random()); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + for ([_]usize{ 0, 2, 4 }) |mode| try writeback(&r, &g, mode); + try readout(&r, &g); + try ngram(&r, &g); + r.report("flashnext-hc: write-backs, the read-out chain, the n-gram gate and conv, stream adds"); +} diff --git a/zig/tests/cuda/flashnext_rig.zig b/zig/tests/cuda/flashnext_rig.zig new file mode 100644 index 0000000000..9d917f5a1c --- /dev/null +++ b/zig/tests/cuda/flashnext_rig.zig @@ -0,0 +1,119 @@ +//! The Flash Next glue tests' shared rig: device buffers, bf16 values, float64 rounding points and tolerances. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); + +pub fn bf16(x: f32) u16 { + const u: u32 = @bitCast(x); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +pub fn bf(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +/// A float64 value at a kernel's bf16 rounding point (through fp32, as the kernel rounds). +pub fn rbf(x: f64) f64 { + return bf(bf16(@floatCast(x))); +} + +pub fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} + +/// `n` bf16 ulps of |want| (one ulp: 2^-7 of it), with a floor for zeros. +pub fn ulps(want: f64, n: f64) f64 { + return n * @abs(want) * std.math.pow(f64, 2, -7) + 1e-6; +} + +/// An fp32 sum's tolerance: `terms` roundings of at most 2^-24 of the sum of |terms|. +pub fn sumTol(abs_sum: f64, terms: usize) f64 { + return abs_sum * @as(f64, @floatFromInt(terms)) * std.math.pow(f64, 2, -24) + 1e-30; +} + +pub const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + rng: std.Random, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + /// The largest |got - want| / tolerance seen, by check name. + worst: std.StringArrayHashMapUnmanaged(f64) = .empty, + + pub fn init(gpu: check.Gpu, a: std.mem.Allocator, rng: std.Random) !Rig { + return .{ .gpu = gpu, .a = a, .s = try cuda.Stream.init(gpu.d, false), .rng = rng }; + } + + pub fn deinit(r: *Rig) void { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + r.s.deinit(); + } + + pub fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + return b.ptr; + } + + /// A device buffer of `bytes`, every byte 0xff (a value the kernel misses reads NaN). + pub fn out(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(bytes, 16)); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0xff, r.s.handle); + return b.ptr; + } + + pub fn back(r: *Rig, comptime T: type, ptr: u64, n: usize) ![]T { + const o = try r.a.alloc(T, n); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(o)); + return o; + } + + /// bf16 values within [-scale, scale). + pub fn bfs(r: *Rig, n: usize, scale: f32) []u16 { + const v = r.a.alloc(u16, n) catch unreachable; + for (v) |*x| x.* = bf16((r.rng.float(f32) * 2 - 1) * scale); + return v; + } + + /// fp32 values within [lo, hi). + pub fn f32s(r: *Rig, n: usize, lo: f32, hi: f32) []f32 { + const v = r.a.alloc(f32, n) catch unreachable; + for (v) |*x| x.* = lo + r.rng.float(f32) * (hi - lo); + return v; + } + + pub fn words(r: *Rig, n: usize) []u32 { + const v = r.a.alloc(u32, n) catch unreachable; + for (v) |*x| x.* = r.rng.int(u32); + return v; + } + + pub fn within(r: *Rig, what: []const u8, got: f64, want: f64, tol: f64) !void { + const q = @abs(got - want) / tol; + if (!(q <= 1)) try check.expect(false, "{s}: within float64 tolerance (got {d}, want {d}, {d:.3} tolerances)", .{ what, got, want, q }); + const gop = try r.worst.getOrPut(r.a, what); + if (!gop.found_existing) gop.value_ptr.* = 0; + gop.value_ptr.* = @max(gop.value_ptr.*, q); + } + + /// Bytes of a launch over fewer rows against the same rows of the widest launch. + pub fn sameRows(r: *Rig, what: []const u8, comptime T: type, ptr: u64, want: []const T) !void { + try check.sameBytes(what, std.mem.sliceAsBytes(try r.back(T, ptr, want.len)), std.mem.sliceAsBytes(want)); + } + + pub fn report(r: *Rig, what: []const u8) void { + var it = r.worst.iterator(); + var line: std.ArrayList(u8) = .empty; + while (it.next()) |e| line.print(r.a, " {s} {d:.3};", .{ e.key_ptr.*, e.value_ptr.* }) catch {}; + check.pass("{s}: every row count's bits equal; worst tolerances used:{s}", .{ what, line.items }); + } +}; + +/// q (4 bits) of an MLX row's value i: word i / 8, nibble i % 8. +pub fn nibble(w: []const u32, row_words: usize, row: usize, i: usize) f64 { + return @floatFromInt((w[row * row_words + i / 8] >> @intCast(4 * (i % 8))) & 0xF); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index d836d45502..6fa70d0753 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -31,6 +31,8 @@ const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); +const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); +const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); const usage = \\usage: tf-cuda-test @@ -54,6 +56,8 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) + \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) \\ window-profile MODEL WIDTHS IDS_FILE a decode window's GPU ms by kernel class (WIDTHS like 1,4,16), beside its graph \\ tree-accept MODEL IDS_FILE COUNT OUT a greedy reply's MTP chain and level-1 runner-up branches at every position (JSON) \\ chunk-costs MODEL IDS_FILE WIDTHS r rows as the decode window graph (r <= 16) and as a prompt chunk (any r) @@ -137,6 +141,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); if (std.mem.eql(u8, cmd, "pocket")) return pocket_tests.run(gpu, if (rest.len > 0) try std.fmt.parseInt(usize, rest[0], 10) else 0); if (std.mem.eql(u8, cmd, "window-profile")) return window_profile.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); From 38c9c63e113804bd1f81b31afc4a351bf445a145 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:14:38 -0500 Subject: [PATCH 045/115] Flash Next CUDA boots on its own glue; a captured set that holds every entry point is used for comparison The engine no longer needs TENSORFOLD_CUDA_KERNELS: every launch Tri made from a captured cubin now runs our kernels (cuda_glue.zig) unless a set is loaded. A set given by TENSORFOLD_CUDA_KERNELS, or found beside the binary, replaces them only when it holds all 15 entry points the pass launches (17 with the MTP head); one without them, Nemotron's say, is passed over with a log line, and the startup line says "own glue kernels" (native/cuda.zig asks the family's Loaded.own_glue; aot.Set.has names a set's functions). flashnext-prompt takes native in place of a set's folder. flashnext-compare loads the engine once and runs a prompt on our glue and on a captured set: every prompt row's scores (top-1 agreement, the largest score difference, KL, each path's perplexity on the prompt's own next tokens), then each path's greedy reply and the token where they part. --- zig/build/cuda.zig | 7 + zig/src/cuda/aot.zig | 6 + .../families/flashnext/cuda_compare_main.zig | 209 ++++++++++++++++++ zig/src/families/flashnext/cuda_lanes.zig | 45 +++- zig/src/families/flashnext/cuda_launch.zig | 30 ++- .../families/flashnext/cuda_prompt_main.zig | 18 +- zig/src/native/cuda.zig | 9 +- 7 files changed, 302 insertions(+), 22 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_compare_main.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 6aca30d30c..d51459bd29 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -172,6 +172,13 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const prompt_count = b.option([]const u8, "flashnext_count", "greedy tokens to emit") orelse "8"; if (embed_model != null and embed_kernels != null and prompt_ids != null) prompt_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.?, prompt_count }); b.step("flashnext-prompt", "Run one Flash Next prompt on the captured CUDA kernels").dependOn(&prompt_run.step); + // not installed: our own glue against a captured set on one loaded engine (scores of every prompt row, replies) + const compare_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_compare_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); + compare_check.addImport("cuda", cuda); + compare_check.addImport("flashnext", mods.flashnext); + const compare_run = b.addRunArtifact(b.addExecutable(.{ .name = "flashnext-compare", .root_module = compare_check })); + if (embed_model != null and embed_kernels != null and prompt_ids != null) compare_run.addArgs(&.{ embed_model.?, embed_kernels.?, prompt_ids.? }); + b.step("flashnext-compare", "Our Flash Next glue against a captured set: every prompt row's scores, perplexity, greedy replies").dependOn(&compare_run.step); // not installed: the native block selection against Triton's lists (capture_sparse.sh's fixtures) const select_check = b.createModule(.{ .root_source_file = b.path("zig/src/families/flashnext/cuda_select_main.zig"), .target = target, .optimize = optimize, .link_libc = true }); select_check.addImport("cuda", cuda); diff --git a/zig/src/cuda/aot.zig b/zig/src/cuda/aot.zig index 64d436f1b0..c957cc9601 100644 --- a/zig/src/cuda/aot.zig +++ b/zig/src/cuda/aot.zig @@ -139,6 +139,12 @@ pub const Set = struct { return error.MissingTritonVariant; } + /// Whether the set holds any variant of `function`. + pub fn has(self: *const Set, function: []const u8) bool { + for (self.variants) |v| if (std.mem.eql(u8, v.spec.@"fn", function)) return true; + return false; + } + /// The smallest value of constexpr `name` at or above `at_least` among `function`'s variants. pub fn smallestConst(self: *const Set, function: []const u8, name: []const u8, at_least: i64) ?i64 { var best: ?i64 = null; diff --git a/zig/src/families/flashnext/cuda_compare_main.zig b/zig/src/families/flashnext/cuda_compare_main.zig new file mode 100644 index 0000000000..c156336d1e --- /dev/null +++ b/zig/src/families/flashnext/cuda_compare_main.zig @@ -0,0 +1,209 @@ +//! Our own glue against a captured set on one loaded engine: every prompt row's scores, then each path's greedy reply. + +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); + +const prompt = flash.prompt; +const vocab = prompt.vocabulary; + +/// Prompt rows a teacher-forced pass takes at once (every row's scores come back). +const rows_at_once = 256; + +fn score(raw: u16) f64 { + return @as(f32, @bitCast(@as(u32, raw) << 16)); +} + +/// What two paths' scores over the same rows say: top-1 agreement, the largest score gap, KL and each path's NLL. +const Tally = struct { + rows: usize = 0, + agree: usize = 0, + max_abs: f64 = 0, + kl: f64 = 0, + kl_max: f64 = 0, + nll_a: f64 = 0, + nll_b: f64 = 0, + scored: usize = 0, + /// Among rows whose top-1 differs, the largest gap between the captured path's best two scores. + tie_gap: f64 = 0, + /// Each row's KL, for its median and the share of rows past 0.1, and each scored row's NLL own minus captured. + kls: std.ArrayList(f64) = .empty, + dnll: std.ArrayList(f64) = .empty, + + fn row(t: *Tally, a: []const u16, b: []const u16, next: ?u32) void { + var best_a: usize = 0; + var best_b: usize = 0; + var max_a = -std.math.inf(f64); + var max_b = -std.math.inf(f64); + for (a, b, 0..) |ra, rb, i| { + const va = score(ra); + const vb = score(rb); + t.max_abs = @max(t.max_abs, @abs(va - vb)); + if (va > max_a) { + max_a = va; + best_a = i; + } + if (vb > max_b) { + max_b = vb; + best_b = i; + } + } + var sum_a: f64 = 0; + var sum_b: f64 = 0; + for (a, b) |ra, rb| { + sum_a += @exp(score(ra) - max_a); + sum_b += @exp(score(rb) - max_b); + } + const lse_a = max_a + @log(sum_a); + const lse_b = max_b + @log(sum_b); + var kl: f64 = 0; + for (a, b) |ra, rb| { + const la = score(ra) - lse_a; + kl += @exp(la) * (la - (score(rb) - lse_b)); + } + t.rows += 1; + t.kl += kl; + t.kls.append(std.heap.page_allocator, kl) catch {}; + t.kl_max = @max(t.kl_max, kl); + if (best_a == best_b) { + t.agree += 1; + } else { + var second = -std.math.inf(f64); + for (a, 0..) |ra, i| if (i != best_a) { + second = @max(second, score(ra)); + }; + t.tie_gap = @max(t.tie_gap, max_a - second); + } + if (next) |n| { + t.nll_a += lse_a - score(a[n]); + t.nll_b += lse_b - score(b[n]); + t.dnll.append(std.heap.page_allocator, (lse_b - score(b[n])) - (lse_a - score(a[n]))) catch {}; + t.scored += 1; + } + } +}; + +/// One line of a tally: agreement, score gaps, KL (mean, median, max, rows past 0.1) and each path's perplexity. +fn report(what: []const u8, t: *Tally) void { + if (t.rows == 0) return; + const n: f64 = @floatFromInt(t.rows); + const s: f64 = @floatFromInt(@max(t.scored, 1)); + std.mem.sort(f64, t.kls.items, {}, std.sort.asc(f64)); + std.mem.sort(f64, t.dnll.items, {}, std.sort.asc(f64)); + var past: usize = 0; + for (t.kls.items) |v| past += @intFromBool(v > 0.1); + const mid = if (t.dnll.items.len > 0) t.dnll.items[t.dnll.items.len / 2] else 0; + std.debug.print("{s}: {d} rows, top-1 agree {d}/{d}, max |score diff| {d:.4}, KL(captured||own) mean {e:.3} median {e:.3} max {e:.3}, {d} rows past 0.1, perplexity captured {d:.4} own {d:.4} over {d} next tokens (median NLL own - captured {e:.3}), largest top-2 gap where top-1 differs {d:.4}\n", .{ what, t.rows, t.agree, t.rows, t.max_abs, t.kl / n, t.kls.items[t.kls.items.len / 2], t.kl_max, past, @exp(t.nll_a / s), @exp(t.nll_b / s), t.scored, mid, t.tie_gap }); +} + +/// Rows from this position on attend through the indexer's picks (512 blocks of 4 keys). +const sparse_from = 2048; + +fn scores(raw: []const u8) []const u16 { + return @alignCast(std.mem.bytesAsSlice(u16, raw)); +} + +fn readIds(gpa: std.mem.Allocator, io: std.Io, path: []const u8) ![]u32 { + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(text); + var ids: std.ArrayList(u32) = .empty; + errdefer ids.deinit(gpa); + var it = std.mem.tokenizeAny(u8, text, ", \t\r\n"); + while (it.next()) |part| try ids.append(gpa, try std.fmt.parseInt(u32, part, 10)); + return ids.toOwnedSlice(gpa); +} + +fn envCount(name: [:0]const u8, default: usize) usize { + const v = std.c.getenv(name) orelse return default; + return std.fmt.parseInt(usize, std.mem.span(v), 10) catch default; +} + +pub fn main(init: std.process.Init) !u8 { + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 4) { + std.debug.print("usage: flashnext-compare ... (COUNT greedy tokens, default 32)\n", .{}); + return 2; + } + const count = envCount("COUNT", 32); + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer set.deinit(); + var glue = try flash.Glue.load(&driver); + defer glue.deinit(); + var stream = try cuda.Stream.init(&driver, true); + defer stream.deinit(); + const captured: flash.Tri = .{ .set = &set, .s = stream }; + const own: flash.Tri = .{ .own = &glue, .s = stream }; + var e = try prompt.Engine.init(gpa, io, &driver, &stream, args[1]); + defer e.deinit(); + var out_a = try cuda.DeviceBuffer.alloc(&driver, rows_at_once * vocab * 2); + defer out_a.free(); + var out_b = try cuda.DeviceBuffer.alloc(&driver, rows_at_once * vocab * 2); + defer out_b.free(); + const host_a = try gpa.alloc(u16, rows_at_once * vocab); + defer gpa.free(host_a); + const host_b = try gpa.alloc(u16, rows_at_once * vocab); + defer gpa.free(host_b); + for (args[3..]) |path| { + const ids = try readIds(gpa, io, path); + defer gpa.free(ids); + var seq_a = try e.newSeq(ids.len + count); + defer seq_a.deinit(gpa); + var seq_b = try e.newSeq(ids.len + count); + defer seq_b.deinit(gpa); + // every prompt row's scores on both paths, the next prompt token's NLL under each; dense rows and sparse ones + var t: Tally = .{}; + var dense: Tally = .{}; + var sparse: Tally = .{}; + var at: usize = 0; + while (at < ids.len) { + const m = @min(rows_at_once, ids.len - at); + try e.forward(flash.Tri, captured, &seq_a, ids[at..][0..m], .{ .logits = .all, .logits_out = out_a.ptr }); + try e.forward(flash.Tri, own, &seq_b, ids[at..][0..m], .{ .logits = .all, .logits_out = out_b.ptr }); + try stream.synchronize(); + try out_a.download(0, std.mem.sliceAsBytes(host_a[0 .. m * vocab])); + try out_b.download(0, std.mem.sliceAsBytes(host_b[0 .. m * vocab])); + for (0..m) |r| { + const next: ?u32 = if (at + r + 1 < ids.len) ids[at + r + 1] else null; + t.row(host_a[r * vocab ..][0..vocab], host_b[r * vocab ..][0..vocab], next); + if (ids.len > sparse_from) (if (at + r < sparse_from) &dense else &sparse).row(host_a[r * vocab ..][0..vocab], host_b[r * vocab ..][0..vocab], next); + } + at += m; + } + var label: [96]u8 = undefined; + report(try std.fmt.bufPrint(&label, "rows {s}", .{std.fs.path.basename(path)}), &t); + report(try std.fmt.bufPrint(&label, "rows {s} before {d}", .{ std.fs.path.basename(path), sparse_from }), &dense); + report(try std.fmt.bufPrint(&label, "rows {s} from {d}", .{ std.fs.path.basename(path), sparse_from }), &sparse); + // each path's own greedy reply from the prompt, and where they part + const reply_a = try gpa.alloc(u32, count); + defer gpa.free(reply_a); + const reply_b = try gpa.alloc(u32, count); + defer gpa.free(reply_b); + var g: Tally = .{}; + var parted: ?usize = null; + for (0..count) |i| { + reply_a[i] = prompt.argmax(seq_a.logits); + reply_b[i] = prompt.argmax(seq_b.logits); + // the same tokens so far: the same context on both paths, until (and at) the token where they part + if (parted == null) g.row(scores(seq_a.logits), scores(seq_b.logits), null); + if (parted == null and reply_a[i] != reply_b[i]) parted = i; + if (i + 1 == count) break; + try e.predict(flash.Tri, captured, &seq_a, reply_a[i..][0..1]); + try e.predict(flash.Tri, own, &seq_b, reply_b[i..][0..1]); + } + std.debug.print("reply {s}: {d} greedy tokens, ", .{ std.fs.path.basename(path), count }); + if (parted) |p| std.debug.print("part at token {d} (top-2 gap on the captured path there {d:.4})", .{ p, g.tie_gap }) else std.debug.print("all equal", .{}); + std.debug.print(", KL up to there mean {e:.3}\n", .{g.kl / @as(f64, @floatFromInt(@max(g.rows, 1)))}); + std.debug.print("tokens captured", .{}); + for (reply_a) |v| std.debug.print(" {d}", .{v}); + std.debug.print("\ntokens own", .{}); + for (reply_b) |v| std.debug.print(" {d}", .{v}); + std.debug.print("\n", .{}); + } + return 0; +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 9b97a94081..ba109285c6 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -7,13 +7,18 @@ const prompt = @import("cuda_prompt.zig"); const weights = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); const snap = @import("cuda_snap.zig"); +const Glue = @import("cuda_glue.zig").Glue; const be = lanes.backend; const ring = 1024; /// Streams a round takes at most (a round with more live streams takes them in turns). const round_streams = 32; -/// `Tri` is the captured-kernel launcher. `open` loads the checkpoint and serves greedy or keyed draws. +/// The Triton entry points the pass launches and the head's two more: a captured set replaces our glue with all. +pub const captured_entries = [_][]const u8{ "_embed", "_hc_writeback", "_qmm_hcdown", "_hc_reduce_act", "_qmm_upmix", "_ple_gate", "_ple_conv", "_attn_prep", "_attn_gate", "_pool", "_scores", "_chunks", "_merge", "_router", "_topk_rows" }; +pub const captured_head_entries = [_][]const u8{ "_rmsnorm", "_add_streams" }; + +/// `Tri` launches our glue, or a captured set's. `open` loads the checkpoint and serves greedy or keyed draws. pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; @@ -57,6 +62,8 @@ pub fn Namespace(comptime Tri: type) type { learn_bytes: usize = 0, /// What the prompt cache's copies are taken from: the backend itself (snap_* look up the owner stream). target: *anyopaque, + /// Whether our own glue runs (no captured set given, or one without this family's entry points). + own_glue: bool = true, }; /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. @@ -72,7 +79,6 @@ pub fn Namespace(comptime Tri: type) type { error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", error.UnsupportedCheckpoint => "Flash Next on CUDA serves MLX affine 4-bit group-32 Qwen3.8-Flash-Next checkpoints; the log names the config field that differs", error.LearnedPocket => "this folder holds Living Weights pockets, which Flash Next on CUDA does not serve yet", - error.CapturedKernelsRequired => "Flash Next on CUDA replays captured Triton kernels: set TENSORFOLD_CUDA_KERNELS to the set's folder", else => null, }; } @@ -95,7 +101,9 @@ pub fn Namespace(comptime Tri: type) type { head: ?*mtp.Head = null, depth: usize = mtp.max_chain, confidence: f64 = mtp.default_confidence, - set: cuda.aot.Set, + /// A captured Triton set (TENSORFOLD_CUDA_KERNELS) for comparison, else null: our own glue. + set: ?cuda.aot.Set = null, + glue: Glue, stream: cuda.Stream, lanes: std.AutoHashMapUnmanaged(*lanes.Stream, Lane) = .empty, /// A round's scores, every row [round_streams, vocab] bf16. @@ -104,7 +112,7 @@ pub fn Namespace(comptime Tri: type) type { next: u64 = 0, fn tri(self: *Cuda) Tri { - return .{ .set = &self.set, .s = self.stream }; + return .{ .set = if (self.set) |*x| x else null, .own = &self.glue, .s = self.stream }; } fn ops(self: *Cuda) Ops { @@ -241,17 +249,18 @@ pub fn Namespace(comptime Tri: type) type { pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: Options) !Loaded { if (o.context > prompt.max_tokens) return error.ContextOverWindow; if (o.segments != 1 or o.slide) return error.UnsupportedOption; - // every launch past our own kernels is a captured Triton cubin: TENSORFOLD_CUDA_KERNELS names the set - const kernels = captured orelse return error.CapturedKernelsRequired; try @import("cuda_config.zig").check(gpa, io, dir); const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); - own.* = .{ .gpa = gpa, .context = o.context, .engine = undefined, .set = undefined, .stream = undefined, .logits = undefined }; + own.* = .{ .gpa = gpa, .context = o.context, .engine = undefined, .glue = undefined, .stream = undefined, .logits = undefined }; own.stream = try cuda.Stream.init(driver, true); errdefer own.stream.deinit(); - own.set = try cuda.aot.Set.load(gpa, io, driver, ctx.device, kernels); - errdefer own.set.deinit(); + own.glue = try Glue.load(driver); + errdefer own.glue.deinit(); + // a captured set replaces our glue, for comparison, when it holds every entry point this engine launches + if (captured) |kernels| own.set = try comparison(gpa, io, driver, ctx, kernels, o.drafts); + errdefer if (own.set) |*x| x.deinit(); own.engine = try prompt.Engine.init(gpa, io, driver, &own.stream, dir); errdefer own.engine.deinit(); // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) @@ -280,9 +289,24 @@ pub fn Namespace(comptime Tri: type) type { .ctx = own, .deinit = release, .target = own, + .own_glue = own.set == null, }; } + /// The captured set at `dir` when it holds every entry point (the head's too when drafting), else null. + fn comparison(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, ctx: *const cuda.Context, dir: []const u8, drafts: bool) !?cuda.aot.Set { + var set = try cuda.aot.Set.load(gpa, io, driver, ctx.device, dir); + const holds = for (captured_entries) |name| { + if (!set.has(name)) break false; + } else for (captured_head_entries) |name| { + if (drafts and !set.has(name)) break false; + } else true; + if (holds) return set; + set.deinit(); + std.log.info("flashnext: the captured set at {s} lacks Flash Next's kernels; its own glue runs", .{dir}); + return null; + } + fn release(p: *anyopaque) void { const own: *Cuda = @ptrCast(@alignCast(p)); var it = own.lanes.valueIterator(); @@ -291,7 +315,8 @@ pub fn Namespace(comptime Tri: type) type { own.logits.free(); if (own.head) |h| h.deinit(); own.engine.deinit(); - own.set.deinit(); + if (own.set) |*x| x.deinit(); + own.glue.deinit(); own.stream.deinit(); own.gpa.destroy(own); } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index e6de83ba1e..020fa63e69 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -1,4 +1,4 @@ -//! Flash Next's Triton launches from the captured cubins: each grid and constexpr matches the Python wrapper. +//! Flash Next's glue launches: our own kernels (cuda_glue.zig), or the captured Triton cubins where a set is loaded. const std = @import("std"); const cuda = @import("cuda"); @@ -51,16 +51,23 @@ pub fn rmsBlock(group: usize) usize { return std.math.ceilPowerOfTwo(usize, group) catch unreachable; } +/// Each launch's grid and constexprs match the Python wrapper's when a captured set is loaded; without one, ours run. pub const Tri = struct { - set: *const aot.Set, + set: ?*const aot.Set = null, + own: ?*const Glue = null, s: cuda.Stream, fn run(t: Tri, name: []const u8, grid: [3]u32, args: []const aot.Arg, consts: []const aot.Const) !void { - try t.set.run(t.s, name, grid, args, consts); + try t.set.?.run(t.s, name, grid, args, consts); + } + + fn glue(t: Tri) ?*const Glue { + return if (t.set == null) t.own.? else null; } /// glue.embed: ids (R,) int32 into (R, copies * dims) bf16, MLX 4-bit group 32. pub fn embed(t: Tri, ids: u64, weight: u64, scales: u64, biases: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + if (t.glue()) |g| return g.embed(t.s, ids, weight, scales, biases, out, dims, copies, rows); try t.run("_embed", embedGrid(rows, dims), &.{ p("IDS", "*i32", ids), p("W", "*i32", weight), @@ -72,6 +79,7 @@ pub const Tri = struct { /// glue.rmsnorm per row and `group` features: bf16 out [rows, dims] and 32-group sums; x rows `x_stride` apart. pub fn rmsnorm(t: Tri, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { + if (t.glue()) |g| return g.rmsnorm(t.s, x, w, out, xs, eps, x_stride, rows, dims, group); try t.run("_rmsnorm", .{ u(rows), u(dims / group), 1 }, &.{ p("X", "*bf16", x), p("W", "*fp32", w), @@ -84,6 +92,7 @@ pub const Tri = struct { /// glue.add_streams: out[r, s] = bf16(e[r] + hs[r, s]), the MTP head's input streams. Program (row, 256 columns). pub fn addStreams(t: Tri, e: u64, hs: u64, out: u64, rows: usize, dims: usize, streams: usize) !void { + if (t.glue()) |g| return g.addStreams(t.s, e, hs, out, rows, dims, streams); try t.run("_add_streams", .{ u(rows), u(dims / 256), 1 }, &.{ p("E", "*bf16", e), p("HS", "*bf16", hs), @@ -93,6 +102,7 @@ pub const Tri = struct { /// glue.attn_gate: one program per row and per query head. pub fn attnGate(t: Tri, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { + if (t.glue()) |g| return g.attnGate(t.s, o, proj, out, xs, rows, pair_width, q_heads, head_dim); try t.run("_attn_gate", .{ u(rows), u(q_heads), 1 }, &.{ p("O", "*bf16", o), p("P", "*bf16", proj), @@ -103,6 +113,7 @@ pub const Tri = struct { /// glue.attn_prep for text positions, bf16 KV (BITS 0, the scale pointers unread); `length` unused at MODE 0. pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { + if (t.glue()) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows); try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ p("P", "*bf16", proj), p("POS0", "*i32", pos), @@ -126,6 +137,7 @@ pub const Tri = struct { /// attention._chunks over 512-key chunks, bf16 KV; `qsa` the sparse form (NCH 5, QSA 1), else dense (NCH 3, QSA 0). pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { + if (t.glue()) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa); try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ p("Q", "*bf16", q), p("KC", "*bf16", kc), @@ -144,6 +156,7 @@ pub const Tri = struct { /// attention._merge: one program per row and per KV head, in attnChunks' form. pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { + if (t.glue()) |g| return g.attnMerge(t.s, po, pm, pl, pos, out, nk, sparse, rows, kv, qsa); try t.run("_merge", .{ u(rows), u(kv), 1 }, &.{ p("PO", "*fp32", po), p("PM", "*fp32", pm), @@ -157,6 +170,7 @@ pub const Tri = struct { /// attention.qsa_pool for text rows (MODE 0): every complete block of 4 indexer keys pooled, normed and roped. pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { + if (t.glue()) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows); try t.run("_pool", .{ u(rows / 4 + 2), 1, 1 }, &.{ p("IKC", "*bf16", ikc), p("POOLED", "*bf16", pooled), @@ -171,6 +185,7 @@ pub const Tri = struct { /// attention._scores: rows' complete blocks scored; `nb` is the row stride (a multiple of 16), `blocks` the grid. pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { + if (t.glue()) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks); try t.run("_scores", .{ u(rows), cdiv(blocks, 64), 1 }, &.{ p("IQ", "*bf16", iq), p("POOLED", "*bf16", pooled), @@ -182,6 +197,7 @@ pub const Tri = struct { /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. pub fn hcReduceAct(t: Tri, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { + if (t.glue()) |g| return g.hcReduceAct(t.s, part, act, xs, inj, sk, rows, streams, low, ndn, has_inj); try t.run("_hc_reduce_act", .{ u(rows), 1, 1 }, &.{ p("PART", "*fp32", part), p("ACT", "*bf16", act), @@ -193,6 +209,7 @@ pub const Tri = struct { /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. pub fn hcUpmix(t: Tri, act: u64, xs: u64, weight: u64, scales: u64, biases: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { + if (t.glue()) |g| return g.hcUpmix(t.s, act, xs, weight, scales, biases, normed, mixed, xs_mixed, rows, n, k, dims, streams); try t.run("_qmm_upmix", .{ cdiv(rows, 16), u(dims / 32), 1 }, &.{ p("X", "*bf16", act), p("XS", "*fp32", xs), @@ -208,6 +225,7 @@ pub const Tri = struct { /// qmm.hc_down for one decode row. N=324 is the layer down stacked with inject. M folds when it is 1. pub fn hcDown(t: Tri, h: u64, pss: u64, scale: u64, normed: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { + if (t.glue()) |g| return g.hcDown(t.s, h, pss, scale, normed, weight, scales, biases, out, part, eps, rows, n, k, dims, nc, streams, sk); try t.run("_qmm_hcdown", .{ cdiv(rows, 16), cdiv(n, 64), u(sk) }, &.{ p("H", "*bf16", h), p("PSS", "*fp32", pss), @@ -225,6 +243,7 @@ pub const Tri = struct { /// glue.hc_writeback. Pointer types follow the tensor dtypes Triton specialized (mode 0..4). pub fn hcWriteback(t: Tri, h: u64, hout: u64, pss: u64, branch: u64, branch_ty: []const u8, inj: u64, y: u64, y_ty: []const u8, wts: u64, wts_ty: []const u8, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { + if (t.glue()) |g| return g.hcWriteback(t.s, h, hout, pss, branch, inj, y, wts, rs, rows, dims, streams, mode, topk, slots, world); try t.run("_hc_writeback", .{ u(rows), u(dims / 256), 1 }, &.{ p("H", "*bf16", h), p("HOUT", "*bf16", hout), @@ -239,6 +258,7 @@ pub const Tri = struct { /// glue.ple_embed: gathered n-gram rows, MLX 4-bit group 32. `scale` is a constexpr. pub fn pleEmbed(t: Tri, weight: u64, scales: u64, biases: u64, out: u64, xs: u64, rows: usize, heads: usize, dh: usize, scale: f32) !void { + if (t.glue() != null) return error.CapturedKernelsRequired; // the row path's n-gram check only try t.run("_ple_embed", .{ u(rows), u(heads), 1 }, &.{ p("W", "*i32", weight), p("S", "*bf16", scales), @@ -250,6 +270,7 @@ pub const Tri = struct { /// glue.ple_gate: one program per row. pub fn pleGate(t: Tri, keys: u64, vals: u64, h: u64, norm_key: u64, norm_query: u64, gated: u64, pss: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + if (t.glue()) |g| return g.pleGate(t.s, keys, vals, h, norm_key, norm_query, gated, pss, eps, rows, dims, streams); try t.run("_ple_gate", .{ u(rows), 1, 1 }, &.{ p("KEYS", "*bf16", keys), p("VALS", "*bf16", vals), @@ -264,6 +285,7 @@ pub const Tri = struct { /// glue.ple_conv. R folds to a constexpr when it is 1. pub fn pleConv(t: Tri, gated: u64, pss: u64, norm_conv: u64, tail: u64, conv_w: u64, h: u64, hout: u64, nrow: u64, eps: f32, rows: usize, dims: usize, streams: usize, taps: usize, dilation: usize) !void { + if (t.glue()) |g| return g.pleConv(t.s, gated, pss, norm_conv, tail, conv_w, h, hout, nrow, eps, rows, dims, streams, taps, dilation); try t.run("_ple_conv", .{ u(rows), u(streams * dims / 512), 1 }, &.{ p("GATED", "*bf16", gated), p("PSS", "*fp32", pss), @@ -290,6 +312,7 @@ pub const Tri = struct { /// moe._topk_rows over `rows` rows of [rows, 513] logits: picks and weights [rows, 11]. pub fn topkMany(t: Tri, logits: u64, pick: u64, wts: u64, rows: usize) !void { + if (t.glue()) |g| return g.topkMany(t.s, logits, pick, wts, rows); try t.run("_topk_rows", .{ u(rows), 1, 1 }, &.{ p("L", "*fp32", logits), p("PICK", "*i32", pick), @@ -299,6 +322,7 @@ pub const Tri = struct { /// moe.router over `rows` rows at the one-row tiles (BM 16, BLOCK_E 32, BK 256), each row keeping its own bits. pub fn routerMany(t: Tri, x: u64, w: u64, out: u64, x_stride: usize, rows: usize) !void { + if (t.glue()) |g| return g.routerMany(t.s, x, w, out, x_stride, rows); try t.run("_router", .{ cdiv(rows, 16), cdiv(513, 32), 1 }, &.{ p("X", "*bf16", x), p("W", "*bf16", w), diff --git a/zig/src/families/flashnext/cuda_prompt_main.zig b/zig/src/families/flashnext/cuda_prompt_main.zig index c5fe7157bb..9ba52f395f 100644 --- a/zig/src/families/flashnext/cuda_prompt_main.zig +++ b/zig/src/families/flashnext/cuda_prompt_main.zig @@ -11,7 +11,7 @@ pub fn main(init: std.process.Init) !u8 { const io = init.io; const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len < 5) { - std.debug.print("usage: flashnext-prompt \n", .{}); + std.debug.print("usage: flashnext-prompt \n", .{}); return 2; } const count = try std.fmt.parseInt(usize, args[4], 10); @@ -29,11 +29,13 @@ pub fn main(init: std.process.Init) !u8 { defer driver.close(); var ctx = try cuda.Context.init(&driver, 0); defer ctx.deinit(); - var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); - defer set.deinit(); + var set: ?cuda.aot.Set = if (std.mem.eql(u8, args[2], "native")) null else try cuda.aot.Set.load(gpa, io, &driver, ctx.device, args[2]); + defer if (set) |*x| x.deinit(); + var glue = try flash.Glue.load(&driver); + defer glue.deinit(); var stream = try cuda.Stream.init(&driver, true); defer stream.deinit(); - const tri: flash.Tri = .{ .set = &set, .s = stream }; + const tri: flash.Tri = .{ .set = if (set) |*x| x else null, .own = &glue, .s = stream }; const got = try gpa.alloc(u32, count); defer gpa.free(got); try flash.row_path.run(flash.Tri, gpa, io, &driver, &stream, tri, args[1], ids.items, count, got); @@ -72,11 +74,13 @@ fn drafted(gpa: std.mem.Allocator, io: std.Io, model: []const u8, kernels: []con defer driver.close(); var ctx = try cuda.Context.init(&driver, 0); defer ctx.deinit(); - var set = try cuda.aot.Set.load(gpa, io, &driver, ctx.device, kernels); - defer set.deinit(); + var set: ?cuda.aot.Set = if (std.mem.eql(u8, kernels, "native")) null else try cuda.aot.Set.load(gpa, io, &driver, ctx.device, kernels); + defer if (set) |*x| x.deinit(); + var glue = try flash.Glue.load(&driver); + defer glue.deinit(); var stream = try cuda.Stream.init(&driver, true); defer stream.deinit(); - const tri: flash.Tri = .{ .set = &set, .s = stream }; + const tri: flash.Tri = .{ .set = if (set) |*x| x else null, .own = &glue, .s = stream }; const same = try flash.mtp.check(flash.Tri, gpa, io, &driver, &stream, tri, model, prompts.items, count); std.debug.print("{s}\n", .{if (same) "drafted == serial on every prompt" else "drafted DIFFERS from serial"}); return if (same) 0 else 1; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 564fd63ee1..8ee688cadd 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -28,6 +28,11 @@ fn cacheRules(comptime F: type) api.prompt_cache.Rules { return .{ .lookahead = r.lookahead, .planned = r.planned, .grid = if (r.planned) F.prompt_rows else 0, .min_gap = r.min_gap, .min_prompt = r.min_prompt }; } +/// Whether a family that may pass over a captured set (one without its entry points) runs its own glue. +fn ownGlue(loaded: anytype) bool { + return if (@hasField(@TypeOf(loaded), "own_glue")) loaded.own_glue else false; +} + /// What the lane host takes from the family: its quiet wait before a warm pass, its lone driver's hand-over budget. const PerFamily = struct { warm_quiet: u64 = 0, lone_min_budget: u32 = 0 }; @@ -368,8 +373,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const snaps = comptime snapshots(F); const cache = if (snaps != null) budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)) else 0; h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", - if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels != null and !ownGlue(loaded)) try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{kernels.?}) else "own glue kernels", }); errdefer gpa.free(h.startup); if (snaps) |vtable| if (cache > 0) { From 235ea01801c9e20dacc134aab61c9d725561972c Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:20:16 -0500 Subject: [PATCH 046/115] Flash Next CUDA: TF_FLASHNEXT_OWNED runs named entry points on our glue beside a captured set With a captured set loaded, TF_FLASHNEXT_OWNED (all, or a list of Triton names such as _chunks,_merge) makes those entry points run our kernels and the rest the set's, so a difference can be put down to one kernel. Glue reads it once at load into a mask; without a set every entry is ours, as before. flashnext-compare takes it on its own side, so it compares the captured set against the captured set with the named entries swapped. --- .../families/flashnext/cuda_compare_main.zig | 3 +- zig/src/families/flashnext/cuda_glue.zig | 29 ++++++++++++ zig/src/families/flashnext/cuda_launch.zig | 46 ++++++++++--------- 3 files changed, 56 insertions(+), 22 deletions(-) diff --git a/zig/src/families/flashnext/cuda_compare_main.zig b/zig/src/families/flashnext/cuda_compare_main.zig index c156336d1e..b1f31d54a2 100644 --- a/zig/src/families/flashnext/cuda_compare_main.zig +++ b/zig/src/families/flashnext/cuda_compare_main.zig @@ -138,7 +138,8 @@ pub fn main(init: std.process.Init) !u8 { var stream = try cuda.Stream.init(&driver, true); defer stream.deinit(); const captured: flash.Tri = .{ .set = &set, .s = stream }; - const own: flash.Tri = .{ .own = &glue, .s = stream }; + // TF_FLASHNEXT_OWNED: only the entries it names run ours on this side, the rest the captured set's + const own: flash.Tri = if (glue.mixed != 0) .{ .set = &set, .own = &glue, .s = stream } else .{ .own = &glue, .s = stream }; var e = try prompt.Engine.init(gpa, io, &driver, &stream, args[1]); defer e.deinit(); var out_a = try cuda.DeviceBuffer.alloc(&driver, rows_at_once * vocab * 2); diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index 0bebfcfa84..1c7683e4dd 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -42,7 +42,30 @@ const Fns = struct { merge: cuda.Function, }; +/// The captured entry points our glue replaces, by Triton name (TF_FLASHNEXT_OWNED lists them). +pub const entries = [_][]const u8{ "_embed", "_rmsnorm", "_add_streams", "_attn_prep", "_attn_gate", "_chunks", "_merge", "_pool", "_scores", "_hc_writeback", "_qmm_hcdown", "_hc_reduce_act", "_qmm_upmix", "_ple_gate", "_ple_conv", "_router", "_topk_rows" }; + +/// The bit of entry `name` in Glue.mixed. +pub fn entryBit(comptime name: []const u8) u32 { + inline for (entries, 0..) |e, i| if (comptime std.mem.eql(u8, e, name)) return 1 << i; + @compileError("no entry point " ++ name); +} + +/// With a captured set loaded, the entries TF_FLASHNEXT_OWNED names ("all", or Triton names) run ours: an A/B switch. +fn mixedFromEnv() u32 { + const v = std.mem.span(std.c.getenv("TF_FLASHNEXT_OWNED") orelse return 0); + if (std.mem.eql(u8, v, "all")) return (1 << entries.len) - 1; + var bits: u32 = 0; + var it = std.mem.tokenizeAny(u8, v, ", "); + while (it.next()) |name| for (entries, 0..) |e, i| if (std.mem.eql(u8, e, name)) { + bits |= @as(u32, 1) << @intCast(i); + }; + return bits; +} + pub const Glue = struct { + /// Entries ours run even with a captured set loaded (TF_FLASHNEXT_OWNED, read once at load). + mixed: u32 = 0, hc: cuda.Module, attn: cuda.Module, emb: cuda.embed_rows.EmbedRows, @@ -53,6 +76,7 @@ pub const Glue = struct { pub fn load(d: *const cuda.Driver) !Glue { if (!cuda.kernels.available) return error.BuiltWithoutKernels; var g: Glue = undefined; + g.mixed = mixedFromEnv(); g.hc = try cuda.Module.load(d, cuda.kernels.flashnext_hc); errdefer g.hc.unload(); g.attn = try cuda.Module.load(d, cuda.kernels.flashnext_attn); @@ -265,6 +289,11 @@ pub const Glue = struct { } }; +test "every entry point has its own bit" { + try std.testing.expectEqual(@as(u32, 1), entryBit("_embed")); + try std.testing.expectEqual(@as(u32, 1) << 16, entryBit("_topk_rows")); +} + test "the attention tiles hold the model's heads" { try std.testing.expect(wt.head_dim == tile_head_dim and wt.q_heads / wt.kv_heads <= tile_heads); } diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 020fa63e69..bcc9ca18ac 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -21,7 +21,8 @@ pub const row_path = @import("cuda_rows.zig"); pub const mtp = @import("cuda_mtp.zig"); pub const native = @import("cuda_lanes.zig").Namespace(Tri); /// Our own glue kernels, which every launch takes when no captured set is loaded -pub const Glue = @import("cuda_glue.zig").Glue; +pub const Glue = glue_.Glue; +const glue_ = @import("cuda_glue.zig"); const p = aot.ptr; @@ -61,13 +62,16 @@ pub const Tri = struct { try t.set.?.run(t.s, name, grid, args, consts); } - fn glue(t: Tri) ?*const Glue { - return if (t.set == null) t.own.? else null; + /// Ours for `name` without a captured set, or with one when TF_FLASHNEXT_OWNED names it; else null: captured. + fn glue(t: Tri, comptime name: []const u8) ?*const Glue { + if (t.set == null) return t.own.?; + const g = t.own orelse return null; + return if (g.mixed & glue_.entryBit(name) != 0) g else null; } /// glue.embed: ids (R,) int32 into (R, copies * dims) bf16, MLX 4-bit group 32. pub fn embed(t: Tri, ids: u64, weight: u64, scales: u64, biases: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { - if (t.glue()) |g| return g.embed(t.s, ids, weight, scales, biases, out, dims, copies, rows); + if (t.glue("_embed")) |g| return g.embed(t.s, ids, weight, scales, biases, out, dims, copies, rows); try t.run("_embed", embedGrid(rows, dims), &.{ p("IDS", "*i32", ids), p("W", "*i32", weight), @@ -79,7 +83,7 @@ pub const Tri = struct { /// glue.rmsnorm per row and `group` features: bf16 out [rows, dims] and 32-group sums; x rows `x_stride` apart. pub fn rmsnorm(t: Tri, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { - if (t.glue()) |g| return g.rmsnorm(t.s, x, w, out, xs, eps, x_stride, rows, dims, group); + if (t.glue("_rmsnorm")) |g| return g.rmsnorm(t.s, x, w, out, xs, eps, x_stride, rows, dims, group); try t.run("_rmsnorm", .{ u(rows), u(dims / group), 1 }, &.{ p("X", "*bf16", x), p("W", "*fp32", w), @@ -92,7 +96,7 @@ pub const Tri = struct { /// glue.add_streams: out[r, s] = bf16(e[r] + hs[r, s]), the MTP head's input streams. Program (row, 256 columns). pub fn addStreams(t: Tri, e: u64, hs: u64, out: u64, rows: usize, dims: usize, streams: usize) !void { - if (t.glue()) |g| return g.addStreams(t.s, e, hs, out, rows, dims, streams); + if (t.glue("_add_streams")) |g| return g.addStreams(t.s, e, hs, out, rows, dims, streams); try t.run("_add_streams", .{ u(rows), u(dims / 256), 1 }, &.{ p("E", "*bf16", e), p("HS", "*bf16", hs), @@ -102,7 +106,7 @@ pub const Tri = struct { /// glue.attn_gate: one program per row and per query head. pub fn attnGate(t: Tri, o: u64, proj: u64, out: u64, xs: u64, rows: usize, pair_width: usize, q_heads: usize, head_dim: usize) !void { - if (t.glue()) |g| return g.attnGate(t.s, o, proj, out, xs, rows, pair_width, q_heads, head_dim); + if (t.glue("_attn_gate")) |g| return g.attnGate(t.s, o, proj, out, xs, rows, pair_width, q_heads, head_dim); try t.run("_attn_gate", .{ u(rows), u(q_heads), 1 }, &.{ p("O", "*bf16", o), p("P", "*bf16", proj), @@ -113,7 +117,7 @@ pub const Tri = struct { /// glue.attn_prep for text positions, bf16 KV (BITS 0, the scale pointers unread); `length` unused at MODE 0. pub fn attnPrep(t: Tri, proj: u64, pos: u64, q_scale: u64, k_scale: u64, i_scale: u64, inv: u64, q: u64, kc: u64, vc: u64, iq: u64, ikc: u64, rope: u64, delta: u64, eps: f32, rows: usize, length: usize) !void { - if (t.glue()) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows); + if (t.glue("_attn_prep")) |g| return g.attnPrep(t.s, proj, pos, q_scale, k_scale, i_scale, inv, q, kc, vc, iq, ikc, eps, rows); try t.run("_attn_prep", .{ u(rows), u(24 + 2 + 4 + 1), 1 }, &.{ p("P", "*bf16", proj), p("POS0", "*i32", pos), @@ -137,7 +141,7 @@ pub const Tri = struct { /// attention._chunks over 512-key chunks, bf16 KV; `qsa` the sparse form (NCH 5, QSA 1), else dense (NCH 3, QSA 0). pub fn attnChunks(t: Tri, q: u64, kc: u64, vc: u64, pos: u64, po: u64, pm: u64, pl: u64, ids: u64, nk: u64, sparse: u64, rows: usize, kv: usize, chunks: usize, qsa: bool) !void { - if (t.glue()) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa); + if (t.glue("_chunks")) |g| return g.attnChunks(t.s, q, kc, vc, pos, po, pm, pl, ids, nk, sparse, rows, kv, chunks, qsa); try t.run("_chunks", .{ u(rows), u(kv), u(chunks) }, &.{ p("Q", "*bf16", q), p("KC", "*bf16", kc), @@ -156,7 +160,7 @@ pub const Tri = struct { /// attention._merge: one program per row and per KV head, in attnChunks' form. pub fn attnMerge(t: Tri, po: u64, pm: u64, pl: u64, pos: u64, out: u64, nk: u64, sparse: u64, rows: usize, kv: usize, qsa: bool) !void { - if (t.glue()) |g| return g.attnMerge(t.s, po, pm, pl, pos, out, nk, sparse, rows, kv, qsa); + if (t.glue("_merge")) |g| return g.attnMerge(t.s, po, pm, pl, pos, out, nk, sparse, rows, kv, qsa); try t.run("_merge", .{ u(rows), u(kv), 1 }, &.{ p("PO", "*fp32", po), p("PM", "*fp32", pm), @@ -170,7 +174,7 @@ pub const Tri = struct { /// attention.qsa_pool for text rows (MODE 0): every complete block of 4 indexer keys pooled, normed and roped. pub fn qsaPool(t: Tri, ikc: u64, pooled: u64, pos: u64, k_scale: u64, inv: u64, eps: f32, rows: usize) !void { - if (t.glue()) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows); + if (t.glue("_pool")) |g| return g.qsaPool(t.s, ikc, pooled, pos, k_scale, inv, eps, rows); try t.run("_pool", .{ u(rows / 4 + 2), 1, 1 }, &.{ p("IKC", "*bf16", ikc), p("POOLED", "*bf16", pooled), @@ -185,7 +189,7 @@ pub const Tri = struct { /// attention._scores: rows' complete blocks scored; `nb` is the row stride (a multiple of 16), `blocks` the grid. pub fn qsaScores(t: Tri, iq: u64, pooled: u64, pos: u64, sc: u64, nb: usize, rows: usize, blocks: usize) !void { - if (t.glue()) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks); + if (t.glue("_scores")) |g| return g.qsaScores(t.s, iq, pooled, pos, sc, nb, rows, blocks); try t.run("_scores", .{ u(rows), cdiv(blocks, 64), 1 }, &.{ p("IQ", "*bf16", iq), p("POOLED", "*bf16", pooled), @@ -197,7 +201,7 @@ pub const Tri = struct { /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. pub fn hcReduceAct(t: Tri, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { - if (t.glue()) |g| return g.hcReduceAct(t.s, part, act, xs, inj, sk, rows, streams, low, ndn, has_inj); + if (t.glue("_hc_reduce_act")) |g| return g.hcReduceAct(t.s, part, act, xs, inj, sk, rows, streams, low, ndn, has_inj); try t.run("_hc_reduce_act", .{ u(rows), 1, 1 }, &.{ p("PART", "*fp32", part), p("ACT", "*bf16", act), @@ -209,7 +213,7 @@ pub const Tri = struct { /// qmm.hc_upmix: the up projection and the stream mix. N is the wide row, K is the low rank. M folds when it is 1. pub fn hcUpmix(t: Tri, act: u64, xs: u64, weight: u64, scales: u64, biases: u64, normed: u64, mixed: u64, xs_mixed: u64, rows: usize, n: usize, k: usize, dims: usize, streams: usize) !void { - if (t.glue()) |g| return g.hcUpmix(t.s, act, xs, weight, scales, biases, normed, mixed, xs_mixed, rows, n, k, dims, streams); + if (t.glue("_qmm_upmix")) |g| return g.hcUpmix(t.s, act, xs, weight, scales, biases, normed, mixed, xs_mixed, rows, n, k, dims, streams); try t.run("_qmm_upmix", .{ cdiv(rows, 16), u(dims / 32), 1 }, &.{ p("X", "*bf16", act), p("XS", "*fp32", xs), @@ -225,7 +229,7 @@ pub const Tri = struct { /// qmm.hc_down for one decode row. N=324 is the layer down stacked with inject. M folds when it is 1. pub fn hcDown(t: Tri, h: u64, pss: u64, scale: u64, normed: u64, weight: u64, scales: u64, biases: u64, out: u64, part: u64, eps: f32, rows: usize, n: usize, k: usize, dims: usize, nc: usize, streams: usize, sk: usize) !void { - if (t.glue()) |g| return g.hcDown(t.s, h, pss, scale, normed, weight, scales, biases, out, part, eps, rows, n, k, dims, nc, streams, sk); + if (t.glue("_qmm_hcdown")) |g| return g.hcDown(t.s, h, pss, scale, normed, weight, scales, biases, out, part, eps, rows, n, k, dims, nc, streams, sk); try t.run("_qmm_hcdown", .{ cdiv(rows, 16), cdiv(n, 64), u(sk) }, &.{ p("H", "*bf16", h), p("PSS", "*fp32", pss), @@ -243,7 +247,7 @@ pub const Tri = struct { /// glue.hc_writeback. Pointer types follow the tensor dtypes Triton specialized (mode 0..4). pub fn hcWriteback(t: Tri, h: u64, hout: u64, pss: u64, branch: u64, branch_ty: []const u8, inj: u64, y: u64, y_ty: []const u8, wts: u64, wts_ty: []const u8, rs: usize, rows: usize, dims: usize, streams: usize, mode: usize, topk: usize, slots: usize, world: usize) !void { - if (t.glue()) |g| return g.hcWriteback(t.s, h, hout, pss, branch, inj, y, wts, rs, rows, dims, streams, mode, topk, slots, world); + if (t.glue("_hc_writeback")) |g| return g.hcWriteback(t.s, h, hout, pss, branch, inj, y, wts, rs, rows, dims, streams, mode, topk, slots, world); try t.run("_hc_writeback", .{ u(rows), u(dims / 256), 1 }, &.{ p("H", "*bf16", h), p("HOUT", "*bf16", hout), @@ -258,7 +262,7 @@ pub const Tri = struct { /// glue.ple_embed: gathered n-gram rows, MLX 4-bit group 32. `scale` is a constexpr. pub fn pleEmbed(t: Tri, weight: u64, scales: u64, biases: u64, out: u64, xs: u64, rows: usize, heads: usize, dh: usize, scale: f32) !void { - if (t.glue() != null) return error.CapturedKernelsRequired; // the row path's n-gram check only + if (t.set == null) return error.CapturedKernelsRequired; // the row path's n-gram check only try t.run("_ple_embed", .{ u(rows), u(heads), 1 }, &.{ p("W", "*i32", weight), p("S", "*bf16", scales), @@ -270,7 +274,7 @@ pub const Tri = struct { /// glue.ple_gate: one program per row. pub fn pleGate(t: Tri, keys: u64, vals: u64, h: u64, norm_key: u64, norm_query: u64, gated: u64, pss: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { - if (t.glue()) |g| return g.pleGate(t.s, keys, vals, h, norm_key, norm_query, gated, pss, eps, rows, dims, streams); + if (t.glue("_ple_gate")) |g| return g.pleGate(t.s, keys, vals, h, norm_key, norm_query, gated, pss, eps, rows, dims, streams); try t.run("_ple_gate", .{ u(rows), 1, 1 }, &.{ p("KEYS", "*bf16", keys), p("VALS", "*bf16", vals), @@ -285,7 +289,7 @@ pub const Tri = struct { /// glue.ple_conv. R folds to a constexpr when it is 1. pub fn pleConv(t: Tri, gated: u64, pss: u64, norm_conv: u64, tail: u64, conv_w: u64, h: u64, hout: u64, nrow: u64, eps: f32, rows: usize, dims: usize, streams: usize, taps: usize, dilation: usize) !void { - if (t.glue()) |g| return g.pleConv(t.s, gated, pss, norm_conv, tail, conv_w, h, hout, nrow, eps, rows, dims, streams, taps, dilation); + if (t.glue("_ple_conv")) |g| return g.pleConv(t.s, gated, pss, norm_conv, tail, conv_w, h, hout, nrow, eps, rows, dims, streams, taps, dilation); try t.run("_ple_conv", .{ u(rows), u(streams * dims / 512), 1 }, &.{ p("GATED", "*bf16", gated), p("PSS", "*fp32", pss), @@ -312,7 +316,7 @@ pub const Tri = struct { /// moe._topk_rows over `rows` rows of [rows, 513] logits: picks and weights [rows, 11]. pub fn topkMany(t: Tri, logits: u64, pick: u64, wts: u64, rows: usize) !void { - if (t.glue()) |g| return g.topkMany(t.s, logits, pick, wts, rows); + if (t.glue("_topk_rows")) |g| return g.topkMany(t.s, logits, pick, wts, rows); try t.run("_topk_rows", .{ u(rows), 1, 1 }, &.{ p("L", "*fp32", logits), p("PICK", "*i32", pick), @@ -322,7 +326,7 @@ pub const Tri = struct { /// moe.router over `rows` rows at the one-row tiles (BM 16, BLOCK_E 32, BK 256), each row keeping its own bits. pub fn routerMany(t: Tri, x: u64, w: u64, out: u64, x_stride: usize, rows: usize) !void { - if (t.glue()) |g| return g.routerMany(t.s, x, w, out, x_stride, rows); + if (t.glue("_router")) |g| return g.routerMany(t.s, x, w, out, x_stride, rows); try t.run("_router", .{ cdiv(rows, 16), cdiv(513, 32), 1 }, &.{ p("X", "*bf16", x), p("W", "*bf16", w), From 7746bb9da4b88c45b6c224d634abc7e52cb665af Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:08:43 -0500 Subject: [PATCH 047/115] Docs: Flash Next on CUDA boots on its own kernels; a captured set is for comparison The recipe, the RUNBOOK line, the CUDA and Flash Next pages and the CHANGELOG say the engine needs no capture: every kernel is built from the repository, TENSORFOLD_CUDA_KERNELS names a captured set only to compare against the Python engine's teacher path, and TF_FLASHNEXT_OWNED swaps single entry points. The recipe lists the new tf-cuda-test cases and flashnext-compare, the semantic settings the config check refuses, measurements on our kernels beside the captured set's, and where our tokens part from the teacher path. The family recipe's lane-kernel row names route_rows.cu. --- CHANGELOG.md | 4 +- RUNBOOK.md | 4 +- docs/recipes/adding-a-zig-family.md | 2 +- docs/recipes/cuda.md | 2 +- docs/recipes/flash-next-cuda.md | 116 ++++++++++++++++------------ docs/recipes/qwen3.8-flash-next.md | 2 +- 6 files changed, 76 insertions(+), 54 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e856dfcb6..e1d5f3ab33 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,9 +4,11 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased -- Qwen3.8 Flash Next serves on CUDA from its MLX 4-bit group-32 checkpoint with the MTP head, through the lane core: drafted windows, shared rounds of streams without drafts, kept prompt states (a drafting stream's too, with the head's cache rows and a lookahead of 1, as Nemotron's) and prompts a chunk at a time between other streams' rounds. Each prompt chunk, verify window and decode round is one batched pass whose kernels give every row the bits it gets alone, so drafted replies equal `"draft": false`, concurrent replies equal their solo runs and resumed prompts equal fresh ones; past 2,048 keys attention reads the 512 blocks the model's indexer picks, up to the 262,144-token window. It replays the Python 0.6.1 engine's Triton kernels from a captured set (`TENSORFOLD_CUDA_KERNELS`; the capture and pack scripts are in `zig/tests/cuda/flashnext`), and 32 greedy tokens equal that engine's teacher path on prompts of 19 to 12,001 tokens. On a GB10 one stream decodes 82.6 tok/s drafted on code and 61.8 on chat, 35.0 without drafts, and 4 sampled streams without drafts 77.3 in all; a cold 12,001-token prompt's first token takes 10.0 s and a 45,030-token one's 37.9 s, and another stream waits at most 1.3 s between tokens while it runs. A checkpoint of another width or shape is refused at load, by its config field. @AdrianBinDC started the port in #476, matching every layer of one decode token to host checks; see [Flash Next on CUDA](docs/recipes/flash-next-cuda.md) (#472). Thanks to @AdrianBinDC. +- Qwen3.8 Flash Next serves on CUDA from its MLX 4-bit group-32 checkpoint with the MTP head, through the lane core: drafted windows, shared rounds of streams without drafts, kept prompt states (a drafting stream's too, with the head's cache rows and a lookahead of 1, as Nemotron's) and prompts a chunk at a time between other streams' rounds. Each prompt chunk, verify window and decode round is one batched pass whose kernels give every row the bits it gets alone, so drafted replies equal `"draft": false`, concurrent replies equal their solo runs and resumed prompts equal fresh ones; past 2,048 keys attention reads the 512 blocks the model's indexer picks, up to the 262,144-token window. Every kernel it launches is built from this repository (`flashnext_hc.cu`, `flashnext_attn.cu`, `flashnext_batch.cu`, and the core's token, RMS, routing and expert rows), so it needs no capture; `TENSORFOLD_CUDA_KERNELS` can name a captured set of the Python 0.6.1 engine's Triton kernels (scripts in `zig/tests/cuda/flashnext`), which then stands in for the glue so replies can be compared with that engine's teacher path. On a GB10 one stream decodes 84.2 tok/s drafted on code and 62.6 on chat, 35.1 without drafts, and 4 sampled streams without drafts 77.9 in all, within 2% of the captured set either way; a cold 12,001-token prompt's first token takes 10.4 s and a 45,030-token prompt 36.5 s, and another stream waits at most 1.3 s between tokens while it runs. 32 greedy tokens equal the Python engine's teacher path on four of five prompts of 19 to 12,001 tokens and part at token 15 on the 3,001-token one (the captured set gives all five); next-token perplexity over 4,096 tokens of this repository's documentation is 6.178 against the captured set's 6.168. A checkpoint of another width, shape or semantic setting (RMS epsilon, RoPE, gate and activation types, layer roles) is refused at load, by its config field. @AdrianBinDC started the port in #476, matching every layer of one decode token to host checks; see [Flash Next on CUDA](docs/recipes/flash-next-cuda.md) (#472). Thanks to @AdrianBinDC. - A Flash Next sequence's buffers are zeroed on the engine's compute stream. A fill with no stream returns before the memset lands and is not ordered against the non-blocking stream the passes run on, so a new sequence's first prompt chunk could read its DeltaNet state before it was zeroed; the drafted check, which zeroes the states between a prompt's serial and drafted runs, once left the serial tokens that way. - CUDA has a lane kernel for affine 4-bit experts over a window's or a prompt chunk's rows (`experts_rows.cu`, `cuda.experts.RowsExperts`): experts.cu's decode sums for each row, two row tiles of 16 a warp, so a warp reads an expert's weights once for 32 pairs. `tf-cuda-test experts-rows` checks it against the decode kernel's one-row launches byte for byte at 1 to 300 rows. +- CUDA has routing rows for every family (`route_rows.cu`, `cuda.route_rows`): router logits on mma tiles, one chain over K from zero so a row's sums never depend on the row count, and softmax top-k with an optional sigmoid-gated shared expert; `tf-cuda-test route-rows` checks both against float64 (`core/route_ref.zig`). The token rows take groups of 32 and a copy count, and the RMS rows gain a form with the stored outputs' group sums (`Norms.rmsSums`); `norm-rows` covers both. +- A Flash Next pass uploads its token ids, positions and n-gram rows from pinned staging the scratch owns until the copies land, instead of allocator memory freed when the call returned (an asynchronous copy from pageable memory is only documented as "might be synchronous"), and a failed zero fill frees the buffer it was filling. - `tools/zig/triton_aot_pack.py` leaves out a constexpr Triton folded from `None` (an unused pointer argument), where it stopped, and takes `--jit` once per capture, merged, so one call packs several captures; `test_aot_manifest.py` covers both. ## 1.0.11 diff --git a/RUNBOOK.md b/RUNBOOK.md index 32a3638f75..f668a07f2b 100644 --- a/RUNBOOK.md +++ b/RUNBOOK.md @@ -123,8 +123,8 @@ The archives require glibc 2.28 or newer and an NVIDIA driver providing `libcuda CUDA serves Nemotron 3.5 Lightning, qualified on GB10 (`linux-aarch64`, greedy and sampled) and on Ampere cards with compute capability 8.6 (`linux-x86_64`, greedy: the RTX 30 series, RTX A6000, A10 and A40), and checked on an RTX PRO 6000 Blackwell (compute capability 12.0, `linux-x86_64`). -Flash Next also serves on CUDA from its MLX 4-bit checkpoint and a captured kernel set named by `TENSORFOLD_CUDA_KERNELS`, -run on a GB10: see [Flash Next on CUDA](docs/recipes/flash-next-cuda.md). +Flash Next also serves on CUDA from its MLX 4-bit checkpoint, on the built-in kernels, run on a GB10: see +[Flash Next on CUDA](docs/recipes/flash-next-cuda.md). On a GB10 arm64 host (an x86_64 host uses `linux-x86_64` in the same commands): diff --git a/docs/recipes/adding-a-zig-family.md b/docs/recipes/adding-a-zig-family.md index 7372013d35..5b100aaa11 100644 --- a/docs/recipes/adding-a-zig-family.md +++ b/docs/recipes/adding-a-zig-family.md @@ -71,7 +71,7 @@ lane kernels go in, check against the engine's own plain output (step 4). | Piece | Shared now | Family copies or gaps | | --- | --- | --- | -| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels. On CUDA, `zig/kernels/cuda/experts_rows.cu` runs affine 4-bit group-32 experts over a window's rows with the decode kernel's sums (Flash Next) | +| Lane kernels | `zig/kernels/metal/ops/qmv.metal` (affine 4-, 6- and 8-bit row matmuls, MLX's arithmetic) and `zig/src/core/frags.zig` (the prompt layout on M1-M4) | Flash Next: `zig/kernels/metal/decode/fn_lane.metal` (6-bit, groups of 32, tensor ops, up to 16 rows). Kimi K3: MXFP4 experts and bf16 projections. Nemotron: generated per-shape kernels. On CUDA, `zig/kernels/cuda/experts_rows.cu` runs affine 4-bit group-32 experts over a window's rows with the decode kernel's sums, and `route_rows.cu` the router logits and softmax top-k (Flash Next) | | GPU-side rounds | the host loop, depth rule and copy proposer in `zig/src/core/lanes`, used by Nemotron | Flash Next: its own host loop, plus `fz_accept`, the ring and event chaining (`zig/src/families/flashnext/replay.zig`). Nemotron on Metal: `zig/src/families/nemotron/gpu_round.zig` for a lone stream. Kimi K3: none | | Several Macs | the MCDMA fabric (`zig/src/fabric`) | Flash Next: `zig/src/families/flashnext/tp.zig` splits rows and holds the whole model on each Mac. Kimi K3: `zig/src/families/kimi_k3/parallel.zig` is a split plan with no link yet | | Prompt path | segments; the prompt cache, used by Flash Next, GLM, Nemotron and the Qwen3.8-27B | per-family prompt kernels and snapshots | diff --git a/docs/recipes/cuda.md b/docs/recipes/cuda.md index 72bcb3550d..aeb28ec541 100644 --- a/docs/recipes/cuda.md +++ b/docs/recipes/cuda.md @@ -8,7 +8,7 @@ They are not the native Zig release's support matrix. The kernel paths, the prom (`--prefill-fp8`, `--no-prefill-fp8`, `--precision`) and the request and memory notes below are that engine's too, and the native binary refuses those flags. -The native binary registers Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), and Flash Next from its MLX 4-bit checkpoint with a captured kernel set, run on a GB10 ([Flash Next on CUDA](flash-next-cuda.md)). +The native binary registers Nemotron on CUDA, qualified on NVIDIA GB10 and on Ampere cards with compute capability 8.6 (the RTX 30 series, RTX A6000, A10 and A40), and Flash Next from its MLX 4-bit checkpoint, run on a GB10 ([Flash Next on CUDA](flash-next-cuda.md)). A user's receipt for 1.0.4 on a bare-metal RTX 3090, at 225 W and 350 W, is in [#592](https://github.com/ashhart/TensorFold/issues/592). GLM-5.3-Flash on CUDA, EXL3 loading and DFlash2 drafting remain Python-engine features, not native upgrade paths. Inspect the installed binary with `tensorfold-native --version` and `tensorfold-native capabilities --json`. diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index 853e255a94..5a3a2ceeee 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -1,18 +1,15 @@ # Qwen3.8 Flash Next on CUDA The native CUDA build serves Flash Next (`qwen4_exp`) on one NVIDIA GPU from an MLX affine 4-bit, group-32 checkpoint -with its MTP head. It replays the Python 0.6.1 engine's Triton kernels from a captured set beside its own CUDA -kernels, so it gives that engine's teacher-path tokens. @AdrianBinDC started the port in +with its MTP head, on kernels built in this repository. A set of the Python 0.6.1 engine's Triton kernels, captured +and packed, can stand in for its glue to compare against that engine's teacher path. @AdrianBinDC started the port in [#476](https://github.com/ashhart/TensorFold/pull/476); it is tracked in [#472](https://github.com/ashhart/TensorFold/issues/472). It has been run on a GB10 (sm_121) only. ## Run -The engine needs a captured kernel set (see [The kernel set](#the-kernel-set)); without `TENSORFOLD_CUDA_KERNELS` it -refuses to load. - ```bash -TENSORFOLD_CUDA_KERNELS=/path/to/flashnext-kernels tensorfold-native serve /path/to/checkpoint --temperature 0 +tensorfold-native serve /path/to/checkpoint --temperature 0 ``` - `--context N` takes up to the model's 262,144 tokens; the default is the whole window. Past 2,048 keys each attention @@ -23,8 +20,15 @@ TENSORFOLD_CUDA_KERNELS=/path/to/flashnext-kernels tensorfold-native serve /path - `TF_FLASHNEXT_DRAFTS` (1 to 6, default 6) and `TF_FLASHNEXT_CONFIDENCE` (default 0.60) set the head's chain. - The engine uploads 73.3 GiB of the checkpoint: the n-gram tables stay mapped from the shards, and the vision tower is not loaded. -- It refuses another quantization or other shapes (config.json is checked at load, and the log names the field), - `--segments` and `--slide`, and a folder whose shared experts carry Living Weights neurons. +- It refuses another quantization, other shapes, and semantic settings its kernels fix (RMS epsilon, RoPE base, rotated + share, type and scaling, the gate and activation types, attention bias, the layer roles, the n-gram layer, kernel + and dilation): config.json is checked at load and the log names the field. It also refuses `--segments`, `--slide` + and a folder whose shared experts carry Living Weights neurons. +- `TENSORFOLD_CUDA_KERNELS=DIR` swaps the glue for a captured set (see [A captured set](#a-captured-set)) when the set + holds every entry point the engine launches; a set without them (Nemotron's, say) is passed over, and the startup + line says which glue runs. With a set loaded, `TF_FLASHNEXT_OWNED` (`all`, or Triton names such as + `_chunks,_merge`) runs those entry points on our kernels and the rest on the set's, to put a difference down to one + kernel. ## What the core runs and what the family runs @@ -35,18 +39,42 @@ The family plugs into the core's pieces and adds no scheduler: - The core's prompt cache keeps and resumes a stream's state at the marks it asks for. The family copies the state (`cuda_snap.zig`), with the MTP head's cache rows for a drafting stream; the head reads one token past a row, so states keep with a lookahead of 1, as Nemotron's do. -- The routed experts over a window's or a chunk's rows run on the core's plan (`cuda/grouped.zig`) and its rows kernel - (`experts_rows.cu`, `cuda.experts.RowsExperts`), which gives every row the decode kernel's sums. -- Captured Triton launches go through `cuda.aot.Set`, so a launch with no exact variant takes the general one. +- Shared row kernels: the token rows (`embed_rows.cu`, groups of 32 and a copy a hyper-connection stream), the MTP + head's RMS rows with their group sums (`norms.cu`), the router logits and softmax top-k (`route_rows.cu`), and the + routed experts over a window's or a chunk's rows on the core's plan (`cuda/grouped.zig`, `experts_rows.cu`). The family supplies its config check, weight loading and packing, the layer graph on one batched pass -(`cuda_batch.zig`), and kernels for what is its own: the Gated DeltaNet chain, the sparse indexer's block selection, -the n-gram gather, the hyper-connection and attention launches, and the MTP head. +(`cuda_batch.zig`), and kernels for what is its own: the hyper-connections and the n-gram gate and conv +(`flashnext_hc.cu`), attention with its gate and the sparse indexer's pooling, scores and block selection +(`flashnext_attn.cu`, `flashnext_batch.cu`), the Gated DeltaNet chain, the n-gram gather and the MTP head. + +Every kernel keeps a row's own sums in a fixed order, so a row's bits are the same at any width: drafted windows equal +serial decoding, rounds equal lone streams, and resumed prompts equal fresh ones. The glue kernels round to bf16 where +the Python engine's do, with IEEE division and square root and a within-one-ulp exp; the router and the 4-bit +hyper-connection dots run on mma tiles with K in ascending order. + +## Checks + +- `tf-cuda-test flashnext-hc`, `flashnext-attn` and `route-rows`, and `norm-rows` for the shared rows: each kernel + against a float64 chain that rounds where it rounds, and every row count's bits equal (tolerances in the test files). +- `zig build flashnext-prompt -Dflashnext_model=DIR -Dflashnext_kernels=native -Dflashnext_ids=FILE + -Dflashnext_count=32`: 32 greedy tokens through the batched pass, against the Python engine's teacher path + (`zig/tests/cuda/flashnext/lr7/results.json`, `results-sparse.json`); a captured set's folder in place of `native` + runs its glue. `TF_FLASHNEXT_DRAFTS=6` decodes each prompt serially and drafted and compares them token for token + (`TF_FLASHNEXT_SEED` for keyed sampled draws). +- `zig build flashnext-compare -Dflashnext_model=DIR -Dflashnext_kernels=SET -Dflashnext_ids=FILE`: our glue against a + captured set on one loaded engine: every prompt row's scores (top-1 agreement, KL, each path's perplexity on the + prompt's own next tokens, dense rows and rows past 2,048 apart), then each path's greedy reply and where they part. + With `TF_FLASHNEXT_OWNED` it compares the set against the set with those entries swapped. +- `zig build flashnext-select -Dflashnext_fixtures=DIR`: the native block selection against Triton's lists. +- `tf-cuda-test experts-rows`: the rows kernel against the decode kernel's one-row launches, byte for byte. +- Served: `tools/bench_concurrent.py --alone --serial` (concurrent equal to solo, drafted equal to `"draft": false`), + and a prompt sent twice (the second resumed from a kept state, equal to the first). -## The kernel set +## A captured set -The set is built from the Python 0.6.1 engine in a container (`IMAGE`, default tag `tensorfold-qwen38:v0.6.1`) with -the scripts in `zig/tests/cuda/flashnext`, then packed with `tools/zig/triton_aot_pack.py`: +A set is built from the Python 0.6.1 engine in a container (`IMAGE`, default tag `tensorfold-qwen38:v0.6.1`) with the +scripts in `zig/tests/cuda/flashnext`, then packed with `tools/zig/triton_aot_pack.py`: ```bash zig/tests/cuda/flashnext/capture.sh WORK/cap /path/to/hf/snapshot # the model's own run, about 78 GiB of GPU @@ -56,49 +84,41 @@ zig/tests/cuda/flashnext/capture_sparse.sh WORK/sparse # th zig/tests/cuda/flashnext/pack.sh KERNELS WORK/cap/out/cap WORK/rows/out WORK/mtp/out WORK/sparse/out ``` -`capture.sh` also writes the prompts and the Python engine's references the checks below read. The set has -76 kernels; `capture_sparse.sh` writes the selection fixtures `flashnext-select` reads. - -## Checks - -- `zig build flashnext-prompt -Dflashnext_model=DIR -Dflashnext_kernels=SET -Dflashnext_ids=FILE -Dflashnext_count=32`: - 32 greedy tokens through the batched pass against the Python engine's teacher path - (`zig/tests/cuda/flashnext/lr7/results.json`, `results-sparse.json`). `TF_FLASHNEXT_PROMPT=rows` runs the row-by-row - reference instead; `TF_FLASHNEXT_DRAFTS=6` decodes each prompt serially and drafted and compares them token for - token (`TF_FLASHNEXT_SEED` for keyed sampled draws). -- `zig build flashnext-select -Dflashnext_fixtures=DIR`: the native block selection against Triton's lists. -- `tf-cuda-test experts-rows`: the rows kernel against the decode kernel's one-row launches, byte for byte. -- Served: `tools/bench_concurrent.py --alone --serial` (concurrent equal to solo, drafted equal to `"draft": false`), - and a prompt sent twice (the second resumed from a kept state, equal to the first). +`capture.sh` also writes the prompts and the Python engine's references the checks read. The set has 76 kernels; +`capture_sparse.sh` writes the selection fixtures `flashnext-select` reads. ## Measurements On a GB10 (DGX Spark, driver 580.178, CUDA 13.0), an MLX 4-bit group-32 Flash Next checkpoint with its MTP head, -the 76-kernel set, greedy unless said, `tools/bench_openai.py` (256 tokens, median of 3) and -`tools/bench_concurrent.py` (256 tokens, aggregate tok/s): - -| | tok/s | -| --- | --- | -| One stream, drafted, code / chat | 82.6 / 61.8 | -| One stream, `--no-drafts`, code / chat | 35.0 / 34.9 | -| `--no-drafts`, 1 / 2 / 4 streams, sampled code | 35.0 / 53.0 / 77.3 | -| drafted, 1 / 2 / 4 streams, greedy code | 81.6 / 80.6 / 80.4 | +greedy unless said, `tools/bench_openai.py` (256 tokens, median of 3) and `tools/bench_concurrent.py` (256 tokens, +aggregate tok/s), on our kernels and, for comparison, on the 76-kernel captured set: + +| | ours | captured | +| --- | --- | --- | +| One stream, drafted, code / chat | 84.2 / 62.6 | 82.5 / 62.1 | +| One stream, `--no-drafts`, code / chat | 35.1 / 35.1 | 35.2 / 35.2 | +| `--no-drafts`, 1 / 2 / 4 streams, sampled code | 35.0 / 53.1 / 77.9 | 35.2 / 53.0 / 77.1 | +| drafted, 1 / 2 / 4 streams, greedy code | 83.8 / 82.7 / 82.2 | 82.3 / 81.1 / 80.7 | +| Cold first token at 3,001 / 12,001 tokens (s) | 2.67 / 10.40 | 2.67 / 10.30 | +| 45,030-token prompt, no cache (s) | 36.5 | 35.8 | Every concurrent reply equalled its solo run, and every drafted one its `"draft": false` run. Drafted streams verify -their windows a pass each, so drafted throughput does not grow with streams yet. A cold prompt's first token takes 2.65 -s at 3,001 tokens, 10.0 s at 12,001 and 37.9 s at 45,030 (a 65,536-token window, no prompt cache). While a 12,001-token -prompt runs, another stream's decode waits at most 1.3 s between tokens (one prompt chunk). Each stream at a -16,384-token window takes 0.74 GiB with drafts, 0.54 GiB without. +their windows a pass each, so drafted throughput does not grow with streams yet. While a 12,001-token prompt runs, +another stream's decode waits at most 1.3 s between tokens (one prompt chunk). Each stream at a 16,384-token window +takes 0.74 GiB with drafts, 0.54 GiB without. + +Our kernels' sums are their own, so tokens part from the captured set's (the Python engine's teacher path) where two +candidates are close: 32 greedy tokens equal it on four of `zig/tests/cuda/flashnext/lr7`'s five prompts and part at +token 15 on the 3,001-token one. Over 4,096 tokens of this repository's documentation, the next-token perplexity is +6.168 on the captured set and 6.178 on ours. ## Not yet -- The set replays the Python engine's Triton kernels for norms, hyper-connections, attention, routing and the head. The - core's shared row kernels (`cuda/norms.cu`) sum in another order, so moving onto them is its own change. - The dense projections run a copy of the Python engine's group-32 lane matmul (`qmm.cu`); the core's `qmm_group` and `lane_gemv` are built for groups of 64. - No CUDA graphs or GPU-side rounds; drafted streams verify their windows one pass each, and the MTP head drafts one stream at a time. -- Prompts run slower than the Python engine's: a cold 12,001-token prompt takes 10.0 s against its 6.7 s. +- Prompts run slower than the Python engine's: a cold 12,001-token prompt takes about 10 s against its 6.7 s. - Draws are made on the host from each row's scores. -- Replies equal the Python engine's teacher path and part from its served path after 10 to 36 tokens, from the - routed-experts kernel's rounding form. +- Our glue's sums are not the captured kernels' bit for bit, so replies part from the Python engine's teacher path + after a while, as that engine's own served path does (10 to 36 tokens, from its routed-experts kernel's rounding). diff --git a/docs/recipes/qwen3.8-flash-next.md b/docs/recipes/qwen3.8-flash-next.md index 60ac2381eb..8da3709b72 100644 --- a/docs/recipes/qwen3.8-flash-next.md +++ b/docs/recipes/qwen3.8-flash-next.md @@ -11,7 +11,7 @@ tensorfold serve TensorFold/Qwen3.8-Flash-Next-MLX-6bit-MTP --name bench --tempe ``` For native serving, see the [Metal runbook](../../RUNBOOK.md#flash-next-and-paired-metal-serving). -Native Flash Next on CUDA serves the MLX 4-bit group-32 checkpoint from a captured kernel set: see +Native Flash Next on CUDA serves the MLX 4-bit group-32 checkpoint on the built-in kernels: see [Flash Next on CUDA](flash-next-cuda.md) ([#472](https://github.com/ashhart/TensorFold/issues/472)). The sections below describe the Python 0.6.6 engine, available at the From e7163c5eee7beb750ee9b99b4ebba1c495a08285 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:19:20 -0500 Subject: [PATCH 048/115] Flash Next CUDA: each layer_types label must be the exact role the kernels run The check compared each label only with linear_attention, so a full-attention layer labelled sliding_attention (or any other word) was accepted and served as full attention. Every label now has to be linear_attention or full_attention, as the layer's place in the stack says. The new test relabels full and linear layers in turn and expects layer_types back; on the old check the full layer relabelled sliding_attention came back accepted. --- zig/src/families/flashnext/cuda_config.zig | 33 ++++++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig index 3d87549150..8fc2652e67 100644 --- a/zig/src/families/flashnext/cuda_config.zig +++ b/zig/src/families/flashnext/cuda_config.zig @@ -110,8 +110,8 @@ pub fn semantics(text: std.json.ObjectMap) ?[]const u8 { if (given(text, "layer_types")) |v| { if (v != .array or v.array.items.len != wt.layers_n) return "layer_types"; for (v.array.items, 0..) |kind, i| { - if (kind != .string) return "layer_types"; - if (std.mem.eql(u8, kind.string, "linear_attention") != wt.linearLayer(i)) return "layer_types"; + const want = if (wt.linearLayer(i)) "linear_attention" else "full_attention"; + if (!holds(kind, .{ .word = want })) return "layer_types"; } } const ple = given(text, "ple_layer_ids") orelse return "ple_layer_ids"; @@ -236,6 +236,35 @@ test "a changed epsilon, RoPE, gate, activation or layer role is refused by name } } +test "each layer's label is exactly the role the kernels run: a sliding or misspelt one is refused" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const gpa = arena.allocator(); + const full = 3; // the first full-attention layer + try std.testing.expect(!wt.linearLayer(full) and wt.linearLayer(0)); + const Case = struct { layer: usize, label: []const u8 }; + const cases = [_]Case{ + .{ .layer = full, .label = "sliding_attention" }, + .{ .layer = full, .label = "attention" }, + .{ .layer = full, .label = "Full_Attention" }, + .{ .layer = wt.layers_n - 1, .label = "sliding_attention" }, + .{ .layer = 0, .label = "full_attention" }, + .{ .layer = 0, .label = "sliding_attention" }, + }; + for (cases) |c| { + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(gpa); + var ple: std.json.Array = .init(gpa); + try served(gpa, &text, &rope_params, &mtp, &kinds, &ple); + try std.testing.expectEqual(@as(?[]const u8, null), semantics(text)); + kinds.items[c.layer] = .{ .string = c.label }; + try text.put(gpa, "layer_types", .{ .array = kinds }); + try std.testing.expectEqualStrings("layer_types", semantics(text) orelse "accepted"); + } +} + test "RoPE's base and type inside rope_parameters, and the MTP layer's base, are checked too" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From 4b46157a7113a41314ac97c8e4a65269a8ca94b1 Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:19:33 -0500 Subject: [PATCH 049/115] Flash Next CUDA: a failed upload or record drains the stream before the staging is reused or freed A pass queues its uploads from the scratch's pinned staging, then records an event behind them, and the next pass waits on that event before it writes the staging again. If a later upload or the record failed, an earlier upload of the same pass stayed queued while the event still marked the previous pass, so the next pass, or the scratch's free, could change or release bytes a copy had yet to read. The staging now notes an upload queued after the last record and, until a record covers it, waits by draining the stream instead of on the event. Scratch.free waits the same way before it frees the staging. The new test runs a clean pass and then one whose second upload, or whose record, fails, against a stand-in queue, and expects nothing queued at the next wait; waiting on the event alone, as before, leaves one copy queued. --- zig/src/families/flashnext/cuda_batch.zig | 9 +- .../families/flashnext/cuda_batch_state.zig | 127 ++++++++++++++++-- 2 files changed, 123 insertions(+), 13 deletions(-) diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 8dcdf58e3c..0686c0f688 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -89,7 +89,7 @@ pub fn Pass(comptime Tri: type) type { stream: cuda.Stream, tri: Tri, k: *const Kernels, - s: *const Scratch, + s: *Scratch, m: Model, seq: SeqState, hooks: Hooks, @@ -152,8 +152,7 @@ pub fn Pass(comptime Tri: type) type { /// A queued copy from pinned host bytes (the scratch's staging), which must stay unchanged until it lands. pub fn upload(self: *Self, dst: u64, bytes: []const u8) !void { - const d = self.stream.d; - try d.check(d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, self.stream.handle), "cuMemcpyHtoDAsync"); + try self.s.upload(dst, bytes); } // ---- the blocks ---------------------------------------------------------------------------------------- @@ -374,7 +373,7 @@ pub fn Pass(comptime Tri: type) type { @memset(pos, 0); for (0..blocks) |b| pos[4 * b] = @intCast(if (self.lanes.len > 0) self.lanes[b].pos else self.pos0 + b * att_rows); try self.upload(s.pos, std.mem.sliceAsBytes(pos)); - try s.uploaded(self.stream); + try s.uploaded(); } /// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. @@ -444,7 +443,7 @@ pub fn Pass(comptime Tri: type) type { if (profile()) std.debug.print("batch: {d} rows' n-gram rows gathered on the host in {d:.1} ms\n", .{ n, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); try self.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); try self.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); - try s.uploaded(self.stream); + try s.uploaded(); if (self.pending != null) try self.writeback(0, n); try self.pleBlock(lw, ids); } diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 6e9fd2e437..c6456a0faf 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -140,6 +140,55 @@ pub const Kernels = struct { } }; +/// The compute stream as a pass's pinned staging uses it: copies from it, an event behind them, and a drain. +pub const Uploads = struct { + stream: cuda.Stream, + event: cuda.Event, + + pub fn copy(u: Uploads, dst: u64, bytes: []const u8) !void { + const d = u.stream.d; + try d.check(d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, u.stream.handle), "cuMemcpyHtoDAsync"); + } + pub fn record(u: Uploads) !void { + try u.event.record(u.stream); + } + pub fn wait(u: Uploads) !void { + try u.event.synchronize(); + } + pub fn drain(u: Uploads) !void { + try u.stream.synchronize(); + } +}; + +/// Pinned bytes' uploads and the fence that says when the bytes may change (`Q`: `Uploads`, or a test's stand-in). +pub fn Staging(comptime Q: type) type { + return struct { + const Self = @This(); + q: Q, + /// An upload was queued after the event's last record (a copy or the record failed): only a drain covers it. + open: bool = false, + + /// Queues a copy from the staged bytes; until `fence` records behind it, `ready` drains the stream. + pub fn upload(g: *Self, dst: u64, bytes: []const u8) !void { + g.open = true; + try g.q.copy(dst, bytes); + } + + /// Records the event behind every upload queued so far. + pub fn fence(g: *Self) !void { + try g.q.record(); + g.open = false; + } + + /// Returns once no queued upload reads the staged bytes: the event's wait, or a drain if one came after it. + pub fn ready(g: *Self) !void { + if (!g.open) return g.q.wait(); + try g.q.drain(); + g.open = false; + } + }; +} + /// Device scratch for chunks of up to `cap` rows. pub const Scratch = struct { cap: usize, @@ -201,8 +250,8 @@ pub const Scratch = struct { dummy: u64 = 0, /// Pinned host copies of what a pass uploads (ids, positions, n-gram rows and sums), owned until the copies land. host: ?cuda.HostBuffer = null, - /// Recorded on the compute stream after a pass's last upload: the next pass waits on it before writing `host`. - staged: ?cuda.Event = null, + /// The uploads from `host` and the event recorded after a pass's last one: the next pass waits before writing it. + staging: ?Staging(Uploads) = null, fn add(s: *Scratch, gpa: std.mem.Allocator, d: *const cuda.Driver, len: usize) !u64 { var b = try cuda.DeviceBuffer.alloc(d, @max(len, 256)); @@ -220,13 +269,18 @@ pub const Scratch = struct { } /// Waits until the last pass's uploads have read `host` (they are queued behind its kernels), so it may change. - pub fn ready(s: *const Scratch) !void { - if (s.staged) |e| try e.synchronize(); + pub fn ready(s: *Scratch) !void { + if (s.staging) |*g| try g.ready(); + } + + /// A queued copy from `host` to the device, which must stay unchanged until it lands. + pub fn upload(s: *Scratch, dst: u64, src: []const u8) !void { + try s.staging.?.upload(dst, src); } /// Marks `host` read once the stream reaches here: every upload queued before this point has copied its bytes. - pub fn uploaded(s: *const Scratch, stream: cuda.Stream) !void { - if (s.staged) |e| try e.record(stream); + pub fn uploaded(s: *Scratch) !void { + if (s.staging) |*g| try g.fence(); } /// The staged ids [n], positions [n] (in 16-byte rows), n-gram rows [n, dims] and their 32-group sums. @@ -331,12 +385,16 @@ pub const Scratch = struct { try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); try stream.synchronize(); s.host = try cuda.HostBuffer.alloc(d, hostAt(cap)[4]); - s.staged = try cuda.Event.init(d, false); + s.staging = .{ .q = .{ .stream = stream, .event = try cuda.Event.init(d, false) } }; return s; } pub fn free(s: *Scratch, gpa: std.mem.Allocator) void { - if (s.staged) |*e| e.deinit(); + if (s.staging) |*g| { + // no queued upload may still read `host`; a wait that fails has lost the context, and its copies with it + g.ready() catch {}; + g.q.event.deinit(); + } if (s.host) |*h| h.free(); for (s.bufs.items) |*b| b.free(); s.bufs.deinit(gpa); @@ -384,3 +442,56 @@ pub const Model = struct { head: Q4, ngram: *const ple.Ngram, }; + +/// A stand-in for the stream: which copies are still queued, and a failure at a chosen copy or at the record. +const FakeQueue = struct { + queued: usize = 0, + covered: usize = 0, + copies: usize = 0, + fail_copy: ?usize = null, + fail_record: bool = false, + + pub fn copy(q: *FakeQueue, _: u64, _: []const u8) !void { + q.copies += 1; + if (q.fail_copy) |k| if (k == q.copies) return error.CopyFailed; + q.queued += 1; + } + pub fn record(q: *FakeQueue) !void { + if (q.fail_record) return error.RecordFailed; + q.covered = q.queued; + } + pub fn wait(q: *FakeQueue) !void { + q.queued -= q.covered; // the event's copies have landed; any queued after it may still be reading + q.covered = 0; + } + pub fn drain(q: *FakeQueue) !void { + q.queued = 0; + q.covered = 0; + } +}; + +test "a pass whose second upload or whose record fails leaves no copy reading the staging when it is reused" { + const Case = struct { fail_copy: ?usize, fail_record: bool }; + for ([_]Case{ .{ .fail_copy = 2, .fail_record = false }, .{ .fail_copy = null, .fail_record = true }, .{ .fail_copy = null, .fail_record = false } }) |c| { + var fq: FakeQueue = .{}; + var g: Staging(*FakeQueue) = .{ .q = &fq }; + // a pass that stages and fences cleanly, then one that fails part way, as `stage()` uploads ids and positions + try g.ready(); + try g.upload(0, "ids"); + try g.upload(0, "pos"); + try g.fence(); + fq.fail_copy = if (c.fail_copy) |k| fq.copies + k else null; + fq.fail_record = c.fail_record; + try g.ready(); + const failed = blk: { + g.upload(0, "ids") catch break :blk true; + g.upload(0, "pos") catch break :blk true; + g.fence() catch break :blk true; + break :blk false; + }; + try std.testing.expectEqual(c.fail_copy != null or c.fail_record, failed); + // the next pass (or `free`) waits before it writes the bytes: nothing may still be queued + try g.ready(); + try std.testing.expectEqual(@as(usize, 0), fq.queued); + } +} From 3f069e3feafd06fbcff7dccc556bda40472e02ee Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:41:48 -0700 Subject: [PATCH 050/115] Add lossless NVFP4 compressed scale planes and TP slices --- zig/src/core/nvfp4_csf.md | 36 ++++++++ zig/src/core/nvfp4_csf.zig | 90 ++++++++++++++++++++ zig/src/core/nvfp4_csf_test.zig | 145 ++++++++++++++++++++++++++++++++ zig/src/core/root.zig | 2 + 4 files changed, 273 insertions(+) create mode 100644 zig/src/core/nvfp4_csf.md create mode 100644 zig/src/core/nvfp4_csf.zig create mode 100644 zig/src/core/nvfp4_csf_test.zig diff --git a/zig/src/core/nvfp4_csf.md b/zig/src/core/nvfp4_csf.md new file mode 100644 index 0000000000..02e2cdf40e --- /dev/null +++ b/zig/src/core/nvfp4_csf.md @@ -0,0 +1,36 @@ +# Lossless NVFP4 scale planes + +`nvfp4_csf.Plane` reads `byte-window4-fixed-stream-u24-exceptions/1` storage. +It preserves the original scale bytes, including signed zeros and NaN encodings; +it performs no floating-point conversion or requantization. + +Each 16-row slab contains 16 byte bases, followed by row-major pairs of four-bit +offsets. The even column occupies the low nibble. A decoded byte is its row base +plus its offset. Bases cannot exceed 240. Exceptions are little-endian uint32 +words: the low 24 bits name a row-major position, and the high byte replaces its +value. Positions must be unique, sorted, and inside the logical plane. + +Rows must be a positive multiple of 16, columns a positive multiple of two, and +the plane can contain at most 2^24 bytes. Stream extents must match exactly. +`Plane.init` validates the entire plane before returning a borrowed view. The +caller keeps the streams alive and immutable while that view is used. + +```zig +const plane = try core.nvfp4_csf.Plane.init(rows, columns, fixed_bytes, exception_bytes); +try plane.decode(rank_scales, row_begin, row_count, column_begin, column_count); +``` + +`decode` writes exactly the selected rectangle in row-major order, with no heap +allocation. Both row and column partitions can end inside a storage slab or a +packed byte. Invalid output extents and overlap with either borrowed stream are +rejected before writing. Sorted exceptions are searched per selected row, so a +rank does not need to materialize another rank's scales. + +This is a host format primitive. It does not change checkpoint admission, +projection arithmetic, CUDA kernels, or the runtime capability advertisement. +Callers must validate safetensors dtypes/shapes and select this exact codec before +constructing the view. + +The public format reference is the [byte-window implementation in b12x](https://github.com/local-inference-lab/b12x/blob/2cc7f66aaef76a9c7cbb660a1efb865cf3f26c88/b12x/_lib/quant/nvfp4_csf.py). +Tests use independently generated synthetic source bytes, handwritten slab +examples, every replacement byte, both partition axes, and malformed streams. diff --git a/zig/src/core/nvfp4_csf.zig b/zig/src/core/nvfp4_csf.zig new file mode 100644 index 0000000000..49d6e24319 --- /dev/null +++ b/zig/src/core/nvfp4_csf.zig @@ -0,0 +1,90 @@ +//! Lossless byte-window4 scale planes, with sorted u24-position/u8-value exceptions. + +const std = @import("std"); + +pub const codec = "byte-window4-fixed-stream-u24-exceptions/1"; +pub const Error = error{ BadGeometry, BadFixedStream, BadExceptions, BadSlice, OverlappingOutput }; + +/// A validated, borrowed checkpoint plane; callers keep both streams immutable until decoding finishes. +pub const Plane = struct { + rows: usize, + columns: usize, + fixed: []const u8, + exceptions: []const u8, + + /// Slabs hold 16 row bases followed by those rows' low-even/high-odd packed offsets. + pub fn init(rows: usize, columns: usize, fixed: []const u8, exceptions: []const u8) Error!Plane { + if (rows == 0 or rows % 16 != 0 or columns == 0 or columns % 2 != 0) return error.BadGeometry; + if (columns > 1 << 24 or rows > (1 << 24) / columns) return error.BadGeometry; + const stride = 16 + 8 * columns; + if (fixed.len != rows / 16 * stride) return error.BadFixedStream; + for (0..rows / 16) |slab| { + for (fixed[slab * stride ..][0..16]) |base| if (base > 240) return error.BadFixedStream; + } + if (exceptions.len % 4 != 0) return error.BadExceptions; + var previous: ?u32 = null; + var offset: usize = 0; + while (offset < exceptions.len) : (offset += 4) { + const position = word(exceptions, offset) & 0xffffff; + if (position >= rows * columns) return error.BadExceptions; + if (previous) |p| if (position <= p) return error.BadExceptions; + previous = position; + } + return .{ .rows = rows, .columns = columns, .fixed = fixed, .exceptions = exceptions }; + } + + /// Decode a rectangular TP slice into exactly sized row-major bytes, without allocating or rounding. + pub fn decode(self: Plane, out: []u8, row: usize, rows: usize, column: usize, columns: usize) Error!void { + if (row > self.rows or rows > self.rows - row or column > self.columns or columns > self.columns - column) + return error.BadSlice; + if (out.len != rows * columns) return error.BadSlice; + if (overlaps(out, self.fixed) or overlaps(out, self.exceptions)) return error.OverlappingOutput; + if (rows == 0 or columns == 0) return; + const stride = 16 + 8 * self.columns; + for (0..rows) |r| { + const source_row = row + r; + const slab = source_row / 16 * stride; + const base = self.fixed[slab + source_row % 16]; + const codes = slab + 16 + source_row % 16 * (self.columns / 2); + for (0..columns) |c| { + const source_column = column + c; + const codes_byte = self.fixed[codes + source_column / 2]; + const shift: u3 = if (source_column % 2 == 0) 0 else 4; + out[r * columns + c] = base + ((codes_byte >> shift) & 15); + } + var i = self.lowerBound(source_row * self.columns + column); + const end = source_row * self.columns + column + columns; + while (i < self.exceptions.len) : (i += 4) { + const entry = word(self.exceptions, i); + const position = entry & 0xffffff; + if (position >= end) break; + out[r * columns + position - (source_row * self.columns + column)] = @intCast(entry >> 24); + } + } + } + + fn lowerBound(self: Plane, position: usize) usize { + var lo: usize = 0; + var hi = self.exceptions.len / 4; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (word(self.exceptions, mid * 4) & 0xffffff < position) lo = mid + 1 else hi = mid; + } + return lo * 4; + } +}; + +fn word(bytes: []const u8, offset: usize) u32 { + return std.mem.readInt(u32, bytes[offset..][0..4], .little); +} + +fn overlaps(out: []u8, source: []const u8) bool { + if (out.len == 0 or source.len == 0) return false; + const a = @intFromPtr(out.ptr); + const b = @intFromPtr(source.ptr); + return if (a <= b) b - a < out.len else a - b < source.len; +} + +test { + _ = @import("nvfp4_csf_test.zig"); +} diff --git a/zig/src/core/nvfp4_csf_test.zig b/zig/src/core/nvfp4_csf_test.zig new file mode 100644 index 0000000000..13e63c88b7 --- /dev/null +++ b/zig/src/core/nvfp4_csf_test.zig @@ -0,0 +1,145 @@ +//! Independent byte fixtures and malformed-storage boundaries for lossless scale decoding. + +const std = @import("std"); +const Plane = @import("nvfp4_csf.zig").Plane; +const expect = std.testing.expectEqualSlices; + +test "handwritten slab preserves nibble order, bases and exception bytes" { + var fixed: [32]u8 = @splat(0); + fixed[0] = 240; + fixed[1] = 1; + fixed[16] = 0xf0; + fixed[17] = 0x21; + const exceptions = [_]u8{ 2, 0, 0, 255, 31, 0, 0, 128 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 2); + try expect(u8, &.{ 240, 255, 255, 3 }, out[0..4]); + try std.testing.expectEqual(@as(u8, 128), out[31]); + var odd: [2]u8 = undefined; + try plane.decode(&odd, 0, 2, 1, 1); + try expect(u8, &.{ 255, 3 }, &odd); +} + +const Fixture = struct { + fixed: []u8, + exceptions: []u8, + source: []u8, + rows: usize, + columns: usize, + + fn init(rows: usize, columns: usize) !Fixture { + const gpa = std.testing.allocator; + const source = try gpa.alloc(u8, rows * columns); + errdefer gpa.free(source); + const fixed = try gpa.alloc(u8, rows + source.len / 2); + errdefer gpa.free(fixed); + const exceptions = try gpa.alloc(u8, ((source.len + 96) / 97) * 4); + errdefer gpa.free(exceptions); + @memset(fixed, 0); + var cursor: usize = 0; + for (0..rows) |row| { + const base: u8 = @intCast((row * 37 + 19) % 241); + const slab = row / 16 * (16 + columns * 8); + fixed[slab + row % 16] = base; + for (0..columns) |column| { + const i = row * columns + column; + const delta: u8 = @intCast((column * 11 + row * 3) % 16); + source[i] = base + delta; + const offset = slab + 16 + row % 16 * (columns / 2) + column / 2; + fixed[offset] |= delta << @as(u3, if (column % 2 == 0) 0 else 4); + if (i % 97 == 0) { + source[i] = @intCast((i / 97) % 256); + const record = @as(u32, @intCast(i)) | (@as(u32, source[i]) << 24); + std.mem.writeInt(u32, exceptions[cursor..][0..4], record, .little); + cursor += 4; + } + } + } + return .{ .fixed = fixed, .exceptions = exceptions, .source = source, .rows = rows, .columns = columns }; + } + + fn deinit(self: Fixture) void { + const gpa = std.testing.allocator; + gpa.free(self.fixed); + gpa.free(self.exceptions); + gpa.free(self.source); + } + + fn check(self: Fixture, row: usize, rows: usize, column: usize, columns: usize) !void { + const out = try std.testing.allocator.alloc(u8, rows * columns); + defer std.testing.allocator.free(out); + @memset(out, 0xd6); + const plane = try Plane.init(self.rows, self.columns, self.fixed, self.exceptions); + try plane.decode(out, row, rows, column, columns); + for (0..rows) |r| try expect(u8, self.source[(row + r) * self.columns + column ..][0..columns], out[r * columns ..][0..columns]); + } +}; + +test "full planes and row/column TP partitions reproduce independent source bytes" { + for ([_][2]usize{ .{ 16, 2 }, .{ 32, 20 }, .{ 640, 160 }, .{ 2560, 40 } }) |shape| { + const fixture = try Fixture.init(shape[0], shape[1]); + defer fixture.deinit(); + try fixture.check(0, shape[0], 0, shape[1]); + for (0..2) |rank| { + try fixture.check(rank * shape[0] / 2, shape[0] / 2, 0, shape[1]); + try fixture.check(0, shape[0], rank * shape[1] / 2, shape[1] / 2); + } + try fixture.check(15, 1, 1, shape[1] - 1); + } +} + +test "all 256 byte values survive exception decoding unchanged" { + const fixed: [144]u8 = @splat(0); + var exceptions: [1024]u8 = undefined; + for (0..256) |i| std.mem.writeInt(u32, exceptions[i * 4 ..][0..4], @as(u32, @intCast(i)) * 0x1000001, .little); + const plane = try Plane.init(16, 16, &fixed, &exceptions); + var out: [256]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 16); + for (out, 0..) |value, i| try std.testing.expectEqual(@as(u8, @intCast(i)), value); +} + +test "the last u24 position is valid and larger planes are refused" { + const columns = 1 << 20; + const fixed = try std.testing.allocator.alloc(u8, 16 + 8 * columns); + defer std.testing.allocator.free(fixed); + @memset(fixed, 0); + const plane = try Plane.init(16, columns, fixed, &.{ 255, 255, 255, 173 }); + var out: [1]u8 = undefined; + try plane.decode(&out, 15, 1, columns - 1, 1); + try expect(u8, &.{173}, &out); + try std.testing.expectError(error.BadGeometry, Plane.init(32, columns, fixed, &.{})); +} + +test "bad geometry, truncated streams, overflow bases and unordered exceptions are refused" { + var fixed: [32]u8 = @splat(0); + const bad = std.testing.expectError; + for ([_][2]usize{ .{ 0, 2 }, .{ 1, 2 }, .{ 16, 0 }, .{ 16, 3 }, .{ 16, 1 << 24 }, .{ 1 << 24, 2 } }) |shape| + try bad(error.BadGeometry, Plane.init(shape[0], shape[1], &fixed, &.{})); + try bad(error.BadFixedStream, Plane.init(16, 2, fixed[0..31], &.{})); + fixed[15] = 241; + try bad(error.BadFixedStream, Plane.init(16, 2, &fixed, &.{})); + fixed[15] = 0; + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{0})); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 32, 0, 0, 1 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 1, 0, 0, 2, 1, 0, 0, 3 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 2, 0, 0, 2, 1, 0, 0, 3 })); +} + +test "invalid slices and aliased storage fail before output changes" { + var fixed: [32]u8 = @splat(0); + var exceptions = [_]u8{ 0, 0, 0, 1 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = @splat(0xd6); + const bad = std.testing.expectError; + try bad(error.BadSlice, plane.decode(&out, 17, 0, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 17, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 3, 0)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 1, 2)); + try bad(error.BadSlice, plane.decode(out[0..31], 0, 16, 0, 2)); + const poison: [32]u8 = @splat(0xd6); + try expect(u8, &poison, &out); + try bad(error.OverlappingOutput, plane.decode(&fixed, 0, 16, 0, 2)); + try bad(error.OverlappingOutput, plane.decode(&exceptions, 0, 2, 0, 2)); + try plane.decode(out[0..0], 16, 0, 2, 0); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 736248a93d..4ffd5c0afb 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -14,6 +14,7 @@ pub const affine4_host = @import("affine4_host.zig"); pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); +pub const nvfp4_csf = @import("nvfp4_csf.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); @@ -25,6 +26,7 @@ test { _ = affine4_fold; _ = shard_edit; _ = quant; + _ = nvfp4_csf; _ = registry; _ = tuning; _ = norms_ref; From 33d6494b63fa8249cb4311bef5b7ecd7696a8539 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 13 Oct 2026 06:37:57 +0900 Subject: [PATCH 051/115] core: admit ModelOpt NVFP4 checkpoints from config.json's quantization block The GLM-5.3-Flash CUDA port reads nvidia/GLM-5.3-Flash-NVFP4 as stored. admit() checks config.json's quantization_config (top level, else text_config's) for quant_method modelopt, quant_algo NVFP4 and 4-bit float weights in groups of 16, and refuses anything else with a reason the caller logs (the Why of nemotron/config.zig), so a test reads the reason instead of failing on an error log. MIXED_PRECISION is refused for now. Which linears are NVFP4 is left to the tensors, not to the block's ignore list: the pinned revision 423acf37 stores the MTP layer's experts in BF16 without listing layer 45 there (da920bb0 added it). The fixture is that revision's config.json. --- .../fixtures/glm53_flash_nvfp4/config.json | 469 ++++++++++++++++++ zig/src/core/modelopt.zig | 112 +++++ zig/src/core/root.zig | 2 + 3 files changed, 583 insertions(+) create mode 100644 zig/src/core/fixtures/glm53_flash_nvfp4/config.json create mode 100644 zig/src/core/modelopt.zig diff --git a/zig/src/core/fixtures/glm53_flash_nvfp4/config.json b/zig/src/core/fixtures/glm53_flash_nvfp4/config.json new file mode 100644 index 0000000000..d0e8e75de4 --- /dev/null +++ b/zig/src/core/fixtures/glm53_flash_nvfp4/config.json @@ -0,0 +1,469 @@ +{ + "architectures": [ + "Glm5NextForConditionalGeneration" + ], + "dtype": "bfloat16", + "image_end_token_id": 154831, + "image_start_token_id": 154830, + "image_token_id": 154854, + "model_type": "glm5_next", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "base_model_ep_plan": { + "layers.*.mlp.experts": "torch_a2a_experts", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.gate": "torch_a2a_router" + }, + "bos_token_id": null, + "dtype": "bfloat16", + "eos_token_id": [ + 154820, + 154827, + 154829 + ], + "first_k_dense_replace": 3, + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 0, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_kpool": 4, + "index_kpool_always_select_tail": true, + "index_kpool_compress": true, + "index_n_heads": 32, + "index_share_for_mtp_iteration": true, + "index_topk": 2048, + "indexer_rope_interleave": true, + "indexer_types": [ + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full", + "full" + ], + "initializer_range": 0.02, + "intermediate_size": 12288, + "kv_lora_rank": 512, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention" + ], + "linear_attn_config": { + "full_attn_layers": [ + 3, + 7, + 11, + 15, + 19, + 23, + 27, + 31, + 35, + 39, + 43 + ], + "gate_lower_bound": -5.0, + "head_dim": 128, + "kda_layers": [ + 0, + 1, + 2, + 4, + 5, + 6, + 8, + 9, + 10, + 12, + 13, + 14, + 16, + 17, + 18, + 20, + 21, + 22, + 24, + 25, + 26, + 28, + 29, + 30, + 32, + 33, + 34, + 36, + 37, + 38, + 40, + 41, + 42, + 44 + ], + "num_heads": 64, + "short_conv_kernel_size": 4 + }, + "linear_conv_kernel_dim": 4, + "linear_head_dim": 128, + "linear_lower_bound": -5.0, + "linear_num_heads": 64, + "max_position_embeddings": 1048576, + "mhc": true, + "mla_use_nope": true, + "mlp_layer_types": [ + "dense", + "dense", + "dense", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse", + "sparse" + ], + "model_type": "glm5_next_text", + "moe_intermediate_size": 2048, + "moe_router_dtype": "float32", + "n_group": 1, + "n_routed_experts": 288, + "n_shared_experts": 1, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts_per_tok": 8, + "num_hidden_layers": 45, + "num_key_value_heads": 64, + "num_nextn_predict_layers": 1, + "output_router_logits": false, + "pad_token_id": 154820, + "q_lora_rank": 1536, + "qk_head_dim": 256, + "qk_nope_head_dim": 256, + "qk_rope_head_dim": 0, + "rms_norm_eps": 1e-05, + "routed_scaling_factor": 2.5, + "router_aux_loss_coef": 0.001, + "scoring_func": "sigmoid", + "swiglu_limit": 10.0, + "tie_word_embeddings": false, + "topk_group": 1, + "topk_method": "noaux_tc", + "use_cache": true, + "v_head_dim": 256, + "vocab_size": 154880 + }, + "tie_word_embeddings": false, + "transformers_version": "5.16.1", + "video_end_token_id": 154833, + "video_start_token_id": 154832, + "video_token_id": 154855, + "vision_config": { + "attention_bias": true, + "attention_dropout": 0.0, + "depth": 24, + "dtype": "bfloat16", + "hidden_act": "silu", + "hidden_size": 1024, + "image_size": 448, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "glm5_next_vision", + "num_heads": 16, + "out_hidden_size": 4096, + "patch_size": 14, + "projection_intermediate_size": 10240, + "rms_norm_eps": 1e-05, + "spatial_merge_size": 2, + "swiglu_limit": 10.0, + "temporal_patch_size": 2 + }, + "quantization_config": { + "config_groups": { + "group_0": { + "input_activations": { + "dynamic": false, + "num_bits": 4, + "type": "float", + "group_size": 16 + }, + "weights": { + "dynamic": false, + "num_bits": 4, + "type": "float", + "group_size": 16 + }, + "targets": [ + "Linear" + ] + } + }, + "ignore": [ + "lm_head", + "model.language_model.embed_tokens", + "model.language_model.layers.0.self_attn*", + "model.language_model.layers.1.self_attn*", + "model.language_model.layers.10.mlp.gate", + "model.language_model.layers.10.mlp.shared_experts*", + "model.language_model.layers.10.self_attn*", + "model.language_model.layers.11.mlp.gate", + "model.language_model.layers.11.mlp.shared_experts*", + "model.language_model.layers.11.self_attn*", + "model.language_model.layers.12.mlp.gate", + "model.language_model.layers.12.mlp.shared_experts*", + "model.language_model.layers.12.self_attn*", + "model.language_model.layers.13.mlp.gate", + "model.language_model.layers.13.mlp.shared_experts*", + "model.language_model.layers.13.self_attn*", + "model.language_model.layers.14.mlp.gate", + "model.language_model.layers.14.mlp.shared_experts*", + "model.language_model.layers.14.self_attn*", + "model.language_model.layers.15.mlp.gate", + "model.language_model.layers.15.mlp.shared_experts*", + "model.language_model.layers.15.self_attn*", + "model.language_model.layers.16.mlp.gate", + "model.language_model.layers.16.mlp.shared_experts*", + "model.language_model.layers.16.self_attn*", + "model.language_model.layers.17.mlp.gate", + "model.language_model.layers.17.mlp.shared_experts*", + "model.language_model.layers.17.self_attn*", + "model.language_model.layers.18.mlp.gate", + "model.language_model.layers.18.mlp.shared_experts*", + "model.language_model.layers.18.self_attn*", + "model.language_model.layers.19.mlp.gate", + "model.language_model.layers.19.mlp.shared_experts*", + "model.language_model.layers.19.self_attn*", + "model.language_model.layers.2.self_attn*", + "model.language_model.layers.20.mlp.gate", + "model.language_model.layers.20.mlp.shared_experts*", + "model.language_model.layers.20.self_attn*", + "model.language_model.layers.21.mlp.gate", + "model.language_model.layers.21.mlp.shared_experts*", + "model.language_model.layers.21.self_attn*", + "model.language_model.layers.22.mlp.gate", + "model.language_model.layers.22.mlp.shared_experts*", + "model.language_model.layers.22.self_attn*", + "model.language_model.layers.23.mlp.gate", + "model.language_model.layers.23.mlp.shared_experts*", + "model.language_model.layers.23.self_attn*", + "model.language_model.layers.24.mlp.gate", + "model.language_model.layers.24.mlp.shared_experts*", + "model.language_model.layers.24.self_attn*", + "model.language_model.layers.25.mlp.gate", + "model.language_model.layers.25.mlp.shared_experts*", + "model.language_model.layers.25.self_attn*", + "model.language_model.layers.26.mlp.gate", + "model.language_model.layers.26.mlp.shared_experts*", + "model.language_model.layers.26.self_attn*", + "model.language_model.layers.27.mlp.gate", + "model.language_model.layers.27.mlp.shared_experts*", + "model.language_model.layers.27.self_attn*", + "model.language_model.layers.28.mlp.gate", + "model.language_model.layers.28.mlp.shared_experts*", + "model.language_model.layers.28.self_attn*", + "model.language_model.layers.29.mlp.gate", + "model.language_model.layers.29.mlp.shared_experts*", + "model.language_model.layers.29.self_attn*", + "model.language_model.layers.3.mlp.gate", + "model.language_model.layers.3.mlp.shared_experts*", + "model.language_model.layers.3.self_attn*", + "model.language_model.layers.30.mlp.gate", + "model.language_model.layers.30.mlp.shared_experts*", + "model.language_model.layers.30.self_attn*", + "model.language_model.layers.31.mlp.gate", + "model.language_model.layers.31.mlp.shared_experts*", + "model.language_model.layers.31.self_attn*", + "model.language_model.layers.32.mlp.gate", + "model.language_model.layers.32.mlp.shared_experts*", + "model.language_model.layers.32.self_attn*", + "model.language_model.layers.33.mlp.gate", + "model.language_model.layers.33.mlp.shared_experts*", + "model.language_model.layers.33.self_attn*", + "model.language_model.layers.34.mlp.gate", + "model.language_model.layers.34.mlp.shared_experts*", + "model.language_model.layers.34.self_attn*", + "model.language_model.layers.35.mlp.gate", + "model.language_model.layers.35.mlp.shared_experts*", + "model.language_model.layers.35.self_attn*", + "model.language_model.layers.36.mlp.gate", + "model.language_model.layers.36.mlp.shared_experts*", + "model.language_model.layers.36.self_attn*", + "model.language_model.layers.37.mlp.gate", + "model.language_model.layers.37.mlp.shared_experts*", + "model.language_model.layers.37.self_attn*", + "model.language_model.layers.38.mlp.gate", + "model.language_model.layers.38.mlp.shared_experts*", + "model.language_model.layers.38.self_attn*", + "model.language_model.layers.39.mlp.gate", + "model.language_model.layers.39.mlp.shared_experts*", + "model.language_model.layers.39.self_attn*", + "model.language_model.layers.4.mlp.gate", + "model.language_model.layers.4.mlp.shared_experts*", + "model.language_model.layers.4.self_attn*", + "model.language_model.layers.40.mlp.gate", + "model.language_model.layers.40.mlp.shared_experts*", + "model.language_model.layers.40.self_attn*", + "model.language_model.layers.41.mlp.gate", + "model.language_model.layers.41.mlp.shared_experts*", + "model.language_model.layers.41.self_attn*", + "model.language_model.layers.42.mlp.gate", + "model.language_model.layers.42.mlp.shared_experts*", + "model.language_model.layers.42.self_attn*", + "model.language_model.layers.43.mlp.gate", + "model.language_model.layers.43.mlp.shared_experts*", + "model.language_model.layers.43.self_attn*", + "model.language_model.layers.44.mlp.gate", + "model.language_model.layers.44.mlp.shared_experts*", + "model.language_model.layers.44.self_attn*", + "model.language_model.layers.5.mlp.gate", + "model.language_model.layers.5.mlp.shared_experts*", + "model.language_model.layers.5.self_attn*", + "model.language_model.layers.6.mlp.gate", + "model.language_model.layers.6.mlp.shared_experts*", + "model.language_model.layers.6.self_attn*", + "model.language_model.layers.7.mlp.gate", + "model.language_model.layers.7.mlp.shared_experts*", + "model.language_model.layers.7.self_attn*", + "model.language_model.layers.8.mlp.gate", + "model.language_model.layers.8.mlp.shared_experts*", + "model.language_model.layers.8.self_attn*", + "model.language_model.layers.9.mlp.gate", + "model.language_model.layers.9.mlp.shared_experts*", + "model.language_model.layers.9.self_attn*", + "model.visual*" + ], + "quant_algo": "NVFP4", + "kv_cache_scheme": { + "dynamic": false, + "num_bits": 8, + "type": "float" + }, + "producer": { + "name": "modelopt", + "version": "0.47.0.dev393+ga4bc45b30.d20260828" + }, + "quant_method": "modelopt" + } +} \ No newline at end of file diff --git a/zig/src/core/modelopt.zig b/zig/src/core/modelopt.zig new file mode 100644 index 0000000000..5f4e0c9e80 --- /dev/null +++ b/zig/src/core/modelopt.zig @@ -0,0 +1,112 @@ +//! ModelOpt's quantization block in config.json: admits NVFP4 weights in groups of 16, refuses the rest by name. +//! Which linears are NVFP4 comes from the tensors (nvfp4.stored), not from "ignore": nvidia/GLM-5.3-Flash-NVFP4 +//! at 423acf37 stores the MTP layer in BF16 without listing it there (da920bb0 added it). + +const std = @import("std"); + +/// The NVFP4 group size ModelOpt writes for this format. +pub const nvfp4_group = 16; + +/// Why the check refused the checkpoint, kept for the caller's log line (tests read it instead). +pub const Why = struct { + buf: [320]u8 = undefined, + len: usize = 0, + + /// Keep the reason, cut short if it is long. + pub fn set(self: *Why, comptime fmt: []const u8, args: anytype) void { + var w: std.Io.Writer = .fixed(&self.buf); + w.print(fmt, args) catch {}; + self.len = w.end; + } + + pub fn text(self: *const Why) []const u8 { + return self.buf[0..self.len]; + } +}; + +/// config.json's quantization_config (top level, else text_config's) checked as ModelOpt NVFP4. +pub fn admit(gpa: std.mem.Allocator, config_json: []const u8, why: *Why) !void { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const root = try std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), config_json, .{}); + const q = block(root) orelse { + why.set("config.json has no ModelOpt quantization_config", .{}); + return error.NotModelOpt; + }; + const algo = text(q.get("quant_algo")) orelse ""; + if (!std.ascii.eqlIgnoreCase(algo, "NVFP4")) { + why.set("ModelOpt {s}: this engine reads ModelOpt NVFP4 (nvidia/GLM-5.3-Flash-NVFP4's layout)", .{if (algo.len > 0) algo else "without quant_algo"}); + return error.UnsupportedQuant; + } + const groups = object(q.get("config_groups")) orelse { + why.set("ModelOpt NVFP4 without config_groups", .{}); + return error.UnsupportedQuant; + }; + var it = groups.iterator(); + while (it.next()) |g| { + const w = if (object(g.value_ptr.*)) |o| object(o.get("weights")) else null; + const bits = if (w) |o| int(o.get("num_bits")) else null; + const size = if (w) |o| int(o.get("group_size")) else null; + const kind = if (w) |o| text(o.get("type")) orelse "" else ""; + if (bits != 4 or size != nvfp4_group or !std.mem.eql(u8, kind, "float")) { + why.set("ModelOpt NVFP4 group {s}: weights {?d}-bit {s} in groups of {?d}, expected 4-bit float in groups of 16", .{ g.key_ptr.*, bits, kind, size }); + return error.UnsupportedQuant; + } + } +} + +fn block(root: std.json.Value) ?std.json.ObjectMap { + const top = object(root) orelse return null; + for ([_]?std.json.ObjectMap{ top, object(top.get("text_config")) }) |source| { + const q = object((source orelse continue).get("quantization_config")) orelse continue; + if (std.ascii.eqlIgnoreCase(text(q.get("quant_method")) orelse "", "modelopt")) return q; + } + return null; +} + +fn object(v: ?std.json.Value) ?std.json.ObjectMap { + const x = v orelse return null; + return if (x == .object) x.object else null; +} + +fn text(v: ?std.json.Value) ?[]const u8 { + const x = v orelse return null; + return if (x == .string) x.string else null; +} + +fn int(v: ?std.json.Value) ?i64 { + const x = v orelse return null; + return if (x == .integer) x.integer else null; +} + +test "nvidia/GLM-5.3-Flash-NVFP4's config is admitted" { + var why: Why = .{}; + try admit(std.testing.allocator, @embedFile("fixtures/glm53_flash_nvfp4/config.json"), &why); +} + +/// admit() must refuse `json` with `err`, its reason naming `word`. +fn expectRefused(json: []const u8, err: anyerror, word: []const u8) !void { + var why: Why = .{}; + try std.testing.expectError(err, admit(std.testing.allocator, json, &why)); + if (std.mem.indexOf(u8, why.text(), word) == null) { + std.debug.print("refusal \"{s}\" does not name \"{s}\"\n", .{ why.text(), word }); + return error.TestUnexpectedResult; + } +} + +test "other ModelOpt formats, and configs without ModelOpt, are refused by name" { + try expectRefused("{\"text_config\": {}}", error.NotModelOpt, "no ModelOpt"); + try expectRefused("{\"quantization_config\": {\"quant_method\": \"exl3\"}}", error.NotModelOpt, "no ModelOpt"); + try expectRefused("{\"quantization_config\": {\"quant_method\": \"modelopt\", \"quant_algo\": \"FP8\"}}", error.UnsupportedQuant, "ModelOpt FP8"); + // MIXED_PRECISION (the AXL re-pack) comes with 3.x + try expectRefused("{\"quantization_config\": {\"quant_method\": \"modelopt\", \"quant_algo\": \"MIXED_PRECISION\"}}", error.UnsupportedQuant, "MIXED_PRECISION"); + try expectRefused( + \\{"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "config_groups": + \\ {"g": {"weights": {"num_bits": 4, "type": "float", "group_size": 32}}}}} + , error.UnsupportedQuant, "groups of 32"); + var why: Why = .{}; + try admit(std.testing.allocator, + \\{"text_config": {"quantization_config": {"quant_method": "modelopt", "quant_algo": "nvfp4", "config_groups": + \\ {"g": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}}}}}} + , &why); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index b6c0d6cf4c..0c1d43adf5 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -19,6 +19,7 @@ pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); pub const act_ref = @import("act_ref.zig"); pub const route_ref = @import("route_ref.zig"); +pub const modelopt = @import("modelopt.zig"); test { _ = safetensors; @@ -39,4 +40,5 @@ test { _ = draft_depth; _ = @import("copy_index.zig"); _ = ids_json; + _ = modelopt; } From 9d4a6fa6434c509d417c2b13c82af6877e88044e Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 13 Oct 2026 06:37:57 +0900 Subject: [PATCH 052/115] core: read ModelOpt NVFP4 linears as stored, with an f32 reference row nvfp4.zig is the format the GLM-5.3-Flash CUDA port's W4A16 linears read: e2m1 codes two to a byte (input 2j in byte j's low nibble) [n, k/2], an e4m3 scale per 16 inputs [n, k/16], and weight_scale_2 as an fp32 scalar that multiplies. dims() checks the four tensors fit together; Linear.read() maps them from a Checkpoint and marks input_scale as used, since the W4A4 math's static input scale is not read by W4A16 but a leftover tensor would refuse the checkpoint; row() gives a row as e2m1 x e4m3 x global in f32, the reference the GPU kernels will be compared with. stored() says a module is NVFP4 by its weight_scale tensor. The e2m1 and e4m3 values come from core.quant's nvfp4 and fp8 (#730), as f32. The fixture is the whole header of nvidia/GLM-5.3-Flash-NVFP4 (147,661 names, dtypes and shapes, gzip'd): every NVFP4 linear reads, 42 MoE layers x 288 experts x 3 plus the three dense MLP layers' 3, all under .mlp., and the MTP layer's experts stay BF16. --- .../glm53_flash_nvfp4/headers.json.gz | Bin 0 -> 434836 bytes zig/src/core/modelopt.zig | 4 +- zig/src/core/nvfp4.zig | 148 ++++++++++++++++++ zig/src/core/nvfp4_fixture_test.zig | 69 ++++++++ zig/src/core/root.zig | 3 + 5 files changed, 221 insertions(+), 3 deletions(-) create mode 100644 zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz create mode 100644 zig/src/core/nvfp4.zig create mode 100644 zig/src/core/nvfp4_fixture_test.zig diff --git a/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz b/zig/src/core/fixtures/glm53_flash_nvfp4/headers.json.gz new file mode 100644 index 0000000000000000000000000000000000000000..db569985ec1dc9183e97e5bb1483a85609e7c9e3 GIT binary patch literal 434836 zcmeFa30RZY);3HALLh+z6$zjufFmdgiiIjB0TibyX^RLIO)MgI5>P9MArq4#7Oi8Q z0#z%mR8(xifecVU0R>Tu7LX7`1Qf!Ykm-M(^!xg5{_A^B+jDGdd+2p_WpT4y&))kY z%Fepiy~um9STpE9q;Yf9hK*5Mw|RcJEoO_^Ro%*hEo*k3-^H^pw0wVGob_{egH>)N&2>+`2!-(Sh-3g4w`*Y}<1Hx3h|`X+nh zz{8_boyMr^^pP6Q_ZV7?y?jF{U)yuu*b;4Q|GH7H4(Hu8NDcX7eQA$G)3!#c*ThJ5 zBlX7CDmH(xjBV&zBh{9zk?NE;VhoxK`BMD{Yt%1~8k_SC>X$VPo*~$%E<4(sG(2c8 zRrfXI8=g7(@O4>+&My9NGGoDHttG?olEKq7NcH^5jLs$e!3*M%FQmrXV#5Kcu|urd z!|xyLF;>L0bs8y8e}~K;mKyKz8rfa(QoX*0ueoE?jqGhU_MR}-NY@YWla3pmhlA}$ z9_-Z(9@I&P*^G*sk&T9{UAjBxA5#rtjWL6N+sp^r8yvXm!x7QHZ z_pikdM-RIVpI`U9Cn?yNc&D@G+SBv=w~ejCV*NaQ!&G|zijm=?2+5byczLi^#53M0 zY)&G-%TCfa#Yi+G{RCAZTR*gh|5>rvNM`GXPOy!GXSVYy2}UJ@*K8bdna>~kTHFeq zNAKQduI441ar<;UqYcR)$s()0+qz;M--{h}F>IMm_g8l;C)xG-Zs&uf6V>1^vFb$3i@k|AO$eaL}T! zVZ^@gpuVb|ccuTn$Myc5c1@ZN=>OZ$|9KhRx~Cc4G`l7j_v@(D+b!~++l1Vk{Nyo6~M$J_Ir9Q3K*S)QthGm7_#)R87-&*k1iN+QW zz2u1|ZgCC&)7#SF!-lrbhqYylE)V_9%HH&gX-TJypT4O%D6OvXoMqIT6SxWjUy-#F zwxj$+1tCe16|W%h6dh}Iw%xr;if@x!Jo(QXpfCED6yGDatl~dftC#l-F1h)bY}jIy zi+OG8s&=-e869y5yMHdz5;5OAg~Zy%#3FC1jI`xOLZRA7t27dBBGk0^1cX9MHSH$> zp~zOPSSutHIH-v~3HMsa=ZO^8{bTo0iV}zsZrF=X=z6ROXCZG(Rvh<_)ljg0*b=NL zHkpz+%Js#H2qIcCMNL`ihpoekwkK2ejdFdlpA$sh$rKG`h99;PE83czojJ<&#byvh zbCR>wl+AwFpRl4G$=UlxxxU!r1koGG*%}JT4_l5E#U*nxN4dV(G=j)0nWLtJ`C)&= ziat%|>>K6!Vvi6+^O89lNtd)fg_ARrsZ zJDj}ra1mx&$?OCEdm@%?WV$}G?ajOuvdGU>Np@$QcHd9-V4e0*kY}+@&)QF(%{o0h zo9szE?a3ljiKnT=O0#n|`*DbK;}k0p=O*soi#X?)U4_VTX5o-IuCnl6udv=Yla&L$ zFub<~`qE@U9Q6MSV!ZBR@rxC%$Wo#rMMgw>QfNA}YKw%093|q3O~gq%cCsD&sL6u< zF+28YJN7}71#!>^Vpi4oI_o&{ng{DAf6lx!dg{|h?-By zvN)tMTSy^JLwD|G9<%)o8GyuK^zU27pOeH(;!WFu>8_MLA}MgFNVei4F7p$_%~^_g z^Jm`(nkP)Tn{l60)^SDY9Tum|^R*5=>FKd5*=~Au#=+>a53VSmhL3`Q=#o6&ouMbE zd8|5ZHywcg{Ui*Rkf4;lOqRpAzIOwn-kMQkdD!1&;mip^shxQ&(nK&Z@Vr;P3?Mi7 z$lrA!tpx#*Af9~$B5;hAh_e>bCw zQ`UAxIVUVmo>!C*8Zpho^YkbfjIPeu8(kK4MX3yr1K~e1iPW97!@Tbb^YKJg7cz=X zq2dqiLln;;?!a``F&{Z9UYEUi(OqfN(A@&rF@mBM5jkIBW0CpT%+FCV9rcVO5aYb> zW;AfhUS3hY5f;~)S7aU9=;^T~+0Lgr<3x1XCs&lu!bibCbZK7E&d|nb9$QY^`2g@= zoJ2Y|}CG>lVA};|S>gVa?IkL&$>Y0jc6 zQ{SQ4e!B;~>TIGbVnPdWNWiv%(`bE@^Ml|7G6<`1;4BiHFbgu1wTl^sXC3x7jJp|P zd4{+Owq($+&LG!hh;PGo6h0R7^I3uCQ)6itJR><^@r1s!EU(l#?ja#AcdX3Ufh&r=xEzk?HD-m+*OCOr~puQx@ie=}f_S#W;S5NM>BlVuM>$M~HuD0&A z!}dm*_m&}g?er?Wg23R99B!qG`BTl*XqGGUuGv)d^9jWFnRgH_SQ}?_1y7K6p4|yB%6uNxTLuPLcEn~Lp6e~S|&^P!l`4=Bgo}`3e;&V zS2b5i*ygTssWwMKH&qd(k%D!@g0787kxZ$kVCmR#&`niDX`(E38|B)Fo+eZ7P?plM zR?tmVl-)>~;Wo;(5h;_it0^<+SV!olD#~u6Y<3&v+K8SdXWyZ0reiVCO;yBcq>$W3 zxi+GfWKK1OM8{5qZmJ?q6D7=TlxrhuOXl36gwe6i&hX{69RpupTTGYNE%2Szc@jE{ zAd1EyjFS-Ed+@8Ce&~uC-*V(u_M(eZ=1^>3>`n;U=bm>v1F z9r>WiO7N^svui|_J3=yn^&#RKiS-TQ`Xpi};yVh_9SQy^&<8VsIOzWu#I$D0SZGgI z8XmC~GiD&m@p~gGn>9NVjW9GYJm&f*mZg9KNn`up~HI>IHSgTa$%# z<5UZUbrWu_6V^?vO%m3Ua!(0sq2Kt_j(O_qCaVPel00=^p8B}Sf;i~^&0|Q-&W=@q z@$sAUYS| zlNTsy;c=?bFc8fNo#f|XbI{JBCgV3GQOgB$9QH3kRemfW+swF+*y)#DZSfo-q`9z` zpi45%fW&{~GQixxBY)psemzO=IpO8Gv1>5pF2TPpxyMEQ7{<2sqA)v(mOn^FV7c*|Lr6i z%bYo`hoVy|_wN^vbB_$zxhfZL)v+vx<{lEAT5^d9BnF@Ns+R%g1|Ipl4umAOb379e z|ITK-RqvZQUemUpPbuF|JZ;&g#bf_b^UCi9Z2K8n;c@+WMIijwCefVAx#E4OFs<^Z z{moA9djofF%gwXtXv(?d7gCsBG6KC{e%T_p=vTZKQV%|}T=o)|`Mp_>RpwW+mzJ4F z%(}7aG}#La^$t#Ks2|3|`KVo%Ji`LIwO&#}pRkki5;&oLejd&T?FcoaU@&@9S!sb1 z`rKWWmjJ@Ady@1DREG0@jX-IUz|6)XzaY*?69l0}t}F+k0{bd9LZV_b6J@h4GCwq< zVlv;8%{I?mZsv(r&zH?Ysognop*Jxe>yO&y%QNhuTk9p|r{TShc?larZ~A$xKWIm- z83lvUo64#Rl=H%S`|=V%`0t;DRFRT4b^FX57t6S%n|%Y0+B#+sgHmp~cw4%7VmlTJ zFHL2|qf62RPYBos#HT)#k1>ym%_Mfjc%(Xrve}y?a=;(Pr%^D;m}Hj(IqMn+`yMT9 z7OBV%f4~akX`Y-Wa`Bb3zTjY&(>6~biEQzQv1t@^vbjj8l@WJx6!U4DCy_+f_`{eq z3Mx5G=u#wOg>$g(w6F=JB1`-MG{zH&iW%6qs`-<_ee20FL;JQge=@v}N8X~=eAfK2 zMsr#mByb_gSe6_tlD655Br@U-8&k85qB6d)K}zgkEBdIL)%GGS?yx>J+aL^%;ynkSOjs79JnI9(|8|nznPRVJ+uFJ1*#z-` z!x?(~kI{Io;sY|Hdv1G@PL>5E2A}t;mjUJm9{IZt6IonKu;6OT#u%D|Rrm>$7ckEz4A9}cn{e19oD7j~#EQs9Ge_~qk?EOzdm&M@D7Jg9I zfb96>@WA)X#1m7C_Z@sPaam03Xc&kNuWRt__~gLA8>m6t>=~Ml4E7B5!$kIs;m046sY7M& zk*R~}(~8~pKY4pu5B6;Q=DG#Qj>(4yK4T`PPc2S9_{4Tuk76_oM2FQa@a>pKXs!ejaDZXA?&Lgu`PENe|vX~d@mZ=PBn@_z?cN%)L{O1SK;UtZ1UK67&Stpis z;D)(_N2raxkDcf|seR0_x{3Ee-3?^Ny2Aq(nTbcI6?^P|vU*vM)!F#)x*NV7>kbTf zt{Me{(Nl}}9(-cAtVek^9)$k^N#Y(y);iY@k%}zp7h-zilClt7PD|-{@_n`H7W4!5 z%OQc={-4u0O4pX$TN5maRdsBmPTDp20GEDxVZZu-eXt*2@j;yyU8fB(+0B3-R;P`v z(=IYu5C{FSn4ixl_Q&WVy8V#D6~g(Xw3plB#5r2dlL{?;ym~`Pk5*b$BOSS)#qY1k zm#FU>HLYumY7@C{FI)GlD#n<2`G$68y|&9o-E7cyNOigz3FC7{pKG(>{!zZJ^C(-_ z5ThLqc*>^g+P4`j`6rFBZPNJ93@_dI#-o}Jv4(G~7#M$1tBsN9yZ2J1+U^VaMxBqQ z->4ddn2}A*MrYm+=2CrUp>)I`VjFcvwX(2THgQhB{?5XJEmt7p@7FY+aZc!s9n{$+7rzp+cA%|E)EJX|0hc^GVX)NFhn&F$>r z>&|nZOD6<~cb)JxUVNbDYK*1z+S?1|5PPrYLrY#`8RP38N3xFU8q*Z%`NIeI)*CxN z7q2ch6z%2zxMpPM9T{PrwRXNdR8bKzGFMWGBCAWu>PTnUf_}2II*_bxovcnKCyY$w zr|i`ZQ~hf?%_$$ftN# zfM6BE1@mAMQDjEdF*56nrLzr;%swM3#*CghQBnE{lAig70<{HASBEM}KgQ5gb|{o4 zCf^pA$>-?|Gx@qOFu|_`X7UwnB1X_hRX~8c$UmzXA}>ug8(xH&zIEOK;hxwP8(FSv z#yRZA?pWmKHpk93YqvVx>d<+srxjL*+O5TP6$Qp#NVG zlQ|a|Xj`;fMnn6$(9l907SAy2N316tOAr*Ipbrk9LU?@?Y)AQ<%AiVk{YKc1!oheT z=4g+_p)_VU_RtYQkL97SnB7)~z87@kQp%Yx@hR{Jhp*%*zhxz_Ta0?BHq80@jljrH zcgNW5E6cgzc{)8HHn&<4QFJB&n_?-P?@HW+{`CxCVj%kW zE#ps0l1a7BJSDs_hxP_WT_S@&(?CY9HKXFx>x9?b6+Y&BpmIi7+;tO^&lhI$ov(rk zewGC=lkd2RvIoFmG!16*oxct9{y-Rr{(VVq2V)Q0=J>gVq^0N%u^h%%ZVQNFEe*Cx zeLH8q$Aq+0ATjv7SG^1{H}J^cw-YE^fpY{17RTAyZ~nMePQDL0WcW-s*{xETOC>_ zyg7^Yy}9l>^RfLc94ac3nU7C7CA{uU{2ucsEd$&l@JQf5GPAImmcj)T+B1|oL3Yfh zs1yJF6ZXFkRQ`@PgvoTDo+GPeY&*ACbcc+`oal6rUR=37yYXDk`T@D z#%y78u9_#T;dq;G3Fcfy$1aiipf@pOg{WkxoDtSrWMcC9!c4xbDwyDhDu9`M$4rzx z00yIJFq1FqHq842VIcbTOycLv8W2=lOkl`R$i>84DNHZ4w<~diuo8=2B6t<2k>JmH zfZFmVCj^z?#d(0-GLsX6jCGOKBQ`n6)*$vm<&3c2>n4(#FU;i2uYw7FxB{5TcicqT z17I+k1~d8cZ^OJl5C)?Eu1TsqTvKgtf3{S(=;oAnW*qo<&xE%YedH<}|MlsmW{ZLS zf!hX7Bf*B`rT9uAD3F5H(^ZfZvJf&0<`Hw|t^+dLQUOOu(G{6)PsTs`&?n?vQRJa)}Y%biF2DYIA-DwsU^44U<7ydHF62UIlWIz|BU(GNoGZ+s%<)0ek>qeU(+>6Br3dSP?iA4gL3PmF;1_gU8 z5*Em$7b13sU(pZuaXu-Bz$mf7_5c`+J_a-SN;ks1KM)3@!6ZCd z_FI^z_XKmC6zQnS?M!FFjDAESlkyh!fvu3{_8T$)iNWaKw~Rk0$@|<7I4|&xhlQ8D zWcHYHiEN!csR5Y~F6*(XUX75>TzV2_^4*VunS4K(n0)I@1ixc2lkfgUn6d}LKy-?U z$+yzP`vYJw`u8Q_1>|O8+Z=M{OnG`frScf-^mzB=fW&uZWLiB8&!No%69dnC<;wta zgOB`Odnu1jEzhkzHNi1CfV*mQre)L29Iq)U(J8v!wI_*vsFD&BlP|)=E%MB4*$w2GnX$W&XQm1k zASEs|EkZJxbP{IrLBG^MOg^Xs9%AxEmV%_dgLJ{-&-c$9{|7h zNm_aZUs`OWXR8s;#2F&yG<1ebwyiK58xlnEg&HpY^V2BDAau#+&85Bew`V)B)ln0)I@1ixc2lMnhNzN%mp3`D1zn0zZuygvZ` z2P7%?%GqGKZu#b=0h7i#epVZmlIEIX>k@-q8L%4{kQhw)*z&<-;U$;sVN}TuW|KAc z1Y!qO5M%zsDKm*J*3mkRg9)6P?6O;)E$3j@(?UE+qKWv4ILsFFz_TLZ3VC)82fK zO<@dZ6?`#m+} zCwtKg+{7XC?va3vNx}<~?4Ox?0T6Q!ruB{dOzV5_Gp&yUeN%lw9>D#X)^{9c@&!(X znS9wMCSQn&;CBRO@&!&Q79L!m_=s9P^6+gbt>e6fzT_l(azE}}-i}*YHR>{32E8Tt zILzdGcPh-}J8NR{tu_(-j=)U5cc;LVJrD+>H=3Ayi%q;g00yIfUlML*bf%{9O|jSW z(&BW8dgjQ3 zWn$X%$ab#!Sn$cdmO#muCMI8qiOF{yX7a6=3KRU^H!=AZnJ9Yz3`QS;nS3jzz`Q>Y z2BKei67CN%)#@e}saIP?g|7T$^W*`igS>rM98lfHXwpP2Oh!-@q-m+w_*bzhm7e5*|~x8pFAZ|zi=;J3-d;W(seFSFm zt(^k%{y-Rr{tJ?n5AP}$)ikLbJm}nt=vkV^+2YI3l83pvFKZeHdN0s{)q#tC#d}$Q zz0*E}_uVu5<-F@kdn?{|PwXp`zW=FWZ_?!-D^}=pOiaEI6O->a%;Z}#73Q;TFfsWS znJ9Yz3`QS;nS4v8z`Q>Y2BKfjB;0E;9}hPTN-j57*z3yEn%DI@E#N`IS23D9^BB_h z^RM7Q){m4^8;2h@NMh^H$BdLLWEZ#L;<-a&#anK)&`T2ZK4mw<{Nli+7QTl;+gEpcu`Z@jBaId5W%{!*j8hZNr z^9=*7&AQ$TRKoz@pneGH&7Fvm7?k!>O<#j_{flm1<)C52+E8fZjZhiJf}X6WSvPML z8VB0Q(vdcM<4}X4n5e1888s*QCFFKSgr?`z8NT*bvoVS{5^ZQUdg)WT42BaiQmqnF z6V!HgF^nq>ItJIIRyZt8C&$TSq$4l$B~nE_fBsyluG47f%;IZms0Ph&Gq=Biq1W(- zwG7_i%R)9+qs=!A{lu8Fl>f9$YP@TF*2%vu@D?Y-g%!NedKIks*Y=2i z^J)7Kp1~~^M>HGO9M$C^hbz|U3M6$kUG}`^A3mzL+<{{}s}uLGc2>*C1{t~EWI=yG zMjnxodt_vt?0%uzt+6#L@gSpQouRf+{V2a3s08D=VRZp=&ZEGxU|c zM{A_x%(%!Iwv=eWIOKLKtXK9k0TG?hV%%`ZK;F^|y%@;}DAk6mX5=m7p%(->)=ecZ zM3OhnRk>V6*%L$`q2-?@D;)h}dnrXh#0VekMO(DJji?IJ|4UZ%`4ba}DYobqj6Z3V z8%so!(Hk)SD*ycn#Iv^O4vfF!C^wckg^d0Hr1?M^<;D{6WHcoCpz^=eGolF1aZk4t%{kVbKqakZixz6h>&a-x?VsEZZ-f!gw#+Z!MD%LGfFc zOgMDkuyz6ZzqgLT z9jmASRU_c^#wseWzB=@)Wtp*xipQ$oRSa1*R#Ex?)po`zDzCZEI&dj1r*i1Qe2h`cnBe z*=~f7c_vXdSGL>C2a#zesA2GRckWXs@lM(?lI<8rO&0Wz*)dMrF%Ft6h=cxE%+Kco zF(*#5FsT_co-WII-zq93H#gw8oz+EFYHHD}vvxubHpHKzM!ZMB?qD89iJUkO&9-0` zc?ot2SFBKaFqIePD?Mn+6&^|tit>V^(u1h1vQq9vDLqD7v~n`fFyK+arW===sVuV> z2~HxGV{2zpeh?f-1rR6sVL!7)S6GW8(QWo0lSZk=Obvw~m}@D_K$X zu|cO9 zH&w9e zuC-TQ7JV~M7hts{YAAAjNg2;Nb>GxLvvFZo4C9e%f6lDhx)GyOSBck>>TlkA9JHyc znFR6LU%QUDfm)LD^@702?Yo0*_Q`V=_;{D6sUGYnJ071K$lLg*b;LoR7tN)tK`r-@ z9Yrl)K&eM9XE5#20fE_}=m5BB;g?>#*S*_q60kt9!47EuQ`gFlITX2}BHtGf@R>V> z3p{CGnX4`d#CFHdAbIcX2YSB1RFG|a#)<}rudK-7j0p*PDf-F*K`1^@=1 z_p0}ju$beG-r~<0KtTgbuEGsN+flLeWIm`(Ub5kTTrXS={%}r?-@K6WG+iUhVN&J6 zz^Ls@Yiv?ma^^3bQ2uB2!l`e|d`U&^$ks4{*eWwnxN92iGqh@tfI&RgjxhMrfaW#m zL-t=9Ixkp8y0nao(|ce-Et0srmO$c({e z>IF9@({`ZLn`Do0MQag})6HVg>LkH6YeoE6JKX>61qAxpe5ZNO%2{D)y=-@2K&kEZ z%vx3o&Go5u&N6JA|I!-El$pZ$ZqRU>0WR~HO>O~V!yMUNbJf=Z-|@2PNdGf(Vl^js zBMt3E@+RTsxU@?1z*>>~eYxyg&h8Le>`apP1iaibtr8u`68XEyiJ6?-MKrVvX%-RR zXpwdcg{l|&x5|hgb8@|DXa|xv9xpdft3(E}gz_(CvXz|OQ)scaB(;CDZ^w-#Nv(NV z6H(KPw?C$C*MIns>aE{?k2;|H@DBB=?&I58@w)fRvV3&Q zAa%c`G2fi^F|Rw}-LXc5 z?Ie2&xj}c@hqyua(tVPG9=7jD3VLi&n8a<09!lc2EO~9#`RkiqF`N|8bpHFKfX7SX zlL8)|e-U@1=aLS)Cb)stkHApV~*&# z4Z*l`@1_U*l3!E3o0Odd|XFG+l4e|O}?+hj(^UW+8%$!714PlN;CFHc-c zDP9Q)0WOWXaJJCAu3=WkCnpEKUz+GNeH08t-}(GW`LdY3XA9%%8UXk&PNL@4kW|;k z(<6UKuTf{Zh!?cYE$6CEW|0TJI7sI`fZi;>Y!O)WtKQ3so{`j*((CQ^-|GfGV?R<| z&?arTQPVxtdq+}KUr{&`2+98~?a4hGkFHzb-7#4>@Y(XjC6waQ&!5Pb_3S77a zAIQKyc2UoT!lagq!Oa7yyCjKK_TB0i)?{h(-ojnFEmH4}b;5y*%M;g7idR57e@lD5 zIU8?YcVkw^x|0JJmnPaz9|Z%^Ri8g8Th^0tHa@QI1_1vPk`%r0+MM;#E_`Ehmvi!P zk)*$h4Vw@KLIQ>4obm&@9O>Gjbr!`4z zsiUG!=u=T3m7^C{(QHHJt;s)O6F}5_sNOVmM-OvrO=MHF+Zz4mFF%B6|847Fo zy{(_?QtdU}J$!9NKF`?3*Inh$US!ah@YlDCRTY}O9>(uJ9U!wc71HW_i0Ic(zMo&z zE#+#PhS`Q8ZI5*1Sq=Xn!$@xCYAg7LE@?VDRt`-oKA+FOsqm3%PEd^<4Gi6YMYry3 zu(3OuS9F(3*PPH9v{H4&_w7x=*GmXS%@;lWrIQu+rM2~JLu1x(tB*D{{;=3{*kY(b zFVycas!P8#J_~LiHnweL=B_i8A(I&dz*FoX1(S;wxJsuggnT8m&Rx*VDtN%>6)^RN5L8$|4F=2Kd!fmt{tqG zj(lWP->Xso$l&SfxgGolY*P0pu7-xt*8VnQ4ZCNJ*wT={g7@PA#nE-TgG2A}ciyg9 zNYv>Dd4taejn|(>zg*9M42}AQWC7EK;v0@bUFa;N^eIFDyx^x2F?1rvGm{1V=OV^{ zi19?kP>asT^e!o?AP-LCr`*ux#PlyIsv!>s@lzJ)H^dBDNbAYHoA~M*dJE~u;Wldi z|8q)`YM(Lv@}M!+$4D+8GNKsJM5HAG!Z1;_VG@e2NS0afC#Y!z6uBgBOc?OzDM_zH z81NIp6~cZWWaTPBv(=1Uh@Eq@UtpfU6kH;+e9$GIjmawfYf*}SC94n{7s|B`<*qYX z(BBZs-5JVVX|f;=`e2NvU$B+3KNwwkK)@J3BOR#;$)1FJzD*cQV+CVM?)_J51CCWf zevgJBSla+f9*Ey1gg90S`SV2%$0{K}EoH0{659Bx<@oN5RYG3Na>`gG0pSEfMq0Q%J0BOzdwF26Pe- z&Y;w>YoQklwd@4+0?qktfL0 ztO&T>Bn}9;zOW*oH!Kb~!w*&j+-edB1YBQO5pWJH4!GG5Rs`H(5(flaUsw_F4OkqI zs z%O1AvpwK+%gQD~5u^a3t(Zb&;X85meu&+)Suw*CmP2#QT$hM_2vE`vaA=jPtnfbsq zrY|XXI;v$eM1O%ad!T9ZkcHG#fAJyZSV*)78Yd4~5C?tS-#g}gZsKMV^@<8VNbE@%;2ZUZ;G7C)kP1$67(tTt?n5@RC`fUW`brJ@YxP-U|TyqOq z9w3<-@?pTW=^+^b*Jn=nCg3}d36%k_qPy&sw5a`hlJ{?Vcz+X6IMX{L;O2Di4+9FP zdM^(sa(n6#P~h?u_cv1o49Tvxdb&Y4XSSlkyz(Z~nKWZODsdyF8vh_iNPCS`0dKm} z9Fc_-+6dyNGwmlt;UwC1M3JN71fqbX*d+v;#eg2N*RwwjSRlA=Kzs06bXsLC>$IbL zWFT)#F5af?cFvsnXQR_gRIF3ZR;_`*o&ih@ME|Y>d0)4YGYQ|gO&B*z7HnRgBU?D0 zv;?VMBqL+1`_1^1##8}cy(*v&x|}C#vLW3-aOZMfm^XeaxIA9qgVfPwjTqGp0n{-7 zP4)d*3-JHll(im!!RUXia^U-u#ycLVJ(rR;wtZ=}Q&Ng3XWP*w zhnS~qvEz}6bjn*64;~9>4xDzRI)r)BQZZhbeUF3PL<^Zi5|QwZ_%wxiV4KKgvz&98 zqX?xfawio{z&l!{DbRsYB9|HR>;oKZ5G}-&BqHJ`VlmyQz^6iCmn=JugMEV*;z$zV z@s5}@1u`&7=#neryv$E@4yZ+oh5!Uc}kN5r4iUqvO$_gg~D18Sg;{63`dq1qDpLw=z z7yZoh(i8MEo!f5E&%8)(pi8>vYUz@{mnxu^p40tq13l-(nH%(+&Ua7HbDl@;qThTr zel7jxQx8DH50v)%3l#r;eI+FxyfdxwD|lzxDn84hbU?6lk-{Hk`LGmU*MvXhH|sS17*o{-rbl#M~c&&O9 zyRoZv?SG>xU{TCl!%e*sujUF1U3p6L491u$;9sRr_#dhQ)=Mr8 zJUmRD#Z7F&Rjpq)Sn_3by`SW1nOzNaOcn4IC?ft#Rlwe!z7zR8wYm_R*jq(5v}dvP zO~XdjOCN~HM&@aog6(hI5N|CO^OE^m<0@lAyHp*w%uuqIJ3K7ax2};I`xv~Rq!V3@ zycR}&f3Uh$S9aL=)BKUFd|fO1Sv|kqsLASB6EVz&X8dZYe64<%oz&khy)8CoL1TYg z8OFipQfiSRo~^%M&seWI(X6|3)Hv{x%o{<)^Xks?hvzaP+U>=#q3u1PFS$_}d{m3lzBI^qH~6{@_xSzi`?bcc*Te=Z_DO2r`1~O$ z_q$-j+fsdBzA@7HK-}A4ob5-~up6XbN-x8NK7N3-K}^@)Xp1ozdJNEvU;c0dL^I5v zc}%Q&oG;bHh;_>CFV1rfjGh}t-X3+7SU=Lm)_=#>K_h;<$PB|rk^$)_DfHnt4j4bV z(5vY)w&xFDs6iX+r!u;hs;VF{zg^!-Au&H^EYFZAP3#}xuj?z=kuQCr6c2jnZQ zw1AcTw!or(fQztzl>CUBU{Sw47O;}v7Fg5|a4jt$CBOZfU{SwR3s}i-3oPmfxaJm+ zk{@dmEb6z<0#@?d0*m?quC)cEK0!);TVPQ?z=cLWsr+f1U{Sxl7O;}v z7Fg5|aB&uplHa9Gu&5uX;Rj0l*@!AEP(f4rLCSS1AFjzFqV zeZO5cu+8p3$?lCy&N5xskDGaD>GzAQrhPf@fNW1l@DzdTuhxa$_p`_(#l9=UBK%O9 z*x2zhtXYD-zp6PxU!trC@0}agyUk=v0Y5psw-x%*WI-JC$6|gypVl{Xl$KGf+}#1E zC%Avb3QEmzKO3H-uv#6qyDH#4TWlKB4lN>Q4+=OIA#YNO5TVwJP0R_XBKxANQzXbK zl5gB3NH@t!0GJB19u z=e|i9t>Fhq`w5d3D*xDOiZ7R_u$HCYLvV{`Lobk=-%9AkoLlq>^kTu4?UPwFTXK(m zD@(5(Q}YX1B5Ik(q8SUT%YLy0!SbwfJRE z!af}@+&V>3gRabB((PtcAV(r7lkg7?3S+(hG(o?%@KJ(+P~rZM6A%1t&+?)#^iE6f zeo;%BtUMSP8TTGTe3hD^8=E7QH zuAQ*f7W_j%uHW~F*CvS<-J7-v(|t>JjHE!JBG)TE!)1Pe@SCNGHh-2Y5KowLH{(91 ztmBH(J1kC_=W88$($iyAvfcFRjDyi-A6!vB4Ic#q(It7lJ3~)S^H_D-ZaM(}`$?LK z?shqiw(e71n#Z}%%V`|%zSyN{g4ODrpX{unTwci|{+uLUPwX65=QK=kTyyLu_ zhH;KlT^eos>~bCx`^-1Aj&?+OELzcJjQUgwTj-9-gO1!C-WC z#@^_%s4GfkcpM1-kx3ZttWD;9mzZBnRJ9_b)>8=hLwgVl=MXnxI%}E792Kw2UcA=Z zlB|Vmi6p(*dIIUCV66?Q$t=QJbw?16R~0i`^NOrP8$CU?B-{Di&1m41y}Y7)BP^~o zuV`mz<1~*gr$@nHbalpw=(10)D4&JLf$(3Oq3vUG4wLH& z0yiY&t|FxNLT{E|wg@iz74OB1H-~0}>AoiOov1jEY-K4ZgjmUvC`4o8x_)tGjSe%Yfosk}022IFm2_FRm(XzaRgiya}9?qxj2mt)rCs8FZ z@1f?9SvoTUc7i`;JhF^fyPL^EySou5AzsBsTAh%Q48<0U%qBC518oS|9V&ZdDRV;b z$eb-G^*Wg{FTpzWrl-gHWIO8JjCxKP)HMM8vtN>zuru`LG!Lj-41~ex>I_KA7wRQ| z{wW9Hzkd?1&n&H1Jbx!3aGY{~ZET9CoA(5lOzh6U&EoZVShTytW?m`-q%-YQ~^jwm=1KuBxvBm_dL^)gK#H*Z3VYFBelB*p)!U|(; z9vCainIk72;?ypsq1{OCws?Om28Tus3tjqTvQIg+zBIHG$=w?7kHJ`@0%L_aS7gKx z&ZQZ&*zqJ+OMCjdKHPZcdvu{^xdQ$7v9_QhG{dBIPJy`ebfSA#LA4@}{NbU%{ zzZu5bEKtSI>5~!%*q7AQSgpOQ9v5LSw>AdE@^h|8i4WMd&#CBsdv^`aUuTXpAcuJ_ zTcxt^*tL(S=q`KrL7e}HxwSqZmX|X}N<75Abc-6>YVX>Mi%^?e>jG3sIeil1c6MzM z72RO({u1XuWRBA!-_c2)+LEq!(V?QW)0n2CuR+Py;95? zxV(_OzJJeb^7=lzm9Qq?xzAxuzI!H3K6hBEZ{iAAlW(d?d=G?y=*uQeK37=5Zz2GL z(XTp5i=Fgd?~SR_dws3*r1$!7ESBCIC|%8H9y}4nXdb!`EFW0VpTACO(o_3Yx=tb6 zx95IEQp?qR=b_XB`mP5R!6Qn4GNavRC9KK!?dPyy9kggjlg}NN>U(1atjRaqB)$j2 zKy-mglg|}a@OuM*!RWsr3B7qPhTEM|L+WjrOOL$OzDkqnMfPiNwMcq)Hj6w6{Y3n- zMPSjddM|Z1cpL8LcWY8-Gm2a7-_UjP*pIpkZX|IFsNI9T-x_L7ntblCCZEkpSd;I} z=dd!}C6gwfD=fYT!eI0YSd-7qq~Hg_uYD4Fqt92`?u6!JJ*_BSWJ%9w>U|Cjzo+fz zb(wDk`zN@sNXIS~tS=vBd32X9Q_#Bp- zt21fxxx-R@o-1HYJ`a=l9tZ=`S4^6GuCRiiCjf)df7c|QyYN4?{^X1rR$Q-8EWFs` zB_Pb=CMY8v$<7++2MiJ}{P~c8Z3CxK_Ypg6q_fg~BE7gfH7fsN;)Z(#A33x8lo{>z zek&8j%tY~W*pfkiMIw1sqIfB6N8w{JKc5wNK6gj(_3b6}z@e6_HJ6e+yXS5vS8EfR zqmhia)OvNi=FS=h7aEMq8*C#O`z!Jd-CDkS=mbP$BS`h#=l3=>1z&p-ER{?3{PV__ z{cK&)6n@KMiB{We=)KU)*N_=6*|-CIsY+b0nXZd7C@cA&Ox9PdG4xXVyQKX_UB!UM zMJ=_NH>^J@)f=?=5>0>e(;8!i50pS`9O-GXk5JvnH*|E7sp>m1hG$j$ks$&@FI>Rr zxG&XyLB80(H`p*uYBW#{mtXE2K6p`8%I7`P>S8U_1Jr@Rqm1EJz4{KXWnJ6(Dp-|| z&yhAU440luHTN}o{;(Rzujec^09mDd+P1AAHiwm9TB-!}C`uZ%U8}2Skg)<6do&zs*G7|1J(-d|G84 z>xrZLM}fSU+*vklcXF=GKO3K3(#L9Yw(9z`)DXCHXc$zyC#@XGb*H5v0$hzRxj%c*&P)=HcSGZLB4P#I(eLg0Um(=EecckwVNTz!7bNTvC$E zq7lngosy+Lw&G%onF@H8Z;BH<))#PP@GRdH2Y9Tn5}xII)(IZ#3%DilEZ?&Z@L1m# zc$V)WCwQzc;L71yzK0y(vA(VFEZ?h6@K|5KErnkGIg@GRdv2Y9S6IK>wj?wg#f==E1sQxsx1*}?8i4V$by5f~ZwyxV%;cP=;Q zJw2NC?CyT{7z6MRpgVd=ylFcy-IcOOBn1u?$yQv%WqyLVIZF|5{_Gn;^B4p051>2# zlNf+fch(N`zAMbf6IETvC^m(PKeP`~Jcqag(^)RVd+`1XZgH zB5G%VNSpNm|?U1O(ca?B4H_GR@@&bC=@j#P89apD`a1Yqf|C3$nX0DR<0n*vkr$%2j5~SoA(j zAT2rPb2;$`PA!Lqrjp$4@%}i>UUO8T$mJcm>>Q_-MMF;~xsS*DTVnR118Jh1adKh` z=Ms|^OD4Gz@DUc6eW<`o!kju8aU-XeMnl_^+;Mn+bIe{O%17w(xlFd0qezx`b=){3 z@#<&|F?)LXiO$E%9DsL?EB{CWoPoPg*{ z|3jFpBkM}qB_o%q+|K*tk&SwL2++l$q4S^S^I=zsr*`v~y&JpV)b->@XM z>q5LFw)OsSacskd_~O|61;c*4s=e`k{L=ephb_+ze{Hh&0RQsY;g++*drTI@L4Pdf z=ksaFtBFnW{QuZH^RT9_?eCKT5($tXVuB)JP@$qEh%Hz#VX%r?g|-%?4uPU#2L;p? z#E^ke5e4f2twW$DCAcLR=aX<@_03sqn2Br)rIqx~?`@20n&%JGJ(P}S$ zxU;gdSZD83(-8M}f4N~8UDlNS2BkarpqgnCBwpO0eA3WHn%^muuRp* z)ywY+0fnc8t*fGO`a6rbrS~7|4)o{-omE2pa5s~y6PF36aUZ=?rW?%Ws-9-EfKHY+ z)9(?J2AtFaR-C!q*iP;?s;+aDUhmm|Ti6BaYX+>&XiH#rADN|BWU>-;CneVe`b-Fq zS1ANL1eKY^(ju0wxv)$qZIV3n5CTs@&R`l79NsO__Z^PX_m;3gA=pP`C1`u07Ja9M zr9+u!jnwT58TT=F>!9q6?t%b#2oCLL>b&&I8cowu&ds~}X5#S4wK6eTGlbKV+Vun6 zVNl)61l*V%i9)cQ%#4mbnjTp4I1238O#>fDK0T`KZqn)_RaKz+29>L;f+UCa+_)Qo zU=266jRi<)k@4K2?2 zC1HAW_XOOdRFx7ruB?s7t)%v+ zG68dL!(@GfC!pA@N|MM276F!wj!caMH#t!4&%I7<+p0fB9ayBwXFgk`D`)m(s=~Rd zKrK<<7qXeF)`y^ZfkxBNkIaJnc*{`^CUX4$HxCTaYP7Nh>KRMvDXV z0NMOPWI$L<7%ugv;N0wB8DWnJEF)|;$Ot#X>OwbLY&OIRG%4>)vX|TAo$wXriS$Hd zD7D8k@D*l>(jg|7@Q4Rw6QNv){gE!{eUy*?!)K^OMi^ zi#$%TPuyjngeyLAWjzf zWF8s+57cCMp#3e2cZ%p(F9$l7Pr#cojrvh*i8*n+A`CgQ3Oa?@v>qNu#2c!h^M+09 z;BiDAS&yj67T^x>eU9M|?BhWeNqIiUOb?vn9U~-{`y4e({@*tF9coxOO@ZtVomyHot0iG`OY1u!OX0Fj@hbMpUDo%lEdNj9KDgAA(up zDaytynP=aGS;F)oWBvSdgRy=vToZyi6dRxpfBF)mGS5+IqG}Y~w<&hTL`3jtW*G_m z5;yuqWBGZW#5(6Rn)5-H7tByNSr4NXc3ptc3YQIN1v(5%IC%|5DT#oOp3C+> zwIX@TqEiQw*-KB|N!}88YC_7k4_sEHd>Z0%AO&${%5g`w@_r-g_PQO9^Wo@9POW{UX9t*xq-e2ki7)Zm{?Dgq?m7As@E) zrFg(jzf}f%Ur*TS7ZIkw_P#4UV5eWCu`j~!7l$kFNRoc-D`GDCO|bn4^V-NN=zsM! z>){cHFtSZ!j(vaq3*Wha^ZP4MQg<->V11Np$l8#b!otq`Q9H-Wk%$CuUo)eKh=9fJ|l9f_r|H$#}2jqN9*E` z*<**=V&ny0?*`j?yILSn^+9rhdqBy-~T~+0Uc4YcRlJCb3wLJ%K{v{o+t9H)|AGc=7#sELt@ssv8tWQpL zOty50!Y}jNN$`sgwEukEohkgY4!OgaiZ6X0n&S_nzNGP@jISOwDzn5-GEWUN@t>CD zu(Le3JPf~NRhKJ$Z=qn~t5>3q}8k=lJ`%mTbcZ)>0gowu2MF|o2 zYSI!SevYk4h-gj2aX9Usi#VJyf8S+e{=VJFdw{e)J~S|WdHe@cOLy&WMW+Ar9Xm1R z@7tDNw4ia7SNE*nw-82p93!efUL5tuo2-q0M4rz1LP~y|5Z^-nF5!#&}4V#q#(?Q%jTfk4Ph`!xzh>`LPQc=XrHcL8M>#Cgs`M>Mtd0WRtZC?DDA3RS&I& z=Q`4|0pO5ib$|a3&SN_rumgW$eKT;z zCq86ascTm2iN!M0{MZlwE?3=GXC{Bq-&dBAxnU8HO~aEL66 z)2E(XEVBxbeP9@mBjSj~vP$TwVMOMU@mS3-_op&c(56{^KeCo_WXR^CxGEvFFKIhf zrhQkMHclwx(wN}W5*qjzQrsyXfOnH=0Q^ZQ(mKSpMV1gn_0`#!BRA`wE3Af}?^?=e4~8_GbbhQ1^H@7pT#H4^~I%L*v*=Rg1m=?K=(&1puHZNTr-9^S9VL25Mp);m>q=_q;=FsEJNHv|834BR=pKD7Uqz0>vU-7PI0ckjH$SDWA)|L)1928>vIb^-{s$BH%0& zIH|*Nhh2+=Y|s@Ya^{UkZ833i@HxpNVLr9M`|_v&GvhdzwnGKd_kvWFkvu%Zh}uVj z9w4J%hz*#Draq?29WpRvgt@Q6zIY*NwI}|JC6;WF2yq48Op=Zv{>Bi0zm4W&Adv^i z7=b_dBQ?K;NEk!>wL`o2kcl2-s{oDAmUcb#^gKKw^T?R+d^KacAnUdh`0KO=sIvLu2|2ikAhqc#lXUt8+%V$>Xuzkg?s6 zF~r}6KSC@VL;StUj>2PzzcIw$7^2>l;qCfYzrT(l{#K4K1x^W7KnnFY4I{Mk{%?a| zWc-{p?qi6*=Q1N64X5BS#NVGVBL=|RTcYYOuDSd5kF1^c8$KJn@8&^MzBwTb`{n&Z zW{$Nw)hcvE`1_FE%88a-!pB3Dt=I4B*AiK>zi5`+altQQ#gpZxdwd-j&U31h`wlcv z##eawaTaV`j#r~zFMHtE1`fz@q)*58@kBX9X$LCeZK(w=AI=O z2&a1YcGM(1zLH6Oz3hSe8R0D_mur*5sq=fR7A8E+52Os7+sa^HcV4cRJFt*`zJIH$ z?g)G0|1uLsyXOX=FH4f6?1}{p_AF+4KkbUtJw7WCyh0Aa50!)|!k^^?9phmi z;bGa&&j_WA=BFJ;%kDqvNEkx|z6d*H+L}SaJFwySH0Bxw!C3f_jL^wg)!R{3wz8+= z#Q&f@a9NAzMQvZa-qULhL$K6IZ3&9R+p~S^C>l zu(JnZ1wu-=r-x;&0MuiIdujrn~_u2%+M~yM_9T%n6R*_BN)V zKTy>HJUqyKVyhc;;tpr4)dDrBs~PY(qlI;UH8jp(4U?n3D*&f+^~^Gjh6w5_C6$52 zQcsq?IZ9X}4G{t?7SP-z)ar>X>cvc54^`dyHdk>?@W_)3lyE=a(W=XLWordgwYQ`d z#n54NL%XTk!7}YSrs9lJ*U72jT$Y>@^s$8V2Uxmp;Bl1Ou#VM6?lSlT8LWmAOn0L;S5SR zN-``1>;ikrw1VuNf}smSU9VVwU!r>w$$l#10A1`S!U?cH@O?1YpaN9%=$;)jr~v(G~m1WD%TC3j3)@KwdML4}2r_V8r z$T%0LZ@}r}aC#~3SXmXd&5nELFysrYOe?!VZDVlnOx9WU4WyOcr2=cYR}bsjWJ6S- zJ7lkrt^ZFz1Q7uw9v=XscS8>VD4NV89yeA+ABG;xRlLc3Vwtsyewa_poAP_+z!23B zK|w%-|1cneTyFncCP|L^Jap&xX!SXC=Owhh9Nl>ht-gise1O)pq1(FAg9>z;8l7hr z^wj7QDQE_&Y(mhds7sWfA5fQVNl#J#vM0?zAF-w;SyQtOBQnmhrXI7V?l+9cJTm?t zsOiWM1!3JB>1L?qcr#xq8eK^gv3X6H>2}gy5prVy89(PHfud83I^mEt4}D^qbO-vx z8(Agj6SH!)=sfqJd6>L6A*;X9*oqjj)$vzmi~Lz+{Jg71{wy*^^m%2rPm)mUXWM6* z9C%>Fv<>Wpw(Gao=Up{fVM_G$$nNjHzmwZ@U5C=LS>#oNvNZzhCwM z{=}}sk-ykPY2~+mfr}V_(7}N@ZAEpm@<@=)gvR}T;VVwvG))PId|&=OJLIIV!_2`? z&bL1bjX|an)h~Z7?hMyC3sd^?8z>XnxB5lA`(*jJJq`}uPLCEPKTgh~k|y3vuMey0 zJP*V0+6*u}H&?eLYx4E<-H}!6&$qXSjz}Y_%L|I*f>%#>bvtHFMx z;lL)8&@E`lSXtrofYRWK36IQmGf(~zdcV9PYw2)8NBBuC@EK|&h8_JoOkcDp$6RzA z#hD#uVWPU^)A4e)20UNVhRZABlhkTM*~4!pQ#PSmttrd+%~q6Is1{4PIsXAs-aF=# zH2OYCYe|$r-YN^qU8A*Tlz;J7nNkiKt;Wm0;DwsVKNKOYkr0ka6$Vt_bQtVML!^<_ z|6C&_WqLsPg7Z%A+gs9hO;|E@m7}rMt#1OP3x04O?_>G9cf7`QN=nQ+W%`n%` z!Vc#8b(oa9CrU-W(JFhd4KUYFK2DV8j!{LQG@@49N-+Y+=Jx`R7&}QGejS3=?-Jpd zRx0{QG__*2{r?1#UM{wj9;=luI5T}V-E#9Tvn7v~IAX?^@7)=>;5)~5%Uck52BMb6 z7)E5w4^^~5NyFGyZ*;;Ui-0;YUb+-@Vq(^A)QL&C^{70%AObqi5&DAkyB(07kumZr z$UHKB&1>9m;?3TXI-)B{(qs_{>t;^}jIQJ2=7+VTS}o+V86j(f-Go1`-o9ooShl4z zrBgY%{$R+V+p?MDwI7`+``7BYO>@C*`;>9}0QNp5&k{`Dr`)*@$k?aMveX^eC(qfZ zIkHclyRU)hV$N@|beV;cSh+0Yw^+MuLOmqo_wawV!Nb4t6_%$%*jQ+UK1_a%ucy<* z=@2&-8b{>+K{blrymQ{$V$-Uhk|eg1cKgRhZ7eWv`_Vx-XHIM?P)Cmrwl#6J$hJ-^ zPsc!C7tgjwhsIC}yuu6wy#3ZV@rX3CdP7w?M3N7Md3nh6e{PaQHFMdI+CC@ux{0cd zH+xD}4W^FGvm@cH&|fUDzYh62yp(Ammx#RNbFJm~n^0JjDDyC}>%%@blc(^{(W0A- zt1lX@ohUCsZ~agdWjY@E0(iE4o8bUhjNt%Sw)M&KbgKS zr!gvPQH18X&)DXiQk_c@SWSxXi(RtOY+Rd-1JiMib27j?8y&pb#MLd^I=eib0DWsc z+rA|<#-^Y+EO@oItJ{8SvUo%qS-qjE3gWI?pP zS{B))u9~;j20st;#V7Wko4=%KAnz0MlE2;%3e}Qgg^u+JyN;79clQ+O8)ZQ{2MMI_$)!~K{YWv}%Xp<3XWcAB$ za#&(BW2yi4&rTJXt4s6FdmZpht9zQ&I>of!|BD?i1w`cR5R)G5wJ2_iHk!4Ri8d1C zXG9xKTaJk~j+5@?Md7b=cu_c$$f}65?a(hZl!CbEU~dywyKL*s@^pw(459X*-)SH) zA2N-o4m0rj_QPO4h<5z?G>P-heS-R$=)&RQu<7<$2Yg(wn5 zr*JJ>M5j#TAB#>Ix4bVpg_Zj83NhDb^9s=uBdeez=wE~?CKbfR1Ya?6U7KwU{gy%A zPz9Yvha~+FbPt(EREGy&@pfIi-x~T21hU@wCpC$AcjHRG?D3PnBBdp#I&HJb%fp8+ zsGaP$pVXe1)oO)bh5llZ{pa{6I^NkSL9v*M4ZMd0=_*u=GcO98xG>2fr#$Fd82)4W zMmi;tY<7^4nroa?mB?RJPWnDf`aXT53nkIY?BKZ6Tx`<9L2?sSP_PWGWj)Rvl=xY|f zj<)gpwk9f?iZXFHSFY$Z8E*9EH#K+}Jo7 z-EDqwAk`;6X<8|Ne=BKmlr%rvZisM!^Ti$7yP9raA-yZ42Br( zUiwB>>x#uurum+qG#*BGFFQE!olpG!X)wZW!s4j*k!eJAXyak8?q&O6;N7S+u@We4 zA{enZBh57O#RB)ITO$*y&k~OfDLNQiTWg5GM#Su1;50^ME#2D%^u7;HCRCiwJT^G) zIOB^(XEV+Hxm3nutCMV1R^i6RRp{|e#%dP59-%8bGd$u3h6h7$%a(pm<)50^!YL!#6bQv{`L9m6{mUuH^r zv!}yMzYh&ezfTNIzk@L9FMC?)?)|OSi)HQkvB)%{I<&FGt9#A90k`Gx|F96Oy5pk3L^WEsaK(kF4d_=$r7{e9kLIO*ppFHr$ zgp6FT_T!FCFAcOW4V(s#(7W!}cv>mZxioMpJdVhZ)%hj=LP}?BXgO?gMs_)sG>M;4~Vyzlz0erkd^pugC>ri0c;vas8Ov_UbkJgxY!l z5Cfn-LSxgUgPnal$U+Uo^#h)Yx%-(Q6{6%pT)#G1oAYY@OefA{AB&PfSU)~TTa~Gcmn57Q#D%E! zTZ9qJhqC*D!-Wz%O`(1pr{6PMt1p!A?%xeuJ*`X19J(fa+Qb6n%%MQt6rp}mo8lVy zGpytjE)t$&<_UTO1-b{Ve_hQN2mW|R#a4qIP1*{EAdd(rEP!&NcfxLF z#eDrKPqkA^Co$nLvmya{qH${JChj=RRF&#?dTJ{ry-d~ZLR+pH_>Vk4vYxuI8*G~h zQWvU0zNZm&Awk~u7K*xXvK-}!rm|h-;csE6e)FP{3_k?gug$)gOoY#oUe^<&}PyBOJrz_ z^Hul;Z6HmcV)x*DC;K*RCjDfI49#)A{k{zwNOP#zBwRVgm%f>lVTlYexN^BKeFMpi zirtMXpX7UXGpWE584_^iO5d{^NHeKeA~7)%TfQ|(YU2wj_u>|kygl*9EU`q3#2Q?= z2khSaoAyS4zj56M#>xoex~;}Ql5yRBW94zCqa=$e3gY7K;P0afXWW3QL?7hrBuMEu(hpAiVV`uJi0mbpkeHZsg z!JAFZ#bGOweWqCAJs?Rev=vji9^Q(1fv!U%dLu45+Qh%=Y?JDnqBS;@$ryHc*fT=Q z7yNUxc-h9l8PRHM`QgM@27=fA$T_)gXV5Xr8UB8pC8sjRbj>x?7Jg!`{WUwDH_0^A=%%bs}3c`Bpv>u=n@k_(o+WuL#rciU)lXrie!W8GXf- z{tddwnf?*F*hxAWeSs=%=0idXXb)tyVGra*`U)X>BQN<+bU3GW^k+$Y+#B5*3t5z!TSv0vx#u{SP* z68(V)kMhQ3;G?ZXcq0|iUnB9>RkYQFDm+#ywi?DtMEoif^pNvJ#(&qGUntS9tt(Kt z1?cQA#2N{VmSAgba+8gYqW#QH&9%ptmnTUnzR@lAUJFUK-uS(i*c$W12y8h&Nh64A|o1Z%o>-P@y1qs^q?l=#RNr)J8p&i;gk2nvH$UHK}{PSx*K4M8p7rlc` zbr6l?JK$IQgsD(DM0vdED(1=*spP{mR|-~hR$WqF-l9+FtRvKSdLD8uoOx0b>UQO) z)d`G4tR!nzwqZoZIo7OW)~x-85t&EEuX#;Jvwy_) zCw&%s?mCpsoij15ueD+2mi=rD>Id;!FGB6PV z{X5pHvP7H{;onK_! z*gRv5%{OmG_Z0qs#j^NC_NB+aYPAfIoz0JpZk*@XJ>|%VbOQpRw=^E-TJ))&a7{{8;10k{R7=#@KxE0-yN*&$C+^ zPhB1_2!KcE9s6q@8UHWTyjE;J$+52`kO(e8ohAf^cXMGCpuQ*yghYW38{RtZ4d&v3k!CaOh{D>VZJd5~})}6iwU$2}`;4%n-=F#2k+2UTW}UDv1t# zVz3mH3H8#GKn6jt-^~JRyWh4Je0hVbyA7!f7lFY&ErG(;SVKWIQw|HZU6OzuGR1z7E(FCgMH+3eTOyHh2YR&qGR=N{ zXA4W;&(x_ug>;c25@3@XAU@bO%+xd&>YIuBp9(cvDp;shvOpEHdl6ULEYN-})Ls*S zkPNb0px6bzKQ-a;tVeov@bK4bDkQ8&X>TVnzGlkaX8vr>-jU1fN*BDp>LEBIg`xdS zb889MB1`{4@V*9+r*NL^(1k@WWnSi2;Gk1`XU3eQXS zOYX_8MfC(ea@O9Co6I6vI_vxYrFMGH$?dpJ2FGFkCq%P|}_moymN8OVkgz*031h5+nJz+y%TNs3rJ@H|($SzQOpBq$yz9dm3(W74)Fz zO3g)Sx6th3Wzd5OTeL@nlbEt|ZivzvELMBNHpWfniJxK1tCFOXe4`&4+>O0p31pKQ zx3}G+(}ruZiQ!+58|A zmoIp$@h%6ASf(!j;;lAwxofn>0zb(6h=jL>cmvRJ%glOsg#Ibz*Z6umbg(k>8a$53 z|I=$g?nix_9(>jF!#z3S4cpH$IJVi*=UD#X5!XmO?4b&=aN1T9g|OlWPMzpA=*2eCz|UbpawVwl2s%(70};kJz{(EO)s{#l|uHgSrt$LvqSo|&~ADW^`PrNM?HdLzM810EBNaQc0Nu^C51pgJ}}TG@?3rv|WbF99GQmyX3TDiC@{Y z6`%TDa$Yg!hWxjyE63cBN7=~upQM~*&!(GTV%_M~X!ZTEabGc zk7+0Wx^DD^YA372EL~kEjrAYD^dbe5NFAjY-hzs)F-xq(p<0dDlqsULW;B~JUxaIg zNQi5oXC`b)f(X}wXUmU@(jJ(y@ug$Z$#3G3lfM?Eg0~@+G7h)_E&uf=M{I}Q&gn0}b|v$g5^*NWxeYCY-0 zyygj$-f?KPn)GH~^Y}^A99k@u)_D&|O5-t4;}>CE;J((9V$NShqUem)5-Ghf?fjw9 zYIFGwUZ|P8bPVM0x;e?(^LqMW2#I*U{c-51crUv%Y$!?md1OiiieXLh=g+gEIH9~P zto-WzCJCU|e~(i*{gL*<*Otl_T{sI1q^&^SR5H2nQ0 zLD0~xHG-hQ@0SUJ2F}f=uI=A73*wcuP?-O?XS$@uu)rf5|(- zTLYCJG9`nD*E1!rkel&*CBaJgww%zrycvtyqj-P9Y+V)T?^C=M_EJW%JpHGJPIyI5jCq$IU&J6F- zG0LQ#ZvlPjlKYCvr5sLu)ETvUK65nI(&)))D=TarN-1F^+$Q&_pUz?xcUpbYx$NtLi6@85+M} zLt~w4Sb4aRqf$hHs=L`tKui^CN@QSLGvwp@kO>Tx2pG@$ICWjRfyqEK7^{(R^t-9F zYb=$6M5lv)MkWL@nY}E+J}xL_vbqHF?0&KUXgYHoRMj!{Hvo+cV1n+tbg?cw3UVo~ zmvBHGBq3DG9*MykPl;Y305!5UGFuwU(%o)ia-^P+1`on8HZy?%vcb~0M4+zxf}^-5 zc{D5mtA2qXIsA)y(N&rj%Ifd3v^7*+MwG526a0pz zi^AoJOP3u1zcgSF`vt?NBNM?MC+25URqLtx#gRV>fm-ntJV)Qp8EhK_HMQcmOZDf8 zU;(SU%N)ZT{^Wk4m8~)7+9GZ>WNSJyodjkzjyS(NrC4 z0eZOi`urFdw-R#TsFAq$^)u?Gg$r`n;}vhNbGn5 zw%O^K866;w0d0IZiPGjlJ01}toHw-9aa@!(*Nlzt5FurM2sL0zl2qv%eFFv}V(nmr zL@`X0K!&z3LLvvINldka5fYbSnglYm#a_etego4aCfH%uV|_3A1}z*B+G4A4zF+$W zu}O36u$!>HCB6*{M})T68#v#?z71@WnH?75U=;h(7mf&Rv6Z;;9A7${G}8_XX&5j2 zo?SR1w8hrr%D?eF%O-hy!{kG;ISfRsGyoA>3_wILyd4!QHcL#yq55{2Er0Z>BgU%j zP;KslGgD{VTRz(vw{YDIM>|Ue(b^@@x;E3=MQkl6PIL*JSerS~MLe;ANX7?}Pi2zv zVluwQsMFS-iaoI3h&v(ZCR!QpGuOPK2G!=BYlly^8q)y&w*E0@VVDb=Y)MHqa>l0G zN2B;{xG;vO2W3T+e<^J*P)#aOT{eu!xS~L%EKq%87?F8ojI2@i^JbH3t@yLY%RTtB zC)9qxpFOc*E1ym|wVzL)Bps!zAg_*$kuyQ&k@4SDLy_1`s`rnL-B@6uxRH0xYuj*Y z_0Oc%Ny?@ETO(@=&3@nF=nr^4{edQ`k9|HddGx7Iw~1 zJ7OZQ{9G!J+b|K zRxmcvA;_37#ciB#e;@Us1%5N{Lrk#^smiF&*lpIXLZe-UPj?oIcG7;_S;*U2`036< zpPe+u3`8DF_Oj6>wXcz>L{?IQF*l=D3y%0C%;20c)Jf&Dwn!!M$#el$`X4DoL> zfqDLjFdByVhnqk=|ExF|;@@oo^ZXHEGz{_o#028`2gSh<|4tK_=Z^@ZVTk{F6Nu;E z5C=p2dre@TKO&5VA^w|8Af7)x4u<$c*31yk9}z~w5dR1hi06Mc4u<&mn7}-LL>LW2 z{68~+c>W{s{RpbR)IJXSZtHX7}*F`&qSlQ&@ix!wn+cIq*YFF5zs~(OYSSoft zfxg*L3RK005gCgLRPq8obRuqfkEWX6HCj3*~El4~;V zbY&#t_9SojvB5vu7(y`YzR*-`B2R6T=StoCZKHdVD=ir7=wtzgW21I z*~<(gGLMX3^O}q(;-1dDl6iB2(^QWeR@z?P;UBRHA7!pSk$1`CQAjEuvL;>`Q(WL> z7HsR~YLR4}R-axOR)xCIE)9*zDe#ILkw#Y2T`i7Tr(H`=imXDW|4S3Ld+r>}mN#>^ zVc2tXf5L2;pF4or=4(F(`{_IO+aN=Xe_t8d>-jG6$~4KKnP(j3m&~TWWumGO-6t3K zpd(fp(Fm$F{7W;Wp2oW%@8Xp)mkYeigO7N@GQN}b>9@nG#$9NC6dLnQf!FpCX=F9s zb@?&tlh@MsM^+)zf4a#v&*UZVxWpzeSzKF`ykwcnaPpG(0z6avK0NhhieKo?S7x5S zzu1h}l!YHYc`Ie%dpn1d7cP5pF?r!)d~7oNosDlLvwc^3!D6=T`tLkMJopOFqiKJco4A%`d>xx5EX++k;yW zX=L@^&;+;{6gsJHmH+xpwKDUARe4wDt_V&W=nGnD_d~OPcmnih`S%hLlOFB0WCWU2 zhe>Z5SFhv+8dnENZ(^$hi~_OHk*S*)z>Dq$a|WS*P==_4(BTA#N?2^b57J>OA#|P| z!Ve)*`%@q^$xL??W4eW+XrF!@8T!;`#7O}El*|mL_;4d zfBkM`Q1Kc4ydAdmBHh+$eR4kw1y&ixH6k)+vC7)sbiF2`&1R01%zc z^{o_=T3Dx^(XfN$IT}Jakg`PYN8*+z3Pq)?oh`KJUh?`P%VGRs)Xsd* z=*FelI=`mEV4dIXK(fwHTzC)gs}~;v;%aEnH~&h(obF}(f$x^Ye_&s_7_u-9i25Pl z)41{QOc)3FT>uOZM5Ynd*9=U-bOSsPkw#WOUz51{$UE%K;>ZaJaknE^By4VqJdhCG zT+@`WK~h8JjLIcmQj={G=Q|0XHaj0k*w*B{BH`29&Jz+o7pFC`Kdw*9X1@bD5(h+m zpI?M+Ts5bA7JuN-lK7?erBwzN;4T9>5Rpb!&xC=3hXP=7ATo`p9`z=1RZ-i3&DWxS z0^)8&4FH>OMa|Jg-xbH|Harj)=@3^&PUz*&b7>|}akbC_sJIfCsjc{_&_Y{LAl|D{ z<TFq5otvAjjvj(7f0>RFN$ew zMx_7dCMu42rTo%lMnqMbOqDTBP}1pfk1aoz$r||jFoSb<5n_pmNssngmI^c}{=)dj za=3GOto2N7vx z^-LH^xHkZ14Q}RgYK!E4UrA3wgfNq)+G8#(Y4`jRkbfnRC=Y+e;fMoA?n(`C-gKBx;80(=DsoWvmItFNj-K+!^LGaMvSxawq2Q@u;U&< zf5HBJlMu^B+>NCyZN-Vgwd#tLK-O@@`ogt?74L}G_N%7R)+zvF=s&7+R_n>H;1PP` z{X4Gtsx=LIY8a7uWIR^$%l%2az^EL&bXJ^|;MME#xnQ@}>>;k7(K`IKh-k%O>Pvyg zr>PvEnaLUKB!c}lh59}%S2c8)#vLXKb$us^ke{(~J5y%@4uMDYMFOzBW~q|Qsg31= z@Xu&zU`L>~$?uVbsS-CqhQ^&hl{5ISo+r#QWhU1{sGkRlx$g?ww}PF`vNBLvB;kU0 zJ-PaM>V8mhqlgQ13H3Wv##GhA5O82Hld8HArGH${)eI4tI^$SQV<6ZUsgFL&Vk zx?!#&bCbA4Fffh@bVZhliLtyu(2vDcZ{=R03Lns}oCXGBU?Cs%e$;R#bM=4(=mzvc zc1s|az&b)xS}}XM>Q4p0BHd8EP@k#$0L+l+qYReDdTZStovUWD)=KkGKkz`(tA8v5 z`^195P8mlF&FxL&ek;yUPn5nnkQ4Lv&<7knQ&s}%Efw`pwvap6q*p9rcB}$b9m2;; z!DinlTcfOog;U;Su=F6%1*x*7uw-x2IL?DKZewT*gI&aDDQ^iBr@iF{yPsQT#=h$Jh0l{%;f1h`wZw`4@c)JuVJ6+qO@vd!^0_)PW!XBZk%cHur%-7 zTggSKKx0FTjncJpG)Q!R1ZP)93HkhX4RRn(TMN@xVO46H_0p@;fp8zz)9?F8PQ3@@A|k@ynZu{2@fI zzhR!gKa=VbzvDBu9rio!a&?E_@#))E`5m8{JkPJpan5AFvN2O*gPk+NkJt@iQ;c*k z&aozo)9cq&Rb3ci{*0(bG^ut)A9si*ppR1_ZT9gg(eNCf9uh3YF&F!tL==_{_(LcGl_S5 z^UgA$zoL^1iq{0Un7D5KFC7qR$I(^@T z;@DDN9k1b+G?O+3?_1nDMy?8FM{-H@nT}0&$hG9>;#VxR_kuoFAiQ$T`YOn;^G`c3 zR}-<-+ePy&C?`<>J?w~a%W+=JcpecA%oTZ?$WM-$8vmc8^RKSN?W1pMtm&e0`=65C zI7h6RN!HA4!-$M?teMBGnfnbRGLMYMYJRyt$E;=?3o=PbcX_lp?_<;TD^8vA`^MVz zOi)Tn@yvW{ei(j*ue}QO5fT5T&p}M$y?9l*3}qB)T7wNv!DkITMCovjoHPUpa;m* z7zWmlW9R$`0qY6K189h2Y5%OVXlVUhbc>bEk=nZpOr}@N+UL7#_2P9tj`uAU8MP}G zcssuB<;s9K=x&A)86(K(3}Ed0E3~EmZ|7!c*T3>P+?s#$`|Ier8i(5jnV`2j1}#K$ zrU#{>x6cTwM#sC6OfX-}CM_IuHU9fujag99Bvhe-eFnRG10L_|u@m}TrimC@-##(kGX z_a}<#(VUR5Gp4F#zORl|u+lv5eapHjdGC)`&dPh=s&0PX`x6?Mi+lxTejp&WI z_vhea>QDfM`XRg6NZOyvuRw=g-{VUXXMs z+%;q~7_qht|2vAs8#1$lC-Sbmx#C_b&>9qKTX)!h{TBf;i?BkA|<4Bsouft;(L3)wO%HX97jEMWVK&>i9Fjv-rd#z-E16Kkm!{6U#gx_Bz)+E0tc)Q!mB;nZ+ zBDH_l2Ml&?$Sb3sL&^~Ui;gcPQ?iNqb8ZCm2%#!Wzg?R|1CU{BFU*;%k8+fM0px zJ+0p*aj-V-9POSa?t3Oew}+JTy{;<1-(OA=ZFo;_SxKwskXDr_xi2qq z;ibo2(5K9=YREr*LzYTZDcZ95W^$iqU>U2p>vR#jZLhdhuF#gomDI?Fkh;D9tZOSl zq>KTg3SNA0K!?-dh&0Oi2n9 z>}b1x`s7-#foEUWWL}ADS7XLxfG7X51oF3KfWR7#?(Q%|-xHK++O?qib_ghcNG(%q zc1wV^`Y7&|`ww*odUPsty$sx{p)tYJ)V+_yS89udpfZr98`>?@U*ysj=yerb21N9$ zQ6F{%OScX*WdfC~$JyM+peD1Myn2`=R1cNV%7DI-!ZLlo1P1ywK^i@lUOz1Na8vb$ zfItb=O*L2uHjzQ#22-!X^#S<;{f$6i%Tb}`QJ}I}rhiOrJ{HUDHkE)~&J6J|tO4{9 zYTHk8N~w$oVn_oRcZs`>`H*>D&?hKk4zR#xbuYKLC6k-NI?C#FUfU~FZxjHhIsH`< zutxvAu)`Bv6USBxFEYQC+^61(dM1G!exn4s2klKlFpbIT*$sIzLsN4%NC0D5;7I{+ z_ZOaDtVB1w^r~3uhm<+T1 z%3+G%0UH?T*8x-fewqve{SaX#O!5241_Jt(OJRy%#$*`ihX_S5#V^AK2Ku$b6u*MW zFwhSXR=^a$0vj0UhvfGmuzpENQiZSly8W1&@BcXsKPSOp1pVMII1v9931Rd|emY(w zsdfs#alCv!zi~qC8h+!%h6H{C<we%qIyGBfR7oJXIrA~L1o>}v({6If3= z?ic(1Ml(~z+&uT`+fF7|XOWIsDxLhc`qdKggvhG(XWO@g#@G~$0NIT)-)=;2#3jRm zS9`mJ_L z7R|({W`xE3E}f2yfYu?aH&oR_B!y6gvLFtbMpPpZcdu2G?eFB%9UoY}k+m~>f$6lb zX6;MdwQ)oJ;QBEc;6Xm3vIb-z!fEO=8Z>fF>!J5ImshoKDEI6Mk5J+2(X{j-iuAL^??0(Raj|d z4kQ5lfJcn5kn;o4jEqo-S8uMOKSojI3#3T`NCEf(pBU~&LNNRg;TX=(fc6`K^IOk* zA6J{gdw-n#1n+%9Z58i*^M+1d5b=~TKL~Oa{e2+d|8NZ@5}WE6ZOISENBf20(PR>Z zWAqqv#g48pio&U4A%LH2wj}t>l`}`UtV_y@E$Tg;b*}Z@v4;q>Ss{{-+^$?a!o3`@ z{Yn6*EMU86C$}PC`_%wWb-?!8ot(OWxa$E4Hv{4t0-oRwr26#X4;RwXD^)Kpn!t7R<;l}0fhttWY6z;;^)^nd|xxuX~(+EhTQy!go>y7j5KUaSp-womRvK3Jr^nHRXIZ z8Mjd{=l8DU{QS_zEi&Av1jNUxj`>@St=#Fk!GEFOsJJgOK64ru9|!0Jf{J=?T2&D2 z4P4Un59u4qF4Ph%Zyogv@!wZtdT?gOT)WxvvErmdM2k~nJy|rpq8|(n0%Ck`M4d1x z8J%RxnM-o%K*}NoCr9U%B0}aEBpb`_Mt@2Wb{c|A8-TjKX_33PrTV9*`QA$;P*3x{ zcd~$bn(uv8^wZOP?|@9y(|oJnfv%sP=6n0PetMd(ZY}^P>S@0G_Yw)z(|k?uvw+|S zMe_}ZcI494Y`VIOtW}G6n!(gCja*69@>M+$Za7pA0fA3*G_BX!8xLu!4`{`wwF&3h z>{oQP$WQA>SIK*1J=7;evi#FPC7Iq!E^MMZOJDh^y!I)p*nL{1AA5KitK3K{D+^+y6=gmy(wgCD;M+uVZJ4r80ungeWkYYKb#6i-vdX16q}+K8c1?7_eB zq^Vlxv%VD_4r_apwd!HnUG_<;Mma1SCJ%NsHxt?YJ<5ST#ek;0Er+eevWpt$t%N|sI9B*e=jcgj0A;hvF zAwYN>5l0EJOkr|_5X<_v4Ssp7c3D*?-PC}BIEFqni)M=W%qW)NzmtYGED&p5KG$M% zWm=a?Ew%t5bNMP7n` zo}#%Jlj8Om;VF-BISK8KEqVd&3&3i^xxt{g;{wBH}dvGt&+co&3)@35nDF&rCl=oZ)|FMmo_McgESD zNXDHZ<0=i#jy{AzoE;-vj5s^~&;i6*>+~u_#w34?Va8;^DhYJ@Zo6)TbZ35xIjbbl zVZ@ob5zK?}|3J;m`)IQT!pVlEIAN>+hjJ#ls0^x2BIgB*5KW`kq*>BqEa@q_5sXh- z($84ZkLpG+55{0k5z*h;@a7c%uMCSE{7)I)bn<_0c#Dj)MqZzV`${jU@a_c_s04p| zWdA1G59Z=h#D>pE$tK6-2J=WbdsO#6+DYhBN!Xsr88%4*H4Wah)B)HVyks@(@wDy= z><9&>0d<8xG?SS{vt4!@_f`s?5ri0I2wiv)le864FkKj7)bVq4^SG(^;g3(cuD~Wq zpt?deY zZ?xbfR@jON`Aj(4BovLiT}_rb|3K=w_qbsqo`22h)UKP%uweE$>wRW|*5 zy*6xj3>$s?v2EepaVHYGhx}I(MSya|Kj{FZ|GB!u?@XihH0J>!kI8v3UbvR?U|ilj z&Vw<6iHK(0Nwq=qXtVpUI})VR3%wQ*>fDi>VVfXqHSCTQas}2%0MkIVKsXrWc`{OD8zZv#OmkE-eD!jGojPvjLR@~YsG2*%Zk#F|9jU3eUk*VnwhSK$4r zF}3t_jPYDBu0P9vX57cq0s?IBjWP;6SK_&5d|e{V1@$utr!x4(1T{Ck5>ae^X^Vj? z@<)Q8P;k~xVo8!LpD3{;NX}YHEOCpi`7Cbm9045O6IZc3}h(( z$+Y_hcBmKQaO-LJ5Vq*ilTa0CE@X*2lPeT@geQ^Qw&5mBL64q*N`MeSI zx#^9hO;bjMqq!Y1>Gw#RC}?A7r^*d(AQ7gF2uE{UVuGtl1PXdQv{U5=brNnGB%S*E9sY-FT9X9+DCfy8W(Z%jnTOEOuoNp`K#q|1lZ3Afv=?`@ zZ)oqYuRQ+r_&}diV9-iu{^d1$HhikSzoX@Bd-v+9xThxu{O}(o1ZBj^vdVcz~BJ^jJ93!}RlGcLIYR6z;(jZ2E(3i6+qg z`H~2Fub*%sGI)i%yVInjme`sR>F^cq#!iz`U@dx+jb){fG~X(T_(CL@{>M!cDv*hj z4#}fMCXuvswBZl>gKas*iO?~JZ!U?xTP3l+5c&2`#}O)GiI&lC&m zM4WeQf$eNlTW56FJkBLM{|Hn`d~_2QU61(AmGrIgvqFO+OI(DZ%$u}nHi2tpcHAtU zWE^80V;foQe=`{UmCMTMggi^L!=sFyjbh@twr<7#KLn#cciA+Rz_l@dpM*$^2fJ?PuKb?{&i zF>L7Iw&tzF2U)~zGWqw--^%4@c;3;r1cARPbd<$qlL3LN#T?hhrE9rmY)(BN*F_ih zlQ+rmTqWkXIxbzq^=8|iYu8$ND7U06ciB*b(3Fvo}E(v{qDrtPyL|IKvt74phw_`Dv>;X$K|3NGuD zm1xMtckpqrZGrXl7lAQx;w_i_>E`20&KGl;pR5!@Pf8vKb0GiUO?5VgUFOH7@wBKI z`$F48P5z6*&{?i4>ulV*%nwVAeWSY8@m+cXebgmcyh&2bPS$tTMQqByU zwo<-(?qH^Tv5M!?z084GlAe5epe2N|y)pqSPqZ(zJJj^$qA=5}+)Z_>4ZCI>m!6sWUE{r_X@E-vE(zR_t`qbZvd8iC7sFk<=`V&JZ=}B%{B;GfPg>?p>>Eh6 zFPw3xX~m*2OjhpZx_ZN|t;eM|Xi=$lg$E&D+(lunBhx_js=9i&uC0fqb9|%TpC&fV zeO4GsL}};+nRp`aa@)>aM)@gh(}3gvW#_|+9K|MJiGSf*x_+iRDjuJrJIcEf>5j6; zQFO=QU$+p+LuFqP$%6^@g;NeSeX*zyofWaEZmwb1#N*QMX;BGwg)xxl?V>*6$TU#B zvTm+h*Tlop(Y{e1R1enYx! z_~FfgVak%VyrM45Fk|opxk2))4TZ(WhYiabS)1$f4ZGGJmtLerrPvob9co&>sBctO z#HzY{x30B^rOv)1(qOe+;Q`3@c2S=wD*{aa7c|-aYe?;+r$dB1lc)2H`?fQa5scuL z(y_#cGVvYg2kf_-1TOpUcO!+OEO`*(|?W=9?WDk}eR{emN z4-L~bUFVxwlBdl>eQ|gAbXiTOYOKbMy|bMk@x4ZIm@TVP_1CEA+L>x@kEX}9nKfLN z(##xw{FJRpRmrt7MJiQY)vTQaL3(_3H;+ABYQ)zJY84_LUtU4f3QqH+5Bc)_^os+5 zJk=t|Zns$#Qa{k^u(n01Ro(8CYetK^h%X}Osz*7)r4Gs;B90Tu!^tnRl|Dlc)enag zG(GIOd^I*jDajF^A2_eNeO~QKe07xlIIKvf)IOsQMyNXYm}EcA(;*+JjdF_OpxQ?9 ztLD{J>T0!OSS#-QhABxULsO&v6wE5A(e(CcrAPUnO;mnF)oOe5wXFTJJyd0Wvs#p_ zE=^{YXKPw2*mn~&i*2PhbN4I!&X#vpX<=xu*HlOIE!NR|*Pp;NpA8p&!Z}Ds^9?Hg zkLLRX`I_ci&C%0*dYVs9^GW^*rKqR*UbBiIdk8X#VEh-@Q+u}{fV#9sz1jcUkx`U66y9zrh-PEyjZm^?LO4VPXG4Mu& z3f8E8ru}Mpr^;}{YC#U>7?V@r;y=e&T1k6mb#5_o-8Ycw%X_kx$OXZ(QQJI%rAUbC zz6N{9>_UcrB5*}+oGBRA3o7V{Ko{89Dfez-YG~OTQJ9c@9c*UYyCyivNEQ%mX4JbR zI0+@5FTkQY+y%*cd!#pSk0j#=+^~5Sh}PNBMB}8Z(R;_a_!^2mqNye)D-1N_F1f+L zUrrUw0Um%<{{Deo5SRw4F;`(c@RW`N1kzwN;5h`Qf$I0ONx|eV35Q?q8@KTG;>mGi z?w(j!JpY>g7`MaS`%(U}z)kwc3rXsXwm@zhMX-$yjkZ8+GbA7k=vq~KSdhA;EcvP^ z^pdFa3OptNvB_0Y1@!O=Jc4;JHqqDU?oY)|&$6+V>HZdR?jHp#$w(SyxqRVO&tl7; zFX2++u1&_7ai*bBR9X+Bc=Wj)v~c7FYr(x}N8@mE_~u0nqlJv+J`5vo#^wbKqxp>G zUJN5o#wriSb{b;=l|hvq;PDm5tLR0V;|9X9wEYI3B9GzHm(lhiJ~4?UxkbXn-}g2! z2pD4mgMfKD5HK7!fr9?t3IYKQ;e&0TK?mE6WG+$B*aGN$+b!r3Bnzl`eavkxbj;0C zhKu3|bjR0R;gf7Y7!4m^L(hUwvfb7lUvq^|vH@W%e0*);Echf_iSGECD}0g-2)Xd_ zHOE=-Nw#v`@ikZYBpVRM!^hV)&4N#|73+?#xxy#efRF1ilgX7rBnhp@=X&C!(6=uX5S zjg*ECU%@b2!C0dk!8mvYBXR{}iEae*VEjK&lQhfF%OY=|zy;-I>w@B7&}LMF0Ynh? zWMp(MQk7qMJmq**7zFV~QA8?z{(Q{6XBtwvVTNPeiwg@PJXr9a8fsB$H(#;7DX z-PgEe6ZDb~`Y;^A%%KtbL-1P+{qJ7|J-h;sU>=MyA9l@p#@;krU-gF2^^t|=a zxSu_*JB_RKd?S~5J5AOXIL-dqGxwv}iJk=xv$uHW+RgslGjB?pljl|2HcUy0PSJkN ziZ#Yg&S{n@#feoLU~zcvs_l{sc}sMv^FSJ`J_(D_SAAVpc2xwXf$BHjgta{V<*2p| zoGUYgPmL;zXh8%=8)Vc5(o3^PXE~QXfqu>*b^iBYE$a&=yI33Ml3WxBNayuAB6otz zEkvHRFcoo?Al%C-*MV?rS2P$qZA*g{-s?BOd|bRv;eF|f2KNzZu==E>YfWPP*RUQu zn0~iSIwImL>-?Lnr+wqa3OhG_bo;%GE1zXWB$V{~H%=PW;|0tPob(%B%X-GgF2=|R zXP0tA#x$3B!w7p9Z^OvRLKDNClY|u<=-p?Cuv?=;>C#}74ti|^fp^exc?hLjq66=M zG+2ES=JB9cJP>#Xrh)4Bv&kf^V-&3enMf2OIFr#|d!`c%tu4@rv<)aO^k?NizZ>gS zwBR(8hSuoq)&eZDe3f9{D1rk*MiTIl)gcD#Q7h9dA>0nq9EM(1l<9E0wK}2ula}O~ zL`ZTNdT9Zsf$BJ&qW#hp1@2BkKpL$6K%11$%vfi-_VY~(JtvH@{=RlYZ2aU{3)?XC zQqO%D&nQ3Ax2BIKaxU4X4Fa{g&MpDt3EWX;)Fn>eZdcw(1k za`Jo;K2MCXmcF68efaiGbOU(Eoqh5vQ*G$g&qV(6GpC<(R!4#n4 zH3j(SH3b+9Jt;e1?8$gd0iJ@vKQB8N{7ct?e*rok@B|G0c}*?k91Xhl5`T+zy&Hdu zb?XH_igmpM{}t;-Tg6<~Wl@EVCQHZqE!QD{r(gtdnH`J(eyxN07U&Q_APrWZfXTpR zQ(*)UOas+_tjX=wl+2EudnuVON>eGB-8=IsnJ;7NDMh`rl$0WI=6f^G-(T!mB_*SG zSv@7=WoAAlqkCB@CF4cNUP?j7*wvJRHm6fC=(pGo=Kp@w@qYn2{_g}#^evtW^M7C( zs9vwb0_W@aKOhZOzv(6=POQw1SuYoCnz4z;)l0o`>9Ge^m%=l3g*zYWa3cSRSN47@e~aD`Pjj@TZWGR z3()a@Ct#w_XDZD9foY(6oem3}ujBuKG+6y!HYrX~mhWy}+h?O>GK#~#l{H%OFSjP7 zs8lDK*9;8oWxa{DoDEZzFzJP_FcD0}PvnMy$E%CND?-C$GIdZvV{3%+tPb>BuH*ks z!TevK9gMbZ&=GwLbo?KX2CGlNSYY5(nEwOQK=p^(q-;a7$KEydA;L*(h6x1~gY%5V zoH(o>V@Yf26zHe)x0?hm`|o!nNbwx+%p9!5J2Mow@XmV&f5tmU4V|hW?~x>Sk|R$^ zDI6)~xNZdF6H-c=lyX=%f_X63*Sx+zX@gc5YvP#kFH1M`)+yICgubAh9F`Q-pjaox zN=aY4D_t#aZPrN6li4yJl{FaAtS!xjd~NgDvR4sQCT#CpqwI2(`Czl0tqDD{rR-|sFY6qK3$udD-ks>>XtrCQ>23+Zyhq-K77dda>C`%DR9Z2X= ziH`E`_VKj`*x&ALQfK+eTBy_kCMyrgP~a7cDxR`cExX5RS=*kStVAkbX#}i%wyNbv zZHHgyu(rLGrC2Q!F%I)&o3&Me$~J0uE?qOoe#ujWs)pFF@xPoGVK3 zG~3jvd~H6J{Un92d=)`gKOI-?C;K&rFAt;2%U(PD7W8S9KDWtAD64r)S5^?!Nnz^N zWbG-kI?M!i_?5EfZebLuwDKC|<5|Vp!Na^y#;C4fwOLo6&{=Jrqfu1t9=<}XU^Cl7 zPY1DoO<+ICkq*Rjc*@x?pPGJe$ZyM}wg&DIVl-mBY9?7$N*9|HL8^Ym!~rfN2HDOBlJ=z%=5=|ysQHudE%8sV^FO4Ca$J69{rVYTDQ1u=A}!KyF5s z93x^sUvwuo_VcBV!?B-#vCuWnbNRX5ah|L8U4=CVZ(VRRUy=4#5`{o-kU#uGK1@jw zao)>c%!%{zL3l%bbmC=4i<^IC^ zsPcsdekkY-`};_hn~NK)-IrVii`d0LY62@Ic{)iy7g(e(`7SK!2d05)T7A9R^K;bSL?vI7`suyl7JlJuJ>8(ZZoWEC2c3NFt(>2BP;)aPdXT)(PA& zc|M5NsRoc4>CWgqWW;=~rlQ9ajebquI z!(!YyE=`QQHVOAKdYTE!Rq!*S!EE1pgJHvz6So7SZwG#NJ23cm%8J{8wA+E7+zuq& zPO-d=H@O`sEgg>G%jD48D3eEC<`L~g#Rc(SAAdb{f3g38tf{R0ask>0AY-j!3l`)7GqdU;Eid6$HEcklEr ziSaIRPd7vL&reT6sg|Uh8TALGCmBiCkj#v`za}LayEzpZI~6(WMlg19Dsp!!ve%7Z z9*jTmn)2L)^8VV!iPq=67&|vjAKh4(ad}Q`ZoI0lw!y03sq}3Ta76_FQ;3_^{OaVMs$ji z>Y`wcr8t!fvUdRL_Qpv)pf`9)fEjGF@jN^E*j&%EQ`=U1p0z)=*Ym7nR;p*lw19li zjOpMv0lnVAJb!;C<@uhO({`nLW;)*9>zQf4Yqe+Q)Y!S6MU!XScovbS1MD4$IvZ zNZyA^>izz5S|rgwntV-3mmwzM9Jw@mRHALV1t%RH;7xLa{;d3lGydtjQI18G?-9%! zO~^*5T!Pb#8Zx7|kBP<_sy-FC8;i4}LG}(n^8R~Cz2659*H0L0eWLbkZ2Z*Qwizwx zwceW?JqZ^5@wHQ9^GI34sFD-3Mhmo!VHB10t;wUt=u1|?!-ldIv<6e54JZ9UFnXg) zz{dnG!OYq$UTEai&b8fC9DF%gxWZ-rG(z4uGwV_DLX_7AuB~Hn`r%--w@bie0vBgC z9&Ox<^lIa9dIag)gVA$b0;~yKGc#-Bc%h-!29E7HLGTyBLVK6_V+eU@Ge)heJu~8w z>qaKynd?u?hz8dVW~8vwfwi-(Gn8eGhz~aKvWvDw3(^t6Xbm|)gXd~6Sb|F7KgO0vW z4(Pyh`!VC?M%z@LQ&}7gKb=EI7m)*+@Z4UE^>AFU((^3S_D)goI>^R`JpTzkulxU@ z_x>NWZ|Jp#-|fH{Bh8jeAixK3eZh;ugjufhVL`vMRjiDj!1Jt(uIvM>jLyJsSQ+i1 zi&+J&cqdjtiyy$=ft1>tCiTGHz$HQ4ub_R_7D{gCtVBvd*Q}o@xjnNgDS5A2dni}? zS}|1Mo4{W0V4lA_lb~CDW4kHWUyZG#-0B(oGv#{M*hI>W&eCrvm)lDh7lG{EJ0;-$ z9cHjVnT=VcV|B9>dVf0h)#M>u*9=ItaYM56V|$>>cj@I$A|8q9f@MWqpp@*_2k{Z`*lB5Nwhit43A zemBIuLTZJ7Il6e(PhFBhpErQzssimw~=}vfd?!e1P|8?5w2BSpJ&`#{9SVK-Zjl3 z2E>&z@$HJ)6!(7UpA}L${8#Yf*&Jp|U0@EQr6fCNZ_AayoT!!`Lv!}1j^J}5HC*T) zb?m~vWx6pFh(i|jl|m17BbWzcea-9pQ@V=RqFD1R^bRF$$m&K-88f~&W)HDXxkQ*? z$|q`rSObr-RLv_QU(uhSmA3}6ht*Vu@*V`SB|!My!-U30KY0O<6~k6)eYN#9Y)Pw~ z97Ao(=WF5bQKx7n+3Ig>`}>Hh&K^j~Z-+*7RQvopDUmL(;Qv6^p4N6!f1*onYWrVB zK;&ONMC5g+QnjaHDZh?-s%C($ou}erRlVog3PlB0)#1t>Xs0rj2|NvVF10J0Jv@t@ zcOL=-Q&k!zUsiFWv(ZlxOVkX^YVJ&!E68SNVs-i z^)|NRW#1gWT*+4sgfiv%&|P@UURtb*<*5~^+M%QLwC~l&c#tP=FP}YMbp^`y(AC%3 zss}Z)Csd6rM_#AhjZ-w#-&W;j>a_f{1@w7iRdk|!)s$kgQtbEHH|lmDRxeFFz#Be~ z&Cx#fVLj;@Okv7G)t^(t2(p`grL5CJNAY5cVixQ82!3w_*jEUd^A;2PxKx>uU$ZDV zJek_PRD+I?do_0?hx4f22cZXw(xz@KD~Bq&rx8aev6}8cXSQ4sXc+BhU}xB7=I}phL z4jUW+p%8X8sx*TQjwGj(|`IyBd|7!3IYX z*wyHXIc#tQge|bE(GxSs;An)O5nyvv>>-wsA~)jP-O=HGUx!Z)9BXwFojc#yZin4r z#QtEP+p{JwvyklD$392f^D~Wgfwt#jH2V^5&lMW$8g0*wXjUFA@)nI*K#MG*x#@)s z{@Le5U>=Nt&yxP7umL2$4++{prvjmIM7*&Ka<+qx1w!M9JXjCZr0`6RacF1I#|}sH zOpj&I&Wt*CIrEcO9}TfUL2{5vyw-)dw-wJRsx%J&NSZeQ5wgJl`{uCt%S zi=P~#Tc4Wi#o9ms*gFu(`(vcu?>8yNBR`)-3PgVHDo8}d4l%tHIT0R(*#_aUYJj`YWoit@*-EP0rdJn89x^%=QEZmcM}hUmVs%Y+UTn2%{I{s z^Sxm)yElUS-RmQ(Wa0-k;Coi z6`Rkm*sA@wL4Ei3Kh!Y%ASl3s2d4Q)P*J~8dmXU zzKqFcs`6^mh9$@k3IvR5a!67>+RSr#nHkC~Ip=p;Y8pe?eLo>te)}UL1pd$bSt2YMl9+ZY2?tzoRoHYOSj4JX_X_ZDy#;!ZhOc zTx!tEURLEGOJ%Gb)_$r^;b|XZ?c}q_{jd7?s)vDmwJ3-Gsu^>b%@^}v3{Yq28|cGU z-=V5Du;ihd9PLR~@psz6)G)qMRMl6c?CzmzmuggOhWZ{))85C%Qd_yrtic+-vV}`$ z(`P{rx>{woA6wps*GhAm83Q%VTDcbD_N0$!0N;t01!m#yQ0@D|lLVZU`5Rj%q4#HKcdSq}%)P70 zet4&oIzjyKsj_Tqsdm@hnr>@WQyI0(Z74TI)o`XIJ7@SnO|JH#9knM~`G3B^gP7O8 z!1;0i=L@_>oIsX^j90Z_Gy)^Gj7!~EOoY?{zY-@n%UUf}#h5ciWyF?A>_@38i{Ac( zqI<*^3j2|b+M-vQQ23ar3S!?*Rk!pF5gCI{?c}HYF$N4nM+;JD^ccfkqtG+cJECzY zS+Ta%O~KAf&|+IyBMLhM|cYhhqcF{ZMZB8 zFf=vs==V6Z5jT~HiZk)B^6=k=J7xifMkXE-5C4t0aYWQX6Av4Y+HJUN7GP*@;xXV+ zyAd~=h=M$x2p%rma7h+m2q{E}JzO^8%!sH1CdJksm$u<9Sb!nMq`2SX(nj2LA_|M; zhN6mh#t5xE#9~sUH_qJ^eZ~TX9nGyUDV_yWhX2wZN#$i!xy-0P#Hezokt)Wh@~BaN zrcvcpBWbBo$wQ;=R-=*uqmr4lPfhNaAlj!2h(^b*Ml%ThzDVhInj7|S6<`URU8Sv4kq0i_DxEt~AV?A;>q( z{3v{{VWwU5T*D$;7bT)-qFz@)udAT5XnMB_5ABwROZI*3lSP{xH0GmY3vbLHW&h(R zhXwls{H8`v2BrpTQvBTyYQ4Kj5+Qt<0F7`*CZWT}3eX058hO9#dAcJj&5{~pNlnp> zV0_Y&dd8A^R5yZoFxJ<+zCX|(`ZNBW)=Vc>7vo4)1DX zVuTV8aWzTEnA?9TD|##6ZyEo$~HS#W#GkKm3HjKSSyDG_;W zXEX!n?`x>~g7zuqP6DERHjQW!I}ja3!udk#sT5-;inDG6V;3iiyA#D;H-dRE{vW6* zx8h`=*US$l7^dL}T!9fPo)|4R$T3;sE~rOn*0xtxRkr6Icgk_TlP zNpYI(9oL&xOBg5G=^3)*g|F#BTU+-jZ9cIbv4@Db@%LeM_x3BIk5_JQxsYeJVvVzt zbBrav8b0W@;R<}B4NL>o#k$jS_VBSdJdg&f-*gkndBY8XBr@jMRs`SKf5@o6g!U_; z)(N>TG~L-e_7~{)lDKfPvK~ElH7nY$>k?xZ@A#3xgZn4`Z##BvFs+TRVVfC z?6fV$()E5~eQ?>UE3g85sZMd;9zOL3rorm!#MIETtyf@udNBQNn+(kL+OXL%%>4Lu zTfS?+_=Nrr|5X#~^1asWIkycPHwgW{{I*1J(ruxf%@3i-wbgOjarcYv0<^rn^?$9NO=YDZNv=*z2`&5PiU@j{0j7cK zk_(Z0RutGf1)Z@Z18K1O18vfJvi4G}v)$~ewxdnWeO`R*>19!}?~rS(z3puy+d>E2 z4OI8{v}d?lC1Ty&^j_oGdj|PnW z7$I}HwqF(t&IZ@|yPzBh)5e;4m>M5MdAV>i#uVdXgD=rs!ij{*ShF=I#s`sJmpB=9 z0^HZZwJt6wJYgEf%)`j|fFaVAV|zp(SP(2cSoEMP|Cgc%J+1eO9=ytbRrH{*)STKZ zNu5D$9(ag$K_Lm#5M~|*#^wfIVs=J9A19?>l99ue_{nO_8jX>;)-#-)afOe2M6Z28 zMh)Pn$uS-(BaFsym}R?zFSt&xeMUz0;HM2@JQPOeYR_<1#>afzG5V!D=C z%*b5jDQ0H$7vXl%YxBscdi=Ck7>^+%jM6ZIX?vweu$*2C`TMp!4$dgDo*oN{0IIiK z!UdU+?}q$+BR^Sr33^fmng2philHZ$MtPA6GmbRj7KYVkxuWW9rgfQfq{fS*kPd~m zN16nS$n##}H>$hO(N9Vr#<14i&-q>+_Yrk&TScf!w2#OWJ1mV_Mk@T&gf06AxK+} zQbHbe-d7~8Bf7t8tzdQ(Bok#1fY|}pd*jmpdjILQ)D2M9HDvoLW7iSq^^Emqc4zq! zq?dP67}XBGa-#1!ua1qtLyE9bXaf!zpdicJqh1}SW znaF=G#;7Tcy_m+L=PxmZl=|-0(?!o-v@%)HJdNt=r*&=RNark#T1_ha94OvP zVLb@lrjRWs01ZJR-Zv3e zF+=kEE)E}5Q;LMHb5y;F&CevIi&?Cyu#0l}JnEZROPwpLr6{LqC@!DEe2DLrch8{b zwO(~)s95BtUP->XE-S*QE`M6rT8`9tX;dJo@cfY`%Y}VcvLd$EHd1Rxggi zI}|z{X)0UP2d4iEnuxAcq|dwfYGye0;$E|+R>dTwOi+^GK>111av1t4{p}`!%l`Y_ zAg@(=cAY;?9!FU_?D^~Y<>YbmRu4qpJz!abcV8OCqebyPheznQ?(gy9D57r^Zy`L6 z$m?rf-z)I`+*g=A?JS{q4YgdWxx{qtowbWtt&D6AHKZ<)sl^fU!N6v%%$2R^KG3J_ z^rNdwBiOQz9tfRFAoJCd^%29%E)KW4%J|xvF1|{Y!*aaK+jBUKt!?KODPG~(>aE(K zeJW*~c1?&(iVQE=<(6CkO!)$18jAL6mwAt0e;FxJe@V%7)aO3!}>e5tvz8H zDVeR^E3?9@hR(A!jaa-&@l@M$fUW5uvQ@|kR-K==!+~}4KAqX`Ab%CX9!TC=lh4}~ z*QXsA%j;Ca>VWCZOvy5S6PEWf6jBE)XrNx<^%8&e8}@6)tDkCjvy?>V;Z}M}U}vCw zf1omq^4vlDb2CFSti8aQg@EG_X#WGa7lCP@`i%<~ z=!w3!uSXy%S+(8jLf#r3VFx%Kfiqz(Pk%Y8Z3E}Z4B=Cw$|712!O;d8wSn}~?9o}y zCFs-L-|YZ22Ct_l`nFyg;j47g(zPbB{_C=PS4X4)ho|^e+XpV>EnU&z?zHVaaZ1t? zeNK}A5eL0f61gT3l2NX@DuOr;>kE-<^+cbZ`umr}qn_x~6McH3@82aA^+exmlJFgo zfqJ6v&C`j$j|J+9zIQ+;>WRL0K_}{oK0VO~(Sjg0_YSZ?J<<36VS#$0Pfzsyr&yq# z=+hH@{|Vx#C;Ie6A504Z%-lP`0`)|n^WQV*KOIi679nj<=I~Z+FfNen^CN$>(ut>!Q4PqwY@N#JiNH(y!N3b zwX;Ama8dm!y)$&8hTXmXQQ+N)D)#;3Z28m3XOpRshl!W(V1wEcX7tD4Y1Y?-4b-~cslM0i6smaG|ffGQD@Z^;%rYQ8+cj~x--kYUK0 zF$ef!QP?fnYDdkA1N^uV;SHG$Su_3sUm_yjlKD7lzBs^+9TDD;nUFOW2iRf}{gw=x z_3{JIGQjYL3{Tcr9)MQC-X)=C$ar6={2VntcYBP$%4HBjZ=h^o(u`E`5c?)_OjT5pCMaY9_GfH&~J03`K&6wJz*v_N51ra z&0dyfcUVJTqfOEFii$9K`Ml=BE}ICZ>_$!V@I4zU8+NXT0R1m?BV?HVTMP-_LWI7B zZUpmS{6A1bF6_YHP<0=nKa)IkVy!F9ahAoYDfvwmxk}LvV)HBdDQOT#8gyJYg7FDy zP?|L8ux)AL}h;ZyozQHdq`&e>!?$ z`X}Gz?438r%lw(O^9C&{)voa1(Wdc>!dkO(H`mo0c5OW_h12%1MZk(hVVJDkRdw}l zU0V-J=lG6DgVjJ2iTccC;Z9dZc_Ox1VzrnOQ5AYYE}Tc?znJC9RQ;R`H2x+%h| z)*hE$q(!CJ7djnkf?lVL%8FQ3m+#iK_OR61cSIVjwktevv?vp~OYFp)u`l!^75dr1BlKJM_t>rw>rm)76&^?A^);{W6>xvBq0IQ6o$GkE z%JAo*&nWm|t6SI(Wb#+y#Sj48@fqdr9t2KRt4}PH(J{+E-kW#cI ziNpQ-hEK&0_s0&GaEJS&hflGF`x(POMh(|R4oB@CW(kdw0}H7em_` zZQ`(j6uM&LtAR513lY^*LcA_Z_E#l1Yo?6TesMsHE@M9vQ6nY9%Q9?`YL}1JRw^ZK z353235_g3{UxtWVBA_qB#BUR{oM9PpeTuetSnh#mqsxd}Bnkd1AS@ow7M2m$OH6`P zi}37GWxOqtY=1DkJDx2nS@RZ2 zu0I&w8_#YnYhEvL4N`gGq1jWnNTB5aVd;3Su#CE10<8wx#E&ZDZ;|BqgCVqs7nSkX zOUOac?pEV4PP^=Y*0fCfLd5ozP_N5){;K2|ngnPse1I6Cye6^OmpqMu5_%~oafZ*` z@)69t@Tz!ovg6+HXFE65?M2n?4bY8Xys~a@c-`Llx=~zTxSbynCM%yv)r?8t{!&Iu%(EpOqS#=*fc!YjA{~p^FLWk9Tro!Wh{Qp>u^c4MB|M|tNbrpU& zvN&`8Eg^nL`7(;w+!fuKE=IvmV3FEyKI?zd5C-OPWSQ{ckxkoRi=B#Rd-dXyJ}Ul+spiHW@}zD zZ}=dRzgOh?y-NH8k19>gW@(=J9c_aC{vDSc5GI?FYn&c?9I z{DjolH_B^jA?N6bG*BJ9C`_2;I={}wt;-xp|C1&H8)oLih}F!@rxmHp%=(CWW@dAO z4Xdctbv3I9dN=9sR|b0hu}gqgey2@>qMnhf6_3B*yL4rns7uo569+KMSiA2e*T~8& zsFW7JQ_`PkQOoQKKRVQuwJ2;NC#3eiQOl+l?ms#r4OD-)x}YxA zt!pWe{_{;3m7z)UhEI7jUz9p9%g@!E9<(Z@L_V)bR#xvJQpK6T?11aN@o50P|MXhU zb}~!y{7w&tBSwPDwcQ_>T(sKs`L zc88kQFA8hP%AH@g+HFJ{tUe(f?;EvvYGK^bCNTY1HeuWjtCls|@@KYJRH(|+nkPza zjlB@0h3a`1XKyUDpC*V!$>Z$OS7EA-x8+AL~1-dOQkOr%PCgsC>%egg; zVk;*Kqat*=ykQ3Ka!1TCL-k`#gH&>!^7|dY#z6HSnnsY~5qW~*@hpCLTehF7ghQX# zhpAu=T}vhom0hK<9uZH$dspaH3ACk!-gbw!v;n#;?Fo333cV4OYLGO^Q>L<-42L_Cc1kjN-6wWsR2n%dH71D%FYRH3I{CAtTEF{DzzjQhK)0m@(qQ#F zu)KRaGYPUDH2ulXj>&zI>A@4q1MWL)qKqGCIpLSZxYicasZ{N1PJ10^M$krLA4vn?pN^Fo?!?t zjrs;l+M@Q-#|&* zp5mSY+V-H@Ib`=!_-oaeTifxUMO^<59~E(H5B^rfbtYaBaf4GaH}Y~^#okClt5KMw zBrCFbaY3|gka#qCNj$!w1T(J1XsftnN~8y-AQ}-Q);KR=;R}i|gl40#zLKnn;;MXF z%i7vrXZPLsYqv0jMx(IalBp5J2lAtf*NU5*mu$lq+`x>hH`@BDBy+dN;{520wc`8E zOE%&Qu3!ki8il=}iRm`8`RQ+_Lk@oq8XnLxp>YJ-2 zjIArgIrm0Mudnieo-BZ#R6|dGx>cMuXNj+kd$#$>hsHIVOWUp$?|e@CbZc$e9Cu%v z+1YUkQdLsZga-p-qatZFHJ$e)?$pyh2?Nq^n{4hqJSvajS;Ay(1A6^pG_sI&z|k^lkGXNF_TeOEudZt6ULhf^xO{o40a#K#hWBLbU#)n`feAJ7PzJnQT;29##rDTO^Iz6y z?eJW^VD1twQ9jSF~jsnI=}(x+E{>W|woP_MF?ZW3^8;BF*$pQ6VVK30E{`q?S^mxT># z-=*DIF5B~Kl=jg7(boAd-3)6Y5c|!4d1n60eh+vBaWM8X<{0kJnVqNP-GOW;{DKn? zO{-Qf-XXjc{`J(J7mK2nE-#7JPm^d`PKtET(^;IMTRe?&KTmzPS3J^J$kO)o=Sw8N zJk=)ECg@*k^YsZMbxhs<=x#-IXlu8OtAi|bwU1cb`;yGb`Y)gMemYvM57l*)^YVA; z`UDc~8BLqNhN-tq7>Urgp3m10JnH0)+=YeyI#oS9ecySeZd9fpuTu|-buBV(|I0zK z=9>?>+Dg5)C7q}1t>fyctc1a6nWkBbM(BnnaWEQjkv2TcnPUI{re)#Q*AR&a;2AhhdS)v(V8jIX7;)aKg##g zkK|X#mrl`%^^atb&oH-dDfg~S4+#Yl^zG65R_>6yM16xg+DO$xLV=A`{X5i=y41gk zcK51O(N5_idLK`&o}?$)*)dBH4ct9!3jSG-?; zA%CEA9+vsVycgZ*$%)kMT_WWg+hRG1rKp~zs<6@s2s#CZcT>pZlu>__lR_2M&Z?Kq6Jh+B+k>_ZlLI>D; z-!^eXKbuJY#QPS)4zEbmr^XKGV@P^xl~Rw2Ht|ZEB(K7O)K6+4)MYxNkWE@_JoFa7!#Tc<%zVL zG`tsnJV`iLNP9|K<}fBS7q%tRZqb%`;m47LbA;JXXfqtfgyzDQ#O#~28D4lBlJH$& z_ETDv!|{9JF#VMLdfS=kw%N(O7?WX+fG5JL$1JB$X3Z@coEiV+q^@9y)i2` z2G|FWw?2s9vB=Y5j-^HV9#!hNqj}@ns>dDe9H+t&jye(AmJ*Ie5}u_JlJW@mstHM* zgrso(HuHldWX*CubKHzbgqo3^g?-k_kFgCPq09H0L{>t+1l1AniUZ?Km5>8L^+tFd zga4Q1csMX73ijiPe*6hYKXNf&sz5QQ@OHV*dUHs5GT~If^k?7XJ-&+m0<;O zFa~qfT<^{Jwy(T5o2pKGZywin(R(xDnbCk$pn01RgTu_BtJ0BegHH{Xb5*Zoa1 zTys2~02nYFO>5T+3JOZwVij0C`7HUsuc%De%)D}gYO z{mP4IzVm)$dA$}Hv0j=_IGP~1$_RKDJ5b~O(E8j8RM=WQ+vDK(<*S z?3Z-X;I0I~VD=j*qPZQg&ZcI)U&NN{y=EDeIalWg2Bi-41gy3FuF)?n0s49Q?TEmp zU-eni;;>7?-lk*B_GL$&`VaB#GyHU>>(6q0mB3KOG>YW{gNOrNJs&p*kGf? zfc=uvZ@|V#APi(f3crwF(s6^k5&(nQZ)K5+R>We`b&j_Z;Z!wX5mYl*Z2v}Nx220R!hCDPQiWp4PxBw=}C zb{cJlJH7@FhQh>b6>WwaejiC#m6&~i7UhnAj0eM%#B4P!$_;;lBrHp0r_m_x_$oXY z3KH2W8pREtLK0RcvJcSKy5oPrgJE(aTTNT*hCfCUrsCko*UN^-*Le8x_4H);sTPMT zYfY4Y?IUF_ zJyHhR!;XO#5b#8%Nk3u(>3BT92&E~h$;iy8jdl*+5PVybAF?y-vZiWVdq(@fg4*QF zW9Or1TWwf>n;DwxxnUZUmFvmA#th5#{A3z4JlAvMB4)%i&rP{rk<&b*rVXIp_oKz3 z-(MjeM!)|NEe`YkdhcP3Z+P|&tS|hT@3Vpy$e9mXp}#YKi=lrXYN3~x;1$HdnDE!m z`Gm7OY8L6~)0`jPyXKx+)l&a-vdw3H@t<9%o83BPCz(UcvP>=8R}oPeR?q^8Xj3kD zgal2&I-4K2#MSMq*jQOsI0go?`Rig87d%!5P4RR#2jKssh#ZHi+30O+Hlq;1^~hd{ zYg6*)P_uF9@+9x67I-)0?^gf|1KEFU8?UKIRVKFTwlvh5G8x4QW4|Cg`BL!1YToU267#Qa*S|Kv%!bDXC^pz-fz_e!krI&!w5R&(E1NBiAXp{9M2(%W)OH zzh40?3}h#*i|xDMu_`Fb(|I)jgW13~Kp4pWLq$H}1VyEdd%8Twb;h;Mq^jffr^eg( z`Nl7~o{qm|YR9C_dOZs8|BRab@9(4~7*li6d&8s|jYw{2_6tH=vfz?yb}+X5ws)uv zem&~4;0IX$WYtnfaL?H~QQx^Y&Xgm)^^bOD*pU zcxpSY+wb*k#Dto$(lX)tPYuS>(j{7+LHfpNuS|${Ph`>y(rW@gYH;SZCWT z%k;8+(C^s83t9y0l!^=SD}y{e$G~7Vpdk;!Z+#J4(u`f+9jJX&IYMBM-{_l7LE2d0 z_jyO4g$~kJa5p~Dc(b!`Pg?o0Hc}F*YMpf9IEoVj`oJwh-3&M3ja!>#nP0Z=K}03= z*#cCYKW_OJ3ekfb1ojPS8TVa?UlnxK(|H2`gV}%RI90AWVWy$)M=kT6GT!E^`iP_y z`y>myX#6VQJp|uf{ zku9azoM?kfVWa%c3hb^+v%}bU7e?SDif|lp5(c{!={HRv*f0H^fdvaT9?J+sQiLW% z84A@}gf9q_hA=%Nz8ez+=f&A9vq|NU({IAOwRp(scLZ|!y=pV-UT|Wv;?n{)^ri}W zlQVCv9C{O&iU5Zt{yytroLqPlQy9j{g*PyTVVqnrPA)8f2K|BAfHyGyV4PeCeJjTw zDkIW;c84^E%WXr3DOalp7h)BHWRgE;Wov~4G+zTL76d!O!}$OE9LI($L&*)n&a4@n zD?Q1Vg3oW5vFqm!WP{vIe{k8CDPnV`I**_a|%6`eI#j zLduJZ&7!c;kSCGPyp9hm#dT=zamDCRZ(l84tXAejyt_!5rkW~IWiqur=jqJR`@ImO zZFK+Br~a3lhGla7h%8U9>6fT}p2lliEY`j3OI^zm<_(Bf=&pY_baixju&P7CmG^VD<5UBr zfs=LIXAee)j%cr!)C^~i1`plUI*)d7^{eji)S}J@ye}{J>c>euB-j0~K4z+-M4I#5 z!ACNEweDLqOc|x9Pl5x-7K2kd>etwBnsIc>kWH z>l^)bEx&pHK5%TQ*1m}5_8(ym6){KLyatPyx{tgDL?mD4@Vl|P#t&LMYxPudlYft? z?vW!eXoaHHMYlZidr4_q-4;9-;`*r-BXut7!qxfAo)%`evv%h#DJj%!^u4kmd3DH0 zBqsys0+ab%4P-vKfy`G3bNq0#Ren|)100X3`P}m68ghT#h)x7i7i!( zlLAuC)Z%!U%y)4zOy=8?20n9O(48l6c3TOj|r-4KO>CS7=_Y;VY_jcofT z=5uVQtT9P0(sdv;4$(Foos{p$&LIm^qa$1bCthfoz_Ob@SV?k z%n#hR6VIPJDYd7me(mJ)^?uyVQ9VR;M%eFL0O4-K9KYsEFy{`;z6q20!Y;woKM)48 zr3OaebeQ^wV^RX9qxzz}zcy1;AR^wQrQ?pOO%~F~)6viOcpta;Cn^Ai!R$Y_jn`Ae zRXmrm1KU$0eMpfLP~j~3MMByZ#MN2yI84WPeA&clw_!40*Cm+D*J2>^xxgGhr<*XD zZ}TOXwFkmLcDaGfHyx(_IRP-3{l|(h?E;r1^`!+^kMH})H)7?CI@1gvyM^yeT$23! zLI7nFSQxn8E8hlC8@%KnI!kMzZ+PId)#DEM*m*h4txD=ST2C2Y?&ix~6tx zVeW01%=hvV%yertkojC-j-UHYn9R5363p5IVIaH0K<1kcQ~%rn7|ebHMI>&4$5Q%! z30OC|Z?|7W?2KshjQe(_3nm^*eLfVhmf8Xw5!m#rK1=HcpE13wO)eBd~ev!=r!#`5$}BmHHR9YF(Kh6`p_mIE6MeHyaUiP z9mrD!$@350hP9v6}i%r>RlMp>eeeyniBJ=h>QQ1CGExb~} z_~t%p-9FI`cpZZqbAG*7;QgsHwe+8n=(~7gUv7YF@|>B0fp)jXVS>+<`-V(v*yqha ze@7#0OuomV7i8BU%FGHrHE~D%jUp|To^@1O(v)kaDlI9>vzAIrvT}`?(h{$Xz$hyb zN=vOmE8j25=bqKHj|3z6{w9vdHe!G;e>K7mdv6L^D3=?^d|P2MpSdke=4-{m96u>c z=JVJ@p5ld9m_p<|lqx0C@=Y5DULZo#mGZa%R^`icp|Ob z3?5DcLODET|8W76iiL@FbJe;P4 zC(<69!NX}l*bGmkJtjiKX>0{Nkyd2}52pd49G*z4BErLIN_Zme7c+P`4G5dziL_sc z&~VzASu|iot;|QEqQ!3_E9xxdr?1NwU7S9LVX<{D@%^^P)r~r~sJ!`cGG zFUCOT%Y-?8s7)}Lk6~c#fiRHGHjw#_z|_BgqKK-DKgasuQdG_7e5(13yNI2h**9^| z5CVo>z*2O18Z_PcTDI{@&hq<3lxGD$xWxoxtEx;`6v`+vE|mQ_K^}pSI5UDVnr~>a zL0MSm)tQ!gW&0*T&#eV5U#yF@x!@6IAiy25bgCNzgV~!ZZ$PtR2D%&w|I01K)W^byz1&w*Fi|xNq48m`=i278{rMX9oQfeL_lue$w-*4yk za~I6Jp5_#I1{S53k3c^!za0_W^ea9KrI(miC4EUq``M&-T-rC%m!@gQOnUID6zL0` zY6m;UK<3MY$$Zez0z~G6e#Z`R{MZIE-w{hv-540m-c(ruO|luNe<1v}7fG!~?Xx|o z=C|PDy|eMiugNn~_)E}2`v7Z!9Gba+sf~Z`ZcvzHN+)92gPny{RPbXb#iS#vmYYBy zxI0byDN+$~67(cm(DLOvrOk!-*r2Oe=M9-K$Ioh0B{VHC&%l)fVIVuqK%qMVQ~#_0 z7|i~Q7KvRMwbnOtyv<>9YEp{B7uGpB_^?ISC;J{Dw+Ih5O~9{5ez}zPIsTcJUobTw z4pqL1|CAtKkBD{RN27&H677zc1>9icKVd{MC_*b@1|cOIlUOMftS=*f%a(u0h?+qW zP9SEOrevcNmk8}BWdVEHcwfd^8l{*-JcPquL?u=U>>f%3HnZ__7*Q08kU-49q+}x# z7YhU@rR3#oxjjR%MD}dpX^`yMpmLM!+0fG+vS-81DWYE0tz)8I^=33hh#+Q|q|i(f zv3xs~BtXr=kI>g@tctaSLpqF$UYNqOE0zR2V&R|DqXw*mYC?u4nWhtBc!J#$@@1C% zF+J*qm2i-dF_KKv3RAdt9+H5sS@;_I+BU0VCE<_?s`414b_&qnFR5WzlzY zi!cq1GrP732Hstejvi*Wl1W!eYtC{RBs5sFHw@eB4WD_geq5wNkvcyY;ND zJ~s~CuwZ7_n>X`)T{{s!asPR=XP4|@Z!3WtUeS9|bzmy@&8PF5`jQuT?XIZq9d3DPD`}*!+7(PIg?thqpTLIPI7~yy%&t(u zz}so&3vLkD_odvp*oG4OuF$Lj7XSvc|Cl$i%WmT1*2lJ}VfmG6`NE66u6)ugPJBza z4b@f+eQ*4BlYol>duWJcd5k>2VP)=>(32dVe|lGQePN_6v#(`uC&?3Z-38+}<-iN> z0gT_~{rqFs;N9xMV3|a7l9@1A2?^%i4c053QpM_#{g58J#-CC8Yu2;W( zgr2XzPg>d+xuLC5##86zb4M##yvGtNdp)o4njncAmuYHKS^tCd8(_wJW7y@94}JgCSGWAH zJfZ4G>y$@$t~||fowAqp{4P_%)U{O)cZvo(we3$>dL32AxlHHww-@PMDcmKSRjr@# zh6j1O+Uqn0JdRdhGEYA^tnc42S|Zc76lp6~Fk3_Qiau7~H~J5q6M9ay>GD2p>!+Hl zy1KP3KR5HPv_*bt#_Vko1%9P{>BuaKG`uk~1ixWf48JjWh`T|e_$au-YjweT)%;9d zT6&)pr&Hi`#|wlHP=6iz4F0k%)+xLrcdVePqI%67UVGFW&oWQauV8=9~Q$KFJ#5e7pkA;#*UsbpVORQm7C+!)FxiGjBijqG1B8^(lZF|9JV=Db=e!F|NR>YcCU+Z*uTIp zV)Kk>vy7kYB<_KeQu{gr)=h0F@eA8|ZYwEy5c+v}9qdWhCTq$8L8BFA3!=%AvP#f6 zfifHM#6n>vctldXl!6>WfJ`SfDF01|bplBM3})Yi zl_Uk<+X6kuE`W$QfJZORo`WodsOs^=AWs)qClEaLA3BKZ1~O%kzuug3-(&-k@)Ljk zI7)^|2ww3ee;rn_#>gfNy80`X^6zv&s!`ckQWfxS+ z3#zXTD~N-!G3VF&qs$6WCe4`4dMI z8fl~xPB795|IM64Bb{*aE2iMR8cw2-PWU!*694^lLOMQemvlBR?NbvvHtl2SY)snw zCN5~zd(v4b)l4=>{sTCH|G8Ab-)BBjs*_S2zj!Cd&NOyel!tGoh0Q+llBBD4b4~4B z@LdZ91&#r6=<*bP6A9ml_|l#BIrd?xiP(}Hhg2=0Ma`uMt%(`L6dE=Wn`jqb7I2x3 z|Cq6M7NyvdcxW8f1(TQ}wDTwn_?nG>pAj{UBD5f8;8SSmM2t`{Dkbk^%im)}O`!+SjWcCWUQS?DK;e@LStQ!sOW(Mt7TgT4{V@D4;|Pl+cJEBOWmqc zeVWRcFoQhIOe@;|MXJFgH0bUk2cLMTHCPRqign|fk~^!z+kpxY4VY# z@ynx|bBjRu-&lk|RZe!}#$AtIsaCp5Tsz}%+NxZt?O=Z)cl*OS|Iy8F%ULqFs*}HB zZBZq2MO%lHx3abjC2tUI8C0%jeLkS{l5B2hM0Rb-82H(H*U9Onvkx`}EsrMT7OiV& z^yu1hXyDz|V_-1bvGl-^rb)}ATXTy*_{|nkYESCbwgkOdhYgHbr5O&72J?av63*HF!(e5$zaz>aZFusw7Pt% zsI<9tc(-y#y6k1;BTbgedt1iBK=!(Z5|6IX zLjx|Wcm0PIvC3SCeHptYU9damX}H4n!7xRNo8;xGbUi>y=R^TF>3_c)U4`^134?6} zoXBgq{D)xw(wJwD+|sE#mDl=lr>B?t(@Omv;T8I+`+MwIN^&aop9ZgEaAVG|_X@Z_ zq{xJn?(JbByV01(kq^DBhOMuVIufq-x(TS$HC$DxX5@S?OLt2o9%&iXYr4+!R6Qga zM_U=KQ?wT`H@xWK)(mnrX#KE$zBcoL{&}i{+Kt}#vR9(H8w}C(@|iDXghM=uLR6=A z(rnXp)bM(zXa{HO6(`m0y=sZhRV(b)b-T;B!$N*S-OQdV><(V^E zg);77on*9G$kLDS5|ni^O>4SDe%pW7U^KTSwOe!FNpb}Ga(CY6yPLGR&K!N6|5Z_M zuZc*P0u}Yw#xt)LigpI8mD}{!#U;FrvZu@h-ZP1=sJ~GkbD7E`NZymwKmvih{$IS* zOSK*QC9pssvlSxv#dG=~fxtWdFVdem4MjiOJz7-9jaGKc^cAA~!OCbc%uzfJ2?XYA zpFYrcWk}weI{Hq&o;UoICt0Mq`astcCsA$DD;nzJA26?KGOuEP3I^X|L=Kjk(kf`8!!5E+p8tHsSIv-@WWK`!f zs`LGOy^W1@J|ms)-)<^wROfpuCH{@Ohfz zO9{h&Zc1+cu2wewI5IrxryrY!(NoFu9fCLA9y-u8;?{Y^;9tyujf(Scz>dWGp=Qfk z<*%>@Y-^lf3)>nK3Spz-69!vj00y)FSdpvtA5spz*fVk2wH5Zs<8K^WR<`KJ>Ek^P zJ>QECNc!D9q;>KwP+P}Q?D&y<3&d6=1z{r7YhFe-Y0Gb=7q&zev~*sA*Hj?J6}D7E zFE7C>h=VcCm}9s<)!Th5C)8vISR{MC6Sy)bZJgzrWrex0XU$R$W(jBWVh&~zXN?~R}tvbWs zHt8S$9)pH8;<%rE}4>vXeQr|cwih*_{Fa76^{(JM7r1&0`{f{(-IzZDx{ zt6&fYviSxF;gtrfU;qZQfg*B94Ac7HdBoYt@>;~%sRz3eXQ|ncET&z6Gcv~k{34*% zAG^d`Q{)rQX7&r>lMKNRu2L&(l~@{TMQKFFuaowUySWA-d1qM`><)Z40(J+!X5i?B z8cc(a!|uR$H^Mf+APi(D8Qg(a8N7r67|i}-MZ|W2Gkq^R1itSpb`1Q~_p(#qKHn?v zO#IIGC+CSZzOUpGucyc-SDoj6=Uep7+#z-jWL3_YRh-Gb72X6WmpXC$I90nE@=fbE$a%M78ka5!) zOOWyQavXBIt-M-LW^e-zG1%f}!6w1bXB)^QSYWXIU1<>H17I-wIP4t^ed2*kf(J!$2~+DJ*L zs&&$Z<0wuD6-_Ea-3&M3ja!=qy8%ClfK7rSFB-@sIMiSgd>pp@eXtQW2?k*x8yau# zyAZ$1U=j?#VD?|MNR?|&m}%(yQOkU%jJNr!J|Zc_KFPu^8o$bS55aerKkakVds77k zcG<(I@*Un!&G07?JLvpq%(YV{Vhg;DSqdBL=a^`>r!2dSjo-isbfE|*5hoF_TQPpQ zLcxl%>>M_JH6w63ML3>lW17N7`#B2j@MUZPTkgeJWJ@VFC)(gr*eJiV0=w(d>@YUo zg%LQ3A{gz$#Ar3vA`RHr6{KaQ(T2yaToak(w-OS#;RJPS4+ z%LqhLgeF8AlN7n$x1DGAOu|;O z2GgVUZQ}|)lVtzE!v8`Kd~PKiAlQr~%Qe0cT)P<(_92$MioU4Ds#rm=Q6NAG?)Y*8^EsK(oOe&azQ7+j!gj$M4ZgsO4F&%P9MlW%X+o+dVaptYr=kFi(?5I~x&qs}3IWNwX zc;!6*=%iQ9v+&(sS34%G_qy8Vd;(VSTj2<+`2EeG;umNT#5)G7_^p@*%k_aUko}24 z#cz>8^ACW*>{nf++?ji(HEsp>Ondcb+%uhV`?+VjQqOb6-R`wq@yk45`@n`?|2{YD zYO|Wo{Zek#k+AED)i()WKD64Du;X5}Q$lQgb*JVVgK*y(gPz|BSkG_0Bdq7Q$)M1; z*r4YJz+m<a4ft!BCXW1}vk3MDOzMCYbEze(5E?_O}B~&L2{g_T2sx0*4-lLv?Rs5DZ!YY0_ z1{J?RgO=SfSjBJYG+3?=gn{fZgNolGgXSLqgV}Fok+RIusvWY>-pQj0oU-W8RZlG? zKeVQ1YBa}WA^rWkxv%0ZXPK(xx`e@Smg-&}b0E2nGw8rlU5jKYly$q6t+x6fStnra zzFYo;3d&<5{ zFWVg=3pb&zRVl7j&-L=`gTAdG-NRpR5^!wbZk+8+X#OR*m(yI8*Sov<=iuI5&EG`! z?$jK!>W$Z>wea_0M&c z2`N2^JE^^+E9I%Cl3pI2sefKV*LRRu3CgE>P4ju4x-}X$`DJN3YNPeLekd7z=fP;V zzk;Xjz01>$h!S`!bl3NDM!V^H#eKR&RVx{yCST`mQ27rV^!$2i^v}FTYKt`5=#D&I zonD*X1OfDPKgHYOH=}0NN1gwdV#u~s_#kPQdc7SBDq)BkCN5Ji6BKkGwt*J)RE2v{VDBr z{|w279@e8?!+%Axrvof~wYpsv(vvSq7M&8cJ8gI=;(R93_(PI@JPjo2r>~0Etusja z_1+OZg#Y7Z%Nlu{q2VB=eo&(6`k19@sf9MzIPs>x>M<)p%M$lR=7lrcyY;fp9^M9f z?vn?)=fO;MqfQ5T^p@y!wY(a4ZD(SACwJvN-h6|oUvsy%?Q4Um-vKnUr$zL|FC(KX zxINyAiU*@NQX!eZ!=f*DX=APQ7pPsGW+>(t^IlY&euZRqG>DXEJQZd<)nwgqcTID& z3MEsdN}AJmR#0^*%x7l2XOngD?wYdb_z?#kUZh;AJ;*H2)8sHa(iQW0_ediSoCfO9 z8J%1-xKvZX)Z}S#ykX}e2}kdYkUH_lBer2@*m>vi$*51w@jiU2hM9IopSoM6XMCln zs!SANX(rSjGV}+6s-P*#(vj2-hQcllWknKxhA#UuQEuZCqofrtA&0o*FIu3j%!ReM zvU!Q}J|A*CImrUujP;?631i4;Dmo17qwqNxPtLVKcVT^O#)L8CX;k!PtWTd$Ks@=V z1-c#UV?8E}Ay1^DBd|V7pZa+6j~3{cSf9yb!Wi;gDtZgnr{9MWPfoKyw_?jEW5O6R zk%|t-mMMG+;>i~*(B0TFn=xSwc@`D@8Mds?2U6~fUqbeD$Dg)9lSslyblLVqxwVf% zPK)P|J>BrOw&-G~IUa}3?Zn;Am$49)i1K=WkLrl5{)=X1ybtC5-yyyp|Hry|vKHjvLl&_q$Sg8ND8-;tsC>0aK5nI_ zv{3-|9l-dH_8@BOjNOp`waa7dhP?ikim@B=%DXnd+fp%hLtc4P259(zeW1U$hm742 zV>hH=Oz-3kv{a1UkhidW-S|Bgq%?hw2Xiz z5YQW{+AiILaL1&PrEbzaChmwd6MipGH{{;SZSJUEr*vDZ+-_JafH-lTvK4x1SV0_& zjXA&GpVs*~Ev69x=l1xXo;dTHfF;TMX6CNTk&j!mc2BME#}@b$Z%ee0nmx#8;{xZ? ziV;C(@=e|oQN>opKTZ=Pr-{FF5F;JLKiZ3t_TumC#7H~w1)A7{CjJ{$Ox4_~lWA*S z(q;PhO%$fti}+cH_wn`b(5~@MpnS&ypsUS<;pnomL^;JL<_XPX33-YK-kgfA zCJHNXh$a&bWvVo_5y_z}l;WDu97-tk3d^A+NO4c_9K|VV>Z5TSe5tqGhrEQGG#MR+ z_5nhv51B(wvPPHTd|G@CE+OYmMsG%gVY$yi4msBveFNvy>=Uqrd~`B80u6>`J^>u^ zQEPN1&ZpI=R;Fs$&Y=a7H2M&H7f$$c10$Z3<&&~boJ>cilW)2z`IxUv?X zf+ge&lhM$#I1rZm6mZBFtkKY;c!l9<^@-sj))Rh+tuuokViObPN*~2dnq1@{U3jQ$ zZ^){RZ2KqUt&d*MScIJ>oN7_GXXWB+4u4B5Y=N5Ht(V{xj8TOxP+J?aT!+>Qs`2^t zPdsz~qvuzIF0{R-wx+!(!#Ur%6LRUjS`fkwzuZ@|ZRAvY`HI^6MaKyFvpYpDAKKCn zi|8NPGFn884{co!i(Egno%1)*obk49ABx<@+s+-|fry|TbV5czwp>+ zJ!S`fvb#F3EiCWb-EpY(+e-u$v@*9ZP@ADhX6z%*5^Xr5B22wpnP)|N*DX6yqzX+)H z$1d^K6mb>LW$eKA)JPvv*4%w0mOFXPNVe+=mHk zzEDgI3UP6ENwl=Oy|0#C*?y^I&f3_r3&rt4A=8~*PLF}X?ANdMZY7x+umjU~$@|eH zMK>~nMWYjs9zc}NA@9IGZ}5I-BY#`Y;<~|tVkTHmrf5wzkSH(t>&+=oO+w5RxA^Ob zic;^^3&mzZPh6b0CR)1R-q*;kd}$aptGH0SGw8{5=dGv5z+m>xeJ3L;w_IxJSQiVz zZ??!=GREd*O-U3iu%d(`yeuhxf&~*O4hT;R1xhfVq-d4OFT@jrJYAe^6D>)%_d!Dn z&}SP+IbQl!PM`rB1owS&A2dV(efDZzrvzaj8$8axeTGvu?u^UC^ii^r@o7F#qKO4|>sW!t-CyqyY2*1B8L>a>M-Bba+~f z4!~gcU$jW;@%n-!7st8N?95HgmOh)~>t|8E=b(Gibh~R9yHY3e&*-C?->R(t*kBzl&#}QIweoqK^sI z%*0Q@%juJl2UrCvdd#Smy_OK7!ztv1N!AFB3=#ax)yKNHb zPptZfbac1X%t36Y+f|}~fgcu9H}vnF4H

uG+PNR(keu)A(i4cXEr+ z4eMui%@zzyTCt1ZRO)!RseJhu7|8x1wH%Y$>^;EOC_F>1&-Q zxzk^=SaN5eVhvL^crt=18@dZ@AK1|A-zR-xA7O&zJq4lPYhe%JuK2wdghH5YJj&XX9)vfI>+wH}@6S;LnwT24p4&3xBJ_|kJxs|#0 z`PVFuzPnD`W@-NE;pA4Y-H)qDqb=*G%;%(4yOz>Q&m3;DSk_yg8;5RKFtcl_V8DFE zE}T=T!{Mf~5a{zwpi^yve!b^HD4j<5ZiRJF^YI^p`9x5wV zaJjY77uD*8%vW)i2CAmKQ#|qmo%8%Ib)-B}JaUN6X(^;h;tf8{gqU^tVpR`Mr?@ZCDBWRqUOKnGy@*%(@TvA_ zw|3A;CD9FcF?rfLvBZ(rwpgqk%;c$_XR#PyozIJF=Y2udwy)Re z-23lJI`z8hfe9B!OQOZvRx6fHnGSK{@?_O=MJVa^~3vRx}KBKx|by^y+q#w5e`RRLM3|6OG<~*i4Bl}b1$!hw|%hpqP9?^ zd8!}U&(wP8`l?4Btl-|huWKX?=WXb%vlBud*0p74){gSD1f$V@K`W7o@Ox&q` z{v{pJ`quq){Uh0@r$)OTjOrp)m3s9pDoz4_*?$^YF=D3OCzU8?&1x7E9-LC z&pNxhy2B~7XIQW4RNasqm+3p}lvLL9yF9K$Uq+RQbl)?3=y^O{H&1=HU95s^iFI2T znmuQlS(2e(mOe9IbBOsu!qw_Wqja^s`s#$2`P{)H63Fs+z)jo@86W3y?()2x6MCGM z>of0O?_^FfQ2)X!^k3YlJ7L3pQpxP{7#hyhJo&mgHGlX>U6KB-BeR=7+UlZP$84xO zX~li^>BIC1S(CUo&oWz=?vzi~Dcm)#(W(lW0w-=J@hYhODNKzSuYR(w-d&^Is!Eh8 z29|0pyjs&o?08B3YCqmhYV&s8IqJYtRW9?zQf(QtJ6#pVQ~8gQbUndad1_rSlJ9Tg zh-@PU`0`gH?6CKykcB2p6_2u1Pa0O}!=~t&3WU6Zg7A7`2mOKalKaRLsTi-DVRAnT z&czn}0562dbn}Gr7N3{{wBoJg5C^==6m%xU7Q~jdC(;x&yaT@6Ovn&|p`3>I!j}_; z@*M&MMOepej)xF?*-Njv*?(lA z+_Q&w&U@!~-rNh`J1_Ei1>QR^d2@gC-uW}1TkIWw#XI4ucf8o!!zhvPCuAM|p+tgJ zWF_K9 z4oYe*+*lD5!i|+;Lt}*jw@OlPz>O6U2C@-_aKGiap|JwMU^Y;s%nJ3s8!Z_1zPmIN z_5K3d1JwIWZ!5H~Uv?1McPS)c@!OSwT7T>k;FaHO5mkfqp>?q{DmFKO$FsG~_SORonfZSuQrm6TivcTg@Hnh_pwVt1tOQ{o`;`~beCPeh@_H>YV!bqv0_#T zw@OYL8Y=({W;5V^OVN#eP-6vzf$TR>G-^5L1N>Rwgg5+SitJ|n6z4;SKN_o{`I}_7c$s*p04jN-p z)aIXdgok=8sXV^@qUJnRJTn892TYv~%L9I3kOy3A5d1p?%L9UNltDG{*X-AC==(}` zVsr&Ix&oK_JC^#>-f9=#zm_YorbN?mp6aE0390qCKbLWje`S{g9f@^*k(5*vquqXuaw_#CS^=4do|am#|cIotgq)3U9Zx%~Ah@WW2I~g~)2%ce z?!5kXX2NKyNGF9v0lQ?{o1_QoN+;eC8BcGee>ik?ba}9<8JbK<;1;7eBUU3LwW3ia zSF=f?xY&2MT_@$1FbiZoRG!CxOw*;&NkV$^B;{6}`!(0jYw4_)M|k(5C7N#iL*`(l zW{5{f_t!lb%6}I9f^|psP*mSLlvk(Q+J91_FJZE}Yh~J@l6>7~<2sx;L*3E3Mkk5K z)X{fD`nu5`3Co`QC{m&u3Ff_|GDoasBh~5p6Lf8Ka(!ph^bSdejfk%2>2^u3KnzS5 z^}5+*^b(cR{f~J&Z`5tUbKy7QweTDFyP-ECweK_YtyP!Lhe=i!SgIT|b!n`20ZG$L z(zqH{F!mv7a!4Abxu%TdnXsQ(zEF40UG4Csos@8rS)KsB7;$*gMcRFysVdd&aUZRa zy=1EH=1<|NHUFy^pp}lgq)XpCNl#s()(hNCs7nZn=64a)C6g5hcO;eLtO$D-MfF`U zMgS1d1#G4jbI2iXco!;KfP-ZKx5A2mK#0IWGJxbwup;1I94rI46;=cULQ@!is=Eh`~WJfB~CeMZg0%SO#z_tOy8%W;jR&uznM)2nhK*K{9|_VMRb7 zMByMAK*lCm5pX{amI2%fD*^%`0SCzd7HonQ0YL#kP#4f#SdA-lOO!YJ{6lXfjUT4w zBTQd7re-@vlZdG~g6TVhsVT$^RA9>QVxG5R%KI_C*4g(E+w8Ifk(_t3?;*FjW(T6a zm`}Th`qYQ^|JeKTuqMy!Z9)PBf+VQOrX+|06(y)tuwufZh)W?>3%CYK6>O`38xcbi z!lo<*6$KRuj99Q*ajSqp5+fK*E}k?R%5?ew~Nwni;1ZYqjGa9uD_= zIL{)G9G>&s_r{M@fd8Hb@!wG?yge$warhcHTo_`_TJAVv%?4bU#+oh8;~K63SqC*; zm5nmLiXr`r(}hW=7k{Cc_vX|@3tscJv)L5u8!(z<);g^hLdGI&mKLHs+If~fashIWTh%iv+5yto3W88oa*oQ8r5 zk9@r<$9eYlsh^~J+*+Yy`L|DxagT2Fn!}j=#yNcDq?KqCvVn_)Il&W=A=9`>HIa0< zy4^`u-tmzR@V}=*{6pygJG#9%d!=2&`3N=LnIMOk0R9`cfo2I^I0nQM<=2Pc1xa;z)83@;vP=WMDMnR|<{C_p(s| ztT|UZ(j2+f*-r#ZM+?Haci<(ncQ z$S4TFVD!7^4D7t&8nCg)QzyyAoM2^N6+MtzZ#1=hu`6vwkS9_M`$!3>B13%3^o9!w zv+XPq0}O6Dn%4;n_>8BgmH14}$%JPGf9c0yq3i(;Kakq*y=43YCkV{D0w6HtEZ%r3 zcIXLdlj+c5w;iF@!Fnejo0l$|ek$(O7|9pQe>W6}o#!GS85KrEKAOSX7Wrs;;qu5w zQ@OKX?U;0#TDzVuz?=sX>issJy$>mgfa3fJz6eW_s-cn7U2xOXD>3y)oJojz7MPi^ zx3Q2;jmY)r!Y;=8DZEO?dQ{;B#(HgTj2c-heTN$PBPHPf@gV{XhaUtbU=sULZ}Cs; zM}3_S*pK>)``M2ME2a?Jg=ZHK+yAnZfXiPS?v%aCBd3ovM$!Bq=6^4a`GmOaMRlN* ze+0`CSOFYs&u8FfY@9VZ9@zjQyGNc6`#^>DfjJqZ@%vNB^!^k2K?2`lr*t5t zJ$HPPu+ccAR$RTk?#A_i<7Le1hl7ocQ>)&W#2`}u2BDWUmQ=Q;M% zp3aDm#DFSF#J5chxNQ`*Qpqzqm$*%c(%@_!M0F9i^mkUWefv5u$Uz!EfD!LQCH?yk zmcwpwxFs0~Z+|-2iJ}Em3|%_5e(=q`}HJi~%sWdL%--@76Bh)%-+$*bMu!oEg_ zh>pzswxO71gu@lpgX8>uGfwkIR>1#Hi3rg8+5VTTfPrMkiC6aw^#~CIF#-9R;T`oi z0?mj+{G;7yr_aM#&f%;zN(<3({=;7zj`a*@c_{59{A0}D_X^yfhZDL@-Q@(gv5uQ{ zS0ii(7VpQ_N+|7t>cn*-nJ7&Bz=uecHWO*1U1-H{b-sK+LKlsl#nZ>p9N9o_LTjs6 zQJWX1j3cAxD*A`XQpy~9$4a(D(k>TXYNv~_#C|gB7@Z?v)j<-!p>lo=y?>^x+D9(L z5BG6~71HVvn=2AL)EFpB9qH z`otKTD2rT1>}}*oPGheOCwR#%A%;G@{96B^@e@}B6?EE5i8Mq@G=d))O(TwX%0v%n z9oxEc@5>@(2kCyCR-!zlha+wj%jiWdUi29JLqqXV9|@WL+)9zxP7^mNuF+rBiTV{Y z%vjR=eq|#c=Wm>A#O(HdFP5xz%z>^5qqCnt4UPR_WLc9^-}iM>yR^WI7W&!P27-b$ z-b&}JkY2B&$wu?#&xtgM%5O9otN4^W(XNwSSNi#Dy3MBqvRok@r5%9Ok+%)q*vup! z6ng#mvwY|{s~ICzJ`heQAN<$^J*b;FN!z|yu+OK0cFxa6u+~;)%&CE6MHN_4fSJ;Q zew3NW4J+!LC5pmQCT7s14@=hv53MS^kDXXTkCs56jI1hb#EyB>1&3w&!DDfpX1XAk zyi`6U{--Q{cu3i|K~b?lPD~t9z=PEYi8FW=yVVJaOL*`-2!hC)SMi-DA^PAX=01SX zuiPa>j4RxU*-AjFqm}%9ElMgs(A87&_th$?{HCIn{C%xTDnHQGW4?&OJy25l>7bd5 zm{~12Yl}&)9+MxHRg1GGAWhKBCCsc=T#&^iSC9EJD(e9*h=9~YGZ!-bT5v`dlUzM! zN0eVJ&WM1Vfo6Wn^lQcWT1;~Fm|anR4{*K&rE+f_lG&|UwwMVu?g~+r5`i1G7 zhXncUK>ttT9H4@80R2B{KL3^j;cc#TY1Xn>@90LR^1r|1K6gF^yjf%m+SC-t+P zEH=hQOH)b4(`xp)?jU)1>BOwh{>)~2><2~TC%Rm3WN`b0YXIp;op#c~>}7N9W1|I$ zDOmk8(_CrJz5wqYNXPdlu=KroCYAX*soNq;=96w|NN+noGrWs}?+9?t*Nr|GdD|Xy zOY_Ycz$F6FZ<}#2F&UYt9k~o=_Yy9maMSe)D_{Xj)RHwt_ZUl!c-f56)FFoQbP%Zf#?r4 z$$gvIG5Ytu_lR6kGUp5HW8WQ}zHY_0vm>Wn%=A!O39JuXHn1BMJ)?!Z9~k=T6ip+2 zo+U$HO(5=0<{X-csYt%T54_6nzOJ-!fd4-8=7IAB{Xc8Ye<+V%MIDRQCA!~D)>b3H zQjsZ(owZ?~X+#^5_BkV&(4DP?OkSZ&NFvUWf|+T7{N5B%jbuXiw>qMD=`JB69M?(> zTcyjZI0tF_?qi^-Qw9FD>V%eE+_;5k5yMIyD_zH{$biIs_eFA}aEUmi1F{^+1YAZG zPJ}})Ml#XNJ5jhKoQDIl0SUU%Q8+Q~e~ZWu4^REf*%pzyhJn{e-Rx|unfeXGRx3K( zc>yvy9x9L^rq}8nN~;C>J1Zeh*Xpm8b`t(S7UO7PH$5@|sq4m=rtXF*WN-y=J1lRM zaa4=4jLW&S_0Xmw>8-t82fCz^nv;g7?g!8J*}ln>b7L!ISJA|iy&XSyNgG$zbgmp4 zUReWomrir9v2Y(+>|W#MKC;2RJiz_cKKJq{cUNMTJ!0>wteXg`SC+lTUZ1R+8vDsO zd(CeGaW|oNd~8h{wyny7zMT!*$%bvIvLFunA9zevL0r{Ped{cfOK#MCLG$%mOS7*n zi7tqhHq^Hm57|__qbpi^G`p;InlBzP7Mb$m;5GZ5&Cu4vuyAMGI^-|T zG;qjt1FgFFimxlnZty`Eh-PdG6nin`11wQkGF$ZHF^hMHAu z?lwcCW_8LY?x0TXmoUyJ9&f%W2HrvgqCpNm00yIf+nqI#|< zqoqgcGcZH&@^^Ga%NDr)1rC{RU|1J_|Le-i8~@40H`q}mY$JX#kZ8PUxatjZ$o{xUqpM8}8ee?dhL}0wPJ&mem z@be2pBbHn?x36Jak3tfkuUKZhYXn_)C{S{fXwxMR^73!^(!Uj=i4)@b6DTtqDvNw=88X?5VJdsLHE5H)>o# zpggZmmed}?5(Vdh*#YCd^=SaPfA?B0bVrmIdYv27t{{Xzu}YlyWi`vLC!m`uO83eb z>TqNKZFR7APQk{8E$Y1selLMt+q+6ay%St3YBlF9?5eJoR%Nu$8ZkXVqx@1=H}Sxq zNbJvjU)Hkyc9-F2w4XZ+Cn9dQ8s3ih`I%vO#I48G%Ob8ls+LH9SRIVYDfqIXS-p4n znUP!0;aldGE=XzHxjMKrr(k76v*RQfj6OB8#4~)$oYJHdZ6N&TCaD@fT*a(w6&TwP zsMUe<#Vre1*Zu@XzuT`2<7q#?G+aZw%{SDc{oH0)6LIT#bwb3oC)I0Z*(xHwjVi{y zGfIBGRdbd6eEzEkc{yP#RQ!AZ3`U<)()O*IQyO)m4TOQ{_cIB#DEPba)=~Dg_G&|E zRc!kswGEH=4)u(EkM zt$=)HzrsWDyP+9DQ89qkt2kYDhv9jMhF=rgkvfT;-LRA=I{HHF=~9J zE+!)Ob;T~0zhrAmU=JaET*$6N(9#7G;o#8Oe40p9K$pI{kCk^PlVz>r3PHaEB&3U_ zNm{)Oi;G!%OIfrix>o`fMF1vt-4srUCr2);u16JEQ;uO-_adC&9pqJd! zho=uzON3`-<;68#X%e`kU&+o#mv+?ChjWJ|ioHcFxehr4KQt|Wj6?m|OTL~X8_HJ% zD1KoHn-vS32x4-xvXPIzk{&>J5(qX2e z?>Ou6RO$5u`WHWqra`j2u{`C2geK*KFH7angD3eT!)GE{#U^8Yh%J%iy}fm!FU>?; ztelG-R$0&=;bJGa*g-B<%55wVS+sN}ho2ypZi=$>HaT<>%yA3nrr4`;sX) zh|*Sh6j4?qOsnHeRimm~>*E3$Q`PtDAQxu6WMB}K4uyZJ9J@ydeNbR2aDo3#HSC^g z&<7Yc#6nP3ppMj$?)m<08Wp>omqT_cNE3R7b zjgMVCLRoQ#xG~G=OB~TV%gKjJ^v!bGgd=Xwa@yid+=_GBmPOc(b0XtXKhhffgO&k) zgAe^BCE&W!M!lJ?;os646oN+%p#NKsAE+M>LV&{_VD~;Agz$gU8OFzhkdFr;J1Tcv zopf07+s-gP9)x^62x-xrbnx==AmoodlKFTLa+&hEHzXyr9XMnG&tBd(?;e;N7!o|! z*{=UdqQ}7%I&=7QkGUTGdToy1?9J1(GdxCDs9To<%1{6N%d7TrkH8!>5;5Gx+#_lZ zutYUR9g!}5(1@Q&%P=5D84%M{7WC5%h?xe&6DkYhp#L%E@B0HaJY`H*?a0h$ zNKu_WFY_*M#TCm%cKUmc>aKaQW;T3k+4n~RSNt%Wum1=-Ig2vpgQ zd|1tLXExH>0Vbc9)uQp@K4Y(GRzAEu!nqWKbbyYUu~7!tG?fMYbOUUr0rrH-f;i~^ z55#mPa!Cl=*><`xPo%C3PhGtN&80G$;Fe~*el^9Fyp|60L{h(&Rz{FIv@gvrvsf3+5w1*y| ztYda)Yj~uSzO@C6IL+EE(#gO&0Jhu^85Q}NG4eb_?vq~@pMPIz@t|KDkFSe|`pF?X z3I7Yl2>Y|IIqeUNt9pfLpJ|fd78V?Ic}mAm<{W$buvlr0U5NKgO&h&5gSfJIcwnXf zRep$y9o|pX_I5(qBnQG^^iE~(oSmvk4ut=Ck|^j#+t<$DYo>n3Kx?J$cm7^0mCDdV zMn^j*A)^zaj@n;G26O$UA>PV7-%1kUg~)4*Lwq!=s?^9vMiTX~ZGPXOcspPm8#^Bj z=?}P2?+i$1H_afoEPh&G<*uvz{hLBfE<5^rXE@n-o-nYgn*@W=J1g(q;M;Es9lGoY z!oNF-IGS`6*_BAr))uFdj%s&hkhFDrb4W*Y9$zLQQEz34ca%hUx&DYwYfe2e2&YAlhHm zym!LDp)TGOfWhc@pX6-2@w~NehrZ3bJVjiZecdtDA+F|CO4}@*`)&sgTVFR$%j|%R^#Xzg~7(>`L!T;32Jn-T|P{@z%SIhMAK$_ngR6Q0=l0M z$>1DN6?p!tlJtZD^wt7|!RVcp_o16`DhhiL{_v9+RY%9p-Tz6pxpt^WkfUpwzDYdB zJ-Wo)UfbLj*}FXQ>ReJ7qCA$-hDJ8S4mjYx(|lU0#x}r&sf*lkLH0&WW8IUwu{g~r z%_wt98L7w*x!KNlzEPop?lB!rTaBn#rnzGo>9ikmja|?jBc{IYNn|Vz5v9S5lyEVJ z{CJ<)1KPSHK78HmEaB+N#5`BT1{|;K%G>?WGR@H^vGrorkBhc!Uu^A znG)1VSxlCk=}tFK=aBA^k$k+bz_4%}bxfjhMaHb6nfr1`7s<#5d{B=ea~O3}9Frww zuB1hpaxh29ymEXH&yYEYIyoMbC1F-Ym_KEcc9D_S@xIDS@Tg;>8doNmv`>AJ`I@NF z$A0D|Ci8oIqGAN!TqTjG%qYK9#*FyXmj^u*KK6^`n9L7qFjeoJa&jcrIXud;)I6n) zv^p4>exHw*odZQc<)v6-HFF}WRHp7Eb?9P<8-oz_*n`T>H|E2^%9-1RCWMv zRQ1^HQgr~HR5k{ln5*mnoUUs117RR~ld1#on6h8+tt6S|B~DDQdmv8CLVHal=G8rz zNW6;nl!yxJY9%5rq$&8G%=68QO?xPbF7*0b5`8uK;Y4(v*XI+_S;-H?f^@HSVnIrN zkgEOfOAp_9+udIoR(f#FyK3CE}YbU}SJY?|z=W$H?iT$C2dp@!n8!`q*QC za{B1A^;p(Or3;odeAZIg@#n2-1l*-+1U#v11Ux%e+1fWl)d&c}K=dY6f8a4?BjEd! z#I$|T7)1!H2)ry3%wk=8v9Ewybq3uwEKDHmYpl+n2m+S)2d<@e3n!%4XHZ&xlEoIH z3~8fm{MPYj%>>#r!@&vBmvm5Y_oF;--Jp0Epg{b4HzbW9irGhAJ{GfqMHq=$OkO?| zGmW)eARONAG$3rG?NW98om4ggrdcW*0l!rBxNTCk_8n6;0)j9YJy+QuXrO8Y1mV9x zlF(-X^=3_DMul2U%QXi>Ba&fMzmAG&*hZ1yE+otR>n4HA{>N^lU5UsZd0bCgNq}vmQuE?SUQvt%$8mv(uB_(6vAq}yvvItd3C>zHd^${b+#4PzpWm>!JbGC%YmcB|>$j9kuB0rK$7Y1|0W07Rq zhQ%B?UqKUA6BXlu97S?Id!is%E>*~}bSb|tpD2+M*`mHc#a;Ia3#`0Vm2LTHA5nIwU9Lxz5IK^@V4`H4NR@}N#JxnR16uwVEf4)++>kYKjy9Tn z`LOTMHlx)g(3<@d3QRS~_xwBZe$?qJSPE*%mu9`zD^ub{4;IK9e4 zUa;kEot(r)69;lVQHs|`iGkjt&^meQvZw6?Vc9-L$QB?l(&+0XEu>?414`9dh3f-HbZWP zOR1??`F!F_f10RK%B>c6)C~ui^~4b$p~G{q6YGd=*)+ih89bkI{zdBeE+BP3^w@LQyYc)86CMP&l9hmXz{xh;@m;z9Q*A6+7H zn=j`k5PSJv*`f^#h-$5)yBVc6C}{uwROh z+_VxYKBr^{CHb*s@tsmoE)mkNu<%bz_Iz*v&gFR?zL#C&&Rnn4|4oiuHVsU4i z<6tYi-BBpCXm=cLb!&GV?Fe8I#~(zoh|ufcVMN2>jVc=i_`VH?LmCdRP+1TM{r`cO zwfW+h1)SS23=c|2;>k}1mryi+PTiGpVL!XDr8-Oky@LDTN})b~$WHR-mKs`>LXG~A zorHt&K#U=$&onJ!{8sSRaiN4A+7T8ktvZQi4-co&hM&}tCD4D9M#JGRRTlI&H5?9Y zIJ{D2K^*iy@R+s16VWSKx4Wl3p^lvL@*TRkjNo5cy+gtG{ZPi+|w=%>#O0vJr>Q2Nr z&#j^(_CK*&5%Eof6(ZuhhjG=^uWI7rs4JnrHFyMH%qc)LY_acM5IJ&cP54?|Y31>@ z^3}mda|%Kx!C>_K-US&Wr#!+@R;3{PM<$`%3I1BT=WcM0lu{e~Lb|6uc#bUOm%2dN z&Zl*`GGJtILhpW_Ym=qrck=b5B~ja`1W+qKtNXCGq~PHQ`%u zrJqB8neYg{lvAM5&^*6)cgDyqk8o3~NiY!o;CNfr>fnT&g3yL$0REGch-h`bf+ABw zKxN#ZC~+>UxNGqfsvt9&JaYOhf%bS6utZ?eZ+k6k+b3eya>Ba}uSiFd$c2JS;}IRV z>)OVI_t}&t=+7l?C))Hq2J=*O`}V!FB1gVo6TS*pT5-IsZ1v!goUo9FWdICD&+na; zG4j1fxRzBZ2m{gYXA;rg_NN2IeGvhiWV|S83a6Owa8N2Z&T1bz`z?)D=k+!=WTXtA zA$fHTpD{6#jn5FjI)={>y^c&~jq^j3Sz~LVzcqLa{+JV{(NJvPyDf6$%9?QK?ack- zZI!DB<8#77C&6I!{N8ODBUe1ap*J}o{1-@)aNjIytngMfcX{E!LP~PqVO<^^h0PSD zSH-O&0P6#n4eZALe$<70Z#GG^zZJD3-@BAF^4){Zxp;3UvcSGKLcKR)p2~v${N4!1 z-iX;M3*w;vG3M|4e(-J47Hhe80w-pYdrM@Bm(NJx2>ICCB53yO6QJdQ?k$lPUap@2t%iJC z1eM@lCP2kE=-v`R6}UkHUC5{JS5~d8&C05Em7%g~y{IZ?zY`KmL@ZxvVRLR$vNLsK zhlLvU`8cv<)#K&r1OB_wo`+XFi>ks;wImU|)6&PE_q^FYPyNijcB{ zw~AU#IiX!~gHq{*_L9-zCfb>i?U5tfPpT~FpBmYoF|z%b%7QrPf8a3V0ap5>Co_EBJf3EuHQDqdy}NW zpOZ3rX+0-pq|t?wGJI(vCuQh{IWbe1pij(vpCt4$hm04p&JG*9U^7RI*R#%!8vA22 z#)P4))8oQONQ4Ngvo9W3*4dyO-=84$_r{sv&y8z`vd)ZY`(w|IYOiOV8PRsZo*nL3 z$VwmT_`vk@%!H=K9Md7wPl#s2rb{`dBc`i~Ge?Ivag4`?cR@cXQ2Bk~xUwn-RRPy3 z>+HFzsvLxY=xeHCeYUbHp8>#N^lwViV8O8yUV?M%1{*uscKw$s*miw4&Jzf|2~h+> zx5MYk`t_&d%KEiYRlm;rBlP4ZFWpC>?dv%_DbatRSVs6cl33zQY%Ft4+_l+x0fD{D~Iuk)1ktL^8? z`t{OrWmyY~$pJ7Jy;fPj+Ni2s5C)<@&?NYz2}EP=ae~Xpqbx50bzb*0du)O@@x!Ebwv?@_3G?t0;}izA;P(rhc^+ z&8Y4D8u4SyC5JnXm1b3L=lbC(`G%HFsO^0kiKEPdaQ`n#c^9^4k?=vehGvbZ&Al3l zBTR>I|EZiDmS7b{$Y-m}E+TNz2Gsrv`)^V*k>+v>jk1b`Np=A+jdy9td0CW7H|K4KW&)c$| z+6H}VX!bK|^K%W|erDJ~|3$^TliRp=Z7BhUme)|*pK9p5W{w@`5Z=x@u#KDlKu7-b zCV%!48&T6v4e9gCJn~tCm*3uC(sz4jK@TT>-qcaL#HI&&Si38#`=lSjMGVVW);I){0B?dUACXGc%ZGJ9V3-pupOi!DDr3;XQ!(OKB1 zWrt>An@;bVg?*B}Y1XXAhgQ#;^@~C6u7vKJJN7>3FWp_1af{-)#WHt6X5);ST?sF5 zc7;Cw+hiDs&d#_MU`%QGJ( zm@m(4r~sx9Oz7RubAeO&@U)Vp+~MJ6PVC{v;-w|SHNyYWP^Lt*2-!_Yhrc{W?- z?#^snQoHMH_s!jVpTAtX`+UZ&U7pQmxx3Fyg2CvTU1wk3><@kZ!_wU#{O2ZVEOzua z&fU`)ePCv3^0XOyU%7kw-n+JJDfJVpJzavIX9BYWC;gV!vVeNN8$a?gB9rs5qTrBx zQE*$A_>xojiCp|p+wx3y6}om;T=&g@z0W(A_H#21QaqPg=FZA&)T`MQ_wr_E==1E! zFc6)UaWKepnOW|vGmUxx{N5+2zaQo=xpJjFMz*Y!VYo!_N@x4ivKwnDR22S7hj6*m z+gMAvmJ^+Qwt5-;c17GW`mM98meH@D*IE{F<#6?95kKN{_i} z*umL1&*!uj(2tX8!UsO|(P}(F@{%YQRdS@gVdMz$vwYgPxT&4kH$W@zWGnh*`HF~% zI$79KgV(Eh@#r9EYm#KU#dFTYEQz6dp zX~asfzpDd@dWX(Q`2^s6w z@CLTXhbTC4K+fO5NhT|LgXPB{Ro@L}f(LXNH*EOxV#SN#ZU>gEzHY34e6-UuWso(N z-!rcICkcKfy{~ZM$S>taD6~YC;+3@L74kBNv5UlRAC^o;=d(m`sH5+BWv<*t0v9(= zDF1m4EtAx-WalEJLj{VV)*Dc_U^mpF*t88#?D&>G+(nnvkTbrQrRFQ>RFn8`Y9OO+OEsjO=nBT7qH)LunkvUF3L96K)Xy6r9N zx1bMPP{gsuPXu37ygo+k%be&mli#3SAPzPt+`J%ZU)}{pwbMkPbRK%EAPz4Im?7;z z$+;SI5tF#H49n)BHw)ryMV$uHGE^pkpBz$!9ZeyYe?spGkR!XL>(LVy86$DaU!iGF zu)?cE=`KZ?59bT{UYX4l#-Q37RLWsS1ER+kcN6wA8kwZV5z&o&6zK6Gh5uFsQ2^sr z7{LfF=FkT?f#<@D!fFs|{dsOiTDEP-E02&@nw55v--5JlM_zfRwCceKe|u(#sf$q( zPT-Dlw?G!?M?XTU??IF~FzrlT1W7oddy2b-M}dCMBjnUz#4da0ep45dBwUAk3f02n zhW?y(q)sqmhdpzzsf%$^mcX6lZed-ZpZW-?vjz+7bNB23B>y4Sl~1QqDt{F8YR zAhLuIudhR9`%E)VM;5HmoV$1KG1yVRRd*K8-k>i$azsIj9J;734$!E9RDRJjtXcC@ zAYwr3ErvWKQqZeD+Nn_v$^3rH*o#iVtJi$a;OJO~!NiNQS~Pq57$H~^9#Nk3Q6%Z# z3q<`pk))7NZ|W3p>Nb@HeSdH2A#dvEDhuME55|ZOIU6Hi9CbF<7M*rB)_HN!*%;MR z&?Nim7R=%lXf;$`;UaiW*?zB6whb_>{z44xvs zkX^L+g-`4|Qofx@wac0V-)o=sHJs|0l>^`Fl=T9>-x)Ut@y!a{*B@nm;~V`d{|uR5 ztLQZ=tOKLhc-+m6UbD`+J9-Vt!z#vgW8TgfSKlKL1A}kq?aPi+yfXmF;6GuWpi1#S zpGjh1jQhqO%NTdkk?v^sbv-{syL%u5qp2%`ETgF(HGc2M-}kG0>l(kS5bmx#O@upv zQ4M$hj28=cU!t}OPF>`u38&iDJKF4tGAOto-|SbJ1o0sFh8F)JnzvwXV8|&q_%Anw zp1oXn$h!rA!RX(nzAJr5+GJ9*#cjv#yMt3sR%Vwh-RK=RGC=Y*`?1;0KLYw;`E`lF zcyD_eiC&sje!PboRlW=_jVcoFA)?Ad%?lxQ_In7II@&qfm_->R-;aj`8zI_Wi1VrN zGDX+h$2nLIxW@62+?82p((uJj~c)K zCvXgi^!~RRzgWBHjAo5QOUAg`d5ywZ+#t0O*>G2|H-{v3w4rI*&}>x}^zCeDPBt`4 zl?8Fo{}}W4{c*C|=#;IiW3c}23D?T$HF>11-tocRNi(=(sCCC@q+&b@ zPh(G<97^jW?^2^`kv210;Hj3*VX*8ec@fS*2(4g5br7O@x6)4X_aNe+N8c$eMP1Bk+HrhNVC%zhA?!drY9uAoBa!f-;B~aOXUM`3v<(sAM&9DNMp8x??C7$n*Nq z&B*c~1k8cSGj&N%!in5dC>9>)^>dn$4MB*b4orI!m)l9WB=?jM3lFw_&NJlkU_=a5 zvYNOQC1oYLlPDI}=k-&Yku^aG4F@LQ)a7(imdL%HVi9y+|7J7tc@QGWfw{!Qr8p@o z$-O?rB8aVj^BM9?Fd_<~@-rzbO!7-~x1(4Xo!3unMpgwO;10|V)3Wp=Kau-YiiPia z{maeBCqak=2j-_HWw(?3lH9L`SopH_FF!+`3PzxV5p_X`y4^}U$=`!Wh8}%~kaqvu zVz!5OS!8rsRH(9`&xJ0`UlygSte{D)%M3?Gq@Goy8;}Onhjus@O{s5!br7<)Bcn|r z{y2z42YUU0EE@D{;~}~n=mi9_lkoqsm{4~=Yjs0?qgXX7L@X`@7TKZYN8t9srlAD~ z8FZhbq3e6aySr-oYQ*CwZSeEc+I_E-T+g5d76}^mj{ekD{_TCSm*FQ^{Ba^Vhsr&o zh(0!?Zy|;hf8!E+yo|N6i@zWS&Y{+Oq-&){I%gtNk1<%YC+_aCK)uY=OR^g_j+5C*}HWLx3X#6Z#X9 zh#I(0jSFqy0RarWZem1av8~-L4Z(FryFvOTb=rm$pdbV5G&>v5OoQBe@h1W+K^TaB z>xCrMHE%`qUiZA{XzH50JJEYR^AIunzp!2r^NqK4bPO;uIH5mbo^PHB#XD$4%$_fT z5HS?bpgYlf)&xaIhq&Ke6}{8-u2ysmMCIojT6DQk*L$m@jct@c+8Bc9q(l*kjnKftKb{y~+)QjoLzP)_ z0pjfYC&e8I1JS4({H{&HGnXlcyh{KWj0P%15C)>(%Or^bt`$>nr`D%LJI=A6nSCEQ zWhK(WwGDG78edkrt>}KnZBxt=`0MRz?K(%6!KfA(dlTI<-Pkpn zZkQ;Ss4`56U*1+b!~&!FGjv;YVm&nAX-pTlGE9bF-a0$PTqFCbx@E}NH3&BhvrN5= z8{v1v$u4A;k-eU7nP%)7xLZo3O9dCR)h}yej&N8ij8*I#li zBV6mNUs5TxRzFZFcYWTQ8|UEF{eN-yJ!X~Ey=pL4>RyvvEp>l5s3~)AK(CTfp9IIs zsL!FRZtn>Vz{p+vI_WDezE(0)h+ijpbrn*$ew~->HNnqHt`mcI-T%PbNATPFJxD%5v~e4jwbJ-QeHs+4{snFSlh%VO?=aZ zit%bHC5?Ua7`3wQiPQF>uPTQ%pkIGhDQ&sOl@@w^`qi&K-WCUasUq)MBKw9ziKVz@eAJ|Lls#?IpD<7kd-AUC14)vKt@i&Mst(Q_63!(=IC z4!vU~TOw(f3oo_P#aLoL8Fh@#5wPmSR?<*8zlPpFQ&#OG7vhKeIKv8Q^@z(Su<5%lq_+F*H#Yxs2G`$dR1J zUKvjCl3Q}1_CfhIs3-8m6+s1^_EI7Z(Grc|M@G|#^ZgIQAk{@1xwIdQO9K|3ni!IDk%k!3&A zr9FYvJ}goGEqjDx;L19D)=}}(FCwZ~G2A41yu3_t;9lJ~+O%ht#EVZR?ysfoaT6#X zRQ4$!{J5PCJ?Qh9IP`#9w{42#0P`mAd7aN;}C9=b5)d;XuWXfV^ zZP;fT(MF_w&PeqVfx>Q$0$m_d*wrY|C6J8LW(K-ITLhh%p$j41q%D!?8^$8c^T-59 z0+`9W?GloNE2LnyS|G2OBASp)lvY_+6z_~nh!}@-LY5<$AyK%*NvD9!E3nqi|xJwG*-i$=n;2l{m?DM8>0;i=whbxF9Fwb0qU%R94a? z*AaOJ&0G?dCB_*!A*+zg&?vveNvN<37aVgrx^?x} z&aT%B~Wz<^dkWUCJ+bx4?Je!L!|CX-fZ;> z3@@6CLD=H#q-wQVl;wVWSexF~3Qu`Tk18w!b zrsZ4QXuE<2bj2tAd8%^?Tjjxt)k9f;(uLva00^r zA_>umWC#=& zV$&oTh%UeEc*r}$QhB&y1;D>Q3DJC`d-PBu$$09}2G#W4R4)rpC% z^2EgOp|T^c>$>vrq3m)Y#e1uz@>Bza!RT7$VaJZ^eEz0T5dI^RaORN?Y7AX>el|nU zr@n)X!|SFd!Acil4r;z?aDHmSdtcVFkTBcM5;4HwrlWbCuz=5adRmFk)SOIsR`8d8 z3>M1%p|S(8|GKh0u2az#aTUb<}hskm2TBws8aSRyd#x4o9mbCHjX3L_#P&ERc| zd^EjqdE}$1+*z=8Ou9_1T~GI+vIDU5y7It)r|JN-Rkp&RYn2^<0oRqSejp4)b5$LH zmdbuXsL|uYNg|*)KY}m9lB8;AlbH;~v>7uFw6#PXDK)2+y4G*|6yyz$O99 z26m%!sox#OS#@_L&)9DTH8l)vmX)FVEOXO+dX?kg=8^lRg>b@8lwN;?Vv z81whN0{5p*+rY~@%5~-Rp&Zh}m__q8ZZvVp&!V ztZd5FZ`2&%e>38iaxcsk8sLOq&Jr3J2`?B34KTtlrw9#@!W|mIN|?|DPiYvY6!RF*~CCYH>yc(u&|)Zdcp_U*Btkc$5uVZw3!U8b=wzsH(&Ryh}r%JNknCgMaC(s;5FYd zuzJ>wutQE+wOT!}NIP@VDn$85b>BZt-B&w%(d^W~=$e!i*3@|cu7^DHbai@HXWPx& z7#LH26$8iygDSmmTVxo_4Gj64LH@YeXlc_&s;kwWD z?ni;&ze^jxZLwmZe&j;)5G_O%RSkp4lW5pIBJua_--*SUe_4TFe~w+B&nj7Y4$Eu=wU2h77QasX%$!K* zG+ZSWMHAnVe(#%gaUF!8uok}D*G;X;LbFDU&lAF)1nNq7yRgK8#oZCo4@=`Wb6mTV zaNxO>uB2xZ%kI?%tBCTf%N%=W9UIZ}4F8~#2)fwUTEm(oRWvMFU}{PWs|X=L9bU8FhrJ0Z> z_J6-2y{4RycAsEL`8F227d1{DJ%w)(J~JiIIEKOFqSlWxzrEZ3u5tVio!+qjQs$Q+ zDSzbS#SRhl?!(VK{QuZ{6S$_1uJKXe8&<=U|_1SYSqQ-;u20R71@U16l|2c11&-WG;@q)@~G`bT+;;c6B*o1rGLt( z_%Xc5VP7&xErs&k+LP7w1q@j?`jvb+OWnc+LNAbMgZ6UOC_|F6n_I}*ACFf{?Rdfj z=aL*oRvrhqhZj{+Mo^-;q%8yJz5=C=vR^Gp@Bph>H6pO8gRZcl?PWjmEouS>8;Ei* z)kqabouXR9sGxuaV%clPZt$qOLOofPtTIsTVJP}U2N}{)#;EYiNO>U?gV?}WRSLv$ z>8fgm@-;RmnA9N!5Ajr^D<2S)eJ{vhPo7HPMG|ZRWfHRVGWab<@v8ZOigLL^SOu~; zt$70m>b#Zobt}bH3)IWMe@+e<+-Dl2saOyIYAS{mP{jsL9rsH;`oPY{Jd<6z^v9P; z0{=sN9fnMY@pP<&%e6X;E*(aZ4nwTt%Y*2CDWslNYGjz)gV&i7$|Ch7Ll=tjRtWEx zM{2KDi^G&?^^i>^nFy{JCpO4j8pYIFfvlZl_8{uKjyr}Q4<@WuFz!Bt-qZ}fQAxrx zZQ9*c(3>8uB)LSNH=dSWU6N2cVY3VJhB4xWQOuW!(ucA9X>Oq{W_LX?n;nn~@Q4?N zF;)6*d^e0ICe9QQgn)}BZWtOSZYrWw->u#4geNB36tM#V7t7pE&@kCk5s&rV+T8p+ zF{e!tTM=-v)Xk5EIXxBeP~T19R_lqmZHm~9fQ#jBwKUA_sfcR*Qoft5Cnm`h5sZL~ zC2qDfOwv?DxqfN8+f7f*byEa1447Evc9Vv=J{3`=Us@i^pXw%lWk&VHxVRuM8Y8$y zG20QPv9WxKoA{|F*pGkl5r658*&o{)?@2UV@p^@Y&cxE=N478h#-iQWSVQz%8x5?F z9^9-221XD5D;fxi9^9p+f`ESeqZQO>F)i9p&-y^rW~B9*apdvV-$!jWv~C!;Wdd?I zY9j_Y`2zys*YOX(hCtxMw;bT|@7;bF{IC)7%Nsl&HbRc-YD^;{W?QypvAkA2vd^J$wiu1vc!JT)!K-ZnK`vg0#tLrAJXv_}c7ADkWzQj`%TM0@A3ea;FzXK+ zAqAyJ;oBxb>t4Co>^U{rp8JLxzJastVKoHn+pt%9{qBd2kaaFDHZxC6hQLm5-#?(f zi%p2%(8-655D2}nK{9xqx_4cX-v7TV0Ri;FI{u@R^HXCILsrcSwP-Lg$V9%?2OL++Qt#CyTlaC? zahBSz5VqRmh#(-b_vM7j_O1gA*8!Wy%21dbxDKGM13n!q;d;3I!RE*Lq552*x=)=n z_v;f@@ki&JUY+&z_#Nx4@48Q((*E6vb%+>qKiz1)eu%4Cqt@pWkb9y&*SR|t^KjgN z?t=M=dE*lEdXD6k93gyrBrobn-p(U=OO6n1ju0js$?HlG=`&>V`n>53IbfqzVih$} z>nlBL^CdY^7@eIHkZz%3GY3!1bbI7+V}w2wzK1B?8_UPJh4RfDXqf4C$O(9a48bo? z*P`8B6NL_*AddFcCMIntNpLoj*>-B-#Y_3IECX8lnrL*;0t^)sXN+jmbHlm8#ZU|a zj|kFp6T6+DVzP}9-Fj}*-Vj4Ev+#%=dTxVmepJk9V?>9Z+tfG2P|PGeVym8;#I2T! zxowQ-*K;#{Lkz_%!XtL;xk=q@shA{V1QeZ!dqWJx7~&C7aH81lCKYqt7}2X&I_(WH z6f+Nx_)@QQ&`m5JTfEw97O!Y1d=F8&Z){D=H>1)pE_O)SyWse3jrvUgQGC>A2I9+6 zpP?Q-jQY%|wj;_9eML9g4?0Wpb`s!!u?=oJ;;==iadZ$e)IAWXjm6;T<2rSUX4*=| zQS@bD5V_ALjkoSb(G4bpd~dMqD|vFq6Pv2;u(K$_e1G1?1x0t5jQie)ioEG;??WX= z7-imv9(dC$ybnD(LVw~-ed--q<4vvg?$SSfc*y|r^w}tp!Rc?742(a0Cu#r{SFxnu zFi!J}?`LoyT>b#@dv6dwwBG5*Q7+T`w&?(&OU_MrB+%MQT!Kc%8ALl-V77gL(*Idd zdSB!j<181;wTI(?(|%J;p6qkq>gVk{K7MtUlg*^Ve=Clk?0PM3^8Hip!G1@$2FVMv zoIhH0I8K~&0&jfY#GU?uOaF&(3zJfiNe0o*X11L=GHOZsguCTh!HdVG=*b>OIpX*q zP=0?5ho8DVaFN~~a^P)Us(;`j-95g6w-KRUCH9DIq>_RUD8F})@}s&1{-CqRF>tL8 zbz$HS+I!{%uGJ1TFL|N0ZEA`C2QK|DpbH&lFPfLN(!gl>5fArr)4(hLvPCS!Ey5X= zq7FaN^{9<0{i3ww+rT4hZA0ha7ELlNH8}hP;o%qKW?70!3cRw^7GaK?hcSGmfB2D( z&WmU_K?!D8;FU$Th^e?mhK8lOho5MB_(hj~T~e|t@W?FNP-7hR5@o+)PcCIhLCvS^ zSMDjK3@Jm)xnaPzDlWuX(l#s|cX-7(4@{JsP6-AeUjYaR6?UE)^-3M@4Qi(`+{!ta zSL7>w1jI0Qkpfi;#II1hV;H3|oRVVlkynJ!e(a(VRH-6x9q-?3?Jm|Hbv| z%PrMre6s!s#PtIj0^LrvV!SC=vh5LdrZA{q>B-iT<&-1Y_Mvqv(aQJ@9H$=aY}HUE zIbE6?Ob>d(ogkM?VQvAf?KrnpEKwFy3oAt<|9!AN7?Kb^_<7x1m_s+b z!?9My$Es$=j@-RdS>k$kU?DZ7H^PvwgTh~wrB@_)5#Qc?=q2_zJA5!c)|VC0)|i_U z;4sv8t}>;F+Ai1mN-k@QN7s6s`mI=Ct-iJniU`pdh&(qQ@Ujg-+8(sTtyy`K(} z7mU`E+m1<>(c2r;X?Y`$Nj(K*^)RE4rLdzbIq{WN;Fn7OA$ustFQw17D@<9xolzjG z=7BA&yx|TJUD@IfexfTiAu5OM!TwWY$$GhH_z!eH0T3Z$&aU4J0v&$t@-TISsgemY=o-IJNUzPSI=OVuYq*OOaYr(K8HHK$&`3 zFARZK4L8YFjpKonQ1air!4iG7SYMq!R>I{>eRZwAnyRnn>tAF)#JA04{Ez`L`34i% z)%Z3d;|Eif@xV|*aRVL*U_8iBwTVXXK)3&K4o!V#9N%O~;y6d0)0oUPOO9xLq912g zf&$T(hQQEc&>H}jm`0&5>B>Sgpf@8~lxZ}YJy~X(5sl880mt#dDSe@s%|zslsfaXm zOvxCO&tAjlgNacZC?C>J!{>W22IaHY@cCe3oCeD0ZKvV$m5o99>@|Eom>8ph@>$zy z_@{1J%}z;)YieoY=2ia+w0Jdf>gh*kmg>!lnPJR5^69b%mKN$P z?G7kC?=9_)d)n$!?T*LV@@Lu|FSQkI+HFvFUWs;_Ts!Ln)?hL|cJlw=3*f*9I0ynR z!)$&HYf$oVg9M`4KpW0b376wHNFbaIv_B1%a6Md(huL%}wJw{VaMr&35kErv@|+We z+Lz~NDz&o~`Z?=leFDJ~-^pFzcZbVxPvCmE{3qEgWVYHj`4EJg41%W&!!RjXi*0K=qC0TEq;X0~l{Vc!@Wy!ADmOTN z4hE(B6M%PvChl98T62+S%*{z3t*wX~Wcp2=RCmffXzi7n@$o_CFz?__|1MN$vaGgFo%^F-SHycd>l!71tSe0*{#jBk92bdw&9V@9myQa<#Ym zn-Nf_cWZGL_>JfemMwXStF_Sy+Mc<<;3U02kX84-;C)hOkF7hbDMn|$H|N-|osuxv>9=73tOk+lB1ikxZPT?i;2OpS z7~;N<+B6<_X540^_(0SKJ@Mxy|9#v&7*7vAco>Ut@Thmt((f!kGczV6PFgW*lZCFy zv-95krPr(`E-`*<^Mx06JPq&835*+eQOgr?*B}rvj;59iC>QB%V!?IU)a$Z)V`Upm zF1;?3T$i02E8%*${K4kO`H?XFByk^GuU%^HJAU#S>r3v-Cf8l}3yq&U@4;5{Z;=R( z%<Or}`c9ie9kF;MK`ngB>XqoH%Fdg$Ijt@rZD{%{znK^93OeZ46 z1`vx+2V(J!!&gk_>#z_sTs{QqgVOUp)gkf|OY})5GTR(&eA#9*u?EUlW~AZs!9>0W z%2#Hn;qyr}P`)}N4WADtwrQYzb%qe1FH@|6@;x!q@cCdOUjyZPVyNNsNi z4<@#0pnOe+5T7qltby`X7-{%?Fp;l;@>LjW_YAQfW9gH&8q~tx4 z3M8cto-kIOP8Fx8DwwI_%+y-6wNdmdW9#`^JQM3x(XS?3@78L;B9o(=rXn@J`0`4R zLRs^m5`wt?EI&O8MQGnzIni;E zUc$huTAXRw8V`C%NR6R9^}Uhlc{o2Oov*%J(rqNJ*1GI$R-k`cJ+8$Jx(O0RC}Q(he<`dHn6VomZk*eJlQzch%<~KRdy(mj-7kIKb#Q#Lr+KT@qNcb# zX|*8}%Us>r%R&w!<$8ZFGCgUo-!|RBn@i435eu}p2AH7@PbX@*evH|s*HgPB{VznO zZ?&*>LhOkx=|c1SwJtgH7wadrj(atY@2s;eI7;-7Bh!=Y)*e1Q_=De6)4_cnTSKjV zC!Bd^?&>`0a6-?JAMU-8>7AiW;rszjz6-XE?PA@5izST{OPsao9)Sb;vJVU34-4Qi zGCpjT34^}BCj|b&;|4}6yFPRGnAm=zHZ*R5<)TUENytMUArsv+ z$%dQ9>lx{KgvMljT#7jzSnF+zu*A(ZHgrSkp%6Nw(dL6CB?kg)9c&S%xVc7#ZhCq~ zIv%0XS>Kgl)&}0RwhcAGSr`~@M(DAh*&kF4tm8HI3Nj*t3q*THlU2l`Hendhd6vbK z)+22ZI=H!7hHm5ZjK+D08CioI%n zquc;plv;a~ZvGXgF?VcY!G-YKt%BuLIldWaT-%!J>Z)F#R)@`!rp3fYb0#pByowUuTU#b)> zmM1SDwRYMofIXLZT|?#H(%CO4i`2gLbe)uUVa)P-062b7Q&Wc z)H2RSDe5EA770q1vr&ZVAa3dxJY=u$67*vT@6GJ6 z&^l*}?itKC)Nt#S;rdp^mM2?Fmy6))zkDaeSN7j12dKPskwAwYTqZgW#nF;nUJG-8 z&U6+_TK^6vgitI-^|J5xp>m1EU8pbF_g|xIh=qJXZ}Ht%f?QQyc9?Em;oR;m(L)PX zh5MTo=bqGLuKg}MY;Rqmg(f-W!&Yj`g5g!V2j9^&CW z^#z}Ts+4Xkmr*nsV&YYD+44>)fel`ufh3xw!yi*Pa_R( z1v_}4unrLNz!YTxM_%o#Zse)zS(2+9rI6I&0bX6FxM#CmowoC}km1M|kwjJ$Q+1A) z`udjgln9yCnP7(U5vy>!OvzO~9_|N+9tM4X&*ZmdJ5fnqTqp=qeh9Y%7BMCmA z>A*+h#g&G#a(y6LM=jQ3JdGmJO7M&L=o#Wllq^=?rT76pD4En_!HD}@U0VehKbK!x zFyytq8(sJVp7t%N{R@VyL9K&Vt&mko@@9=A>R8WNQlqtEVrG8S3B4a&%w}nF{05#@ z<6pIdu@Dxp#C{w;X?}@*n2ETwQI8g7uZ@>|TcW?uPVCkQfdP#-TN-b+A1mRqgYjmf z@n$NwV#^BdBWsJeg+8vv82VaaoEt)JQ&+}t%Fk^hb zM7~E$)uG+itlRSf6(_SbWGVyg;`dIlJ&=J3iSXcnx{(#HBce{jx_gjn@ zJig@RL~(@n)-`6L@lZ%kF#)q!ucu}#Mdp8R!X(zrY2_Ya2Sbu3NFHY`bl5LET+xWR zIBn3wW8bz%Deo^~QbknV)DxR@R#2iS6Rcab?zv_T>38Wx*P8onLb#_u6pWyo?V%g? z;@)KW{VqBmoc#B1BmUk=Cz7n_g?6D{dKC}G;c+;rcG%{?Pf`3JEs~9`moAX~SK=kSxp z6D-rYu^xxMx{@@Z`(>7`%i5fTGTHk|n8cfkGt~(rL-bMcmgrk^LPrr5XUDY~i%qrn zkxD4~dQfuTo9$2-1@8JAl)Y^a%k$e8ZjgxIRKu(T|B4vE-vA_}x@2}~okLsGqx2C+ z%}y9Ze}N2HYG$o{#u#}VdRYGR*^tDPT+x5^Rml}S@fRgmbYFc|as|P6kIK_|`f*gA z_LR4R@7@i34-WQ&-SSVs0KVOmH2^Xbtv+*29d|VJ ztvo}COJKH^Jqovc99PS8e5hSi>NvhkSk6nLn|~k$z5`0)A4v)%%H&=CW#r?GkF>;7 zmn_%%YJ%U%aTj#s7G-v7nHx`cF`jN|Jl%e*gv$=b(}~8@?KCBSMikl;&5v!?l?4uI z5&0LOO9FIh1YPP7OJ+o)X}OpJQSuuM}N^yyS;-$8p7C3w3q*SDy52D;EW}!)`I5AMN^rol9dVAyR zsmLk5h%aZw955-JmKfN!^d`+>y`Awq7INlAn-~__2!3F?f;QG0QX^4`)APa zKdfqhR|a2)cR}*UeOT3^VEp{|v_hABSk-=5)uxobfx`b2mc}1eweJ~j|Mpew-cLi; zyQfW@b`p~ici3|O)T}II(9%by?x!&AF()pXAlGOg@HG1r*(Hy%o#D4#r|e%*7f}2) zT6G*zpxFl#l>I^&!7%x>c&2diQZI-xf;m-%r?Xt7h__0#^7)! zBGxm;99Qa>7>IPYU2lfFi#ANv*Snz;TM=#ESmL)M5c#of2o4v6GCZSuI8!@zS#)$p z31)d9-@;a$%)V70R-mW|!1MY{3nRk{DSe}DMkJ2FzCe?s(KE^%SK?;?A+K#i#^GW> z)EO{76O4Jth;HU!y2<=OLWm3-qePui#b>HxJQ?Qc9KXk8WIJKK7<+dVm8#Uep^B-X zo3H2iT_q#y2q9hA7%A$EJU$bMSw@dG_-y-k-AOl{RXebzFw#9gWnkFg zTjKZiR*MvkM0{(fNg<7TDlepNlXmy~Gec*Wgu^rOuCLqV&^y{FmG?<&M9~SkhVt zQO@)O<%GV`?tb>oW`BESSpci0N8lW}_fpC~FB33JG*yQfR<=D!CZ5V#_?R!YY*&G(28xYC;ct zKzS-_GbfQc>?;9)`ED`+$CKBOHf$xYAMRd8UN3FWSY6S`!2RDAUWK;>t%YSM?eyl_#qAOJ4k{2t7(7> zc~t}wIKpDc8-#d{DkYBy3Mur_4lJHAT*zgstI7uka$H0NikP9Q?AMsS%=!)%X5#v)*+dYPCu$Mxlk5t>U z7!sA}wGwPQ28AfTz|+yB_GmIgP2k7}qRGG@sE`rJYA9cDj5Dbl%CIYfk^pmLuYIBB z>dvtwz(a8quvbL8Nb!&I1($H-vQc`Gi1;#P^XKZpyI?i_zUDt3!*qxYG|Pv{t^)*4 zJef&;ZAI+I({B5!S91Va`H$Q{dCEzM8~DA-I5L>Al-J0~V)a8|klS;~bagXJHOPs4 z#PtpzuHvb$GgLY4?ZR(b6@;o${cX#Y%?TE{FL}gC>mR9mEmI`rctHZG44nt`qzo34 zDhOcgkao6xW|6j)Hql|B*ex-pw@CB)vM44rz|A4gEXn^wjg7Z(a<*jrlS8^4MKK=- zxH;!#UhuE2u)zr@CrQUwP<5dO2Lei``62&WjrI z$^<4mG)VSd69liVth{KS;b0&0=e(>MpkyjEg8{>GXaap^- z3%6&j`10%TteqP6IydY!9V_1yZkg!&;Bc#LeGa%lFVC^^-w;0g-0Y0*iY@%LmC$zx z`r<+5k3tPTT52B%7+ha71NwGAAGfCbPlUC&*8g5DTnAl7J@HqfL z-|wjnl6}_%!nKv$E5})DW{`56=$>O0kh$HD6<3e5R*zLcEs}ah4v6o=7!}5}Q)?&Kf zM9s6z^v6_;J{8fP4)Uju|#wJ8C_ADRL4IX=zo$>;c{l!epD{z+)^U0$* zNlEt~L0vgvbpr%nmqcW9w7OH4$ZbNe1q)YXe)-B^|Ak;)-o63lCllzLnwJ+}IKGlj z%ZUgOj?3v~^7kiOGH*!dt9ogxz-hdCPL2Pq8`3S;zobp=b!r`P;*U&gA91?U?sTnv z+V!$~=GFJis_)|aS5)u28gb(`_wtOBx0ii3;=K;%y*CpPvu1|#k52S%oPof%Eh=j{ZNsBzRk05i=@_i?>Qc;>I`n6g3tJ|j_H9*vcoYgaZD3jk`9h3#WDrh zq)|luk(!OkvhfVU|xx z%Aaz2DCb`!bGf5%o9C9ph z2^J39EgZZoa8?!$`W6mCb6lmlLyoyaf;n!xxr3LvgOxc>-`tu*?V_$vF0ex>c@)D0Akiy#p(7@gNGuhLes*U&VnkS7VZ)yRPke1!nF(^x&Cto zFR}>w{VC~Xs0VYoB)Sq$j#P7bz(5s&sO|}4MXFkY!8rqAzGLTjMt#3$4P~*_{o5g{ z-V}!VP$~H-&t)pm0aTw#CD%enZWI|QtVn6KkW2%evF^;9!3qNCw3XLI{p6$tTg&FSzQ5_o`fh!Z5YJFF5j& zBYVtJkK#Fu84Pu>s>+g4Q?G8WQeEXL>toq!X9jXFL%~%+!0%FDx!p`zWeP{({OrJ7 zV*fb~SdhX}UE(x2jn=b9Tbu!zh@7XP$~ty$VeV|ounWLea}e?Sd5WgvU;d*9S7(qF~D2U ztmWr{PNlB^2{>d=2GSTZ0ew3|E+m7}?cm@G@P*`(x|>3mt|3Xoz`-2GYms``SN4FT zY)6AUds5hZMV@XodGw-sNCpC%8P`a3g+M*nt>&q?-heo`RYan|6BQsLy3H>|s<0a&y;GnN#`M}H#oK*GhqSdy5JUMSN+tFvQC z@qFJ;WNVOc@n5kd3Ey|2%ml6eI+hd1CwwAXiG+(^#B#)Z!a^Aet^OvK6VIo7BHM(7 zi`!y35_5mLcZj7s~X}>dUc^*W^!R&@f=))>wvwPhKd4 zM&YFon#6!tOTjokW51YWE@Mpte39x(J+M3$Oe_Vv_&LcbOK5A`EXU!SMv+Bc4X)Y) zLA%jD5ldgjmd~6nF+IxL3LF+xnsdzMTb@wX({WPidERTFNGME;2rdr8(t;mWrqP0{ z!?2OT4GC$H^cH(8o!*hNTm&8TT|HK2!Q`OjBIro3@mLAh!{z@18-i;m_8Xw*1UX$& zZ%Gd-%bOsJmlK^^D|3|Xp?Fq5`TS5o^iaT=u@Wwy9SUF$1*DFZa6Mdx+Y~*wKN?xk zWd9_xsMUTvy`bHG1^sTP@EAS6M_5INw}Q*?9^iVo{5Q9$J4Fnt&8d;cEyTNam>j0{ zB>Cbbw@iq%604d~u_WN_Yz;Va18%)CT*s=|d@LImT>Z;{j<;e--ir8mH*UbZgFF5E z#`4xq2rYfsaMjLSa@SCOGChl%cSjz-fYkgN`qTHtF}z4@1w(GYNgRIy<)gL3u8FpX z#SIy1yFk3uARJL7ZNtCn7c@t1PHrzn!h|zF;H2!IRc@`*2{-DPGs0U^o2|SF~7>>_d9hBaBdO z%V&v3AlAiCJe6POguJJ*5Q;mn2Z5f$8Vq#%KWtgJf6VDUczaKyc z1rGm!-7^1OV^4@YJ3ZP!92X~Pvm*vS68b=i{qQbAfqrCXb_#wN`edP`WPdn)rS5Z% z-%#!+9KXT($s9j%?l9YLAisea(0A;6Vn7dsdiaS!0;h`nU11|{9}d4fGWfpe_-L}6 zNaWkk1_qAs(k0~{bovv@ak<l-7DN~}37Tf=jPmbS@?CF;7Dl1x(Wvdg)9^Fp&=>P9 z{Fb^~ggEvzxjVmpFCngY+OT_V{=$@mBAkR)QlC9)kUe^Jtb7KOZ)K0RWsfF~m2f>= z{$TUt{Ip%?hDI*vv-=m`M`_qBY$tk*mKLRqM$Q&ripRTj2)6^BaV(aoygrhChvG9L zzejXx$l0!FUCO#Bgf`X#W=aC;#Q_(^0qey9GsOY*qJWE{fVHB4X`%q`Kr5PBA~2@J z6`4vsYXz$*amA*ip8WDuZUTF%k|xGG$x<+Ymyy~M0cOR5$Wq1+e3H8uzgUKz33wT* zWA#bJmq4?=Zqk)P1-gu_)+x@1F8iTNCAv(l)+sCpqeK$YkwO3s5J_Epw7JZ~NR5es zi;aA=i_F7NT@p*`=KGq<))=WL#lXdvd|wyY8bkGiSWXw8U@lu}q@EZ97aREm7uiZf zby+N@n@=&9Z8B0%iGhnR`4kt~CPQ^aETfB$Hr{rETfwbt;PL~ z)RSZ2;!8fX6!$k&Ka2&RjV)hyYL>6~8Jgv5j-F;!d#e-_^BG@=1$$VFiQ1CY@hRhI zAv{a*%+E)WcN4yMb?7(Vik=Vwfu~ZIgr}^~l+Yi5pXHU|XrJ&DZ%v72zcTGDF65iq z5ZSLyzhoD7ndTCU`lP+=dt#{#^snU4t>7}e2e=+C{}KmG^&^)!I1@fIhDb*RDJbqH2C>UX`u^{c1X67XFbxglV?3D3FleMT8{Fp z@xcoP_S?Q~Yb7MIPO)6+Zr+smQ zSM7*F4=R0RVl9>fue>ti?J1+-nWW*?^E4F6!`_-DEnNM}d8bWfp1j8H!cCFEc`1~U z3w`8tZq zrcUgAVAm?qHmErY^cy(w46mVvM=)epz8nu%llOqd#OJx_zxJnOB2aKhCv>v!G>^%$>9Qn-qK zsVb#Fo=oI4ciB6sLUUQ7;f!iJyBT_3{vu-;I!kVp)+Q&EP6D#o9*0c->H;E@zTYGmh|8X zG+Hj7!*T2s3ISNh@caEPH&zmWvio@plx5d^35v4wc?;xaN4X1RvT(wJQQ1z;+p|W) zb5X;rtH!XJN5|HHFg4s7mbn_PhFSk5J6U;&SUPky+ZPkLsLRBK^+b-SuKa@dYB2G1 z%2`?{gC!csD2`<3qeqA3H`p)hh0cl)5^JC{qL_C76yCNK)jANrd^9F4LaV%(-x83% zi$MJzYu$uu>WlXpiJ20i6fF*;oZB20 z1lH`~OJ)?^jjtRP+g6}j+vAr@V$!Hu9!2~a0qKSWYCP7u4AsPs_Y%iUp=#+A@l{*W zfd$k^tSuYW+8V!H6q6PfT~?UcvW42WAmjkn_CBihRs8aS7;KnUY$3mROZuw?)ZJL? zJE*3nc(4ALvj??23i%hdq(5Ci-HNsS4%ONizq~I7doWtNFf(!swP-=e7OeQzhEjaS zof|K{SOG zH07C0JGiTI0_A0*t8>u~eVBK!`~4?wRD-Ouh0V7{=5lWrdmtNdqoVbF0bMOOKeg5R z=$6i%E{c?7>Yjlyr1wiVmrnf5_53HfH%8UYI|rGISNUYP|j=Srai(P#(7Uishso~Z?ce0VXy)v)wY4Qxt&DhZc!2Oz6 z&m~n(`~AK0dWu%fD86EoUO$Dqqw4kThx?jh9Cu5YcRuymBf=%J? z$qo1rVYh@$NW2%79%HXhrPC8C!|37m31ReuUS470UybU8h3}!H*DvPoXnuXWdSBCF z$K6uqomD=~E;+kT*DtQ#(JZ*V>PLQ?z5hn0fG1A7_`G8K>(fRPU%1-1c}=4YPK0?5|}W zbn$UslrsaS{>yigl_Dw|-r~wZjfN@45jOK1Oo~Edx2K3@lLK7J1!2^8F_w6g3z=6v zSmPgd_pbE&mU;cOtzIbSEia0{^YI^!ckCZz9$e*9=#sM~y?z0A$BWmutM)Z5bL_8U z9(467T=a$-Zhg9bLG_Log4-+iHNn;YfKD!dpIowgEyXg*bmJ@JYtK*YweJTt7}e}q z8Wvk_3L77M*gwz4Mh6k?RrPIYP(prQ?Z(PjnC#)$ zf61|5>Yyp%dbsTI!DeiJE*-kS?_C>cg>*jOXi(v^>|o4y!RKZS^eqkXT3#rk&*A{> z&3Wn`G?_Ni4+3q+81l|Aa-{qv#0!_d;E{R<=w$*nV^lo~I;w;O@L`xEu$L=S6tJKp zJY>d58f93dDwL49sx@~Q0JnoO6BwP}H{^VsUnd@e`cd5i!K z31syk*x}ERrEtL6fK(W%ddLAg3x$#LDscEEhcTST14Eoi9mg2MI=l@+0&(~N6voJr zf7C36&LZ)fL7*g8KEhJ^(j7VSz65rGDBwj(t{IQiPk#1dB(LjBWdUWMx&zP61ABJx zz;X1w^i{$IRmbe4VD^kBOY4?g8iyI_M!^rY_@?aQQWfCA^ zIbUFHz$iEeOh}tvf?@%EXCBDCDBHJDg&*zn?Y^uYNhIYz1KTQ9KU55{TqVN&-;f4g zl7#)hhE-~`Yf`w)y(wU&i#&&uoM0p`^`JjEfp>v|0Dl$@<jAx$^>wH8)v6152J>O zRfoa$HqI9^!c28$Io=)GIAI{$&`yAP9O(yf0!AHJj;BKVCJcU?uM;pQj?^H|!Kky# zS? z$8j3O1sL@=IK~g+6pT8koJ56o zOhA4c*$J?TV>E~%?=P2g+@ZY_Xf))-z_A3zLowvMR|06z#@H!iVc=rrVX&=@@ddPn zLT{2GA9zhBYKyZXQ%WM%wpj9KJ|E%V_B!qA&{t~A$&ziUpg7Q{51O9jT?P(!BrGK5 z$ye2IBk0gJ$#Cm@ndigGmF(5hmai#mBk%S}5Afyzv=0$~zaH&Ki0s7P0D4l%cO~^! z^tQ6RnX-5#F`~6HPuac~&v}Ote=qFMM;0{OH$@il?5EKS1oj*0cZI?XdVa64`M)~i zFOAq%o6{j3|B>Or@ZefW)iB;>USo+=~{)k#FEMFY<0E;#ZN1yRaib&v|m= zVE#%vt&(sDo6X`A62b8jSD4){bC;wQn z$mDrs&mG_OBhLznK@B;36miS&g~IX|wDz-T^4~1~m<>I+QnWyxJfFnwv>yQmzUDO! zmFLpwFDMI?^4-w$@)yaaR%gkvmE`JhY(@0~wrziIIoq~3{}PedeJqwp{Ktj+(Rp?g zXan(F2T7YfQSe9@AWyQx(}V(aWas4+(t8W{+rp_G%5~+`4%W}$)QWRe?An3+7Gi7P zv0`HDe`B~`Wu)M~=;COyl1S&ZR9OIhfkyes3Su$<8)#+yTtexPv~Tx=BMU7$ceu(aIQcL-nhDVXIm zj4#^^y7`Xa%XUFmqxdo!Sn4Cg%aTBIA2}Xi;`c*`)EHOV(B^M}Do3A2M4I1i_O zSPphUVSRwF5ol5d7aPSS7r@t$G%&J5V2QzhE<=_ijJRrxRv$=-idf$qTR(I42>+(A zHhUDAJ$hxVgv&RxNBP;KN5)EiN9&c*J*D<@i2;Mj=S5Q@FZGmfW#uWiwFy5&_Pr%8 zvfGXm+Tou9lqIu>5)_s`Ttav?`cex0D8{BJWxMg<<~@Sgdo8}5i;3q=!CiX<<Ao{by;) zH}wA?o1x3T4ntSNh;2{3$nrQ#R=B{#ht?O5CW#YAU4bI(+e6jhVQ?_((y>!&N5&R? zFg4s77Bm-)R#aq*jw>pzr=XP;w?)U56-8(?06d@^2Y{#lu;9L-@5qhAxw-V9dMFSi z-j#E|6MF~fJBjB>j@i&@f2H6)*;(#->AeVI@UsMpBI$EhVXxiUNb#K%^5{9}XW-u} zxKE^vJHWc%Z*n$r=!Wn5;j0a`7juyB+0ef3M&^*->7jq;MrYxGE)2p6tpg z?6P015?^EShBK<@^uJPY-y_1OKo1nOtsJ^b_*B_*lkim0cO}JF(Vm&Ym9ILbNt2yB zc7AJ-CVTds6`Iu2)noavkBuE8gQ?-x^E64d?|5snXT#Mn>p!8BWZr0!4lAO~&W9#T z<~j`C^Ct?P>ht6>57wb>fo~+tJGj&DdMx{25EiL+atXJA{gs48z)lX~wsKcNioNoS zyp#e3JbVon)%N?s<^EQk$oBH?4nP}b0(J^smW<*=wA=3i#HkQqC1X2XT?74b`6;aM z-=2{DZP*_6#*z4IzSzjD&fJg4S7f_ND|?IA4@f>l_C*ifSfR-l4vR7T)yO_Eoy8l- zD2$}%qea8=Z`t?ih4zXrBCEACqL3E;{v!L(17)vo93Lp_^FfQshHe}el>s#q#4-fo zxL9V04ld%)4@jRxpeAB%AD~*>;vb7+{Hf6vMVUqcR0JU;8f#mEYURa07R9^>(-IW& z`?sVIEubF4S{I?3TH-4QV%CLeT`J^PZ%Kc>fVvlJdl%K(9RIjK=EcG2FA6ii-9r6g zLC6lQ?QK-+%lOBAG5!a&EDQO`Thbpapax;BZ=#xB#8>vlJU^ht*tJfq3GA!frVHNM z@V_pj>ATj|YC`*~1-wC+_nL_j66o9vHggv;M_ZrBhH)<(@wA znG1LG%hL|Ryn{RaCj=qB-4k)maq`P^rq1*reNnIMz=eC?3epaA-8df5a45qme6_!A z&isq@XCLo4CcLe?uc_T}_c`Vrs?R3tocZVL&sOd@*7=4SW*z@k)7Or>4>9kA_-uly z|Hw|3<#ool6y0ZH??{Kj0~)!xsPG^U?u{!I^R&ktdqZL`o?j01oO@gpJtf)Pwr%BzK=u2)fWzJ)4h0txL9U*s)^56eXAyr zo+SD^0`s6q*q??7bNXG61r;_#^Pr4KY0l^{Db0>DGDx#YIXJlAQZy{tZ%VSwnQ^iH z_~RW3!rMCgn%W%u)0qdUKF-!TGtSo^uiTN)`Gy*19s5;NmSaDac`(Gs8K(ZrcjAHP z#rGD5)3F1hA+3PBCAn7AHa+glHi=CU^>C#~{4U0V!cy7njE)O=1EG6`Ulwj?l0P{< z<>37r5nC_Z=t1pk>Tv8o&pdd@r_eTM%f9S*gIP>6+H`H)z>zpm; z>u*%{zxqLniM^j%U8EP$W7mbcfP#_s=Ep==~>sab2M9WmP6(b@B$W|8& zRWY@UVr7YdEl#z9g%%Z(2neYQ1VjXcq*Y{ob>AW^u4~Rb%t8)^t$xs z&Yhd*ESCs}=YRkHjtf~^`l|h6b3~Uu*-6)KuGbysJTG73mvBd(a3^*L9>Le!KXc)Z zJK+g;`r`2j*Wvuvo^S6@xOq`!%Rer!n7C`+`ALUz4h%l=uKbT79c>>r9R7G0;(m?q zs3~K*n<-NJU_0ihCSk*M2B~wngz44;g__{D(!Ed4!}M+gmf%z{JqHat?3!HD2#bML z1Fqj773oJ@Y>h0}Rk#(ii?1W0cEhKARI0t4axH>igPF^;O%#`y$#d7miUsC664RrT zn+h=upaP8)0rrNyPvLSN6kv61g=;ty8z>QA{Y)9wxQl;{g!PY=mee!3{R8$9n!vpL z>po>f^5rtju{+W{#IT9c9ei^ek;g(L=O~ejrGh)M{35STkA?K2*BQA?thq^wJ%)Cn z&oG8IVx6hHew(8?R)9sKlVp`eo<;^q>|AH%3@e?uPVWf!T>*NrXe-~*E^ssz&|Ns6 z(A0uydqmhxgzG@0Mp9u6Lwy}%CcBG*u85pB5wy2Z+G~&w7s|T9Z40XCVz3u6O8#&| zFz_7z>o2`hyl4!U;*~mq6i<9}xqE1mU{e8BEpq)MUB}>bDRi$3I!Xlf?T(TYC?)0T zk+*it^oWv&BLx0L|KYttG%b*k^bMR&9e!MhZVqJZ{6-BpeShWOc!zI)kKeh2Z%^Uh z*u%FU;5PtF-wA$2KHvDcXLpb)ts*OVx+SSUy|w%zZ4WU>ERuf0qoX@lV}Ye3l=ou5 z2PMPjEmT5ejd1+}#1G*)KC1GoSWURYRK{&<@B}faMb$ zRI@_i&5pY}Al3T?Mx*FyPh~Z<*#Yfja zlLYG|&v+wThjX|msfabtwXKX*>PnKa=AqlL zr&L~5UWCLvE|%`R7BtOq%U`^`D~8Io?~bT2sQV-*z1CHK@mg1%z^$v6%)lM3HsID( zv{!EW2kRz4Ww`s8Wz33{CC}nUS8y5Xh8N*B(4OP5Z_zgKxUIC$ca<3yw?34;-OjO(@b~OIcC%dA&Z~VNF8)E&zBYV#Ac%Q!_4(2jlbuV*YawO z*G}sMZl|?+2JVFQH?Kj~Qm>sBREE2c;VxUNpT+I8;4;*Ge~Vlx@*DGFY zFsKZ7kM<%}k!bB>UJgd{;G#TYE$y8J-$Tx7W)l&Kt`qg8l;N5}}hG z?Q8k8oS&zQPQzaPZ?Xzt7s!}coof*6PCk&N(rAK0Y8fX2Jn-o*nhbh z;-d>VX7O8*68@|N{8$2)p>9A4?mWeQSrNT}J&245`y`3?D$AX{l`{*w9^-X$$zK69qy<(gOmJ=##SLlo)y)a&%M z#%p4A!gB6UDL_;K48Y!9X6QP_Uh1{hqUMp5R*SE%gL(jHO{aR7O+ppi7 zOfk4#?+hm}?AH{Cy^g>TAi^RS8ixRA8=GhOJOZ&u4$^rFS#g&nv`g}tS`w<36n9EO zJ0+hXl2Al)S|ypMlH@8Sl%!7675^jbxvYKU2+30%n(bAiy4&1HD^Uos1~nMO0*bcj z`{1-#4JNU`Q{G;oI9JQq_lajoPCsMcMbF;cLB_r|@Ji3v=kV;!F*1x(u_f8V4EsWc z79Sv6iBFa7^~0yi=HpXkt;Vrr&qRExEW-~UAe)a*mFe)QvcKR{Wf>C$UFwKAV5&@} z#RtgV#HY#*`r%V$^YN*&RwIe*`5iu0mgR>Jkj=-Z%5?Zt*<1KjS=I!BM$MW7rpg3b ze1L2fK2;|4!>7vT<5Oj=Mk3ks6h2i3p0B|G*?fGeOovaEt-+_tGADp(IOZHMRfaXH z1up|~yeDWMoSs(;lI-HA`5i+|wc&rDe{Ov01L6?p2~|Nq`7vAZ*JBIA@&_w7D!_lp z_awnrUYD#gDEF}b!ma4`pO4IRzS=CjPG{QcsuYpnKSc6w3;C!w!uc@^Iom=$;*D?} z&i@N~ZpQM}*NVuFRDTAqi!=-E-_5*WDCq+~6jRAIQ#;*rv7>fuLb>=x9Cseg)7DWd z1TX62xS=%E(SeTLnI&$h6TCw6kV0$4#tVY#81yNkClb6cT@c7)&_E)Z8H$RBj|;9> zV{?eefPbB|eaL@^t}^&fkhPoqXEHj><~L+*BW8RUqv%g5MSsR4;J)6?m#0uyq!hh| z$M4mDK|Q+V;_cei25!|SBI)qBaQaAXP`165K;d^zGH~q4RM>7Aa$$bk`^mL> zmErEY=FH#oY=*j_MINVE(w89tBk5}?6l?kxB*2!QOd;DHA0pEnj!%Gy@*e)rKjILh za-O5Ii0)lk@jzsxyXrG>w1vRoC)QE7I@)G1q#9BrYTwUzK<>@3y)=AC$Vb{VQ+&QR z!N#r4CnA=qXXxyf*q2ik>i6mK*M(Ha0mxe#ww!+OPSmhf_P;n2<7C-+4LcL9Jw{=g{W*Mg zBWtTuy@nwhpk=yITyzYczu)vefB%aW;4;)b+OFEVBh2IW=H1NWHeEdPxV3pB^Y}>H zDvH=rvzQ__!-maUyw9xP`Z&+PpIN7PA5P)&PiXHvDre#^GfW*zw^hYv+l~I>`0j)v zZtW*Tso5l?n{Ud+56r6i9CNd?IhVQFp*zgnY;Wc>H{05_QCQZRXbNjY_#FPs3LI|0 zGwWvWGwV_O$xS!|e`W;+H{jtEE<@dKd!Jd4;7_=Z)FMWWG+qfd%xZZQhS8tIw_Ra% z*wxb+f*zUzMRULx%YS}DbXC+_ysricYs9mhFV{sbm##Fm?Gd$b&Ua{<`Tun`)c1DP z&U}@(N4I*J_z%v_4M9otqvmlZL^a$CJ(HfN%ye|R=`tma?Ls$ER0HO9yo5pp%hJfB zOX_U99G2%88!AW@Nuhe(XaqZ|HA-EfLj(MZV9`cf@EV^S(z6|L6|HCG= z=T&xPR{89!QO2%W8@7`w$_N?lzC}YmHMt5wYPdOT!)8)N2_d7+w`kDk)m@eT z<+DGGGFHu6xq(z#L=d(3elg(VX+7J?Sn%QvQQdDnF(v=^rMva)N$_q3c-INu?JccL zxHEfUoH2XBh86yyg%dKG#wfNt-c!0b;f^slt|5C$^i6X}O2SeFORM$`lZk09W^@-h zs=lc_nKeIx)jO#dr^Lp$*X4{e-76D%YaU36Om`m$nP7@>8{# z1{Ni(lT!y8q5>qJn-?Tn^3QGXV)HU^njWDSr|JFCOVeBF<>QHPnjTz+yL)iR-f%EZ z(}T-D;UWj4>j>q^i3_P)`*b4}#o|bkrD}g}RHG&&C4QNRrQ7(hT+0-)+A&Twt<&KM zAR`W)a+=D~77(D+T~+#b=q^e&da-#KI8D#L7pLhh@S^5cdHHxEoTdkt;qD$BvS$s( zX?k$^=Uqg%E_?f-)kX1D3lB@>EEcG%bYL1zb_^eA7;6us85 zUp3h0*hSX&IyR{WdmL|(hq`sIss=Q=1^f+OY+eRV(;M51)AZ(eX?iQYd^{0O(}T-! zcMlHPGX>)`J-Ga^FCwH%2`%gqbXEtVd%R-%?#KvZ)r8!zDotdPX}kzgZybdiGG&&L z%cETnVJg*K^)0(fO3gpnqCG@(HWeM$-KS5r)-rZ{Wse-$4~#rzq#W{?fx~i;PR)G6Z}{kpqupnzZ>cb zb3Co7UKminh>|#0$8BiLZkl6z^RRGQZ>gxSRMCw`BAoY>ih4^G8a#fl{%g;-_X@f{ zqTa+kMXPt`{9$nOJ3&u}Ol})@?eM$1E|g5l3z#t^bsa&653*#=GYqb&tqd~@sh*xJ zgg4R?D|0G=LJd1x?iwd@d6X^$>$8YuC<*B?%EL4R2;U>5^4)DrX|v2`s1=0SBHBDA zMy4#;h%8qyqm69F-^J(wtdG@L$gnbGPG^>aj`e0yrD&5#6M_MNio>2gnhUz-|^WTxA&Ae@1{=$6^%BA*4z@&JqP?J}Bc`L@uE@obD>Oxr)f1 zp#p(vHPa>ISqp7lRM{znVTuhM6jYpNnBoQe`WA-q#fi2cK@(;}Jo@)hox zhvU$b5nAD8oHNU~6WmCu0}=ou4d?!Pd!5B{MukhQ6@= zacuBI2H5Ha9=(6(53t1cxx(EPen&)HG}bmru)^llPXApW0oZ? z7R}_5y}*z+)bwh%@`it&d^PKes!sA0H`u|o)D4=|g;`Ew-WS+Pe-|+DZXx3ezJ&rY z`0j%F{#Gqc0wzXa3x#@@!3Tl)7NDs@*%f>X1!C~s1@mpHMv{Po5!gb37<>?zZvnbj zD7b=ep+F42yI{U|t7S>Rz6hMWU?&dZE8LGW`Ji;GokI2m?FSnw|Goq7ue&C9cZTkB zr{^BKCLXZSJNY%k8gGPiu)$peirxs<;T-O%;qVID*yX$erQF3UP_sk$1wBFlQqyKVE-MM6kIrlCJ)yP-WjYiNaA=6S-YHBE#Sd zQR!e<-6%WKvO4!l=jh4{q;m!|)zUdbi zElmgc8x=+89H}!I<=t{N+I4_}JWTex(w@1ux%yni+ZJ^eJ*K%&?>zJ_rD90a!t0fa zn2(nIu8Y1=S-8!$?_G*yVA5N>{?8)bww&kCxphI?9L5=smi?}@3;b$D|1r`!ta+aw zqcF^JroX|c=|;*ribYXP0~e!gc5p4za~Q(tkop&ACxOD z0$)fc%q zZzr-v*YyYa6F&LP;*Gv<2I{u6RgxK#-&^&EaYHxeY`Z*udV=C&uQO9vE|~kv5qGY) zK3yvRDr}6tGEr2?S&zj0hVr)Ectt*s9y$!`P9X^=q*DE~g!&5e4e3PgBHPP@&9?}7 zv2#V?(XWS6l1^BncUz*5dLx`4vqWcGqK|kZT!-^td%nFt*&|cWDN=d+2dlS;H@en$ z#A@lsZ6>MW34W<@ZyAanMmmgU7cZ$?VX)TK(Pb^+WlqZqq+Q$OLMObas++Eg`yW-^ zB-I3hs_t=BoKsn6R;KqW-_R(B+LZ3ch@rVj(}%b9rZ;H*L3_v0t9p{2&80gx0pmgN zc)9#Z?9FPx&V$nhD1nImt{TjfH3hiG<`rTmca2c?0moitn_yy!i})TJTdXxOiUT}j zE9BT8b~z|^VA%T~$pm4MA|zeBt5^Jqb%fLa&x`v2;SZ<>;nd#)G=vA#awR>E`$Sj4 zZz+%28>Hc#VV~=t+D?3*v;EHWYvgIxmrht_?Y7K1>Wy#?bhg=+Sx3CF;)enBzkAK_ z<}w3_`QG?#{?0N(6?o~5a2?M73wo#$oM6`fF)(s$dXc}ex@eX?)lDx~$;(__J1EGI zM$xl?d}E0INqZ8|2?Ca%xI>%l?#7 z_GdiO;Cy)sZAD7iYk2%#{TI}udqZ5TU2Wr5oe-&p$1SHP%7f0@TPIMYos(=F;4%mc z{~)m*r0^RR<=o>eOZo~VZX_KjS*+>XkT_d9pR&#FIE-v{IC6h!GQo|q*#nW`?y6H_ zvW1YqPrN~W-O(1vple9ssD0nVBolmIJf}CIk6U|^D7H*3ptEnsj(2r#p)mAwfuQ&D zBPA2ewT;uz6Ym6#cWKv(XZNoT;#Hj#<(Snm^!RH+qT@%|vcpE|aQ<_ds?)n@n^65n z$njxK3_qsc;O9(V#t7~2i*@%Fi^bsS_ue>K7*h-9FtQe|!})#oWNve$9%jbbQ~Ba; zw$yFRIBROOIBvwWg1OaVqRSqTa`8u~Fh}J~jPBg?oe4BT=Rrt0GwCA)A%Elsb(%vrNw&Q+QsJtK7l#dLggnz(5oKgl^pTI$BBmbc zc_q8WRm-O^jsE_8b`5LLuAae=_0U>S6#rfW#x39S&LYAZ)%%jJ{-A2`7pB4-jm!CA z)u4mUP|{*u=-(*01LoDdgwGTamd0dySeLnNwy9F}QE^DB zCnv7c2V@%X*^!_X6Ug0#up zzdibReUm16BfZE^n>sEmBwrqD>H3H)zFD1ZHD+6nd*fv&KbdW9$hIbVBV32`pP*;T zHnqb~wYAgX3D|}>^vbO&hbkaWsjFAM-=ULA9+2SjK}kQgQQ5_ZR+q(BEjp}{k0UHh zoE~(xyESK;^zW0(cG&+4zIOg8Oogv}lohB`=#A#FDdtcGrBB@yDmYG445O>aVON{R zCL5+jP&AK}C=*s%LITYD)0pad%8PlCEru#jPV{Ahe^R0^b+EH-MyNEE)MB)g$lFY+&J3u5N2Cr)rywHqO(nFc#o3c@e4^h_IBdl68}bJr*RCVN`>#;T@CA$`Vb0x?hv5uCFPhvsuIlh}yY`%4eI)1P;M<#o zAt6uQS#?uiF6yR<1)PygCeNcQ!%WS6GFNW|lW)?oFiTqtL-oFbdu51^>Hy}iq+dqG zy340|E^)YSikLUzFfiQeFFd!9bQxlEZN;359`QDGwa%gT71;7Z0p!oQ&SENa7g1%N z-X?|A(QBqkv0y2N$n3*sWT=IxMi9e*0z-ucy{(4jX(<$VhD8j1=}}68i6!XGS;HSb zU!<`iXcM*VY&oNM4B}}~RGeU89RkaBM|B~?rjl}^WasvvNo5Rx_!9veUInn>twC-Z zF!eSD|2N{Z5yP5oRswY}2I45OZab451JQ_BF z^fv(w;Yc5552p&;p9<_87F2mu0W5eU8A&x5?`HDa ziFz#fjf%Id3_*EetXsgA)56_Mwn4OTx!WATBwGZYT@l33F#TcHiqB+v6^ox;UrvB>&-5A@{>TX3{rsI(67r$lVgi#L<3I`(LFu-jUkhlXmWq+Eb)A z_DJmqqzxZSwI`$%`BH%Pdn}^jsbre@Rjuik|9N%->j}yVI)yja#kK2}xW}sB#ukUQ z3ML-5Qnc0JWfFMV30~SMtxG&((c{?SP_3X-J55L`U@1{QpH8+JYAHJF6squg!7klQNg;XXS@-v z!})(f&!;VD)iGw&Fd>AWQbXNpZ}ycHBe5+m`!42sy{6m^3?IPnB7jZ;L^uaF6=04Q z=ruru>u?VDq>D$=tEHSfw8JiYG7@Yo58`XSA&MQ2C9>rHJ`n)?DSdRx7rYV9U(zXq zb;@VF5w64eBi{2VJ}n-hG>?0qo^n0tCC9A*adB4+jceZ%(P2;zJTSmtMI&tyudJDt z0zg!>OT02A&B?o=|IE1omBw~&WkP-NRok3U*~KUL5n*R3t$*@FeKysFa? zIhJvI>AUKMBFB~IDfI4;dAcwkUDylW29!h&^S&v+wThx4DS$DW29HbCEizUq!ixQ&GHc5`9`@5 zg^7p|q3~zmCyvUQRM|MSdlX9J)EDVodF)G$4jehM*TzNj6&)7S89H={@bN~O1sUzU zMV;ns2EKvaI6`+fGt5Ac^Q@&JinXSM&cF3bV=O}8#D771tj%=Q?n0sr{w^N3lC{;Y z#E~PyOcXtatT&tS?=}B9bn*vzfA?0VJyQI{QMsJTALjktgTVNMyubg}fer&aR08Ms zJ77}tL%hGF>k5OtZURGw2+LgPHl$a{@4^Awj5M-m3O?BeaR4DW-~_BwO_>8G+pwFsh!5EI0zDs;eyIjiZAJG) z{=v#RIEPhsRX;quc5f*lED?lrSS|?H;rtQrx%obSHow?cc*}HIxkS zxaHJYHj|R>0k*uA9r}!m>EQ_)MehP~{SZB;JqZy1fGh7~%HBzr9N{}PNaq}Bm4+{r z+kZ#&%p_~x57pa%tZ(=-(yf+Vi{WPqu|E<8w`qPZ{T3uxuL0(KmvIUEbOyJ#kLd3G zhnN8rFKogo_8D)4^O93+%_;UlZ-nb`{)qQvs&Ha=q9W+5LpxJk+`XF3tvVo58r1}P zyi8bR+XKn>U{&9J;rxDt*w6SXqO~aBLbM~MO^DX2+=Xau>iuHyzvMBo!Lh}Q%0G$& z>*2LVTTckD{S#hnCfQ0T;V0grzV6g5md);ql%rLD7Y7V$=JMmOh$8Iwrq|IP5Jzv4 zcz9|Ion0S0-f3JWp3|QY%d0IW=2}e}dV)$$L^peJ`bTk`K77&-9_RNXZip^fJX^oI zkXKbK${iV3Ku^3GG~U^|kV5aBRLHe|WPl}VaDJb1H5>f7aQ;)Q(>a|VQ)}>Zq(9G~ z_tJV%{Ym7wrMiI6zs>S%4|AFPiT@C0sP-OAKlEO zxTyeq<)e;VTrJzKjZ~mjN5pg!VJv^7f;!FC)69Z%F1W^krq_ zb?M3xT|V^|_p5HgWbsUUl7w_UEPCjXq9u-oPVYw$B)8lOI!-@o8J8~_!`suhdZGBF zSga&r;6fEQ#@(=!SQw8;VI4Erg5#-=Yzp+LX$5lG*==Hoi$+d74z( zL=fqGzp(hcno{X6nf+n3@n!0UlcWj-A*0W?$m~;-ytuu>xH-Dvk5p}Kl(wV!c9d4B z>y6T?nypbhzsy z^cj;>=_;S?h%!27ZQznBB!mo=FV*0a%3fSvZfuQe=%2M=C#m8hA)~{Ws`nYguH=-@ zz8Gb^Gi&8mQt4L&(H-BJLq7X=Rfd$${v^tHZPtdrkt)s;GFpMW!^d~mV%H^OQdGmG zS^r17_kCes1Mbb+de%tJnr!(|_&@SZ^W2n#R}>Li^kpcaj6@n@W;d^CBW1 zCj4LSt#Jj565k3MKcL-`GyB}?zQn4_hjUuh(?jAn35kyu{_m3j#+P%0cA=O3a$8KA zMG1eE69*dB284ZXelF2+QV+}6;QS|$v{~Wx;rxD&)liVUaY2z^U#c!NESJ=4?*AbA zctgqb6z;23Kb`6MFj#E}=RcD=Q9UKs&{;Po*MO+4IR;f-Z;nCPcsobmA*s#Lw?8O} z9ju#! zFr)EgdYxIAY<@{W>1oUe<$YwUH(BOLVYM10V!Jb|57Pi&;$w9T_s4aYTrt~>Ljw1WGY+f59WKjeaGeQ; zq4JiIffv3Jv0c-hcRXp*uQ7B%9}|spdnkIVW~}k5^I3*|lgwB}xAOQ7T?BSO(Cv>7 zXL{Ci2U%^*tK!@EU*wxeYni-gumxS?Y!^#C7~N%N@m+mXflbWBdaGF;&lz)5rboj( zC9W{KBAwO9VG4h#NA587QC%&_7c45;9l74!Qdk6TgQHBqKX1yj6+P)8vwi2+7R8k0$~Q?4*QNYV*2o_h3@qDvyP!36 z7HwdoWUb-1^(pJ)ZGa;=DeE#LgoD}8Iy z>JLuLu>~%ozPwW<)Q0*sCi~}`|M!)?RcZb4Y5g11`lGzb+O+;PY5hyn`gz|)@NQgM zKWJN(=8R8sZcKAVd6Tti&NXSyrD;yycM-fBm*xa*tI~So(|R|i^+tJ0yu6bl)Y1u!yvg~lXI`=K()0%>Ds6$6z=a_Dwe;UM z@_v^ue?MQiE}y+TUmlVl`(l2ie?EzjZ*%352J&p}dGhOdW#xI{XY#_c^4w252Zar> zvb}WI97~(xq=N6uvl6r=$P%&J$2n+gh-K`jyJlM2MBuF_c-9zg3FT$6+s`>D2XBGT znrvyKg13nO{4B6w8F*_C-Sre$Fk*XzXggK0T|tZ|M#LKq$Fu6=FE7^DE%OV$7RI

TW7R#rmJe|iL zmon{StRs1{n50XleVoVCs*nrv@%k_JaQ1bQ zr(Tdhq0jq(Gq#iLDvL#jj`BIX5mkE`$ucD2a>gU7OJ(7MLq}3Lhg9SfW$f;u^GTdB zs_5i0Qpb>F7iYIpwXRHlYv_xeoF|mzg=O*^LwP$m!Va>3S?raequV);b*NlllPZQJ z+c<~XRd>D)FC99vl{3DbeCccU*`f3AarWPd{^V=YsUgW`&SQ5}oUi5CLtkv->}#Xf zf624uaRE2gGoTP-pVO>F%qA9uhzBD)Baud`STG`2$ed-Oac;en$Y7hO+%$w~Ym&Vz zkl9_~iF~#Rq2xGu-TXR^=hFd!^6?S=g=&RthZj>r>qc)i* r;*M;+rD0_+TiVZNri+m$KI%fI=5VQZQr literal 0 HcmV?d00001 diff --git a/zig/src/core/modelopt.zig b/zig/src/core/modelopt.zig index 5f4e0c9e80..da37ad7b5e 100644 --- a/zig/src/core/modelopt.zig +++ b/zig/src/core/modelopt.zig @@ -1,6 +1,4 @@ -//! ModelOpt's quantization block in config.json: admits NVFP4 weights in groups of 16, refuses the rest by name. -//! Which linears are NVFP4 comes from the tensors (nvfp4.stored), not from "ignore": nvidia/GLM-5.3-Flash-NVFP4 -//! at 423acf37 stores the MTP layer in BF16 without listing it there (da920bb0 added it). +//! ModelOpt NVFP4 config check (groups of 16); tensors say which linears are NVFP4, not ignore (423acf37's BF16 MTP). const std = @import("std"); diff --git a/zig/src/core/nvfp4.zig b/zig/src/core/nvfp4.zig new file mode 100644 index 0000000000..9452970477 --- /dev/null +++ b/zig/src/core/nvfp4.zig @@ -0,0 +1,148 @@ +//! NVFP4 as ModelOpt stores it: e2m1 codes two to a byte, an e4m3 scale per 16 inputs, an fp32 scale per tensor. + +const std = @import("std"); +const ckpt = @import("checkpoint.zig"); +const quant = @import("quant/quant.zig"); +const DType = ckpt.DType; + +/// Inputs sharing one e4m3 scale. +pub const group = quant.nvfp4.group; + +/// An e2m1 code's value as f32 (every e2m1 value is one exactly). +fn e2m1(code: u8) f32 { + return @floatCast(quant.nvfp4.e2m1(code)); +} + +/// An e4m3 byte's value as f32 (every e4m3 value is one exactly; 0x7f and 0xff are NaN). +fn e4m3(b: u8) f32 { + return @floatCast(quant.fp8.e4m3(b)); +} + +/// A stored tensor's dtype and shape, as a header lists it. +pub const Meta = struct { + dtype: DType, + shape: []const usize, +}; + +pub const Dims = struct { n: usize, k: usize }; + +/// The linear's [n, k] from codes U8 [n, k/2], scales F8_E4M3 [n, k/16] and fp32 scalars; else error.BadNvfp4. +pub fn dims(weight: Meta, scale: Meta, scale2: Meta, input_scale: ?Meta) !Dims { + if (weight.dtype != .u8 or weight.shape.len != 2) return error.BadNvfp4; + if (scale.dtype != .f8_e4m3 or scale.shape.len != 2) return error.BadNvfp4; + const n = weight.shape[0]; + const k = 2 * weight.shape[1]; + if (k % group != 0 or scale.shape[0] != n or scale.shape[1] * group != k) return error.BadNvfp4; + if (!scalar(scale2)) return error.BadNvfp4; + if (input_scale) |s| if (!scalar(s)) return error.BadNvfp4; + return .{ .n = n, .k = k }; +} + +fn scalar(m: Meta) bool { + return m.dtype == .f32 and (m.shape.len == 0 or (m.shape.len == 1 and m.shape[0] == 1)); +} + +/// Whether `module` (a linear's name without ".weight") is stored as NVFP4: its tensors say, not config.json. +pub fn stored(ck: *const ckpt.Checkpoint, module: []const u8) bool { + var buf: [256]u8 = undefined; + const name = std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module}) catch return false; + for (ck.files.items) |*f| if (f.names.contains(name)) return true; + return false; +} + +/// One NVFP4 linear's stored bytes, as the checkpoint maps them. +pub const Linear = struct { + n: usize, + k: usize, + codes: []const u8, // [n, k/2]: input 2j in byte j's low nibble + scales: []const u8, // [n, k/16] e4m3 + global: f32, // weight_scale_2 + + /// `module`'s four tensors, marked as used; input_scale is the W4A4 math's static scale, unread by W4A16. + pub fn read(ck: *ckpt.Checkpoint, module: []const u8) !Linear { + var buf: [256]u8 = undefined; + const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); + const s = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module})); + const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); + const a = try ck.get(try std.fmt.bufPrint(&buf, "{s}.input_scale", .{module})); + const d = dims(meta(&w), meta(&s), meta(&g), meta(&a)) catch |e| { + std.log.err("{s}: not ModelOpt NVFP4 ({t} {any}, {t} {any})", .{ module, w.dtype, w.shape[0..w.rank], s.dtype, s.shape[0..s.rank] }); + return e; + }; + return .{ .n = d.n, .k = d.k, .codes = w.bytes, .scales = s.bytes, .global = std.mem.bytesToValue(f32, g.bytes[0..4]) }; + } + + /// Row r as f32, e2m1 x e4m3 x global in that order: the reference the GPU kernels are compared with. + pub fn row(l: Linear, r: usize, out: []f32) void { + std.debug.assert(r < l.n and out.len == l.k); + const codes = l.codes[r * l.k / 2 ..][0 .. l.k / 2]; + const scales = l.scales[r * l.k / group ..][0 .. l.k / group]; + for (out, 0..) |*o, i| { + const byte = codes[i / 2]; + const code = if (i % 2 == 0) byte & 0xf else byte >> 4; + o.* = e2m1(code) * e4m3(scales[i / group]) * l.global; + } + } +}; + +fn meta(t: *const ckpt.Tensor) Meta { + return .{ .dtype = t.dtype, .shape = t.shape[0..t.rank] }; +} + +test "dims refuse what is not ModelOpt NVFP4" { + const w: Meta = .{ .dtype = .u8, .shape = &.{ 4, 16 } }; + const s: Meta = .{ .dtype = .f8_e4m3, .shape = &.{ 4, 2 } }; + const g: Meta = .{ .dtype = .f32, .shape = &.{} }; + const d = try dims(w, s, g, g); + try std.testing.expectEqual(Dims{ .n = 4, .k = 32 }, d); + try std.testing.expectError(error.BadNvfp4, dims(.{ .dtype = .bf16, .shape = &.{ 4, 16 } }, s, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, .{ .dtype = .f8_e4m3, .shape = &.{ 4, 3 } }, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, .{ .dtype = .u8, .shape = &.{ 4, 2 } }, g, g)); + try std.testing.expectError(error.BadNvfp4, dims(w, s, .{ .dtype = .f32, .shape = &.{2} }, null)); +} + +test "a linear read from a checkpoint folder, row by row" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + // n = 2, k = 32: codes 2 x 16 bytes, scales 2 x 2, two fp32 scalars + const header = + \\{"l.weight":{"dtype":"U8","shape":[2,16],"data_offsets":[0,32]}, + \\ "l.weight_scale":{"dtype":"F8_E4M3","shape":[2,2],"data_offsets":[32,36]}, + \\ "l.weight_scale_2":{"dtype":"F32","shape":[],"data_offsets":[36,40]}, + \\ "l.input_scale":{"dtype":"F32","shape":[],"data_offsets":[40,44]}} + ; + var data: [8 + header.len + 44]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + const body = data[8 + header.len ..]; + for (body[0..32], 0..) |*b, i| b.* = @intCast((i * 7 + 3) & 0xff); // every nibble value appears + @memcpy(body[32..36], &[_]u8{ 0x38, 0x40, 0x30, 0xb8 }); // 1.0, 2.0, 0.5, -1.0 + @memcpy(body[36..40], std.mem.asBytes(&@as(f32, 0.25))); + @memcpy(body[40..44], std.mem.asBytes(&@as(f32, 3.0))); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &data }); + const dir = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer gpa.free(dir); + var ck = try ckpt.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + try std.testing.expect(stored(&ck, "l")); + try std.testing.expect(!stored(&ck, "m")); + const l = try Linear.read(&ck, "l"); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); + try std.testing.expectEqual(@as(f32, 0.25), l.global); + var out: [32]f32 = undefined; + l.row(0, &out); // byte 0 is 3: input 0 is code 3 (1.5), input 1 code 0; inputs 16.. take the second scale (2.0) + try std.testing.expectEqual(@as(f32, 1.5 * 1.0 * 0.25), out[0]); + try std.testing.expectEqual(@as(f32, 0.0), out[1]); + try std.testing.expectEqual(e2m1(body[8] & 0xf) * 2.0 * 0.25, out[16]); + for (0..2) |r| { + l.row(r, &out); + for (out, 0..) |v, i| { + const byte = body[r * 16 + i / 2]; + const code = if (i % 2 == 0) byte & 0xf else byte >> 4; + const scale = e4m3(body[32 + r * 2 + i / 16]); + try std.testing.expectEqual(e2m1(code) * scale * 0.25, v); + } + } +} diff --git a/zig/src/core/nvfp4_fixture_test.zig b/zig/src/core/nvfp4_fixture_test.zig new file mode 100644 index 0000000000..9f4bfa3694 --- /dev/null +++ b/zig/src/core/nvfp4_fixture_test.zig @@ -0,0 +1,69 @@ +//! nvidia/GLM-5.3-Flash-NVFP4's whole header: every linear is NVFP4 or BF16 as nvfp4.dims reads it. + +const std = @import("std"); +const nvfp4 = @import("nvfp4.zig"); +const DType = @import("safetensors.zig").DType; + +const Meta = nvfp4.Meta; + +/// name -> [dtype, shape] from the gzip'd fixture, kept in `arena`. +fn headers(arena: std.mem.Allocator) !std.json.ObjectMap { + var in: std.Io.Reader = .fixed(@embedFile("fixtures/glm53_flash_nvfp4/headers.json.gz")); + var gz: std.compress.flate.Decompress = .init(&in, .gzip, &.{}); + var out: std.Io.Writer.Allocating = .init(arena); + _ = try gz.reader.streamRemaining(&out.writer); + const v = try std.json.parseFromSliceLeaky(std.json.Value, arena, out.written(), .{}); + return v.object; +} + +fn meta(arena: std.mem.Allocator, h: std.json.ObjectMap, name: []const u8) !?Meta { + const v = h.get(name) orelse return null; + const pair = v.array.items; + const dims = pair[1].array.items; + const shape = try arena.alloc(usize, dims.len); + for (dims, shape) |d, *s| s.* = @intCast(d.integer); + return .{ .dtype = DType.parse(pair[0].string) orelse return error.UnsupportedDType, .shape = shape }; +} + +fn at(arena: std.mem.Allocator, h: std.json.ObjectMap, module: []const u8, suffix: []const u8) !?Meta { + return meta(arena, h, try std.fmt.allocPrint(arena, "{s}.{s}", .{ module, suffix })); +} + +test "every NVFP4 linear of the checkpoint reads as ModelOpt stores it, and the MTP layer stays BF16" { + var state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer state.deinit(); + const a = state.allocator(); + const h = try headers(a); + try std.testing.expectEqual(@as(usize, 147661), h.count()); + var quantized: usize = 0; + var experts: usize = 0; + var it = h.iterator(); + while (it.next()) |kv| { + const name = kv.key_ptr.*; + if (std.mem.endsWith(u8, name, ".weight_scale")) { + // no scale without its codes + try std.testing.expect((try at(a, h, name[0 .. name.len - ".weight_scale".len], "weight")) != null); + continue; + } + if (!std.mem.endsWith(u8, name, ".weight")) continue; + const module = name[0 .. name.len - ".weight".len]; + const scale = try at(a, h, module, "weight_scale") orelse continue; + const d = try nvfp4.dims((try meta(a, h, name)).?, scale, (try at(a, h, module, "weight_scale_2")).?, try at(a, h, module, "input_scale")); + quantized += 1; + try std.testing.expect(std.mem.indexOf(u8, module, ".mlp.") != null); + try std.testing.expect(std.mem.indexOf(u8, module, ".layers.45.") == null); + const down = std.mem.endsWith(u8, module, "down_proj"); + if (std.mem.indexOf(u8, module, ".experts.") != null) { + experts += 1; + try std.testing.expectEqual(if (down) nvfp4.Dims{ .n = 4096, .k = 2048 } else nvfp4.Dims{ .n = 2048, .k = 4096 }, d); + } else { + try std.testing.expectEqual(if (down) nvfp4.Dims{ .n = 4096, .k = 12288 } else nvfp4.Dims{ .n = 12288, .k = 4096 }, d); + } + } + try std.testing.expectEqual(@as(usize, 42 * 288 * 3 + 3 * 3), quantized); + try std.testing.expectEqual(@as(usize, 42 * 288 * 3), experts); + // the MTP layer's routed experts are BF16 as stored (a family may re-pack them for drafting) + const mtp = (try meta(a, h, "model.language_model.layers.45.mlp.experts.0.gate_proj.weight")).?; + try std.testing.expectEqual(DType.bf16, mtp.dtype); + try std.testing.expect((try meta(a, h, "model.language_model.layers.45.mlp.experts.0.gate_proj.weight_scale")) == null); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 0c1d43adf5..c5480e9dce 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -19,6 +19,7 @@ pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); pub const act_ref = @import("act_ref.zig"); pub const route_ref = @import("route_ref.zig"); +pub const nvfp4 = @import("nvfp4.zig"); pub const modelopt = @import("modelopt.zig"); test { @@ -40,5 +41,7 @@ test { _ = draft_depth; _ = @import("copy_index.zig"); _ = ids_json; + _ = nvfp4; _ = modelopt; + _ = @import("nvfp4_fixture_test.zig"); } From 798d38fcfa955f084614bac6639725d3fbe34217 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:35:04 +0900 Subject: [PATCH 053/115] core: the ModelOpt NVFP4 linear's test decodes with core/quant's reference, not an f32 row of its own --- zig/src/core/nvfp4.zig | 49 +++++++++++------------------------------- 1 file changed, 12 insertions(+), 37 deletions(-) diff --git a/zig/src/core/nvfp4.zig b/zig/src/core/nvfp4.zig index 9452970477..ed6948a003 100644 --- a/zig/src/core/nvfp4.zig +++ b/zig/src/core/nvfp4.zig @@ -8,16 +8,6 @@ const DType = ckpt.DType; /// Inputs sharing one e4m3 scale. pub const group = quant.nvfp4.group; -/// An e2m1 code's value as f32 (every e2m1 value is one exactly). -fn e2m1(code: u8) f32 { - return @floatCast(quant.nvfp4.e2m1(code)); -} - -/// An e4m3 byte's value as f32 (every e4m3 value is one exactly; 0x7f and 0xff are NaN). -fn e4m3(b: u8) f32 { - return @floatCast(quant.fp8.e4m3(b)); -} - /// A stored tensor's dtype and shape, as a header lists it. pub const Meta = struct { dtype: DType, @@ -71,18 +61,6 @@ pub const Linear = struct { }; return .{ .n = d.n, .k = d.k, .codes = w.bytes, .scales = s.bytes, .global = std.mem.bytesToValue(f32, g.bytes[0..4]) }; } - - /// Row r as f32, e2m1 x e4m3 x global in that order: the reference the GPU kernels are compared with. - pub fn row(l: Linear, r: usize, out: []f32) void { - std.debug.assert(r < l.n and out.len == l.k); - const codes = l.codes[r * l.k / 2 ..][0 .. l.k / 2]; - const scales = l.scales[r * l.k / group ..][0 .. l.k / group]; - for (out, 0..) |*o, i| { - const byte = codes[i / 2]; - const code = if (i % 2 == 0) byte & 0xf else byte >> 4; - o.* = e2m1(code) * e4m3(scales[i / group]) * l.global; - } - } }; fn meta(t: *const ckpt.Tensor) Meta { @@ -101,7 +79,7 @@ test "dims refuse what is not ModelOpt NVFP4" { try std.testing.expectError(error.BadNvfp4, dims(w, s, .{ .dtype = .f32, .shape = &.{2} }, null)); } -test "a linear read from a checkpoint folder, row by row" { +test "a linear read from a checkpoint folder decodes as core/quant's reference" { const gpa = std.testing.allocator; const io = std.testing.io; var tmp = std.testing.tmpDir(.{}); @@ -131,18 +109,15 @@ test "a linear read from a checkpoint folder, row by row" { const l = try Linear.read(&ck, "l"); try std.testing.expectEqual(@as(usize, 0), ck.unused()); try std.testing.expectEqual(@as(f32, 0.25), l.global); - var out: [32]f32 = undefined; - l.row(0, &out); // byte 0 is 3: input 0 is code 3 (1.5), input 1 code 0; inputs 16.. take the second scale (2.0) - try std.testing.expectEqual(@as(f32, 1.5 * 1.0 * 0.25), out[0]); - try std.testing.expectEqual(@as(f32, 0.0), out[1]); - try std.testing.expectEqual(e2m1(body[8] & 0xf) * 2.0 * 0.25, out[16]); - for (0..2) |r| { - l.row(r, &out); - for (out, 0..) |v, i| { - const byte = body[r * 16 + i / 2]; - const code = if (i % 2 == 0) byte & 0xf else byte >> 4; - const scale = e4m3(body[32 + r * 2 + i / 16]); - try std.testing.expectEqual(e2m1(code) * scale * 0.25, v); - } - } + var out: [64]f64 = undefined; + try quant.nvfp4.reference(.{ + .weight = .{ .dtype = .u8, .rank = 2, .shape = .{ l.n, l.k / 2, 1, 1, 1 }, .bytes = l.codes }, + .scale = .{ .dtype = .f8_e4m3, .rank = 2, .shape = .{ l.n, l.k / group, 1, 1, 1 }, .bytes = l.scales }, + .global = &[_]f32{l.global}, + }, &out); + // byte 0 is 3: input 0 is code 3 (1.5), input 1 code 0; inputs 16.. take the second scale (2.0) + try std.testing.expectEqual(@as(f64, 1.5 * 1.0 * 0.25), out[0]); + try std.testing.expectEqual(@as(f64, 0.0), out[1]); + try std.testing.expectEqual(quant.nvfp4.e2m1(body[8] & 0xf) * 2.0 * 0.25, out[16]); + try std.testing.expectEqual(quant.nvfp4.e2m1(body[16 + 15] >> 4) * -1.0 * 0.25, out[32 + 31]); // row 1, scale -1.0 } From d06009dcd7c8101eeb7faefd8ae83e371109eaff Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:41:48 -0700 Subject: [PATCH 054/115] Add lossless NVFP4 compressed scale planes and TP slices --- zig/src/core/nvfp4_csf.md | 36 ++++++++ zig/src/core/nvfp4_csf.zig | 90 ++++++++++++++++++++ zig/src/core/nvfp4_csf_test.zig | 145 ++++++++++++++++++++++++++++++++ zig/src/core/root.zig | 2 + 4 files changed, 273 insertions(+) create mode 100644 zig/src/core/nvfp4_csf.md create mode 100644 zig/src/core/nvfp4_csf.zig create mode 100644 zig/src/core/nvfp4_csf_test.zig diff --git a/zig/src/core/nvfp4_csf.md b/zig/src/core/nvfp4_csf.md new file mode 100644 index 0000000000..02e2cdf40e --- /dev/null +++ b/zig/src/core/nvfp4_csf.md @@ -0,0 +1,36 @@ +# Lossless NVFP4 scale planes + +`nvfp4_csf.Plane` reads `byte-window4-fixed-stream-u24-exceptions/1` storage. +It preserves the original scale bytes, including signed zeros and NaN encodings; +it performs no floating-point conversion or requantization. + +Each 16-row slab contains 16 byte bases, followed by row-major pairs of four-bit +offsets. The even column occupies the low nibble. A decoded byte is its row base +plus its offset. Bases cannot exceed 240. Exceptions are little-endian uint32 +words: the low 24 bits name a row-major position, and the high byte replaces its +value. Positions must be unique, sorted, and inside the logical plane. + +Rows must be a positive multiple of 16, columns a positive multiple of two, and +the plane can contain at most 2^24 bytes. Stream extents must match exactly. +`Plane.init` validates the entire plane before returning a borrowed view. The +caller keeps the streams alive and immutable while that view is used. + +```zig +const plane = try core.nvfp4_csf.Plane.init(rows, columns, fixed_bytes, exception_bytes); +try plane.decode(rank_scales, row_begin, row_count, column_begin, column_count); +``` + +`decode` writes exactly the selected rectangle in row-major order, with no heap +allocation. Both row and column partitions can end inside a storage slab or a +packed byte. Invalid output extents and overlap with either borrowed stream are +rejected before writing. Sorted exceptions are searched per selected row, so a +rank does not need to materialize another rank's scales. + +This is a host format primitive. It does not change checkpoint admission, +projection arithmetic, CUDA kernels, or the runtime capability advertisement. +Callers must validate safetensors dtypes/shapes and select this exact codec before +constructing the view. + +The public format reference is the [byte-window implementation in b12x](https://github.com/local-inference-lab/b12x/blob/2cc7f66aaef76a9c7cbb660a1efb865cf3f26c88/b12x/_lib/quant/nvfp4_csf.py). +Tests use independently generated synthetic source bytes, handwritten slab +examples, every replacement byte, both partition axes, and malformed streams. diff --git a/zig/src/core/nvfp4_csf.zig b/zig/src/core/nvfp4_csf.zig new file mode 100644 index 0000000000..49d6e24319 --- /dev/null +++ b/zig/src/core/nvfp4_csf.zig @@ -0,0 +1,90 @@ +//! Lossless byte-window4 scale planes, with sorted u24-position/u8-value exceptions. + +const std = @import("std"); + +pub const codec = "byte-window4-fixed-stream-u24-exceptions/1"; +pub const Error = error{ BadGeometry, BadFixedStream, BadExceptions, BadSlice, OverlappingOutput }; + +/// A validated, borrowed checkpoint plane; callers keep both streams immutable until decoding finishes. +pub const Plane = struct { + rows: usize, + columns: usize, + fixed: []const u8, + exceptions: []const u8, + + /// Slabs hold 16 row bases followed by those rows' low-even/high-odd packed offsets. + pub fn init(rows: usize, columns: usize, fixed: []const u8, exceptions: []const u8) Error!Plane { + if (rows == 0 or rows % 16 != 0 or columns == 0 or columns % 2 != 0) return error.BadGeometry; + if (columns > 1 << 24 or rows > (1 << 24) / columns) return error.BadGeometry; + const stride = 16 + 8 * columns; + if (fixed.len != rows / 16 * stride) return error.BadFixedStream; + for (0..rows / 16) |slab| { + for (fixed[slab * stride ..][0..16]) |base| if (base > 240) return error.BadFixedStream; + } + if (exceptions.len % 4 != 0) return error.BadExceptions; + var previous: ?u32 = null; + var offset: usize = 0; + while (offset < exceptions.len) : (offset += 4) { + const position = word(exceptions, offset) & 0xffffff; + if (position >= rows * columns) return error.BadExceptions; + if (previous) |p| if (position <= p) return error.BadExceptions; + previous = position; + } + return .{ .rows = rows, .columns = columns, .fixed = fixed, .exceptions = exceptions }; + } + + /// Decode a rectangular TP slice into exactly sized row-major bytes, without allocating or rounding. + pub fn decode(self: Plane, out: []u8, row: usize, rows: usize, column: usize, columns: usize) Error!void { + if (row > self.rows or rows > self.rows - row or column > self.columns or columns > self.columns - column) + return error.BadSlice; + if (out.len != rows * columns) return error.BadSlice; + if (overlaps(out, self.fixed) or overlaps(out, self.exceptions)) return error.OverlappingOutput; + if (rows == 0 or columns == 0) return; + const stride = 16 + 8 * self.columns; + for (0..rows) |r| { + const source_row = row + r; + const slab = source_row / 16 * stride; + const base = self.fixed[slab + source_row % 16]; + const codes = slab + 16 + source_row % 16 * (self.columns / 2); + for (0..columns) |c| { + const source_column = column + c; + const codes_byte = self.fixed[codes + source_column / 2]; + const shift: u3 = if (source_column % 2 == 0) 0 else 4; + out[r * columns + c] = base + ((codes_byte >> shift) & 15); + } + var i = self.lowerBound(source_row * self.columns + column); + const end = source_row * self.columns + column + columns; + while (i < self.exceptions.len) : (i += 4) { + const entry = word(self.exceptions, i); + const position = entry & 0xffffff; + if (position >= end) break; + out[r * columns + position - (source_row * self.columns + column)] = @intCast(entry >> 24); + } + } + } + + fn lowerBound(self: Plane, position: usize) usize { + var lo: usize = 0; + var hi = self.exceptions.len / 4; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (word(self.exceptions, mid * 4) & 0xffffff < position) lo = mid + 1 else hi = mid; + } + return lo * 4; + } +}; + +fn word(bytes: []const u8, offset: usize) u32 { + return std.mem.readInt(u32, bytes[offset..][0..4], .little); +} + +fn overlaps(out: []u8, source: []const u8) bool { + if (out.len == 0 or source.len == 0) return false; + const a = @intFromPtr(out.ptr); + const b = @intFromPtr(source.ptr); + return if (a <= b) b - a < out.len else a - b < source.len; +} + +test { + _ = @import("nvfp4_csf_test.zig"); +} diff --git a/zig/src/core/nvfp4_csf_test.zig b/zig/src/core/nvfp4_csf_test.zig new file mode 100644 index 0000000000..13e63c88b7 --- /dev/null +++ b/zig/src/core/nvfp4_csf_test.zig @@ -0,0 +1,145 @@ +//! Independent byte fixtures and malformed-storage boundaries for lossless scale decoding. + +const std = @import("std"); +const Plane = @import("nvfp4_csf.zig").Plane; +const expect = std.testing.expectEqualSlices; + +test "handwritten slab preserves nibble order, bases and exception bytes" { + var fixed: [32]u8 = @splat(0); + fixed[0] = 240; + fixed[1] = 1; + fixed[16] = 0xf0; + fixed[17] = 0x21; + const exceptions = [_]u8{ 2, 0, 0, 255, 31, 0, 0, 128 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 2); + try expect(u8, &.{ 240, 255, 255, 3 }, out[0..4]); + try std.testing.expectEqual(@as(u8, 128), out[31]); + var odd: [2]u8 = undefined; + try plane.decode(&odd, 0, 2, 1, 1); + try expect(u8, &.{ 255, 3 }, &odd); +} + +const Fixture = struct { + fixed: []u8, + exceptions: []u8, + source: []u8, + rows: usize, + columns: usize, + + fn init(rows: usize, columns: usize) !Fixture { + const gpa = std.testing.allocator; + const source = try gpa.alloc(u8, rows * columns); + errdefer gpa.free(source); + const fixed = try gpa.alloc(u8, rows + source.len / 2); + errdefer gpa.free(fixed); + const exceptions = try gpa.alloc(u8, ((source.len + 96) / 97) * 4); + errdefer gpa.free(exceptions); + @memset(fixed, 0); + var cursor: usize = 0; + for (0..rows) |row| { + const base: u8 = @intCast((row * 37 + 19) % 241); + const slab = row / 16 * (16 + columns * 8); + fixed[slab + row % 16] = base; + for (0..columns) |column| { + const i = row * columns + column; + const delta: u8 = @intCast((column * 11 + row * 3) % 16); + source[i] = base + delta; + const offset = slab + 16 + row % 16 * (columns / 2) + column / 2; + fixed[offset] |= delta << @as(u3, if (column % 2 == 0) 0 else 4); + if (i % 97 == 0) { + source[i] = @intCast((i / 97) % 256); + const record = @as(u32, @intCast(i)) | (@as(u32, source[i]) << 24); + std.mem.writeInt(u32, exceptions[cursor..][0..4], record, .little); + cursor += 4; + } + } + } + return .{ .fixed = fixed, .exceptions = exceptions, .source = source, .rows = rows, .columns = columns }; + } + + fn deinit(self: Fixture) void { + const gpa = std.testing.allocator; + gpa.free(self.fixed); + gpa.free(self.exceptions); + gpa.free(self.source); + } + + fn check(self: Fixture, row: usize, rows: usize, column: usize, columns: usize) !void { + const out = try std.testing.allocator.alloc(u8, rows * columns); + defer std.testing.allocator.free(out); + @memset(out, 0xd6); + const plane = try Plane.init(self.rows, self.columns, self.fixed, self.exceptions); + try plane.decode(out, row, rows, column, columns); + for (0..rows) |r| try expect(u8, self.source[(row + r) * self.columns + column ..][0..columns], out[r * columns ..][0..columns]); + } +}; + +test "full planes and row/column TP partitions reproduce independent source bytes" { + for ([_][2]usize{ .{ 16, 2 }, .{ 32, 20 }, .{ 640, 160 }, .{ 2560, 40 } }) |shape| { + const fixture = try Fixture.init(shape[0], shape[1]); + defer fixture.deinit(); + try fixture.check(0, shape[0], 0, shape[1]); + for (0..2) |rank| { + try fixture.check(rank * shape[0] / 2, shape[0] / 2, 0, shape[1]); + try fixture.check(0, shape[0], rank * shape[1] / 2, shape[1] / 2); + } + try fixture.check(15, 1, 1, shape[1] - 1); + } +} + +test "all 256 byte values survive exception decoding unchanged" { + const fixed: [144]u8 = @splat(0); + var exceptions: [1024]u8 = undefined; + for (0..256) |i| std.mem.writeInt(u32, exceptions[i * 4 ..][0..4], @as(u32, @intCast(i)) * 0x1000001, .little); + const plane = try Plane.init(16, 16, &fixed, &exceptions); + var out: [256]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 16); + for (out, 0..) |value, i| try std.testing.expectEqual(@as(u8, @intCast(i)), value); +} + +test "the last u24 position is valid and larger planes are refused" { + const columns = 1 << 20; + const fixed = try std.testing.allocator.alloc(u8, 16 + 8 * columns); + defer std.testing.allocator.free(fixed); + @memset(fixed, 0); + const plane = try Plane.init(16, columns, fixed, &.{ 255, 255, 255, 173 }); + var out: [1]u8 = undefined; + try plane.decode(&out, 15, 1, columns - 1, 1); + try expect(u8, &.{173}, &out); + try std.testing.expectError(error.BadGeometry, Plane.init(32, columns, fixed, &.{})); +} + +test "bad geometry, truncated streams, overflow bases and unordered exceptions are refused" { + var fixed: [32]u8 = @splat(0); + const bad = std.testing.expectError; + for ([_][2]usize{ .{ 0, 2 }, .{ 1, 2 }, .{ 16, 0 }, .{ 16, 3 }, .{ 16, 1 << 24 }, .{ 1 << 24, 2 } }) |shape| + try bad(error.BadGeometry, Plane.init(shape[0], shape[1], &fixed, &.{})); + try bad(error.BadFixedStream, Plane.init(16, 2, fixed[0..31], &.{})); + fixed[15] = 241; + try bad(error.BadFixedStream, Plane.init(16, 2, &fixed, &.{})); + fixed[15] = 0; + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{0})); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 32, 0, 0, 1 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 1, 0, 0, 2, 1, 0, 0, 3 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 2, 0, 0, 2, 1, 0, 0, 3 })); +} + +test "invalid slices and aliased storage fail before output changes" { + var fixed: [32]u8 = @splat(0); + var exceptions = [_]u8{ 0, 0, 0, 1 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = @splat(0xd6); + const bad = std.testing.expectError; + try bad(error.BadSlice, plane.decode(&out, 17, 0, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 17, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 3, 0)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 1, 2)); + try bad(error.BadSlice, plane.decode(out[0..31], 0, 16, 0, 2)); + const poison: [32]u8 = @splat(0xd6); + try expect(u8, &poison, &out); + try bad(error.OverlappingOutput, plane.decode(&fixed, 0, 16, 0, 2)); + try bad(error.OverlappingOutput, plane.decode(&exceptions, 0, 2, 0, 2)); + try plane.decode(out[0..0], 16, 0, 2, 0); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index c5480e9dce..281977f3f3 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -14,6 +14,7 @@ pub const affine4_host = @import("affine4_host.zig"); pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); +pub const nvfp4_csf = @import("nvfp4_csf.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); @@ -28,6 +29,7 @@ test { _ = affine4_fold; _ = shard_edit; _ = quant; + _ = nvfp4_csf; _ = registry; _ = tuning; _ = norms_ref; From 2868a9a6251ad493052bc4a1e06c0138221f9fec Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:57:42 -0700 Subject: [PATCH 055/115] cuda: refuse tensor parallelism before loading unranked families --- zig/src/native/cuda_refusal.zig | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/zig/src/native/cuda_refusal.zig b/zig/src/native/cuda_refusal.zig index b9105ff59d..0aa8a00c7c 100644 --- a/zig/src/native/cuda_refusal.zig +++ b/zig/src/native/cuda_refusal.zig @@ -12,10 +12,23 @@ pub fn backend(o: api.Open) ?[]const u8 { /// A family takes --drafter only if it says so (`takes_drafter`); null when `o` may open `F`. pub fn family(comptime F: type, o: api.Open) ?[]const u8 { + const ranked = @hasDecl(F, "Options") and @hasField(F.Options, "rank"); + if (o.tp != 1 and !ranked) return "this CUDA family has no tensor-parallel ranks: serve without --tp"; if (o.drafter != null and !(@hasDecl(F, "takes_drafter") and F.takes_drafter)) return "--drafter loads a separate drafter, and this CUDA family drafts with its own MTP head: serve without --drafter"; return null; } +test "a family must consume rank options before tensor parallelism is admitted" { + const Single = struct { + pub const Options = struct { context: usize }; + }; + const Ranked = struct { + pub const Options = struct { context: usize, rank: usize }; + }; + try std.testing.expectStringStartsWith(family(Single, .{ .dir = "d", .model_type = "m", .tp = 2 }) orelse "", "this CUDA family has no tensor-parallel ranks"); + try std.testing.expect(family(Ranked, .{ .dir = "d", .model_type = "m", .tp = 2 }) == null); +} + test "--learn is refused whatever the family, and --slide or a learned-disk size alone is not" { try std.testing.expectStringStartsWith(backend(.{ .dir = "d", .model_type = "nemotron_h", .learn = "/x" }) orelse "", "--learn learns prompt prefixes on Macs only"); try std.testing.expect(backend(.{ .dir = "d", .model_type = "nemotron_h", .slide = true, .learn_gib = 8, .learn_min_free_gib = 2 }) == null); From d631b1582f80e216aa532eb51d5867a0ea689b5c Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:57:50 -0700 Subject: [PATCH 056/115] core: read exact rank-local NVFP4 tensors with plain or compressed scales --- zig/src/core/nvfp4_slice.zig | 105 ++++++++++++++++++++++ zig/src/core/nvfp4_slice_test.zig | 141 ++++++++++++++++++++++++++++++ zig/src/core/root.zig | 2 + 3 files changed, 248 insertions(+) create mode 100644 zig/src/core/nvfp4_slice.zig create mode 100644 zig/src/core/nvfp4_slice_test.zig diff --git a/zig/src/core/nvfp4_slice.zig b/zig/src/core/nvfp4_slice.zig new file mode 100644 index 0000000000..5f3dc0fbc7 --- /dev/null +++ b/zig/src/core/nvfp4_slice.zig @@ -0,0 +1,105 @@ +//! Plain or CSF-compressed ModelOpt scales and packed codes, materialized only for one tensor-parallel rectangle. + +const std = @import("std"); +const ckpt = @import("checkpoint.zig"); +const csf = @import("nvfp4_csf.zig"); +const Tensor = ckpt.Tensor; +const Span = @import("quant/slice.zig").Span; + +pub const Scales = union(enum) { plain: Tensor, csf: struct { fixed: Tensor, exceptions: Tensor } }; +pub const Source = struct { weight: Tensor, scales: Scales, global: Tensor, input: ?Tensor = null }; +pub const Dims = struct { n: usize, k: usize }; + +/// Owned rank-local bytes: low-even FP4 nibbles, row-major E4M3 scales and the unchanged FP32 global scale. +pub const Linear = struct { + n: usize, + k: usize, + codes: []u8, + scales: []u8, + global: f32, + + pub fn deinit(self: *Linear, gpa: std.mem.Allocator) void { + gpa.free(self.codes); + gpa.free(self.scales); + self.* = undefined; + } +}; + +/// Validate a projection's original geometry before allocating rank-local storage. +pub fn dimensions(source: Source) !Dims { + const w = source.weight; + if (w.rank != 2 or w.dtype != .u8 or w.shape[0] == 0 or w.shape[1] == 0) return error.BadNvfp4; + const k = std.math.mul(usize, w.shape[1], 2) catch return error.BadNvfp4; + const count = std.math.mul(usize, w.shape[0], w.shape[1]) catch return error.BadNvfp4; + if (w.bytes.len != count or k % 16 != 0) return error.BadNvfp4; + const d: Dims = .{ .n = w.shape[0], .k = k }; + if (!scalar(source.global)) return error.BadNvfp4; + if (source.input) |a| if (!scalar(a)) return error.BadNvfp4; + switch (source.scales) { + .plain => |s| if (!s.is(.f8_e4m3, &.{ d.n, d.k / 16 }) or s.bytes.len != count / 8) return error.BadNvfp4, + .csf => |s| { + if (s.exceptions.dtype != .u32 or s.exceptions.rank != 1) return error.BadNvfp4; + if (s.exceptions.shape[0] > std.math.maxInt(usize) / 4 or s.exceptions.bytes.len != s.exceptions.shape[0] * 4) + return error.BadNvfp4; + if (!s.fixed.is(.u8, &.{ d.n / 16, 16 + 8 * (d.k / 16) })) return error.BadNvfp4; + _ = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + }, + } + return d; +} + +/// Rows and input channels are logical weight axes; input cuts must preserve complete groups of 16. +pub fn slice(gpa: std.mem.Allocator, source: Source, rows: Span, inputs: Span) !Linear { + const d = try dimensions(source); + if (rows.from >= rows.to or rows.to > d.n or inputs.from >= inputs.to or inputs.to > d.k) + return error.BadNvfp4Slice; + if (inputs.from % 16 != 0 or inputs.to % 16 != 0) return error.BadNvfp4Slice; + const n = rows.to - rows.from; + const k = inputs.to - inputs.from; + const codes = try gpa.alloc(u8, n * (k / 2)); + errdefer gpa.free(codes); + const scales = try gpa.alloc(u8, n * (k / 16)); + errdefer gpa.free(scales); + for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 2) + inputs.from / 2; + @memcpy(codes[row * (k / 2) ..][0 .. k / 2], source.weight.bytes[begin..][0 .. k / 2]); + } + switch (source.scales) { + .plain => |s| for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 16) + inputs.from / 16; + @memcpy(scales[row * (k / 16) ..][0 .. k / 16], s.bytes[begin..][0 .. k / 16]); + }, + .csf => |s| { + const plane = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + try plane.decode(scales, rows.from, n, inputs.from / 16, k / 16); + }, + } + return .{ .n = n, .k = k, .codes = codes, .scales = scales, .global = @bitCast(std.mem.readInt(u32, source.global.bytes[0..4], .little)) }; +} + +/// Read exact stored tensors, consuming optional calibration and refusing ambiguous plain/compressed scales. +pub fn read(ck: *ckpt.Checkpoint, gpa: std.mem.Allocator, module: []const u8, rows: Span, inputs: Span) !Linear { + var buf: [512]u8 = undefined; + const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); + const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); + const input_name = try std.fmt.bufPrint(&buf, "{s}.input_scale", .{module}); + const a: ?Tensor = if (ck.has(input_name)) try ck.get(input_name) else null; + const fixed_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_fixed", .{module}); + const fixed: ?Tensor = if (ck.has(fixed_name)) try ck.get(fixed_name) else null; + const exception_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_exceptions", .{module}); + const exceptions: ?Tensor = if (ck.has(exception_name)) try ck.get(exception_name) else null; + if ((fixed != null) != (exceptions != null)) return error.BadNvfp4; + const plain_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module}); + const plain = ck.has(plain_name); + if (plain == (fixed != null)) return error.BadNvfp4; + const scales: Scales = if (plain) .{ .plain = try ck.get(plain_name) } else .{ .csf = .{ .fixed = fixed.?, .exceptions = exceptions.? } }; + return slice(gpa, .{ .weight = w, .scales = scales, .global = g, .input = a }, rows, inputs); +} + +fn scalar(t: Tensor) bool { + return t.dtype == .f32 and t.bytes.len == 4 and (t.rank == 0 or (t.rank == 1 and t.shape[0] == 1)); +} + +test { + _ = @import("nvfp4_slice_test.zig"); +} diff --git a/zig/src/core/nvfp4_slice_test.zig b/zig/src/core/nvfp4_slice_test.zig new file mode 100644 index 0000000000..1ca0a128df --- /dev/null +++ b/zig/src/core/nvfp4_slice_test.zig @@ -0,0 +1,141 @@ +//! Rank-local NVFP4 bytes against the independent FP64 decoder and real safetensors headers. + +const std = @import("std"); +const reader = @import("nvfp4_slice.zig"); +const st = @import("safetensors.zig"); +const ckpt = @import("checkpoint.zig"); +const quant = @import("quant/quant.zig"); + +const Fixture = struct { + codes: [512]u8 = undefined, + scales: [64]u8 = undefined, + fixed: [48]u8 = @splat(0), + exceptions: [4]u8 = .{ 3, 0, 0, 0x40 }, + global: [4]u8 = undefined, + + fn init() Fixture { + var f: Fixture = .{}; + std.mem.writeInt(u32, &f.global, @bitCast(@as(f32, 0.25)), .little); + for (&f.codes, 0..) |*b, i| b.* = @truncate(i * 29 + 3); + for (0..16) |row| { + const base: u8 = @intCast(0x30 + row % 8); + f.fixed[row] = base; + f.fixed[16 + row * 2] = 0x10; + f.fixed[17 + row * 2] = 0x32; + for (0..4) |column| f.scales[row * 4 + column] = base + @as(u8, @intCast(column)); + } + f.scales[3] = 0x40; + return f; + } + + fn source(f: *const Fixture, compressed: bool) reader.Source { + return .{ + .weight = tensor(.u8, &.{ 16, 32 }, &f.codes), + .scales = if (compressed) .{ .csf = .{ + .fixed = tensor(.u8, &.{ 1, 48 }, &f.fixed), + .exceptions = tensor(.u32, &.{1}, &f.exceptions), + } } else .{ .plain = tensor(.f8_e4m3, &.{ 16, 4 }, &f.scales) }, + .global = tensor(.f32, &.{}, &f.global), + }; + } +}; + +fn tensor(dtype: st.DType, shape: []const usize, bytes: []const u8) st.Tensor { + var out: st.Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(out.shape[0..shape.len], shape); + return out; +} + +fn reference(linear: reader.Linear, out: []f64) !void { + try quant.nvfp4.reference(.{ + .weight = tensor(.u8, &.{ linear.n, linear.k / 2 }, linear.codes), + .scale = tensor(.f8_e4m3, &.{ linear.n, linear.k / 16 }, linear.scales), + .global = &.{linear.global}, + }, out); +} + +test "plain and compressed rows/columns preserve every FP64 decoded value" { + const f = Fixture.init(); + const a = std.testing.allocator; + var whole = try reader.slice(a, f.source(false), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 64 }); + defer whole.deinit(a); + var expected: [1024]f64 = undefined; + try reference(whole, &expected); + for ([_]bool{ false, true }) |compressed| { + for (0..2) |axis| for (0..2) |rank| { + const r0 = if (axis == 0) rank * 8 else 0; + const rn: usize = if (axis == 0) 8 else 16; + const c0 = if (axis == 1) rank * 32 else 0; + const cn: usize = if (axis == 1) 32 else 64; + var part = try reader.slice(a, f.source(compressed), .{ .from = r0, .to = r0 + rn }, .{ .from = c0, .to = c0 + cn }); + defer part.deinit(a); + var out: [512]f64 = undefined; + try reference(part, &out); + for (0..rn) |row| try std.testing.expectEqualSlices(f64, expected[(r0 + row) * 64 + c0 ..][0..cn], out[row * cn ..][0..cn]); + }; + } +} + +test "bad compressed tensor metadata and cuts cannot become a rank projection" { + const f = Fixture.init(); + var source = f.source(true); + const bad = std.testing.expectError; + source.scales.csf.exceptions.dtype = .i32; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(true); + source.scales.csf.fixed.shape[1] -= 1; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(false); + source.global.dtype = .bf16; + try bad(error.BadNvfp4, reader.dimensions(source)); + const a = std.testing.allocator; + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 17 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 8, .to = 8 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 1, .to = 32 })); +} + +test "allocation failures free every rank-local buffer" { + const f = Fixture.init(); + for (0..3) |fail_at| { + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = fail_at }); + const a = failing.allocator(); + if (reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 32 })) |value| { + var owned = value; + owned.deinit(a); + } else |e| try std.testing.expectEqual(error.OutOfMemory, e); + try std.testing.expectEqual(failing.allocated_bytes, failing.freed_bytes); + } +} + +test "checkpoint reading consumes CSF streams and preserves optional calibration" { + const f = Fixture.init(); + const io = std.testing.io; + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"l.weight":{"dtype":"U8","shape":[16,32],"data_offsets":[0,512]}, + \\ "l.weight_scale.nvfp4_csf_fixed":{"dtype":"U8","shape":[1,48],"data_offsets":[512,560]}, + \\ "l.weight_scale.nvfp4_csf_exceptions":{"dtype":"U32","shape":[1],"data_offsets":[560,564]}, + \\ "l.weight_scale_2":{"dtype":"F32","shape":[],"data_offsets":[564,568]}, + \\ "l.input_scale":{"dtype":"F32","shape":[1],"data_offsets":[568,572]}} + ; + var image: [8 + header.len + 572]u8 = undefined; + std.mem.writeInt(u64, image[0..8], header.len, .little); + @memcpy(image[8..][0..header.len], header); + const data = image[8 + header.len ..]; + @memcpy(data[0..512], &f.codes); + @memcpy(data[512..560], &f.fixed); + @memcpy(data[560..564], &f.exceptions); + @memcpy(data[564..568], &f.global); + @memcpy(data[568..572], &f.global); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &image }); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + var ck = try ckpt.Checkpoint.openModel(a, io, dir); + defer ck.close(); + var part = try reader.read(&ck, a, "l", .{ .from = 8, .to = 16 }, .{ .from = 32, .to = 64 }); + defer part.deinit(a); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); + for (0..8) |r| try std.testing.expectEqualSlices(u8, f.scales[(8 + r) * 4 + 2 ..][0..2], part.scales[r * 2 ..][0..2]); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 281977f3f3..3c125aa71f 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -15,6 +15,7 @@ pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); pub const nvfp4_csf = @import("nvfp4_csf.zig"); +pub const nvfp4_slice = @import("nvfp4_slice.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); @@ -30,6 +31,7 @@ test { _ = shard_edit; _ = quant; _ = nvfp4_csf; + _ = nvfp4_slice; _ = registry; _ = tuning; _ = norms_ref; From 0a0bb8634af8bdd6fc66c8d45ab00b6b00004dc1 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:57:50 -0700 Subject: [PATCH 057/115] core: read exact rank-local NVFP4 tensors with plain or compressed scales --- zig/src/core/nvfp4_slice.md | 29 ++++++ zig/src/core/nvfp4_slice.zig | 105 ++++++++++++++++++++++ zig/src/core/nvfp4_slice_test.zig | 141 ++++++++++++++++++++++++++++++ zig/src/core/root.zig | 2 + 4 files changed, 277 insertions(+) create mode 100644 zig/src/core/nvfp4_slice.md create mode 100644 zig/src/core/nvfp4_slice.zig create mode 100644 zig/src/core/nvfp4_slice_test.zig diff --git a/zig/src/core/nvfp4_slice.md b/zig/src/core/nvfp4_slice.md new file mode 100644 index 0000000000..b3d4355835 --- /dev/null +++ b/zig/src/core/nvfp4_slice.md @@ -0,0 +1,29 @@ +# Rank-local NVFP4 checkpoint reads + +`nvfp4_slice.read(checkpoint, allocator, module, rows, inputs)` reads a logical +rectangle without expanding the whole projection. The result owns its packed +FP4 code bytes and row-major E4M3 scale bytes; call `deinit` with the same +allocator. `global` preserves the stored FP32 tensor scale. There is no numeric +conversion or requantization. + +The source weight is U8 `[N, K/2]`. Plain scales are F8_E4M3 `[N, K/16]`. +Alternatively, the source may contain the fixed U8 and U32 exception tensors +described by [the CSF codec](nvfp4_csf.md). Both representations cannot be present +at once, and a partial compressed pair is refused. Tensor scale and optional +input calibration must be FP32 scalars. Calibration is consumed as checkpoint +metadata; this reader does not implement an activation quantizer. + +Row and input spans are half-open, nonempty and bounded by the original tensor. +Input cuts must align to groups of 16. Row cuts can cross a compressed 16-row +slab. Source geometry, dtype, byte extents and complete compressed streams are +validated before allocation. Allocation failures release every intermediate +buffer. The result contains only the requested code and scale rectangle. + +`dimensions(source)` and `slice(allocator, source, rows, inputs)` also accept +borrowed tensor views for callers with their own checkpoint mapping. The caller +owns those mappings; the returned buffers are independent of them. + +Tests compare plain and compressed row/column partitions with the independent +FP64 decoder, exercise allocator failures and malformed inputs, and read a real +synthetic safetensors file through the checkpoint reader. These are host format +checks, not tensor-parallel execution or GPU qualification. diff --git a/zig/src/core/nvfp4_slice.zig b/zig/src/core/nvfp4_slice.zig new file mode 100644 index 0000000000..5f3dc0fbc7 --- /dev/null +++ b/zig/src/core/nvfp4_slice.zig @@ -0,0 +1,105 @@ +//! Plain or CSF-compressed ModelOpt scales and packed codes, materialized only for one tensor-parallel rectangle. + +const std = @import("std"); +const ckpt = @import("checkpoint.zig"); +const csf = @import("nvfp4_csf.zig"); +const Tensor = ckpt.Tensor; +const Span = @import("quant/slice.zig").Span; + +pub const Scales = union(enum) { plain: Tensor, csf: struct { fixed: Tensor, exceptions: Tensor } }; +pub const Source = struct { weight: Tensor, scales: Scales, global: Tensor, input: ?Tensor = null }; +pub const Dims = struct { n: usize, k: usize }; + +/// Owned rank-local bytes: low-even FP4 nibbles, row-major E4M3 scales and the unchanged FP32 global scale. +pub const Linear = struct { + n: usize, + k: usize, + codes: []u8, + scales: []u8, + global: f32, + + pub fn deinit(self: *Linear, gpa: std.mem.Allocator) void { + gpa.free(self.codes); + gpa.free(self.scales); + self.* = undefined; + } +}; + +/// Validate a projection's original geometry before allocating rank-local storage. +pub fn dimensions(source: Source) !Dims { + const w = source.weight; + if (w.rank != 2 or w.dtype != .u8 or w.shape[0] == 0 or w.shape[1] == 0) return error.BadNvfp4; + const k = std.math.mul(usize, w.shape[1], 2) catch return error.BadNvfp4; + const count = std.math.mul(usize, w.shape[0], w.shape[1]) catch return error.BadNvfp4; + if (w.bytes.len != count or k % 16 != 0) return error.BadNvfp4; + const d: Dims = .{ .n = w.shape[0], .k = k }; + if (!scalar(source.global)) return error.BadNvfp4; + if (source.input) |a| if (!scalar(a)) return error.BadNvfp4; + switch (source.scales) { + .plain => |s| if (!s.is(.f8_e4m3, &.{ d.n, d.k / 16 }) or s.bytes.len != count / 8) return error.BadNvfp4, + .csf => |s| { + if (s.exceptions.dtype != .u32 or s.exceptions.rank != 1) return error.BadNvfp4; + if (s.exceptions.shape[0] > std.math.maxInt(usize) / 4 or s.exceptions.bytes.len != s.exceptions.shape[0] * 4) + return error.BadNvfp4; + if (!s.fixed.is(.u8, &.{ d.n / 16, 16 + 8 * (d.k / 16) })) return error.BadNvfp4; + _ = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + }, + } + return d; +} + +/// Rows and input channels are logical weight axes; input cuts must preserve complete groups of 16. +pub fn slice(gpa: std.mem.Allocator, source: Source, rows: Span, inputs: Span) !Linear { + const d = try dimensions(source); + if (rows.from >= rows.to or rows.to > d.n or inputs.from >= inputs.to or inputs.to > d.k) + return error.BadNvfp4Slice; + if (inputs.from % 16 != 0 or inputs.to % 16 != 0) return error.BadNvfp4Slice; + const n = rows.to - rows.from; + const k = inputs.to - inputs.from; + const codes = try gpa.alloc(u8, n * (k / 2)); + errdefer gpa.free(codes); + const scales = try gpa.alloc(u8, n * (k / 16)); + errdefer gpa.free(scales); + for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 2) + inputs.from / 2; + @memcpy(codes[row * (k / 2) ..][0 .. k / 2], source.weight.bytes[begin..][0 .. k / 2]); + } + switch (source.scales) { + .plain => |s| for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 16) + inputs.from / 16; + @memcpy(scales[row * (k / 16) ..][0 .. k / 16], s.bytes[begin..][0 .. k / 16]); + }, + .csf => |s| { + const plane = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + try plane.decode(scales, rows.from, n, inputs.from / 16, k / 16); + }, + } + return .{ .n = n, .k = k, .codes = codes, .scales = scales, .global = @bitCast(std.mem.readInt(u32, source.global.bytes[0..4], .little)) }; +} + +/// Read exact stored tensors, consuming optional calibration and refusing ambiguous plain/compressed scales. +pub fn read(ck: *ckpt.Checkpoint, gpa: std.mem.Allocator, module: []const u8, rows: Span, inputs: Span) !Linear { + var buf: [512]u8 = undefined; + const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); + const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); + const input_name = try std.fmt.bufPrint(&buf, "{s}.input_scale", .{module}); + const a: ?Tensor = if (ck.has(input_name)) try ck.get(input_name) else null; + const fixed_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_fixed", .{module}); + const fixed: ?Tensor = if (ck.has(fixed_name)) try ck.get(fixed_name) else null; + const exception_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_exceptions", .{module}); + const exceptions: ?Tensor = if (ck.has(exception_name)) try ck.get(exception_name) else null; + if ((fixed != null) != (exceptions != null)) return error.BadNvfp4; + const plain_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module}); + const plain = ck.has(plain_name); + if (plain == (fixed != null)) return error.BadNvfp4; + const scales: Scales = if (plain) .{ .plain = try ck.get(plain_name) } else .{ .csf = .{ .fixed = fixed.?, .exceptions = exceptions.? } }; + return slice(gpa, .{ .weight = w, .scales = scales, .global = g, .input = a }, rows, inputs); +} + +fn scalar(t: Tensor) bool { + return t.dtype == .f32 and t.bytes.len == 4 and (t.rank == 0 or (t.rank == 1 and t.shape[0] == 1)); +} + +test { + _ = @import("nvfp4_slice_test.zig"); +} diff --git a/zig/src/core/nvfp4_slice_test.zig b/zig/src/core/nvfp4_slice_test.zig new file mode 100644 index 0000000000..1ca0a128df --- /dev/null +++ b/zig/src/core/nvfp4_slice_test.zig @@ -0,0 +1,141 @@ +//! Rank-local NVFP4 bytes against the independent FP64 decoder and real safetensors headers. + +const std = @import("std"); +const reader = @import("nvfp4_slice.zig"); +const st = @import("safetensors.zig"); +const ckpt = @import("checkpoint.zig"); +const quant = @import("quant/quant.zig"); + +const Fixture = struct { + codes: [512]u8 = undefined, + scales: [64]u8 = undefined, + fixed: [48]u8 = @splat(0), + exceptions: [4]u8 = .{ 3, 0, 0, 0x40 }, + global: [4]u8 = undefined, + + fn init() Fixture { + var f: Fixture = .{}; + std.mem.writeInt(u32, &f.global, @bitCast(@as(f32, 0.25)), .little); + for (&f.codes, 0..) |*b, i| b.* = @truncate(i * 29 + 3); + for (0..16) |row| { + const base: u8 = @intCast(0x30 + row % 8); + f.fixed[row] = base; + f.fixed[16 + row * 2] = 0x10; + f.fixed[17 + row * 2] = 0x32; + for (0..4) |column| f.scales[row * 4 + column] = base + @as(u8, @intCast(column)); + } + f.scales[3] = 0x40; + return f; + } + + fn source(f: *const Fixture, compressed: bool) reader.Source { + return .{ + .weight = tensor(.u8, &.{ 16, 32 }, &f.codes), + .scales = if (compressed) .{ .csf = .{ + .fixed = tensor(.u8, &.{ 1, 48 }, &f.fixed), + .exceptions = tensor(.u32, &.{1}, &f.exceptions), + } } else .{ .plain = tensor(.f8_e4m3, &.{ 16, 4 }, &f.scales) }, + .global = tensor(.f32, &.{}, &f.global), + }; + } +}; + +fn tensor(dtype: st.DType, shape: []const usize, bytes: []const u8) st.Tensor { + var out: st.Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(out.shape[0..shape.len], shape); + return out; +} + +fn reference(linear: reader.Linear, out: []f64) !void { + try quant.nvfp4.reference(.{ + .weight = tensor(.u8, &.{ linear.n, linear.k / 2 }, linear.codes), + .scale = tensor(.f8_e4m3, &.{ linear.n, linear.k / 16 }, linear.scales), + .global = &.{linear.global}, + }, out); +} + +test "plain and compressed rows/columns preserve every FP64 decoded value" { + const f = Fixture.init(); + const a = std.testing.allocator; + var whole = try reader.slice(a, f.source(false), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 64 }); + defer whole.deinit(a); + var expected: [1024]f64 = undefined; + try reference(whole, &expected); + for ([_]bool{ false, true }) |compressed| { + for (0..2) |axis| for (0..2) |rank| { + const r0 = if (axis == 0) rank * 8 else 0; + const rn: usize = if (axis == 0) 8 else 16; + const c0 = if (axis == 1) rank * 32 else 0; + const cn: usize = if (axis == 1) 32 else 64; + var part = try reader.slice(a, f.source(compressed), .{ .from = r0, .to = r0 + rn }, .{ .from = c0, .to = c0 + cn }); + defer part.deinit(a); + var out: [512]f64 = undefined; + try reference(part, &out); + for (0..rn) |row| try std.testing.expectEqualSlices(f64, expected[(r0 + row) * 64 + c0 ..][0..cn], out[row * cn ..][0..cn]); + }; + } +} + +test "bad compressed tensor metadata and cuts cannot become a rank projection" { + const f = Fixture.init(); + var source = f.source(true); + const bad = std.testing.expectError; + source.scales.csf.exceptions.dtype = .i32; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(true); + source.scales.csf.fixed.shape[1] -= 1; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(false); + source.global.dtype = .bf16; + try bad(error.BadNvfp4, reader.dimensions(source)); + const a = std.testing.allocator; + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 17 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 8, .to = 8 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 1, .to = 32 })); +} + +test "allocation failures free every rank-local buffer" { + const f = Fixture.init(); + for (0..3) |fail_at| { + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = fail_at }); + const a = failing.allocator(); + if (reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 32 })) |value| { + var owned = value; + owned.deinit(a); + } else |e| try std.testing.expectEqual(error.OutOfMemory, e); + try std.testing.expectEqual(failing.allocated_bytes, failing.freed_bytes); + } +} + +test "checkpoint reading consumes CSF streams and preserves optional calibration" { + const f = Fixture.init(); + const io = std.testing.io; + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"l.weight":{"dtype":"U8","shape":[16,32],"data_offsets":[0,512]}, + \\ "l.weight_scale.nvfp4_csf_fixed":{"dtype":"U8","shape":[1,48],"data_offsets":[512,560]}, + \\ "l.weight_scale.nvfp4_csf_exceptions":{"dtype":"U32","shape":[1],"data_offsets":[560,564]}, + \\ "l.weight_scale_2":{"dtype":"F32","shape":[],"data_offsets":[564,568]}, + \\ "l.input_scale":{"dtype":"F32","shape":[1],"data_offsets":[568,572]}} + ; + var image: [8 + header.len + 572]u8 = undefined; + std.mem.writeInt(u64, image[0..8], header.len, .little); + @memcpy(image[8..][0..header.len], header); + const data = image[8 + header.len ..]; + @memcpy(data[0..512], &f.codes); + @memcpy(data[512..560], &f.fixed); + @memcpy(data[560..564], &f.exceptions); + @memcpy(data[564..568], &f.global); + @memcpy(data[568..572], &f.global); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &image }); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + var ck = try ckpt.Checkpoint.openModel(a, io, dir); + defer ck.close(); + var part = try reader.read(&ck, a, "l", .{ .from = 8, .to = 16 }, .{ .from = 32, .to = 64 }); + defer part.deinit(a); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); + for (0..8) |r| try std.testing.expectEqualSlices(u8, f.scales[(8 + r) * 4 + 2 ..][0..2], part.scales[r * 2 ..][0..2]); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 4ffd5c0afb..9febe4107b 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -15,6 +15,7 @@ pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); pub const nvfp4_csf = @import("nvfp4_csf.zig"); +pub const nvfp4_slice = @import("nvfp4_slice.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); @@ -27,6 +28,7 @@ test { _ = shard_edit; _ = quant; _ = nvfp4_csf; + _ = nvfp4_slice; _ = registry; _ = tuning; _ = norms_ref; From f1f7f9abd0ebc20929cea20d1817426183584ba9 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:08:12 -0700 Subject: [PATCH 058/115] cuda: serve losslessly packed MXFP8 projections through native qlinear --- zig/kernels/cuda/qmmf.cu | 17 +++- zig/src/cuda/mxfp8.md | 27 +++++++ zig/src/cuda/mxfp8.zig | 41 ++++++++++ zig/src/cuda/qlinear.zig | 9 ++- zig/src/cuda/qmmf.zig | 14 ++-- zig/src/cuda/root.zig | 2 + zig/tests/cuda/main.zig | 3 + zig/tests/cuda/mxfp8_tests.zig | 143 +++++++++++++++++++++++++++++++++ 8 files changed, 247 insertions(+), 9 deletions(-) create mode 100644 zig/src/cuda/mxfp8.md create mode 100644 zig/src/cuda/mxfp8.zig create mode 100644 zig/tests/cuda/mxfp8_tests.zig diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 4703273410..027b8ef630 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -32,6 +32,11 @@ __device__ __forceinline__ uint32_t fp8pair(uint32_t w) { return r; } +// UE8M0 has no zero encoding: byte 0 is 2^-127, byte 255 is NaN. +__device__ __forceinline__ float e8m0f(uint8_t b) { + return __uint_as_float(b == 0 ? 0x00400000u : b == 255 ? 0x7FC00000u : uint32_t(b) << 23); +} + __device__ __forceinline__ uint32_t comp(const uint4& v, int c) { return c == 0 ? v.x : c == 1 ? v.y : c == 2 ? v.z : v.w; } __device__ __forceinline__ float e4m3f(uint8_t b) { @@ -193,8 +198,8 @@ __global__ void __launch_bounds__(WM * WN * 32) qmmf_kernel( s0 = e4m3f(static_cast(sq[j][0] >> (8 * blk))); s1 = e4m3f(static_cast(sq[j][1] >> (8 * blk))); } else { - s0 = __int_as_float(static_cast((sq[j][0] >> (8 * blk)) & 0xFFu) << 23); - s1 = __int_as_float(static_cast((sq[j][1] >> (8 * blk)) & 0xFFu) << 23); + s0 = e8m0f(static_cast(sq[j][0] >> (8 * blk))); + s1 = e8m0f(static_cast(sq[j][1] >> (8 * blk))); } #pragma unroll for (int i = 0; i < T::MT; ++i) @@ -293,6 +298,14 @@ __global__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); } +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); + // The instances the Zig launchers load by name (the Python host instantiates these through its launch switch). template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); diff --git a/zig/src/cuda/mxfp8.md b/zig/src/cuda/mxfp8.md new file mode 100644 index 0000000000..5878df43eb --- /dev/null +++ b/zig/src/cuda/mxfp8.md @@ -0,0 +1,27 @@ +# Native MXFP8 projections + +`cuda.mxfp8` packs row-major E4M3 code bytes and one UE8M0 scale byte per +32 input channels for `qlinear.Weight.mxfp8`. Both prompt and decode calls use +the lane matmul, with BF16 inputs, FP32 accumulation and BF16 output. Packing +permutes bytes; it does not requantize weights or expand them to BF16. + +The logical shape is `[N, K]`, with positive dimensions and K divisible by 64. +`packCodes` writes `padded(N) * K` bytes; `packScales` writes +`padded(N) * K/32` bytes. Rows are padded to 128; padded codes are zero and +padded scales are one. `weight` returns the borrowed device view. The caller +owns its device buffers and the lane's split-K scratch. + +UE8M0 bytes 0 through 254 mean `2^(byte - 127)`; byte 255 is NaN. In +particular, byte zero is a nonzero subnormal FP32 scale. See the +[CUDA datatype contract](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-math-api/cuda_math_api/struct____nv__fp8__e8m0.html). + +`tf-cuda-test mxfp8` checks all 256 scale encodings through the production +matmul, then compares synthetic projections with an independent FP64 +calculation. It also checks strided activations, partial output tiles, +row counts 1/7/17/33/256, prompt/decode dispatch, and cluster versus explicit +split reduction. Host tests cover scale permutation, padding, shared-memory +sizes, symbols and format refusals. + +This primitive does not admit new checkpoint families or certify model +correctness, tensor parallelism or performance. Rebuild the qmmf fatbin with +the native runtime; previous fatbins lack the MXFP8 instantiations. diff --git a/zig/src/cuda/mxfp8.zig b/zig/src/cuda/mxfp8.zig new file mode 100644 index 0000000000..7c57aa50df --- /dev/null +++ b/zig/src/cuda/mxfp8.zig @@ -0,0 +1,41 @@ +//! E4M3 codes with one UE8M0 byte per 32 inputs, repacked losslessly for the lane matmul. +const std = @import("std"); +const fp8 = @import("fp8.zig"); +const qmmf = @import("qmmf.zig"); + +pub const group = 32; +pub const padded = fp8.padded; +pub const packCodes = fp8.packCodes; + +/// Row-major [n, k/32] bytes to [npad/64][k/64][64][2]; padded rows have scale 1. +pub fn packScales(out: []u8, scales: []const u8, n: usize, k: usize) void { + const npad = padded(n); + const kg = k / qmmf.group; + std.debug.assert(k % qmmf.group == 0 and out.len == npad * (k / group) and scales.len == n * (k / group)); + for (0..npad / 64) |t| for (0..kg) |g| for (0..64) |r| for (0..2) |h| { + const row = t * 64 + r; + out[((t * kg + g) * 64 + r) * 2 + h] = if (row < n) scales[row * (k / group) + g * 2 + h] else 127; + }; +} + +pub fn weight(codes: u64, scales: u64, n: usize, k: usize) qmmf.Weight { + return .{ .mode = .mxfp8, .codes = codes, .scales = scales, .n = @intCast(n), .k = @intCast(k), .npad = @intCast(padded(n)) }; +} + +test "every scale byte lands in its independent 32-input group, including padded rows" { + const a = std.testing.allocator; + const n = 130; + const k = 256; + const scales = try a.alloc(u8, n * (k / group)); + defer a.free(scales); + for (scales, 0..) |*v, i| v.* = @truncate(i * 13 + 7); + const out = try a.alloc(u8, padded(n) * (k / group)); + defer a.free(out); + packScales(out, scales, n, k); + for (0..padded(n)) |row| for (0..k / group) |g| { + const at = (((row / 64) * (k / 64) + g / 2) * 64 + row % 64) * 2 + g % 2; + try std.testing.expectEqual(if (row < n) scales[row * (k / group) + g] else @as(u8, 127), out[at]); + }; + const w = weight(1, 2, n, k); + try std.testing.expect(w.mode == .mxfp8 and w.npad == 256 and w.scale == 1.0); +} diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..b924132f79 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -33,6 +33,7 @@ pub const Affine4 = struct { pub const Weight = union(enum) { affine4: Affine4, fp8g: qmmf.Weight, + mxfp8: qmmf.Weight, // e4m3 with UE8M0 scale bytes per 32 inputs nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values @@ -40,14 +41,14 @@ pub const Weight = union(enum) { pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, - .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .fp8g, .mxfp8, .nvfp4, .fp8, .bf16 => |q| q.n, }; } pub fn inputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.k, - .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .fp8g, .mxfp8, .nvfp4, .fp8, .bf16 => |q| q.k, }; } }; @@ -70,6 +71,7 @@ pub const Linear = struct { return a.decode(s, in.x, in.sums, q, out, rows); }, .fp8g => |q| return l.onLane(s, q, .fp8g, in, out, rows), + .mxfp8 => |q| return l.onLane(s, q, .mxfp8, in, out, rows), .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), @@ -85,6 +87,7 @@ pub const Linear = struct { return a.prompt(s, in.x, q, out, rows); }, .fp8g => |q| return l.onLane(s, q, .fp8g, in, out, rows), + .mxfp8 => |q| return l.onLane(s, q, .mxfp8, in, out, rows), .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), @@ -330,5 +333,7 @@ test "a lane weight whose mode is not its format's is refused before any launch" const fp4: qmmf.Weight = .{ .mode = .fp4, .codes = 0, .scales = 0, .n = 64, .k = 64, .npad = 128 }; const s: Stream = undefined; try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .fp8g = fp4 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .mxfp8 = fp4 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatMismatch, l.prompt(s, .{ .mxfp8 = fp4 }, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, .{ .nvfp4 = fp4 }, .{ .x = 0 }, 0, 1)); } diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 054e2c34bb..3e2b4168a7 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -1,4 +1,4 @@ -//! The lane matmul (qmmf.cu) for FP4, FP8, FP8G and BF16 weights, with its kernel, tile and slices. +//! The lane matmul (qmmf.cu) for FP4, FP8, MXFP8, FP8G and BF16 weights, with its kernel, tile and slices. const std = @import("std"); const driver = @import("driver.zig"); const module = @import("module.zig"); @@ -17,9 +17,9 @@ const bn = 64; const threads = 128; // WM 1 x WN 4 warps /// The packed formats (qmmf.cu Mode): FP4 with e4m3 scales a 16, e4m3 with one scale or fp32 a 64, and bf16. -pub const Mode = enum(u3) { fp4 = 0, fp8 = 1, fp8g = 3, bf16 = 4 }; +pub const Mode = enum(u3) { fp4 = 0, fp8 = 1, mxfp8 = 2, fp8g = 3, bf16 = 4 }; -const modes = [_]Mode{ .fp4, .fp8, .fp8g, .bf16 }; +const modes = [_]Mode{ .fp4, .fp8, .mxfp8, .fp8g, .bf16 }; fn modeIndex(mode: Mode) usize { return std.mem.indexOfScalar(Mode, &modes, mode).?; @@ -29,10 +29,10 @@ fn modeIndex(mode: Mode) usize { pub fn sharedBytes(mode: Mode, bm: u32) u32 { const w: u32 = switch (mode) { .fp4 => bn * group / 2, - .fp8, .fp8g => bn * group, + .fp8, .mxfp8, .fp8g => bn * group, .bf16 => bn * group * 2, }; - const s: u32 = if (mode == .fp4 or mode == .fp8g) bn * 4 else 0; + const s: u32 = if (mode == .fp4 or mode == .fp8g) bn * 4 else if (mode == .mxfp8) bn * 2 else 0; const stage = (bm * group * 2 + w + s + 127) / 128 * 128; const partials = (bm / 16) * 2 * 4 * threads * 4; return @max(4 * stage, partials); @@ -186,6 +186,9 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { try std.testing.expectEqual(@as(u32, 25600), sharedBytes(.fp8g, 16)); try std.testing.expectEqual(@as(u32, 33792), sharedBytes(.fp8g, 32)); try std.testing.expectEqual(@as(u32, 50176), sharedBytes(.fp8g, 64)); + try std.testing.expectEqual(@as(u32, 25088), sharedBytes(.mxfp8, 16)); + try std.testing.expectEqual(@as(u32, 33280), sharedBytes(.mxfp8, 32)); + try std.testing.expectEqual(@as(u32, 49664), sharedBytes(.mxfp8, 64)); try std.testing.expectEqual(@as(u32, 17408), sharedBytes(.fp4, 16)); try std.testing.expectEqual(@as(u32, 25600), sharedBytes(.fp4, 32)); try std.testing.expectEqual(@as(u32, 41984), sharedBytes(.fp4, 64)); @@ -209,6 +212,7 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.mxfp8, 16, true, false)); try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false)); try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true)); } diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..ada7a002ce 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -26,6 +26,7 @@ pub const aot = @import("aot.zig"); pub const kernels = @import("kernels.zig"); pub const qmmf = @import("qmmf.zig"); pub const fp8 = @import("fp8.zig"); +pub const mxfp8 = @import("mxfp8.zig"); pub const exlr8 = @import("exlr8.zig"); pub const nvfp4 = @import("nvfp4.zig"); pub const segments = @import("segments.zig"); @@ -51,6 +52,7 @@ test { _ = segments; _ = @import("qmmf.zig"); _ = @import("fp8.zig"); + _ = @import("mxfp8.zig"); _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..2d6b6bd85f 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -16,6 +16,7 @@ const prompt_cuts = @import("prompt_cuts.zig"); const chunk_costs = @import("chunk_costs.zig"); const grouped_tests = @import("grouped_tests.zig"); const fp8_tests = @import("fp8_tests.zig"); +const mxfp8_tests = @import("mxfp8_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); @@ -61,6 +62,7 @@ const usage = \\ fp8-lane

block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes @@ -122,6 +124,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); diff --git a/zig/tests/cuda/mxfp8_tests.zig b/zig/tests/cuda/mxfp8_tests.zig new file mode 100644 index 0000000000..5417d3d1fd --- /dev/null +++ b/zig/tests/cuda/mxfp8_tests.zig @@ -0,0 +1,143 @@ +//! Native MXFP8: lossless packing, independent FP64 products, strided rows and row-count invariance. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +fn bf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +fn tobf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +pub fn run(gpu: check.Gpu) !void { + try scaleRange(gpu); + const a = gpu.gpa; + const n = 70; // a partial tile and padded rows + const k = 1024; // split-K, both cluster and explicit reduction + const m = 256; // fused prompt reduction + const ldx = k + 64; + const npad = cuda.mxfp8.padded(n); + var prng = std.Random.DefaultPrng.init(0x8e80_5eed); + const rng = prng.random(); + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + for (codes) |*v| { + v.* = rng.int(u8); + if (v.* & 0x7f == 0x7f) v.* ^= 1; + } + const scales = try a.alloc(u8, n * (k / 32)); + defer a.free(scales); + for (scales) |*v| v.* = 121 + rng.uintLessThan(u8, 8); + // A dedicated row makes UE8M0 byte 0 observable in a normal BF16 result. + @memset(codes[0..k], 0x38); // E4M3 1 + @memset(scales[0 .. k / 32], 0); + const packed_codes = try a.alloc(u8, npad * k); + defer a.free(packed_codes); + cuda.mxfp8.packCodes(packed_codes, codes, n, k); + const packed_scales = try a.alloc(u8, npad * (k / 32)); + defer a.free(packed_scales); + cuda.mxfp8.packScales(packed_scales, scales, n, k); + var dw = try cuda.DeviceBuffer.fromHost(gpu.d, packed_codes); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(gpu.d, packed_scales); + defer ds.free(); + const w = cuda.mxfp8.weight(dw.ptr, ds.ptr, n, k); + const x = try a.alloc(u16, m * ldx); + defer a.free(x); + @memset(x, 0x7fc0); // stride padding is poison, not another activation + for (0..m) |row| { + for (x[row * ldx ..][0..k]) |*v| v.* = tobf(rng.float(f32) * 2); + } + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(x)); + defer dx.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 2); + defer dy.free(); + var part = try cuda.DeviceBuffer.alloc(gpu.d, 8 * m * n * 4); + defer part.free(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var baseline: ?[]u8 = null; + defer if (baseline) |b| a.free(b); + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters and (try gpu.ctx.capability()) >= 90; + const lin: cuda.qlinear.Linear = .{ .lane = &lane }; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + for ([_]bool{ false, true }) |prompt| { + try dy.fill8(0xff, s.handle); + const in: cuda.qlinear.Rows = .{ .x = dx.ptr, .ldx = ldx, .part = part.ptr }; + if (prompt) try lin.prompt(s, .{ .mxfp8 = w }, in, dy.ptr, rows) else try lin.decode(s, .{ .mxfp8 = w }, in, dy.ptr, rows); + try s.synchronize(); + const got = try a.alloc(u8, rows * n * 2); + defer a.free(got); + try dy.download(0, got); + if (baseline) |b| { + try check.sameBytes("MXFP8 row-count/dispatch/cluster invariance", got, b[0..got.len]); + } else baseline = try a.dupe(u8, got); + for (0..rows) |row| for (0..n) |col| { + var ref: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const scale = std.math.ldexp(@as(f64, 1), @as(i32, scales[col * (k / 32) + i / 32]) - 127); + const v = bf(x[row * ldx + i]) * core.quant.fp8.e4m3(codes[col * k + i]) * scale; + ref += v; + magnitude += @abs(v); + } + const actual = bf(std.mem.readInt(u16, got[(row * n + col) * 2 ..][0..2], .little)); + const tolerance = @abs(ref) / 128 + magnitude / 4096 + 0x1p-133; + try check.expect(@abs(actual - ref) <= tolerance, "MXFP8 FP64 row {d}, column {d}: {e} versus {e}, tolerance {e}", .{ row, col, actual, ref, tolerance }); + }; + } + } + } + check.pass("MXFP8: strided rows 1,7,17,33,256 equal across dispatch and reduction paths; FP64 including UE8M0 byte zero", .{}); +} + +/// Every UE8M0 encoding, through the production matmul: each dot product is exactly one times its scale. +fn scaleRange(gpu: check.Gpu) !void { + const a = gpu.gpa; + const n = 256; + const k = 64; + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + @memset(codes, 0x38); + var scales: [n * 2]u8 = undefined; + for (0..n) |r| @memset(scales[r * 2 ..][0..2], @intCast(r)); + const tiled_codes = try a.alloc(u8, n * k); + defer a.free(tiled_codes); + cuda.mxfp8.packCodes(tiled_codes, codes, n, k); + var tiles: [n * 2]u8 = undefined; + cuda.mxfp8.packScales(&tiles, &scales, n, k); + var dw = try cuda.DeviceBuffer.fromHost(gpu.d, tiled_codes); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(gpu.d, &tiles); + defer ds.free(); + var x: [k]u16 = @splat(0x3c80); // 1/64: 64 equal terms sum to one + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&x)); + defer dx.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, n * 2); + defer dy.free(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + try lane.matmul(s, dx.ptr, k, 1, cuda.mxfp8.weight(dw.ptr, ds.ptr, n, k), dy.ptr, 0); + try s.synchronize(); + var got: [n * 2]u8 = undefined; + try dy.download(0, &got); + for (0..n) |r| { + const bits = std.mem.readInt(u16, got[r * 2 ..][0..2], .little); + if (r == 255) { + try check.expect(std.math.isNan(bf(bits)), "UE8M0 reserved byte produces NaN", .{}); + } else { + const want: u16 = if (r == 0) 0x40 else @intCast(r << 7); + try check.expect(bits == want, "UE8M0 {d}: BF16 bits {x} versus {x}", .{ r, bits, want }); + } + } + check.pass("MXFP8: all 256 UE8M0 encodings preserve the specified value, including subnormal and NaN", .{}); +} From 17ce05406e280d568cc2f314f73033fac168eef4 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:08:12 -0700 Subject: [PATCH 059/115] cuda: serve losslessly packed MXFP8 projections through native qlinear --- zig/kernels/cuda/qmmf.cu | 17 +++- zig/src/cuda/mxfp8.md | 27 +++++++ zig/src/cuda/mxfp8.zig | 41 ++++++++++ zig/src/cuda/qlinear.zig | 9 ++- zig/src/cuda/qmmf.zig | 14 ++-- zig/src/cuda/root.zig | 2 + zig/tests/cuda/main.zig | 3 + zig/tests/cuda/mxfp8_tests.zig | 143 +++++++++++++++++++++++++++++++++ 8 files changed, 247 insertions(+), 9 deletions(-) create mode 100644 zig/src/cuda/mxfp8.md create mode 100644 zig/src/cuda/mxfp8.zig create mode 100644 zig/tests/cuda/mxfp8_tests.zig diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 4703273410..027b8ef630 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -32,6 +32,11 @@ __device__ __forceinline__ uint32_t fp8pair(uint32_t w) { return r; } +// UE8M0 has no zero encoding: byte 0 is 2^-127, byte 255 is NaN. +__device__ __forceinline__ float e8m0f(uint8_t b) { + return __uint_as_float(b == 0 ? 0x00400000u : b == 255 ? 0x7FC00000u : uint32_t(b) << 23); +} + __device__ __forceinline__ uint32_t comp(const uint4& v, int c) { return c == 0 ? v.x : c == 1 ? v.y : c == 2 ? v.z : v.w; } __device__ __forceinline__ float e4m3f(uint8_t b) { @@ -193,8 +198,8 @@ __global__ void __launch_bounds__(WM * WN * 32) qmmf_kernel( s0 = e4m3f(static_cast(sq[j][0] >> (8 * blk))); s1 = e4m3f(static_cast(sq[j][1] >> (8 * blk))); } else { - s0 = __int_as_float(static_cast((sq[j][0] >> (8 * blk)) & 0xFFu) << 23); - s1 = __int_as_float(static_cast((sq[j][1] >> (8 * blk)) & 0xFFu) << 23); + s0 = e8m0f(static_cast(sq[j][0] >> (8 * blk))); + s1 = e8m0f(static_cast(sq[j][1] >> (8 * blk))); } #pragma unroll for (int i = 0; i < T::MT; ++i) @@ -293,6 +298,14 @@ __global__ void reduce_kernel(const float* __restrict__ part, void* __restrict__ else reinterpret_cast<__nv_bfloat16*>(out)[i] = __float2bfloat16_rn(acc); } +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); + // The instances the Zig launchers load by name (the Python host instantiates these through its launch switch). template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); diff --git a/zig/src/cuda/mxfp8.md b/zig/src/cuda/mxfp8.md new file mode 100644 index 0000000000..5878df43eb --- /dev/null +++ b/zig/src/cuda/mxfp8.md @@ -0,0 +1,27 @@ +# Native MXFP8 projections + +`cuda.mxfp8` packs row-major E4M3 code bytes and one UE8M0 scale byte per +32 input channels for `qlinear.Weight.mxfp8`. Both prompt and decode calls use +the lane matmul, with BF16 inputs, FP32 accumulation and BF16 output. Packing +permutes bytes; it does not requantize weights or expand them to BF16. + +The logical shape is `[N, K]`, with positive dimensions and K divisible by 64. +`packCodes` writes `padded(N) * K` bytes; `packScales` writes +`padded(N) * K/32` bytes. Rows are padded to 128; padded codes are zero and +padded scales are one. `weight` returns the borrowed device view. The caller +owns its device buffers and the lane's split-K scratch. + +UE8M0 bytes 0 through 254 mean `2^(byte - 127)`; byte 255 is NaN. In +particular, byte zero is a nonzero subnormal FP32 scale. See the +[CUDA datatype contract](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-math-api/cuda_math_api/struct____nv__fp8__e8m0.html). + +`tf-cuda-test mxfp8` checks all 256 scale encodings through the production +matmul, then compares synthetic projections with an independent FP64 +calculation. It also checks strided activations, partial output tiles, +row counts 1/7/17/33/256, prompt/decode dispatch, and cluster versus explicit +split reduction. Host tests cover scale permutation, padding, shared-memory +sizes, symbols and format refusals. + +This primitive does not admit new checkpoint families or certify model +correctness, tensor parallelism or performance. Rebuild the qmmf fatbin with +the native runtime; previous fatbins lack the MXFP8 instantiations. diff --git a/zig/src/cuda/mxfp8.zig b/zig/src/cuda/mxfp8.zig new file mode 100644 index 0000000000..7c57aa50df --- /dev/null +++ b/zig/src/cuda/mxfp8.zig @@ -0,0 +1,41 @@ +//! E4M3 codes with one UE8M0 byte per 32 inputs, repacked losslessly for the lane matmul. +const std = @import("std"); +const fp8 = @import("fp8.zig"); +const qmmf = @import("qmmf.zig"); + +pub const group = 32; +pub const padded = fp8.padded; +pub const packCodes = fp8.packCodes; + +/// Row-major [n, k/32] bytes to [npad/64][k/64][64][2]; padded rows have scale 1. +pub fn packScales(out: []u8, scales: []const u8, n: usize, k: usize) void { + const npad = padded(n); + const kg = k / qmmf.group; + std.debug.assert(k % qmmf.group == 0 and out.len == npad * (k / group) and scales.len == n * (k / group)); + for (0..npad / 64) |t| for (0..kg) |g| for (0..64) |r| for (0..2) |h| { + const row = t * 64 + r; + out[((t * kg + g) * 64 + r) * 2 + h] = if (row < n) scales[row * (k / group) + g * 2 + h] else 127; + }; +} + +pub fn weight(codes: u64, scales: u64, n: usize, k: usize) qmmf.Weight { + return .{ .mode = .mxfp8, .codes = codes, .scales = scales, .n = @intCast(n), .k = @intCast(k), .npad = @intCast(padded(n)) }; +} + +test "every scale byte lands in its independent 32-input group, including padded rows" { + const a = std.testing.allocator; + const n = 130; + const k = 256; + const scales = try a.alloc(u8, n * (k / group)); + defer a.free(scales); + for (scales, 0..) |*v, i| v.* = @truncate(i * 13 + 7); + const out = try a.alloc(u8, padded(n) * (k / group)); + defer a.free(out); + packScales(out, scales, n, k); + for (0..padded(n)) |row| for (0..k / group) |g| { + const at = (((row / 64) * (k / 64) + g / 2) * 64 + row % 64) * 2 + g % 2; + try std.testing.expectEqual(if (row < n) scales[row * (k / group) + g] else @as(u8, 127), out[at]); + }; + const w = weight(1, 2, n, k); + try std.testing.expect(w.mode == .mxfp8 and w.npad == 256 and w.scale == 1.0); +} diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..b924132f79 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -33,6 +33,7 @@ pub const Affine4 = struct { pub const Weight = union(enum) { affine4: Affine4, fp8g: qmmf.Weight, + mxfp8: qmmf.Weight, // e4m3 with UE8M0 scale bytes per 32 inputs nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values @@ -40,14 +41,14 @@ pub const Weight = union(enum) { pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, - .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .fp8g, .mxfp8, .nvfp4, .fp8, .bf16 => |q| q.n, }; } pub fn inputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.k, - .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .fp8g, .mxfp8, .nvfp4, .fp8, .bf16 => |q| q.k, }; } }; @@ -70,6 +71,7 @@ pub const Linear = struct { return a.decode(s, in.x, in.sums, q, out, rows); }, .fp8g => |q| return l.onLane(s, q, .fp8g, in, out, rows), + .mxfp8 => |q| return l.onLane(s, q, .mxfp8, in, out, rows), .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), @@ -85,6 +87,7 @@ pub const Linear = struct { return a.prompt(s, in.x, q, out, rows); }, .fp8g => |q| return l.onLane(s, q, .fp8g, in, out, rows), + .mxfp8 => |q| return l.onLane(s, q, .mxfp8, in, out, rows), .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), @@ -330,5 +333,7 @@ test "a lane weight whose mode is not its format's is refused before any launch" const fp4: qmmf.Weight = .{ .mode = .fp4, .codes = 0, .scales = 0, .n = 64, .k = 64, .npad = 128 }; const s: Stream = undefined; try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .fp8g = fp4 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatMismatch, l.decode(s, .{ .mxfp8 = fp4 }, .{ .x = 0 }, 0, 1)); + try std.testing.expectError(error.FormatMismatch, l.prompt(s, .{ .mxfp8 = fp4 }, .{ .x = 0 }, 0, 1)); try std.testing.expectError(error.FormatNotLoaded, l.prompt(s, .{ .nvfp4 = fp4 }, .{ .x = 0 }, 0, 1)); } diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 054e2c34bb..3e2b4168a7 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -1,4 +1,4 @@ -//! The lane matmul (qmmf.cu) for FP4, FP8, FP8G and BF16 weights, with its kernel, tile and slices. +//! The lane matmul (qmmf.cu) for FP4, FP8, MXFP8, FP8G and BF16 weights, with its kernel, tile and slices. const std = @import("std"); const driver = @import("driver.zig"); const module = @import("module.zig"); @@ -17,9 +17,9 @@ const bn = 64; const threads = 128; // WM 1 x WN 4 warps /// The packed formats (qmmf.cu Mode): FP4 with e4m3 scales a 16, e4m3 with one scale or fp32 a 64, and bf16. -pub const Mode = enum(u3) { fp4 = 0, fp8 = 1, fp8g = 3, bf16 = 4 }; +pub const Mode = enum(u3) { fp4 = 0, fp8 = 1, mxfp8 = 2, fp8g = 3, bf16 = 4 }; -const modes = [_]Mode{ .fp4, .fp8, .fp8g, .bf16 }; +const modes = [_]Mode{ .fp4, .fp8, .mxfp8, .fp8g, .bf16 }; fn modeIndex(mode: Mode) usize { return std.mem.indexOfScalar(Mode, &modes, mode).?; @@ -29,10 +29,10 @@ fn modeIndex(mode: Mode) usize { pub fn sharedBytes(mode: Mode, bm: u32) u32 { const w: u32 = switch (mode) { .fp4 => bn * group / 2, - .fp8, .fp8g => bn * group, + .fp8, .mxfp8, .fp8g => bn * group, .bf16 => bn * group * 2, }; - const s: u32 = if (mode == .fp4 or mode == .fp8g) bn * 4 else 0; + const s: u32 = if (mode == .fp4 or mode == .fp8g) bn * 4 else if (mode == .mxfp8) bn * 2 else 0; const stage = (bm * group * 2 + w + s + 127) / 128 * 128; const partials = (bm / 16) * 2 * 4 * threads * 4; return @max(4 * stage, partials); @@ -186,6 +186,9 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { try std.testing.expectEqual(@as(u32, 25600), sharedBytes(.fp8g, 16)); try std.testing.expectEqual(@as(u32, 33792), sharedBytes(.fp8g, 32)); try std.testing.expectEqual(@as(u32, 50176), sharedBytes(.fp8g, 64)); + try std.testing.expectEqual(@as(u32, 25088), sharedBytes(.mxfp8, 16)); + try std.testing.expectEqual(@as(u32, 33280), sharedBytes(.mxfp8, 32)); + try std.testing.expectEqual(@as(u32, 49664), sharedBytes(.mxfp8, 64)); try std.testing.expectEqual(@as(u32, 17408), sharedBytes(.fp4, 16)); try std.testing.expectEqual(@as(u32, 25600), sharedBytes(.fp4, 32)); try std.testing.expectEqual(@as(u32, 41984), sharedBytes(.fp4, 64)); @@ -209,6 +212,7 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.mxfp8, 16, true, false)); try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false)); try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true)); } diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index dd9b45046c..0483821247 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -26,6 +26,7 @@ pub const aot = @import("aot.zig"); pub const kernels = @import("kernels.zig"); pub const qmmf = @import("qmmf.zig"); pub const fp8 = @import("fp8.zig"); +pub const mxfp8 = @import("mxfp8.zig"); pub const exlr8 = @import("exlr8.zig"); pub const nvfp4 = @import("nvfp4.zig"); pub const segments = @import("segments.zig"); @@ -53,6 +54,7 @@ test { _ = segments; _ = @import("qmmf.zig"); _ = @import("fp8.zig"); + _ = @import("mxfp8.zig"); _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 6fa70d0753..c4a8291323 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -16,6 +16,7 @@ const prompt_cuts = @import("prompt_cuts.zig"); const chunk_costs = @import("chunk_costs.zig"); const grouped_tests = @import("grouped_tests.zig"); const fp8_tests = @import("fp8_tests.zig"); +const mxfp8_tests = @import("mxfp8_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); @@ -69,6 +70,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes @@ -130,6 +132,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); diff --git a/zig/tests/cuda/mxfp8_tests.zig b/zig/tests/cuda/mxfp8_tests.zig new file mode 100644 index 0000000000..5417d3d1fd --- /dev/null +++ b/zig/tests/cuda/mxfp8_tests.zig @@ -0,0 +1,143 @@ +//! Native MXFP8: lossless packing, independent FP64 products, strided rows and row-count invariance. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); + +fn bf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +fn tobf(v: f32) u16 { + const u: u32 = @bitCast(v); + return @truncate((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +pub fn run(gpu: check.Gpu) !void { + try scaleRange(gpu); + const a = gpu.gpa; + const n = 70; // a partial tile and padded rows + const k = 1024; // split-K, both cluster and explicit reduction + const m = 256; // fused prompt reduction + const ldx = k + 64; + const npad = cuda.mxfp8.padded(n); + var prng = std.Random.DefaultPrng.init(0x8e80_5eed); + const rng = prng.random(); + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + for (codes) |*v| { + v.* = rng.int(u8); + if (v.* & 0x7f == 0x7f) v.* ^= 1; + } + const scales = try a.alloc(u8, n * (k / 32)); + defer a.free(scales); + for (scales) |*v| v.* = 121 + rng.uintLessThan(u8, 8); + // A dedicated row makes UE8M0 byte 0 observable in a normal BF16 result. + @memset(codes[0..k], 0x38); // E4M3 1 + @memset(scales[0 .. k / 32], 0); + const packed_codes = try a.alloc(u8, npad * k); + defer a.free(packed_codes); + cuda.mxfp8.packCodes(packed_codes, codes, n, k); + const packed_scales = try a.alloc(u8, npad * (k / 32)); + defer a.free(packed_scales); + cuda.mxfp8.packScales(packed_scales, scales, n, k); + var dw = try cuda.DeviceBuffer.fromHost(gpu.d, packed_codes); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(gpu.d, packed_scales); + defer ds.free(); + const w = cuda.mxfp8.weight(dw.ptr, ds.ptr, n, k); + const x = try a.alloc(u16, m * ldx); + defer a.free(x); + @memset(x, 0x7fc0); // stride padding is poison, not another activation + for (0..m) |row| { + for (x[row * ldx ..][0..k]) |*v| v.* = tobf(rng.float(f32) * 2); + } + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(x)); + defer dx.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, m * n * 2); + defer dy.free(); + var part = try cuda.DeviceBuffer.alloc(gpu.d, 8 * m * n * 4); + defer part.free(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var baseline: ?[]u8 = null; + defer if (baseline) |b| a.free(b); + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters and (try gpu.ctx.capability()) >= 90; + const lin: cuda.qlinear.Linear = .{ .lane = &lane }; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + for ([_]bool{ false, true }) |prompt| { + try dy.fill8(0xff, s.handle); + const in: cuda.qlinear.Rows = .{ .x = dx.ptr, .ldx = ldx, .part = part.ptr }; + if (prompt) try lin.prompt(s, .{ .mxfp8 = w }, in, dy.ptr, rows) else try lin.decode(s, .{ .mxfp8 = w }, in, dy.ptr, rows); + try s.synchronize(); + const got = try a.alloc(u8, rows * n * 2); + defer a.free(got); + try dy.download(0, got); + if (baseline) |b| { + try check.sameBytes("MXFP8 row-count/dispatch/cluster invariance", got, b[0..got.len]); + } else baseline = try a.dupe(u8, got); + for (0..rows) |row| for (0..n) |col| { + var ref: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const scale = std.math.ldexp(@as(f64, 1), @as(i32, scales[col * (k / 32) + i / 32]) - 127); + const v = bf(x[row * ldx + i]) * core.quant.fp8.e4m3(codes[col * k + i]) * scale; + ref += v; + magnitude += @abs(v); + } + const actual = bf(std.mem.readInt(u16, got[(row * n + col) * 2 ..][0..2], .little)); + const tolerance = @abs(ref) / 128 + magnitude / 4096 + 0x1p-133; + try check.expect(@abs(actual - ref) <= tolerance, "MXFP8 FP64 row {d}, column {d}: {e} versus {e}, tolerance {e}", .{ row, col, actual, ref, tolerance }); + }; + } + } + } + check.pass("MXFP8: strided rows 1,7,17,33,256 equal across dispatch and reduction paths; FP64 including UE8M0 byte zero", .{}); +} + +/// Every UE8M0 encoding, through the production matmul: each dot product is exactly one times its scale. +fn scaleRange(gpu: check.Gpu) !void { + const a = gpu.gpa; + const n = 256; + const k = 64; + const codes = try a.alloc(u8, n * k); + defer a.free(codes); + @memset(codes, 0x38); + var scales: [n * 2]u8 = undefined; + for (0..n) |r| @memset(scales[r * 2 ..][0..2], @intCast(r)); + const tiled_codes = try a.alloc(u8, n * k); + defer a.free(tiled_codes); + cuda.mxfp8.packCodes(tiled_codes, codes, n, k); + var tiles: [n * 2]u8 = undefined; + cuda.mxfp8.packScales(&tiles, &scales, n, k); + var dw = try cuda.DeviceBuffer.fromHost(gpu.d, tiled_codes); + defer dw.free(); + var ds = try cuda.DeviceBuffer.fromHost(gpu.d, &tiles); + defer ds.free(); + var x: [k]u16 = @splat(0x3c80); // 1/64: 64 equal terms sum to one + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&x)); + defer dx.free(); + var dy = try cuda.DeviceBuffer.alloc(gpu.d, n * 2); + defer dy.free(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + try lane.matmul(s, dx.ptr, k, 1, cuda.mxfp8.weight(dw.ptr, ds.ptr, n, k), dy.ptr, 0); + try s.synchronize(); + var got: [n * 2]u8 = undefined; + try dy.download(0, &got); + for (0..n) |r| { + const bits = std.mem.readInt(u16, got[r * 2 ..][0..2], .little); + if (r == 255) { + try check.expect(std.math.isNan(bf(bits)), "UE8M0 reserved byte produces NaN", .{}); + } else { + const want: u16 = if (r == 0) 0x40 else @intCast(r << 7); + try check.expect(bits == want, "UE8M0 {d}: BF16 bits {x} versus {x}", .{ r, bits, want }); + } + } + check.pass("MXFP8: all 256 UE8M0 encodings preserve the specified value, including subnormal and NaN", .{}); +} From f6a0055f5d84dfe9f61342ef4923cd6c233bc037 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:23:49 -0700 Subject: [PATCH 060/115] Expose FP32 lane projection outputs before residual reduction --- zig/kernels/cuda/qmmf.cu | 36 +++++++ zig/src/cuda/qmmf.md | 19 ++++ zig/src/cuda/qmmf.zig | 40 ++++++-- zig/tests/cuda/main.zig | 3 + zig/tests/cuda/projection_f32_tests.zig | 127 ++++++++++++++++++++++++ 5 files changed, 214 insertions(+), 11 deletions(-) create mode 100644 zig/src/cuda/qmmf.md create mode 100644 zig/tests/cuda/projection_f32_tests.zig diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 027b8ef630..f34fc02118 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -335,5 +335,41 @@ template __global__ void qmmf_kernel( template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); } diff --git a/zig/src/cuda/qmmf.md b/zig/src/cuda/qmmf.md new file mode 100644 index 0000000000..82f5c5a1ad --- /dev/null +++ b/zig/src/cuda/qmmf.md @@ -0,0 +1,19 @@ +# FP32 projection outputs + +`cuda.qmmf.Lane.matmulF32` has the same inputs, packed weights, stride, +shape-driven split-K plan and scratch contract as `matmul`. Its output is +`[M, N]` FP32, requiring `M * N * 4` bytes. It preserves FP32 output for a +later residual or tensor-rank sum; the caller chooses when to round to BF16. + +All five lane formats are supported: NVFP4, tensor-scaled FP8, MXFP8, +block-scaled FP8 and BF16. The change exposes the existing kernel template's +FP32 output path. It does not change weight packing or accumulation order. +Rebuild the qmmf fatbin together with the native runtime: older fatbins do +not contain the FP32 instances or reduction symbol. + +`tf-cuda-test projection-f32` compares each format against an independent +FP64 calculation, checks conversion to the existing BF16 result, and varies +row count, activation stride, partial output tiles, fused reduction, +cluster reduction and explicit split reduction. The fixture is synthetic; +model correctness and multi-device collective behavior require separate +qualification. diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 3e2b4168a7..adfddf7124 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -66,9 +66,9 @@ pub fn plan(m: usize, n: usize, k: usize, clusters: bool) Plan { } /// qmmf.cu's instance names, per mode: row tiles 16, 32, 64 (each plain, then with a cluster), then the fused kernel. -fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool) [:0]const u8 { - const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb0ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; - return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(cluster), @intFromBool(fused) }); +fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool, comptime f32_out: bool) [:0]const u8 { + const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb{d}ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; + return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(f32_out), @intFromBool(cluster), @intFromBool(fused) }); } pub const reduce_symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif"; @@ -78,22 +78,30 @@ pub const Lane = struct { mod: Module, fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused reduce: Function, + f32_fns: [modes.len][7]Function, + reduce_f32: Function, clusters: bool, // sm_90 on: slices meet in a cluster's shared memory pub fn load(d: *const Driver, capability_major: u32) !Lane { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.qmmf); errdefer mod.unload(); - var l: Lane = .{ .mod = mod, .fns = undefined, .reduce = try mod.function(reduce_symbol), .clusters = capability_major >= 9 }; + var l: Lane = .{ .mod = mod, .fns = undefined, .f32_fns = undefined, .reduce = try mod.function(reduce_symbol), .reduce_f32 = try mod.function("_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif"), .clusters = capability_major >= 9 }; inline for (modes, 0..) |mode, mi| { inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { - const f = try mod.function(symbol(mode, bm, cl, false)); + const f = try mod.function(symbol(mode, bm, cl, false, false)); try f.allowDynamicShared(sharedBytes(mode, bm)); l.fns[mi][ti * 2 + ci] = f; + const fp32_fn = try mod.function(symbol(mode, bm, cl, false, true)); + try fp32_fn.allowDynamicShared(sharedBytes(mode, bm)); + l.f32_fns[mi][ti * 2 + ci] = fp32_fn; }; - const f = try mod.function(symbol(mode, 64, false, true)); + const f = try mod.function(symbol(mode, 64, false, true, false)); try f.allowDynamicShared(sharedBytes(mode, 64)); l.fns[mi][6] = f; + const fp32_fn = try mod.function(symbol(mode, 64, false, true, true)); + try fp32_fn.allowDynamicShared(sharedBytes(mode, 64)); + l.f32_fns[mi][6] = fp32_fn; } return l; } @@ -110,8 +118,17 @@ pub const Lane = struct { /// y [m, n] bf16 = x [m, k] bf16 (rows `ldx` apart) times the weight, fp32 sums; `part` holds partBytes. pub fn matmul(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulTo(s, x, ldx, m, w, y, part, false); + } + + /// The same accumulation and reduction, preserving FP32 outputs for a later residual or tensor-rank sum. + pub fn matmulF32(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulTo(s, x, ldx, m, w, y, part, true); + } + + fn matmulTo(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64, f32_out: bool) !void { const p = plan(m, w.n, w.k, l.clusters); - const set = l.fns[modeIndex(w.mode)]; + const set = (if (f32_out) l.f32_fns else l.fns)[modeIndex(w.mode)]; const ti: usize = if (p.bm == 16) 0 else if (p.bm == 32) 1 else 2; const f = if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; const rows_t = (m + p.bm - 1) / p.bm; @@ -144,7 +161,7 @@ pub const Lane = struct { r.add(@as(i64, @intCast(total))); r.add(@as(i32, @intCast(p.sk))); r.add(w.scale); - try launch_.launch(l.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); + try launch_.launch(if (f32_out) l.reduce_f32 else l.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); } } }; @@ -212,7 +229,8 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.mxfp8, 16, true, false)); - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false)); - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.bf16, 64, false, true, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.mxfp8, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true, false)); } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 2d6b6bd85f..ac5249e1d6 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -17,6 +17,7 @@ const chunk_costs = @import("chunk_costs.zig"); const grouped_tests = @import("grouped_tests.zig"); const fp8_tests = @import("fp8_tests.zig"); const mxfp8_tests = @import("mxfp8_tests.zig"); +const projection_f32_tests = @import("projection_f32_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); @@ -62,6 +63,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) @@ -124,6 +126,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); diff --git a/zig/tests/cuda/projection_f32_tests.zig b/zig/tests/cuda/projection_f32_tests.zig new file mode 100644 index 0000000000..85b1160365 --- /dev/null +++ b/zig/tests/cuda/projection_f32_tests.zig @@ -0,0 +1,127 @@ +//! FP32 lane output before residual/rank combination: independent FP64, rounding and row/reduction invariance. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const reference = @import("quant_ref_tests.zig"); +const Rig = reference.RigT; +const bf = reference.helpers.bfOf; +const bf16 = reference.helpers.bf16Of; +const Mode = cuda.qmmf.Mode; + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0xf32_5eed); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + for ([_]Mode{ .fp4, .fp8, .mxfp8, .fp8g, .bf16 }) |mode| try projection(&r, &lane, mode); +} + +fn projection(r: *Rig, lane: *cuda.qmmf.Lane, mode: Mode) !void { + const n = 70; + const k = 1024; + const m = 256; + const ldx = k + 64; + const npad = cuda.fp8.padded(n); + const a = r.a; + const codes = try a.alloc(u8, n * k); + for (codes) |*v| v.* = reference.helpers.fp8CodeOf(r); + const scales = try a.alloc(u8, n * (k / 16)); + for (scales) |*v| v.* = if (mode == .mxfp8) 121 + r.rng.uintLessThan(u8, 8) else 0x20 + r.rng.uintLessThan(u8, 0x28); + const inv = try a.alloc(f32, ((n + 127) / 128) * (k / 128)); + for (inv) |*v| v.* = r.rng.float(f32) * 0.01 + 0.001; + const w16 = try r.bfs(n * k, 0.25); + const packed16 = try a.alloc(u16, npad * k); + const tiled_codes = std.mem.sliceAsBytes(packed16); + const global: f32 = 0.0123; + var w: cuda.qmmf.Weight = undefined; + switch (mode) { + .bf16 => { + cuda.qmmf.packBf16(packed16, w16, n, k); + w = cuda.modelopt.bf16Weight(try r.dev(tiled_codes), n, k); + }, + .fp4 => { + const words = try a.alloc(u32, npad * k / 8); + cuda.nvfp4.packWords(words, codes[0 .. n * k / 2], n, k); + const tiles = try a.alloc(u8, npad * (k / 16)); + cuda.nvfp4.packScales(tiles, scales, n, k); + w = cuda.nvfp4.weight(try r.dev(std.mem.sliceAsBytes(words)), try r.dev(tiles), global, n, k); + }, + .fp8, .fp8g, .mxfp8 => { + cuda.fp8.packCodes(tiled_codes[0 .. npad * k], codes, n, k); + const dw = try r.dev(tiled_codes[0 .. npad * k]); + w = switch (mode) { + .fp8 => cuda.modelopt.fp8Weight(dw, global, n, k), + .fp8g => blk: { + const tiles = try a.alloc(f32, npad * (k / 64)); + cuda.fp8.packScales(tiles, inv, n, k); + break :blk cuda.fp8.weight(dw, try r.dev(std.mem.sliceAsBytes(tiles)), n, k); + }, + .mxfp8 => blk: { + const tiles = try a.alloc(u8, npad * (k / 32)); + cuda.mxfp8.packScales(tiles, scales[0 .. n * k / 32], n, k); + break :blk cuda.mxfp8.weight(dw, try r.dev(tiles), n, k); + }, + else => unreachable, + }; + }, + } + const x = try a.alloc(u16, m * ldx); + @memset(x, 0x7fc0); + for (0..m) |row| { + for (x[row * ldx ..][0..k]) |*v| v.* = bf16((r.rng.float(f32) - 0.5) * 2); + } + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const out = try r.zeros(m * n * 4); + const out_bf16 = try r.zeros(m * n * 2); + const part = try r.zeros(8 * m * n * 4); + const want = try a.alloc(f64, m * n); + const tol = try a.alloc(f64, m * n); + for (0..m) |row| for (0..n) |col| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const code = codes[col * k + i]; + const weight: f64 = switch (mode) { + .bf16 => bf(w16[col * k + i]), + .fp8 => core.quant.fp8.e4m3(code) * global, + .fp8g => core.quant.fp8.e4m3(code) * inv[(col / 128) * (k / 128) + i / 128], + .mxfp8 => core.quant.fp8.e4m3(code) * std.math.ldexp(@as(f64, 1), @as(i32, scales[col * (k / 32) + i / 32]) - 127), + .fp4 => blk: { + const byte = codes[col * (k / 2) + i / 2]; + break :blk core.quant.nvfp4.e2m1(if (i % 2 == 0) byte & 15 else byte >> 4) * core.quant.fp8.e4m3(scales[col * (k / 16) + i / 16]) * global; + }, + }; + const v = bf(x[row * ldx + i]) * weight; + sum += v; + magnitude += @abs(v); + } + want[row * n + col] = sum; + tol[row * n + col] = magnitude * 2e-6 + @abs(sum) * 1e-6 + 1e-30; + }; + var baseline: ?[]u8 = null; + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters and (try r.gpu.ctx.capability()) >= 90; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + try lane.matmulF32(r.s, dx, ldx, rows, w, out, part); + try r.close("FP32 projection", out, true, want[0 .. rows * n], tol[0 .. rows * n]); + const got = try r.back(out, rows * n * 4); + if (baseline) |b| try check.sameBytes("FP32 row/reduction invariance", got, b[0..got.len]) else baseline = got; + try lane.matmul(r.s, dx, ldx, rows, w, out_bf16, part); + const rounded = try r.back(out_bf16, rows * n * 2); + for (0..rows * n) |i| { + const v = std.mem.readInt(u32, got[i * 4 ..][0..4], .little); + try check.expect(bf16(@bitCast(v)) == std.mem.readInt(u16, rounded[i * 2 ..][0..2], .little), "FP32 rounds once to the existing BF16 result", .{}); + } + } + } + check.pass("{t}: FP32 outputs preserve FP64 accuracy, BF16 rounding and row/reduction invariance", .{mode}); +} From 1f17be1d0195fada4b82ce747d6d113467c3f708 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:29:20 -0700 Subject: [PATCH 061/115] Read native NVFP4 n-gram rows directly from mapped checkpoint shards --- zig/src/core/nvfp4_rows.zig | 57 +++++++++ zig/src/core/nvfp4_rows_test.zig | 77 ++++++++++++ zig/src/core/root.zig | 2 + zig/src/families/flashnext/cuda_ple.zig | 66 ++++++++--- .../flashnext/cuda_ple_nvfp4_test.zig | 111 ++++++++++++++++++ 5 files changed, 297 insertions(+), 16 deletions(-) create mode 100644 zig/src/core/nvfp4_rows.zig create mode 100644 zig/src/core/nvfp4_rows_test.zig create mode 100644 zig/src/families/flashnext/cuda_ple_nvfp4_test.zig diff --git a/zig/src/core/nvfp4_rows.zig b/zig/src/core/nvfp4_rows.zig new file mode 100644 index 0000000000..5a03b7b2c8 --- /dev/null +++ b/zig/src/core/nvfp4_rows.zig @@ -0,0 +1,57 @@ +//! Borrowed NVFP4 table rows: decode requested mapped rows to BF16 without expanding the table. +const std = @import("std"); +const Tensor = @import("safetensors.zig").Tensor; +const quant = @import("quant/quant.zig"); + +pub const Error = error{ BadNvfp4Table, BadNvfp4Rows, AliasedNvfp4Rows }; + +pub const Table = struct { + codes: []const u8, + scales: []const u8, + rows: usize, + width: usize, + global: f32, + + /// Borrow exact row-major U8 codes and E4M3 scales per 16 values; the caller retains the maps. + pub fn init(weight: Tensor, scales: Tensor, global: f32) Error!Table { + if (weight.rank != 2 or weight.dtype != .u8 or weight.dim(0) == 0 or weight.dim(1) == 0) + return error.BadNvfp4Table; + const width = std.math.mul(usize, weight.dim(1), 2) catch return error.BadNvfp4Table; + const size = std.math.mul(usize, weight.dim(0), weight.dim(1)) catch return error.BadNvfp4Table; + if (width % 16 != 0 or weight.bytes.len != size or !std.math.isFinite(global)) return error.BadNvfp4Table; + if (!scales.is(.f8_e4m3, &.{ weight.dim(0), width / 16 }) or scales.bytes.len != size / 8) + return error.BadNvfp4Table; + return .{ .codes = weight.bytes, .scales = scales.bytes, .rows = weight.dim(0), .width = width, .global = global }; + } + + /// Token-major requested rows, with FP32 products rounded once to BF16; malformed requests write nothing. + pub fn gather(t: Table, ids: []const usize, out: []u16) Error!void { + const count = std.math.mul(usize, ids.len, t.width) catch return error.BadNvfp4Rows; + if (out.len != count) return error.BadNvfp4Rows; + for (ids) |id| if (id >= t.rows) return error.BadNvfp4Rows; + const bytes = std.mem.sliceAsBytes(out); + if (overlaps(bytes, t.codes) or overlaps(bytes, t.scales) or overlaps(bytes, std.mem.sliceAsBytes(ids))) + return error.AliasedNvfp4Rows; + for (ids, 0..) |id, row| { + for (0..t.width) |col| { + const code = t.codes[id * (t.width / 2) + col / 2]; + const nib = if (col % 2 == 0) code & 15 else code >> 4; + const a: f32 = @floatCast(quant.nvfp4.e2m1(nib)); + const b: f32 = @floatCast(quant.fp8.e4m3(t.scales[id * (t.width / 16) + col / 16])); + const product: f32 = a * b; + out[row * t.width + col] = quant.convert.bf16(product * t.global); + } + } + } +}; + +fn overlaps(a: []const u8, b: []const u8) bool { + if (a.len == 0 or b.len == 0) return false; + const x = @intFromPtr(a.ptr); + const y = @intFromPtr(b.ptr); + return if (x <= y) y - x < a.len else x - y < b.len; +} + +test { + _ = @import("nvfp4_rows_test.zig"); +} diff --git a/zig/src/core/nvfp4_rows_test.zig b/zig/src/core/nvfp4_rows_test.zig new file mode 100644 index 0000000000..284f13fc57 --- /dev/null +++ b/zig/src/core/nvfp4_rows_test.zig @@ -0,0 +1,77 @@ +const std = @import("std"); +const Table = @import("nvfp4_rows.zig").Table; +const Tensor = @import("safetensors.zig").Tensor; + +fn tensor(dtype: @import("safetensors.zig").DType, rows: usize, cols: usize, bytes: []const u8) Tensor { + return .{ .dtype = dtype, .rank = 2, .shape = .{ rows, cols, 1, 1, 1 }, .bytes = bytes }; +} + +fn rounded(value: f64) u16 { + const v: f32 = @floatCast(value); + if (std.math.isNan(v)) return 0x7fc0; + const bits: u32 = @bitCast(v); + const high = bits >> 16; + const low = bits & 65535; + return @intCast(high + @intFromBool(low > 32768 or (low == 32768 and high % 2 == 1))); +} + +fn scale(byte: u8) f64 { + const mag = byte & 127; + if (mag == 127) return std.math.nan(f64); + const e: i32 = @as(i32, mag >> 3) - 7; + const value = if (mag < 8) @as(f64, @floatFromInt(mag)) / 512 else + std.math.ldexp(1 + @as(f64, @floatFromInt(mag & 7)) / 8, e); + return if (byte & 128 != 0) -value else value; +} + +test "mapped NVFP4 rows preserve every nibble and scale encoding in reordered repeated gathers" { + var codes: [256 * 8]u8 = undefined; + var scales: [256]u8 = undefined; + for (0..256) |row| { + scales[row] = @intCast(row); + for (0..8) |col| codes[row * 8 + col] = @intCast(col * 2 + ((col * 2 + 1) << 4)); + } + const t = try Table.init(tensor(.u8, 256, 8, &codes), tensor(.f8_e4m3, 256, 1, &scales), 1.375); + var ids: [258]usize = undefined; + for (0..256) |i| ids[i] = 255 - i; + ids[256] = 17; + ids[257] = 17; + var out: [258 * 16]u16 = undefined; + try t.gather(&ids, &out); + const values = [_]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + for (ids, 0..) |id, row| for (0..16) |col| { + const sign: f64 = if (col >= 8) -1 else 1; + const want = rounded(sign * values[col % 8] * scale(@intCast(id)) * 1.375); + const actual = out[row * 16 + col]; + if (id & 127 == 127) { + try std.testing.expect(actual & 0x7fc0 == 0x7fc0); + } else try std.testing.expectEqual(want, actual); + }; + try t.gather(&.{}, &.{}); +} + +test "a bad row or destination refuses before writing and maps are never aliased" { + var storage: [128]u8 align(8) = @splat(0x22); + const scales = [_]u8{ 0x38, 0x40 }; + const t = try Table.init(tensor(.u8, 2, 8, storage[0..16]), tensor(.f8_e4m3, 2, 1, &scales), 1); + var out: [32]u16 = @splat(0xbeef); + try std.testing.expectError(error.BadNvfp4Rows, t.gather(&.{ 0, 2 }, &out)); + for (out) |v| try std.testing.expectEqual(@as(u16, 0xbeef), v); + try std.testing.expectError(error.BadNvfp4Rows, t.gather(&.{0}, &out)); + try std.testing.expectError(error.AliasedNvfp4Rows, t.gather(&.{0}, std.mem.bytesAsSlice(u16, storage[0..32]))); + var ids: [4]usize = .{ 0, 0, 0, 0 }; + try std.testing.expectError(error.AliasedNvfp4Rows, t.gather(ids[0..1], std.mem.bytesAsSlice(u16, std.mem.sliceAsBytes(&ids)))); +} + +test "invalid table dtype shape extent and scale cannot create a borrowed view" { + const codes: [16]u8 = @splat(0); + const scales = [_]u8{0x38}; + const w = tensor(.u8, 1, 16, &codes); + const s = tensor(.f8_e4m3, 1, 1, &scales); + try std.testing.expectError(error.BadNvfp4Table, Table.init(w, s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, 1, 7, codes[0..7]), s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.bf16, 1, 8, &codes), s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, 1, 8, &codes), s, 1)); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, 1, 8, codes[0..8]), s, std.math.inf(f32))); + try std.testing.expectError(error.BadNvfp4Table, Table.init(tensor(.u8, std.math.maxInt(usize), 8, &codes), s, 1)); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 3c125aa71f..19ae58e821 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -15,6 +15,7 @@ pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); pub const nvfp4_csf = @import("nvfp4_csf.zig"); +pub const nvfp4_rows = @import("nvfp4_rows.zig"); pub const nvfp4_slice = @import("nvfp4_slice.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); @@ -32,6 +33,7 @@ test { _ = quant; _ = nvfp4_csf; _ = nvfp4_slice; + _ = nvfp4_rows; _ = registry; _ = tuning; _ = norms_ref; diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index 96348db971..459f56905d 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -203,10 +203,15 @@ pub fn weightName(buf: []u8, shard: usize, nested: bool) ![]const u8 { } pub fn partName(buf: []u8, shard: usize, nested: bool, part: []const u8) ![]const u8 { + return partNameFor(buf, shard, nested, part, false); +} + +fn partNameFor(buf: []u8, shard: usize, nested: bool, part: []const u8, modelopt: bool) ![]const u8 { + const prefix = if (modelopt) "model.language_model" else "language_model.model"; return if (nested) - std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.{s}", .{ layer, shard, part }) + std.fmt.bufPrint(buf, "{s}.layers.{d}.ple.ple_embedding.ngram_embedding.shards.{d}.{s}", .{ prefix, layer, shard, part }) else - std.fmt.bufPrint(buf, "language_model.model.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.{s}", .{ layer, shard, part }); + std.fmt.bufPrint(buf, "{s}.layers.{d}.ple.ple_embedding.ngram_embedding.shard_{d}.{s}", .{ prefix, layer, shard, part }); } pub fn findOpen(opened: []Open, path: []const u8) ?*Open { @@ -230,8 +235,11 @@ pub fn copyRow(file: *const Open, shard: usize, nested: bool, local: usize, dh: } pub fn tableScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap) !f32 { - const name = "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale"; - const file = wm.get(name) orelse return 1; + return readScale(gpa, io, dir, wm, "language_model.model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale", false); +} + +fn readScale(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, wm: std.json.ObjectMap, name: []const u8, required: bool) !f32 { + const file = wm.get(name) orelse return if (required) error.MissingTensor else 1; const path = try std.fs.path.join(gpa, &.{ dir, file.string }); defer gpa.free(path); var opened = try Open.load(gpa, io, path); @@ -333,6 +341,8 @@ pub const Ngram = struct { w: [max_shards][]const u8 = undefined, s: [max_shards][]const u8 = undefined, b: [max_shards][]const u8 = undefined, + format: enum { mlx, nvfp4 } = .mlx, + tables: [max_shards]core.nvfp4_rows.Table = undefined, pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { const cfg = try readSpec(gpa, io, model_dir); @@ -356,16 +366,22 @@ pub const Ngram = struct { const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); defer parsed.deinit(); const wm = parsed.value.object.get("weight_map").?.object; - g.scale = try tableScale(gpa, io, model_dir, wm); + const first = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_0.weight"; + const modelopt = wm.contains(first); + g.format = if (modelopt) .nvfp4 else .mlx; + g.scale = if (modelopt) + try readScale(gpa, io, model_dir, wm, "model.language_model.layers.1.ple.ple_embedding.ngram_embedding.weight_scale_2", true) + else + try tableScale(gpa, io, model_dir, wm); const dh = g.hash.dh; var key: [180]u8 = undefined; for (0..g.hash.shards) |shard| { var nested = false; - const mapped = if (wm.get(try weightName(&key, shard, false))) |v| + const mapped = if (wm.get(try partNameFor(&key, shard, false, "weight", modelopt))) |v| v.string else blk: { nested = true; - break :blk (wm.get(try weightName(&key, shard, true)) orelse return error.MissingTensor).string; + break :blk (wm.get(try partNameFor(&key, shard, true, "weight", modelopt)) orelse return error.MissingTensor).string; }; const path = try std.fs.path.join(gpa, &.{ model_dir, mapped }); defer gpa.free(path); @@ -376,13 +392,23 @@ pub const Ngram = struct { try g.files.append(gpa, o); break :file &g.files.items[g.files.items.len - 1]; }; - const w = try file.tensor(try partName(&key, shard, nested, "weight"), .u32); - const s = try file.tensor(try partName(&key, shard, nested, "scales"), .bf16); - const b = try file.tensor(try partName(&key, shard, nested, "biases"), .bf16); - if (w.dim(1) != dh / 8 or s.dim(1) != dh / 32 or b.dim(1) != dh / 32) return error.UnexpectedTensor; - g.w[shard] = w.bytes; - g.s[shard] = s.bytes; - g.b[shard] = b.bytes; + const w = try file.tensor(try partNameFor(&key, shard, nested, "weight", modelopt), if (modelopt) .u8 else .u32); + const expected_rows: usize = @intCast(g.hash.starts[shard + 1] - g.hash.starts[shard]); + if (w.rank != 2 or w.dim(0) != expected_rows) return error.UnexpectedTensor; + if (modelopt) { + const sc = try file.tensor(try partNameFor(&key, shard, nested, "weight_scale", true), .f8_e4m3); + const table = try core.nvfp4_rows.Table.init(w, sc, g.scale); + if (table.width != dh) return error.UnexpectedTensor; + g.tables[shard] = table; + } else { + const sc = try file.tensor(try partNameFor(&key, shard, nested, "scales", false), .bf16); + const bias = try file.tensor(try partNameFor(&key, shard, nested, "biases", false), .bf16); + if (!w.is(.u32, &.{ expected_rows, dh / 8 }) or !sc.is(.bf16, &.{ expected_rows, dh / 32 }) or + !bias.is(.bf16, &.{ expected_rows, dh / 32 })) return error.UnexpectedTensor; + g.w[shard] = w.bytes; + g.s[shard] = sc.bytes; + g.b[shard] = bias.bytes; + } } return g; } @@ -402,7 +428,11 @@ pub const Ngram = struct { const dh = g.hash.dh; for (lo..hi) |i| { const loc = try g.hash.locate(ids[i]); - try embed.dequant(g.w[loc.shard], g.s[loc.shard], g.b[loc.shard], dh, loc.local, out[i * dh ..][0..dh]); + const row = out[i * dh ..][0..dh]; + switch (g.format) { + .mlx => try embed.dequant(g.w[loc.shard], g.s[loc.shard], g.b[loc.shard], dh, loc.local, row), + .nvfp4 => try g.tables[loc.shard].gather(&.{loc.local}, row), + } } } @@ -449,7 +479,7 @@ pub const Ngram = struct { } for (0..tokens.len) |t| { const row = out[t * w ..][0..w]; - if (g.scale != 1) { + if (g.format == .mlx and g.scale != 1) { for (row) |*v| { const f: f32 = @bitCast(@as(u32, v.*) << 16); v.* = toBf16(f * g.scale); @@ -484,3 +514,7 @@ test "a fresh window hashes token 7 into sixteen rows" { try std.testing.expectEqual(@as(usize, 1), loc.shard); try std.testing.expectEqual(@as(usize, 427641), loc.local); } + +test { + _ = @import("cuda_ple_nvfp4_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_ple_nvfp4_test.zig b/zig/src/families/flashnext/cuda_ple_nvfp4_test.zig new file mode 100644 index 0000000000..ee81d7a7b4 --- /dev/null +++ b/zig/src/families/flashnext/cuda_ple_nvfp4_test.zig @@ -0,0 +1,111 @@ +//! Synthetic mapped shards exercise the production gather, threading and chunk-history boundaries. +const std = @import("std"); +const core = @import("core"); +const ple = @import("cuda_ple.zig"); +const a = std.testing.allocator; +const io = std.testing.io; +const base = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"; + +fn fixture(dir: std.Io.Dir, include_scale: bool) !void { + const config = + \\{"text_config":{"num_hidden_layers":2,"hidden_size":32,"vocab_size":100,"ple_layer_ids":[2], + \\ "ngram_size":2,"heads_per_ngram":2,"ngram_vocab_size_base":13,"make_ngram_vocab_size_divisible_by":2, + \\ "split_ngram_parts":2,"ple_embed_dim":32,"eos_token_id":99}} + ; + try dir.writeFile(io, .{ .sub_path = "config.json", .data = config }); + const h = try ple.Hash.init(.{ .vocab = 100, .ngram = 2, .per = 2, .base = 13, .divisor = 2, .shards = 2, .seed = 1234, .eos = 99, .ple_dim = 32, .ple_index = 0 }); + var header: std.Io.Writer.Allocating = .init(a); + defer header.deinit(); + var index: std.Io.Writer.Allocating = .init(a); + defer index.deinit(); + var data: std.ArrayList(u8) = .empty; + defer data.deinit(a); + try header.writer.writeByte('{'); + try index.writer.writeAll("{\"weight_map\":{"); + for (0..2) |shard| { + const rows: usize = @intCast(h.starts[shard + 1] - h.starts[shard]); + for ([_][]const u8{ "weight", "weight_scale" }, 0..) |part, component| { + const cols: usize = if (component == 0) 8 else 1; + const dtype = if (component == 0) "U8" else "F8_E4M3"; + const begin = data.items.len; + for (0..rows) |row| { + const id: usize = @intCast(h.starts[shard] + @as(i64, @intCast(row))); + const code: u8 = if (component == 0) @intCast(id % 16 + ((id + 3) % 16) * 16) else 0x38; + try data.appendNTimes(a, code, cols); + } + if (begin > 0) { + try header.writer.writeByte(','); + try index.writer.writeByte(','); + } + try header.writer.print("\"{s}.shard_{d}.{s}\":{{\"dtype\":\"{s}\",\"shape\":[{d},{d}],\"data_offsets\":[{d},{d}]}}", .{ base, shard, part, dtype, rows, cols, begin, data.items.len }); + try index.writer.print("\"{s}.shard_{d}.{s}\":\"table.safetensors\"", .{ base, shard, part }); + } + } + if (include_scale) { + const begin = data.items.len; + const bytes = std.mem.toBytes(@as(u32, @bitCast(@as(f32, 0.375)))); + try data.appendSlice(a, &bytes); + try header.writer.print(",\"{s}.weight_scale_2\":{{\"dtype\":\"F32\",\"shape\":[],\"data_offsets\":[{d},{d}]}}", .{ base, begin, data.items.len }); + try index.writer.print(",\"{s}.weight_scale_2\":\"table.safetensors\"", .{base}); + } + try header.writer.writeByte('}'); + try index.writer.writeAll("}}"); + var image: std.Io.Writer.Allocating = .init(a); + defer image.deinit(); + const size = std.mem.toBytes(std.mem.nativeToLittle(u64, header.written().len)); + try image.writer.writeAll(&size); + try image.writer.writeAll(header.written()); + try image.writer.writeAll(data.items); + try dir.writeFile(io, .{ .sub_path = "table.safetensors", .data = image.written() }); + try dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.written() }); +} + +test "NVFP4 mapped PLE rows equal token decode across prompt chunks and threaded gathers" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture(tmp.dir, true); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(path); + var g = try ple.Ngram.init(a, io, path); + defer g.deinit(); + try std.testing.expectEqual(@as(usize, 1), g.files.items.len); + var tokens: [64]u32 = undefined; + for (&tokens, 0..) |*v, i| v.* = @intCast((i * 7) % 100); + var full: [64 * 32]u16 = undefined; + var xs_full: [64]f32 = undefined; + var history = [_]i64{99}; + try g.rows(&history, &tokens, &full, &xs_full); + var chunked: [64 * 32]u16 = undefined; + var xs_chunked: [64]f32 = undefined; + var chunk_hist = [_]i64{99}; + try g.rows(&chunk_hist, tokens[0..17], chunked[0 .. 17 * 32], xs_chunked[0..17]); + try g.rows(&chunk_hist, tokens[17..], chunked[17 * 32 ..], xs_chunked[17..]); + try std.testing.expectEqualSlices(u16, &full, &chunked); + try std.testing.expectEqualSlices(f32, &xs_full, &xs_chunked); + var single_hist = [_]i64{99}; + const values = [_]f32{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 }; + for (tokens, 0..) |token, row| { + var ids: [2]i64 = undefined; + try ple.lastIds(g.hash, &single_hist, token, &ids); + var one: [32]u16 = undefined; + var xs: [1]f32 = undefined; + try g.rows(&single_hist, &.{token}, &one, &xs); + try std.testing.expectEqualSlices(u16, full[row * 32 ..][0..32], &one); + for (ids, 0..) |id, head| for (0..16) |col| { + const code: usize = @intCast(@mod(id + @as(i64, if (col % 2 == 0) 0 else 3), 16)); + const v = values[code % 8] * @as(f32, if (code < 8) 1 else -1) * 0.375; + try std.testing.expectEqual(core.quant.convert.bf16(v), one[head * 16 + col]); + }; + } + try std.testing.expectEqualSlices(i64, &history, &single_hist); + try std.testing.expectEqualSlices(i64, &history, &chunk_hist); +} + +test "NVFP4 PLE requires its global table scale" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try fixture(tmp.dir, false); + const path = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(path); + try std.testing.expectError(error.MissingTensor, ple.Ngram.init(a, io, path)); +} From 3d56b2ac7487cd4ca7ec3341deddaff41895916a Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:23:49 -0700 Subject: [PATCH 062/115] Expose FP32 lane projection outputs before residual reduction --- zig/kernels/cuda/qmmf.cu | 36 +++++++ zig/src/cuda/qmmf.md | 19 ++++ zig/src/cuda/qmmf.zig | 40 ++++++-- zig/tests/cuda/main.zig | 3 + zig/tests/cuda/projection_f32_tests.zig | 127 ++++++++++++++++++++++++ 5 files changed, 214 insertions(+), 11 deletions(-) create mode 100644 zig/src/cuda/qmmf.md create mode 100644 zig/tests/cuda/projection_f32_tests.zig diff --git a/zig/kernels/cuda/qmmf.cu b/zig/kernels/cuda/qmmf.cu index 027b8ef630..f34fc02118 100644 --- a/zig/kernels/cuda/qmmf.cu +++ b/zig/kernels/cuda/qmmf.cu @@ -335,5 +335,41 @@ template __global__ void qmmf_kernel( template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void qmmf_kernel(const __nv_bfloat16* __restrict__, const unsigned char* __restrict__, const uint8_t* __restrict__, float, void* __restrict__, float* __restrict__, int, int, int, int, int, int, int); +template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); template __global__ void reduce_kernel(const float* __restrict__, void* __restrict__, long long, int, float); } diff --git a/zig/src/cuda/qmmf.md b/zig/src/cuda/qmmf.md new file mode 100644 index 0000000000..82f5c5a1ad --- /dev/null +++ b/zig/src/cuda/qmmf.md @@ -0,0 +1,19 @@ +# FP32 projection outputs + +`cuda.qmmf.Lane.matmulF32` has the same inputs, packed weights, stride, +shape-driven split-K plan and scratch contract as `matmul`. Its output is +`[M, N]` FP32, requiring `M * N * 4` bytes. It preserves FP32 output for a +later residual or tensor-rank sum; the caller chooses when to round to BF16. + +All five lane formats are supported: NVFP4, tensor-scaled FP8, MXFP8, +block-scaled FP8 and BF16. The change exposes the existing kernel template's +FP32 output path. It does not change weight packing or accumulation order. +Rebuild the qmmf fatbin together with the native runtime: older fatbins do +not contain the FP32 instances or reduction symbol. + +`tf-cuda-test projection-f32` compares each format against an independent +FP64 calculation, checks conversion to the existing BF16 result, and varies +row count, activation stride, partial output tiles, fused reduction, +cluster reduction and explicit split reduction. The fixture is synthetic; +model correctness and multi-device collective behavior require separate +qualification. diff --git a/zig/src/cuda/qmmf.zig b/zig/src/cuda/qmmf.zig index 3e2b4168a7..adfddf7124 100644 --- a/zig/src/cuda/qmmf.zig +++ b/zig/src/cuda/qmmf.zig @@ -66,9 +66,9 @@ pub fn plan(m: usize, n: usize, k: usize, clusters: bool) Plan { } /// qmmf.cu's instance names, per mode: row tiles 16, 32, 64 (each plain, then with a cluster), then the fused kernel. -fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool) [:0]const u8 { - const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb0ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; - return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(cluster), @intFromBool(fused) }); +fn symbol(comptime mode: Mode, comptime bm: u32, comptime cluster: bool, comptime fused: bool, comptime f32_out: bool) [:0]const u8 { + const name = "_ZN7tf_qmmf11qmmf_kernelILi{d}ELi{d}ELi64ELi1ELi4ELi4ELb{d}ELb{d}ELb{d}EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii"; + return std.fmt.comptimePrint(name, .{ @backingInt(mode), bm, @intFromBool(f32_out), @intFromBool(cluster), @intFromBool(fused) }); } pub const reduce_symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif"; @@ -78,22 +78,30 @@ pub const Lane = struct { mod: Module, fns: [modes.len][7]Function, // per mode: lane16, lane16c, lane32, lane32c, lane64, lane64c, fused reduce: Function, + f32_fns: [modes.len][7]Function, + reduce_f32: Function, clusters: bool, // sm_90 on: slices meet in a cluster's shared memory pub fn load(d: *const Driver, capability_major: u32) !Lane { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.qmmf); errdefer mod.unload(); - var l: Lane = .{ .mod = mod, .fns = undefined, .reduce = try mod.function(reduce_symbol), .clusters = capability_major >= 9 }; + var l: Lane = .{ .mod = mod, .fns = undefined, .f32_fns = undefined, .reduce = try mod.function(reduce_symbol), .reduce_f32 = try mod.function("_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif"), .clusters = capability_major >= 9 }; inline for (modes, 0..) |mode, mi| { inline for (row_tiles, 0..) |bm, ti| inline for (.{ false, true }, 0..) |cl, ci| { - const f = try mod.function(symbol(mode, bm, cl, false)); + const f = try mod.function(symbol(mode, bm, cl, false, false)); try f.allowDynamicShared(sharedBytes(mode, bm)); l.fns[mi][ti * 2 + ci] = f; + const fp32_fn = try mod.function(symbol(mode, bm, cl, false, true)); + try fp32_fn.allowDynamicShared(sharedBytes(mode, bm)); + l.f32_fns[mi][ti * 2 + ci] = fp32_fn; }; - const f = try mod.function(symbol(mode, 64, false, true)); + const f = try mod.function(symbol(mode, 64, false, true, false)); try f.allowDynamicShared(sharedBytes(mode, 64)); l.fns[mi][6] = f; + const fp32_fn = try mod.function(symbol(mode, 64, false, true, true)); + try fp32_fn.allowDynamicShared(sharedBytes(mode, 64)); + l.f32_fns[mi][6] = fp32_fn; } return l; } @@ -110,8 +118,17 @@ pub const Lane = struct { /// y [m, n] bf16 = x [m, k] bf16 (rows `ldx` apart) times the weight, fp32 sums; `part` holds partBytes. pub fn matmul(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulTo(s, x, ldx, m, w, y, part, false); + } + + /// The same accumulation and reduction, preserving FP32 outputs for a later residual or tensor-rank sum. + pub fn matmulF32(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64) !void { + return l.matmulTo(s, x, ldx, m, w, y, part, true); + } + + fn matmulTo(l: Lane, s: Stream, x: u64, ldx: usize, m: usize, w: Weight, y: u64, part: u64, f32_out: bool) !void { const p = plan(m, w.n, w.k, l.clusters); - const set = l.fns[modeIndex(w.mode)]; + const set = (if (f32_out) l.f32_fns else l.fns)[modeIndex(w.mode)]; const ti: usize = if (p.bm == 16) 0 else if (p.bm == 32) 1 else 2; const f = if (p.fused) set[6] else set[ti * 2 + @intFromBool(p.cluster)]; const rows_t = (m + p.bm - 1) / p.bm; @@ -144,7 +161,7 @@ pub const Lane = struct { r.add(@as(i64, @intCast(total))); r.add(@as(i32, @intCast(p.sk))); r.add(w.scale); - try launch_.launch(l.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); + try launch_.launch(if (f32_out) l.reduce_f32 else l.reduce, .{ .grid = .{ .x = @intCast((total + 255) / 256), .y = 1, .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &r); } } }; @@ -212,7 +229,8 @@ test "tile shared bytes, slices and row tiles as the Python host picks them" { } test "instance names are qmmf.cu's" { - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.mxfp8, 16, true, false)); - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false)); - try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.bf16, 64, false, true, true)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.mxfp8, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp8g, 16, true, false, false)); + try std.testing.expectEqualStrings("_ZN7tf_qmmf11qmmf_kernelILi0ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", symbol(.fp4, 64, false, true, false)); } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index c4a8291323..97a58caa1a 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -17,6 +17,7 @@ const chunk_costs = @import("chunk_costs.zig"); const grouped_tests = @import("grouped_tests.zig"); const fp8_tests = @import("fp8_tests.zig"); const mxfp8_tests = @import("mxfp8_tests.zig"); +const projection_f32_tests = @import("projection_f32_tests.zig"); const fp8_experts_tests = @import("fp8_experts_tests.zig"); const nvfp4_experts_tests = @import("nvfp4_experts_tests.zig"); const quant_ref_tests = @import("quant_ref_tests.zig"); @@ -70,6 +71,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) @@ -132,6 +134,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); diff --git a/zig/tests/cuda/projection_f32_tests.zig b/zig/tests/cuda/projection_f32_tests.zig new file mode 100644 index 0000000000..85b1160365 --- /dev/null +++ b/zig/tests/cuda/projection_f32_tests.zig @@ -0,0 +1,127 @@ +//! FP32 lane output before residual/rank combination: independent FP64, rounding and row/reduction invariance. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const reference = @import("quant_ref_tests.zig"); +const Rig = reference.RigT; +const bf = reference.helpers.bfOf; +const bf16 = reference.helpers.bf16Of; +const Mode = cuda.qmmf.Mode; + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0xf32_5eed); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + for ([_]Mode{ .fp4, .fp8, .mxfp8, .fp8g, .bf16 }) |mode| try projection(&r, &lane, mode); +} + +fn projection(r: *Rig, lane: *cuda.qmmf.Lane, mode: Mode) !void { + const n = 70; + const k = 1024; + const m = 256; + const ldx = k + 64; + const npad = cuda.fp8.padded(n); + const a = r.a; + const codes = try a.alloc(u8, n * k); + for (codes) |*v| v.* = reference.helpers.fp8CodeOf(r); + const scales = try a.alloc(u8, n * (k / 16)); + for (scales) |*v| v.* = if (mode == .mxfp8) 121 + r.rng.uintLessThan(u8, 8) else 0x20 + r.rng.uintLessThan(u8, 0x28); + const inv = try a.alloc(f32, ((n + 127) / 128) * (k / 128)); + for (inv) |*v| v.* = r.rng.float(f32) * 0.01 + 0.001; + const w16 = try r.bfs(n * k, 0.25); + const packed16 = try a.alloc(u16, npad * k); + const tiled_codes = std.mem.sliceAsBytes(packed16); + const global: f32 = 0.0123; + var w: cuda.qmmf.Weight = undefined; + switch (mode) { + .bf16 => { + cuda.qmmf.packBf16(packed16, w16, n, k); + w = cuda.modelopt.bf16Weight(try r.dev(tiled_codes), n, k); + }, + .fp4 => { + const words = try a.alloc(u32, npad * k / 8); + cuda.nvfp4.packWords(words, codes[0 .. n * k / 2], n, k); + const tiles = try a.alloc(u8, npad * (k / 16)); + cuda.nvfp4.packScales(tiles, scales, n, k); + w = cuda.nvfp4.weight(try r.dev(std.mem.sliceAsBytes(words)), try r.dev(tiles), global, n, k); + }, + .fp8, .fp8g, .mxfp8 => { + cuda.fp8.packCodes(tiled_codes[0 .. npad * k], codes, n, k); + const dw = try r.dev(tiled_codes[0 .. npad * k]); + w = switch (mode) { + .fp8 => cuda.modelopt.fp8Weight(dw, global, n, k), + .fp8g => blk: { + const tiles = try a.alloc(f32, npad * (k / 64)); + cuda.fp8.packScales(tiles, inv, n, k); + break :blk cuda.fp8.weight(dw, try r.dev(std.mem.sliceAsBytes(tiles)), n, k); + }, + .mxfp8 => blk: { + const tiles = try a.alloc(u8, npad * (k / 32)); + cuda.mxfp8.packScales(tiles, scales[0 .. n * k / 32], n, k); + break :blk cuda.mxfp8.weight(dw, try r.dev(tiles), n, k); + }, + else => unreachable, + }; + }, + } + const x = try a.alloc(u16, m * ldx); + @memset(x, 0x7fc0); + for (0..m) |row| { + for (x[row * ldx ..][0..k]) |*v| v.* = bf16((r.rng.float(f32) - 0.5) * 2); + } + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const out = try r.zeros(m * n * 4); + const out_bf16 = try r.zeros(m * n * 2); + const part = try r.zeros(8 * m * n * 4); + const want = try a.alloc(f64, m * n); + const tol = try a.alloc(f64, m * n); + for (0..m) |row| for (0..n) |col| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const code = codes[col * k + i]; + const weight: f64 = switch (mode) { + .bf16 => bf(w16[col * k + i]), + .fp8 => core.quant.fp8.e4m3(code) * global, + .fp8g => core.quant.fp8.e4m3(code) * inv[(col / 128) * (k / 128) + i / 128], + .mxfp8 => core.quant.fp8.e4m3(code) * std.math.ldexp(@as(f64, 1), @as(i32, scales[col * (k / 32) + i / 32]) - 127), + .fp4 => blk: { + const byte = codes[col * (k / 2) + i / 2]; + break :blk core.quant.nvfp4.e2m1(if (i % 2 == 0) byte & 15 else byte >> 4) * core.quant.fp8.e4m3(scales[col * (k / 16) + i / 16]) * global; + }, + }; + const v = bf(x[row * ldx + i]) * weight; + sum += v; + magnitude += @abs(v); + } + want[row * n + col] = sum; + tol[row * n + col] = magnitude * 2e-6 + @abs(sum) * 1e-6 + 1e-30; + }; + var baseline: ?[]u8 = null; + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters and (try r.gpu.ctx.capability()) >= 90; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + try lane.matmulF32(r.s, dx, ldx, rows, w, out, part); + try r.close("FP32 projection", out, true, want[0 .. rows * n], tol[0 .. rows * n]); + const got = try r.back(out, rows * n * 4); + if (baseline) |b| try check.sameBytes("FP32 row/reduction invariance", got, b[0..got.len]) else baseline = got; + try lane.matmul(r.s, dx, ldx, rows, w, out_bf16, part); + const rounded = try r.back(out_bf16, rows * n * 2); + for (0..rows * n) |i| { + const v = std.mem.readInt(u32, got[i * 4 ..][0..4], .little); + try check.expect(bf16(@bitCast(v)) == std.mem.readInt(u16, rounded[i * 2 ..][0..2], .little), "FP32 rounds once to the existing BF16 result", .{}); + } + } + } + check.pass("{t}: FP32 outputs preserve FP64 accuracy, BF16 rounding and row/reduction invariance", .{mode}); +} From 32be3e53fc584c39049a4fd55eb27cc445990798 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:41:08 -0700 Subject: [PATCH 063/115] Preserve dense connection weights and FP32 normalization in native CUDA passes --- zig/kernels/cuda/flashnext_hc.cu | 40 +++++ zig/src/families/flashnext/cuda_batch.zig | 18 +++ .../families/flashnext/cuda_batch_state.zig | 8 + zig/src/families/flashnext/cuda_glue.zig | 21 +++ zig/src/families/flashnext/cuda_launch.zig | 12 ++ .../families/flashnext/cuda_projection.zig | 71 +++++++++ zig/src/families/flashnext/cuda_prompt_wt.zig | 94 ++++++++--- zig/tests/cuda/flashnext_dense_hc_tests.zig | 146 ++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 9 files changed, 393 insertions(+), 20 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_projection.zig create mode 100644 zig/tests/cuda/flashnext_dense_hc_tests.zig diff --git a/zig/kernels/cuda/flashnext_hc.cu b/zig/kernels/cuda/flashnext_hc.cu index 3c85e54762..277dd6563b 100644 --- a/zig/kernels/cuda/flashnext_hc.cu +++ b/zig/kernels/cuda/flashnext_hc.cu @@ -392,3 +392,43 @@ extern "C" __global__ void fn_add_streams(const __nv_bfloat16* __restrict__ E, c OUT[at] = __float2bfloat16_rn(__fadd_rn(e, bf(HS[at]))); } } + +// Dense HC uses the same stored stream sums and FP32 gamma before one BF16 normalization rounding. +extern "C" __global__ void fn_hc_normalize(const __nv_bfloat16* H, const float* PSS, const float* SCALE, + __nv_bfloat16* OUT, float eps, int D, int SS) { + __shared__ float rinv; + const int r = blockIdx.x, stream = blockIdx.y, nc = D / 256; + if (threadIdx.x == 0) { + float total = 0.0f; + for (int c = 0; c < nc; ++c) total = __fadd_rn(total, PSS[(static_cast(r) * nc + c) * SS + stream]); + rinv = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(__fdiv_rn(total, static_cast(D)), eps))); + } + __syncthreads(); + for (int d = threadIdx.x; d < D; d += blockDim.x) { + const int col = stream * D + d; + const long long at = static_cast(r) * SS * D + col; + OUT[at] = __float2bfloat16_rn(__fmul_rn(__fmul_rn(bf(H[at]), rinv), SCALE[col])); + } +} + +// Dense up-projection values are already BF16: sigmoid, product and stream mean keep their original rounding seams. +extern "C" __global__ void fn_hc_dense_mix(const __nv_bfloat16* UP, const __nv_bfloat16* NORMED, + __nv_bfloat16* MIXED, float* XSM, int D, int SS) { + __shared__ float row[32]; + const int r = blockIdx.x, group = blockIdx.y, d = group * 32 + threadIdx.x; + float total = 0.0f; + for (int stream = 0; stream < SS; ++stream) { + const long long at = (static_cast(r) * SS + stream) * D + d; + const float value = rbf(__fmul_rn(bsig(bf(UP[at])), bf(NORMED[at]))); + total = __fadd_rn(total, value); + } + const float mixed = rbf(__fdiv_rn(total, static_cast(SS))); + MIXED[static_cast(r) * D + d] = __float2bfloat16_rn(mixed); + row[threadIdx.x] = mixed; + __syncthreads(); + if (threadIdx.x == 0) { + float sum = 0.0f; + for (int i = 0; i < 32; ++i) sum = __fadd_rn(sum, row[i]); + XSM[static_cast(r) * (D / 32) + group] = sum; + } +} diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 0686c0f688..52d4aadd58 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -104,6 +104,14 @@ pub fn Pass(comptime Tri: type) type { /// qmm.cu's lane tile over m rows: out bf16 [m, n] (sk 1) or the unreduced fp32 slices [sk, m, n] in part. pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !void { + if (w.native) |native| { + if (sk == 1) return self.k.native_lane.matmul(self.stream, x, w.k, m, native, out, self.s.native_part); + // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. + try self.k.native_lane.matmulF32(self.stream, x, w.k, m, native, part, self.s.native_part); + const bytes = m * w.n * 4; + if (sk > 1) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), self.stream.handle), "cuMemsetD8Async"); + return; + } const bm: usize = 16; const bn: usize = 64; const k = w.k; @@ -184,6 +192,16 @@ pub fn Pass(comptime Tri: type) type { const norm = s.norm + r0 * wide * 2; const act = s.act + r0 * low * 2; const xs_act = s.xs_act + r0 * (low / 32) * 4; + if (hc_w.down.native) |down| { + const up = hc_w.up.native orelse return error.UnexpectedTensor; + try self.tri.hcNormalize(h, pss, hc_w.scale.ptr, norm, eps, n, dims, streams); + try self.k.native_lane.matmulF32(self.stream, norm, wide, n, down, s.hc_part, s.native_part); + const inj_ptr = if (hc_w.has_inj == 1) inj else act; + try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, 1, n, streams, low, hc_w.n_down, hc_w.has_inj); + try self.k.native_lane.matmul(self.stream, act, low, n, up, s.hc_up, s.native_part); + try self.tri.hcDenseMix(s.hc_up, norm, s.mixed + r0 * dims * 2, s.xsm + r0 * (dims / 32) * 4, n, dims, streams); + return; + } try self.tri.hcDown(h, pss, hc_w.scale.ptr, norm, hc_w.down.w.ptr, hc_w.down.s.ptr, hc_w.down.b.ptr, s.dn, s.hc_part, eps, n, hc_w.n_down, wide, dims, dims / 256, streams, hc_sk); const inj_ptr = if (hc_w.has_inj == 1) inj else act; try self.tri.hcReduceAct(s.hc_part, act, xs_act, inj_ptr, hc_sk, n, streams, low, hc_w.n_down, hc_w.has_inj); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index c6456a0faf..05f74fffa8 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -98,6 +98,7 @@ pub const Kernels = struct { router: cuda.grouped.Router, moe: cuda.experts.RowsExperts, lane: cuda.Function, + native_lane: cuda.qmmf.Lane, front: cuda.Function, chain: cuda.Function, chain_to: cuda.Function, @@ -122,6 +123,8 @@ pub const Kernels = struct { try k.lane.allowDynamicShared(qmm_smem); k.router = try cuda.grouped.Router.resolve(k.experts); k.moe = try cuda.experts.RowsExperts.resolve(k.rows); + k.native_lane = try cuda.qmmf.Lane.load(d, (try d.capability(0)) / 10); + errdefer k.native_lane.unload(); k.sum = try k.batch.function("fn_sum_slices"); k.shift = try k.batch.function("fn_shift_window"); k.select = try k.batch.function("fn_qsa_select"); @@ -133,6 +136,7 @@ pub const Kernels = struct { } pub fn deinit(k: *Kernels) void { + k.native_lane.unload(); k.batch.unload(); k.rows.unload(); k.experts.unload(); @@ -200,6 +204,8 @@ pub const Scratch = struct { norm: u64 = 0, dn: u64 = 0, hc_part: u64 = 0, + hc_up: u64 = 0, + native_part: u64 = 0, act: u64 = 0, xs_act: u64 = 0, inj: u64 = 0, @@ -332,6 +338,8 @@ pub const Scratch = struct { s.norm = try s.add(gpa, d, cap * wide * 2); s.dn = try s.add(gpa, d, cap * (low + 4) * 2); s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); + s.hc_up = try s.add(gpa, d, cap * wide * 2); + s.native_part = try s.add(gpa, d, 8 * cap * dims * 4); s.act = try s.add(gpa, d, cap * low * 2); s.xs_act = try s.add(gpa, d, cap * (low / 32) * 4); s.inj = try s.add(gpa, d, cap * streams * 2); diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index 1c7683e4dd..7348ea6aea 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -27,6 +27,8 @@ const Fns = struct { wb0: cuda.Function, wb_moe: cuda.Function, wb_slices: cuda.Function, + normalize: cuda.Function, + dense_mix: cuda.Function, down: cuda.Function, reduce_act: cuda.Function, upmix: cuda.Function, @@ -91,6 +93,8 @@ pub const Glue = struct { .wb0 = try g.hc.function("fn_hc_writeback0"), .wb_moe = try g.hc.function("fn_hc_writeback_moe"), .wb_slices = try g.hc.function("fn_hc_writeback_slices"), + .normalize = try g.hc.function("fn_hc_normalize"), + .dense_mix = try g.hc.function("fn_hc_dense_mix"), .down = try g.hc.function("fn_hc_down"), .reduce_act = try g.hc.function("fn_hc_reduce_act"), .upmix = try g.hc.function("fn_hc_upmix"), @@ -204,6 +208,23 @@ pub const Glue = struct { try go(g.f.scores, s, .{ rows, cdiv(blocks, 64), 1 }, 256, 0, &a); } + pub fn hcNormalize(g: *const Glue, s: cuda.Stream, h: u64, pss: u64, scale: u64, out: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + var args: cuda.Args = .{}; + for ([_]u64{ h, pss, scale, out }) |v| args.add(v); + args.add(eps); + args.add(int(dims)); + args.add(int(streams)); + try cuda.launch.launch(g.f.normalize, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(streams) }, .block = .{ .x = 256 } }, s, &args); + } + + pub fn hcDenseMix(g: *const Glue, s: cuda.Stream, up: u64, norm: u64, mixed: u64, sums: u64, rows: usize, dims: usize, streams: usize) !void { + var args: cuda.Args = .{}; + for ([_]u64{ up, norm, mixed, sums }) |v| args.add(v); + args.add(int(dims)); + args.add(int(streams)); + try cuda.launch.launch(g.f.dense_mix, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(dims / 32) }, .block = .{ .x = 32 } }, s, &args); + } + pub fn hcReduceAct(g: *const Glue, s: cuda.Stream, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { const cols = low + has_inj * streams; if (cols > 1024 or sk > 32) return error.RowShape; diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index bcc9ca18ac..c162105317 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); +pub const projection = @import("cuda_projection.zig"); pub const hc = @import("cuda_hc.zig"); pub const qmm = @import("cuda_qmm.zig"); pub const gdn = @import("cuda_gdn.zig"); @@ -199,6 +200,17 @@ pub const Tri = struct { }, &.{ ci("HI", 4), ci("DI", 128), ci("RATIO", 4), ci("TOP", 512), ci("BB", 64) }); } + /// Dense connections always use native glue; a captured affine projection cannot replace these operators. + pub fn hcNormalize(t: Tri, h: u64, pss: u64, scale: u64, out: u64, eps: f32, rows: usize, dims: usize, streams: usize) !void { + const g = t.own orelse return error.NativeGlueRequired; + return g.hcNormalize(t.s, h, pss, scale, out, eps, rows, dims, streams); + } + + pub fn hcDenseMix(t: Tri, up: u64, norm: u64, mixed: u64, sums: u64, rows: usize, dims: usize, streams: usize) !void { + const g = t.own orelse return error.NativeGlueRequired; + return g.hcDenseMix(t.s, up, norm, mixed, sums, rows, dims, streams); + } + /// glue.hc_reduce_act: split-K sum fused with hc_act. M folds to a constexpr when it is 1. pub fn hcReduceAct(t: Tri, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { if (t.glue("_hc_reduce_act")) |g| return g.hcReduceAct(t.s, part, act, xs, inj, sk, rows, streams, low, ndn, has_inj); diff --git a/zig/src/families/flashnext/cuda_projection.zig b/zig/src/families/flashnext/cuda_projection.zig new file mode 100644 index 0000000000..c0b96adb3b --- /dev/null +++ b/zig/src/families/flashnext/cuda_projection.zig @@ -0,0 +1,71 @@ +//! Owned affine or native lane projections; packed device views borrow these exact buffers. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +pub const Projection = struct { + w: cuda.DeviceBuffer, + s: cuda.DeviceBuffer, + b: cuda.DeviceBuffer, + n: usize, + k: usize, + native: ?cuda.qmmf.Weight = null, + + pub fn free(self: *Projection) void { + self.w.free(); + self.s.free(); + self.b.free(); + } + + pub fn empty(d: *const cuda.Driver) Projection { + const z: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + return .{ .w = z, .s = z, .b = z, .n = 0, .k = 0 }; + } +}; + +/// Repack the checkpoint's BF16 bits without rounding or widening the resident weight. +pub fn dense(gpa: std.mem.Allocator, driver: *const cuda.Driver, t: core.safetensors.Tensor) !Projection { + if (t.rank != 2 or t.dtype != .bf16 or t.dim(0) == 0 or t.dim(1) == 0 or t.dim(1) % 64 != 0) + return error.UnexpectedTensor; + const n = t.dim(0); + const k = t.dim(1); + const count = try std.math.mul(usize, n, k); + if (t.bytes.len != try std.math.mul(usize, count, 2)) return error.UnexpectedTensor; + const input = try gpa.alloc(u16, count); + defer gpa.free(input); + for (input, 0..) |*v, i| v.* = std.mem.readInt(u16, t.bytes[i * 2 ..][0..2], .little); + const tiled = try gpa.alloc(u16, cuda.nvfp4.padded(n) * k); + defer gpa.free(tiled); + cuda.qmmf.packBf16(tiled, input, n, k); + var p = Projection.empty(driver); + p.w = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(tiled)); + p.n = n; + p.k = k; + p.native = cuda.modelopt.bf16Weight(p.w.ptr, n, k); + return p; +} + +/// E4M3 codes and UE8M0 group-32 scales are permuted only; calibration metadata is not applied as another scale. +pub fn mxfp8(gpa: std.mem.Allocator, driver: *const cuda.Driver, w: core.safetensors.Tensor, s: core.safetensors.Tensor) !Projection { + if (w.rank != 2 or w.dtype != .f8_e4m3 or w.dim(0) == 0 or w.dim(1) == 0 or w.dim(1) % 64 != 0) + return error.UnexpectedTensor; + const n = w.dim(0); + const k = w.dim(1); + const count = try std.math.mul(usize, n, k); + if (w.bytes.len != count or !s.is(.u8, &.{ n, k / 32 }) or s.bytes.len != count / 32) + return error.UnexpectedTensor; + const codes = try gpa.alloc(u8, cuda.mxfp8.padded(n) * k); + defer gpa.free(codes); + const scales = try gpa.alloc(u8, codes.len / 32); + defer gpa.free(scales); + cuda.mxfp8.packCodes(codes, w.bytes, n, k); + cuda.mxfp8.packScales(scales, s.bytes, n, k); + var p = Projection.empty(driver); + p.w = try cuda.DeviceBuffer.fromHost(driver, codes); + errdefer p.w.free(); + p.s = try cuda.DeviceBuffer.fromHost(driver, scales); + p.n = n; + p.k = k; + p.native = cuda.mxfp8.weight(p.w.ptr, p.s.ptr, n, k); + return p; +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index 127afe2128..bea01330d0 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -109,7 +109,9 @@ pub const Store = struct { self.nheld = 0; } - pub fn hold(self: *Store, key: []const u8) !Tensor { + pub fn hold(self: *Store, requested: []const u8) !Tensor { + var mapped: [512]u8 = undefined; + const key = try self.resolveKey(requested, &mapped); const file = try self.shard(key); for (self.held[0..self.nheld]) |*h| { if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; @@ -136,6 +138,23 @@ pub const Store = struct { self.rel = &.{}; } + pub fn resolveKey(self: *const Store, requested: []const u8, buf: []u8) ![]const u8 { + const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + if (wm.object.contains(requested)) return requested; + const aliases = .{ + .{ "language_model.model.", "model.language_model." }, + .{ "language_model.lm_head.", "lm_head." }, + .{ "language_model.mtp.", "mtp." }, + }; + inline for (aliases) |pair| { + if (std.mem.startsWith(u8, requested, pair[0])) { + const name = try std.fmt.bufPrint(buf, "{s}{s}", .{ pair[1], requested[pair[0].len..] }); + if (wm.object.contains(name)) return name; + } + } + return error.MissingTensor; + } + fn shard(self: *Store, key: []const u8) ![]const u8 { const wm = self.parsed.value.object.get("weight_map") orelse return error.MissingTensor; const v = wm.object.get(key) orelse { @@ -154,7 +173,9 @@ pub const Store = struct { self.file = try Shard.load(self.gpa, self.io, self.path); } - pub fn tensor(self: *Store, key: []const u8) !Tensor { + pub fn tensor(self: *Store, requested: []const u8) !Tensor { + var mapped: [512]u8 = undefined; + const key = try self.resolveKey(requested, &mapped); try self.use(key); return self.file.?.get(key) orelse { std.debug.print("absent {s}\n", .{key}); @@ -179,24 +200,7 @@ pub const Face = struct { routed: bool, }; -pub const Q4 = struct { - w: cuda.DeviceBuffer, - s: cuda.DeviceBuffer, - b: cuda.DeviceBuffer, - n: usize, - k: usize, - - pub fn free(self: *Q4) void { - self.w.free(); - self.s.free(); - self.b.free(); - } - - pub fn empty(d: *cuda.Driver) Q4 { - const z: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; - return .{ .w = z, .s = z, .b = z, .n = 0, .k = 0 }; - } -}; +pub const Q4 = @import("cuda_projection.zig").Projection; pub const HcW = struct { scale: cuda.DeviceBuffer, @@ -254,6 +258,8 @@ pub fn cat(gpa: std.mem.Allocator, parts: []const []const u8) ![]u8 { pub fn hostOnly(name: []const u8) bool { return std.mem.indexOf(u8, name, ".ple_embedding.ngram_embedding.") != null or std.mem.startsWith(u8, name, "language_model.mtp.") or + std.mem.startsWith(u8, name, "mtp.") or + std.mem.startsWith(u8, name, "model.visual.") or std.mem.startsWith(u8, name, "vision_tower."); } @@ -284,6 +290,7 @@ pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { pub fn asF32(gpa: std.mem.Allocator, t: Tensor) ![]f32 { if (t.rank != 1) return error.UnexpectedTensor; const out = try gpa.alloc(f32, t.dim(0)); + errdefer gpa.free(out); if (t.dtype == .f32 and t.bytes.len == out.len * 4) { for (out, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); } else if (t.dtype == .bf16 and t.bytes.len == out.len * 2) { @@ -330,6 +337,9 @@ pub fn join3(gpa: std.mem.Allocator, store: *Store, buf: []u8, prefix: []const u pub fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { var buf: [180]u8 = undefined; + const down_name = try std.fmt.bufPrint(&buf, "{s}.input_mix_weight_down.weight", .{prefix}); + const down_tensor = try store.tensor(down_name); + if (down_tensor.dtype == .bf16) return loadDenseHc(gpa, driver, store, prefix, inject); const down = try join3(gpa, store, &buf, prefix, ".input_mix_weight_down"); defer gpa.free(down.w); defer gpa.free(down.s); @@ -379,6 +389,33 @@ pub fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefi return .{ .scale = scale_b, .down = down_q, .up = up_q, .n_down = n, .has_inj = @intFromBool(inject) }; } +fn loadDenseHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefix: []const u8, inject: bool) !HcW { + const projection = @import("cuda_projection.zig"); + var buf: [200]u8 = undefined; + const n = low + if (inject) streams else @as(usize, 0); + const data = try gpa.alloc(u8, n * wide * 2); + defer gpa.free(data); + const down = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.input_mix_weight_down.weight", .{prefix})); + if (!down.is(.bf16, &.{ low, wide }) or down.bytes.len != low * wide * 2) return error.UnexpectedTensor; + @memcpy(data[0..down.bytes.len], down.bytes); + if (inject) { + const inj = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.block_inject_weight.weight", .{prefix})); + if (!inj.is(.bf16, &.{ streams, wide }) or inj.bytes.len != streams * wide * 2) return error.UnexpectedTensor; + @memcpy(data[low * wide * 2 ..], inj.bytes); + } + var dq = try projection.dense(gpa, driver, .{ .dtype = .bf16, .rank = 2, .shape = .{ n, wide, 1, 1, 1 }, .bytes = data }); + errdefer dq.free(); + const up = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.input_mix_weight_up.weight", .{prefix})); + if (!up.is(.bf16, &.{ wide, low })) return error.UnexpectedTensor; + var uq = try projection.dense(gpa, driver, up); + errdefer uq.free(); + const gamma = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix}))); + defer gpa.free(gamma); + if (gamma.len != wide) return error.UnexpectedTensor; + const scale = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(gamma)); + return .{ .scale = scale, .down = dq, .up = uq, .n_down = n, .has_inj = @intFromBool(inject) }; +} + pub fn stitch(gpa: std.mem.Allocator, store: *Store, prefix: []const u8, names: []const []const u8) !struct { w: []u8, s: []u8, b: []u8, n: usize } { var buf: [180]u8 = undefined; var ws: [6][]u8 = undefined; @@ -408,3 +445,20 @@ test "the n-gram tables, the draft head and the vision tower stay off the device try std.testing.expect(!hostOnly("language_model.model.layers.0.mlp.switch_mlp.up_proj.weight")); try std.testing.expect(!hostOnly("language_model.lm_head.weight")); } + +test "canonical checkpoint aliases preserve exact stored namespace and prefer a direct key" { + const a = std.testing.allocator; + const parsed = try std.json.parseFromSlice(std.json.Value, a, + \\{"weight_map":{"model.language_model.layers.0.x.weight":"a","lm_head.weight":"b","mtp.x":"c","language_model.model.direct":"d","model.language_model.direct":"e"}} + , .{}); + var store: Store = .{ .gpa = a, .io = std.testing.io, .dir = "", .parsed = parsed }; + defer store.deinit(); + var buf: [512]u8 = undefined; + try std.testing.expectEqualStrings("model.language_model.layers.0.x.weight", try store.resolveKey("language_model.model.layers.0.x.weight", &buf)); + try std.testing.expectEqualStrings("lm_head.weight", try store.resolveKey("language_model.lm_head.weight", &buf)); + try std.testing.expectEqualStrings("mtp.x", try store.resolveKey("language_model.mtp.x", &buf)); + try std.testing.expectEqualStrings("language_model.model.direct", try store.resolveKey("language_model.model.direct", &buf)); + try std.testing.expectError(error.MissingTensor, store.resolveKey("language_model.model.missing", &buf)); + try std.testing.expect(hostOnly("model.visual.weight")); + try std.testing.expect(hostOnly("mtp.layers.0.weight")); +} diff --git a/zig/tests/cuda/flashnext_dense_hc_tests.zig b/zig/tests/cuda/flashnext_dense_hc_tests.zig new file mode 100644 index 0000000000..7afa56cc9d --- /dev/null +++ b/zig/tests/cuda/flashnext_dense_hc_tests.zig @@ -0,0 +1,146 @@ +//! Dense hyper-connection normalization, projections and mixing against an independent FP64 chain. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const bf = rig.bf; +const rbf = rig.rbf; +const d = 2560; +const streams = 4; +const wide = d * streams; +const low = 320; +const ndn = low + streams; +const most = 17; + +fn tensor(values: []const u16, n: usize, k: usize) core.safetensors.Tensor { + return .{ .dtype = .bf16, .rank = 2, .shape = .{ n, k, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(values) }; +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var prng = std.Random.DefaultPrng.init(0xde05e); + var r = try rig.Rig.init(gpu, arena.allocator(), prng.random()); + defer r.deinit(); + var g = try flash.Glue.load(gpu.d); + defer g.deinit(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + const h = r.bfs(most * wide, 1); + const gamma = r.f32s(wide, 0.5, 1.5); + const dw = r.bfs(ndn * wide, 0.01); + const uw = r.bfs(wide * low, 0.03); + var down = try flash.projection.dense(r.a, gpu.d, tensor(dw, ndn, wide)); + defer down.free(); + var up = try flash.projection.dense(r.a, gpu.d, tensor(uw, wide, low)); + defer up.free(); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const dg = try r.dev(std.mem.sliceAsBytes(gamma)); + const pss = try r.out(most * (d / 256) * streams * 4); + try g.hcWriteback(r.s, dh, dh, pss, 0, 0, 0, 0, most * d, most, d, streams, 0, 1, 1, 1); + const sums = try r.back(f32, pss, most * (d / 256) * streams); + const norm = try r.out(most * wide * 2); + const projected = try r.out(most * ndn * 4); + const act = try r.out(most * low * 2); + const xs = try r.out(most * (low / 32) * 4); + const inj = try r.out(most * streams * 2); + const up_out = try r.out(most * wide * 2); + const mixed = try r.out(most * d * 2); + const xsm = try r.out(most * (d / 32) * 4); + const part = try r.out(8 * most * ndn * 4); + var norms: []u16 = undefined; + var dots: []f32 = undefined; + var acts: []u16 = undefined; + var injs: []u16 = undefined; + var ups: []u16 = undefined; + var mixes: []u16 = undefined; + var mix_sums: []f32 = undefined; + for ([_]usize{ most, 1, 7 }) |rows| { + try g.hcNormalize(r.s, dh, pss, dg, norm, 1e-6, rows, d, streams); + try lane.matmulF32(r.s, norm, wide, rows, down.native.?, projected, part); + try g.hcReduceAct(r.s, projected, act, xs, inj, 1, rows, streams, low, ndn, 1); + try lane.matmul(r.s, act, low, rows, up.native.?, up_out, part); + try g.hcDenseMix(r.s, up_out, norm, mixed, xsm, rows, d, streams); + if (rows == most) { + norms = try r.back(u16, norm, most * wide); + dots = try r.back(f32, projected, most * ndn); + acts = try r.back(u16, act, most * low); + injs = try r.back(u16, inj, most * streams); + ups = try r.back(u16, up_out, most * wide); + mixes = try r.back(u16, mixed, most * d); + mix_sums = try r.back(f32, xsm, most * (d / 32)); + } else { + try r.sameRows("dense HC normalize", u16, norm, norms[0 .. rows * wide]); + try r.sameRows("dense HC down FP32", f32, projected, dots[0 .. rows * ndn]); + try r.sameRows("dense HC activation", u16, act, acts[0 .. rows * low]); + try r.sameRows("dense HC inject", u16, inj, injs[0 .. rows * streams]); + try r.sameRows("dense HC up BF16", u16, up_out, ups[0 .. rows * wide]); + try r.sameRows("dense HC mix", u16, mixed, mixes[0 .. rows * d]); + try r.sameRows("dense HC group sums", f32, xsm, mix_sums[0 .. rows * (d / 32)]); + } + } + for (0..most) |row| { + for (0..streams) |st| { + var total: f64 = 0; + for (0..d / 256) |c| total += sums[(row * (d / 256) + c) * streams + st]; + const inv = 1 / @sqrt(total / d + 1e-6); + for (0..d) |col| { + const at = row * wide + st * d + col; + const want = bf(h[at]) * inv * gamma[st * d + col]; + try r.within("dense norm FP32 gamma", bf(norms[at]), want, rig.ulps(want, 0.5) + @abs(want) * 0x1p-20); + } + } + for (0..ndn) |col| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..wide) |i| { + const v = bf(norms[row * wide + i]) * bf(dw[col * wide + i]); + sum += v; + magnitude += @abs(v); + } + try r.within("dense down FP32", dots[row * ndn + col], sum, magnitude * 2e-6 + @abs(sum) * 1e-6 + 1e-30); + const v = rbf(rbf(dots[row * ndn + col]) / streams); + if (col < low) { + const want = rbf(v * rig.sigmoid(v)); + try r.within("dense activation", bf(acts[row * low + col]), want, rig.ulps(want, 1)); + } else { + const want = rbf(2 * rbf(rig.sigmoid(v))); + try r.within("dense inject", bf(injs[row * streams + col - low]), want, rig.ulps(want, 1)); + } + } + for (0..wide) |col| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..low) |i| { + const v = bf(acts[row * low + i]) * bf(uw[col * low + i]); + sum += v; + magnitude += @abs(v); + } + try r.within("dense up BF16", bf(ups[row * wide + col]), sum, rig.ulps(sum, 0.5) + magnitude * 2e-6); + } + for (0..d / 32) |group| { + var sum: f64 = 0; + var magnitude: f64 = 0; + for (0..32) |i| { + const col = group * 32 + i; + var total: f64 = 0; + var terms: f64 = 0; + for (0..streams) |st| { + const at = row * wide + st * d + col; + const term = rbf(rbf(rig.sigmoid(bf(ups[at]))) * bf(norms[at])); + total += term; + terms += @abs(term); + } + const want = rbf(total / streams); + const got = bf(mixes[row * d + col]); + try r.within("dense stream mix", got, want, rig.ulps(terms / streams, 1) + rig.ulps(want, 0.5)); + sum += got; + magnitude += @abs(got); + } + try r.within("dense group sums", mix_sums[row * (d / 32) + group], sum, rig.sumTol(magnitude, 32)); + } + } + r.report("dense HC: FP32 gamma, BF16 weights, FP32 down and one-round up; rows 1/7/17"); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 97a58caa1a..6a324a358b 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -33,6 +33,7 @@ const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); +const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); @@ -58,6 +59,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) \\ window-profile MODEL WIDTHS IDS_FILE a decode window's GPU ms by kernel class (WIDTHS like 1,4,16), beside its graph @@ -147,6 +149,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); if (std.mem.eql(u8, cmd, "pocket")) From 21871114a52e82588d3e7f9deb886b2b52dfb524 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:50:30 -0700 Subject: [PATCH 064/115] Support exact BF16 token rows and dense vocabulary projections --- zig/kernels/cuda/embed_rows.cu | 13 +++++++ zig/src/cuda/embed_rows.zig | 14 +++++++- zig/src/families/flashnext/cuda_batch.zig | 6 +++- .../families/flashnext/cuda_batch_state.zig | 1 + zig/src/families/flashnext/cuda_embed.zig | 11 ++++++ zig/src/families/flashnext/cuda_glue.zig | 4 +++ zig/src/families/flashnext/cuda_launch.zig | 5 +++ zig/src/families/flashnext/cuda_mtp_ops.zig | 6 +++- zig/src/families/flashnext/cuda_prompt.zig | 34 ++++++++++++------- zig/src/families/flashnext/cuda_rows.zig | 1 + zig/tests/cuda/norm_rows_tests.zig | 19 +++++++++-- 11 files changed, 97 insertions(+), 17 deletions(-) diff --git a/zig/kernels/cuda/embed_rows.cu b/zig/kernels/cuda/embed_rows.cu index ca5f2602d2..664e6a41ec 100644 --- a/zig/kernels/cuda/embed_rows.cu +++ b/zig/kernels/cuda/embed_rows.cu @@ -43,3 +43,16 @@ extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_bf16(const int* _ if (8 * c >= d) return; *reinterpret_cast(out + r * d + 8 * c) = *reinterpret_cast(w + static_cast(ids[r]) * d + 8 * c); } + +// Exact BF16 rows repeated across residual streams; no floating-point conversion. +extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_bf16_copies(const int* __restrict__ ids, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ out, int d, + int copies) { + const int64_t r = blockIdx.x; + const int c = blockIdx.y * NT + threadIdx.x; + if (8 * c >= d) return; + const uint4 value = *reinterpret_cast(w + static_cast(ids[r]) * d + 8 * c); + for (int k = 0; k < copies; ++k) + *reinterpret_cast(out + (r * copies + k) * d + 8 * c) = value; +} diff --git a/zig/src/cuda/embed_rows.zig b/zig/src/cuda/embed_rows.zig index 973b658f7c..1fc5cda5b7 100644 --- a/zig/src/cuda/embed_rows.zig +++ b/zig/src/cuda/embed_rows.zig @@ -16,12 +16,13 @@ pub const EmbedRows = struct { mod: Module, affine4_fn: Function, bf16_fn: Function, + bf16_copies_fn: Function, pub fn load(d: *const Driver) !EmbedRows { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.embed_rows); errdefer mod.unload(); - return .{ .mod = mod, .affine4_fn = try mod.function("tf_embed_rows_affine4"), .bf16_fn = try mod.function("tf_embed_rows_bf16") }; + return .{ .mod = mod, .affine4_fn = try mod.function("tf_embed_rows_affine4"), .bf16_fn = try mod.function("tf_embed_rows_bf16"), .bf16_copies_fn = try mod.function("tf_embed_rows_bf16_copies") }; } pub fn unload(g: *EmbedRows) void { @@ -52,4 +53,15 @@ pub const EmbedRows = struct { a.add(@as(i32, @intCast(d))); try launch_.launch(g.bf16_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); } + + /// Copy exact BF16 bits into out [rows, copies, d], including signed zeros and NaN payloads. + pub fn bf16Copies(g: *const EmbedRows, s: Stream, ids: u64, w: u64, out: u64, rows: usize, d: usize, copies: usize) !void { + if (d == 0 or d % 8 != 0 or d > std.math.maxInt(i32) or copies == 0 or copies > std.math.maxInt(i32) or + rows > std.math.maxInt(u32) or out % 16 != 0 or w % 16 != 0) return error.RowShape; + if (rows == 0) return; + var a: launch_.Args = .{}; + for ([_]u64{ ids, w, out }) |v| a.add(v); + for ([_]usize{ d, copies }) |v| a.add(@as(i32, @intCast(v))); + try launch_.launch(g.bf16_copies_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((d / 8 + threads - 1) / threads), .z = 1 }, .block = .{ .x = threads } }, s, &a); + } }; diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 52d4aadd58..c35ac83f9c 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -451,7 +451,11 @@ pub fn Pass(comptime Tri: type) type { var thread: ?std.Thread = try std.Thread.spawn(.{}, Gather.main, .{&gather}); defer if (thread) |th| th.join(); try self.stage(ids); - try self.tri.embed(s.ids, self.m.table_w, self.m.table_s, self.m.table_b, s.h, dims, streams, n); + if (self.m.table_dense) { + try self.tri.embedDense(s.ids, self.m.table_w, s.h, dims, streams, n); + } else { + try self.tri.embed(s.ids, self.m.table_w, self.m.table_s, self.m.table_b, s.h, dims, streams, n); + } for (self.m.layers, 0..) |*lw, layer| { const keep: LayerKeep = if (self.hooks.keep) |kp| kp.layers[layer] else .{}; if (layer == ple_layer) { diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 05f74fffa8..03650116ab 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -445,6 +445,7 @@ pub const Model = struct { table_w: u64, table_s: u64, table_b: u64, + table_dense: bool = false, layers: []const Layer, mixer: HcW, head: Q4, diff --git a/zig/src/families/flashnext/cuda_embed.zig b/zig/src/families/flashnext/cuda_embed.zig index 0b31f38a69..453ba86271 100644 --- a/zig/src/families/flashnext/cuda_embed.zig +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -82,6 +82,7 @@ pub const Table = struct { b: cuda.DeviceBuffer, rows: usize, dims: usize, + dense: bool = false, /// Copies the three tensors. `dims` is hidden size (eight values per stored word). pub fn upload(d: *const cuda.Driver, weight: Tensor, scales: Tensor, biases: Tensor) !Table { @@ -100,6 +101,16 @@ pub const Table = struct { return .{ .w = w, .s = s, .b = b, .rows = rows, .dims = dims }; } + /// Own a dense table without converting any BF16 values. + pub fn uploadDense(d: *const cuda.Driver, weight: Tensor) !Table { + if (weight.dtype != .bf16 or weight.rank != 2 or weight.dim(0) == 0 or weight.dim(1) == 0 or weight.dim(1) % 8 != 0) + return error.UnexpectedTensor; + const count = try std.math.mul(usize, weight.dim(0), weight.dim(1)); + if (weight.bytes.len != try std.math.mul(usize, count, 2)) return error.UnexpectedTensor; + const zero: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + return .{ .w = try cuda.DeviceBuffer.fromHost(d, weight.bytes), .s = zero, .b = zero, .rows = weight.dim(0), .dims = weight.dim(1), .dense = true }; + } + pub fn deinit(self: *Table) void { self.w.free(); self.s.free(); diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index 7348ea6aea..efa304177d 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -141,6 +141,10 @@ pub const Glue = struct { try g.emb.affine4Groups(s, ids, w, sc, b, out, rows, dims, 32, copies); } + pub fn embedDense(g: *const Glue, s: cuda.Stream, ids: u64, w: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + try g.emb.bf16Copies(s, ids, w, out, rows, dims, copies); + } + /// One RMS norm over each row's `dims` features (rows `x_stride` apart), bf16 out and 32-group sums (core norms). pub fn rmsnorm(g: *const Glue, s: cuda.Stream, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { if (group != dims) return error.RowShape; // the head norms each row whole diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index c162105317..65e9024d06 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -82,6 +82,11 @@ pub const Tri = struct { }, &.{ ci("D", dims), ci("S_COPIES", copies) }); } + pub fn embedDense(t: Tri, ids: u64, w: u64, out: u64, dims: usize, copies: usize, rows: usize) !void { + const g = t.own orelse return error.NativeGlueRequired; + try g.embedDense(t.s, ids, w, out, dims, copies, rows); + } + /// glue.rmsnorm per row and `group` features: bf16 out [rows, dims] and 32-group sums; x rows `x_stride` apart. pub fn rmsnorm(t: Tri, x: u64, w: u64, out: u64, xs: u64, eps: f32, x_stride: usize, rows: usize, dims: usize, group: usize) !void { if (t.glue("_rmsnorm")) |g| return g.rmsnorm(t.s, x, w, out, xs, eps, x_stride, rows, dims, group); diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 80e08aeb06..3c13f1b568 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -120,7 +120,11 @@ pub fn Ops(comptime Tri: type) type { p.pending = null; try p.stage(next[at..][0..m]); // the embedding branch: the token's row (stream 0 of the four copies), normed, fc_embedding - try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, m); + if (e.table.dense) { + try tri.embedDense(s.ids, e.table.w.ptr, s.ple_keys, dims, streams, m); + } else { + try tri.embed(s.ids, e.table.w.ptr, e.table.s.ptr, e.table.b.ptr, s.ple_keys, dims, streams, m); + } try tri.rmsnorm(s.ple_keys, h.norm_e.ptr, s.mixed, s.xsm, eps, wide, m, dims, dims); try p.lane(s.mixed, s.xsm, h.fc_e, s.dummy, s.part, m, fc_sk); try p.sumSlices(s.part, s.branch, true, m * dims, fc_sk); diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 17fd7463bb..032f7b596a 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -149,20 +149,30 @@ pub const Engine = struct { e.store = try Store.open(gpa, io, e.model_dir); errdefer e.store.deinit(); const weight = try e.store.tensor(embed.weight_name); - const scales = try e.store.tensor(embed.scales_name); - const biases = try e.store.tensor(embed.biases_name); - e.table = try embed.Table.upload(driver, weight, scales, biases); + if (weight.dtype == .bf16) { + e.table = try embed.Table.uploadDense(driver, weight); + } else { + const scales = try e.store.tensor(embed.scales_name); + const biases = try e.store.tensor(embed.biases_name); + e.table = try embed.Table.upload(driver, weight, scales, biases); + } errdefer e.table.deinit(); e.eos = try ple.eosOf(gpa, io, e.model_dir); e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); errdefer e.mixer.free(); - var buf: [80]u8 = undefined; - const head_w = try join3(gpa, &e.store, &buf, "language_model.lm_head", ""); - defer gpa.free(head_w.w); - defer gpa.free(head_w.s); - defer gpa.free(head_w.b); - if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; - e.head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); + const head_tensor = try e.store.tensor("language_model.lm_head.weight"); + if (head_tensor.dtype == .bf16) { + if (!head_tensor.is(.bf16, &.{ vocab, dims })) return error.UnexpectedTensor; + e.head_q = try @import("cuda_projection.zig").dense(gpa, driver, head_tensor); + } else { + var buf: [80]u8 = undefined; + const head_w = try join3(gpa, &e.store, &buf, "language_model.lm_head", ""); + defer gpa.free(head_w.w); + defer gpa.free(head_w.s); + defer gpa.free(head_w.b); + if (head_w.w.len != vocab * (dims / 8) * 4) return error.UnexpectedTensor; + e.head_q = try uploadQ4(gpa, driver, head_w.w, head_w.s, head_w.b, vocab, dims / 8); + } errdefer e.head_q.free(); var pack_mod = try cuda.Module.load(driver, cuda.kernels.experts_pack); defer pack_mod.unload(); @@ -276,7 +286,7 @@ pub const Engine = struct { .tri = tri, .k = &e.batch_k, .s = s, - .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, .seq = seq.view(scores), .hooks = hooks, }; @@ -314,7 +324,7 @@ pub const Engine = struct { .tri = tri, .k = &e.batch_k, .s = s, - .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, .seq = rows[0].seq, .lanes = rows, .hooks = .{ .logits = .all, .logits_out = out }, diff --git a/zig/src/families/flashnext/cuda_rows.zig b/zig/src/families/flashnext/cuda_rows.zig index e005b1bde2..89095b53ee 100644 --- a/zig/src/families/flashnext/cuda_rows.zig +++ b/zig/src/families/flashnext/cuda_rows.zig @@ -70,6 +70,7 @@ pub const Rows = struct { /// Every row of `ids` from seq.fed one fused decode step; the last row's scores land in seq.logits. pub fn predict(r: *Rows, e: *prompt.Engine, comptime Tri: type, tri: Tri, seq: *prompt.Seq, ids: []const u32) !void { + if (e.table.dense or e.head_q.native != null) return error.UnsupportedQuantization; if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; // dense attention only: the reference stops where the indexer would start choosing if (seq.fed + ids.len > weights.dense_nch * 512) return error.UnexpectedTensor; diff --git a/zig/tests/cuda/norm_rows_tests.zig b/zig/tests/cuda/norm_rows_tests.zig index ec95d7fe36..9eb6005c30 100644 --- a/zig/tests/cuda/norm_rows_tests.zig +++ b/zig/tests/cuda/norm_rows_tests.zig @@ -179,7 +179,8 @@ fn embed(r: *Rig, d: usize, group: usize, copies: usize) !void { for (w) |*x| x.* = r.rng.int(u32); const sc = r.bfs(vocab * d / group, 0.1); const bi = r.bfs(vocab * d / group, 0.5); - const tab = r.bfs(vocab * d, 1); + const tab = try r.a.alloc(u16, vocab * d); + for (tab, 0..) |*x, i| x.* = @truncate(i); // every BF16 bit pattern, including NaN payloads const ids = try r.a.alloc(i32, most); for (ids) |*x| x.* = @intCast(r.rng.uintLessThan(usize, vocab)); const did = try r.dev(std.mem.sliceAsBytes(ids)); @@ -189,7 +190,21 @@ fn embed(r: *Rig, d: usize, group: usize, copies: usize) !void { const ds = try r.dev(std.mem.sliceAsBytes(sc)); const db = try r.dev(std.mem.sliceAsBytes(bi)); if (group == 64 and copies == 1) try e.affine4(r.s, did, dw, ds, db, o4, most, d) else try e.affine4Groups(r.s, did, dw, ds, db, o4, most, d, group, copies); - try e.bf16(r.s, did, try r.dev(std.mem.sliceAsBytes(tab)), o16, most, d); + const dt = try r.dev(std.mem.sliceAsBytes(tab)); + try e.bf16(r.s, did, dt, o16, most, d); + const oc = try r.out(most * copies * d * 2); + try e.bf16Copies(r.s, did, dt, oc, 0, d, copies); + try std.testing.expectError(error.RowShape, e.bf16Copies(r.s, did, dt, oc, most, d, 0)); + try std.testing.expectError(error.RowShape, e.bf16Copies(r.s, did, dt, oc + 2, most, d, copies)); + try std.testing.expectError(error.RowShape, e.bf16Copies(r.s, did, dt, oc, most, d - 1, copies)); + for (counts) |rows| { + try e.bf16Copies(r.s, did, dt, oc, rows, d, copies); + const got = try r.back(u16, oc, rows * copies * d); + for (0..rows) |row| { + const token: usize = @intCast(ids[row]); + for (0..copies) |c| try check.sameBytes("embed BF16 exact copies", std.mem.sliceAsBytes(got[(row * copies + c) * d ..][0..d]), std.mem.sliceAsBytes(tab[token * d ..][0..d])); + } + } const g4 = try r.back(u16, o4, most * copies * d); const g16 = try r.back(u16, o16, most * d); for (0..most) |row| { From 97715a28beec893ade236b3f793937acf3dfb6db Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:51:45 -0700 Subject: [PATCH 065/115] Load mixed dense and microscaled projections without requantization --- zig/src/families/flashnext/cuda_launch.zig | 1 + .../flashnext/cuda_projection_load.zig | 118 ++++++++++++++++++ .../flashnext/cuda_projection_load_test.zig | 63 ++++++++++ .../families/flashnext/cuda_prompt_load.zig | 61 +++------ 4 files changed, 202 insertions(+), 41 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_projection_load.zig create mode 100644 zig/src/families/flashnext/cuda_projection_load_test.zig diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 65e9024d06..2037bb07d4 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -370,6 +370,7 @@ test { _ = native; _ = @import("cuda_snap.zig"); _ = @import("cuda_config.zig"); + _ = @import("cuda_projection_load.zig"); _ = row_path; _ = @import("cuda_mtp_check.zig"); } diff --git a/zig/src/families/flashnext/cuda_projection_load.zig b/zig/src/families/flashnext/cuda_projection_load.zig new file mode 100644 index 0000000000..4fa63bb581 --- /dev/null +++ b/zig/src/families/flashnext/cuda_projection_load.zig @@ -0,0 +1,118 @@ +//! Projection loading preserves dense BF16 bits and MXFP8 codes/scales; affine layouts keep their existing packer. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const projection = @import("cuda_projection.zig"); + +pub const Micro = struct { + n: usize, + k: usize, + codes: []u8, + scales: []u8, + + pub fn deinit(self: *Micro, a: std.mem.Allocator) void { + a.free(self.codes); + a.free(self.scales); + self.* = undefined; + } +}; + +/// Own each plane before the store advances to a possibly different mapped shard. +pub fn readMicro(a: std.mem.Allocator, store: anytype, prefix: []const u8) !Micro { + var buf: [512]u8 = undefined; + const w = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); + if (w.dtype != .f8_e4m3 or w.rank != 2 or w.dim(0) == 0 or w.dim(1) == 0 or w.dim(1) % 64 != 0) + return error.UnexpectedTensor; + const n = w.dim(0); + const k = w.dim(1); + const count = try std.math.mul(usize, n, k); + if (w.bytes.len != count) return error.UnexpectedTensor; + const codes = try a.dupe(u8, w.bytes); + errdefer a.free(codes); + const s = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{prefix})); + if (!s.is(.u8, &.{ n, k / 32 }) or s.bytes.len != count / 32) return error.UnexpectedTensor; + return .{ .n = n, .k = k, .codes = codes, .scales = try a.dupe(u8, s.bytes) }; +} + +/// Concatenate output rows of compatible projections without converting or recalibrating either plane. +pub fn readStitched(a: std.mem.Allocator, store: anytype, prefix: []const u8, names: []const []const u8) !Micro { + if (names.len == 0 or names.len > 6) return error.UnexpectedTensor; + var parts: [6]Micro = undefined; + var held: usize = 0; + defer for (parts[0..held]) |*part| part.deinit(a); + var n: usize = 0; + var k: usize = 0; + var buf: [512]u8 = undefined; + for (names) |name| { + parts[held] = try readMicro(a, store, try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, name })); + held += 1; + const part = parts[held - 1]; + if (k != 0 and k != part.k) return error.UnexpectedTensor; + k = part.k; + n = try std.math.add(usize, n, part.n); + } + const count = try std.math.mul(usize, n, k); + const codes = try a.alloc(u8, count); + errdefer a.free(codes); + const scales = try a.alloc(u8, count / 32); + var row: usize = 0; + for (parts[0..held]) |part| { + @memcpy(codes[row * k ..][0..part.codes.len], part.codes); + @memcpy(scales[row * (k / 32) ..][0..part.scales.len], part.scales); + row += part.n; + } + return .{ .n = n, .k = k, .codes = codes, .scales = scales }; +} + +fn uploadMicro(a: std.mem.Allocator, d: *const cuda.Driver, value: Micro) !wt.Q4 { + return projection.mxfp8(a, d, .{ .dtype = .f8_e4m3, .rank = 2, .shape = .{ value.n, value.k, 1, 1, 1 }, .bytes = value.codes }, .{ .dtype = .u8, .rank = 2, .shape = .{ value.n, value.k / 32, 1, 1, 1 }, .bytes = value.scales }); +} + +pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, n: usize, k: usize) !wt.Q4 { + var buf: [512]u8 = undefined; + const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); + switch (t.dtype) { + .bf16 => { + if (!t.is(.bf16, &.{ n, k })) return error.UnexpectedTensor; + return projection.dense(a, d, t); + }, + .f8_e4m3 => { + var value = try readMicro(a, store, prefix); + defer value.deinit(a); + if (value.n != n or value.k != k) return error.UnexpectedTensor; + return uploadMicro(a, d, value); + }, + .u32 => { + if (!t.is(.u32, &.{ n, k / 8 })) return error.UnexpectedTensor; + const planes = try wt.join3(a, store, &buf, prefix, ""); + defer a.free(planes.w); + defer a.free(planes.s); + defer a.free(planes.b); + return wt.uploadQ4(a, d, planes.w, planes.s, planes.b, n, k / 8); + }, + else => return error.UnsupportedQuantization, + } +} + +pub fn loadStitched(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !wt.Q4 { + if (names.len == 0) return error.UnexpectedTensor; + var buf: [512]u8 = undefined; + const first = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, names[0] })); + if (first.dtype == .f8_e4m3) { + var value = try readStitched(a, store, prefix, names); + defer value.deinit(a); + if (value.n != n or value.k != k) return error.UnexpectedTensor; + return uploadMicro(a, d, value); + } + if (first.dtype != .u32) return error.UnsupportedQuantization; + const planes = try wt.stitch(a, store, prefix, names); + defer a.free(planes.w); + defer a.free(planes.s); + defer a.free(planes.b); + if (planes.w.len != n * (k / 8) * 4) return error.UnexpectedTensor; + return wt.uploadQ4(a, d, planes.w, planes.s, planes.b, n, k / 8); +} + +test { + _ = @import("cuda_projection_load_test.zig"); +} diff --git a/zig/src/families/flashnext/cuda_projection_load_test.zig b/zig/src/families/flashnext/cuda_projection_load_test.zig new file mode 100644 index 0000000000..a76e08c304 --- /dev/null +++ b/zig/src/families/flashnext/cuda_projection_load_test.zig @@ -0,0 +1,63 @@ +//! A store that invalidates its previous view on every read, like advancing mapped shards. +const std = @import("std"); +const core = @import("core"); +const source = @import("cuda_projection_load.zig"); +const Tensor = core.safetensors.Tensor; + +const Fake = struct { + scratch: [1024]u8 = undefined, + bad: enum { none, weight, scale, extent, width, missing } = .none, + + pub fn tensor(self: *Fake, name: []const u8) !Tensor { + const part: usize = if (std.mem.startsWith(u8, name, "root.a.")) 0 else if (std.mem.startsWith(u8, name, "root.b.")) 1 else return error.MissingTensor; + const n: usize = if (part == 0) 3 else 5; + const k: usize = if (self.bad == .width and part == 1) 128 else 64; + const scale = std.mem.endsWith(u8, name, ".weight_scale"); + if (scale and self.bad == .missing) return error.MissingTensor; + if (!scale and !std.mem.endsWith(u8, name, ".weight")) return error.MissingTensor; + const cols = if (scale) k / 32 else k; + const len = n * cols; + @memset(&self.scratch, 0xee); + for (self.scratch[0..len], 0..) |*v, i| v.* = @truncate(i * 17 + part * 53 + @as(usize, if (scale) 129 else 7)); + return .{ + .dtype = if (scale) (if (self.bad == .scale) .f8_e4m3 else .u8) else (if (self.bad == .weight) .bf16 else .f8_e4m3), + .rank = 2, + .shape = .{ n, cols, 1, 1, 1 }, + .bytes = self.scratch[0 .. len - @as(usize, if (self.bad == .extent) 1 else 0)], + }; + } +}; + +fn trial(a: std.mem.Allocator) !void { + var fake: Fake = .{}; + var value = try source.readStitched(a, &fake, "root", &.{ ".a", ".b" }); + defer value.deinit(a); + try std.testing.expectEqual(@as(usize, 8), value.n); + try std.testing.expectEqual(@as(usize, 64), value.k); + for (0..2) |part| { + const start: usize = if (part == 0) 0 else 3; + const n: usize = if (part == 0) 3 else 5; + for (0..n * 64) |i| try std.testing.expectEqual(@as(u8, @truncate(i * 17 + part * 53 + 7)), value.codes[start * 64 + i]); + for (0..n * 2) |i| try std.testing.expectEqual(@as(u8, @truncate(i * 17 + part * 53 + 129)), value.scales[start * 2 + i]); + } +} + +test "MXFP8 fusion owns each shard plane and preserves exact output row order" { + try trial(std.testing.allocator); +} + +test "MXFP8 fusion cleans every partial allocation on failure" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, trial, .{}); +} + +test "MXFP8 source refuses incompatible shapes, dtypes, extents and missing scales" { + const a = std.testing.allocator; + for ([_]@FieldType(Fake, "bad"){ .weight, .scale, .extent, .width }) |bad| { + var fake: Fake = .{ .bad = bad }; + try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{ ".a", ".b" })); + } + var fake: Fake = .{ .bad = .missing }; + try std.testing.expectError(error.MissingTensor, source.readMicro(a, &fake, "root.a")); + try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{})); + try std.testing.expectError(error.UnexpectedTensor, source.readStitched(a, &fake, "root", &.{ ".a", ".a", ".a", ".a", ".a", ".a", ".a" })); +} diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index 7626925352..9f2c1d3b9d 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const embed = @import("cuda_embed.zig"); const gdn = @import("cuda_gdn.zig"); const wt = @import("cuda_prompt_wt.zig"); +const projections = @import("cuda_projection_load.zig"); const layers_n = wt.layers_n; const dims = wt.dims; @@ -24,10 +25,8 @@ const Q4 = wt.Q4; const HcW = wt.HcW; const linearLayer = wt.linearLayer; const asF32 = wt.asF32; -const uploadQ4 = wt.uploadQ4; const join3 = wt.join3; const loadHc = wt.loadHc; -const stitch = wt.stitch; const OwnedFace = struct { face: Face, @@ -217,14 +216,20 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S const gate_copy = try gpa.alloc(u8, gate.bytes.len); defer gpa.free(gate_copy); @memcpy(gate_copy, gate.bytes); - const sg = try join3(gpa, store, &buf, mlp_root, ".shared_expert_gate"); - defer gpa.free(sg.w); - defer gpa.free(sg.s); - defer gpa.free(sg.b); const router = try gpa.alloc(u16, (experts_n + 1) * dims); defer gpa.free(router); @memcpy(std.mem.sliceAsBytes(router[0 .. experts_n * dims]), gate_copy); - try embed.dequant(sg.w, sg.s, sg.b, dims, 0, router[experts_n * dims ..][0..dims]); + const sg_tensor = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.shared_expert_gate.weight", .{mlp_root})); + if (sg_tensor.dtype == .bf16) { + if (!sg_tensor.is(.bf16, &.{ 1, dims }) or sg_tensor.bytes.len != dims * 2) return error.UnexpectedTensor; + @memcpy(std.mem.sliceAsBytes(router[experts_n * dims ..]), sg_tensor.bytes); + } else { + const sg = try join3(gpa, store, &buf, mlp_root, ".shared_expert_gate"); + defer gpa.free(sg.w); + defer gpa.free(sg.s); + defer gpa.free(sg.b); + try embed.dequant(sg.w, sg.s, sg.b, dims, 0, router[experts_n * dims ..][0..dims]); + } var router_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(router)); errdefer router_b.free(); @@ -242,6 +247,8 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S var iks_b = deadBuf(driver); var key_q = Q4.empty(driver); var val_q = Q4.empty(driver); + errdefer key_q.free(); + errdefer val_q.free(); var nk_b = deadBuf(driver); var nq_b = deadBuf(driver); var nc_b = deadBuf(driver); @@ -264,17 +271,8 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S @memcpy(pre_buf[0..pre.len], pre); const lp = pre_buf[0..pre.len]; const parts = [_][]const u8{ ".in_proj_qkv", ".in_proj_z", ".in_proj_b", ".in_proj_a" }; - const stitched = try stitch(gpa, store, lp, &parts); - defer gpa.free(stitched.w); - defer gpa.free(stitched.s); - defer gpa.free(stitched.b); - if (stitched.w.len != gdn.proj_width * (dims / 8) * 4) return error.UnexpectedTensor; - proj_q = try uploadQ4(gpa, driver, stitched.w, stitched.s, stitched.b, gdn.proj_width, dims / 8); - const out = try join3(gpa, store, &buf, lp, ".out_proj"); - defer gpa.free(out.w); - defer gpa.free(out.s); - defer gpa.free(out.b); - out_q = try uploadQ4(gpa, driver, out.w, out.s, out.b, dims, gdn.value_dim / 8); + proj_q = try projections.loadStitched(gpa, driver, store, lp, &parts, gdn.proj_width, dims); + out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.out_proj", .{lp}), dims, gdn.value_dim); const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); @@ -294,17 +292,8 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S @memcpy(pre_buf[0..pre.len], pre); const ap = pre_buf[0..pre.len]; const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; - const stitched = try stitch(gpa, store, ap, &parts); - defer gpa.free(stitched.w); - defer gpa.free(stitched.s); - defer gpa.free(stitched.b); - if (stitched.w.len != proj_n * (dims / 8) * 4) return error.UnexpectedTensor; - proj_q = try uploadQ4(gpa, driver, stitched.w, stitched.s, stitched.b, proj_n, dims / 8); - const out = try join3(gpa, store, &buf, ap, ".o_proj"); - defer gpa.free(out.w); - defer gpa.free(out.s); - defer gpa.free(out.b); - out_q = try uploadQ4(gpa, driver, out.w, out.s, out.b, dims, out_k / 8); + proj_q = try projections.loadStitched(gpa, driver, store, ap, &parts, proj_n, dims); + out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.o_proj", .{ap}), dims, out_k); const qn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); defer gpa.free(qn); const kn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); @@ -324,16 +313,8 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S var pre_buf: [96]u8 = undefined; @memcpy(pre_buf[0..pre.len], pre); const pp = pre_buf[0..pre.len]; - const key = try join3(gpa, store, &buf, pp, ".key_proj"); - defer gpa.free(key.w); - defer gpa.free(key.s); - defer gpa.free(key.b); - const val = try join3(gpa, store, &buf, pp, ".value_proj"); - defer gpa.free(val.w); - defer gpa.free(val.s); - defer gpa.free(val.b); - key_q = try uploadQ4(gpa, driver, key.w, key.s, key.b, wide, dims / 8); - val_q = try uploadQ4(gpa, driver, val.w, val.s, val.b, dims, dims / 8); + key_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.key_proj", .{pp}), wide, dims); + val_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.value_proj", .{pp}), dims, dims); const nk = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); defer gpa.free(nk); const nq = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); @@ -348,8 +329,6 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); } - errdefer key_q.free(); - errdefer val_q.free(); var shared: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; errdefer for (&shared) |*f| f.free(gpa); From 2661d284cf6fa4a46a2c72ea1f4f91fc90b60535 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:06:30 -0700 Subject: [PATCH 066/115] Load exact rank-local NVFP4 routed expert partitions --- zig/src/core/nvfp4_slice.zig | 2 +- zig/src/families/flashnext/cuda_launch.zig | 2 + .../flashnext/cuda_modelopt_experts.zig | 134 +++++++++++++++ zig/src/families/flashnext/cuda_prompt_wt.zig | 12 ++ zig/tests/cuda/flashnext_nvfp4_tp_tests.zig | 161 ++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 6 files changed, 313 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/flashnext/cuda_modelopt_experts.zig create mode 100644 zig/tests/cuda/flashnext_nvfp4_tp_tests.zig diff --git a/zig/src/core/nvfp4_slice.zig b/zig/src/core/nvfp4_slice.zig index 5f3dc0fbc7..c76e3ed7a5 100644 --- a/zig/src/core/nvfp4_slice.zig +++ b/zig/src/core/nvfp4_slice.zig @@ -78,7 +78,7 @@ pub fn slice(gpa: std.mem.Allocator, source: Source, rows: Span, inputs: Span) ! } /// Read exact stored tensors, consuming optional calibration and refusing ambiguous plain/compressed scales. -pub fn read(ck: *ckpt.Checkpoint, gpa: std.mem.Allocator, module: []const u8, rows: Span, inputs: Span) !Linear { +pub fn read(ck: anytype, gpa: std.mem.Allocator, module: []const u8, rows: Span, inputs: Span) !Linear { var buf: [512]u8 = undefined; const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 2037bb07d4..1367aec3a0 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,6 +6,7 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const projection = @import("cuda_projection.zig"); +pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); pub const hc = @import("cuda_hc.zig"); pub const qmm = @import("cuda_qmm.zig"); pub const gdn = @import("cuda_gdn.zig"); @@ -371,6 +372,7 @@ test { _ = @import("cuda_snap.zig"); _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); + _ = @import("cuda_modelopt_experts.zig"); _ = row_path; _ = @import("cuda_mtp_check.zig"); } diff --git a/zig/src/families/flashnext/cuda_modelopt_experts.zig b/zig/src/families/flashnext/cuda_modelopt_experts.zig new file mode 100644 index 0000000000..553a6be9a9 --- /dev/null +++ b/zig/src/families/flashnext/cuda_modelopt_experts.zig @@ -0,0 +1,134 @@ +//! Rank-local intermediate-channel TP for NVFP4 routed experts; the shared expert retains its own projection format. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const ex = cuda.experts.Nvfp4Experts; +const slice = core.nvfp4_slice; + +pub const Cut = struct { + from: usize, + to: usize, + + pub fn init(total_width: usize, world: usize, rank: usize) !Cut { + if (world == 0 or rank >= world or total_width == 0 or total_width % world != 0) return error.UnevenExpertSplit; + const part = total_width / world; + if (part % 32 != 0) return error.UnevenExpertSplit; + return .{ .from = rank * part, .to = (rank + 1) * part }; + } + + pub fn width(self: Cut) usize { + return self.to - self.from; + } +}; + +pub const Owned = struct { + up: cuda.DeviceBuffer, + down: cuda.DeviceBuffer, + up_scale: cuda.DeviceBuffer, + down_scale: cuda.DeviceBuffer, + width: usize, + dims: usize, + count: usize, + + pub fn free(self: *Owned) void { + self.up.free(); + self.down.free(); + self.up_scale.free(); + self.down_scale.free(); + self.* = undefined; + } + + pub fn view(self: *const Owned) cuda.experts.Layer { + return .{ .format = .nvfp4, .up = self.up.ptr, .down = self.down.ptr, .up_scale = self.up_scale.ptr, .down_scale = self.down_scale.ptr, .width = self.width, .dims = self.dims, .experts = self.count }; + } +}; + +/// The reader owns only the selected rectangle, not a decoded full expert stack. +/// The store is an exclusive loader scratch here: every borrowed plane is released before returning. +pub fn readPart(a: std.mem.Allocator, store: anytype, name: []const u8, cut: Cut, down: bool, dims: usize, full_width: usize) !slice.Linear { + defer store.release(); + var buf: [512]u8 = undefined; + const w = try store.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{name})); + const n = if (down) dims else full_width; + const k = if (down) full_width else dims; + if (!w.is(.u8, &.{ n, k / 2 })) return error.UnexpectedTensor; + var linear = try slice.read(store, a, name, .{ .from = if (down) 0 else cut.from, .to = if (down) dims else cut.to }, .{ .from = if (down) cut.from else 0, .to = if (down) cut.to else dims }); + errdefer linear.deinit(a); + if (!std.math.isFinite(linear.global) or linear.global <= 0) return error.UnsupportedQuantization; + return linear; +} + +/// Pack and upload one expert at a time, bounding host scratch independently of expert count. +/// Activations remain BF16 (W4A16); input_scale is validated as metadata by the source reader. +pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !Owned { + return loadShape(a, d, store, prefix, world, rank, wt.dims, wt.moe_w, wt.experts_n); +} + +pub fn loadShape(a: std.mem.Allocator, d: *const cuda.Driver, store: anytype, prefix: []const u8, world: usize, rank: usize, dims: usize, full_width: usize, count: usize) !Owned { + const cut = try Cut.init(full_width, world, rank); + if (dims == 0 or dims % 32 != 0 or count == 0) return error.UnexpectedTensor; + const width = cut.width(); + const up_bytes = try packedBytes(width, dims, 2); + const down_bytes = try packedBytes(dims, width, 1); + const up_total = std.math.mul(usize, count, up_bytes) catch return error.UnexpectedTensor; + const down_total = std.math.mul(usize, count, down_bytes) catch return error.UnexpectedTensor; + const globals_count = std.math.mul(usize, count, 2) catch return error.UnexpectedTensor; + var up = try cuda.DeviceBuffer.alloc(d, up_total); + errdefer up.free(); + var down = try cuda.DeviceBuffer.alloc(d, down_total); + errdefer down.free(); + const up_globals = try a.alloc(f32, globals_count); + defer a.free(up_globals); + const down_globals = try a.alloc(f32, count); + defer a.free(down_globals); + const packed_up = try a.alloc(u32, up_bytes / 4); + defer a.free(packed_up); + const packed_down = try a.alloc(u32, down_bytes / 4); + defer a.free(packed_down); + var buf: [512]u8 = undefined; + for (0..count) |i| { + var gate = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.gate_proj", .{ prefix, i }), cut, false, dims, full_width); + defer gate.deinit(a); + var value = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.up_proj", .{ prefix, i }), cut, false, dims, full_width); + defer value.deinit(a); + try ex.packGateUp(a, packed_up, .{ gate.codes, gate.scales }, .{ value.codes, value.scales }, width, dims); + try up.upload(i * up_bytes, std.mem.sliceAsBytes(packed_up)); + up_globals[i * 2] = gate.global; + up_globals[i * 2 + 1] = value.global; + var output = try readPart(a, store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.down_proj", .{ prefix, i }), cut, true, dims, full_width); + defer output.deinit(a); + ex.packOne(packed_down, output.codes, output.scales, dims, width); + try down.upload(i * down_bytes, std.mem.sliceAsBytes(packed_down)); + down_globals[i] = output.global; + } + var us = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(up_globals)); + errdefer us.free(); + const ds = try cuda.DeviceBuffer.fromHost(d, std.mem.sliceAsBytes(down_globals)); + return .{ .up = up, .down = down, .up_scale = us, .down_scale = ds, .width = width, .dims = dims, .count = count }; +} + +test "intermediate-channel cuts cover the source once and preserve 32-input kernel groups" { + const a = try Cut.init(640, 2, 0); + const b = try Cut.init(640, 2, 1); + try std.testing.expectEqual(@as(usize, 0), a.from); + try std.testing.expectEqual(@as(usize, 320), a.to); + try std.testing.expectEqual(a.to, b.from); + try std.testing.expectEqual(@as(usize, 640), b.to); + try std.testing.expectEqual(a.width(), b.width()); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 0, 0)); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 2, 2)); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 3, 0)); + try std.testing.expectError(error.UnevenExpertSplit, Cut.init(640, 32, 0)); +} + +fn packedBytes(n: usize, k: usize, parts: usize) !usize { + var size = std.math.mul(usize, n / ex.cols, k / 32) catch return error.UnexpectedTensor; + for ([_]usize{ parts, ex.words, 4 }) |factor| size = std.math.mul(usize, size, factor) catch return error.UnexpectedTensor; + return size; +} + +test "expert packed extents reject overflow before device allocation" { + try std.testing.expectEqual(ex.bytes(1, 320, 2560, 2), try packedBytes(320, 2560, 2)); + try std.testing.expectError(error.UnexpectedTensor, packedBytes(std.math.maxInt(usize), std.math.maxInt(usize), 2)); +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index bea01330d0..5828a6ab8a 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -183,6 +183,18 @@ pub const Store = struct { }; } + /// Check the exact or compatibility name without mapping its shard. + pub fn has(self: *const Store, requested: []const u8) bool { + var mapped: [512]u8 = undefined; + _ = self.resolveKey(requested, &mapped) catch return false; + return true; + } + + /// Checkpoint-reader adapter: the borrowed tensor stays mapped until release(). + pub fn get(self: *Store, key: []const u8) !Tensor { + return self.hold(key); + } + pub fn copyOf(self: *Store, key: []const u8) ![]u8 { const t = try self.tensor(key); const out = try self.gpa.alloc(u8, t.bytes.len); diff --git a/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig new file mode 100644 index 0000000000..2feac586ef --- /dev/null +++ b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig @@ -0,0 +1,161 @@ +//! Intermediate-channel NVFP4 partitions: exact source cuts, owned upload and FP32 expert partials against FP64. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const d = 64; +const width = 64; +const experts = 2; +const slots = 3; +const most = 17; +const rank_width = width / 2; +const fx = cuda.experts.Nvfp4Experts; + +const Source = struct { + codes: [experts][3][d * width / 2]u8 = undefined, + scales: [experts][3][d * width / 16]u8 = undefined, + globals: [experts][3][4]u8 = undefined, + input: [4]u8 = undefined, + releases: usize = 0, + + fn init() Source { + var s: Source = .{}; + for (0..experts) |e| for (0..3) |p| { + for (&s.codes[e][p], 0..) |*v, i| v.* = @truncate(i * 13 + e * 37 + p * 19); + for (&s.scales[e][p], 0..) |*v, i| v.* = @intCast(0x20 + (i % 3) * 8); + std.mem.writeInt(u32, &s.globals[e][p], @bitCast(@as(f32, if (p == 2) 0.25 else 0.125)), .little); + }; + std.mem.writeInt(u32, &s.input, @bitCast(@as(f32, 1)), .little); + return s; + } + + pub fn has(_: *Source, name: []const u8) bool { + return std.mem.endsWith(u8, name, ".weight") or std.mem.endsWith(u8, name, ".weight_scale") or + std.mem.endsWith(u8, name, ".weight_scale_2") or std.mem.endsWith(u8, name, ".input_scale"); + } + + pub fn get(s: *Source, name: []const u8) !core.safetensors.Tensor { + const marker = std.mem.indexOf(u8, name, ".experts.") orelse return error.MissingTensor; + const tail = name[marker + ".experts.".len ..]; + const end = std.mem.indexOfScalar(u8, tail, '.') orelse return error.MissingTensor; + const e = try std.fmt.parseInt(usize, tail[0..end], 10); + if (e >= experts) return error.MissingTensor; + const p: usize = if (std.mem.indexOf(u8, tail, ".gate_proj.") != null) 0 else if (std.mem.indexOf(u8, tail, ".up_proj.") != null) 1 else if (std.mem.indexOf(u8, tail, ".down_proj.") != null) 2 else return error.MissingTensor; + if (std.mem.endsWith(u8, name, ".weight")) return tensor(.u8, &.{ width, d / 2 }, &s.codes[e][p]); + if (std.mem.endsWith(u8, name, ".weight_scale")) return tensor(.f8_e4m3, &.{ width, d / 16 }, &s.scales[e][p]); + if (std.mem.endsWith(u8, name, ".weight_scale_2")) return tensor(.f32, &.{}, &s.globals[e][p]); + if (std.mem.endsWith(u8, name, ".input_scale")) return tensor(.f32, &.{}, &s.input); + return error.MissingTensor; + } + + pub fn release(s: *Source) void { + s.releases += 1; + } +}; + +fn tensor(dtype: core.safetensors.DType, shape: []const usize, bytes: []const u8) core.safetensors.Tensor { + var t: core.safetensors.Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +fn weight(s: *const Source, e: usize, p: usize, n: usize, k: usize) f64 { + const byte = s.codes[e][p][n * (d / 2) + k / 2]; + const code: u4 = if (k % 2 == 0) @truncate(byte) else @truncate(byte >> 4); + const scale = s.scales[e][p][n * (d / 16) + k / 16]; + const global: f32 = @bitCast(std.mem.readInt(u32, &s.globals[e][p], .little)); + return core.quant.nvfp4.e2m1(code) * core.quant.fp8.e4m3(scale) * global; +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0x2f4); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + var source = Source.init(); + var left = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 2, 0, d, width, experts); + defer left.free(); + var right = try flash.modelopt_experts.loadShape(a, gpu.d, &source, "model.mlp", 2, 1, d, width, experts); + defer right.free(); + try check.expect(source.releases == 2 * experts * 3, "every source projection released after its owned rank cut", .{}); + const x = r.bfs(most * d, 0.125); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const picks = try a.alloc(i32, most * slots); + for (0..most) |row| { + picks[row * slots] = @intCast(row % experts); + picks[row * slots + 1] = @intCast((row + 1) % experts); + picks[row * slots + 2] = experts; // shared expert is a separate format, skipped here + } + const dp = try r.dev(std.mem.sliceAsBytes(picks)); + const plan: cuda.grouped.Plan = .{ + .members = try r.out(most * slots * 4), + .items = try r.out((most * slots + experts + 1) * 12), + .counts = try r.out(8), + .rank = try r.out(most * slots * 4), + .hist = try r.out((experts + 1) * 4), + }; + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try cuda.grouped.Router.resolve(mod); + var kernels = try fx.load(gpu.d, @intCast(try gpu.ctx.attribute(.multiprocessor_count))); + defer kernels.unload(); + const act = try r.out(most * slots * rank_width * 2); + const out = try r.out(most * slots * d * 4); + const combined = try a.alloc(f32, most * slots * d); + @memset(combined, 0); + const reference_full = try a.alloc(f64, combined.len); + @memset(reference_full, 0); + const reference_abs = try a.alloc(f64, combined.len); + @memset(reference_abs, 0); + for ([_]*const flash.modelopt_experts.Owned{ &left, &right }, 0..) |owned, rank| { + var all: []f32 = undefined; + for ([_]usize{ most, 1, 7 }) |rows| { + try gpu.d.check(gpu.d.api.cuMemsetD8Async(out, 0, most * slots * d * 4, r.s.handle), "cuMemsetD8Async"); + const items = cuda.grouped.maxItems(rows * slots, experts + 1, 16); + try router.route(r.s, dp, rows * slots, experts + 1, 16, plan); + try kernels.run(r.s, .up, .{ .x = dx, .stride = d, .slots = slots }, owned.view(), plan, act, items, experts); + try kernels.run(r.s, .down_f32, .{ .x = act, .stride = rank_width }, owned.view(), plan, out, items, experts); + if (rows == most) all = try r.back(f32, out, most * slots * d) else try r.sameRows("NVFP4 rank output rows", f32, out, all[0 .. rows * slots * d]); + } + for (0..most) |row| for (0..slots) |slot| { + const e: usize = @intCast(picks[row * slots + slot]); + if (e == experts) { + for (all[(row * slots + slot) * d ..][0..d]) |v| try check.expect(v == 0, "shared slot untouched by NVFP4", .{}); + continue; + } + var activation: [rank_width]f64 = undefined; + for (&activation, 0..) |*v, j| { + const channel = rank * rank_width + j; + var gate: f64 = 0; + var up: f64 = 0; + for (0..d) |k| { + gate += rig.bf(x[row * d + k]) * weight(&source, e, 0, channel, k); + up += rig.bf(x[row * d + k]) * weight(&source, e, 1, channel, k); + } + gate = rig.rbf(gate); + up = rig.rbf(up); + v.* = rig.rbf(rig.rbf(gate / (1 + @exp(-gate))) * up); + } + for (0..d) |j| { + var want: f64 = 0; + var mag: f64 = 0; + for (activation, 0..) |v, k| { + const term = v * weight(&source, e, 2, j, rank * rank_width + k); + want += term; + mag += @abs(term); + } + const at = (row * slots + slot) * d + j; + try r.within("NVFP4 rank FP32 down", all[at], want, rig.sumTol(mag, rank_width) + 1e-6); + combined[at] += all[at]; + reference_full[at] += want; + reference_abs[at] += mag; + } + }; + } + for (combined, reference_full, reference_abs) |got, want, mag| try r.within("NVFP4 summed rank partials", got, want, rig.sumTol(mag, width) + 2e-6); + r.report("NVFP4 TP cuts: both intermediate halves, exact source reader and packing, independent FP64 reference, FP32 partials and untouched shared slot; rows 1/7/17 equal"); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 6a324a358b..54ce9fd6f4 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -33,6 +33,7 @@ const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); +const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); @@ -59,6 +60,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) @@ -149,6 +151,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); From 041cc0a44efb7bbef02f78bb72a913b9097ea6e9 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:12:27 -0700 Subject: [PATCH 067/115] Run the mixed-format shared expert separately from routed partitions --- zig/kernels/cuda/flashnext_hc.cu | 15 +++ zig/src/families/flashnext/cuda_batch.zig | 19 +++ .../families/flashnext/cuda_batch_state.zig | 17 +++ zig/src/families/flashnext/cuda_glue.zig | 19 +++ .../families/flashnext/cuda_mixed_experts.zig | 36 +++++ .../families/flashnext/cuda_prompt_load.zig | 84 ++++++++---- zig/tests/cuda/flashnext_shared_tests.zig | 125 ++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 8 files changed, 290 insertions(+), 28 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_mixed_experts.zig create mode 100644 zig/tests/cuda/flashnext_shared_tests.zig diff --git a/zig/kernels/cuda/flashnext_hc.cu b/zig/kernels/cuda/flashnext_hc.cu index 277dd6563b..1ea5886167 100644 --- a/zig/kernels/cuda/flashnext_hc.cu +++ b/zig/kernels/cuda/flashnext_hc.cu @@ -432,3 +432,18 @@ extern "C" __global__ void fn_hc_dense_mix(const __nv_bfloat16* UP, const __nv_b XSM[static_cast(r) * (D / 32) + group] = sum; } } + +// Shared expert uses its own BF16 projection outputs and rounding boundaries. +extern "C" __global__ void fn_shared_swiglu(const __nv_bfloat16* GATE, const __nv_bfloat16* UP, + __nv_bfloat16* ACT, int W) { + const int r = blockIdx.x, j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= W) return; + const long long at = static_cast(r) * W + j; + ACT[at] = __float2bfloat16_rn(__fmul_rn(bsilu(bf(GATE[at])), bf(UP[at]))); +} + +// Write only the final reserved slot; routed slots retain their FP32 partial sums. +extern "C" __global__ void fn_shared_scatter(const float* SHARED, float* Y, int D, int SLOTS) { + const int r = blockIdx.x, j = blockIdx.y * blockDim.x + threadIdx.x; + if (j < D) Y[(static_cast(r) * SLOTS + SLOTS - 1) * D + j] = SHARED[static_cast(r) * D + j]; +} diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index c35ac83f9c..f45879407b 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -369,6 +369,25 @@ pub fn Pass(comptime Tri: type) type { try self.tri.topkMany(s.router_l, s.picks, s.wts, n); const e = experts_n + 1; const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; + if (lw.native_experts) |*native| { + const owned = &native.experts; + try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); + try self.k.router.route(self.stream, s.picks, pairs, e, 16, plan); + const items = cuda.grouped.maxItems(pairs, e, 16); + try self.k.nvfp4.run(self.stream, .up, .{ .x = s.mixed, .stride = dims, .slots = slots }, owned.view(), plan, s.moe_act, items, experts_n); + try self.k.nvfp4.run(self.stream, .down_f32, .{ .x = s.moe_act, .stride = owned.width }, owned.view(), plan, s.y, items, experts_n); + if (native.rank == 0) { + const glue = self.tri.own orelse return error.NativeGlueRequired; + try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[0].native.?, s.shared_gate, s.native_part); + try self.k.native_lane.matmul(self.stream, s.mixed, dims, n, native.shared[1].native.?, s.shared_up, s.native_part); + try glue.sharedSwiGLU(self.stream, s.shared_gate, s.shared_up, s.shared_act, n, moe_w); + try self.k.native_lane.matmulF32(self.stream, s.shared_act, moe_w, n, native.shared[2].native.?, s.shared_out, s.native_part); + try glue.sharedScatter(self.stream, s.shared_out, s.y, n, dims, slots); + } + // Multi-rank admission remains refused until the owning engine provides its collective. + if (native.world != 1) return error.TensorParallelUnavailable; + return; + } try self.k.router.route(self.stream, s.picks, pairs, e, rows_tile, plan); const items = cdiv(pairs, rows_tile) + e; try self.k.moe.run(self.stream, true, s.mixed, dims, slots, lw.experts_up.ptr, dims / 32, moe_w / 32, plan, s.moe_act, moe_w, items); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 03650116ab..e0e2296792 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -99,6 +99,7 @@ pub const Kernels = struct { moe: cuda.experts.RowsExperts, lane: cuda.Function, native_lane: cuda.qmmf.Lane, + nvfp4: cuda.experts.Nvfp4Experts, front: cuda.Function, chain: cuda.Function, chain_to: cuda.Function, @@ -125,6 +126,12 @@ pub const Kernels = struct { k.moe = try cuda.experts.RowsExperts.resolve(k.rows); k.native_lane = try cuda.qmmf.Lane.load(d, (try d.capability(0)) / 10); errdefer k.native_lane.unload(); + var device: c_int = 0; + var sms: c_int = 0; + try d.check(d.api.cuDeviceGet(&device, 0), "cuDeviceGet"); + try d.check(d.api.cuDeviceGetAttribute(&sms, .multiprocessor_count, device), "cuDeviceGetAttribute"); + k.nvfp4 = try cuda.experts.Nvfp4Experts.load(d, @intCast(sms)); + errdefer k.nvfp4.unload(); k.sum = try k.batch.function("fn_sum_slices"); k.shift = try k.batch.function("fn_shift_window"); k.select = try k.batch.function("fn_qsa_select"); @@ -136,6 +143,7 @@ pub const Kernels = struct { } pub fn deinit(k: *Kernels) void { + k.nvfp4.unload(); k.native_lane.unload(); k.batch.unload(); k.rows.unload(); @@ -241,6 +249,10 @@ pub const Scratch = struct { rank: u64 = 0, hist: u64 = 0, moe_act: u64 = 0, + shared_gate: u64 = 0, + shared_up: u64 = 0, + shared_act: u64 = 0, + shared_out: u64 = 0, y: u64 = 0, branch: u64 = 0, ple_emb: u64 = 0, @@ -318,6 +330,7 @@ pub const Scratch = struct { n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; n += att_rows * (idw + 2) * 4; n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; + n += cap * moe_w * 6 + cap * dims * 4; n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + half * 4 + 4096; @@ -373,6 +386,10 @@ pub const Scratch = struct { s.counts = try s.add(gpa, d, 8); s.rank = try s.add(gpa, d, pairs * 4); s.hist = try s.add(gpa, d, nblk * (experts_n + 1) * 4); + s.shared_gate = try s.add(gpa, d, cap * moe_w * 2); + s.shared_up = try s.add(gpa, d, cap * moe_w * 2); + s.shared_act = try s.add(gpa, d, cap * moe_w * 2); + s.shared_out = try s.add(gpa, d, cap * dims * 4); s.moe_act = try s.add(gpa, d, pairs * moe_w * 2); s.y = try s.add(gpa, d, pairs * dims * 4); s.branch = try s.add(gpa, d, cap * dims * 2); diff --git a/zig/src/families/flashnext/cuda_glue.zig b/zig/src/families/flashnext/cuda_glue.zig index efa304177d..30debfbb8a 100644 --- a/zig/src/families/flashnext/cuda_glue.zig +++ b/zig/src/families/flashnext/cuda_glue.zig @@ -29,6 +29,8 @@ const Fns = struct { wb_slices: cuda.Function, normalize: cuda.Function, dense_mix: cuda.Function, + shared_swiglu: cuda.Function, + shared_scatter: cuda.Function, down: cuda.Function, reduce_act: cuda.Function, upmix: cuda.Function, @@ -95,6 +97,8 @@ pub const Glue = struct { .wb_slices = try g.hc.function("fn_hc_writeback_slices"), .normalize = try g.hc.function("fn_hc_normalize"), .dense_mix = try g.hc.function("fn_hc_dense_mix"), + .shared_swiglu = try g.hc.function("fn_shared_swiglu"), + .shared_scatter = try g.hc.function("fn_shared_scatter"), .down = try g.hc.function("fn_hc_down"), .reduce_act = try g.hc.function("fn_hc_reduce_act"), .upmix = try g.hc.function("fn_hc_upmix"), @@ -229,6 +233,21 @@ pub const Glue = struct { try cuda.launch.launch(g.f.dense_mix, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(dims / 32) }, .block = .{ .x = 32 } }, s, &args); } + pub fn sharedSwiGLU(g: *const Glue, s: cuda.Stream, gate: u64, up: u64, act: u64, rows: usize, width: usize) !void { + var a: cuda.Args = .{}; + ptrs(&a, &.{ gate, up, act }); + ints(&a, &.{width}); + try go(g.f.shared_swiglu, s, .{ rows, cdiv(width, 256), 1 }, 256, 0, &a); + } + + pub fn sharedScatter(g: *const Glue, s: cuda.Stream, shared: u64, out: u64, rows: usize, dims: usize, slots: usize) !void { + if (slots == 0) return error.RowShape; + var a: cuda.Args = .{}; + ptrs(&a, &.{ shared, out }); + ints(&a, &.{ dims, slots }); + try go(g.f.shared_scatter, s, .{ rows, cdiv(dims, 256), 1 }, 256, 0, &a); + } + pub fn hcReduceAct(g: *const Glue, s: cuda.Stream, part: u64, act: u64, xs: u64, inj: u64, sk: usize, rows: usize, streams: usize, low: usize, ndn: usize, has_inj: usize) !void { const cols = low + has_inj * streams; if (cols > 1024 or sk > 32) return error.RowShape; diff --git a/zig/src/families/flashnext/cuda_mixed_experts.zig b/zig/src/families/flashnext/cuda_mixed_experts.zig new file mode 100644 index 0000000000..b26225731d --- /dev/null +++ b/zig/src/families/flashnext/cuda_mixed_experts.zig @@ -0,0 +1,36 @@ +//! Routed NVFP4 intermediate partitions and a separately owned MXFP8 shared expert. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const routed = @import("cuda_modelopt_experts.zig"); +const projections = @import("cuda_projection_load.zig"); + +pub const Owned = struct { + experts: routed.Owned, + shared: [3]wt.Q4, + rank: usize, + world: usize, + + pub fn free(self: *Owned) void { + self.experts.free(); + for (&self.shared) |*p| p.free(); + } +}; + +pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !Owned { + var e = try routed.load(a, d, store, prefix, world, rank); + errdefer e.free(); + var shared: [3]wt.Q4 = @splat(wt.Q4.empty(d)); + errdefer for (&shared) |*p| p.free(); + // The shared branch participates exactly once in the sum across ranks. + if (rank == 0) { + var buf: [256]u8 = undefined; + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }, 0..) |name, i| { + const n = if (i == 2) wt.dims else wt.moe_w; + const k = if (i == 2) wt.moe_w else wt.dims; + shared[i] = try projections.load(a, d, store, try std.fmt.bufPrint(&buf, "{s}.shared_expert.{s}", .{ prefix, name }), n, k); + if (shared[i].native == null) return error.UnsupportedQuantization; + } + } + return .{ .experts = e, .shared = shared, .rank = rank, .world = world }; +} diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index 9f2c1d3b9d..c2b9dde3ca 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const embed = @import("cuda_embed.zig"); const gdn = @import("cuda_gdn.zig"); const wt = @import("cuda_prompt_wt.zig"); +const mixed = @import("cuda_mixed_experts.zig"); const projections = @import("cuda_projection_load.zig"); const layers_n = wt.layers_n; @@ -78,6 +79,7 @@ pub const Layer = struct { /// Packed gate and up for every routed expert plus the shared expert at index 512. experts_up: cuda.DeviceBuffer, experts_down: cuda.DeviceBuffer, + native_experts: ?mixed.Owned = null, pub fn free(self: *Layer, gpa: std.mem.Allocator) void { _ = gpa; @@ -100,6 +102,7 @@ pub const Layer = struct { self.nq.free(); self.nc.free(); self.ple_cw.free(); + if (self.native_experts) |*e| e.free(); self.experts_up.free(); self.experts_down.free(); } @@ -192,13 +195,21 @@ fn stackUp(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, g } pub fn loadLayer(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, layer: usize) !Layer { + return loadLayerRank(gpa, driver, stream, store, pack_fn, layer, 1, 0); +} + +pub fn loadLayerRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, layer: usize, world: usize, rank: usize) !Layer { var buf: [80]u8 = undefined; const base = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}", .{layer}); - return loadLayerAt(gpa, driver, stream, store, pack_fn, base, linearLayer(layer), layer == ple_layer); + return loadLayerAtRank(gpa, driver, stream, store, pack_fn, base, linearLayer(layer), layer == ple_layer, world, rank); } /// The decoder layer under `base` (the MTP head's "language_model.mtp.layers.0": full attention, no n-gram branch). pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, base: []const u8, linear: bool, with_ple: bool) !Layer { + return loadLayerAtRank(gpa, driver, stream, store, pack_fn, base, linear, with_ple, 1, 0); +} + +fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.Stream, store: *Store, pack_fn: cuda.Function, base: []const u8, linear: bool, with_ple: bool, world: usize, rank: usize) !Layer { var buf: [200]u8 = undefined; var base_buf: [80]u8 = undefined; @memcpy(base_buf[0..base.len], base); @@ -330,33 +341,49 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); } - var shared: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; - errdefer for (&shared) |*f| f.free(gpa); - // a folder Living Weights taught holds a wider shared expert (its 640 neurons, then the pocket's) - const width = (try store.hold(try std.fmt.bufPrint(&buf, "{s}.shared_expert.gate_proj.weight", .{mlp_root}))).dim(0); - if (width != moe_w) return error.LearnedPocket; - shared[0] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.gate_proj", false, moe_w, dims); - shared[1] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.up_proj", false, moe_w, dims); - shared[2] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.down_proj", false, dims, moe_w); - var routed: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; - errdefer for (&routed) |*f| f.free(gpa); - routed[0] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.gate_proj", true, moe_w, dims); - routed[1] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.up_proj", true, moe_w, dims); - routed[2] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.down_proj", true, dims, moe_w); - var gate_p = try packAll(driver, stream, pack_fn, routed[0], shared[0]); - errdefer gate_p.free(); - var up_p = try packAll(driver, stream, pack_fn, routed[1], shared[1]); - errdefer up_p.free(); - var down_p = try packAll(driver, stream, pack_fn, routed[2], shared[2]); - errdefer down_p.free(); - for (&shared) |*f| f.free(gpa); - for (&routed) |*f| f.free(gpa); - shared = .{ blankFace(), blankFace(), blankFace() }; - routed = .{ blankFace(), blankFace(), blankFace() }; - var experts_up = try stackUp(gpa, driver, stream, gate_p, up_p); + var native_experts: ?mixed.Owned = null; + errdefer if (native_experts) |*e| e.free(); + var experts_up = deadBuf(driver); errdefer experts_up.free(); - gate_p.free(); - up_p.free(); + var experts_down = deadBuf(driver); + errdefer experts_down.free(); + if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { + native_experts = try mixed.load(gpa, driver, store, mlp_root, world, rank); + } else { + if (world != 1 or rank != 0) return error.UnsupportedQuantization; + const legacy = legacy_block: { + var shared: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&shared) |*f| f.free(gpa); + // a folder Living Weights taught holds a wider shared expert (its 640 neurons, then the pocket's) + const width = (try store.hold(try std.fmt.bufPrint(&buf, "{s}.shared_expert.gate_proj.weight", .{mlp_root}))).dim(0); + if (width != moe_w) return error.LearnedPocket; + shared[0] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.gate_proj", false, moe_w, dims); + shared[1] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.up_proj", false, moe_w, dims); + shared[2] = try keepFace(gpa, store, &buf, mlp_root, ".shared_expert.down_proj", false, dims, moe_w); + var routed: [3]OwnedFace = .{ blankFace(), blankFace(), blankFace() }; + errdefer for (&routed) |*f| f.free(gpa); + routed[0] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.gate_proj", true, moe_w, dims); + routed[1] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.up_proj", true, moe_w, dims); + routed[2] = try keepFace(gpa, store, &buf, mlp_root, ".switch_mlp.down_proj", true, dims, moe_w); + var gate_p = try packAll(driver, stream, pack_fn, routed[0], shared[0]); + errdefer gate_p.free(); + var up_p = try packAll(driver, stream, pack_fn, routed[1], shared[1]); + errdefer up_p.free(); + var down_p = try packAll(driver, stream, pack_fn, routed[2], shared[2]); + errdefer down_p.free(); + for (&shared) |*f| f.free(gpa); + for (&routed) |*f| f.free(gpa); + shared = .{ blankFace(), blankFace(), blankFace() }; + routed = .{ blankFace(), blankFace(), blankFace() }; + var combined_up = try stackUp(gpa, driver, stream, gate_p, up_p); + errdefer combined_up.free(); + gate_p.free(); + up_p.free(); + break :legacy_block .{ combined_up, down_p }; + }; + experts_up = legacy[0]; + experts_down = legacy[1]; + } store.release(); return .{ .attn_hc = attn_hc, @@ -380,6 +407,7 @@ pub fn loadLayerAt(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S .nc = nc_b, .ple_cw = pcw_b, .experts_up = experts_up, - .experts_down = down_p, + .experts_down = experts_down, + .native_experts = native_experts, }; } diff --git a/zig/tests/cuda/flashnext_shared_tests.zig b/zig/tests/cuda/flashnext_shared_tests.zig new file mode 100644 index 0000000000..1077d0e4dc --- /dev/null +++ b/zig/tests/cuda/flashnext_shared_tests.zig @@ -0,0 +1,125 @@ +//! Independent stage-wise FP64 references for the separate MXFP8 shared expert and its reserved slot. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const d = 2560; +const w = 640; +const most = 17; +const slots = 11; + +const Weight = struct { + codes: []u8, + scales: []u8, + n: usize, + k: usize, + + fn init(a: std.mem.Allocator, n: usize, k: usize, salt: usize) !Weight { + const codes = try a.alloc(u8, n * k); + const scales = try a.alloc(u8, n * k / 32); + for (codes, 0..) |*v, i| v.* = @intCast(0x20 + ((i * 13 + salt * 7) % 32) + (if ((i + salt) % 3 == 0) @as(usize, 128) else 0)); + for (scales, 0..) |*v, i| v.* = @intCast(121 + (i + salt) % 3); + return .{ .codes = codes, .scales = scales, .n = n, .k = k }; + } + + fn value(self: Weight, n: usize, k: usize) f64 { + const exponent: i32 = @as(i32, self.scales[n * (self.k / 32) + k / 32]) - 127; + return core.quant.fp8.e4m3(self.codes[n * self.k + k]) * std.math.pow(f64, 2, @floatFromInt(exponent)); + } + + fn upload(self: Weight, a: std.mem.Allocator, driver: *const cuda.Driver) !flash.projection.Projection { + return flash.projection.mxfp8(a, driver, tensor(.f8_e4m3, self.codes, self.n, self.k), tensor(.u8, self.scales, self.n, self.k / 32)); + } +}; + +fn tensor(dtype: core.safetensors.DType, bytes: []const u8, n: usize, k: usize) core.safetensors.Tensor { + return .{ .dtype = dtype, .rank = 2, .shape = .{ n, k, 1, 1, 1 }, .bytes = bytes }; +} + +fn projectionReference(r: *rig.Rig, name: []const u8, x: []const u16, weight: Weight, out: anytype, rows: usize, comptime bf16: bool) !void { + for (0..rows) |row| for (0..weight.n) |n| { + var total: f64 = 0; + var magnitude: f64 = 0; + for (0..weight.k) |k| { + const term = rig.bf(x[row * weight.k + k]) * weight.value(n, k); + total += term; + magnitude += @abs(term); + } + const at = row * weight.n + n; + if (bf16) { + try r.within(name, rig.bf(out[at]), rig.rbf(total), rig.ulps(rig.rbf(total), 1)); + } else try r.within(name, out[at], total, rig.sumTol(magnitude, weight.k) + 1e-6); + }; +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0x5a4ed); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + defer lane.unload(); + const gw = try Weight.init(a, w, d, 1); + const uw = try Weight.init(a, w, d, 2); + const dw = try Weight.init(a, d, w, 3); + var gate = try gw.upload(a, gpu.d); + defer gate.free(); + var up = try uw.upload(a, gpu.d); + defer up.free(); + var down = try dw.upload(a, gpu.d); + defer down.free(); + const x = r.bfs(most * d, 0.125); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const g = try r.out(most * w * 2); + const u = try r.out(most * w * 2); + const act = try r.out(most * w * 2); + const out = try r.out(most * d * 4); + const temp = try r.out(8 * most * d * 4); + const sentinel = try a.alloc(f32, most * slots * d); + @memset(sentinel, 0.375); + const y = try r.dev(std.mem.sliceAsBytes(sentinel)); + var all_g: []u16 = undefined; + var all_u: []u16 = undefined; + var all_act: []u16 = undefined; + var all_out: []f32 = undefined; + for ([_]usize{ most, 1, 7 }) |rows| { + try lane.matmul(r.s, dx, d, rows, gate.native.?, g, temp); + try lane.matmul(r.s, dx, d, rows, up.native.?, u, temp); + try glue.sharedSwiGLU(r.s, g, u, act, rows, w); + try lane.matmulF32(r.s, act, w, rows, down.native.?, out, temp); + try glue.sharedScatter(r.s, out, y, rows, d, slots); + if (rows == most) { + all_g = try r.back(u16, g, most * w); + all_u = try r.back(u16, u, most * w); + all_act = try r.back(u16, act, most * w); + all_out = try r.back(f32, out, most * d); + } else { + try r.sameRows("shared gate rows", u16, g, all_g[0 .. rows * w]); + try r.sameRows("shared up rows", u16, u, all_u[0 .. rows * w]); + try r.sameRows("shared activation rows", u16, act, all_act[0 .. rows * w]); + try r.sameRows("shared FP32 output rows", f32, out, all_out[0 .. rows * d]); + } + } + try projectionReference(&r, "shared MXFP8 gate", x, gw, all_g, most, true); + try projectionReference(&r, "shared MXFP8 up", x, uw, all_u, most, true); + for (all_g, all_u, all_act) |gate_value, up_value, actual| { + const f = rig.bf(gate_value); + const want = rig.rbf(rig.rbf(f / (1 + @exp(-f))) * rig.bf(up_value)); + try r.within("shared SwiGLU rounding", rig.bf(actual), want, rig.ulps(want, 1)); + } + try projectionReference(&r, "shared MXFP8 FP32 down", all_act, dw, all_out, most, false); + const scattered = try r.back(f32, y, most * slots * d); + for (scattered, 0..) |actual, i| { + const slot = (i / d) % slots; + const want = if (slot == slots - 1) all_out[(i / (slots * d)) * d + i % d] else @as(f32, 0.375); + try check.expect(actual == want, "shared scatter preserves routed slots and copies exact FP32 values", .{}); + } + try std.testing.expectError(error.RowShape, glue.sharedScatter(r.s, out, y, most, d, 0)); + r.report("MXFP8 shared expert, BF16 projection/activation seams, FP32 down, reserved-slot scatter; rows 1/7/17 equal"); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 54ce9fd6f4..f2fe435be5 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -33,6 +33,7 @@ const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); +const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); @@ -60,6 +61,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) @@ -151,6 +153,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); From 359eef387d94bb1ca1cbdce80ca09d9ad29505d4 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:20:38 -0700 Subject: [PATCH 068/115] Bound CUDA rank control and NCCL collective failures --- zig/src/cuda/nccl.zig | 33 +++ .../families/flashnext/cuda_collective.zig | 87 ++++++++ zig/src/families/flashnext/cuda_launch.zig | 3 + zig/src/families/flashnext/cuda_rank_wire.zig | 207 ++++++++++++++++++ zig/tests/cuda/flashnext_rank_tests.zig | 84 +++++++ zig/tests/cuda/main.zig | 3 + 6 files changed, 417 insertions(+) create mode 100644 zig/src/families/flashnext/cuda_collective.zig create mode 100644 zig/src/families/flashnext/cuda_rank_wire.zig create mode 100644 zig/tests/cuda/flashnext_rank_tests.zig diff --git a/zig/src/cuda/nccl.zig b/zig/src/cuda/nccl.zig index f742d62549..35240f8808 100644 --- a/zig/src/cuda/nccl.zig +++ b/zig/src/cuda/nccl.zig @@ -10,6 +10,31 @@ pub const UniqueId = extern struct { internal: [128]u8 }; pub const DataType = enum(c_int) { i8 = 0, u8 = 1, i32 = 2, u32 = 3, i64 = 4, u64 = 5, f16 = 6, f32 = 7, f64 = 8, bf16 = 9 }; pub const RedOp = enum(c_int) { sum = 0, prod = 1, max = 2, min = 3, avg = 4 }; +/// NCCL 2.30 public ncclConfig_v23000 ABI. New nonblocking rank owners explicitly require 2.30 or later. +pub const Config = extern struct { + size: usize = @sizeOf(Config), + magic: u32 = 0xcafebeef, + version: u32 = 23000, + blocking: c_int = 0, + cga_cluster_size: c_int = std.math.minInt(c_int), + min_ctas: c_int = std.math.minInt(c_int), + max_ctas: c_int = std.math.minInt(c_int), + net_name: ?[*:0]const u8 = null, + split_share: c_int = std.math.minInt(c_int), + traffic_class: c_int = std.math.minInt(c_int), + comm_name: ?[*:0]const u8 = null, + collnet_enable: c_int = std.math.minInt(c_int), + cta_policy: c_int = std.math.minInt(c_int), + shrink_share: c_int = std.math.minInt(c_int), + nvls_ctas: c_int = std.math.minInt(c_int), + n_channels_per_net_peer: c_int = std.math.minInt(c_int), + nvlink_centric_sched: c_int = std.math.minInt(c_int), + graph_usage_mode: c_int = std.math.minInt(c_int), + num_rma_ctx: c_int = std.math.minInt(c_int), + max_p2p_peers: c_int = std.math.minInt(c_int), + graph_stream_ordering: c_int = std.math.minInt(c_int), +}; + pub const Error = error{ LibraryUnavailable, MissingSymbol, NcclFailed }; const R = Result; @@ -69,3 +94,11 @@ pub const Library = struct { return v; } }; + +test "nonblocking communicator config preserves the public ABI defaults" { + const c: Config = .{}; + try std.testing.expectEqual(@as(usize, 96), @sizeOf(Config)); + try std.testing.expectEqual(@as(usize, 32), @offsetOf(Config, "net_name")); + try std.testing.expectEqual(@as(usize, 48), @offsetOf(Config, "comm_name")); + try std.testing.expect(c.size == 96 and c.blocking == 0 and c.max_ctas == std.math.minInt(c_int)); +} diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig new file mode 100644 index 0000000000..a52924a586 --- /dev/null +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -0,0 +1,87 @@ +//! One two-rank NCCL owner. Nonblocking initialization and collectives observe the control peer and a deadline. +const std = @import("std"); +const cuda = @import("cuda"); +const wire = @import("cuda_rank_wire.zig"); +const nccl = cuda.nccl; + +const Advanced = struct { + ncclCommInitRankConfig: *const fn (*nccl.Comm, c_int, nccl.UniqueId, c_int, *nccl.Config) callconv(.c) nccl.Result, + ncclCommAbort: *const fn (nccl.Comm) callconv(.c) nccl.Result, + ncclCommGetAsyncError: *const fn (nccl.Comm, *nccl.Result) callconv(.c) nccl.Result, +}; + +pub const Collective = struct { + lib: nccl.Library, + api: Advanced, + comm: nccl.Comm = null, + link: wire.Link, + timeout_ms: u64, + failed: bool = false, + + /// `link` has completed the source/settings handshake; it stays owned by the caller. + pub fn init(link: wire.Link, rank: usize, timeout_ms: u64) !Collective { + if (rank > 1 or timeout_ms == 0) return error.BadRank; + var c: Collective = .{ .lib = try nccl.Library.open(), .api = undefined, .link = link, .timeout_ms = timeout_ms }; + errdefer c.lib.close(); + if (try c.lib.version() < 23000) return error.NcclVersion; + const info = @typeInfo(Advanced).@"struct"; + inline for (info.field_names, info.field_types) |name, T| @field(c.api, name) = c.lib.lib.lookup(T, name) orelse return error.MissingSymbol; + var id: nccl.UniqueId = undefined; + const until = wire.deadline(timeout_ms); + if (rank == 0) { + try c.lib.check(c.lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try link.send(std.mem.asBytes(&id), until); + } else try link.receive(std.mem.asBytes(&id), until); + var config: nccl.Config = .{}; + errdefer c.abort(); + try c.submitted(c.api.ncclCommInitRankConfig(&c.comm, 2, id, @intCast(rank), &config)); + try c.wait(null, until); + return c; + } + + pub fn abort(c: *Collective) void { + c.failed = true; + if (c.comm) |comm| { + _ = c.api.ncclCommAbort(comm); + c.comm = null; + } + } + + pub fn deinit(c: *Collective) void { + // Abort is local and also releases an idle communicator if its peer exited first. + c.abort(); + c.lib.close(); + } + + fn submitted(c: *Collective, result: nccl.Result) !void { + if (result == 0 or result == 7) return; + return c.lib.check(result, "NCCL nonblocking operation"); + } + + fn wait(c: *Collective, stream: ?cuda.Stream, until: u64) !void { + errdefer c.abort(); + while (true) { + try c.link.checkPeer(); + if (wire.now() >= until) return error.RankTimeout; + var state: nccl.Result = 0; + try c.lib.check(c.api.ncclCommGetAsyncError(c.comm, &state), "ncclCommGetAsyncError"); + if (state != 0 and state != 7) return c.lib.check(state, "NCCL asynchronous operation"); + if (state == 0) { + if (stream) |s| { + if (try s.done()) return; + } else return; + } + wire.pause(); + } + } + + /// FP32 sum is complete on the compute stream on return, or the communicator is permanently aborted. + pub fn sum(c: *Collective, stream: cuda.Stream, values: u64, count: usize) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + if (count == 0) return; + errdefer c.abort(); + try c.link.checkPeer(); + try c.submitted(c.lib.api.ncclAllReduce(values, values, count, .f32, .sum, c.comm, stream.handle)); + try c.wait(stream, wire.deadline(c.timeout_ms)); + } +}; diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 1367aec3a0..186902c090 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,6 +6,8 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const projection = @import("cuda_projection.zig"); +pub const collective = @import("cuda_collective.zig"); +pub const rank_wire = @import("cuda_rank_wire.zig"); pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); pub const hc = @import("cuda_hc.zig"); pub const qmm = @import("cuda_qmm.zig"); @@ -372,6 +374,7 @@ test { _ = @import("cuda_snap.zig"); _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); + _ = @import("cuda_rank_wire.zig"); _ = @import("cuda_modelopt_experts.zig"); _ = row_path; _ = @import("cuda_mtp_check.zig"); diff --git a/zig/src/families/flashnext/cuda_rank_wire.zig b/zig/src/families/flashnext/cuda_rank_wire.zig new file mode 100644 index 0000000000..2e0c4fe361 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_wire.zig @@ -0,0 +1,207 @@ +//! Bounded Linux TCP control channel. Tensor bytes travel over NCCL; this channel owns command order and peer EOF. +const std = @import("std"); +const builtin = @import("builtin"); +const p = std.posix; +const net = std.Io.net; + +pub const max_payload = 1 << 20; +pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack }; +pub const Header = struct { + kind: Kind, + sequence: u64 = 0, + object: u64 = 0, + arg: u64 = 0, + bytes: u32 = 0, + flags: u32 = 0, + + pub fn encode(h: Header) ![40]u8 { + if (h.bytes > max_payload) return error.BadRankFrame; + var out: [40]u8 = @splat(0); + @memcpy(out[0..4], "TFR1"); + std.mem.writeInt(u32, out[4..8], @backingInt(h.kind), .little); + std.mem.writeInt(u64, out[8..16], h.sequence, .little); + std.mem.writeInt(u64, out[16..24], h.object, .little); + std.mem.writeInt(u64, out[24..32], h.arg, .little); + std.mem.writeInt(u32, out[32..36], h.bytes, .little); + std.mem.writeInt(u32, out[36..40], h.flags, .little); + return out; + } + + pub fn decode(raw: *const [40]u8) !Header { + if (!std.mem.eql(u8, raw[0..4], "TFR1")) return error.BadRankFrame; + const h: Header = .{ + .kind = std.enums.fromInt(Kind, std.mem.readInt(u32, raw[4..8], .little)) orelse return error.BadRankFrame, + .sequence = std.mem.readInt(u64, raw[8..16], .little), + .object = std.mem.readInt(u64, raw[16..24], .little), + .arg = std.mem.readInt(u64, raw[24..32], .little), + .bytes = std.mem.readInt(u32, raw[32..36], .little), + .flags = std.mem.readInt(u32, raw[36..40], .little), + }; + if (h.bytes > max_payload) return error.BadRankFrame; + return h; + } +}; + +pub fn now() u64 { + var ts: std.c.timespec = undefined; + if (std.c.clock_gettime(.MONOTONIC, &ts) != 0 or ts.sec < 0) return 0; + return @as(u64, @intCast(ts.sec)) * std.time.ns_per_ms * 1000 + @as(u64, @intCast(ts.nsec)); +} + +pub fn deadline(ms: u64) u64 { + return now() +| (ms *| std.time.ns_per_ms); +} + +pub fn pause() void { + const ts: std.c.timespec = .{ .sec = 0, .nsec = 50_000 }; + _ = std.c.nanosleep(&ts, null); +} + +fn wait(fd: p.socket_t, events: i16, until: ?u64) !void { + while (true) { + const current = now(); + if (until) |end| if (current >= end) return error.RankTimeout; + const ms: i32 = if (until) |end| @intCast(@min(1000, (end - current + std.time.ns_per_ms - 1) / std.time.ns_per_ms)) else 1000; + var fds = [_]p.pollfd{.{ .fd = fd, .events = events, .revents = 0 }}; + const count = try p.poll(&fds, ms); + if (count == 0) continue; + if (fds[0].revents & events != 0) return; + if (fds[0].revents & (p.POLL.ERR | p.POLL.HUP | p.POLL.NVAL) != 0) return error.RankDisconnected; + } +} + +pub const Link = struct { + fd: p.socket_t, + + pub fn close(l: Link) void { + _ = std.c.shutdown(l.fd, p.SHUT.RDWR); + _ = std.c.close(l.fd); + } + + /// Exactly two rank processes. Retry only connection establishment, never an accepted command. + pub fn open(io: std.Io, host: []const u8, port: u16, rank: usize, timeout_ms: u64) !Link { + if (builtin.os.tag != .linux) return error.RankPlatform; + if (rank > 1 or port == 0) return error.BadRank; + const address = net.IpAddress.parse(host, port) catch try net.IpAddress.resolve(io, host, port); + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&address, &storage); + const family: u32 = if (address == .ip4) p.AF.INET else p.AF.INET6; + const until = deadline(timeout_ms); + if (rank == 0) { + const fd = try socket(family); + defer _ = std.c.close(fd); + const one: c_int = 1; + try p.setsockopt(fd, p.SOL.SOCKET, p.SO.REUSEADDR, std.mem.asBytes(&one)); + if (std.c.bind(fd, &storage.any, len) != 0) return error.RankBind; + if (std.c.listen(fd, 1) != 0) return error.RankListen; + while (true) { + try wait(fd, p.POLL.IN, until); + const peer = std.c.accept4(fd, null, null, p.SOCK.NONBLOCK | p.SOCK.CLOEXEC); + switch (p.errno(peer)) { + .SUCCESS => return .{ .fd = peer }, + .INTR, .AGAIN => continue, + else => return error.RankAccept, + } + } + } + while (now() < until) { + const fd = try socket(family); + const rc = std.c.connect(fd, &storage.any, len); + if (rc == 0) return .{ .fd = fd }; + const e = p.errno(rc); + if (e == .INPROGRESS or e == .INTR) { + wait(fd, p.POLL.OUT, until) catch |err| { + _ = std.c.close(fd); + if (err == error.RankTimeout) return err; + pause(); + continue; + }; + var status: c_int = 0; + var size: p.socklen_t = @sizeOf(c_int); + if (std.c.getsockopt(fd, p.SOL.SOCKET, p.SO.ERROR, &status, &size) == 0 and status == 0) return .{ .fd = fd }; + } + _ = std.c.close(fd); + pause(); + } + return error.RankTimeout; + } + + pub fn send(l: Link, bytes: []const u8, until: u64) !void { + var at: usize = 0; + while (at < bytes.len) { + if (now() >= until) return error.RankTimeout; + const rc = std.c.send(l.fd, bytes.ptr + at, bytes.len - at, p.MSG.NOSIGNAL); + switch (p.errno(rc)) { + .SUCCESS => { + if (rc == 0) return error.RankDisconnected; + at += @intCast(rc); + }, + .INTR => continue, + .AGAIN => try wait(l.fd, p.POLL.OUT, until), + else => return error.RankDisconnected, + } + } + } + + pub fn receive(l: Link, bytes: []u8, until: ?u64) !void { + var at: usize = 0; + while (at < bytes.len) { + if (until) |end| if (now() >= end) return error.RankTimeout; + const rc = std.c.recv(l.fd, bytes.ptr + at, bytes.len - at, 0); + switch (p.errno(rc)) { + .SUCCESS => { + if (rc == 0) return error.RankDisconnected; + at += @intCast(rc); + }, + .INTR => continue, + .AGAIN => try wait(l.fd, p.POLL.IN, until), + else => return error.RankDisconnected, + } + } + } + + pub fn sendHeader(l: Link, h: Header, until: u64) !void { + const raw = try h.encode(); + try l.send(&raw, until); + } + + pub fn receiveHeader(l: Link, until: ?u64) !Header { + var raw: [40]u8 = undefined; + // An idle follower can wait indefinitely; a partial frame cannot. + try l.receive(raw[0..1], until); + try l.receive(raw[1..], until orelse deadline(30_000)); + return Header.decode(&raw); + } + + /// A non-consuming EOF check while a GPU collective progresses. A future control frame stays queued. + pub fn checkPeer(l: Link) !void { + var byte: u8 = 0; + const rc = std.c.recv(l.fd, &byte, 1, p.MSG.PEEK | p.MSG.DONTWAIT); + if (rc == 0) return error.RankDisconnected; + if (rc < 0) switch (p.errno(rc)) { + .AGAIN, .INTR => {}, + else => return error.RankDisconnected, + }; + } +}; + +fn socket(family: u32) !p.socket_t { + const fd = std.c.socket(family, p.SOCK.STREAM | p.SOCK.NONBLOCK | p.SOCK.CLOEXEC, 0); + if (p.errno(fd) != .SUCCESS) return error.RankSocket; + return fd; +} + +test "rank frames are versioned, bounded and exact" { + const h: Header = .{ .kind = .forward, .sequence = 700, .object = 19, .arg = 42, .bytes = max_payload, .flags = 3 }; + var raw = try h.encode(); + try std.testing.expectEqualDeep(h, try Header.decode(&raw)); + raw[0] = 0; + try std.testing.expectError(error.BadRankFrame, Header.decode(&raw)); + raw = try h.encode(); + @memset(raw[4..8], 0xff); + try std.testing.expectError(error.BadRankFrame, Header.decode(&raw)); + raw = try h.encode(); + @memset(raw[32..36], 0xff); + try std.testing.expectError(error.BadRankFrame, Header.decode(&raw)); + try std.testing.expectError(error.BadRankFrame, (Header{ .kind = .forward, .bytes = max_payload + 1 }).encode()); +} diff --git a/zig/tests/cuda/flashnext_rank_tests.zig b/zig/tests/cuda/flashnext_rank_tests.zig new file mode 100644 index 0000000000..dbbd504fc6 --- /dev/null +++ b/zig/tests/cuda/flashnext_rank_tests.zig @@ -0,0 +1,84 @@ +//! Real two-process/two-device NCCL sums and bounded disconnect/timeout behavior over the framed control link. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const wire = flash.rank_wire; + +pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { + const missing = std.mem.eql(u8, mode, "missing"); + const started = wire.now(); + if (missing) { + try std.testing.expectError(error.RankTimeout, wire.Link.open(gpu.io, "127.0.0.1", port, rank, 200)); + try check.expect(wire.now() - started < 3 * std.time.ns_per_s, "missing peer startup is bounded", .{}); + check.pass("rank {d}: missing peer startup timed out within bound", .{rank}); + return; + } + const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 30_000); + defer link.close(); + const until = wire.deadline(30_000); + const hello: wire.Header = .{ .kind = .hello, .sequence = 1, .object = 77, .arg = 123, .bytes = 7 }; + if (rank == 0) { + const raw = try hello.encode(); + try link.send(raw[0..3], until); + try link.send(raw[3..], until); + try link.send("payload", until); + const ack = try link.receiveHeader(until); + try check.expect(ack.kind == .ack and ack.sequence == 1 and ack.bytes == 0, "exact hello acknowledgment", .{}); + } else { + try std.testing.expectEqualDeep(hello, try link.receiveHeader(until)); + var payload: [7]u8 = undefined; + try link.receive(&payload, until); + try check.sameBytes("fragmented control frame payload", &payload, "payload"); + try link.sendHeader(.{ .kind = .ack, .sequence = 1 }, until); + } + var collective = try flash.collective.Collective.init(link, rank, if (std.mem.eql(u8, mode, "timeout")) 200 else 30_000); + defer collective.deinit(); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + if (std.mem.eql(u8, mode, "eof") or std.mem.eql(u8, mode, "timeout")) { + if (rank == 1) { + // A deliberate peer omission: the owner either disconnects or stays alive without entering the collective. + if (std.mem.eql(u8, mode, "eof")) { + _ = std.c.shutdown(link.fd, std.posix.SHUT.RDWR); + } else { + const delay: std.c.timespec = .{ .sec = 0, .nsec = 600_000_000 }; + _ = std.c.nanosleep(&delay, null); + } + check.pass("rank 1: injected {s}", .{mode}); + return; + } + var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&@as(f32, 1))); + defer values.free(); + const start = wire.now(); + collective.sum(stream, values.ptr, 1) catch |err| { + try check.expect(err == error.RankDisconnected or err == error.RankTimeout or err == error.NcclFailed, "failure is surfaced ({s})", .{@errorName(err)}); + try check.expect(collective.failed and collective.comm == null, "failed communicator is aborted and cannot be reused", .{}); + try check.expect(wire.now() - start < 3 * std.time.ns_per_s, "failed collective terminates within bound", .{}); + try std.testing.expectError(error.RankDisconnected, collective.sum(stream, values.ptr, 1)); + check.pass("rank 0: {s} surfaced as {s}, aborted, and refused reuse within bound", .{ mode, @errorName(err) }); + return; + }; + return error.ExpectedRankFailure; + } + for ([_]usize{ 0, 1, 17, 4099, 1_000_003 }, 0..) |count, seq| { + const host = try gpu.gpa.alloc(f32, @max(1, count)); + defer gpu.gpa.free(host); + for (host, 0..) |*v, i| v.* = @floatFromInt(i % 127 + rank); + var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer values.free(); + try collective.sum(stream, values.ptr, count); + try values.download(0, std.mem.sliceAsBytes(host)); + for (host[0..count], 0..) |actual, i| try check.expect(actual == @as(f32, @floatFromInt(2 * (i % 127) + 1)), "rank {d}: exact sum at {d}", .{ rank, i }); + // Acknowledge completion before either peer can close or start its next reduction. + const h: wire.Header = .{ .kind = .ready, .sequence = seq + 2, .arg = count }; + if (rank == 0) { + try link.sendHeader(h, wire.deadline(30_000)); + try std.testing.expectEqualDeep(h, try link.receiveHeader(wire.deadline(30_000))); + } else { + try std.testing.expectEqualDeep(h, try link.receiveHeader(wire.deadline(30_000))); + try link.sendHeader(h, wire.deadline(30_000)); + } + } + check.pass("rank {d}: two-process NCCL FP32 sums exact for empty, scalar, tail and 1,000,003 values; fragmented framed control preserved", .{rank}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index f2fe435be5..92ab44b213 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -33,6 +33,7 @@ const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); +const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); @@ -61,6 +62,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 @@ -153,6 +155,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum"); if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); From 567be14f2756047462660d8bc7aa7c92514ad567 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:46:07 -0700 Subject: [PATCH 069/115] Add bounded Flash Next rank replay and checkpoint norm conversion --- zig/src/families/flashnext/cuda_batch.zig | 6 +- .../families/flashnext/cuda_batch_state.zig | 2 + zig/src/families/flashnext/cuda_launch.zig | 2 + .../flashnext/cuda_modelopt_experts.zig | 6 +- zig/src/families/flashnext/cuda_norms.zig | 27 ++ zig/src/families/flashnext/cuda_prompt.zig | 16 +- .../families/flashnext/cuda_prompt_load.zig | 14 +- zig/src/families/flashnext/cuda_prompt_wt.zig | 131 +++++- .../families/flashnext/cuda_rank_engine.zig | 377 ++++++++++++++++++ zig/src/families/flashnext/cuda_snap.zig | 1 + zig/tests/cuda/flashnext_model_rank_tests.zig | 77 ++++ zig/tests/cuda/main.zig | 3 + 12 files changed, 642 insertions(+), 20 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_norms.zig create mode 100644 zig/src/families/flashnext/cuda_rank_engine.zig create mode 100644 zig/tests/cuda/flashnext_model_rank_tests.zig diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index f45879407b..621523de75 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -384,8 +384,10 @@ pub fn Pass(comptime Tri: type) type { try self.k.native_lane.matmulF32(self.stream, s.shared_act, moe_w, n, native.shared[2].native.?, s.shared_out, s.native_part); try glue.sharedScatter(self.stream, s.shared_out, s.y, n, dims, slots); } - // Multi-rank admission remains refused until the owning engine provides its collective. - if (native.world != 1) return error.TensorParallelUnavailable; + if (native.world != 1) { + const collective = self.m.collective orelse return error.TensorParallelUnavailable; + try collective.sum(self.stream, s.y, pairs * dims); + } return; } try self.k.router.route(self.stream, s.picks, pairs, e, rows_tile, plan); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index e0e2296792..72a62dc897 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -323,6 +323,7 @@ pub const Scratch = struct { var n: usize = 0; n += cap * 4 + (cap + 1) * 16; n += cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; + n += cap * wide * 2 + 8 * cap * dims * 4; n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; @@ -459,6 +460,7 @@ pub const Lane = struct { seq: SeqState, pos: usize }; /// The engine's resident weights the pass reads. pub const Model = struct { + collective: ?*@import("cuda_collective.zig").Collective = null, table_w: u64, table_s: u64, table_b: u64, diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 186902c090..b2c54b1fe0 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,6 +6,7 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const projection = @import("cuda_projection.zig"); +pub const rank_engine = @import("cuda_rank_engine.zig"); pub const collective = @import("cuda_collective.zig"); pub const rank_wire = @import("cuda_rank_wire.zig"); pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); @@ -374,6 +375,7 @@ test { _ = @import("cuda_snap.zig"); _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); + _ = @import("cuda_rank_engine.zig"); _ = @import("cuda_rank_wire.zig"); _ = @import("cuda_modelopt_experts.zig"); _ = row_path; diff --git a/zig/src/families/flashnext/cuda_modelopt_experts.zig b/zig/src/families/flashnext/cuda_modelopt_experts.zig index 553a6be9a9..b6ec493320 100644 --- a/zig/src/families/flashnext/cuda_modelopt_experts.zig +++ b/zig/src/families/flashnext/cuda_modelopt_experts.zig @@ -44,8 +44,7 @@ pub const Owned = struct { } }; -/// The reader owns only the selected rectangle, not a decoded full expert stack. -/// The store is an exclusive loader scratch here: every borrowed plane is released before returning. +/// Own only the selected rectangle; the exclusive loader store releases every borrowed plane before return. pub fn readPart(a: std.mem.Allocator, store: anytype, name: []const u8, cut: Cut, down: bool, dims: usize, full_width: usize) !slice.Linear { defer store.release(); var buf: [512]u8 = undefined; @@ -59,8 +58,7 @@ pub fn readPart(a: std.mem.Allocator, store: anytype, name: []const u8, cut: Cut return linear; } -/// Pack and upload one expert at a time, bounding host scratch independently of expert count. -/// Activations remain BF16 (W4A16); input_scale is validated as metadata by the source reader. +/// Pack one expert at a time with BF16 activations (W4A16); input_scale is validated as metadata. pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, world: usize, rank: usize) !Owned { return loadShape(a, d, store, prefix, world, rank, wt.dims, wt.moe_w, wt.experts_n); } diff --git a/zig/src/families/flashnext/cuda_norms.zig b/zig/src/families/flashnext/cuda_norms.zig new file mode 100644 index 0000000000..f257ac56ca --- /dev/null +++ b/zig/src/families/flashnext/cuda_norms.zig @@ -0,0 +1,27 @@ +//! Classify the checkpoint-wide centered or direct RMSNorm convention, rejecting ambiguous mixtures. +const std = @import("std"); + +pub fn offset(comptime layers_n: usize, input: [layers_n]f64) !f32 { + comptime std.debug.assert(layers_n >= 2 and layers_n % 2 == 0); + var means = input; + var above: usize = 0; + for (means) |mean| { + if (!std.math.isFinite(mean)) return error.InvalidNormConvention; + if (mean > 0.5) above += 1; + } + std.mem.sort(f64, &means, {}, std.sort.asc(f64)); + const median = (means[layers_n / 2 - 1] + means[layers_n / 2]) / 2; + if (above * 10 >= layers_n * 9 and median >= 0.75 and median <= 1.5) return 0; + if (above * 10 <= layers_n and median >= -0.5 and median <= 0.25) return 1; + return error.InvalidNormConvention; +} + +test "centered and direct norms have distinct offsets; mixed and nonfinite conventions fail" { + try std.testing.expectEqual(@as(f32, 1), try offset(48, @splat(-0.1))); + try std.testing.expectEqual(@as(f32, 0), try offset(48, @splat(1.05))); + var mixed: [48]f64 = @splat(1.0); + @memset(mixed[0 .. 48 / 2], 0); + try std.testing.expectError(error.InvalidNormConvention, offset(48, mixed)); + mixed[0] = std.math.nan(f64); + try std.testing.expectError(error.InvalidNormConvention, offset(48, mixed)); +} diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 032f7b596a..91c316f91f 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -133,9 +133,15 @@ pub const Engine = struct { scores: ?cuda.DeviceBuffer = null, /// Rows a batched chunk takes at most (TF_FLASHNEXT_CHUNK; batch.default_chunk by default). chunk: usize, + collective: ?*@import("cuda_collective.zig").Collective = null, pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { + return initRank(gpa, io, driver, stream, model_dir, 1, 0); + } + + pub fn initRank(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8, world: usize, rank: usize) !*Engine { if (!cuda.kernels.available) return error.BuiltWithoutKernels; + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; const e = try gpa.create(Engine); errdefer gpa.destroy(e); e.gpa = gpa; @@ -143,11 +149,13 @@ pub const Engine = struct { e.driver = driver; e.stream = stream; e.n_layers = 0; + e.collective = null; errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); e.model_dir = try gpa.dupe(u8, model_dir); errdefer gpa.free(e.model_dir); e.store = try Store.open(gpa, io, e.model_dir); errdefer e.store.deinit(); + try e.store.detectNorms(); const weight = try e.store.tensor(embed.weight_name); if (weight.dtype == .bf16) { e.table = try embed.Table.uploadDense(driver, weight); @@ -178,8 +186,9 @@ pub const Engine = struct { defer pack_mod.unload(); const pack_fn = try pack_mod.function(weights.pack_symbol); for (0..layers_n) |i| { - e.layers[i] = try loadLayer(gpa, driver, stream, &e.store, pack_fn, i); + e.layers[i] = try load.loadLayerRank(gpa, driver, stream, &e.store, pack_fn, i, world, rank); e.n_layers += 1; + std.log.info("flashnext: rank {d}/{d} loaded layer {d}/{d}", .{ rank + 1, world, i + 1, layers_n }); } e.store.release(); e.store.closeFile(); @@ -286,7 +295,7 @@ pub const Engine = struct { .tri = tri, .k = &e.batch_k, .s = s, - .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .collective = e.collective }, .seq = seq.view(scores), .hooks = hooks, }; @@ -324,7 +333,7 @@ pub const Engine = struct { .tri = tri, .k = &e.batch_k, .s = s, - .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram }, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = e.layers[0..e.n_layers], .mixer = e.mixer, .head = e.head_q, .ngram = &e.ngram, .collective = e.collective }, .seq = rows[0].seq, .lanes = rows, .hooks = .{ .logits = .all, .logits_out = out }, @@ -341,6 +350,7 @@ pub const Engine = struct { /// One stream's caches. `fed` tokens have already gone through the model. pub const Seq = struct { + rank_id: u64 = 0, conv: []cuda.DeviceBuffer, state: []cuda.DeviceBuffer, kc: []cuda.DeviceBuffer, diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index c2b9dde3ca..94df19ae5c 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -305,13 +305,13 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S const parts = [_][]const u8{ ".q_proj", ".k_proj", ".v_proj", ".indexer.index_qk_proj" }; proj_q = try projections.loadStitched(gpa, driver, store, ap, &parts, proj_n, dims); out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.o_proj", .{ap}), dims, out_k); - const qn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); + const qn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.q_norm.weight", .{ap}))); defer gpa.free(qn); - const kn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); + const kn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.k_norm.weight", .{ap}))); defer gpa.free(kn); - const iqn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); + const iqn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.q_layernorm.weight", .{ap}))); defer gpa.free(iqn); - const ikn = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.k_layernorm.weight", .{ap}))); + const ikn = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.indexer.k_layernorm.weight", .{ap}))); defer gpa.free(ikn); if (qn.len != head_dim or kn.len != head_dim or iqn.len != index_dim or ikn.len != index_dim) return error.UnexpectedTensor; qs_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(qn)); @@ -326,11 +326,11 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S const pp = pre_buf[0..pre.len]; key_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.key_proj", .{pp}), wide, dims); val_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.value_proj", .{pp}), dims, dims); - const nk = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); + const nk = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); defer gpa.free(nk); - const nq = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); + const nq = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); defer gpa.free(nq); - const nc = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_conv.weight", .{pp}))); + const nc = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_conv.weight", .{pp}))); defer gpa.free(nc); if (nk.len != wide or nq.len != wide or nc.len != wide) return error.UnexpectedTensor; nk_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nk)); diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index 5828a6ab8a..6e54d694ed 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -86,6 +86,9 @@ pub const Store = struct { file: ?Shard = null, held: [12]Held = undefined, nheld: usize = 0, + released: [2]Held = undefined, + nreleased: usize = 0, + norm_offset: f32 = 0, pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); @@ -95,18 +98,63 @@ pub const Store = struct { return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}) }; } + /// Infer one checkpoint-wide convention from every decoder HC norm, never from the final mixer's outlier. + pub fn detectNorms(self: *Store) !void { + var means: [layers_n]f64 = undefined; + var buf: [160]u8 = undefined; + for (0..layers_n) |i| { + const name = try std.fmt.bufPrint(&buf, "language_model.model.layers.{d}.attn_hyper_connection.hc_norm.weight", .{i}); + const values = try asF32(self.gpa, try self.tensor(name)); + defer self.gpa.free(values); + if (values.len != wide) return error.UnexpectedTensor; + var sum: f64 = 0; + for (values) |v| { + if (!std.math.isFinite(v)) return error.InvalidNormConvention; + sum += v; + } + means[i] = sum / @as(f64, @floatFromInt(values.len)); + } + self.norm_offset = try @import("cuda_norms.zig").offset(layers_n, means); + } + + pub fn normScale(self: *const Store, gpa: std.mem.Allocator, t: Tensor) ![]f32 { + const values = try asF32(gpa, t); + errdefer gpa.free(values); + for (values) |*v| { + if (!std.math.isFinite(v.*)) return error.InvalidNormConvention; + v.* += self.norm_offset; + } + return values; + } + pub fn deinit(self: *Store) void { self.release(); + self.closeReleased(); self.closeFile(); self.parsed.deinit(); } + /// Borrowers expire here; two indexed shards cover separate weight and calibration planes. pub fn release(self: *Store) void { - for (self.held[0..self.nheld]) |*h| { + for (self.held[0..self.nheld]) |h| { + if (self.nreleased == self.released.len) { + self.released[0].file.close(); + self.gpa.free(self.released[0].rel); + self.released[0] = self.released[1]; + self.nreleased -= 1; + } + self.released[self.nreleased] = h; + self.nreleased += 1; + } + self.nheld = 0; + } + + fn closeReleased(self: *Store) void { + for (self.released[0..self.nreleased]) |*h| { h.file.close(); self.gpa.free(h.rel); } - self.nheld = 0; + self.nreleased = 0; } pub fn hold(self: *Store, requested: []const u8) !Tensor { @@ -117,6 +165,15 @@ pub const Store = struct { if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; } if (self.nheld == self.held.len) return error.UnexpectedTensor; + for (self.released[0..self.nreleased], 0..) |cached, i| { + if (std.mem.eql(u8, cached.rel, file)) { + self.held[self.nheld] = cached; + self.nreleased -= 1; + if (i < self.nreleased) self.released[i] = self.released[self.nreleased]; + self.nheld += 1; + return self.held[self.nheld - 1].file.get(key) orelse return error.MissingTensor; + } + } const rel = try self.gpa.dupe(u8, file); errdefer self.gpa.free(rel); const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); @@ -361,7 +418,7 @@ pub fn loadHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, prefi defer gpa.free(up.s); defer gpa.free(up.b); const gamma = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix})); - const scale = try asF32(gpa, gamma); + const scale = try store.normScale(gpa, gamma); defer gpa.free(scale); if (scale.len != wide) return error.UnexpectedTensor; var n: usize = low; @@ -421,7 +478,7 @@ fn loadDenseHc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *Store, pref if (!up.is(.bf16, &.{ wide, low })) return error.UnexpectedTensor; var uq = try projection.dense(gpa, driver, up); errdefer uq.free(); - const gamma = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix}))); + const gamma = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.hc_norm.weight", .{prefix}))); defer gpa.free(gamma); if (gamma.len != wide) return error.UnexpectedTensor; const scale = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(gamma)); @@ -474,3 +531,69 @@ test "canonical checkpoint aliases preserve exact stored namespace and prefer a try std.testing.expect(hostOnly("model.visual.weight")); try std.testing.expect(hostOnly("mtp.layers.0.weight")); } + +test "released expert planes reuse two indexed shards without invalidating held peers" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + const header = "{\"x\":{\"dtype\":\"U8\",\"shape\":[1],\"data_offsets\":[0,1]}}"; + var data: [8 + header.len + 1]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + data[data.len - 1] = 17; + try tmp.dir.writeFile(io, .{ .sub_path = "a", .data = &data }); + data[data.len - 1] = 29; + try tmp.dir.writeFile(io, .{ .sub_path = "b", .data = &data }); + try tmp.dir.writeFile(io, .{ .sub_path = "c", .data = &data }); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{\"x\":\"a\"}}" }); + var store = try Store.open(a, io, dir); + defer store.deinit(); + const first = try store.hold("x"); + const address = first.bytes.ptr; + store.release(); + try std.testing.expectEqual(@as(usize, 0), store.nheld); + try std.testing.expectEqual(@as(usize, 1), store.nreleased); + // Removing the path proves the next borrow uses the existing mapping and parsed header. + try tmp.dir.deleteFile(io, "a"); + const again = try store.hold("x"); + try std.testing.expectEqual(address, again.bytes.ptr); + try std.testing.expectEqual(@as(u8, 17), again.bytes[0]); + const wm = &store.parsed.value.object.getPtr("weight_map").?.object; + wm.getPtr("x").?.* = .{ .string = "b" }; + const other = try store.hold("x"); + try std.testing.expectEqual(@as(u8, 29), other.bytes[0]); + try std.testing.expectEqual(@as(u8, 17), again.bytes[0]); + store.release(); + try std.testing.expectEqual(@as(usize, 2), store.nreleased); + try tmp.dir.deleteFile(io, "b"); + wm.getPtr("x").?.* = .{ .string = "a" }; + try std.testing.expectEqual(@as(u8, 17), (try store.hold("x")).bytes[0]); + wm.getPtr("x").?.* = .{ .string = "b" }; + try std.testing.expectEqual(@as(u8, 29), (try store.hold("x")).bytes[0]); + store.release(); + store.release(); + try std.testing.expectEqual(@as(usize, 2), store.nreleased); + wm.getPtr("x").?.* = .{ .string = "c" }; + _ = try store.hold("x"); + store.release(); + try std.testing.expectEqual(@as(usize, 2), store.nreleased); + for (store.released[0..store.nreleased]) |h| try std.testing.expect(!std.mem.eql(u8, "a", h.rel)); +} + +test "centered scales add one in FP32 while direct scales preserve the original bits" { + _ = @import("cuda_norms.zig"); + const raw = [_]f32{ -0.25, 0.0001, 2.75 }; + const t: Tensor = .{ .dtype = .f32, .rank = 1, .shape = .{ 3, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&raw) }; + var store: Store = undefined; + store.norm_offset = 1; + const scales = try store.normScale(std.testing.allocator, t); + defer std.testing.allocator.free(scales); + for (raw, scales) |v, scale| try std.testing.expectEqual(v + @as(f32, 1), scale); + store.norm_offset = 0; + const direct = try store.normScale(std.testing.allocator, t); + defer std.testing.allocator.free(direct); + try std.testing.expectEqualSlices(f32, &raw, direct); +} diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig new file mode 100644 index 0000000000..ee00b4f793 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -0,0 +1,377 @@ +//! Leader-owned operations replayed on the other CUDA rank; sequence and cache identities survive map relocation. +const std = @import("std"); +const cuda = @import("cuda"); +const prompt = @import("cuda_prompt.zig"); +const snaps = @import("cuda_snap.zig"); +const wire = @import("cuda_rank_wire.zig"); +const Collective = @import("cuda_collective.zig").Collective; +const wt = @import("cuda_prompt_wt.zig"); + +pub const Identity = extern struct { + config: [32]u8, + index: [32]u8, + program: [32]u8, + checkpoint: [32]u8, + context: u64, + chunk: u64, +}; + +fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 28)); + defer a.free(bytes); + var out: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(bytes, &out, .{}); + return out; +} + +pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usize, chunk: usize) !Identity { + const digest = std.mem.span(std.c.getenv("TENSORFOLD_CHECKPOINT_SHA256") orelse return error.NoCheckpointIdentity); + if (digest.len != 64) return error.NoCheckpointIdentity; + var checkpoint: [32]u8 = undefined; + _ = std.fmt.hexToBytes(&checkpoint, digest) catch return error.NoCheckpointIdentity; + const config = try std.fs.path.join(a, &.{ dir, "config.json" }); + defer a.free(config); + const index = try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" }); + defer a.free(index); + const program = try std.process.executablePathAlloc(io, a); + defer a.free(program); + return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk }; +} + +pub const Session = struct { + gpa: std.mem.Allocator, + rank: usize, + link: wire.Link, + collective: Collective, + engine: *prompt.Engine = undefined, + next: u64 = 1, + next_object: u64 = 1, + seqs: std.AutoHashMapUnmanaged(u64, *prompt.Seq) = .empty, + saved: std.AutoHashMapUnmanaged(u64, *snaps.Snap) = .empty, + stopped: bool = false, + + pub fn init(a: std.mem.Allocator, io: std.Io, host: []const u8, port: u16, rank: usize, mine: Identity) !*Session { + const s = try a.create(Session); + errdefer a.destroy(s); + const link = try wire.Link.open(io, host, port, rank, 60_000); + errdefer link.close(); + const until = wire.deadline(60_000); + var peer: Identity = undefined; + const hello: wire.Header = .{ .kind = .hello, .arg = rank, .bytes = @sizeOf(Identity) }; + if (rank == 0) { + try link.sendHeader(hello, until); + try link.send(std.mem.asBytes(&mine), until); + const h = try link.receiveHeader(until); + if (h.kind != .hello or h.sequence != 0 or h.arg != 1 or h.bytes != @sizeOf(Identity) or h.flags != 0 or h.object != 0) return error.RankIdentity; + try link.receive(std.mem.asBytes(&peer), until); + } else { + const h = try link.receiveHeader(until); + if (h.kind != .hello or h.sequence != 0 or h.arg != 0 or h.bytes != @sizeOf(Identity) or h.flags != 0 or h.object != 0) return error.RankIdentity; + try link.receive(std.mem.asBytes(&peer), until); + try link.sendHeader(hello, until); + try link.send(std.mem.asBytes(&mine), until); + } + if (!sameIdentity(mine, peer)) return error.RankIdentity; + s.* = .{ .gpa = a, .rank = rank, .link = link, .collective = try Collective.init(link, rank, 30_000) }; + return s; + } + + pub fn poison(s: *Session) void { + s.stopped = true; + s.collective.abort(); + _ = std.c.shutdown(s.link.fd, std.posix.SHUT.RDWR); + } + + pub fn deinit(s: *Session) void { + if (s.rank == 0 and !s.stopped) { + const h = s.command(.{ .kind = .stop }, &.{}) catch null; + if (h) |sent| s.ack(sent) catch s.poison(); + } + s.collective.deinit(); + var qi = s.seqs.valueIterator(); + while (qi.next()) |q| { + q.*.deinit(s.gpa); + s.gpa.destroy(q.*); + } + s.seqs.deinit(s.gpa); + var si = s.saved.valueIterator(); + while (si.next()) |saved| snaps.drop(s.gpa, s.engine.stream.*, saved.*); + s.saved.deinit(s.gpa); + s.link.close(); + s.gpa.destroy(s); + } + + /// Every rank's host budgets the same number of streams and cache bytes, using the least local room. + pub fn room(s: *Session, local: u64) !u64 { + const until = wire.deadline(600_000); + errdefer s.poison(); + const mine: wire.Header = .{ .kind = .ready, .arg = local }; + const peer = if (s.rank == 0) peer: { + try s.link.sendHeader(mine, until); + break :peer try s.link.receiveHeader(until); + } else peer: { + const h = try s.link.receiveHeader(until); + try s.link.sendHeader(mine, until); + break :peer h; + }; + if (peer.kind != .ready or peer.sequence != 0 or peer.object != 0 or peer.bytes != 0 or peer.flags != 0) return error.BadRankFrame; + return @min(local, peer.arg); + } + + fn command(s: *Session, input: wire.Header, payload: []const u8) !wire.Header { + if (s.rank != 0 or s.stopped or s.collective.failed) return error.RankDisconnected; + if (payload.len > wire.max_payload or input.bytes != payload.len) return error.BadRankFrame; + errdefer s.poison(); + var h = input; + h.sequence = s.next; + s.next = std.math.add(u64, s.next, 1) catch return error.BadRankFrame; + const until = wire.deadline(30_000); + try s.link.sendHeader(h, until); + try s.link.send(payload, until); + return h; + } + + fn ack(s: *Session, sent: wire.Header) !void { + errdefer s.poison(); + const h = try s.link.receiveHeader(wire.deadline(30_000)); + if (h.kind != .ack or h.sequence != sent.sequence or h.object != sent.object or h.bytes != 0 or h.arg != 0) return error.BadRankFrame; + if (h.flags == 1) return error.OutOfDeviceMemory; + if (h.flags != 0) return error.RemoteRank; + } + + pub fn create(s: *Session, cap: usize) !prompt.Seq { + var q = try s.engine.newSeq(cap); + errdefer q.deinit(s.gpa); + const id = s.next_object; + s.next_object = std.math.add(u64, id, 1) catch return error.BadRankFrame; + const h = try s.command(.{ .kind = .create, .object = id, .arg = cap }, &.{}); + try s.ack(h); + q.rank_id = id; + return q; + } + + pub fn drop(s: *Session, q: *prompt.Seq) void { + if (!s.stopped) { + const h = s.command(.{ .kind = .drop_seq, .object = q.rank_id }, &.{}) catch { + s.poison(); + q.deinit(s.gpa); + return; + }; + s.ack(h) catch s.poison(); + } + q.deinit(s.gpa); + } + + pub fn forward(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, ids: []const u32, hooks: prompt.batch.Hooks) !void { + if (ids.len == 0 or ids.len > wire.max_payload / 4 or q.fed + ids.len > q.cap or hooks.keep != null or hooks.streams != 0 or hooks.logits == .all) return error.UnsupportedRankOperation; + try tokenIds(ids); + const h = try s.command(.{ .kind = .forward, .object = q.rank_id, .bytes = @intCast(ids.len * 4), .flags = @backingInt(hooks.logits) }, std.mem.sliceAsBytes(ids)); + errdefer s.poison(); + try s.engine.forward(Tri, tri, q, ids, hooks); + try s.ack(h); + } + + pub fn round(s: *Session, comptime Tri: type, tri: Tri, seqs: []const *prompt.Seq, ids: []const u32, out: u64) !void { + if (seqs.len == 0 or seqs.len != ids.len or seqs.len > 32) return error.UnsupportedRankOperation; + try tokenIds(ids); + var raw: [32 * 12]u8 = undefined; + for (seqs, ids, 0..) |q, token, i| { + if (q.rank_id == 0 or q.fed >= q.cap) return error.UnexpectedTensor; + for (seqs[0..i]) |prior| if (prior.rank_id == q.rank_id) return error.DuplicateRankSequence; + std.mem.writeInt(u64, raw[i * 12 ..][0..8], q.rank_id, .little); + std.mem.writeInt(u32, raw[i * 12 + 8 ..][0..4], token, .little); + } + const bytes = raw[0 .. seqs.len * 12]; + const h = try s.command(.{ .kind = .round, .bytes = @intCast(bytes.len) }, bytes); + errdefer s.poison(); + try s.engine.round(Tri, tri, seqs, ids, out); + try s.ack(h); + } + + pub fn save(s: *Session, q: *prompt.Seq) !*snaps.Snap { + const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); + errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); + kept.rank_id = s.next_object; + s.next_object = std.math.add(u64, s.next_object, 1) catch return error.BadRankFrame; + const h = try s.command(.{ .kind = .save, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + try s.ack(h); + return kept; + } + + pub fn restore(s: *Session, q: *prompt.Seq, kept: *const snaps.Snap) !void { + if (kept.at > q.cap) return error.SequenceTooShort; + const h = try s.command(.{ .kind = .restore, .object = q.rank_id, .arg = kept.rank_id }, &.{}); + errdefer s.poison(); + try snaps.restore(s.engine.stream.*, q, null, kept); + try s.ack(h); + } + + pub fn dropSaved(s: *Session, kept: *snaps.Snap) void { + if (!s.stopped) { + const h = s.command(.{ .kind = .drop_snap, .object = kept.rank_id }, &.{}) catch { + s.poison(); + snaps.drop(s.gpa, s.engine.stream.*, kept); + return; + }; + s.ack(h) catch s.poison(); + } + snaps.drop(s.gpa, s.engine.stream.*, kept); + } + + pub fn follow(s: *Session, comptime Tri: type, tri: Tri, logits: u64) !void { + if (s.rank != 1) return error.BadRank; + errdefer s.poison(); + const payload = try s.gpa.alloc(u8, wire.max_payload); + defer s.gpa.free(payload); + while (true) { + const h = try s.link.receiveHeader(null); + if (h.sequence != s.next) return error.BadRankFrame; + s.next = std.math.add(u64, s.next, 1) catch return error.BadRankFrame; + try s.link.receive(payload[0..h.bytes], wire.deadline(30_000)); + const stop = s.apply(Tri, tri, logits, h, payload[0..h.bytes]) catch |err| { + s.link.sendHeader(.{ .kind = .ack, .sequence = h.sequence, .object = h.object, .flags = 2 }, wire.deadline(30_000)) catch {}; + return err; + }; + try s.link.sendHeader(.{ .kind = .ack, .sequence = h.sequence, .object = h.object }, wire.deadline(30_000)); + if (stop) { + s.stopped = true; + return; + } + } + } + + fn apply(s: *Session, comptime Tri: type, tri: Tri, logits: u64, h: wire.Header, payload: []const u8) !bool { + try shape(h); + switch (h.kind) { + .create => { + if (s.seqs.contains(h.object)) return error.BadRankFrame; + const q = try s.gpa.create(prompt.Seq); + errdefer s.gpa.destroy(q); + q.* = try s.engine.newSeq(@intCast(h.arg)); + errdefer q.deinit(s.gpa); + q.rank_id = h.object; + try s.seqs.put(s.gpa, h.object, q); + }, + .forward => { + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const ids = try s.gpa.alloc(u32, payload.len / 4); + defer s.gpa.free(ids); + for (ids, 0..) |*v, i| v.* = std.mem.readInt(u32, payload[i * 4 ..][0..4], .little); + try tokenIds(ids); + try s.engine.forward(Tri, tri, q, ids, .{ .logits = @fromBackingInt(@intCast(h.flags)) }); + }, + .round => { + var seqs: [32]*prompt.Seq = undefined; + var ids: [32]u32 = undefined; + const n = payload.len / 12; + try roundPayload(payload); + for (0..n) |i| { + seqs[i] = s.seqs.get(std.mem.readInt(u64, payload[i * 12 ..][0..8], .little)) orelse return error.UnknownRankSequence; + ids[i] = std.mem.readInt(u32, payload[i * 12 + 8 ..][0..4], .little); + if (seqs[i].fed >= seqs[i].cap) return error.UnexpectedTensor; + } + try s.engine.round(Tri, tri, seqs[0..n], ids[0..n], logits); + }, + .drop_seq => { + const entry = s.seqs.fetchRemove(h.object) orelse return error.UnknownRankSequence; + entry.value.deinit(s.gpa); + s.gpa.destroy(entry.value); + }, + .save => { + if (s.saved.contains(h.arg)) return error.BadRankFrame; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); + errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); + try s.saved.put(s.gpa, h.arg, kept); + }, + .restore => { + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const kept = s.saved.get(h.arg) orelse return error.UnknownRankSnapshot; + try snaps.restore(s.engine.stream.*, q, null, kept); + }, + .drop_snap => { + const entry = s.saved.fetchRemove(h.object) orelse return error.UnknownRankSnapshot; + snaps.drop(s.gpa, s.engine.stream.*, entry.value); + }, + .stop => return true, + else => return error.BadRankFrame, + } + return false; + } +}; + +fn tokenIds(ids: []const u32) !void { + for (ids) |id| if (id >= wt.vocab) return error.InvalidToken; +} + +fn roundPayload(payload: []const u8) !void { + if (payload.len == 0 or payload.len % 12 != 0 or payload.len / 12 > 32) return error.BadRankFrame; + var seen: [32]u64 = undefined; + for (0..payload.len / 12) |i| { + const id = std.mem.readInt(u64, payload[i * 12 ..][0..8], .little); + const token = std.mem.readInt(u32, payload[i * 12 + 8 ..][0..4], .little); + if (id == 0) return error.BadRankFrame; + for (seen[0..i]) |prior| if (prior == id) return error.DuplicateRankSequence; + try tokenIds(&.{token}); + seen[i] = id; + } +} + +fn sameIdentity(a: Identity, b: Identity) bool { + return std.mem.eql(u8, std.mem.asBytes(&a), std.mem.asBytes(&b)); +} + +/// Refuse malformed operations before looking up an owner or touching the GPU. +pub fn shape(h: wire.Header) !void { + switch (h.kind) { + .create => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .forward => if (h.object == 0 or h.arg != 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame, + .round => if (h.object != 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 32) return error.BadRankFrame, + .save, .restore => if (h.object == 0 or h.arg == 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .drop_seq, .drop_snap => if (h.object == 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + .stop => if (h.object != 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, + else => return error.BadRankFrame, + } +} + +test "rank identity binds checkpoint, config, index, program and execution settings" { + const base: Identity = .{ .config = @splat(1), .index = @splat(2), .program = @splat(3), .checkpoint = @splat(4), .context = 4096, .chunk = 64 }; + try std.testing.expect(sameIdentity(base, base)); + inline for (.{ "config", "index", "program", "checkpoint" }) |name| { + var other = base; + @field(other, name)[0] ^= 1; + try std.testing.expect(!sameIdentity(base, other)); + } + var other = base; + other.context += 1; + try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.chunk += 1; + try std.testing.expect(!sameIdentity(base, other)); +} + +test "rank operations refuse invalid geometry and unsupported hooks before execution" { + try shape(.{ .kind = .create, .object = 1, .arg = 4096 }); + try shape(.{ .kind = .forward, .object = 1, .bytes = 4 }); + try shape(.{ .kind = .round, .bytes = 24 }); + for ([_]wire.Header{ + .{ .kind = .hello }, .{ .kind = .create, .arg = 4096 }, .{ .kind = .create, .object = 1, .arg = prompt.max_tokens + 1 }, + .{ .kind = .forward, .object = 1, .bytes = 3 }, .{ .kind = .forward, .object = 1, .bytes = 4, .flags = 1 }, .{ .kind = .round, .bytes = 33 * 12 }, + .{ .kind = .round, .bytes = 13 }, .{ .kind = .stop, .object = 1 }, .{ .kind = .save, .object = 1 }, + .{ .kind = .drop_seq, .object = 1, .bytes = 4 }, + }) |h| try std.testing.expectError(error.BadRankFrame, shape(h)); +} + +test "round payload rejects duplicate owners and out-of-vocabulary tokens before GPU execution" { + var bytes: [24]u8 = @splat(0); + std.mem.writeInt(u64, bytes[0..8], 9, .little); + std.mem.writeInt(u64, bytes[12..20], 10, .little); + try roundPayload(&bytes); + std.mem.writeInt(u64, bytes[12..20], 9, .little); + try std.testing.expectError(error.DuplicateRankSequence, roundPayload(&bytes)); + std.mem.writeInt(u64, bytes[12..20], 10, .little); + std.mem.writeInt(u32, bytes[20..24], wt.vocab, .little); + try std.testing.expectError(error.InvalidToken, roundPayload(&bytes)); + try std.testing.expectError(error.InvalidToken, tokenIds(&.{ 0, wt.vocab })); + try std.testing.expectError(error.BadRankFrame, roundPayload(bytes[0..13])); +} diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index f3f4a72cdc..0e551809b3 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -24,6 +24,7 @@ pub fn bytes(at: usize, head: bool) usize { /// One kept state; `head`: it holds the MTP head's cache over the same rows, which a drafting stream needs. pub const Snap = struct { + rank_id: u64 = 0, buf: cuda.DeviceBuffer, at: usize, history: [2]i64, diff --git a/zig/tests/cuda/flashnext_model_rank_tests.zig b/zig/tests/cuda/flashnext_model_rank_tests.zig new file mode 100644 index 0000000000..ebb937ae71 --- /dev/null +++ b/zig/tests/cuda/flashnext_model_rank_tests.zig @@ -0,0 +1,77 @@ +//! Full-checkpoint experiment: two-rank prompt, snapshot replay and concurrent decode agree byte for byte. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const prompt = flash.prompt; +const ranks = flash.rank_engine; + +pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: []const u8) !void { + const context = 256; + const chunk = 32; + const source = try std.Io.Dir.cwd().readFileAlloc(gpu.io, ids_path, gpu.gpa, .limited(1 << 20)); + defer gpu.gpa.free(source); + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpu.gpa); + var pieces = std.mem.tokenizeAny(u8, source, " \r\n\t"); + while (pieces.next()) |word| try ids.append(gpu.gpa, try std.fmt.parseInt(u32, word, 10)); + if (ids.items.len == 0 or ids.items.len > context / 2) return error.BadPrompt; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + const identity = try ranks.identity(gpu.gpa, gpu.io, dir, context, chunk); + const session = try ranks.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const engine = prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir, 2, rank) catch |err| { + session.poison(); + session.deinit(); + return err; + }; + session.engine = engine; + defer { + session.deinit(); + engine.deinit(); + } + engine.collective = &session.collective; + if (engine.chunk != chunk) return error.RankIdentity; + const room = try session.room((try gpu.ctx.memInfo()).free); + try check.expect(room > prompt.seqBytes(context) * 3, "both ranks have room for experimental caches", .{}); + var logits = try cuda.DeviceBuffer.alloc(gpu.d, 2 * prompt.vocabulary * 2); + defer logits.free(); + if (rank == 1) { + try session.follow(flash.Tri, tri, logits.ptr); + check.pass("rank 1 replay completed", .{}); + return; + } + var first = try session.create(context); + defer session.drop(&first); + try session.forward(flash.Tri, tri, &first, ids.items, .{}); + check.pass("two-rank prompt completed ({d} tokens)", .{ids.items.len}); + for (0..prompt.vocabulary) |i| { + const value: f32 = @bitCast(@as(u32, std.mem.readInt(u16, first.logits[i * 2 ..][0..2], .little)) << 16); + try check.expect(std.math.isFinite(value), "finite initial logits at {d}", .{i}); + } + const kept = try session.save(&first); + defer session.dropSaved(kept); + const token = prompt.argmax(first.logits); + try session.forward(flash.Tri, tri, &first, &.{token}, .{}); + const expected = try gpu.gpa.dupe(u8, first.logits); + defer gpu.gpa.free(expected); + try session.restore(&first, kept); + try session.forward(flash.Tri, tri, &first, &.{token}, .{}); + try check.sameBytes("two-rank snapshot restore and replay", first.logits, expected); + try session.restore(&first, kept); + var second = try session.create(context); + defer session.drop(&second); + try session.forward(flash.Tri, tri, &second, ids.items, .{}); + try session.round(flash.Tri, tri, &.{ &first, &second }, &.{ token, token }, logits.ptr); + try check.sameBytes("two-rank concurrent first stream", first.logits, expected); + try check.sameBytes("two-rank concurrent second stream", second.logits, expected); + for (0..32) |_| { + const next = prompt.argmax(first.logits); + std.debug.print("TOKEN {d}\n", .{next}); + try session.round(flash.Tri, tri, &.{&first}, &.{next}, logits.ptr); + } + check.pass("two-rank full-model prompt, snapshot replay, concurrent byte equality and 32 decode rows", .{}); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 92ab44b213..a3e6bf56c2 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -33,6 +33,7 @@ const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); +const flashnext_model_rank_tests = @import("flashnext_model_rank_tests.zig"); const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); @@ -62,6 +63,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) @@ -155,6 +157,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum"); if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); From aa4015ee36699e53dbfcb8d9d8d959c5ffddd2a8 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:08:48 -0700 Subject: [PATCH 070/115] Serve mixed-format Flash Next ranks with bounded admission and health --- zig/build/cuda.zig | 2 +- zig/src/core/engine_api.zig | 7 ++ zig/src/core/lane_host.zig | 3 +- zig/src/core/lane_host_status.zig | 1 + zig/src/core/lane_host_test.zig | 36 ++++++ zig/src/families/flashnext/cuda_config.zig | 104 +++++++++++++++++- zig/src/families/flashnext/cuda_lanes.zig | 76 +++++++++++-- zig/src/families/flashnext/cuda_launch.zig | 1 + .../families/flashnext/cuda_rank_engine.zig | 9 ++ .../families/flashnext/cuda_rank_memory.zig | 78 +++++++++++++ zig/src/native/cuda.zig | 40 ++++--- zig/src/native/cuda_family_memory.zig | 41 +++++++ zig/src/server/chat.zig | 2 +- zig/src/server/openai.zig | 4 +- zig/src/server/status_routes.zig | 5 +- zig/src/server/stream_preflight_test.zig | 27 +++++ 16 files changed, 397 insertions(+), 39 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_rank_memory.zig create mode 100644 zig/src/native/cuda_family_memory.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 0804c11093..e359895847 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -258,7 +258,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, mods.heat, mods.tokenizer, native, stagger(b, host, .debug), vision_build.module(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.flashnext, mods.heat, mods.tokenizer, native_modules.api, native, stagger(b, host, .debug), vision_build.module(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index e9e7d0a540..4ee72ac0a8 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -263,7 +263,14 @@ pub const MemorySource = struct { read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, }; +/// Thread-safe backend readiness; no GPU context or blocking device work is allowed on HTTP threads. +pub const HealthSource = struct { + ctx: *anyopaque, + read: *const fn (*anyopaque) bool, +}; + pub const Status = struct { + healthy: bool = true, /// Requests in prefill or decode, and those waiting for a lane. running: u32 = 0, waiting: u32 = 0, diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index e53d328c2e..ebc26dab81 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -43,6 +43,7 @@ pub const LaneHost = struct { lone_job: ?*Job = null, // the job that driver holds now cache: ?*pc.Store = null, // kept prompt states (engine thread only); the backend restores and saves them decoded_rows: bool = false, // the backend prefills Request.decode_spans with decoded rows' arithmetic, cache or not + health: ?api.HealthSource = null, memory: ?api.MemorySource = null, // the backend's memory counts; null: Engine.memory reports none explain: ?api.Explain = null, // the backend's words for a request it refuses; null: the error's name learner: ?api.Learner = null, // the family's Sliding Weights learner; null: learn requests are refused @@ -143,7 +144,7 @@ pub const LaneHost = struct { job.* = .{ .host = h, .id = id, .request = request, .sink = sink }; h.mutex.lockUncancelable(h.io); defer h.mutex.unlock(h.io); - if (h.closing) { + if (h.closing or (if (h.health) |source| !source.read(source.ctx) else false)) { h.gpa.destroy(job); return error.Closed; } diff --git a/zig/src/core/lane_host_status.zig b/zig/src/core/lane_host_status.zig index a92bf19a0c..9932fe8271 100644 --- a/zig/src/core/lane_host_status.zig +++ b/zig/src/core/lane_host_status.zig @@ -20,6 +20,7 @@ pub fn read(ctx: *anyopaque, out: *api.Status, stream_tokens: []u32) void { const n = @min(stream_tokens.len, h.live_tokens.items.len); @memcpy(stream_tokens[0..n], h.live_tokens.items[0..n]); out.* = .{ + .healthy = if (h.health) |source| source.read(source.ctx) else true, .running = @intCast(h.admitted.items.len), .waiting = @intCast(h.queued.items.len), .decode_tokens_per_second = @as(f64, @floatFromInt(tokens)) / 2.0, diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index 20a91348f1..b5460c9a59 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -562,3 +562,39 @@ test { _ = @import("lane_host_lone_test.zig"); _ = @import("lane_host_admit_test.zig"); } + +test "an unhealthy backend fails health and refuses a request before queue admission" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{}, 1, 0); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + const Health = struct { + available: bool = true, + fn read(ctx: *anyopaque) bool { + return @as(*@This(), @ptrCast(@alignCast(ctx))).available; + } + fn event(_: *anyopaque, _: Id, _: *const Event) void { + @panic("a refused request must never reach its sink"); + } + }; + var health: Health = .{}; + const engine = host.engine(); + var status: api.Status = .{}; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + host.health = .{ .ctx = &health, .read = Health.read }; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + health.available = false; + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); + const request: Request = .{ .prompt = &.{1}, .max_tokens = 1 }; + try std.testing.expectError(error.Closed, engine.submit(1, &request, .{ .ctx = &health, .event = Health.event })); + try std.testing.expectEqual(@as(usize, 0), host.queued.items.len); + try std.testing.expectEqual(@as(usize, 0), host.admitted.items.len); +} diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig index 8fc2652e67..e14cfab98f 100644 --- a/zig/src/families/flashnext/cuda_config.zig +++ b/zig/src/families/flashnext/cuda_config.zig @@ -97,10 +97,46 @@ pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u if (int(q, "bits") != 4) return "quantization.bits"; if (int(q, "group_size") != 32) return "quantization.group_size"; if (q.get("mode")) |m| if (m != .string or !std.mem.eql(u8, m.string, "affine")) return "quantization.mode"; + return geometry(text); +} + +fn geometry(text: std.json.ObjectMap) ?[]const u8 { for (shapes) |w| if (int(text, w.name) != w.value) return w.name; return semantics(text); } +pub const Mode = enum { affine, modelopt }; + +fn projection(q: std.json.ObjectMap, name: []const u8, algo: []const u8, group: i64) bool { + const p = object(q, name) orelse return false; + return holds(p.get("quant_algo") orelse return false, .{ .word = algo }) and int(p, "group_size") == group; +} + +fn modeloptMismatch(text: std.json.ObjectMap, quant: std.json.ObjectMap) ?[]const u8 { + if (!holds(quant.get("quant_method") orelse return "quantization_config.quant_method", .{ .word = "modelopt" })) return "quantization_config.quant_method"; + if (!holds(quant.get("quant_algo") orelse return "quantization_config.quant_algo", .{ .word = "MIXED_PRECISION" })) return "quantization_config.quant_algo"; + if (given(quant, "kv_cache_quant_algo") != null) return "quantization_config.kv_cache_quant_algo"; + if (geometry(text)) |field| return field; + const entries = object(quant, "quantized_layers") orelse return "quantization_config.quantized_layers"; + var buf: [256]u8 = undefined; + for (0..wt.layers_n) |layer| { + const routed = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts", .{layer}) catch unreachable; + if (!projection(entries, routed, "NVFP4", 16)) return "quantization_config.routed_experts"; + const p = object(entries, routed).?; + if (given(p, "weight_scale_encoding")) |encoding| if (!holds(encoding, .{ .word = "csf" })) return "quantization_config.weight_scale_encoding"; + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |part| { + const name = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.shared_expert.{s}", .{ layer, part }) catch unreachable; + if (!projection(entries, name, "MXFP8", 32)) return "quantization_config.shared_expert"; + } + const parts: []const []const u8 = if (wt.linearLayer(layer)) &.{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj" } else &.{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }; + for (parts) |part| { + const name = std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.{s}.{s}", .{ layer, if (wt.linearLayer(layer)) "linear_attn" else "self_attn", part }) catch unreachable; + if (!projection(entries, name, "MXFP8", 32)) return "quantization_config.attention_projection"; + } + } + return null; +} + /// The first semantic setting the kernels would ignore: eps, RoPE, gates, activations, layer roles, the head's layer. pub fn semantics(text: std.json.ObjectMap) ?[]const u8 { const eps = float(given(text, "rms_norm_eps") orelse return "rms_norm_eps") orelse return "rms_norm_eps"; @@ -142,6 +178,10 @@ fn rope(text: std.json.ObjectMap) ?[]const u8 { /// Reads `dir`/config.json; error.UnsupportedCheckpoint names the field in a log line when one differs. pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { + _ = try mode(gpa, io, dir); +} + +pub fn mode(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Mode { const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); defer gpa.free(path); const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); @@ -152,11 +192,12 @@ pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { const root = parsed.value.object; const text = if (root.get("text_config")) |v| (if (v == .object) v.object else root) else root; const quant = if (root.get("quantization")) |v| (if (v == .object) v.object else null) else null; - var field = mismatch(text, quant); + const selected: Mode = if (quant != null) .affine else .modelopt; + var field = if (selected == .affine) mismatch(text, quant) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; if (field == null) if (given(root, "tie_word_embeddings")) |v| if (!holds(v, .{ .flag = false })) { field = "tie_word_embeddings"; }; - std.log.err("flashnext: config.json's {s} differs from the MLX 4-bit group-32 Flash Next the CUDA path serves", .{field orelse return}); + std.log.err("flashnext: config.json's {s} differs from the supported Flash Next CUDA layout", .{field orelse return selected}); return error.UnsupportedCheckpoint; } @@ -295,3 +336,62 @@ test "RoPE's base and type inside rope_parameters, and the MTP layer's base, are try std.testing.expectEqualStrings(c.refused, semantics(text) orelse "accepted"); } } + +test "mixed precision requires every routed, shared and attention declaration and refuses altered formats" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var text: std.json.ObjectMap = .empty; + var rope_params: std.json.ObjectMap = .empty; + var mtp: std.json.ObjectMap = .empty; + var kinds: std.json.Array = .init(a); + var ple: std.json.Array = .init(a); + try served(a, &text, &rope_params, &mtp, &kinds, &ple); + var quant: std.json.ObjectMap = .empty; + try quant.put(a, "quant_method", .{ .string = "modelopt" }); + try quant.put(a, "quant_algo", .{ .string = "MIXED_PRECISION" }); + var entries: std.json.ObjectMap = .empty; + for (0..wt.layers_n) |layer| { + const routed = try std.fmt.allocPrint(a, "model.language_model.layers.{d}.mlp.experts", .{layer}); + var nv: std.json.ObjectMap = .empty; + try nv.put(a, "quant_algo", .{ .string = "NVFP4" }); + try nv.put(a, "group_size", .{ .integer = 16 }); + try nv.put(a, "weight_scale_encoding", .{ .string = "csf" }); + try entries.put(a, routed, .{ .object = nv }); + const prefix = if (wt.linearLayer(layer)) "linear_attn" else "self_attn"; + const parts: []const []const u8 = if (wt.linearLayer(layer)) &.{ "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj" } else &.{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }; + for (parts) |part| try addMxfp8(a, &entries, layer, prefix, part); + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |part| try addMxfp8(a, &entries, layer, "mlp.shared_expert", part); + } + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expect(modeloptMismatch(text, quant) == null); + const cases = [_]struct { name: []const u8, field: []const u8, value: std.json.Value, refused: []const u8 }{ + .{ .name = "model.language_model.layers.47.mlp.experts", .field = "group_size", .value = .{ .integer = 32 }, .refused = "quantization_config.routed_experts" }, + .{ .name = "model.language_model.layers.0.mlp.experts", .field = "weight_scale_encoding", .value = .{ .string = "other" }, .refused = "quantization_config.weight_scale_encoding" }, + .{ .name = "model.language_model.layers.17.mlp.shared_expert.down_proj", .field = "quant_algo", .value = .{ .string = "NVFP4" }, .refused = "quantization_config.shared_expert" }, + .{ .name = "model.language_model.layers.0.linear_attn.in_proj_b", .field = "group_size", .value = .{ .integer = 16 }, .refused = "quantization_config.attention_projection" }, + .{ .name = "model.language_model.layers.47.self_attn.indexer.index_qk_proj", .field = "quant_algo", .value = .{ .string = "FP8" }, .refused = "quantization_config.attention_projection" }, + }; + for (cases) |c| { + const entry = entries.getPtr(c.name).?; + const original = entry.object.get(c.field).?; + try entry.object.put(a, c.field, c.value); + try std.testing.expectEqualStrings(c.refused, modeloptMismatch(text, quant) orelse "accepted"); + try entry.object.put(a, c.field, original); + } + try quant.put(a, "kv_cache_quant_algo", .{ .string = "FP8" }); + try std.testing.expectEqualStrings("quantization_config.kv_cache_quant_algo", modeloptMismatch(text, quant).?); + try quant.put(a, "kv_cache_quant_algo", .null); + try std.testing.expect(modeloptMismatch(text, quant) == null); + _ = entries.swapRemove("model.language_model.layers.47.self_attn.o_proj"); + try quant.put(a, "quantized_layers", .{ .object = entries }); + try std.testing.expectEqualStrings("quantization_config.attention_projection", modeloptMismatch(text, quant).?); +} + +fn addMxfp8(a: std.mem.Allocator, entries: *std.json.ObjectMap, layer: usize, prefix: []const u8, part: []const u8) !void { + var p: std.json.ObjectMap = .empty; + try p.put(a, "quant_algo", .{ .string = "MXFP8" }); + try p.put(a, "group_size", .{ .integer = 32 }); + const name = try std.fmt.allocPrint(a, "model.language_model.layers.{d}.{s}.{s}", .{ layer, prefix, part }); + try entries.put(a, name, .{ .object = p }); +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index ba109285c6..0d0cb0f42f 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -7,6 +7,8 @@ const prompt = @import("cuda_prompt.zig"); const weights = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); const snap = @import("cuda_snap.zig"); +const ranks = @import("cuda_rank_engine.zig"); +const config = @import("cuda_config.zig"); const Glue = @import("cuda_glue.zig").Glue; const be = lanes.backend; @@ -22,7 +24,7 @@ pub const captured_head_entries = [_][]const u8{ "_rmsnorm", "_add_streams" }; pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; - pub const formats: []const []const u8 = &.{"mlx-q4g32"}; + pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed" }; // the model's window: past 2,048 keys the block indexer picks each row's 512 best blocks of 4 (cuda_batch.zig) pub const default_context: i64 = @intCast(prompt.max_tokens); // a prompt runs on the batched pass in chunks of this many rows, one a prefill_step: no staggered segments @@ -42,7 +44,7 @@ pub fn Namespace(comptime Tri: type) type { pub const min_prompt: u32 = 512; }; - pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551 }; /// One stream's prompt-and-reply driver, as native/cuda.zig types it; this family has none. pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; @@ -51,6 +53,7 @@ pub fn Namespace(comptime Tri: type) type { backend: be.Backend, facts: lanes.Model, rows: u32, + prompt_rows: u32, /// Device bytes each admitted stream allocates: its caches (and with the head, its cache and window). stream_bytes: usize, /// Every stream allocates its own sequence; none rides on the engine's. @@ -64,6 +67,7 @@ pub fn Namespace(comptime Tri: type) type { target: *anyopaque, /// Whether our own glue runs (no captured set given, or one without this family's entry points). own_glue: bool = true, + follow: ?*const fn (*anyopaque) anyerror!void = null, }; /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. @@ -72,12 +76,36 @@ pub fn Namespace(comptime Tri: type) type { return main + mtp.headBytes(); } + pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + const mode = try config.mode(gpa, io, dir); + if (mode == .modelopt) { + if (drafts) return error.ModelOptDraftsUnavailable; + return @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); + } + if (world != 1) return error.UnsupportedQuantization; + const main = try weights.deviceBytes(gpa, io, dir); + return main + if (drafts) mtp.headBytes() else @as(u64, 0); + } + + pub fn healthy(ptr: *anyopaque) bool { + if (of(ptr).ranks) |r| return r.healthy(); + return true; + } + + pub fn rankRoom(ptr: *anyopaque, local: u64) !u64 { + if (of(ptr).ranks) |r| return r.room(local); + return local; + } + /// A request or load this engine refuses, in words; null: none of its own. pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", - error.UnsupportedCheckpoint => "Flash Next on CUDA serves MLX affine 4-bit group-32 Qwen3.8-Flash-Next checkpoints; the log names the config field that differs", + error.UnsupportedCheckpoint => "the checkpoint differs from the supported Flash Next CUDA layout; the log names its config field", + error.ModelOptDraftsUnavailable => "ModelOpt Flash Next serves without its MTP head so far: use --no-drafts", + error.RankDisconnected, error.RankTimeout, error.RemoteRank, error.BadRankFrame => "a tensor-parallel rank failed; restart both ranks from the same checkpoint before retrying", error.LearnedPocket => "this folder holds Living Weights pockets, which Flash Next on CUDA does not serve yet", else => null, }; @@ -98,6 +126,7 @@ pub fn Namespace(comptime Tri: type) type { gpa: std.mem.Allocator, context: usize, engine: *prompt.Engine, + ranks: ?*ranks.Session = null, head: ?*mtp.Head = null, depth: usize = mtp.max_chain, confidence: f64 = mtp.default_confidence, @@ -194,7 +223,7 @@ pub fn Namespace(comptime Tri: type) type { } // a window's rows past the reply's last token (its drafts) still need cache slots const cap = @min(ids.len + s.max_new + (if (drafting) mtp.max_rows else 0), prompt.max_tokens); - gop.value_ptr.* = .{ .seq = self.engine.newSeq(cap) catch |err| { + gop.value_ptr.* = .{ .seq = (if (self.ranks) |r| r.create(cap) else self.engine.newSeq(cap)) catch |err| { self.lanes.removeByPtr(gop.key_ptr); return err; } }; @@ -209,7 +238,7 @@ pub fn Namespace(comptime Tri: type) type { const kept: *snap.Snap = @ptrCast(@alignCast(saved)); if (kept.at != s.reuse.at or kept.at == 0 or kept.at >= ids.len) { s.reuse_failed = true; - } else if (snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { + } else if (if (self.ranks) |r| r.restore(&l.seq, kept) else snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { l.filled = kept.at; s.cached = @intCast(kept.at); } else |err| switch (err) { @@ -232,6 +261,8 @@ pub fn Namespace(comptime Tri: type) type { const last = stop == ids.len; if (l.draft) |*d| { try self.ops().prefillChunk(&l.seq, d, ids, at, stop); + } else if (self.ranks) |r| { + try r.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none }); } else try self.engine.forward(Tri, self.tri(), &l.seq, ids[at..stop], .{ .logits = if (last) .last else .none }); l.filled = if (last) null else stop; if (!last and std.mem.indexOfScalar(u32, s.reuse.marks, @intCast(stop)) != null) { @@ -249,7 +280,10 @@ pub fn Namespace(comptime Tri: type) type { pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: Options) !Loaded { if (o.context > prompt.max_tokens) return error.ContextOverWindow; if (o.segments != 1 or o.slide) return error.UnsupportedOption; - try @import("cuda_config.zig").check(gpa, io, dir); + if ((o.tp != 1 and o.tp != 2) or o.rank >= o.tp) return error.BadRank; + const mode = try config.mode(gpa, io, dir); + if (mode == .modelopt and o.drafts) return error.ModelOptDraftsUnavailable; + if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); @@ -261,7 +295,20 @@ pub fn Namespace(comptime Tri: type) type { // a captured set replaces our glue, for comparison, when it holds every entry point this engine launches if (captured) |kernels| own.set = try comparison(gpa, io, driver, ctx, kernels, o.drafts); errdefer if (own.set) |*x| x.deinit(); - own.engine = try prompt.Engine.init(gpa, io, driver, &own.stream, dir); + if (o.tp == 2) { + const configured_chunk = envUsize("TF_FLASHNEXT_CHUNK", prompt.batch.default_chunk); + const id = try ranks.identity(gpa, io, dir, o.context, if (configured_chunk == 0) prompt.batch.default_chunk else configured_chunk); + own.ranks = try ranks.Session.init(gpa, io, o.master orelse "127.0.0.1", o.master_port, o.rank, id); + } + errdefer if (own.ranks) |r| { + r.poison(); + r.deinit(); + }; + own.engine = try prompt.Engine.initRank(gpa, io, driver, &own.stream, dir, o.tp, o.rank); + if (own.ranks) |r| { + r.engine = own.engine; + own.engine.collective = &r.collective; + } errdefer own.engine.deinit(); // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) _ = try own.engine.reserve(@min(o.context, own.engine.chunk)); @@ -285,11 +332,13 @@ pub fn Namespace(comptime Tri: type) type { .backend = own.backend(), .facts = own.facts(), .rows = if (own.head != null) mtp.max_rows else 2, + .prompt_rows = @intCast(own.engine.chunk), .stream_bytes = stream_bytes, .ctx = own, .deinit = release, .target = own, .own_glue = own.set == null, + .follow = if (o.rank != 0) follow else null, }; } @@ -307,6 +356,11 @@ pub fn Namespace(comptime Tri: type) type { return null; } + fn follow(ptr: *anyopaque) anyerror!void { + const self = of(ptr); + return self.ranks.?.follow(Tri, self.tri(), self.logits.ptr); + } + fn release(p: *anyopaque) void { const own: *Cuda = @ptrCast(@alignCast(p)); var it = own.lanes.valueIterator(); @@ -314,6 +368,7 @@ pub fn Namespace(comptime Tri: type) type { own.lanes.deinit(own.gpa); own.logits.free(); if (own.head) |h| h.deinit(); + if (own.ranks) |r| r.deinit(); own.engine.deinit(); if (own.set) |*x| x.deinit(); own.glue.deinit(); @@ -323,7 +378,7 @@ pub fn Namespace(comptime Tri: type) type { fn dropLane(self: *Cuda, l: *Lane) void { if (l.draft) |*d| d.deinit(); - l.seq.deinit(self.gpa); + if (self.ranks) |r| r.drop(&l.seq) else l.seq.deinit(self.gpa); } fn of(ptr: *anyopaque) *Cuda { @@ -399,7 +454,7 @@ pub fn Namespace(comptime Tri: type) type { var at: usize = 0; while (at < plain) : (at += round_streams) { const n = @min(round_streams, plain - at); - try self.engine.round(Tri, self.tri(), seqs[at..][0..n], ids[at..][0..n], self.logits.ptr); + if (self.ranks) |r| try r.round(Tri, self.tri(), seqs[at..][0..n], ids[at..][0..n], self.logits.ptr) else try self.engine.round(Tri, self.tri(), seqs[at..][0..n], ids[at..][0..n], self.logits.ptr); } for (seqs[0..plain], at_window[0..plain]) |q, wi| out[wi].sampled[0] = try self.tokenAt(q, windows[wi].stream, windows[wi].positions[0]); } @@ -478,6 +533,7 @@ pub fn Namespace(comptime Tri: type) type { const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); const l = try self.lane(s); if (l.seq.fed != at) return error.NotAtMark; + if (self.ranks) |r| return try r.save(&l.seq); return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); } @@ -485,11 +541,13 @@ pub fn Namespace(comptime Tri: type) type { const self = of(ptr); const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); const l = try self.lane(s); + if (self.ranks) |r| return r.restore(&l.seq, @ptrCast(@alignCast(saved))); try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); } fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { const self = of(ptr); + if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); } }; diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index b2c54b1fe0..178d2081c1 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -376,6 +376,7 @@ test { _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); + _ = @import("cuda_rank_memory.zig"); _ = @import("cuda_rank_wire.zig"); _ = @import("cuda_modelopt_experts.zig"); _ = row_path; diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index ee00b4f793..de549deff6 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -49,6 +49,7 @@ pub const Session = struct { seqs: std.AutoHashMapUnmanaged(u64, *prompt.Seq) = .empty, saved: std.AutoHashMapUnmanaged(u64, *snaps.Snap) = .empty, stopped: bool = false, + available: std.atomic.Value(bool) = .init(true), pub fn init(a: std.mem.Allocator, io: std.Io, host: []const u8, port: u16, rank: usize, mine: Identity) !*Session { const s = try a.create(Session); @@ -76,8 +77,15 @@ pub const Session = struct { return s; } + pub fn healthy(s: *const Session) bool { + if (!s.available.load(.acquire)) return false; + s.link.checkPeer() catch return false; + return true; + } + pub fn poison(s: *Session) void { s.stopped = true; + s.available.store(false, .release); s.collective.abort(); _ = std.c.shutdown(s.link.fd, std.posix.SHUT.RDWR); } @@ -235,6 +243,7 @@ pub const Session = struct { try s.link.sendHeader(.{ .kind = .ack, .sequence = h.sequence, .object = h.object }, wire.deadline(30_000)); if (stop) { s.stopped = true; + s.available.store(false, .release); return; } } diff --git a/zig/src/families/flashnext/cuda_rank_memory.zig b/zig/src/families/flashnext/cuda_rank_memory.zig new file mode 100644 index 0000000000..642510a680 --- /dev/null +++ b/zig/src/families/flashnext/cuda_rank_memory.zig @@ -0,0 +1,78 @@ +//! Conservative rank-local checkpoint bytes: replicated dense operators, partitioned routed weights, one shared branch. +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const Entry = core.safetensors.Entry; + +pub fn estimate(a: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize) !u64 { + if ((world != 1 and world != 2) or rank >= world) return error.BadRank; + var store = try wt.Store.open(a, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(a); + var wi = wm.object.iterator(); + while (wi.next()) |item| try files.put(a, item.value_ptr.string, {}); + var total: u64 = 0; + var fi = files.keyIterator(); + while (fi.next()) |name| { + const path = try std.fs.path.join(a, &.{ dir, name.* }); + defer a.free(path); + var file = try core.safetensors.File.open(a, io, path); + defer file.close(io); + var tensors = file.names.iterator(); + while (tensors.next()) |tensor| { + total = try std.math.add(u64, total, try tensorBytes(tensor.key_ptr.*, tensor.value_ptr.*, world, rank)); + } + } + return total; +} + +fn tensorBytes(name: []const u8, e: Entry, world: usize, rank: usize) !u64 { + const bytes: u64 = e.end - e.begin; + if (wt.hostOnly(name)) return 0; + if (rank != 0 and std.mem.indexOf(u8, name, ".mlp.shared_expert.") != null) return 0; + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null) { + if (std.mem.endsWith(u8, name, ".weight_scale_2")) return bytes; + if (!std.mem.endsWith(u8, name, ".weight")) return 0; + if (e.dtype != .u8 or e.rank != 2 or bytes % (8 * world) != 0) return error.UnexpectedTensor; + // NVFP4 device blocks hold 128 code words and 16 scale words; global scales stay replicated. + return try std.math.mul(u64, bytes / (8 * world), 9); + } + if (std.mem.indexOf(u8, name, ".weight_scale") != null or std.mem.endsWith(u8, name, ".input_scale")) return 0; + if (e.dtype == .f8_e4m3 and e.rank == 2 and std.mem.endsWith(u8, name, ".weight")) { + const codes = try std.math.mul(u64, cuda.mxfp8.padded(e.dim(0)), e.dim(1)); + return try std.math.add(u64, codes, codes / 32); + } + if (e.dtype == .bf16 and e.rank == 2 and std.mem.endsWith(u8, name, ".weight")) { + if (std.mem.indexOf(u8, name, ".mlp.gate.") != null or std.mem.indexOf(u8, name, ".shared_expert_gate.") != null) return bytes; + return try std.math.mul(u64, try std.math.mul(u64, cuda.nvfp4.padded(e.dim(0)), e.dim(1)), 2); + } + if (e.dtype == .bf16 and centered(name)) return try std.math.mul(u64, bytes, 2); + return bytes; +} + +fn centered(name: []const u8) bool { + for ([_][]const u8{ ".hc_norm.weight", ".q_norm.weight", ".k_norm.weight", ".q_layernorm.weight", ".k_layernorm.weight", ".norm_key.weight", ".norm_query.weight", ".norm_conv.weight" }) |suffix| { + if (std.mem.endsWith(u8, name, suffix)) return true; + } + return false; +} + +test "rank estimates partition routed storage while retaining replicated and rank-zero shared operators" { + const e: Entry = .{ .dtype = .u8, .rank = 2, .shape = .{ 640, 1280, 1, 1, 1 }, .begin = 0, .end = 640 * 1280 }; + const name = "model.language_model.layers.0.mlp.experts.0.gate_proj.weight"; + try std.testing.expectEqual(@as(u64, cuda.experts.Nvfp4Experts.bytes(1, 320, 2560, 1)), try tensorBytes(name, e, 2, 0)); + try std.testing.expectEqual(try tensorBytes(name, e, 2, 0), try tensorBytes(name, e, 2, 1)); + var shared = e; + shared.dtype = .f8_e4m3; + shared.shape = .{ 640, 2560, 1, 1, 1 }; + shared.end = 640 * 2560; + const shared_name = "model.language_model.layers.0.mlp.shared_expert.gate_proj.weight"; + try std.testing.expectEqual(@as(u64, 640 * 2560 * 33 / 32), try tensorBytes(shared_name, shared, 2, 0)); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(shared_name, shared, 2, 1)); + const norm: Entry = .{ .dtype = .bf16, .rank = 1, .shape = .{ 256, 1, 1, 1, 1 }, .begin = 0, .end = 512 }; + try std.testing.expectEqual(@as(u64, 1024), try tensorBytes("model.language_model.layers.3.self_attn.q_norm.weight", norm, 2, 1)); + try std.testing.expectEqual(@as(u64, 512), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", norm, 2, 1)); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index a947de3beb..dedb098700 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -10,6 +10,9 @@ const lanes = @import("lanes"); const nemotron = @import("nemotron"); const flashnext = @import("flashnext"); const Allocator = std.mem.Allocator; +const family_memory = @import("cuda_family_memory.zig"); +const learnBytes = family_memory.learnBytes; +const toGib = family_memory.toGib; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); @@ -147,11 +150,6 @@ fn pool(a: Allocator, io: std.Io, ctx: *const cuda.Context, problem: *[]const u8 return .{ .free = free, .total = total, .reserve = reserve, .limit = limit, .unified = unified }; } -fn readMemory(_: ?*anyopaque, reset_peak: bool) ?api.Memory { - const u = cuda.usage(reset_peak); - return .{ .active = u.device, .cache = 0, .peak = u.peak, .pinned = u.host }; -} - /// A lone driver returns false at completion or true when yield hands it to the lane core. const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; @@ -220,15 +218,16 @@ const Host = struct { } /// The round loop and the lane host over the family's backend, its thread started. - fn serve(h: *Host, io: std.Io, facts: lanes.Model, rows: u32, info: api.Info, explain: ?api.Explain, per: PerFamily) !void { + fn serve(h: *Host, io: std.Io, facts: lanes.Model, rows: u32, info: api.Info, explain: ?api.Explain, per: PerFamily, health: ?api.HealthSource) !void { h.cfg = try lanes.Config.init(h.gpa, facts, rows, rows - 1); errdefer h.cfg.deinit(h.gpa); h.clock = .{ .io = io }; h.core = lanes.Engine.init(h.gpa, &h.cfg, h.backend(), h.clock.clock()); errdefer h.core.deinit(); h.host = api.LaneHost.init(h.gpa, io, &h.core, info); - h.host.memory = if (h.gpu != null) .{ .read = readMemory } else null; + h.host.memory = if (h.gpu != null) .{ .read = family_memory.read } else null; h.host.explain = explain; + h.host.health = health; h.host.cache = h.store; h.host.warm_quiet = if (info.warm_turns) per.warm_quiet else 0; if (h.lone != null) h.host.lone = .{ .ctx = h, .run = loneRun, .sampled = true, .logprobs = h.inner.vtable.first_row != null, .structures = h.inner.vtable.first_masked != null, .min_budget = per.lone_min_budget, .resumes = true }; @@ -372,7 +371,10 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; const before = try pool(a, io, &g.ctx, problem) orelse return null; // a family that keeps part of its checkpoint off the device says what it uploads; else every shard's bytes - const weights = if (@hasDecl(F, "deviceWeightBytes")) F.deviceWeightBytes(a, io, o.dir) else weightBytes(io, o.dir); + const weights = family_memory.weights(F, a, io, o, weightBytes) catch |err| { + problem.* = try std.fmt.allocPrint(a, "the checkpoint placement is unsupported ({s})", .{@errorName(err)}); + return null; + }; const held0 = cuda.usage(false).device; if (weights > before.room(held0)) { problem.* = try std.fmt.allocPrint(a, "the checkpoint's {d:.1} GiB of weights do not fit the {d:.1} GiB the CUDA memory budget grants ({d:.1} GiB free less a {d:.1} GiB reserve{s}); free device memory or adjust TENSORFOLD_MEMORY_RESERVE_GIB / TENSORFOLD_CUDA_MEMORY_LIMIT_GB", .{ toGib(weights), toGib(before.room(held0)), toGib(before.free), toGib(before.reserve), if (before.limit != null) ", under TENSORFOLD_CUDA_MEMORY_LIMIT_GB" else "" }); @@ -398,7 +400,10 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O defer if (!opened) loaded.deinit(loaded.ctx); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| learnBytes(loaded); + const room = family_memory.room(F, loaded.ctx, after.room(model) -| learnBytes(loaded)) catch |err| { + problem.* = try std.fmt.allocPrint(a, "the ranks could not agree on available memory ({s})", .{@errorName(err)}); + return null; + }; const free = loaded.free_streams; const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { @@ -417,8 +422,9 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; const snaps = comptime snapshots(F); const cache = if (snaps != null) budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)) else 0; + const prompt_rows = if (@hasField(@TypeOf(loaded), "prompt_rows")) loaded.prompt_rows else F.prompt_rows; h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learnBytes(loaded) > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learnBytes(loaded))}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learnBytes(loaded) > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learnBytes(loaded))}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels != null and !ownGlue(loaded)) try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{kernels.?}) else "own glue kernels", }); errdefer gpa.free(h.startup); @@ -454,7 +460,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null and @hasDecl(F, "warm_quiet"), .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }, comptime perFamily(F)); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null and @hasDecl(F, "warm_quiet"), .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }, comptime perFamily(F), if (@hasDecl(F, "healthy")) .{ .ctx = loaded.ctx, .read = F.healthy } else null); opened = true; const follow_loop: ?*const fn (*anyopaque) anyerror!void = if (@hasField(@TypeOf(loaded), "follow")) loaded.follow else null; if (follow_loop) |f| { @@ -464,16 +470,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } -/// A family's bytes kept for --slide (0 for families without a learner). -fn learnBytes(loaded: anytype) u64 { - return if (@hasField(@TypeOf(loaded), "learn_bytes")) loaded.learn_bytes else 0; -} - -fn toGib(bytes: u64) f64 { - return @as(f64, @floatFromInt(bytes)) / gib; -} - test { + _ = family_memory; _ = budget; // cuda_memory.zig's own tests, which no other test binary reaches _ = refusals; _ = lane_backend; @@ -529,7 +527,7 @@ test "a cancel during a long prompt stops it, and the next request's reply is un const h = try gpa.create(Host); defer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = null, .family = &target, .release = Nothing.release, .inner = target.backend(), .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined }; - try h.serve(std.testing.io, .{ .exact_width = 8, .mtp = true, .speculate = true, .drafts = 7, .hidden_rows = true }, 8, .{ .lanes = 2 }, null, .{}); + try h.serve(std.testing.io, .{ .exact_width = 8, .mtp = true, .speculate = true, .drafts = 7, .hidden_rows = true }, 8, .{ .lanes = 2 }, null, .{}, null); defer { h.host.stop(); h.core.deinit(); diff --git a/zig/src/native/cuda_family_memory.zig b/zig/src/native/cuda_family_memory.zig new file mode 100644 index 0000000000..3f7cec80ab --- /dev/null +++ b/zig/src/native/cuda_family_memory.zig @@ -0,0 +1,41 @@ +//! Family-aware checkpoint and stream admission; rank coordination runs before the host admits any stream. +const std = @import("std"); +const cuda = @import("cuda"); +const api = @import("engine_api"); + +pub fn weights(comptime F: type, a: std.mem.Allocator, io: std.Io, o: api.Open, comptime fallback: anytype) !u64 { + if (@hasDecl(F, "placementWeightBytes")) return F.placementWeightBytes(a, io, o.dir, o.tp, o.rank, o.drafts); + if (@hasDecl(F, "deviceWeightBytes")) return F.deviceWeightBytes(a, io, o.dir); + return fallback(io, o.dir); +} + +pub fn room(comptime F: type, ptr: *anyopaque, local: u64) !u64 { + if (@hasDecl(F, "rankRoom")) return F.rankRoom(ptr, local); + return local; +} + +pub fn read(_: ?*anyopaque, reset_peak: bool) ?api.Memory { + const u = cuda.usage(reset_peak); + return .{ .active = u.device, .cache = 0, .peak = u.peak, .pinned = u.host }; +} + +pub fn learnBytes(loaded: anytype) u64 { + return if (@hasField(@TypeOf(loaded), "learn_bytes")) loaded.learn_bytes else 0; +} + +pub fn toGib(bytes: u64) f64 { + return @as(f64, @floatFromInt(bytes)) / @as(f64, 1 << 30); +} + +test "rank room is agreed before stream admission and preserves a single-rank family's pool" { + const Single = struct {}; + const Pair = struct { + pub fn rankRoom(_: *anyopaque, local: u64) !u64 { + return @min(local, 19); + } + }; + var context: u8 = 0; + try std.testing.expectEqual(@as(u64, 30), try room(Single, &context, 30)); + try std.testing.expectEqual(@as(u64, 19), try room(Pair, &context, 30)); + try std.testing.expectEqual(@as(u64, 7), try room(Pair, &context, 7)); +} diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 63eebea817..884ebcb2da 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -217,7 +217,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an defer if (srv.keepalive) |k| k.end(); srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| return switch (e) { error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), - error.Closed => cx.fail(.other, "the scheduler is closed", .{}), + error.Closed => cx.fail(.capacity, "the engine is unavailable; retry after it restarts", .{}), error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), }; release(srv, preparing); // a background request waits only while a foreground one prepares diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index aaf23610ae..cdfa58b433 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -344,10 +344,10 @@ const Run = struct { const a = r.a; const kind_ok = cx.kind != .other and cx.kind != .server; const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; - r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); + r.out.vt.reply(r.out.ctx, if (kind_ok) cx.status() else 500, wrapError(a, body) catch return); } - /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. + /// An unsent reply: request 400, unavailable 503, failure 500, nothing when the client left. fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { switch (e) { error.Cancelled => {}, diff --git a/zig/src/server/status_routes.zig b/zig/src/server/status_routes.zig index 89ee5c8488..a4056d9503 100644 --- a/zig/src/server/status_routes.zig +++ b/zig/src/server/status_routes.zig @@ -11,6 +11,9 @@ const Allocator = std.mem.Allocator; /// With keys on, health names nothing about the model. pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { + var status: api.Status = .{}; + srv.engine.status(&status, &.{}); + if (!status.healthy) return conn.sendJson(503, "{\"status\": \"unavailable\"}"); if (srv.keys) |k| if (k.enabled()) return conn.sendJson(200, "{\"status\": \"ok\"}"); const o = try json.newObject(a); try o.put(a, "status", .{ .string = "ok" }); @@ -20,8 +23,6 @@ pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { try o.put(a, "model_ids", .{ .array = ids }); try o.put(a, "max_batch_size", try json.intValue(a, srv.info.lanes)); if (srv.info.policy.len > 0) try o.put(a, "policy", .{ .string = srv.info.policy }); - var status: api.Status = .{}; - srv.engine.status(&status, &.{}); try o.put(a, "warming", .{ .bool = status.warming }); const memory = try json.newObject(a); if (srv.engine.memory(std.mem.indexOf(u8, conn.path, "reset_peak=1") != null)) |mem| { diff --git a/zig/src/server/stream_preflight_test.zig b/zig/src/server/stream_preflight_test.zig index 601573da5f..25e4024d07 100644 --- a/zig/src/server/stream_preflight_test.zig +++ b/zig/src/server/stream_preflight_test.zig @@ -457,3 +457,30 @@ test "unsupported budgets leave zero unset and thinking-disabled requests valid" } try std.testing.expectEqual(@as(usize, 3), backend.submissions); } + +test "closed or busy admission is service unavailable across chat, messages and responses" { + const ClosedEngine = struct { + reason: api.SubmitError = error.Closed, + fn engine(e: *@This()) api.Engine { + return .{ .ctx = e, .vtable = &.{ .info = HiEngine.info, .submit = submit, .cancel = HiEngine.cancel, .status = HiEngine.status, .memory = HiEngine.memory } }; + } + fn submit(ctx: *anyopaque, _: api.Id, _: *const api.Request, _: api.Sink) api.SubmitError!void { + return @as(*@This(), @ptrCast(@alignCast(ctx))).reason; + } + }; + const cases = [_]struct { route: Route, body: []const u8 }{ + .{ .route = .chat, .body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":1}" }, + .{ .route = .messages, .body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":1}" }, + .{ .route = .responses, .body = "{\"model\":\"test-model\",\"input\":\"x\",\"max_output_tokens\":1}" }, + }; + var engine: ClosedEngine = .{}; + for ([_]api.SubmitError{ error.Closed, error.Busy }) |reason| { + engine.reason = reason; + for (cases) |c| { + const seen = try sendWith(engine.engine(), c.route, c.body, .{}); + try std.testing.expectEqual(@as(?u16, 503), seen.status); + try std.testing.expect(!seen.opened); + try std.testing.expectEqual(@as(i64, 0), seen.preparing); + } + } +} From a801b9a35572b700c2f7c06d38fc67365e5efb39 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:13:32 -0700 Subject: [PATCH 071/115] Report unavailable engine health and preserve admission HTTP 503 --- zig/build/cuda.zig | 2 +- zig/src/core/engine_api.zig | 7 +++++ zig/src/core/lane_host.zig | 3 +- zig/src/core/lane_host_status.zig | 1 + zig/src/core/lane_host_test.zig | 36 ++++++++++++++++++++++++ zig/src/server/chat.zig | 2 +- zig/src/server/openai.zig | 4 +-- zig/src/server/status_routes.zig | 5 ++-- zig/src/server/stream_preflight_test.zig | 27 ++++++++++++++++++ 9 files changed, 80 insertions(+), 7 deletions(-) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..a46d2e846e 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -210,7 +210,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native_modules.api, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 903dc4d40b..bcd3f8f3b8 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -244,7 +244,14 @@ pub const MemorySource = struct { read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, }; +/// Thread-safe backend readiness; no GPU context or blocking device work is allowed on HTTP threads. +pub const HealthSource = struct { + ctx: *anyopaque, + read: *const fn (*anyopaque) bool, +}; + pub const Status = struct { + healthy: bool = true, /// Requests in prefill or decode, and those waiting for a lane. running: u32 = 0, waiting: u32 = 0, diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4394fdac05..f213633df3 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -43,6 +43,7 @@ pub const LaneHost = struct { lone_job: ?*Job = null, // the job that driver holds now cache: ?*pc.Store = null, // kept prompt states (engine thread only); the backend restores and saves them decoded_rows: bool = false, // the backend prefills Request.decode_spans with decoded rows' arithmetic, cache or not + health: ?api.HealthSource = null, memory: ?api.MemorySource = null, // the backend's memory counts; null: Engine.memory reports none explain: ?api.Explain = null, // the backend's words for a request it refuses; null: the error's name learner: ?api.Learner = null, // the family's Sliding Weights learner; null: learn requests are refused @@ -143,7 +144,7 @@ pub const LaneHost = struct { job.* = .{ .host = h, .id = id, .request = request, .sink = sink }; h.mutex.lockUncancelable(h.io); defer h.mutex.unlock(h.io); - if (h.closing) { + if (h.closing or (if (h.health) |source| !source.read(source.ctx) else false)) { h.gpa.destroy(job); return error.Closed; } diff --git a/zig/src/core/lane_host_status.zig b/zig/src/core/lane_host_status.zig index a92bf19a0c..9932fe8271 100644 --- a/zig/src/core/lane_host_status.zig +++ b/zig/src/core/lane_host_status.zig @@ -20,6 +20,7 @@ pub fn read(ctx: *anyopaque, out: *api.Status, stream_tokens: []u32) void { const n = @min(stream_tokens.len, h.live_tokens.items.len); @memcpy(stream_tokens[0..n], h.live_tokens.items[0..n]); out.* = .{ + .healthy = if (h.health) |source| source.read(source.ctx) else true, .running = @intCast(h.admitted.items.len), .waiting = @intCast(h.queued.items.len), .decode_tokens_per_second = @as(f64, @floatFromInt(tokens)) / 2.0, diff --git a/zig/src/core/lane_host_test.zig b/zig/src/core/lane_host_test.zig index 20a91348f1..b5460c9a59 100644 --- a/zig/src/core/lane_host_test.zig +++ b/zig/src/core/lane_host_test.zig @@ -562,3 +562,39 @@ test { _ = @import("lane_host_lone_test.zig"); _ = @import("lane_host_admit_test.zig"); } + +test "an unhealthy backend fails health and refuses a request before queue admission" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{}, 1, 0); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 1 }); + const Health = struct { + available: bool = true, + fn read(ctx: *anyopaque) bool { + return @as(*@This(), @ptrCast(@alignCast(ctx))).available; + } + fn event(_: *anyopaque, _: Id, _: *const Event) void { + @panic("a refused request must never reach its sink"); + } + }; + var health: Health = .{}; + const engine = host.engine(); + var status: api.Status = .{}; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + host.health = .{ .ctx = &health, .read = Health.read }; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + health.available = false; + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); + const request: Request = .{ .prompt = &.{1}, .max_tokens = 1 }; + try std.testing.expectError(error.Closed, engine.submit(1, &request, .{ .ctx = &health, .event = Health.event })); + try std.testing.expectEqual(@as(usize, 0), host.queued.items.len); + try std.testing.expectEqual(@as(usize, 0), host.admitted.items.len); +} diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 7c702ef432..2f1e2f8a09 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -191,7 +191,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an defer if (srv.keepalive) |k| k.end(); srv.engine.submit(id, &request, .{ .ctx = &box, .event = Mailbox.onEvent }) catch |e| return switch (e) { error.Busy => cx.fail(.capacity, "the engine is busy; retry shortly", .{}), - error.Closed => cx.fail(.other, "the scheduler is closed", .{}), + error.Closed => cx.fail(.capacity, "the engine is unavailable; retry after it restarts", .{}), error.InvalidSpans => cx.fail(.request, "the prompt span layout is invalid", .{}), }; release(srv, preparing); // a background request waits only while a foreground one prepares diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index b8782aaf96..8d79e58330 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -339,10 +339,10 @@ const Run = struct { const a = r.a; const kind_ok = cx.kind != .other and cx.kind != .server; const body = (if (kind_ok) errorBody(a, cx, field) else errorOther(a, cx.message)) catch return; - r.out.vt.reply(r.out.ctx, if (kind_ok) 400 else 500, wrapError(a, body) catch return); + r.out.vt.reply(r.out.ctx, if (kind_ok) cx.status() else 500, wrapError(a, body) catch return); } - /// A reply that ended before anything was sent: a refusal's 400, a failure's 500, nothing when the client left. + /// An unsent reply: request 400, unavailable 503, failure 500, nothing when the client left. fn unsent(r: *Run, cx: *Cx, e: chat.Failure, field: []const u8) void { switch (e) { error.Cancelled => {}, diff --git a/zig/src/server/status_routes.zig b/zig/src/server/status_routes.zig index 89ee5c8488..a4056d9503 100644 --- a/zig/src/server/status_routes.zig +++ b/zig/src/server/status_routes.zig @@ -11,6 +11,9 @@ const Allocator = std.mem.Allocator; /// With keys on, health names nothing about the model. pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { + var status: api.Status = .{}; + srv.engine.status(&status, &.{}); + if (!status.healthy) return conn.sendJson(503, "{\"status\": \"unavailable\"}"); if (srv.keys) |k| if (k.enabled()) return conn.sendJson(200, "{\"status\": \"ok\"}"); const o = try json.newObject(a); try o.put(a, "status", .{ .string = "ok" }); @@ -20,8 +23,6 @@ pub fn health(srv: *Server, conn: *Conn, a: Allocator) !void { try o.put(a, "model_ids", .{ .array = ids }); try o.put(a, "max_batch_size", try json.intValue(a, srv.info.lanes)); if (srv.info.policy.len > 0) try o.put(a, "policy", .{ .string = srv.info.policy }); - var status: api.Status = .{}; - srv.engine.status(&status, &.{}); try o.put(a, "warming", .{ .bool = status.warming }); const memory = try json.newObject(a); if (srv.engine.memory(std.mem.indexOf(u8, conn.path, "reset_peak=1") != null)) |mem| { diff --git a/zig/src/server/stream_preflight_test.zig b/zig/src/server/stream_preflight_test.zig index 601573da5f..25e4024d07 100644 --- a/zig/src/server/stream_preflight_test.zig +++ b/zig/src/server/stream_preflight_test.zig @@ -457,3 +457,30 @@ test "unsupported budgets leave zero unset and thinking-disabled requests valid" } try std.testing.expectEqual(@as(usize, 3), backend.submissions); } + +test "closed or busy admission is service unavailable across chat, messages and responses" { + const ClosedEngine = struct { + reason: api.SubmitError = error.Closed, + fn engine(e: *@This()) api.Engine { + return .{ .ctx = e, .vtable = &.{ .info = HiEngine.info, .submit = submit, .cancel = HiEngine.cancel, .status = HiEngine.status, .memory = HiEngine.memory } }; + } + fn submit(ctx: *anyopaque, _: api.Id, _: *const api.Request, _: api.Sink) api.SubmitError!void { + return @as(*@This(), @ptrCast(@alignCast(ctx))).reason; + } + }; + const cases = [_]struct { route: Route, body: []const u8 }{ + .{ .route = .chat, .body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":1}" }, + .{ .route = .messages, .body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":1}" }, + .{ .route = .responses, .body = "{\"model\":\"test-model\",\"input\":\"x\",\"max_output_tokens\":1}" }, + }; + var engine: ClosedEngine = .{}; + for ([_]api.SubmitError{ error.Closed, error.Busy }) |reason| { + engine.reason = reason; + for (cases) |c| { + const seen = try sendWith(engine.engine(), c.route, c.body, .{}); + try std.testing.expectEqual(@as(?u16, 503), seen.status); + try std.testing.expect(!seen.opened); + try std.testing.expectEqual(@as(i64, 0), seen.preparing); + } + } +} From 09830ff681e82189121d99f2d404de29191b1148 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:47:52 -0700 Subject: [PATCH 072/115] perf: queue TP work and disable delayed control frames --- zig/src/families/flashnext/cuda_batch.zig | 11 ++++++++--- zig/src/families/flashnext/cuda_collective.zig | 10 ++++++++-- zig/src/families/flashnext/cuda_prompt.zig | 13 +++++++++---- zig/src/families/flashnext/cuda_rank_wire.zig | 13 ++++++++++++- zig/tests/cuda/flashnext_rank_tests.zig | 16 +++++++++++++--- 5 files changed, 50 insertions(+), 13 deletions(-) diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 621523de75..28c927b8c4 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -452,7 +452,7 @@ pub fn Pass(comptime Tri: type) type { const n = ids.len; const t0 = std.Io.Clock.awake.now(self.io).toNanoseconds(); defer if (profile()) { - self.stream.synchronize() catch {}; + self.complete() catch {}; std.debug.print("batch: {d} rows at {d} in {d:.1} ms\n", .{ n, pos0, @as(f64, @floatFromInt(std.Io.Clock.awake.now(self.io).toNanoseconds() - t0)) / 1e6 }); }; if (n == 0 or n > s.cap) return error.UnexpectedTensor; @@ -505,6 +505,11 @@ pub fn Pass(comptime Tri: type) type { try self.head(logits_host); } + fn complete(self: *Self) !void { + if (self.m.collective) |c| return c.finish(self.stream); + try self.stream.synchronize(); + } + fn head(self: *Self, logits_host: ?[]u8) !void { const s = self.s; const n = self.rows; @@ -522,7 +527,7 @@ pub fn Pass(comptime Tri: type) type { const out = if (self.hooks.logits_out != 0) self.hooks.logits_out else return error.UnexpectedTensor; try self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); if (logits_host) |lh| { - try self.stream.synchronize(); + try self.complete(); const d = self.stream.d; try d.check(d.api.cuMemcpyDtoH_v2(lh.ptr, out + (n - 1) * vocab * 2, vocab * 2), "cuMemcpyDtoH"); } @@ -549,7 +554,7 @@ pub fn Pass(comptime Tri: type) type { try self.readout(self.m.mixer, 0, 1); try self.lane(s.mixed, s.xsm, self.m.head, s.logits, s.dummy, 1, 1); if (logits_host) |lh| { - try self.stream.synchronize(); + try self.complete(); const d = self.stream.d; try d.check(d.api.cuMemcpyDtoH_v2(lh.ptr, s.logits, vocab * 2), "cuMemcpyDtoH"); } diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index a52924a586..a2aae8dc97 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -75,13 +75,19 @@ pub const Collective = struct { } } - /// FP32 sum is complete on the compute stream on return, or the communicator is permanently aborted. + /// Wait at a host-visible boundary while retaining peer, asynchronous-error and timeout checks. + pub fn finish(c: *Collective, stream: cuda.Stream) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + try c.wait(stream, wire.deadline(c.timeout_ms)); + } + + /// FP32 sum is enqueued on the compute stream; later kernels on that stream observe its result. pub fn sum(c: *Collective, stream: cuda.Stream, values: u64, count: usize) !void { if (c.failed or c.comm == null) return error.RankDisconnected; if (count == 0) return; errdefer c.abort(); try c.link.checkPeer(); try c.submitted(c.lib.api.ncclAllReduce(values, values, count, .f32, .sum, c.comm, stream.handle)); - try c.wait(stream, wire.deadline(c.timeout_ms)); + try c.wait(null, wire.deadline(c.timeout_ms)); } }; diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 91c316f91f..71be96b5c7 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -257,12 +257,17 @@ pub const Engine = struct { return e.forward(Tri, tri, seq, ids, .{}); } + fn complete(e: *Engine) !void { + if (e.collective) |c| return c.finish(e.stream.*); + try e.stream.synchronize(); + } + /// Block scores for a `cap`-token sequence's chunks of up to `rows` rows, reallocated only to grow. pub fn reserveScores(e: *Engine, cap: usize, rows: usize) !u64 { const need = batch.scoreBytes(cap, rows); if (e.scores) |*b| { if (b.len >= need) return b.ptr; - try e.stream.synchronize(); + try e.complete(); b.free(); e.scores = null; } @@ -274,7 +279,7 @@ pub const Engine = struct { pub fn reserve(e: *Engine, rows: usize) !*batch.Scratch { if (e.batch_s) |*bs| { if (bs.cap >= rows) return bs; - try e.stream.synchronize(); + try e.complete(); bs.free(e.gpa); e.batch_s = null; } @@ -308,7 +313,7 @@ pub const Engine = struct { try p.chunk(ids[at..][0..n], seq.fed + at, if (last and hooks.logits != .none) seq.logits else null); at += n; } - if (hooks.logits == .none) try e.stream.synchronize(); + if (hooks.logits == .none) try e.complete(); seq.fed += ids.len; } @@ -339,7 +344,7 @@ pub const Engine = struct { .hooks = .{ .logits = .all, .logits_out = out }, }; try p.chunk(ids, 0, null); - try e.stream.synchronize(); + try e.complete(); const d = e.stream.d; for (seqs, 0..) |q, r| { try d.check(d.api.cuMemcpyDtoH_v2(q.logits.ptr, out + r * vocab * 2, vocab * 2), "cuMemcpyDtoH"); diff --git a/zig/src/families/flashnext/cuda_rank_wire.zig b/zig/src/families/flashnext/cuda_rank_wire.zig index 2e0c4fe361..e93f02c853 100644 --- a/zig/src/families/flashnext/cuda_rank_wire.zig +++ b/zig/src/families/flashnext/cuda_rank_wire.zig @@ -98,7 +98,11 @@ pub const Link = struct { try wait(fd, p.POLL.IN, until); const peer = std.c.accept4(fd, null, null, p.SOCK.NONBLOCK | p.SOCK.CLOEXEC); switch (p.errno(peer)) { - .SUCCESS => return .{ .fd = peer }, + .SUCCESS => { + errdefer _ = std.c.close(peer); + try noDelay(peer); + return .{ .fd = peer }; + }, .INTR, .AGAIN => continue, else => return error.RankAccept, } @@ -188,9 +192,16 @@ pub const Link = struct { fn socket(family: u32) !p.socket_t { const fd = std.c.socket(family, p.SOCK.STREAM | p.SOCK.NONBLOCK | p.SOCK.CLOEXEC, 0); if (p.errno(fd) != .SUCCESS) return error.RankSocket; + errdefer _ = std.c.close(fd); + try noDelay(fd); return fd; } +fn noDelay(fd: p.socket_t) !void { + const one: c_int = 1; + try p.setsockopt(fd, p.IPPROTO.TCP, p.TCP.NODELAY, std.mem.asBytes(&one)); +} + test "rank frames are versioned, bounded and exact" { const h: Header = .{ .kind = .forward, .sequence = 700, .object = 19, .arg = 42, .bytes = max_payload, .flags = 3 }; var raw = try h.encode(); diff --git a/zig/tests/cuda/flashnext_rank_tests.zig b/zig/tests/cuda/flashnext_rank_tests.zig index dbbd504fc6..f7bdcd15c9 100644 --- a/zig/tests/cuda/flashnext_rank_tests.zig +++ b/zig/tests/cuda/flashnext_rank_tests.zig @@ -16,6 +16,9 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { } const link = try wire.Link.open(gpu.io, "127.0.0.1", port, rank, 30_000); defer link.close(); + var no_delay: c_int = 0; + var option_len: std.posix.socklen_t = @sizeOf(c_int); + try check.expect(std.c.getsockopt(link.fd, std.posix.IPPROTO.TCP, std.posix.TCP.NODELAY, &no_delay, &option_len) == 0 and no_delay == 1, "both connected rank sockets disable Nagle", .{}); const until = wire.deadline(30_000); const hello: wire.Header = .{ .kind = .hello, .sequence = 1, .object = 77, .arg = 123, .bytes = 7 }; if (rank == 0) { @@ -51,7 +54,7 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.asBytes(&@as(f32, 1))); defer values.free(); const start = wire.now(); - collective.sum(stream, values.ptr, 1) catch |err| { + completeSum(&collective, stream, values.ptr, 1) catch |err| { try check.expect(err == error.RankDisconnected or err == error.RankTimeout or err == error.NcclFailed, "failure is surfaced ({s})", .{@errorName(err)}); try check.expect(collective.failed and collective.comm == null, "failed communicator is aborted and cannot be reused", .{}); try check.expect(wire.now() - start < 3 * std.time.ns_per_s, "failed collective terminates within bound", .{}); @@ -68,8 +71,10 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { var values = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); defer values.free(); try collective.sum(stream, values.ptr, count); + try collective.sum(stream, values.ptr, count); + try collective.finish(stream); try values.download(0, std.mem.sliceAsBytes(host)); - for (host[0..count], 0..) |actual, i| try check.expect(actual == @as(f32, @floatFromInt(2 * (i % 127) + 1)), "rank {d}: exact sum at {d}", .{ rank, i }); + for (host[0..count], 0..) |actual, i| try check.expect(actual == @as(f32, @floatFromInt(4 * (i % 127) + 2)), "rank {d}: exact sum at {d}", .{ rank, i }); // Acknowledge completion before either peer can close or start its next reduction. const h: wire.Header = .{ .kind = .ready, .sequence = seq + 2, .arg = count }; if (rank == 0) { @@ -80,5 +85,10 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { try link.sendHeader(h, wire.deadline(30_000)); } } - check.pass("rank {d}: two-process NCCL FP32 sums exact for empty, scalar, tail and 1,000,003 values; fragmented framed control preserved", .{rank}); + check.pass("rank {d}: two-process chained NCCL FP32 sums exact for empty, scalar, tail and 1,000,003 values; fragmented framed control preserved", .{rank}); +} + +fn completeSum(c: *flash.collective.Collective, stream: cuda.Stream, values: u64, count: usize) !void { + try c.sum(stream, values, count); + try c.finish(stream); } From e2ebb0d3c22574c27c7579feed7ff33ddfd9364f Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:05:53 -0700 Subject: [PATCH 073/115] feat: add bounded opt-in tensor-rank decode graph replay --- zig/src/families/flashnext/cuda_batch.zig | 13 +-- zig/src/families/flashnext/cuda_lanes.zig | 1 + zig/src/families/flashnext/cuda_launch.zig | 1 + zig/src/families/flashnext/cuda_prompt.zig | 13 ++- .../families/flashnext/cuda_rank_engine.zig | 8 +- .../families/flashnext/cuda_round_graph.zig | 98 +++++++++++++++++++ 6 files changed, 126 insertions(+), 8 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_round_graph.zig diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 28c927b8c4..f1b436898e 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -99,6 +99,7 @@ pub fn Pass(comptime Tri: type) type { pos0: usize = 0, /// The MoE write-back the next layer (or the head) applies first, and the layer it came from. pending: ?usize = null, + prepared: bool = false, // ---- launches ---------------------------------------------------------------------------------------- @@ -402,7 +403,7 @@ pub fn Pass(comptime Tri: type) type { pub fn stage(self: *Self, ids: []const u32) !void { const s = self.s; const n = ids.len; - try s.ready(); + if (!self.prepared) try s.ready(); const id32 = s.hostIds(n); for (id32, ids) |*o, v| o.* = @intCast(v); try self.upload(s.ids, std.mem.sliceAsBytes(id32)); @@ -412,7 +413,7 @@ pub fn Pass(comptime Tri: type) type { @memset(pos, 0); for (0..blocks) |b| pos[4 * b] = @intCast(if (self.lanes.len > 0) self.lanes[b].pos else self.pos0 + b * att_rows); try self.upload(s.pos, std.mem.sliceAsBytes(pos)); - try s.uploaded(); + if (!self.prepared) try s.uploaded(); } /// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. @@ -465,11 +466,11 @@ pub fn Pass(comptime Tri: type) type { self.pos0 = pos0; self.pending = null; // the gather writes the scratch's pinned rows, so the last pass's uploads must have read them first - try s.ready(); + if (!self.prepared) try s.ready(); const emb = s.hostEmb(n); const xs = s.hostXs(n); var gather: Gather = .{ .ngram = self.m.ngram, .history = self.seq.history[0..], .ids = ids, .emb = emb, .xs = xs, .io = self.io, .lanes = self.lanes }; - var thread: ?std.Thread = try std.Thread.spawn(.{}, Gather.main, .{&gather}); + var thread: ?std.Thread = if (self.prepared) null else try std.Thread.spawn(.{}, Gather.main, .{&gather}); defer if (thread) |th| th.join(); try self.stage(ids); if (self.m.table_dense) { @@ -480,13 +481,13 @@ pub fn Pass(comptime Tri: type) type { for (self.m.layers, 0..) |*lw, layer| { const keep: LayerKeep = if (self.hooks.keep) |kp| kp.layers[layer] else .{}; if (layer == ple_layer) { - thread.?.join(); + if (thread) |th| th.join(); thread = null; if (gather.failed) |err| return err; if (profile()) std.debug.print("batch: {d} rows' n-gram rows gathered on the host in {d:.1} ms\n", .{ n, @as(f64, @floatFromInt(gather.ns)) / 1e6 }); try self.upload(s.ple_emb, std.mem.sliceAsBytes(emb)); try self.upload(s.ple_xs, std.mem.sliceAsBytes(xs)); - try s.uploaded(); + if (!self.prepared) try s.uploaded(); if (self.pending != null) try self.writeback(0, n); try self.pleBlock(lw, ids); } diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 0d0cb0f42f..a8bdbd52a3 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -366,6 +366,7 @@ pub fn Namespace(comptime Tri: type) type { var it = own.lanes.valueIterator(); while (it.next()) |l| dropLane(own, l); own.lanes.deinit(own.gpa); + own.engine.round_graph.clear(); own.logits.free(); if (own.head) |h| h.deinit(); if (own.ranks) |r| r.deinit(); diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 178d2081c1..b770dda8eb 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -376,6 +376,7 @@ test { _ = @import("cuda_config.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); + _ = @import("cuda_round_graph.zig"); _ = @import("cuda_rank_memory.zig"); _ = @import("cuda_rank_wire.zig"); _ = @import("cuda_modelopt_experts.zig"); diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 71be96b5c7..9a3478aadc 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -110,6 +110,8 @@ pub fn pooledBytes(cap: usize) usize { return (cap + weights.index_ratio - 1) / weights.index_ratio * index_dim * 2; } +const round_graph = @import("cuda_round_graph.zig"); + /// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. pub const Engine = struct { gpa: std.mem.Allocator, @@ -133,6 +135,7 @@ pub const Engine = struct { scores: ?cuda.DeviceBuffer = null, /// Rows a batched chunk takes at most (TF_FLASHNEXT_CHUNK; batch.default_chunk by default). chunk: usize, + round_graph: round_graph.Cache = .{}, collective: ?*@import("cuda_collective.zig").Collective = null, pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { @@ -150,6 +153,7 @@ pub const Engine = struct { e.stream = stream; e.n_layers = 0; e.collective = null; + e.round_graph = .{}; errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); e.model_dir = try gpa.dupe(u8, model_dir); errdefer gpa.free(e.model_dir); @@ -204,6 +208,7 @@ pub const Engine = struct { } pub fn deinit(e: *Engine) void { + e.round_graph.clear(); if (e.batch_s) |*bs| bs.free(e.gpa); if (e.scores) |*b| b.free(); e.batch_k.deinit(); @@ -219,6 +224,7 @@ pub const Engine = struct { /// A fresh cache. `cap` is how many tokens it can hold, prompt plus generated. pub fn newSeq(e: *Engine, cap: usize) !Seq { + e.round_graph.clear(); if (cap == 0) return error.UnexpectedTensor; const logits = try e.gpa.alloc(u8, vocab * 2); errdefer e.gpa.free(logits); @@ -268,6 +274,7 @@ pub const Engine = struct { if (e.scores) |*b| { if (b.len >= need) return b.ptr; try e.complete(); + e.round_graph.clear(); b.free(); e.scores = null; } @@ -280,6 +287,7 @@ pub const Engine = struct { if (e.batch_s) |*bs| { if (bs.cap >= rows) return bs; try e.complete(); + e.round_graph.clear(); bs.free(e.gpa); e.batch_s = null; } @@ -343,7 +351,10 @@ pub const Engine = struct { .lanes = rows, .hooks = .{ .logits = .all, .logits_out = out }, }; - try p.chunk(ids, 0, null); + if (e.collective != null and round_graph.enabled()) { + const key = try round_graph.Key.init(seqs, s.ids, scores, out); + try e.round_graph.run(key, &p, ids); + } else try p.chunk(ids, 0, null); try e.complete(); const d = e.stream.d; for (seqs, 0..) |q, r| { diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index de549deff6..d0eff5c3dd 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -14,6 +14,7 @@ pub const Identity = extern struct { checkpoint: [32]u8, context: u64, chunk: u64, + decode_graph: u64 = 0, }; fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { @@ -35,7 +36,7 @@ pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usiz defer a.free(index); const program = try std.process.executablePathAlloc(io, a); defer a.free(program); - return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk }; + return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()) }; } pub const Session = struct { @@ -159,6 +160,7 @@ pub const Session = struct { } pub fn drop(s: *Session, q: *prompt.Seq) void { + s.engine.round_graph.clear(); if (!s.stopped) { const h = s.command(.{ .kind = .drop_seq, .object = q.rank_id }, &.{}) catch { s.poison(); @@ -282,6 +284,7 @@ pub const Session = struct { try s.engine.round(Tri, tri, seqs[0..n], ids[0..n], logits); }, .drop_seq => { + s.engine.round_graph.clear(); const entry = s.seqs.fetchRemove(h.object) orelse return error.UnknownRankSequence; entry.value.deinit(s.gpa); s.gpa.destroy(entry.value); @@ -357,6 +360,9 @@ test "rank identity binds checkpoint, config, index, program and execution setti other = base; other.chunk += 1; try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.decode_graph = 1; + try std.testing.expect(!sameIdentity(base, other)); } test "rank operations refuse invalid geometry and unsupported hooks before execution" { diff --git a/zig/src/families/flashnext/cuda_round_graph.zig b/zig/src/families/flashnext/cuda_round_graph.zig new file mode 100644 index 0000000000..e367e31dd9 --- /dev/null +++ b/zig/src/families/flashnext/cuda_round_graph.zig @@ -0,0 +1,98 @@ +//! A single bounded decode graph, recaptured when sequence owners, scratch or attention geometry change. +const std = @import("std"); +const cuda = @import("cuda"); +const dims = 2560; + +pub fn enabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_DECODE_GRAPH") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub const Key = struct { + count: usize, + owners: [32]u64 = @splat(0), + chunks: [32]u32 = @splat(0), + score_tiles: [32]u32 = @splat(0), + scratch: u64, + scores: u64, + output: u64, + + pub fn init(seqs: anytype, scratch: u64, scores: u64, output: u64) !Key { + if (seqs.len == 0 or seqs.len > 32) return error.UnexpectedTensor; + var key: Key = .{ .count = seqs.len, .scratch = scratch, .scores = scores, .output = output }; + for (seqs, 0..) |q, i| { + if (q.rank_id == 0) return error.UnsupportedRankOperation; + key.owners[i] = q.rank_id; + const end = q.fed + 1; + key.chunks[i] = @intCast((@min(end, 2051) + 511) / 512); + key.score_tiles[i] = if (end / 4 > 512) @intCast(((end + 3) / 4 + 63) / 64) else 0; + } + return key; + } +}; + +pub const Cache = struct { + exec: ?cuda.graph.Exec = null, + key: ?Key = null, + + pub fn clear(c: *Cache) void { + if (c.exec) |*e| e.deinit(); + c.* = .{}; + } + + /// CPU histories and pinned inputs advance once per round, outside stream capture and replay. + fn prepare(pass: anytype, ids: []const u32) !void { + const s = pass.s; + try s.ready(); + for (s.hostIds(ids.len), ids) |*out, token| out.* = @intCast(token); + const positions = s.hostPos(ids.len * 4); + @memset(positions, 0); + for (pass.lanes, 0..) |lane, i| { + positions[i * 4] = @intCast(lane.pos); + try pass.m.ngram.rows(lane.seq.history[0..], ids[i..][0..1], s.hostEmb(ids.len)[i * dims ..][0..dims], s.hostXs(ids.len)[i * (dims / 32) ..][0 .. dims / 32]); + } + } + + pub fn run(c: *Cache, key: Key, pass: anytype, ids: []const u32) !void { + try prepare(pass, ids); + if (c.key == null or !std.meta.eql(c.key.?, key)) { + c.clear(); + try cuda.graph.beginCapture(pass.stream, .thread_local); + pass.prepared = true; + pass.chunk(ids, 0, null) catch |err| { + if (cuda.graph.endCapture(pass.stream)) |g| { + var discarded = g; + discarded.deinit(); + } else |_| {} + return err; + }; + var graph = try cuda.graph.endCapture(pass.stream); + defer graph.deinit(); + c.exec = try graph.instantiate(); + errdefer c.clear(); + try c.exec.?.upload(pass.stream); + c.key = key; + } + try c.exec.?.launchOn(pass.stream); + } +}; + +test "decode graph keys invalidate on owner order, attention seams and score tile seams" { + const Seq = struct { rank_id: u64, fed: usize }; + var first: Seq = .{ .rank_id = 1, .fed = 10 }; + var second: Seq = .{ .rank_id = 2, .fed = 10 }; + const base = try Key.init(@as([]const *Seq, &.{ &first, &second }), 1, 2, 3); + try std.testing.expect(!std.meta.eql(base, try Key.init(@as([]const *Seq, &.{ &second, &first }), 1, 2, 3))); + first.fed = 511; + const before = try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3); + first.fed = 512; + try std.testing.expect(!std.meta.eql(before, try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3))); + first.fed = 2048; + const dense = try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3); + first.fed = 2051; + try std.testing.expect(!std.meta.eql(dense, try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3))); + first.fed = 2303; + const score = try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3); + first.fed = 2304; + try std.testing.expect(!std.meta.eql(score, try Key.init(@as([]const *Seq, &.{&first}), 1, 2, 3))); +} From 64bd35b439f1b6d88ffbe68bd32c4e51b6288893 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:25:31 -0700 Subject: [PATCH 074/115] Load native ModelOpt MTP projections and exact draft vocabulary rows --- zig/src/families/flashnext/cuda_mtp.zig | 36 +++-------------- zig/src/families/flashnext/cuda_mtp_load.zig | 40 +++++++++++++++++++ .../flashnext/cuda_projection_load.zig | 14 +++++++ 3 files changed, 60 insertions(+), 30 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_mtp_load.zig diff --git a/zig/src/families/flashnext/cuda_mtp.zig b/zig/src/families/flashnext/cuda_mtp.zig index 73cd614672..f548647ad0 100644 --- a/zig/src/families/flashnext/cuda_mtp.zig +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -159,10 +159,11 @@ pub const Head = struct { h.gpa = gpa; var store = try wt.Store.open(gpa, io, dir); defer store.deinit(); + try store.detectNorms(); var buf: [120]u8 = undefined; - const ne = try wt.asF32(gpa, try store.tensor("language_model.mtp.pre_fc_norm_embedding.weight")); + const ne = try store.normScale(gpa, try store.tensor("language_model.mtp.pre_fc_norm_embedding.weight")); defer gpa.free(ne); - const nh = try wt.asF32(gpa, try store.tensor("language_model.mtp.pre_fc_norm_hidden.weight")); + const nh = try store.normScale(gpa, try store.tensor("language_model.mtp.pre_fc_norm_hidden.weight")); defer gpa.free(nh); if (ne.len != dims or nh.len != wide) return error.UnexpectedTensor; h.norm_e = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(ne)); @@ -184,28 +185,7 @@ pub const Head = struct { errdefer h.layer.free(gpa); h.ids = try draftIds(gpa); errdefer gpa.free(h.ids); - { - const head_w = try wt.join3(gpa, &store, &buf, "language_model.lm_head", ""); - defer gpa.free(head_w.w); - defer gpa.free(head_w.s); - defer gpa.free(head_w.b); - const wrow = (dims / 8) * 4; - const grow = (dims / 32) * 2; - if (head_w.w.len != vocab * wrow or head_w.s.len != vocab * grow) return error.UnexpectedTensor; - const n = h.ids.len; - const w = try gpa.alloc(u8, n * wrow); - defer gpa.free(w); - const s = try gpa.alloc(u8, n * grow); - defer gpa.free(s); - const b = try gpa.alloc(u8, n * grow); - defer gpa.free(b); - for (h.ids, 0..) |id, j| { - @memcpy(w[j * wrow ..][0..wrow], head_w.w[id * wrow ..][0..wrow]); - @memcpy(s[j * grow ..][0..grow], head_w.s[id * grow ..][0..grow]); - @memcpy(b[j * grow ..][0..grow], head_w.b[id * grow ..][0..grow]); - } - h.draft = try wt.uploadQ4(gpa, driver, w, s, b, n, dims / 8); - } + h.draft = try @import("cuda_mtp_load.zig").draft(gpa, driver, &store, h.ids); errdefer h.draft.free(); h.module = try cuda.Module.load(driver, cuda.kernels.flashnext_mtp); errdefer h.module.unload(); @@ -214,12 +194,8 @@ pub const Head = struct { } fn fc(gpa: std.mem.Allocator, driver: *cuda.Driver, store: *wt.Store, buf: []u8, name: []const u8) !Q4 { - const t = try wt.join3(gpa, store, buf, name, ""); - defer gpa.free(t.w); - defer gpa.free(t.s); - defer gpa.free(t.b); - if (t.w.len != dims * (dims / 8) * 4) return error.UnexpectedTensor; - return wt.uploadQ4(gpa, driver, t.w, t.s, t.b, dims, dims / 8); + _ = buf; + return @import("cuda_projection_load.zig").load(gpa, driver, store, name, dims, dims); } pub fn deinit(h: *Head) void { diff --git a/zig/src/families/flashnext/cuda_mtp_load.zig b/zig/src/families/flashnext/cuda_mtp_load.zig new file mode 100644 index 0000000000..11bf42ccc3 --- /dev/null +++ b/zig/src/families/flashnext/cuda_mtp_load.zig @@ -0,0 +1,40 @@ +//! The draft vocabulary borrows exact rows of the checkpoint's dense or affine language head. +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("cuda_prompt_wt.zig"); +const projection = @import("cuda_projection.zig"); + +pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []const u32) !wt.Q4 { + const tensor = try store.tensor("language_model.lm_head.weight"); + if (tensor.dtype == .bf16) { + if (!tensor.is(.bf16, &.{ wt.vocab, wt.dims })) return error.UnexpectedTensor; + const rows = try a.alloc(u8, ids.len * wt.dims * 2); + defer a.free(rows); + for (ids, 0..) |id, i| { + if (id >= wt.vocab) return error.InvalidToken; + @memcpy(rows[i * wt.dims * 2 ..][0 .. wt.dims * 2], tensor.bytes[id * wt.dims * 2 ..][0 .. wt.dims * 2]); + } + return projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ ids.len, wt.dims, 1, 1, 1 }, .bytes = rows }); + } + var buf: [120]u8 = undefined; + const head = try wt.join3(a, store, &buf, "language_model.lm_head", ""); + defer a.free(head.w); + defer a.free(head.s); + defer a.free(head.b); + const wrow = wt.dims / 2; + const grow = wt.dims / 32 * 2; + if (head.w.len != wt.vocab * wrow or head.s.len != wt.vocab * grow) return error.UnexpectedTensor; + const w = try a.alloc(u8, ids.len * wrow); + defer a.free(w); + const s = try a.alloc(u8, ids.len * grow); + defer a.free(s); + const b = try a.alloc(u8, ids.len * grow); + defer a.free(b); + for (ids, 0..) |id, i| { + if (id >= wt.vocab) return error.InvalidToken; + @memcpy(w[i * wrow ..][0..wrow], head.w[id * wrow ..][0..wrow]); + @memcpy(s[i * grow ..][0..grow], head.s[id * grow ..][0..grow]); + @memcpy(b[i * grow ..][0..grow], head.b[id * grow ..][0..grow]); + } + return wt.uploadQ4(a, d, w, s, b, ids.len, wt.dims / 8); +} diff --git a/zig/src/families/flashnext/cuda_projection_load.zig b/zig/src/families/flashnext/cuda_projection_load.zig index 4fa63bb581..c0c37f1b79 100644 --- a/zig/src/families/flashnext/cuda_projection_load.zig +++ b/zig/src/families/flashnext/cuda_projection_load.zig @@ -104,6 +104,20 @@ pub fn loadStitched(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, pre if (value.n != n or value.k != k) return error.UnexpectedTensor; return uploadMicro(a, d, value); } + if (first.dtype == .bf16) { + const bytes = try a.alloc(u8, try std.math.mul(usize, try std.math.mul(usize, n, k), 2)); + defer a.free(bytes); + var at: usize = 0; + for (names) |name| { + const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, name })); + if (t.dtype != .bf16 or t.rank != 2 or t.dim(1) != k or t.bytes.len != t.dim(0) * k * 2 or t.bytes.len > bytes.len - at) + return error.UnexpectedTensor; + @memcpy(bytes[at..][0..t.bytes.len], t.bytes); + at += t.bytes.len; + } + if (at != bytes.len) return error.UnexpectedTensor; + return projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ n, k, 1, 1, 1 }, .bytes = bytes }); + } if (first.dtype != .u32) return error.UnsupportedQuantization; const planes = try wt.stitch(a, store, prefix, names); defer a.free(planes.w); From 1967f90acd58b4ad4041a4440b6b5fb1cc55082e Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:25:31 -0700 Subject: [PATCH 075/115] Mirror MTP verification windows and prefix commits across CUDA ranks --- zig/src/families/flashnext/cuda_lanes.zig | 44 +++++++--- zig/src/families/flashnext/cuda_mtp_ops.zig | 22 ++++- .../families/flashnext/cuda_rank_engine.zig | 79 +++++++++++++++++- zig/src/families/flashnext/cuda_rank_wire.zig | 2 +- zig/tests/cuda/flashnext_window_tests.zig | 80 +++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 6 files changed, 216 insertions(+), 14 deletions(-) create mode 100644 zig/tests/cuda/flashnext_window_tests.zig diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index a8bdbd52a3..bcb9a5bb73 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -80,8 +80,9 @@ pub fn Namespace(comptime Tri: type) type { if ((world != 1 and world != 2) or rank >= world) return error.BadRank; const mode = try config.mode(gpa, io, dir); if (mode == .modelopt) { - if (drafts) return error.ModelOptDraftsUnavailable; - return @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); + const main = try @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); + // Dense draft projections and vocabulary rows exceed the affine head's estimate. + return main + if (drafts and rank == 0) mtp.headBytes() + (1 << 30) else @as(u64, 0); } if (world != 1) return error.UnsupportedQuantization; const main = try weights.deviceBytes(gpa, io, dir); @@ -104,7 +105,7 @@ pub fn Namespace(comptime Tri: type) type { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", error.UnsupportedCheckpoint => "the checkpoint differs from the supported Flash Next CUDA layout; the log names its config field", - error.ModelOptDraftsUnavailable => "ModelOpt Flash Next serves without its MTP head so far: use --no-drafts", + error.TensorRankDraftSnapshotUnavailable => "tensor-rank MTP snapshots are unavailable: disable the prompt cache for this configuration", error.RankDisconnected, error.RankTimeout, error.RemoteRank, error.BadRankFrame => "a tensor-parallel rank failed; restart both ranks from the same checkpoint before retrying", error.LearnedPocket => "this folder holds Living Weights pockets, which Flash Next on CUDA does not serve yet", else => null, @@ -145,7 +146,27 @@ pub fn Namespace(comptime Tri: type) type { } fn ops(self: *Cuda) Ops { - return .{ .e = self.engine, .tri = self.tri(), .h = self.head.? }; + return .{ .e = self.engine, .tri = self.tri(), .h = self.head.?, .main = if (self.ranks != null) .{ + .ptr = self, + .forward = mainForward, + .verify = mainVerify, + .rollback = mainRollback, + } else null }; + } + + fn mainForward(ptr: *anyopaque, q: *prompt.Seq, ids: []const u32, hooks: prompt.batch.Hooks) anyerror!void { + const self = of(ptr); + return self.ranks.?.forward(Tri, self.tri(), q, ids, hooks); + } + + fn mainVerify(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) anyerror!void { + const self = of(ptr); + return self.ranks.?.verify(Tri, self.tri(), q, d, ids); + } + + fn mainRollback(ptr: *anyopaque, q: *prompt.Seq, d: *mtp.Draft, keep: usize) anyerror!void { + const self = of(ptr); + return self.ranks.?.keep(Tri, self.tri(), q, d, keep); } fn backend(self: *Cuda) be.Backend { @@ -282,7 +303,6 @@ pub fn Namespace(comptime Tri: type) type { if (o.segments != 1 or o.slide) return error.UnsupportedOption; if ((o.tp != 1 and o.tp != 2) or o.rank >= o.tp) return error.BadRank; const mode = try config.mode(gpa, io, dir); - if (mode == .modelopt and o.drafts) return error.ModelOptDraftsUnavailable; if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); @@ -297,7 +317,8 @@ pub fn Namespace(comptime Tri: type) type { errdefer if (own.set) |*x| x.deinit(); if (o.tp == 2) { const configured_chunk = envUsize("TF_FLASHNEXT_CHUNK", prompt.batch.default_chunk); - const id = try ranks.identity(gpa, io, dir, o.context, if (configured_chunk == 0) prompt.batch.default_chunk else configured_chunk); + var id = try ranks.identity(gpa, io, dir, o.context, if (configured_chunk == 0) prompt.batch.default_chunk else configured_chunk); + id.mtp = @intFromBool(o.drafts); own.ranks = try ranks.Session.init(gpa, io, o.master orelse "127.0.0.1", o.master_port, o.rank, id); } errdefer if (own.ranks) |r| { @@ -313,7 +334,7 @@ pub fn Namespace(comptime Tri: type) type { // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) _ = try own.engine.reserve(@min(o.context, own.engine.chunk)); // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) - if (o.drafts) { + if (o.drafts and (o.tp == 1 or o.rank == 0)) { own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); own.depth = @max(1, @min(mtp.max_chain, envUsize("TF_FLASHNEXT_DRAFTS", mtp.max_chain))); if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; @@ -480,7 +501,9 @@ pub fn Namespace(comptime Tri: type) type { const l = try self.lane(r.stream); const d = if (l.draft) |*x| x else return error.NoDraftHead; // the window's kept rows: rollback is idempotent, so a later keep finds it done - if (r.rows) |path| try self.ops().rollback(&l.seq, d, path.len) else d.settle(&l.seq); + if (r.rows) |path| try self.ops().rollback(&l.seq, d, path.len) else { + if (d.open) try self.ops().rollback(&l.seq, d, d.rows); + } if (r.position != l.seq.fed + 1) return error.PositionMismatch; // the window's rows past the cache's room cannot run: fewer drafts there const room = l.seq.cap - l.seq.fed - 1; @@ -534,7 +557,10 @@ pub fn Namespace(comptime Tri: type) type { const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); const l = try self.lane(s); if (l.seq.fed != at) return error.NotAtMark; - if (self.ranks) |r| return try r.save(&l.seq); + if (self.ranks) |r| { + if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; + return try r.save(&l.seq); + } return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); } diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 3c13f1b568..0e36c57869 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -54,15 +54,29 @@ fn download(stream: cuda.Stream, dst: []u8, src: u64) !void { try d.check(d.api.cuMemcpyDtoH_v2(dst.ptr, src, dst.len), "cuMemcpyDtoH"); } +/// An owner can mirror main-model mutations to a peer without exposing device pointers on the control wire. +pub const Main = struct { + ptr: *anyopaque, + forward: *const fn (*anyopaque, *prompt.Seq, []const u32, batch.Hooks) anyerror!void, + verify: *const fn (*anyopaque, *prompt.Seq, *Draft, []const u32) anyerror!void, + rollback: *const fn (*anyopaque, *prompt.Seq, *Draft, usize) anyerror!void, +}; + /// The head's work over a sequence: prompt pass, verify windows, rollback, forward and chains; `Tri` launches kernels. pub fn Ops(comptime Tri: type) type { return struct { e: *prompt.Engine, tri: Tri, h: *Head, + main: ?Main = null, const Self = @This(); + fn forwardMain(o: Self, seq: *prompt.Seq, ids: []const u32, hooks: batch.Hooks) !void { + if (o.main) |m| return m.forward(m.ptr, seq, ids, hooks); + return o.e.forward(Tri, o.tri, seq, ids, hooks); + } + fn stream(o: Self) cuda.Stream { return o.e.stream.*; } @@ -72,7 +86,7 @@ pub fn Ops(comptime Tri: type) type { if (stop > d.cap or stop <= at or stop - at > d.rows_cap) return error.PromptTooLong; if (seq.fed != at or d.len != at or d.drafted != 0) return error.PositionMismatch; const last = stop == ids.len; - try o.e.forward(Tri, o.tri, seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr }); + try o.forwardMain(seq, ids[at..stop], .{ .logits = if (last) .last else .none, .streams = d.streams.ptr }); const head_to = @min(stop, ids.len - 1); if (head_to > at) try o.forwardHead(d, ids[at + 1 .. head_to + 1], d.streams.ptr, false); if (last) d.prompt_last = d.streams.ptr + (stop - 1 - at) * wide * 2; @@ -108,7 +122,7 @@ pub fn Ops(comptime Tri: type) type { .tri = tri, .k = &e.batch_k, .s = s, - .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, + .m = .{ .table_w = e.table.w.ptr, .table_s = e.table.s.ptr, .table_b = e.table.b.ptr, .table_dense = e.table.dense, .layers = (&h.layer)[0..1], .mixer = h.mixer, .head = h.draft, .ngram = &e.ngram }, .seq = .{ .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, .hooks = .{ .logits = .none }, }; @@ -212,6 +226,8 @@ pub fn Ops(comptime Tri: type) type { /// One forward over the window `ids` from seq.fed: each row's scores in d.logits, its streams in d.streams. pub fn verify(o: Self, seq: *prompt.Seq, d: *Draft, ids: []const u32) !void { + if (o.main) |m| return m.verify(m.ptr, seq, d, ids); + o.e.round_graph.clear(); const n = ids.len; if (n == 0 or n > max_rows) return error.WindowTooWide; if (seq.fed + n > seq.cap) return error.PromptTooLong; @@ -244,6 +260,8 @@ pub fn Ops(comptime Tri: type) type { /// Keep the window's first `keep` rows (Python forward.commit) from its saved rows; a repeat is a no-op. pub fn rollback(o: Self, seq: *prompt.Seq, d: *Draft, keep: usize) !void { + if (o.main) |m| return m.rollback(m.ptr, seq, d, keep); + o.e.round_graph.clear(); if (keep == 0 or keep > d.rows) return error.CannotRewind; if (keep == d.rows) return d.settle(seq); if (!d.open) return error.CannotRewind; diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index d0eff5c3dd..b8852aa595 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -6,6 +6,7 @@ const snaps = @import("cuda_snap.zig"); const wire = @import("cuda_rank_wire.zig"); const Collective = @import("cuda_collective.zig").Collective; const wt = @import("cuda_prompt_wt.zig"); +const mtp = @import("cuda_mtp.zig"); pub const Identity = extern struct { config: [32]u8, @@ -15,6 +16,7 @@ pub const Identity = extern struct { context: u64, chunk: u64, decode_graph: u64 = 0, + mtp: u64 = 0, }; fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { @@ -49,6 +51,8 @@ pub const Session = struct { next_object: u64 = 1, seqs: std.AutoHashMapUnmanaged(u64, *prompt.Seq) = .empty, saved: std.AutoHashMapUnmanaged(u64, *snaps.Snap) = .empty, + windows: std.AutoHashMapUnmanaged(u64, mtp.Draft) = .empty, + drafting: bool = false, stopped: bool = false, available: std.atomic.Value(bool) = .init(true), @@ -74,7 +78,7 @@ pub const Session = struct { try link.send(std.mem.asBytes(&mine), until); } if (!sameIdentity(mine, peer)) return error.RankIdentity; - s.* = .{ .gpa = a, .rank = rank, .link = link, .collective = try Collective.init(link, rank, 30_000) }; + s.* = .{ .gpa = a, .rank = rank, .link = link, .collective = try Collective.init(link, rank, 30_000), .drafting = mine.mtp == 1 }; return s; } @@ -96,6 +100,9 @@ pub const Session = struct { const h = s.command(.{ .kind = .stop }, &.{}) catch null; if (h) |sent| s.ack(sent) catch s.poison(); } + var wi = s.windows.valueIterator(); + while (wi.next()) |d| d.deinit(); + s.windows.deinit(s.gpa); s.collective.deinit(); var qi = s.seqs.valueIterator(); while (qi.next()) |q| { @@ -173,7 +180,7 @@ pub const Session = struct { } pub fn forward(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, ids: []const u32, hooks: prompt.batch.Hooks) !void { - if (ids.len == 0 or ids.len > wire.max_payload / 4 or q.fed + ids.len > q.cap or hooks.keep != null or hooks.streams != 0 or hooks.logits == .all) return error.UnsupportedRankOperation; + if (ids.len == 0 or ids.len > wire.max_payload / 4 or q.fed + ids.len > q.cap or hooks.keep != null or hooks.logits == .all) return error.UnsupportedRankOperation; try tokenIds(ids); const h = try s.command(.{ .kind = .forward, .object = q.rank_id, .bytes = @intCast(ids.len * 4), .flags = @backingInt(hooks.logits) }, std.mem.sliceAsBytes(ids)); errdefer s.poison(); @@ -198,6 +205,36 @@ pub const Session = struct { try s.ack(h); } + /// Both ranks keep their own recurrent window; only the leader consumes draft scores and hidden rows. + pub fn verify(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, ids: []const u32) !void { + if (!s.drafting or ids.len == 0 or ids.len > mtp.max_rows or q.fed + ids.len > q.cap) return error.UnsupportedRankOperation; + try tokenIds(ids); + const h = try s.command(.{ .kind = .verify, .object = q.rank_id, .bytes = @intCast(ids.len * 4) }, std.mem.sliceAsBytes(ids)); + errdefer s.poison(); + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.verify(q, d, ids); + try s.ack(h); + } + + pub fn keep(s: *Session, comptime Tri: type, tri: Tri, q: *prompt.Seq, d: *mtp.Draft, count: usize) !void { + if (!s.drafting or count == 0 or count > d.rows) return error.CannotRewind; + const h = try s.command(.{ .kind = .keep, .object = q.rank_id, .arg = count }, &.{}); + errdefer s.poison(); + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.rollback(q, d, count); + try s.collective.finish(s.engine.stream.*); + try s.ack(h); + } + + fn window(s: *Session, q: *prompt.Seq) !*mtp.Draft { + const item = try s.windows.getOrPut(s.gpa, q.rank_id); + if (!item.found_existing) { + errdefer _ = s.windows.remove(q.rank_id); + item.value_ptr.* = try mtp.Draft.init(s.gpa, s.engine.driver, s.engine.stream.*, q.cap, 1, s.engine.chunk); + } + return item.value_ptr; + } + pub fn save(s: *Session, q: *prompt.Seq) !*snaps.Snap { const kept = try snaps.save(s.gpa, s.engine.driver, s.engine.stream.*, q, null); errdefer snaps.drop(s.gpa, s.engine.stream.*, kept); @@ -283,8 +320,30 @@ pub const Session = struct { } try s.engine.round(Tri, tri, seqs[0..n], ids[0..n], logits); }, + .verify => { + if (!s.drafting) return error.UnsupportedRankOperation; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + var ids: [mtp.max_rows]u32 = undefined; + const n = payload.len / 4; + for (ids[0..n], 0..) |*id, i| id.* = std.mem.readInt(u32, payload[i * 4 ..][0..4], .little); + try tokenIds(ids[0..n]); + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.verify(q, try s.window(q), ids[0..n]); + }, + .keep => { + if (!s.drafting) return error.UnsupportedRankOperation; + const q = s.seqs.get(h.object) orelse return error.UnknownRankSequence; + const d = s.windows.getPtr(h.object) orelse return error.CannotRewind; + const ops: mtp.Ops(Tri) = .{ .e = s.engine, .tri = tri, .h = undefined }; + try ops.rollback(q, d, @intCast(h.arg)); + try s.collective.finish(s.engine.stream.*); + }, .drop_seq => { s.engine.round_graph.clear(); + if (s.windows.fetchRemove(h.object)) |entry| { + var d = entry.value; + d.deinit(); + } const entry = s.seqs.fetchRemove(h.object) orelse return error.UnknownRankSequence; entry.value.deinit(s.gpa); s.gpa.destroy(entry.value); @@ -338,6 +397,8 @@ pub fn shape(h: wire.Header) !void { switch (h.kind) { .create => if (h.object == 0 or h.arg == 0 or h.arg > prompt.max_tokens or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .forward => if (h.object == 0 or h.arg != 0 or h.bytes == 0 or h.bytes % 4 != 0 or (h.flags != @backingInt(prompt.batch.Logits.last) and h.flags != @backingInt(prompt.batch.Logits.none))) return error.BadRankFrame, + .verify => if (h.object == 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 4 != 0 or h.bytes / 4 > mtp.max_rows) return error.BadRankFrame, + .keep => if (h.object == 0 or h.arg == 0 or h.arg > mtp.max_rows or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .round => if (h.object != 0 or h.arg != 0 or h.flags != 0 or h.bytes == 0 or h.bytes % 12 != 0 or h.bytes / 12 > 32) return error.BadRankFrame, .save, .restore => if (h.object == 0 or h.arg == 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, .drop_seq, .drop_snap => if (h.object == 0 or h.arg != 0 or h.bytes != 0 or h.flags != 0) return error.BadRankFrame, @@ -363,12 +424,26 @@ test "rank identity binds checkpoint, config, index, program and execution setti other = base; other.decode_graph = 1; try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.mtp = 1; + try std.testing.expect(!sameIdentity(base, other)); } test "rank operations refuse invalid geometry and unsupported hooks before execution" { try shape(.{ .kind = .create, .object = 1, .arg = 4096 }); try shape(.{ .kind = .forward, .object = 1, .bytes = 4 }); try shape(.{ .kind = .round, .bytes = 24 }); + try shape(.{ .kind = .verify, .object = 1, .bytes = mtp.max_rows * 4 }); + try shape(.{ .kind = .keep, .object = 1, .arg = mtp.max_rows }); + for ([_]wire.Header{ + .{ .kind = .verify, .object = 1, .bytes = (mtp.max_rows + 1) * 4 }, + .{ .kind = .verify, .object = 1, .bytes = 3 }, + .{ .kind = .verify, .bytes = 4 }, + .{ .kind = .verify, .object = 1, .bytes = 4, .flags = 1 }, + .{ .kind = .keep, .object = 1 }, + .{ .kind = .keep, .object = 1, .arg = mtp.max_rows + 1 }, + .{ .kind = .keep, .object = 1, .arg = 1, .bytes = 4 }, + }) |bad| try std.testing.expectError(error.BadRankFrame, shape(bad)); for ([_]wire.Header{ .{ .kind = .hello }, .{ .kind = .create, .arg = 4096 }, .{ .kind = .create, .object = 1, .arg = prompt.max_tokens + 1 }, .{ .kind = .forward, .object = 1, .bytes = 3 }, .{ .kind = .forward, .object = 1, .bytes = 4, .flags = 1 }, .{ .kind = .round, .bytes = 33 * 12 }, diff --git a/zig/src/families/flashnext/cuda_rank_wire.zig b/zig/src/families/flashnext/cuda_rank_wire.zig index e93f02c853..afad725e94 100644 --- a/zig/src/families/flashnext/cuda_rank_wire.zig +++ b/zig/src/families/flashnext/cuda_rank_wire.zig @@ -5,7 +5,7 @@ const p = std.posix; const net = std.Io.net; pub const max_payload = 1 << 20; -pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack }; +pub const Kind = enum(u32) { hello = 1, ready, create, forward, round, drop_seq, save, restore, drop_snap, stop, ack, verify, keep }; pub const Header = struct { kind: Kind, sequence: u64 = 0, diff --git a/zig/tests/cuda/flashnext_window_tests.zig b/zig/tests/cuda/flashnext_window_tests.zig new file mode 100644 index 0000000000..9e16c517a8 --- /dev/null +++ b/zig/tests/cuda/flashnext_window_tests.zig @@ -0,0 +1,80 @@ +//! A two-rank verification window agrees with serial rows and rolls back at partial and full acceptance seams. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const prompt = flash.prompt; +const mtp = flash.mtp; + +pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: []const u8) !void { + const context = 256; + const source = try std.Io.Dir.cwd().readFileAlloc(gpu.io, ids_path, gpu.gpa, .limited(1 << 20)); + defer gpu.gpa.free(source); + var ids: std.ArrayList(u32) = .empty; + defer ids.deinit(gpu.gpa); + var words = std.mem.tokenizeAny(u8, source, " \r\n\t"); + while (words.next()) |word| try ids.append(gpu.gpa, try std.fmt.parseInt(u32, word, 10)); + if (ids.items.len == 0 or ids.items.len > 128) return error.BadPrompt; + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + var identity = try flash.rank_engine.identity(gpu.gpa, gpu.io, dir, context, 32); + identity.mtp = 1; + const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const e = prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, dir, 2, rank) catch |err| { + session.poison(); + session.deinit(); + return err; + }; + session.engine = e; + e.collective = &session.collective; + defer { + session.deinit(); + e.deinit(); + } + if (e.chunk != 32) return error.RankIdentity; + var logits = try cuda.DeviceBuffer.alloc(gpu.d, 2 * prompt.vocabulary * 2); + defer logits.free(); + if (rank == 1) { + try session.follow(flash.Tri, tri, logits.ptr); + check.pass("rank 1 verification windows and prefix commits completed", .{}); + return; + } + var serial = try session.create(context); + defer session.drop(&serial); + try session.forward(flash.Tri, tri, &serial, ids.items, .{}); + const base = try session.save(&serial); + defer session.dropSaved(base); + const forced = [_]u32{ 37, 44, 102, 1, 320, 17, 19, 71 }; + const row_bytes = prompt.vocabulary * 2; + const expected = try gpu.gpa.alloc(u8, forced.len * row_bytes); + defer gpu.gpa.free(expected); + for (forced, 0..) |token, i| { + try session.forward(flash.Tri, tri, &serial, &.{token}, .{}); + @memcpy(expected[i * row_bytes ..][0..row_bytes], serial.logits); + } + var window = try session.create(context); + defer session.drop(&window); + var draft = try mtp.Draft.init(gpu.gpa, e.driver, stream, context, 1, 32); + defer draft.deinit(); + const rows = try gpu.gpa.alloc(u8, expected.len); + defer gpu.gpa.free(rows); + for ([_]usize{ 1, 3, 7, 8 }) |keep| { + try session.restore(&window, base); + try session.verify(flash.Tri, tri, &window, &draft, &forced); + try stream.synchronize(); + try e.driver.check(e.driver.api.cuMemcpyDtoH_v2(rows.ptr, draft.logits.ptr, rows.len), "cuMemcpyDtoH"); + try check.sameBytes("each two-rank window row equals its serial teacher-forced row", rows, expected); + try session.keep(flash.Tri, tri, &window, &draft, keep); + try check.expect(window.fed == ids.items.len + keep, "kept prefix position is exact", .{}); + try session.restore(&serial, base); + for (forced[0..keep]) |token| try session.forward(flash.Tri, tri, &serial, &.{token}, .{}); + try check.expect(std.mem.eql(i64, &window.history, &serial.history), "kept host n-gram history matches serial", .{}); + try session.forward(flash.Tri, tri, &serial, &.{13}, .{}); + try session.forward(flash.Tri, tri, &window, &.{13}, .{}); + try check.sameBytes("rejected suffix cannot change the next row after prefix commit", window.logits, serial.logits); + check.pass("two-rank forced window: all rows and next row byte-exact at keep={d}/{d}", .{ keep, forced.len }); + } +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index a3e6bf56c2..3219eb0c47 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -34,6 +34,7 @@ const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); const flashnext_model_rank_tests = @import("flashnext_model_rank_tests.zig"); +const flashnext_window_tests = @import("flashnext_window_tests.zig"); const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); @@ -63,6 +64,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-window-rank MODEL RANK PORT IDS_FILE two-rank serial/window and rollback parity \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 @@ -157,6 +159,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum"); if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); From 69e335e75e3920d6142ad24c1f4c1b2facc896aa Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:54:51 -0700 Subject: [PATCH 076/115] Add opt-in affine proposal head with independent arithmetic checks --- .../families/flashnext/cuda_draft_affine.zig | 91 +++++++++++++++++++ zig/src/families/flashnext/cuda_lanes.zig | 4 +- zig/src/families/flashnext/cuda_launch.zig | 2 + zig/src/families/flashnext/cuda_mtp_load.zig | 15 +++ .../cuda/flashnext_draft_affine_tests.zig | 61 +++++++++++++ zig/tests/cuda/main.zig | 3 + 6 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/flashnext/cuda_draft_affine.zig create mode 100644 zig/tests/cuda/flashnext_draft_affine_tests.zig diff --git a/zig/src/families/flashnext/cuda_draft_affine.zig b/zig/src/families/flashnext/cuda_draft_affine.zig new file mode 100644 index 0000000000..c6efcf823d --- /dev/null +++ b/zig/src/families/flashnext/cuda_draft_affine.zig @@ -0,0 +1,91 @@ +//! An optional group-32 affine proposal head; the main checkpoint and its verifier remain unchanged. +const std = @import("std"); + +pub fn bf16(v: f32) u16 { + const bits: u32 = @bitCast(v); + return @truncate((bits +% 0x7fff +% ((bits >> 16) & 1)) >> 16); +} + +fn value(bytes: []const u8, i: usize) f32 { + return @bitCast(@as(u32, std.mem.readInt(u16, bytes[i * 2 ..][0..2], .little)) << 16); +} + +fn code(x: f32) u8 { + const bounded = std.math.clamp(x, 0, 15); + const lo = @floor(bounded); + const integer: u8 = @intFromFloat(lo); + const fraction = bounded - lo; + return integer + @intFromBool(fraction > 0.5 or (fraction == 0.5 and integer & 1 == 1)); +} + +/// Codes use the FP32 range before its BF16 table rounding, matching the Python proposal-head recipe. +pub fn row(input: []const u8, words: []u8, scales: []u8, biases: []u8) !void { + const n = input.len / 2; + if (n == 0 or n % 32 != 0 or input.len != n * 2 or words.len != n / 2 or scales.len != n / 16 or biases.len != scales.len) + return error.UnexpectedTensor; + for (0..n / 32) |group| { + const at = group * 32; + var lo = std.math.inf(f32); + var hi = -std.math.inf(f32); + for (0..32) |i| { + const v = value(input, at + i); + if (!std.math.isFinite(v)) return error.UnexpectedTensor; + lo = @min(lo, v); + hi = @max(hi, v); + } + const scale = @max((hi - lo) / 15.0, 1e-8); + if (!std.math.isFinite(scale)) return error.UnexpectedTensor; + std.mem.writeInt(u16, scales[group * 2 ..][0..2], bf16(scale), .little); + std.mem.writeInt(u16, biases[group * 2 ..][0..2], bf16(lo), .little); + for (0..16) |j| { + const low = code((value(input, at + j * 2) - lo) / scale); + const high = code((value(input, at + j * 2 + 1) - lo) / scale); + words[group * 16 + j] = low | (high << 4); + } + } +} + +test "proposal codes have known nearest-even ties and a flat row stays on its exact bias" { + var input: [128]u8 = undefined; + for (0..64) |i| std.mem.writeInt(u16, input[i * 2 ..][0..2], bf16(if (i < 32) @as(f32, @floatFromInt(@min(i, 30))) / 2 else 7), .little); + var words: [32]u8 = undefined; + var scales: [4]u8 = undefined; + var biases: [4]u8 = undefined; + try row(&input, &words, &scales, &biases); + try std.testing.expectEqual(@as(u8, 0x00), words[0]); + try std.testing.expectEqual(@as(u8, 0x21), words[1]); + try std.testing.expectEqual(@as(u8, 0x22), words[2]); + try std.testing.expectEqual(@as(u16, 0x3f80), std.mem.readInt(u16, scales[0..2], .little)); + try std.testing.expectEqual(@as(u16, 0x40e0), std.mem.readInt(u16, biases[2..4], .little)); + for (words[16..]) |b| try std.testing.expectEqual(@as(u8, 0), b); + std.mem.writeInt(u16, input[0..2], 0x7f80, .little); + try std.testing.expectError(error.UnexpectedTensor, row(&input, &words, &scales, &biases)); +} + +test "proposal approximation is bounded per value against an independent FP64 dequantizer" { + var input: [128]u8 = undefined; + for (0..64) |i| std.mem.writeInt(u16, input[i * 2 ..][0..2], bf16(@sin(@as(f32, @floatFromInt(i)) * 0.37) * 0.05), .little); + var words: [32]u8 = undefined; + var scales: [4]u8 = undefined; + var biases: [4]u8 = undefined; + try row(&input, &words, &scales, &biases); + for (0..2) |g| { + var lo = std.math.inf(f64); + var hi = -std.math.inf(f64); + for (0..32) |i| { + const v: f64 = value(&input, g * 32 + i); + lo = @min(lo, v); + hi = @max(hi, v); + } + const scale = (hi - lo) / 15; + const stored_s: f64 = value(&scales, g); + const stored_b: f64 = value(&biases, g); + const bound = scale * 0.50001 + 15 * @abs(stored_s - scale) + @abs(stored_b - lo); + for (0..32) |i| { + const at = g * 32 + i; + const q = (words[at / 2] >> @as(u3, @intCast((at % 2) * 4))) & 15; + const back = @as(f64, @floatFromInt(q)) * stored_s + stored_b; + try std.testing.expect(@abs(back - value(&input, at)) <= bound); + } + } +} diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index bcb9a5bb73..869a72712a 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -348,7 +348,9 @@ pub fn Namespace(comptime Tri: type) type { own.lanes = .empty; own.next = 0; var stream_bytes = prompt.seqBytes(window); - if (own.head) |h| stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); + if (own.head) |h| { + stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); + } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); return .{ .backend = own.backend(), .facts = own.facts(), diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index b770dda8eb..0683c45bfd 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -11,6 +11,7 @@ pub const collective = @import("cuda_collective.zig"); pub const rank_wire = @import("cuda_rank_wire.zig"); pub const modelopt_experts = @import("cuda_modelopt_experts.zig"); pub const hc = @import("cuda_hc.zig"); +pub const draft_affine = @import("cuda_draft_affine.zig"); pub const qmm = @import("cuda_qmm.zig"); pub const gdn = @import("cuda_gdn.zig"); pub const mlp = @import("cuda_mlp.zig"); @@ -382,6 +383,7 @@ test { _ = @import("cuda_modelopt_experts.zig"); _ = row_path; _ = @import("cuda_mtp_check.zig"); + _ = @import("cuda_draft_affine.zig"); } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_mtp_load.zig b/zig/src/families/flashnext/cuda_mtp_load.zig index 11bf42ccc3..dc1a226d31 100644 --- a/zig/src/families/flashnext/cuda_mtp_load.zig +++ b/zig/src/families/flashnext/cuda_mtp_load.zig @@ -5,6 +5,9 @@ const wt = @import("cuda_prompt_wt.zig"); const projection = @import("cuda_projection.zig"); pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []const u32) !wt.Q4 { + const mode = if (std.c.getenv("TF_FLASHNEXT_DRAFT_HEAD")) |v| std.mem.span(v) else "bf16"; + const affine = std.mem.eql(u8, mode, "affine4"); + if (!affine and !std.mem.eql(u8, mode, "bf16")) return error.BadDraftHead; const tensor = try store.tensor("language_model.lm_head.weight"); if (tensor.dtype == .bf16) { if (!tensor.is(.bf16, &.{ wt.vocab, wt.dims })) return error.UnexpectedTensor; @@ -14,6 +17,7 @@ pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []con if (id >= wt.vocab) return error.InvalidToken; @memcpy(rows[i * wt.dims * 2 ..][0 .. wt.dims * 2], tensor.bytes[id * wt.dims * 2 ..][0 .. wt.dims * 2]); } + if (affine) return quantized(a, d, rows, ids.len); return projection.dense(a, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ ids.len, wt.dims, 1, 1, 1 }, .bytes = rows }); } var buf: [120]u8 = undefined; @@ -38,3 +42,14 @@ pub fn draft(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, ids: []con } return wt.uploadQ4(a, d, w, s, b, ids.len, wt.dims / 8); } + +fn quantized(a: std.mem.Allocator, d: *cuda.Driver, rows: []const u8, n: usize) !wt.Q4 { + const words = try a.alloc(u8, n * wt.dims / 2); + defer a.free(words); + const scales = try a.alloc(u8, n * wt.dims / 16); + defer a.free(scales); + const biases = try a.alloc(u8, scales.len); + defer a.free(biases); + for (0..n) |i| try @import("cuda_draft_affine.zig").row(rows[i * wt.dims * 2 ..][0 .. wt.dims * 2], words[i * wt.dims / 2 ..][0 .. wt.dims / 2], scales[i * wt.dims / 16 ..][0 .. wt.dims / 16], biases[i * wt.dims / 16 ..][0 .. wt.dims / 16]); + return wt.uploadQ4(a, d, words, scales, biases, n, wt.dims / 8); +} diff --git a/zig/tests/cuda/flashnext_draft_affine_tests.zig b/zig/tests/cuda/flashnext_draft_affine_tests.zig new file mode 100644 index 0000000000..a160eb8e41 --- /dev/null +++ b/zig/tests/cuda/flashnext_draft_affine_tests.zig @@ -0,0 +1,61 @@ +//! The proposal projection's quantizer, packing and GPU arithmetic against an independent FP64 reference. +const std = @import("std"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const rig = @import("flashnext_rig.zig"); +const k = 2560; +const n = 137; +const most = 7; + +fn half(bytes: []const u8, i: usize) f64 { + return rig.bf(std.mem.readInt(u16, bytes[i * 2 ..][0..2], .little)); +} + +pub fn run(gpu: check.Gpu) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var rng = std.Random.DefaultPrng.init(0xaff14e); + var r = try rig.Rig.init(gpu, a, rng.random()); + defer r.deinit(); + const source = r.bfs(n * k, 0.125); + const words = try a.alloc(u8, n * k / 2); + const scales = try a.alloc(u8, n * k / 16); + const biases = try a.alloc(u8, scales.len); + for (0..n) |i| try flash.draft_affine.row(std.mem.sliceAsBytes(source[i * k ..][0..k]), words[i * k / 2 ..][0 .. k / 2], scales[i * k / 16 ..][0 .. k / 16], biases[i * k / 16 ..][0 .. k / 16]); + var layout = try flash.qmm.pack(a, words, scales, biases, n, k / 8); + defer layout.deinit(a); + const dw = try r.dev(layout.weight); + const ds = try r.dev(layout.scales); + const db = try r.dev(layout.biases); + const x = r.bfs(most * k, 0.25); + const sums = try a.alloc(f32, most * k / 32); + for (sums, 0..) |*v, g| { + v.* = 0; + for (x[g * 32 ..][0..32]) |h| v.* += @floatCast(rig.bf(h)); + } + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const dsum = try r.dev(std.mem.sliceAsBytes(sums)); + const out = try r.out(most * n * 2); + var widest: []u16 = undefined; + for ([_]usize{ most, 1, 3 }) |rows| { + try flash.qmm.matmul(gpu.d, r.s, dx, dsum, dw, ds, db, out, rows, n, k); + if (rows == most) widest = try r.back(u16, out, most * n) else try r.sameRows("affine proposal rows", u16, out, widest[0 .. rows * n]); + } + for (0..most) |row| for (0..n) |col| { + var total: f64 = 0; + var magnitude: f64 = 0; + for (0..k) |i| { + const at = col * k + i; + const code = (words[at / 2] >> @as(u3, @intCast((i % 2) * 4))) & 15; + const g = col * (k / 32) + i / 32; + const weight = @as(f64, @floatFromInt(code)) * half(scales, g) + half(biases, g); + const term = rig.bf(x[row * k + i]) * weight; + total += term; + magnitude += @abs(term); + } + const want = rig.rbf(total); + try r.within("affine proposal FP64", rig.bf(widest[row * n + col]), want, rig.ulps(want, 1) + rig.sumTol(magnitude, k)); + }; + r.report("group-32 proposal quantization and affine projection, every column; rows 1/3/7 equal"); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 3219eb0c47..3063654595 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -36,6 +36,7 @@ const route_rows_tests = @import("route_rows_tests.zig"); const flashnext_model_rank_tests = @import("flashnext_model_rank_tests.zig"); const flashnext_window_tests = @import("flashnext_window_tests.zig"); const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); +const flashnext_draft_affine_tests = @import("flashnext_draft_affine_tests.zig"); const flashnext_shared_tests = @import("flashnext_shared_tests.zig"); const flashnext_nvfp4_tp_tests = @import("flashnext_nvfp4_tp_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); @@ -67,6 +68,7 @@ const usage = \\ flashnext-window-rank MODEL RANK PORT IDS_FILE two-rank serial/window and rollback parity \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes + \\ flashnext-draft-affine proposal-head affine projection against FP64 and row-count invariance \\ flashnext-shared separate MXFP8 shared expert and reserved-slot scatter against FP64 \\ flashnext-nvfp4-tp rank-local expert cuts and FP32 partials against FP64 (one GPU) \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 @@ -162,6 +164,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum"); + if (std.mem.eql(u8, cmd, "flashnext-draft-affine")) return flashnext_draft_affine_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-shared")) return flashnext_shared_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-nvfp4-tp")) return flashnext_nvfp4_tp_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); From a7af69b331ffe9b3e7dd0b67e2e1c7f38d4ceee2 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:58:03 -0700 Subject: [PATCH 077/115] Allow bounded draft chains through the full verification window --- zig/src/families/flashnext/cuda_lanes.zig | 4 ++-- zig/src/families/flashnext/cuda_mtp.zig | 12 ++++++++++-- zig/src/families/flashnext/cuda_mtp_check.zig | 2 +- zig/tests/cuda/flashnext_window_tests.zig | 4 ++-- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 869a72712a..059284e9d9 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -129,7 +129,7 @@ pub fn Namespace(comptime Tri: type) type { engine: *prompt.Engine, ranks: ?*ranks.Session = null, head: ?*mtp.Head = null, - depth: usize = mtp.max_chain, + depth: usize = mtp.default_depth, confidence: f64 = mtp.default_confidence, /// A captured Triton set (TENSORFOLD_CUDA_KERNELS) for comparison, else null: our own glue. set: ?cuda.aot.Set = null, @@ -336,7 +336,7 @@ pub fn Namespace(comptime Tri: type) type { // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) if (o.drafts and (o.tp == 1 or o.rank == 0)) { own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); - own.depth = @max(1, @min(mtp.max_chain, envUsize("TF_FLASHNEXT_DRAFTS", mtp.max_chain))); + own.depth = @max(1, @min(mtp.max_chain, envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth))); if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; } errdefer if (own.head) |h| h.deinit(); diff --git a/zig/src/families/flashnext/cuda_mtp.zig b/zig/src/families/flashnext/cuda_mtp.zig index f548647ad0..8c3d03e9ac 100644 --- a/zig/src/families/flashnext/cuda_mtp.zig +++ b/zig/src/families/flashnext/cuda_mtp.zig @@ -27,8 +27,10 @@ const HcW = wt.HcW; pub const Sampling = lanes.sampling.Sampling; -/// Drafts a round at most (Python DEPTH). -pub const max_chain: usize = 6; +/// Drafts a verify window can hold beside its pending token. +pub const max_chain: usize = max_rows - 1; +/// Preserve the existing recipe unless an explicit depth is selected. +pub const default_depth: usize = 6; /// The live recipe's --mtp-confidence (Python's CONFIDENCE default is 0.7). pub const default_confidence: f64 = 0.6; /// A verify window's rows at most: the pending token and a chain, or a copied continuation. @@ -267,3 +269,9 @@ test "a keyed draw over the top-k columns is the draw over every column" { } try std.testing.expectEqual(@as(u32, @intCast(argmaxRow(&row))), try drawRow(gpa, &row, null, 0, null)); } + +test "draft chains fit one verify window while the default stays six" { + try std.testing.expectEqual(@as(usize, 16), max_chain + 1); + try std.testing.expectEqual(@as(usize, 6), default_depth); + try std.testing.expect(default_depth <= max_chain); +} diff --git a/zig/src/families/flashnext/cuda_mtp_check.zig b/zig/src/families/flashnext/cuda_mtp_check.zig index aea0f37daf..435cf6c9da 100644 --- a/zig/src/families/flashnext/cuda_mtp_check.zig +++ b/zig/src/families/flashnext/cuda_mtp_check.zig @@ -102,7 +102,7 @@ pub fn check(comptime Tri: type, gpa: std.mem.Allocator, io: std.Io, driver: *cu const h = try Head.init(gpa, io, driver, stream, model_dir); defer h.deinit(); const o: Ops(Tri) = .{ .e = e, .tri = tri, .h = h }; - const depth = envUsize("TF_FLASHNEXT_DRAFTS", max_chain); + const depth = envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth); const confidence = envF64("TF_FLASHNEXT_CONFIDENCE", default_confidence); const sm: ?Sampling = if (std.c.getenv("TF_FLASHNEXT_SEED")) |seed| .{ .seed = std.fmt.parseInt(u64, std.mem.span(seed), 10) catch 0, diff --git a/zig/tests/cuda/flashnext_window_tests.zig b/zig/tests/cuda/flashnext_window_tests.zig index 9e16c517a8..999d773cd0 100644 --- a/zig/tests/cuda/flashnext_window_tests.zig +++ b/zig/tests/cuda/flashnext_window_tests.zig @@ -47,7 +47,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] try session.forward(flash.Tri, tri, &serial, ids.items, .{}); const base = try session.save(&serial); defer session.dropSaved(base); - const forced = [_]u32{ 37, 44, 102, 1, 320, 17, 19, 71 }; + const forced = [_]u32{ 37, 44, 102, 1, 320, 17, 19, 71, 99, 4, 80, 122, 7, 25, 61, 94 }; const row_bytes = prompt.vocabulary * 2; const expected = try gpu.gpa.alloc(u8, forced.len * row_bytes); defer gpu.gpa.free(expected); @@ -61,7 +61,7 @@ pub fn run(gpu: check.Gpu, dir: []const u8, rank: usize, port: u16, ids_path: [] defer draft.deinit(); const rows = try gpu.gpa.alloc(u8, expected.len); defer gpu.gpa.free(rows); - for ([_]usize{ 1, 3, 7, 8 }) |keep| { + for ([_]usize{ 1, 3, 7, 15, 16 }) |keep| { try session.restore(&window, base); try session.verify(flash.Tri, tri, &window, &draft, &forced); try stream.synchronize(); From c325d764acb19004bbac0e6f1584fd10d84bf0f8 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:01:15 -0700 Subject: [PATCH 078/115] Expose bounded NVFP4 expert CTA geometry without changing tile arithmetic --- zig/kernels/cuda/nvfp4_experts.cu | 6 ++ zig/src/cuda/experts.zig | 40 ++++++- zig/tests/cuda/flashnext_nvfp4_tp_tests.zig | 110 ++++++++++---------- 3 files changed, 98 insertions(+), 58 deletions(-) diff --git a/zig/kernels/cuda/nvfp4_experts.cu b/zig/kernels/cuda/nvfp4_experts.cu index 0a92b042f4..b80e8f1f6d 100644 --- a/zig/kernels/cuda/nvfp4_experts.cu +++ b/zig/kernels/cuda/nvfp4_experts.cu @@ -314,6 +314,12 @@ template __global__ void nvfp4_expert_kernel<2, 2, 4>(const __nv_bfloat16* __res template __global__ void nvfp4_expert_kernel<1, 0, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 3, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 1, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); + +// One warp per CTA exposes more independent expert tiles for narrow rank-local projections. +template __global__ void nvfp4_expert_kernel<2, 2, 1>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void nvfp4_expert_kernel<1, 0, 1>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void nvfp4_expert_kernel<1, 3, 1>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void nvfp4_expert_kernel<1, 1, 1>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); } // The prompt form's entries, 4 row tiles a warp in one row of WN warps: relu^2 up, SwiGLU gate-up, bf16 sums down. diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 16b641ef3a..6dc0ddfd7e 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -305,6 +305,7 @@ pub const Nvfp4Experts = struct { fns: [4]Function, // by Epi, then relu^2 up resident: [4]usize, // blocks the grid may hold: per SM times SMs prompt_fns: [3]Function, // relu^2 up, SwiGLU gate-up, bf16 sums down + warps: usize = 4, pub const cols = 32; // output columns a block pub const words = 144; // int32 a (32 columns, 32 inputs) block: 128 code words, then 16 of e4m3 scales @@ -327,14 +328,36 @@ pub const Nvfp4Experts = struct { return stages * stage; } + pub fn decodeWarps(value: ?[]const u8) !usize { + const v = value orelse return 4; + if (std.mem.eql(u8, v, "1")) return 1; + if (std.mem.eql(u8, v, "4")) return 4; + return error.BadExpertWarps; + } + pub fn load(d: *const Driver, sms: usize) !Nvfp4Experts { + const value = if (std.c.getenv("TF_NVFP4_DECODE_WARPS")) |v| std.mem.span(v) else null; + return loadWarps(d, sms, try decodeWarps(value)); + } + + /// Every warp owns independent tiles; changing CTA grouping preserves each tile's arithmetic. + pub fn loadWarps(d: *const Driver, sms: usize, warps: usize) !Nvfp4Experts { + if (warps != 1 and warps != 4) return error.BadExpertWarps; if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.nvfp4_experts); errdefer mod.unload(); - var e: Nvfp4Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .prompt_fns = undefined }; + var e: Nvfp4Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .prompt_fns = undefined, .warps = warps }; for ([_][:0]const u8{ symbols.gate_up, symbols.down_f32, symbols.down_bf16, symbols.up_relu2 }, 0..) |name, i| { - e.fns[i] = try mod.function(name); - e.resident[i] = @max(1, try e.fns[i].occupancy(128, 0)) * sms; + var narrow: [256]u8 = undefined; + const selected = if (warps == 1) blk: { + const marker = std.mem.indexOf(u8, name, "ELi4E") orelse return error.MissingSymbol; + @memcpy(narrow[0..name.len], name); + narrow[marker + 3] = '1'; + narrow[name.len] = 0; + break :blk narrow[0..name.len :0]; + } else name; + e.fns[i] = try mod.function(selected); + e.resident[i] = @max(1, try e.fns[i].occupancy(@intCast(warps * 32), 0)) * sms; } for (symbols.prompt, &e.prompt_fns, [3]usize{ 1, 2, 1 }, prompt_wn) |name, *f, parts, wn| { f.* = try mod.function(name); @@ -353,7 +376,7 @@ pub const Nvfp4Experts = struct { const kg = (if (gate) l.dims else l.width) / 32; const nb = (if (gate) l.width else l.dims) / cols; const ei: usize = if (gate and l.relu2) 3 else @backingInt(epi); - const grid = @min((items * nb + 3) / 4, e.resident[ei]); + const grid = @min((items * nb + e.warps - 1) / e.warps, e.resident[ei]); if (grid < 1) return; var a: launch_.Args = .{}; a.add(in.x); @@ -367,7 +390,7 @@ pub const Nvfp4Experts = struct { a.add(@as(i32, @intCast(if (gate) l.width else l.dims))); a.add(if (gate) l.limit else @as(f32, 0.0)); a.add(skip); - try launch_.launch(e.fns[ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &a); + try launch_.launch(e.fns[ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = @intCast(e.warps * 32), .y = 1, .z = 1 } }, s, &a); } /// Prompt rows on a plan of at most `items` items of 64 pairs: up (relu^2 or SwiGLU by the layer), bf16 sums down. @@ -510,3 +533,10 @@ test "a format that is not loaded, or an option it was not built with, is refuse try std.testing.expectError(error.EpilogueNotBuilt, h.prompt(s, .down_f32, .{ .x = 0, .stride = 64 }, la, p, 0, 1, -1)); try std.testing.expectError(error.OptionNotBuilt, h.decode(s, .up, .{ .x = 0, .stride = 64 }, la, p, 0, 1, 3)); } + +test "NVFP4 decode CTA geometry is explicit, bounded and defaults to four warps" { + try std.testing.expectEqual(@as(usize, 4), try Nvfp4Experts.decodeWarps(null)); + try std.testing.expectEqual(@as(usize, 1), try Nvfp4Experts.decodeWarps("1")); + try std.testing.expectEqual(@as(usize, 4), try Nvfp4Experts.decodeWarps("4")); + for ([_][]const u8{ "0", "2", "8", "", "garbage" }) |v| try std.testing.expectError(error.BadExpertWarps, Nvfp4Experts.decodeWarps(v)); +} diff --git a/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig index 2feac586ef..7f9b81f1b4 100644 --- a/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig +++ b/zig/tests/cuda/flashnext_nvfp4_tp_tests.zig @@ -101,61 +101,65 @@ pub fn run(gpu: check.Gpu) !void { var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); defer mod.unload(); const router = try cuda.grouped.Router.resolve(mod); - var kernels = try fx.load(gpu.d, @intCast(try gpu.ctx.attribute(.multiprocessor_count))); - defer kernels.unload(); - const act = try r.out(most * slots * rank_width * 2); - const out = try r.out(most * slots * d * 4); - const combined = try a.alloc(f32, most * slots * d); - @memset(combined, 0); - const reference_full = try a.alloc(f64, combined.len); - @memset(reference_full, 0); - const reference_abs = try a.alloc(f64, combined.len); - @memset(reference_abs, 0); - for ([_]*const flash.modelopt_experts.Owned{ &left, &right }, 0..) |owned, rank| { - var all: []f32 = undefined; - for ([_]usize{ most, 1, 7 }) |rows| { - try gpu.d.check(gpu.d.api.cuMemsetD8Async(out, 0, most * slots * d * 4, r.s.handle), "cuMemsetD8Async"); - const items = cuda.grouped.maxItems(rows * slots, experts + 1, 16); - try router.route(r.s, dp, rows * slots, experts + 1, 16, plan); - try kernels.run(r.s, .up, .{ .x = dx, .stride = d, .slots = slots }, owned.view(), plan, act, items, experts); - try kernels.run(r.s, .down_f32, .{ .x = act, .stride = rank_width }, owned.view(), plan, out, items, experts); - if (rows == most) all = try r.back(f32, out, most * slots * d) else try r.sameRows("NVFP4 rank output rows", f32, out, all[0 .. rows * slots * d]); - } - for (0..most) |row| for (0..slots) |slot| { - const e: usize = @intCast(picks[row * slots + slot]); - if (e == experts) { - for (all[(row * slots + slot) * d ..][0..d]) |v| try check.expect(v == 0, "shared slot untouched by NVFP4", .{}); - continue; + var prior: [2]?[]f32 = .{ null, null }; + for ([_]usize{ 4, 1 }) |warps| { + var kernels = try fx.loadWarps(gpu.d, @intCast(try gpu.ctx.attribute(.multiprocessor_count)), warps); + defer kernels.unload(); + const act = try r.out(most * slots * rank_width * 2); + const out = try r.out(most * slots * d * 4); + const combined = try a.alloc(f32, most * slots * d); + @memset(combined, 0); + const reference_full = try a.alloc(f64, combined.len); + @memset(reference_full, 0); + const reference_abs = try a.alloc(f64, combined.len); + @memset(reference_abs, 0); + for ([_]*const flash.modelopt_experts.Owned{ &left, &right }, 0..) |owned, rank| { + var all: []f32 = undefined; + for ([_]usize{ most, 1, 7 }) |rows| { + try gpu.d.check(gpu.d.api.cuMemsetD8Async(out, 0, most * slots * d * 4, r.s.handle), "cuMemsetD8Async"); + const items = cuda.grouped.maxItems(rows * slots, experts + 1, 16); + try router.route(r.s, dp, rows * slots, experts + 1, 16, plan); + try kernels.run(r.s, .up, .{ .x = dx, .stride = d, .slots = slots }, owned.view(), plan, act, items, experts); + try kernels.run(r.s, .down_f32, .{ .x = act, .stride = rank_width }, owned.view(), plan, out, items, experts); + if (rows == most) all = try r.back(f32, out, most * slots * d) else try r.sameRows("NVFP4 rank output rows", f32, out, all[0 .. rows * slots * d]); } - var activation: [rank_width]f64 = undefined; - for (&activation, 0..) |*v, j| { - const channel = rank * rank_width + j; - var gate: f64 = 0; - var up: f64 = 0; - for (0..d) |k| { - gate += rig.bf(x[row * d + k]) * weight(&source, e, 0, channel, k); - up += rig.bf(x[row * d + k]) * weight(&source, e, 1, channel, k); + if (prior[rank]) |ref| try check.sameBytes("NVFP4 CTA grouping leaves every rank output byte-exact", std.mem.sliceAsBytes(all), std.mem.sliceAsBytes(ref)) else prior[rank] = all; + for (0..most) |row| for (0..slots) |slot| { + const e: usize = @intCast(picks[row * slots + slot]); + if (e == experts) { + for (all[(row * slots + slot) * d ..][0..d]) |v| try check.expect(v == 0, "shared slot untouched by NVFP4", .{}); + continue; } - gate = rig.rbf(gate); - up = rig.rbf(up); - v.* = rig.rbf(rig.rbf(gate / (1 + @exp(-gate))) * up); - } - for (0..d) |j| { - var want: f64 = 0; - var mag: f64 = 0; - for (activation, 0..) |v, k| { - const term = v * weight(&source, e, 2, j, rank * rank_width + k); - want += term; - mag += @abs(term); + var activation: [rank_width]f64 = undefined; + for (&activation, 0..) |*v, j| { + const channel = rank * rank_width + j; + var gate: f64 = 0; + var up: f64 = 0; + for (0..d) |k| { + gate += rig.bf(x[row * d + k]) * weight(&source, e, 0, channel, k); + up += rig.bf(x[row * d + k]) * weight(&source, e, 1, channel, k); + } + gate = rig.rbf(gate); + up = rig.rbf(up); + v.* = rig.rbf(rig.rbf(gate / (1 + @exp(-gate))) * up); } - const at = (row * slots + slot) * d + j; - try r.within("NVFP4 rank FP32 down", all[at], want, rig.sumTol(mag, rank_width) + 1e-6); - combined[at] += all[at]; - reference_full[at] += want; - reference_abs[at] += mag; - } - }; + for (0..d) |j| { + var want: f64 = 0; + var mag: f64 = 0; + for (activation, 0..) |v, k| { + const term = v * weight(&source, e, 2, j, rank * rank_width + k); + want += term; + mag += @abs(term); + } + const at = (row * slots + slot) * d + j; + try r.within("NVFP4 rank FP32 down", all[at], want, rig.sumTol(mag, rank_width) + 1e-6); + combined[at] += all[at]; + reference_full[at] += want; + reference_abs[at] += mag; + } + }; + } + for (combined, reference_full, reference_abs) |got, want, mag| try r.within("NVFP4 summed rank partials", got, want, rig.sumTol(mag, width) + 2e-6); + r.report("NVFP4 TP cuts: both intermediate halves, exact source reader and packing, independent FP64 reference, FP32 partials and untouched shared slot; rows 1/7/17 equal"); } - for (combined, reference_full, reference_abs) |got, want, mag| try r.within("NVFP4 summed rank partials", got, want, rig.sumTol(mag, width) + 2e-6); - r.report("NVFP4 TP cuts: both intermediate halves, exact source reader and packing, independent FP64 reference, FP32 partials and untouched shared slot; rows 1/7/17 equal"); } From e71e400dbbaa58ad15e1a1f91935c8c4fdcd91ff Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:10:54 -0700 Subject: [PATCH 079/115] Export fixed-history full score rows for independent model fidelity checks --- zig/tests/cuda/flashnext_fidelity_tests.zig | 88 +++++++++++++++++++++ zig/tests/cuda/main.zig | 3 + 2 files changed, 91 insertions(+) create mode 100644 zig/tests/cuda/flashnext_fidelity_tests.zig diff --git a/zig/tests/cuda/flashnext_fidelity_tests.zig b/zig/tests/cuda/flashnext_fidelity_tests.zig new file mode 100644 index 0000000000..44affe1389 --- /dev/null +++ b/zig/tests/cuda/flashnext_fidelity_tests.zig @@ -0,0 +1,88 @@ +//! Emit full BF16 score rows at fixed histories for comparison with an independent model implementation. +const std = @import("std"); +const cuda = @import("cuda"); +const flash = @import("flashnext"); +const check = @import("check.zig"); +const context = 8192; +const chunk = 32; + +const Case = struct { id: []const u8, prompt: []const u32, continuation: []const u32 }; + +fn validate(cases: []const Case) !void { + if (cases.len == 0 or cases.len > 64) return error.BadFixture; + for (cases, 0..) |case, i| { + if (case.id.len == 0 or case.id.len > 64 or case.prompt.len == 0 or case.continuation.len > 64 or case.prompt.len + case.continuation.len > context) return error.BadFixture; + for (case.id) |c| if (!std.ascii.isAlphanumeric(c) and c != '-') return error.BadFixture; + for (cases[0..i]) |prior| if (std.mem.eql(u8, prior.id, case.id)) return error.BadFixture; + for (case.prompt) |id| if (id >= flash.prompt.vocabulary) return error.InvalidToken; + for (case.continuation) |id| if (id >= flash.prompt.vocabulary) return error.InvalidToken; + } +} + +pub fn run(gpu: check.Gpu, model: []const u8, rank: usize, port: u16, fixture: []const u8, output: []const u8) !void { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, fixture, gpu.gpa, .limited(1 << 22)); + defer gpu.gpa.free(text); + const parsed = try std.json.parseFromSlice([]const Case, gpu.gpa, text, .{}); + defer parsed.deinit(); + try validate(parsed.value); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + const identity = try flash.rank_engine.identity(gpu.gpa, gpu.io, model, context, chunk); + const session = try flash.rank_engine.Session.init(gpu.gpa, gpu.io, "127.0.0.1", port, rank, identity); + const e = flash.prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, model, 2, rank) catch |err| { + session.poison(); + session.deinit(); + return err; + }; + session.engine = e; + e.collective = &session.collective; + defer { + session.deinit(); + e.deinit(); + } + if (e.chunk != chunk) return error.RankIdentity; + var logits = try cuda.DeviceBuffer.alloc(gpu.d, flash.prompt.vocabulary * 4); + defer logits.free(); + if (rank == 1) { + try session.follow(flash.Tri, tri, logits.ptr); + check.pass("fidelity follower completed every fixed-history case", .{}); + return; + } + for (parsed.value) |case| { + var seq = try session.create(context); + defer session.drop(&seq); + const row_bytes = flash.prompt.vocabulary * 2; + const rows = try gpu.gpa.alloc(u8, (case.continuation.len + 1) * row_bytes); + defer gpu.gpa.free(rows); + try session.forward(flash.Tri, tri, &seq, case.prompt, .{}); + @memcpy(rows[0..row_bytes], seq.logits); + for (case.continuation, 0..) |token, i| { + try session.forward(flash.Tri, tri, &seq, &.{token}, .{}); + @memcpy(rows[(i + 1) * row_bytes ..][0..row_bytes], seq.logits); + } + const name = try std.fmt.allocPrint(gpu.gpa, "{s}.bf16", .{case.id}); + defer gpu.gpa.free(name); + const path = try std.fs.path.join(gpu.gpa, &.{ output, name }); + defer gpu.gpa.free(path); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = rows }); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(rows, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + check.pass("fixed history {s}: prompt={d} score-rows={d} sha256={s}", .{ case.id, case.prompt.len, case.continuation.len + 1, &hex }); + } +} + +test "fixed histories reject unsafe output identities and invalid token or context bounds" { + const good: Case = .{ .id = "short-1", .prompt = &.{1}, .continuation = &.{2} }; + try validate(&.{good}); + var bad = good; + bad.id = "../bad"; + try std.testing.expectError(error.BadFixture, validate(&.{bad})); + try std.testing.expectError(error.BadFixture, validate(&.{ good, good })); + bad = good; + bad.prompt = &.{flash.prompt.vocabulary}; + try std.testing.expectError(error.InvalidToken, validate(&.{bad})); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 3063654595..6d027146d7 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -34,6 +34,7 @@ const recapture_tests = @import("recapture_tests.zig"); const experts_rows_tests = @import("experts_rows_tests.zig"); const route_rows_tests = @import("route_rows_tests.zig"); const flashnext_model_rank_tests = @import("flashnext_model_rank_tests.zig"); +const flashnext_fidelity_tests = @import("flashnext_fidelity_tests.zig"); const flashnext_window_tests = @import("flashnext_window_tests.zig"); const flashnext_rank_tests = @import("flashnext_rank_tests.zig"); const flashnext_draft_affine_tests = @import("flashnext_draft_affine_tests.zig"); @@ -65,6 +66,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-fidelity-rank MODEL RANK PORT FIXTURE OUTPUT fixed-history full score rows \\ flashnext-window-rank MODEL RANK PORT IDS_FILE two-rank serial/window and rollback parity \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) \\ flashnext-rank RANK PORT [sum|eof|timeout|missing] two-process collective and failure probes @@ -161,6 +163,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-fidelity-rank")) return flashnext_fidelity_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-rank")) return flashnext_rank_tests.run(gpu, try std.fmt.parseInt(usize, try arg(rest, 0), 10), try std.fmt.parseInt(u16, try arg(rest, 1), 10), if (rest.len > 2) rest[2] else "sum"); From 961e1d38f681d13eec7074c587d012fefbafa4fd Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:37:24 -0700 Subject: [PATCH 080/115] Experiment with exact vocabulary-column projection across CUDA ranks --- zig/src/families/flashnext/cuda_batch.zig | 14 +++++- .../families/flashnext/cuda_batch_state.zig | 4 +- .../families/flashnext/cuda_collective.zig | 10 ++++ zig/src/families/flashnext/cuda_head.zig | 48 +++++++++++++++++++ zig/src/families/flashnext/cuda_launch.zig | 1 + zig/src/families/flashnext/cuda_prompt.zig | 3 +- .../families/flashnext/cuda_rank_engine.zig | 6 ++- zig/tests/cuda/flashnext_rank_tests.zig | 27 +++++++++++ 8 files changed, 108 insertions(+), 5 deletions(-) create mode 100644 zig/src/families/flashnext/cuda_head.zig diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index f1b436898e..0930545bd8 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -511,6 +511,16 @@ pub fn Pass(comptime Tri: type) type { try self.stream.synchronize(); } + /// Column tiles keep their own reduction order; each gathered row restores global vocabulary order. + fn projectHead(self: *Self, out: u64, n: usize) !void { + const s = self.s; + if (self.m.head.n == vocab) return self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); + if (self.m.head.n != vocab / 2 or n > s.cap) return error.UnexpectedTensor; + const c = self.m.collective orelse return error.TensorParallelUnavailable; + try self.lane(s.mixed, s.xsm, self.m.head, s.head_local, s.dummy, n, 1); + for (0..n) |r| try c.gather(self.stream, s.head_local + r * vocab, out + r * vocab * 2, vocab / 2); + } + fn head(self: *Self, logits_host: ?[]u8) !void { const s = self.s; const n = self.rows; @@ -526,7 +536,7 @@ pub fn Pass(comptime Tri: type) type { try self.copy(self.hooks.streams, s.h, n * wide * 2); try self.readout(self.m.mixer, 0, n); const out = if (self.hooks.logits_out != 0) self.hooks.logits_out else return error.UnexpectedTensor; - try self.lane(s.mixed, s.xsm, self.m.head, out, s.dummy, n, 1); + try self.projectHead(out, n); if (logits_host) |lh| { try self.complete(); const d = self.stream.d; @@ -553,7 +563,7 @@ pub fn Pass(comptime Tri: type) type { } if (self.hooks.streams == 0) try self.writeback(0, 1); try self.readout(self.m.mixer, 0, 1); - try self.lane(s.mixed, s.xsm, self.m.head, s.logits, s.dummy, 1, 1); + try self.projectHead(s.logits, 1); if (logits_host) |lh| { try self.complete(); const d = self.stream.d; diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 72a62dc897..f837db0887 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -264,6 +264,7 @@ pub const Scratch = struct { ple_pss: u64 = 0, ple_nrow: u64 = 0, logits: u64 = 0, + head_local: u64 = 0, inv: u64 = 0, dummy: u64 = 0, /// Pinned host copies of what a pass uploads (ids, positions, n-gram rows and sums), owned until the copies land. @@ -334,7 +335,7 @@ pub const Scratch = struct { n += cap * moe_w * 6 + cap * dims * 4; n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; - n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + half * 4 + 4096; + n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + cap * vocab + half * 4 + 4096; return n; } @@ -403,6 +404,7 @@ pub const Scratch = struct { s.ple_pss = try s.add(gpa, d, cap * streams * 4); s.ple_nrow = try s.add(gpa, d, cap * wide * 2); s.logits = try s.add(gpa, d, vocab * 2); + s.head_local = try s.add(gpa, d, cap * vocab); var inv: [half]f32 = undefined; for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); s.inv = try s.add(gpa, d, half * 4); diff --git a/zig/src/families/flashnext/cuda_collective.zig b/zig/src/families/flashnext/cuda_collective.zig index a2aae8dc97..c453cd900e 100644 --- a/zig/src/families/flashnext/cuda_collective.zig +++ b/zig/src/families/flashnext/cuda_collective.zig @@ -81,6 +81,16 @@ pub const Collective = struct { try c.wait(stream, wire.deadline(c.timeout_ms)); } + /// Disjoint BF16 vocabulary columns are concatenated without an arithmetic reduction. + pub fn gather(c: *Collective, stream: cuda.Stream, local: u64, out: u64, count: usize) !void { + if (c.failed or c.comm == null) return error.RankDisconnected; + if (count == 0) return; + errdefer c.abort(); + try c.link.checkPeer(); + try c.submitted(c.lib.api.ncclAllGather(local, out, count, .bf16, c.comm, stream.handle)); + try c.wait(null, wire.deadline(c.timeout_ms)); + } + /// FP32 sum is enqueued on the compute stream; later kernels on that stream observe its result. pub fn sum(c: *Collective, stream: cuda.Stream, values: u64, count: usize) !void { if (c.failed or c.comm == null) return error.RankDisconnected; diff --git a/zig/src/families/flashnext/cuda_head.zig b/zig/src/families/flashnext/cuda_head.zig new file mode 100644 index 0000000000..4793a99b60 --- /dev/null +++ b/zig/src/families/flashnext/cuda_head.zig @@ -0,0 +1,48 @@ +//! Optional vocabulary-column partitioning: each rank retains the checkpoint bits for its own output rows. +const std = @import("std"); +const Tensor = @import("core").safetensors.Tensor; + +pub fn setting(value: ?[]const u8) !bool { + const v = value orelse return false; + if (std.mem.eql(u8, v, "0")) return false; + if (std.mem.eql(u8, v, "1")) return true; + return error.BadHeadSharding; +} + +pub fn enabled() !bool { + return setting(if (std.c.getenv("TF_FLASHNEXT_SHARD_HEAD")) |v| std.mem.span(v) else null); +} + +pub fn shard(t: Tensor, world: usize, rank: usize) !Tensor { + if (world != 2 or rank >= world or t.rank != 2 or t.dtype != .bf16 or t.dim(0) == 0 or t.dim(0) % 256 != 0) + return error.BadHeadSharding; + const bytes = try std.math.mul(usize, try std.math.mul(usize, t.dim(0), t.dim(1)), 2); + if (t.bytes.len != bytes) return error.UnexpectedTensor; + var out = t; + out.shape[0] /= world; + out.bytes = t.bytes[rank * (bytes / world) ..][0 .. bytes / world]; + return out; +} + +test "head partitions preserve every BF16 byte in global vocabulary order" { + var bytes: [256 * 64 * 2]u8 = undefined; + for (&bytes, 0..) |*v, i| v.* = @truncate(i * 17 + i / 251); + const t: Tensor = .{ .dtype = .bf16, .rank = 2, .shape = .{ 256, 64, 0, 0, 0 }, .bytes = &bytes }; + const a = try shard(t, 2, 0); + const b = try shard(t, 2, 1); + try std.testing.expectEqual(@as(usize, 128), a.dim(0)); + try std.testing.expectEqualSlices(u8, bytes[0 .. bytes.len / 2], a.bytes); + try std.testing.expectEqualSlices(u8, bytes[bytes.len / 2 ..], b.bytes); + try std.testing.expectError(error.BadHeadSharding, shard(t, 1, 0)); + try std.testing.expectError(error.BadHeadSharding, shard(t, 2, 2)); + var bad = t; + bad.shape[0] = 255; + try std.testing.expectError(error.BadHeadSharding, shard(bad, 2, 0)); + bad = t; + bad.bytes = bytes[1..]; + try std.testing.expectError(error.UnexpectedTensor, shard(bad, 2, 0)); + try std.testing.expect(!try setting(null)); + try std.testing.expect(!try setting("0")); + try std.testing.expect(try setting("1")); + try std.testing.expectError(error.BadHeadSharding, setting("2")); +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index 0683c45bfd..05143c866f 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -384,6 +384,7 @@ test { _ = row_path; _ = @import("cuda_mtp_check.zig"); _ = @import("cuda_draft_affine.zig"); + _ = @import("cuda_head.zig"); } test "embed grid is one program per row and per group of 32" { diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index 9a3478aadc..71a6ae74cf 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -175,8 +175,9 @@ pub const Engine = struct { const head_tensor = try e.store.tensor("language_model.lm_head.weight"); if (head_tensor.dtype == .bf16) { if (!head_tensor.is(.bf16, &.{ vocab, dims })) return error.UnexpectedTensor; - e.head_q = try @import("cuda_projection.zig").dense(gpa, driver, head_tensor); + e.head_q = try @import("cuda_projection.zig").dense(gpa, driver, if (try @import("cuda_head.zig").enabled()) try @import("cuda_head.zig").shard(head_tensor, world, rank) else head_tensor); } else { + if (try @import("cuda_head.zig").enabled()) return error.BadHeadSharding; var buf: [80]u8 = undefined; const head_w = try join3(gpa, &e.store, &buf, "language_model.lm_head", ""); defer gpa.free(head_w.w); diff --git a/zig/src/families/flashnext/cuda_rank_engine.zig b/zig/src/families/flashnext/cuda_rank_engine.zig index b8852aa595..ff7db8e5ae 100644 --- a/zig/src/families/flashnext/cuda_rank_engine.zig +++ b/zig/src/families/flashnext/cuda_rank_engine.zig @@ -17,6 +17,7 @@ pub const Identity = extern struct { chunk: u64, decode_graph: u64 = 0, mtp: u64 = 0, + head_sharded: u64 = 0, }; fn fileHash(a: std.mem.Allocator, io: std.Io, path: []const u8) ![32]u8 { @@ -38,7 +39,7 @@ pub fn identity(a: std.mem.Allocator, io: std.Io, dir: []const u8, context: usiz defer a.free(index); const program = try std.process.executablePathAlloc(io, a); defer a.free(program); - return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()) }; + return .{ .config = try fileHash(a, io, config), .index = try fileHash(a, io, index), .program = try fileHash(a, io, program), .checkpoint = checkpoint, .context = context, .chunk = chunk, .decode_graph = @intFromBool(@import("cuda_round_graph.zig").enabled()), .head_sharded = @intFromBool(try @import("cuda_head.zig").enabled()) }; } pub const Session = struct { @@ -427,6 +428,9 @@ test "rank identity binds checkpoint, config, index, program and execution setti other = base; other.mtp = 1; try std.testing.expect(!sameIdentity(base, other)); + other = base; + other.head_sharded = 1; + try std.testing.expect(!sameIdentity(base, other)); } test "rank operations refuse invalid geometry and unsupported hooks before execution" { diff --git a/zig/tests/cuda/flashnext_rank_tests.zig b/zig/tests/cuda/flashnext_rank_tests.zig index f7bdcd15c9..6bbe109dfb 100644 --- a/zig/tests/cuda/flashnext_rank_tests.zig +++ b/zig/tests/cuda/flashnext_rank_tests.zig @@ -85,6 +85,33 @@ pub fn run(gpu: check.Gpu, rank: usize, port: u16, mode: []const u8) !void { try link.sendHeader(h, wire.deadline(30_000)); } } + for ([_]usize{ 0, 1, 17, 124_160 }) |count| { + const host = try gpu.gpa.alloc(u16, @max(1, count)); + defer gpu.gpa.free(host); + for (host, 0..) |*v, i| v.* = @truncate(i * 31 + rank * 0x8001); + var local = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer local.free(); + var output = try cuda.DeviceBuffer.alloc(gpu.d, @max(2, count * 4)); + defer output.free(); + try collective.gather(stream, local.ptr, output.ptr, count); + try collective.finish(stream); + const all = try gpu.gpa.alloc(u16, count * 2); + defer gpu.gpa.free(all); + if (count > 0) try output.download(0, std.mem.sliceAsBytes(all)); + for (0..2) |owner| for (0..count) |i| { + const expected: u16 = @truncate(i * 31 + owner * 0x8001); + try check.expect(all[owner * count + i] == expected, "rank {d}: gathered BF16 bits at {d}/{d}", .{ rank, owner, i }); + }; + } + const done: wire.Header = .{ .kind = .ready, .sequence = 99 }; + if (rank == 0) { + try link.sendHeader(done, wire.deadline(30_000)); + try std.testing.expectEqualDeep(done, try link.receiveHeader(wire.deadline(30_000))); + } else { + try std.testing.expectEqualDeep(done, try link.receiveHeader(wire.deadline(30_000))); + try link.sendHeader(done, wire.deadline(30_000)); + } + check.pass("rank {d}: BF16 all-gather preserves every bit and rank column order for empty, scalar, tail and vocabulary-half rows", .{rank}); check.pass("rank {d}: two-process chained NCCL FP32 sums exact for empty, scalar, tail and 1,000,003 values; fragmented framed control preserved", .{rank}); } From 8f014e1db00ec75ffc441fc6c2a88134b2b2da72 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:41:26 -0700 Subject: [PATCH 081/115] Add single-rank score export to isolate full-model fidelity --- zig/tests/cuda/flashnext_fidelity_tests.zig | 37 +++++++++++++++++++++ zig/tests/cuda/main.zig | 2 ++ 2 files changed, 39 insertions(+) diff --git a/zig/tests/cuda/flashnext_fidelity_tests.zig b/zig/tests/cuda/flashnext_fidelity_tests.zig index 44affe1389..d3c189fb5f 100644 --- a/zig/tests/cuda/flashnext_fidelity_tests.zig +++ b/zig/tests/cuda/flashnext_fidelity_tests.zig @@ -75,6 +75,43 @@ pub fn run(gpu: check.Gpu, model: []const u8, rank: usize, port: u16, fixture: [ } } +pub fn runSingle(gpu: check.Gpu, model: []const u8, fixture: []const u8, output: []const u8) !void { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, fixture, gpu.gpa, .limited(1 << 22)); + defer gpu.gpa.free(text); + const parsed = try std.json.parseFromSlice([]const Case, gpu.gpa, text, .{}); + defer parsed.deinit(); + try validate(parsed.value); + var stream = try cuda.Stream.init(gpu.d, true); + defer stream.deinit(); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + const tri: flash.Tri = .{ .s = stream, .own = &glue }; + const e = try flash.prompt.Engine.init(gpu.gpa, gpu.io, @constCast(gpu.d), &stream, model); + defer e.deinit(); + for (parsed.value) |case| { + var seq = try e.newSeq(context); + defer seq.deinit(gpu.gpa); + const row_bytes = flash.prompt.vocabulary * 2; + const rows = try gpu.gpa.alloc(u8, (case.continuation.len + 1) * row_bytes); + defer gpu.gpa.free(rows); + try e.forward(flash.Tri, tri, &seq, case.prompt, .{}); + @memcpy(rows[0..row_bytes], seq.logits); + for (case.continuation, 0..) |token, i| { + try e.forward(flash.Tri, tri, &seq, &.{token}, .{}); + @memcpy(rows[(i + 1) * row_bytes ..][0..row_bytes], seq.logits); + } + const name = try std.fmt.allocPrint(gpu.gpa, "{s}.bf16", .{case.id}); + defer gpu.gpa.free(name); + const path = try std.fs.path.join(gpu.gpa, &.{ output, name }); + defer gpu.gpa.free(path); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = rows }); + var digest: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(rows, &digest, .{}); + const hex = std.fmt.bytesToHex(digest, .lower); + check.pass("fixed history {s}: prompt={d} score-rows={d} sha256={s}", .{ case.id, case.prompt.len, case.continuation.len + 1, &hex }); + } +} + test "fixed histories reject unsafe output identities and invalid token or context bounds" { const good: Case = .{ .id = "short-1", .prompt = &.{1}, .continuation = &.{2} }; try validate(&.{good}); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 6d027146d7..be3f6ed8ff 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -66,6 +66,7 @@ const usage = \\ pocket [reps] the pocket's two kernels against tf_pocket_rows byte for byte; reps > 0 times them too \\ experts-rows experts_rows.cu's windows and chunks against the decode kernel's one-row launches \\ route-rows route_rows.cu's router logits and softmax top-k against float64 (synthetic rows) + \\ flashnext-fidelity-single MODEL FIXTURE OUTPUT single-rank reference isolation \\ flashnext-fidelity-rank MODEL RANK PORT FIXTURE OUTPUT fixed-history full score rows \\ flashnext-window-rank MODEL RANK PORT IDS_FILE two-rank serial/window and rollback parity \\ flashnext-model-rank MODEL RANK PORT IDS_FILE explicit full-model replay experiment (no drafts) @@ -163,6 +164,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glue")) return glue_tests.run(gpu); if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-fidelity-single")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); if (std.mem.eql(u8, cmd, "flashnext-fidelity-rank")) return flashnext_fidelity_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); From c611adeb28c894e1769d177c029f6a57b13dd2c8 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:53:34 -0700 Subject: [PATCH 082/115] Disable Nagle on accepted HTTP streaming connections --- zig/src/server/listener.zig | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/zig/src/server/listener.zig b/zig/src/server/listener.zig index 8eb7a27df6..a4d0896e9b 100644 --- a/zig/src/server/listener.zig +++ b/zig/src/server/listener.zig @@ -53,8 +53,15 @@ pub const Listener = struct { const e = posix.errno(rc); switch (e) { .SUCCESS => { + const fd: posix.socket_t = @intCast(rc); + const one: c_int = 1; + posix.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, std.mem.asBytes(&one)) catch |err| { + _ = posix.system.close(fd); + log.line("could not configure accepted TCP connection: {s}", .{@errorName(err)}); + continue; + }; nofile_logged.store(false, .release); - return .{ .fd = @intCast(rc), .peer = Threaded.addressFromPosix(&storage) }; + return .{ .fd = fd, .peer = Threaded.addressFromPosix(&storage) }; }, .INTR, .AGAIN, .CONNABORTED, .MFILE, .NFILE, .NOBUFS, .NOMEM => { if (e != .INTR) { @@ -132,3 +139,25 @@ test "peer text" { try std.testing.expectEqualStrings("127.0.0.1", peerText(&buf, .{ .ip4 = .loopback(0) })); try std.testing.expectEqualStrings("::1", peerText(&buf, .{ .ip6 = .loopback(0) })); } + +test "accepted IPv4 and IPv6 TCP streams disable Nagle for incremental responses" { + for ([_]net.IpAddress{ .{ .ip4 = .loopback(0) }, .{ .ip6 = .loopback(0) } }) |address| { + const listener = try Listener.open(address); + defer listener.close(); + const rc = posix.system.socket(if (address == .ip4) posix.AF.INET else posix.AF.INET6, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const client: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(client); + const peer: net.IpAddress = if (address == .ip4) .{ .ip4 = .loopback(listener.port()) } else .{ .ip6 = .loopback(listener.port()) }; + var storage: Threaded.PosixAddress = undefined; + const len = Threaded.addressToPosix(&peer, &storage); + if (posix.errno(posix.system.connect(client, &storage.any, len)) != .SUCCESS) return error.ConnectFailed; + const stop = std.atomic.Value(bool).init(false); + const accepted = listener.accept(&stop) orelse return error.AcceptFailed; + defer _ = posix.system.close(accepted.fd); + var value: c_int = 0; + var size: posix.socklen_t = @sizeOf(c_int); + try std.testing.expect(std.c.getsockopt(accepted.fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &value, &size) == 0); + try std.testing.expect(value != 0); + } +} From 7b0cd5aede94908ea31aa42d6dfae231c86a4d43 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:53:34 -0700 Subject: [PATCH 083/115] Disable Nagle on accepted HTTP streaming connections --- zig/src/server/listener.zig | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/zig/src/server/listener.zig b/zig/src/server/listener.zig index 8eb7a27df6..a4d0896e9b 100644 --- a/zig/src/server/listener.zig +++ b/zig/src/server/listener.zig @@ -53,8 +53,15 @@ pub const Listener = struct { const e = posix.errno(rc); switch (e) { .SUCCESS => { + const fd: posix.socket_t = @intCast(rc); + const one: c_int = 1; + posix.setsockopt(fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, std.mem.asBytes(&one)) catch |err| { + _ = posix.system.close(fd); + log.line("could not configure accepted TCP connection: {s}", .{@errorName(err)}); + continue; + }; nofile_logged.store(false, .release); - return .{ .fd = @intCast(rc), .peer = Threaded.addressFromPosix(&storage) }; + return .{ .fd = fd, .peer = Threaded.addressFromPosix(&storage) }; }, .INTR, .AGAIN, .CONNABORTED, .MFILE, .NFILE, .NOBUFS, .NOMEM => { if (e != .INTR) { @@ -132,3 +139,25 @@ test "peer text" { try std.testing.expectEqualStrings("127.0.0.1", peerText(&buf, .{ .ip4 = .loopback(0) })); try std.testing.expectEqualStrings("::1", peerText(&buf, .{ .ip6 = .loopback(0) })); } + +test "accepted IPv4 and IPv6 TCP streams disable Nagle for incremental responses" { + for ([_]net.IpAddress{ .{ .ip4 = .loopback(0) }, .{ .ip6 = .loopback(0) } }) |address| { + const listener = try Listener.open(address); + defer listener.close(); + const rc = posix.system.socket(if (address == .ip4) posix.AF.INET else posix.AF.INET6, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const client: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(client); + const peer: net.IpAddress = if (address == .ip4) .{ .ip4 = .loopback(listener.port()) } else .{ .ip6 = .loopback(listener.port()) }; + var storage: Threaded.PosixAddress = undefined; + const len = Threaded.addressToPosix(&peer, &storage); + if (posix.errno(posix.system.connect(client, &storage.any, len)) != .SUCCESS) return error.ConnectFailed; + const stop = std.atomic.Value(bool).init(false); + const accepted = listener.accept(&stop) orelse return error.AcceptFailed; + defer _ = posix.system.close(accepted.fd); + var value: c_int = 0; + var size: posix.socklen_t = @sizeOf(c_int); + try std.testing.expect(std.c.getsockopt(accepted.fd, posix.IPPROTO.TCP, posix.TCP.NODELAY, &value, &size) == 0); + try std.testing.expect(value != 0); + } +} From 94cefdd8a8a8f72a9a2ddeeeb15e1f2fa103b0ec Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:14:11 -0700 Subject: [PATCH 084/115] Prefer shared plain rounds when admitting concurrent decode streams --- zig/src/families/flashnext/cuda_lanes.zig | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 059284e9d9..89d6095007 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -236,7 +236,8 @@ pub fn Namespace(comptime Tri: type) type { if (ids.len == 0 or ids.len + s.max_new > self.context) return error.PromptTooLong; const gop = try self.lanes.getOrPut(self.gpa, s); if (gop.found_existing) dropLane(self, gop.value_ptr); - const drafting = self.head != null and s.drafts; + // Admit shared plain rounds when four or more streams are active. + const drafting = self.head != null and s.drafts and self.lanes.count() < 4; // without the head's drafts, one row a stream a round: only a drafted window can rewind its states if (!drafting) { s.drafts = false; From 50c4f56ff5622829d3067e5bd77a173fa139611f Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:13:42 -0700 Subject: [PATCH 085/115] Exercise FP32 packed projections through numerical contract gates --- zig/tests/cuda/contract_projection_f32.zig | 157 +++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 8 ++ 2 files changed, 165 insertions(+) create mode 100644 zig/tests/cuda/contract_projection_f32.zig diff --git a/zig/tests/cuda/contract_projection_f32.zig b/zig/tests/cuda/contract_projection_f32.zig new file mode 100644 index 0000000000..26d32ac369 --- /dev/null +++ b/zig/tests/cuda/contract_projection_f32.zig @@ -0,0 +1,157 @@ +//! Packed projections writing FP32, with independent FP64 products and mirrored accumulator stress. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; +const Rig = @import("contract_rig.zig").Rig; + +pub const modes = [_]cuda.qmmf.Mode{ .bf16, .fp8, .fp8g, .mxfp8 }; +const served = [_]c.Shape{ + .{ .rows = 1, .n = 70, .k = 1024 }, + .{ .rows = 17, .n = 70, .k = 1024 }, + .{ .rows = 33, .n = 70, .k = 1024 }, + .{ .rows = 256, .n = 70, .k = 1024 }, + .{ .rows = 16, .n = 1280, .k = 2560 }, + .{ .rows = 16, .n = 2560, .k = 640 }, +}; + +fn depth(shape: c.Shape) usize { + return shape.k / 16 + 24; +} + +pub const Product = struct { + rig: Rig, + lane: *const cuda.qmmf.Lane, + mode: cuda.qmmf.Mode, + clusters: bool, + d: ?Dense = null, + weight: cuda.qmmf.Weight = undefined, + x: u64 = 0, + out: u64 = 0, + part: u64 = 0, + + pub fn deinit(self: *Product) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Product) g.Kernel { + const name = switch (self.mode) { + .bf16 => if (self.clusters) "cuda BF16 fp32 clustered" else "cuda BF16 fp32 reduce", + .fp8 => if (self.clusters) "cuda FP8 fp32 clustered" else "cuda FP8 fp32 reduce", + .fp8g => if (self.clusters) "cuda block FP8 fp32 clustered" else "cuda block FP8 fp32 reduce", + .mxfp8 => if (self.clusters) "cuda MXFP8 fp32 clustered" else "cuda MXFP8 fp32 reduce", + else => unreachable, + }; + return .{ .spec = .{ + .name = name, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 16 }, .{ .size = 64 }, .{ .size = 512, .step = 128, .min = 128 } }, + .served = &served, + }, .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Product); + + pub fn prepare(self: *Product, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + self.part = 0; + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 128 != 0) return false; + const d = try Dense.init(self.rig.gpu.gpa, shape, 1); + self.d = d; + d.fillX(fill, seed, 16); + try self.prepareWeight(shape, fill, seed); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + const bytes = self.lane.partBytes(shape.rows, self.weight); + if (bytes > 0) self.part = try self.rig.scratch(bytes); + self.out = try self.rig.output(shape.rows * shape.n * 4); + return true; + } + + fn prepareWeight(self: *Product, shape: c.Shape, fill: g.Fill, seed: u64) !void { + var arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const d = self.d.?; + const n = shape.n; + const k = shape.k; + const npad = cuda.fp8.padded(n); + var prng = std.Random.DefaultPrng.init(seed ^ 0xf32); + const random = prng.random(); + const codes = try a.alloc(u8, n * k); + const words = try a.alloc(u16, n * k); + const ue8 = try a.alloc(u8, n * k / 32); + for (ue8) |*v| v.* = if (fill == .stress) 127 else 123 + random.uintLessThan(u8, 8); + const inv = try a.alloc(f32, ((n + 127) / 128) * (k / 128)); + for (inv) |*v| v.* = if (fill == .stress) 1 else std.math.ldexp(@as(f32, 1), @as(i32, random.uintLessThan(u8, 8)) - 4); + for (0..n) |j| for (0..k) |i| { + const at = j * k + i; + codes[at] = if (fill == .stress) 0x38 + random.uintLessThan(u8, 24) else random.uintLessThan(u8, 0x7f); + if (fill != .stress and random.boolean()) codes[at] |= 0x80; + words[at] = c.toBf16(if (fill == .stress) @floatFromInt(1 + random.uintLessThan(u8, 7)) else random.float(f32) * 4 - 2); + }; + if (fill == .stress) for (0..n) |j| for (k / 2..k) |i| { + codes[j * k + i] = codes[j * k + i - k / 2] | 0x80; + words[j * k + i] = words[j * k + i - k / 2] | 0x8000; + }; + const scale: f32 = if (fill == .stress) 1 else 0.125; + for (0..n) |j| for (0..k) |i| { + const value = core.quant.fp8.e4m3(codes[j * k + i]); + d.w[j * k + i] = switch (self.mode) { + .bf16 => c.fromBf16(words[j * k + i]), + .fp8 => value * scale, + .fp8g => value * inv[(j / 128) * (k / 128) + i / 128], + .mxfp8 => value * std.math.ldexp(@as(f64, 1), @as(i32, ue8[j * (k / 32) + i / 32]) - 127), + else => unreachable, + }; + }; + if (self.mode == .bf16) { + const tiles = try a.alloc(u16, npad * k); + cuda.qmmf.packBf16(tiles, words, n, k); + self.weight = cuda.modelopt.bf16Weight(try self.rig.input(std.mem.sliceAsBytes(tiles)), n, k); + return; + } + const tiles = try a.alloc(u8, npad * k); + cuda.fp8.packCodes(tiles, codes, n, k); + const dw = try self.rig.input(tiles); + self.weight = switch (self.mode) { + .fp8 => cuda.modelopt.fp8Weight(dw, scale, n, k), + .mxfp8 => blk: { + const scales = try a.alloc(u8, npad * k / 32); + cuda.mxfp8.packScales(scales, ue8, n, k); + break :blk cuda.mxfp8.weight(dw, try self.rig.input(scales), n, k); + }, + .fp8g => blk: { + const scales = try a.alloc(f32, npad * k / 64); + cuda.fp8.packScales(scales, inv, n, k); + break :blk cuda.fp8.weight(dw, try self.rig.input(std.mem.sliceAsBytes(scales)), n, k); + }, + else => unreachable, + }; + } + + pub fn run(self: *Product) !void { + const shape = self.d.?.shape; + try self.lane.matmulOut(self.rig.s, self.x, shape.k, shape.rows, self.weight, self.out, .f32, self.part); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Product) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *Product, want: []f64, magnitude: []f64) !void { + try self.d.?.reference(want, magnitude); + } + + pub fn simulate(self: *Product, accumulator: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(accumulator, order, out); + } +}; diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index c89c814c6c..e1f58336a2 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -8,6 +8,7 @@ const rows_k = @import("contract_rows.zig"); const router_k = @import("contract_router.zig"); const mm_k = @import("contract_mm.zig"); const mx_k = @import("contract_mxfp8.zig"); +const wide_k = @import("contract_projection_f32.zig"); const rank_k = @import("contract_rank.zig"); const contract = core.contract; @@ -93,6 +94,13 @@ pub fn run(gpu: check.Gpu) !void { }; defer for (&mx) |*m| m.deinit(); for (&mx) |*m| try kernels.append(a, m.kernel()); + var wide: [wide_k.modes.len * 2]wide_k.Product = undefined; + for (wide_k.modes, 0..) |mode, i| { + wide[i * 2] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane, .mode = mode, .clusters = true }; + wide[i * 2 + 1] = .{ .rig = .{ .gpu = gpu, .s = s }, .lane = &lane_reduce, .mode = mode, .clusters = false }; + } + defer for (&wide) |*m| m.deinit(); + for (&wide) |*m| try kernels.append(a, m.kernel()); var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum }; defer rank.deinit(); try kernels.append(a, rank.kernel()); From 37578da9e2f739c6981999b6db0b4a61f572b206 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:16:59 -0700 Subject: [PATCH 086/115] Cover served dense and attention projection dimensions --- zig/tests/cuda/contract_projection_f32.zig | 2 ++ 1 file changed, 2 insertions(+) diff --git a/zig/tests/cuda/contract_projection_f32.zig b/zig/tests/cuda/contract_projection_f32.zig index 26d32ac369..1e172dec30 100644 --- a/zig/tests/cuda/contract_projection_f32.zig +++ b/zig/tests/cuda/contract_projection_f32.zig @@ -16,6 +16,8 @@ const served = [_]c.Shape{ .{ .rows = 256, .n = 70, .k = 1024 }, .{ .rows = 16, .n = 1280, .k = 2560 }, .{ .rows = 16, .n = 2560, .k = 640 }, + .{ .rows = 7, .n = 2560, .k = 10240 }, + .{ .rows = 7, .n = 13952, .k = 2560 }, }; fn depth(shape: c.Shape) usize { From a008b5cd8157dba2827b254d3add4fe9a5f4b8e1 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 04:12:04 -0700 Subject: [PATCH 087/115] test(cuda): register MXFP8 launch resource limits --- zig/src/core/contract/registry.zig | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..18d8928632 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -37,6 +37,14 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out fused slices", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 1 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 128 }, .{ .regs = 94, .static = 0, .bound = 128 }, .{ .regs = 94, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280 }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf reduce_kernel", .symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif", .image = "qmmf", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "qmmf reduce_kernel", .symbol = "_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif", .image = "qmmf", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "rank_sum tf_rank_sum", .symbol = "tf_rank_sum", .image = "rank_sum", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit vision_kernel<80, 8> (RADIO)", .symbol = vision80_symbol, .image = "prefill_attention", .threads = 256, .dynamic = 45056 }, .usage = .{ .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 122, .static = 0, .bound = 256 }, .{ .regs = 122, .static = 0, .bound = 256 } } }, From 5bb739249784564c9b8ec591c94b8c4aa8fbd5ad Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 04:42:24 -0700 Subject: [PATCH 088/115] test(cuda): register FP32 projection launch resources --- zig/src/core/contract/registry.zig | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 18d8928632..a14c3f58a8 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -45,6 +45,34 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "qmmf MXFP8 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf MXFP8 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb0ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf reduce_kernel", .symbol = "_ZN7tf_qmmf13reduce_kernelILb0EEEvPKfPvxif", .image = "qmmf", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 50, .static = 0, .bound = 128 }, .{ .regs = 50, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 24576, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 32768 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 32768, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 126, .static = 0, .bound = 128 }, .{ .regs = 126, .static = 0, .bound = 128 }, .{ .regs = 126, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8 fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi1ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 179, .static = 0, .bound = 128 }, .{ .regs = 179, .static = 0, .bound = 128 }, .{ .regs = 179, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 }, .{ .regs = 60, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25088, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 69, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280 }, .usage = .{ .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33280, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 }, .{ .regs = 103, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf MXFP8 fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi2ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49664, .opt_in = true }, .usage = .{ .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 229, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 }, .{ .regs = 158, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25600 }, .usage = .{ .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 92, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33792 }, .usage = .{ .{ .regs = 116, .static = 0, .bound = 128 }, .{ .regs = 116, .static = 0, .bound = 128 }, .{ .regs = 116, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 33792, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 50176, .opt_in = true }, .usage = .{ .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 }, .{ .regs = 108, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 50176, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 122, .static = 0, .bound = 128 }, .{ .regs = 122, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP8G fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi3ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 50176, .opt_in = true }, .usage = .{ .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 196, .static = 0, .bound = 128 }, .{ .regs = 155, .static = 0, .bound = 128 }, .{ .regs = 155, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 16-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 40960 }, .usage = .{ .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 16-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi16ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 40960, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 32-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi32ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 93, .static = 0, .bound = 128 }, .{ .regs = 93, .static = 0, .bound = 128 }, .{ .regs = 93, .static = 0, .bound = 128 }, .{ .regs = 55, .static = 0, .bound = 128 }, .{ .regs = 55, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 32-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi32ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 49152, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 64-row tile", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 124, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 64-row tile, cluster of 8", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb1ELb0EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 65536, .opt_in = true, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf BF16 fp32 out 64-row tile, fused slices", .symbol = "_ZN7tf_qmmf11qmmf_kernelILi4ELi64ELi64ELi1ELi4ELi4ELb1ELb0ELb1EEEvPK13__nv_bfloat16PKhS5_fPvPfiiiiiii", .image = "qmmf", .threads = 128, .dynamic = 65536, .opt_in = true }, .usage = .{ .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 }, .{ .regs = 140, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf reduce_kernel", .symbol = "_ZN7tf_qmmf13reduce_kernelILb1EEEvPKfPvxif", .image = "qmmf", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "rank_sum tf_rank_sum", .symbol = "tf_rank_sum", .image = "rank_sum", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, .{ .launch = .{ .name = "vit vision_kernel<80, 8> (RADIO)", .symbol = vision80_symbol, .image = "prefill_attention", .threads = 256, .dynamic = 45056 }, .usage = .{ .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 143, .static = 0, .bound = 256 }, .{ .regs = 122, .static = 0, .bound = 256 }, .{ .regs = 122, .static = 0, .bound = 256 } } }, From cb2c74984d4017b780f4e792001e1c3bffd99c59 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 06:07:35 -0700 Subject: [PATCH 089/115] cuda: refresh compiler-derived launch resource records --- docs/kernel-contract.md | 4 ++++ zig/src/core/contract/registry.zig | 12 ++++++------ 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index e625bcb2fd..0b80446a42 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -108,3 +108,7 @@ CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router tile, lane_gemv and its split-K form (both bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles with fp32 out (K slices met in a cluster, fused, or by `reduce_kernel`), and the reduce-scatter's `tf_rank_sum`. + +## CUDA compiler resource baseline + +The registered usage arrays were re-derived with NVCC 13.3.73 (`cuda_13.3.r13.3/compiler.38244171_0`) for SM 80, 86, 89, 120 and 121. `kernel-contract-limits` reads each compiled image and reports its register count, static shared memory and thread bound. Source changes or a different compiler may change these values; re-derive the arrays and run the live audit rather than changing the audit tolerance. Launch dimensions, dynamic shared memory, device limits and planted controls remain independent checks. Only SM 120 was executed on hardware in this update; the other four records come from their compiled images, not hardware execution. diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index b79337b5c5..028a30b511 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -20,20 +20,20 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK /// CUDA kernels as launched; tf-cuda-test kernel-contract-limits re-derives the usage records. pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, - .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 } } }, + .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 } } }, .{ .launch = .{ .name = "router tf_nemo_router", .symbol = "tf_nemo_router", .image = "nemotron_route", .threads = 256, .dynamic = 28800 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "router tf_nemo_topk", .symbol = "tf_nemo_topk", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "router tf_nemo_topk", .symbol = "tf_nemo_topk", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "qmm_group cluster of 8", .symbol = group_symbol, .image = "qmm_group", .threads = 128, .dynamic = 35328, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmm_group one CTA a tile", .symbol = group_symbol, .image = "qmm_group", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "lane_gemv", .symbol = gemv_symbol, .image = "lane_gemv", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "lane_gemv split-K", .symbol = "tf_lane_gemv_split", .image = "lane_gemv", .threads = 128, .dynamic = 17664 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 44, .static = 128, .bound = 128 }, .{ .regs = 44, .static = 128, .bound = 128 } } }, + .{ .launch = .{ .name = "lane_gemv", .symbol = gemv_symbol, .image = "lane_gemv", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "lane_gemv split-K", .symbol = "tf_lane_gemv_split", .image = "lane_gemv", .threads = 128, .dynamic = 17664 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 72, .static = 128, .bound = 128 }, .{ .regs = 72, .static = 128, .bound = 128 } } }, .{ .launch = .{ .name = "prompt16 FP4", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 0, 0 }), .image = "prompt16", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "prompt16 FP8", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 1, 0 }), .image = "prompt16", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 218, .static = 0, .bound = 128 }, .{ .regs = 218, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "prompt16 FP8", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 1, 0 }), .image = "prompt16", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "prompt16 FP4 fp32 out", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 0, 1 }), .image = "prompt16", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 16-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 16, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 17408 }, .usage = .{ .{ .regs = 88, .static = 0, .bound = 128 }, .{ .regs = 88, .static = 0, .bound = 128 }, .{ .regs = 96, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 16-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 16, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 17408, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600 }, .usage = .{ .{ .regs = 106, .static = 0, .bound = 128 }, .{ .regs = 106, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 63, .static = 0, .bound = 128 }, .{ .regs = 63, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 61, .static = 0, .bound = 128 }, .{ .regs = 61, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out fused slices", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 1 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 } } }, From fe461bb37250734006dcab38402c5ba90eac463c Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 06:13:51 -0700 Subject: [PATCH 090/115] test(cuda): expose writeback double rounding against FP64 --- zig/tests/cuda/flashnext_fidelity_tests.zig | 6 ++- zig/tests/cuda/flashnext_hc_tests.zig | 46 +++++++++++++++++++++ zig/tests/cuda/main.zig | 3 +- 3 files changed, 52 insertions(+), 3 deletions(-) diff --git a/zig/tests/cuda/flashnext_fidelity_tests.zig b/zig/tests/cuda/flashnext_fidelity_tests.zig index d3c189fb5f..3f3c793ce1 100644 --- a/zig/tests/cuda/flashnext_fidelity_tests.zig +++ b/zig/tests/cuda/flashnext_fidelity_tests.zig @@ -75,7 +75,7 @@ pub fn run(gpu: check.Gpu, model: []const u8, rank: usize, port: u16, fixture: [ } } -pub fn runSingle(gpu: check.Gpu, model: []const u8, fixture: []const u8, output: []const u8) !void { +pub fn runSingle(gpu: check.Gpu, model: []const u8, fixture: []const u8, output: []const u8, serial: bool) !void { const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, fixture, gpu.gpa, .limited(1 << 22)); defer gpu.gpa.free(text); const parsed = try std.json.parseFromSlice([]const Case, gpu.gpa, text, .{}); @@ -94,7 +94,9 @@ pub fn runSingle(gpu: check.Gpu, model: []const u8, fixture: []const u8, output: const row_bytes = flash.prompt.vocabulary * 2; const rows = try gpu.gpa.alloc(u8, (case.continuation.len + 1) * row_bytes); defer gpu.gpa.free(rows); - try e.forward(flash.Tri, tri, &seq, case.prompt, .{}); + if (serial) { + for (case.prompt) |token| try e.forward(flash.Tri, tri, &seq, &.{token}, .{}); + } else try e.forward(flash.Tri, tri, &seq, case.prompt, .{}); @memcpy(rows[0..row_bytes], seq.logits); for (case.continuation, 0..) |token, i| { try e.forward(flash.Tri, tri, &seq, &.{token}, .{}); diff --git a/zig/tests/cuda/flashnext_hc_tests.zig b/zig/tests/cuda/flashnext_hc_tests.zig index bfca6aff72..1c0767849d 100644 --- a/zig/tests/cuda/flashnext_hc_tests.zig +++ b/zig/tests/cuda/flashnext_hc_tests.zig @@ -324,6 +324,51 @@ fn ngram(r: *Rig, g: *const flash.Glue) !void { }; } +fn writebackRounding(r: *Rig, g: *const flash.Glue) !void { + const n: usize = 256; + const m: usize = 17; + const ss: usize = 4; + const h = r.bfs(m * ss * n, 1); + const inject = r.bfs(m * ss, 0.125); + const values = r.bfs(m * n, 1); + h[0] = 0x3b34; + inject[0] = 0x3c9a; + values[0] = 0xbe16; + const branch = try r.a.alloc(f32, m * n); + for (values, branch) |v, *b| b.* = @floatCast(bf(v)); + const weights = try r.a.alloc(f32, m); + @memset(weights, 1); + const dh = try r.dev(std.mem.sliceAsBytes(h)); + const di = try r.dev(std.mem.sliceAsBytes(inject)); + const db = try r.dev(std.mem.sliceAsBytes(branch)); + const dw = try r.dev(std.mem.sliceAsBytes(weights)); + for ([_]usize{ 2, 4 }) |mode| { + const output = try r.out(m * ss * n * 2); + const sums = try r.out(m * ss * 4); + try g.hcWriteback(r.s, dh, output, sums, db, di, db, dw, m * n, m, n, ss, mode, 0, 1, 1); + const got = try r.back(u16, output, h.len); + var worst: f64 = 0; + var improved: usize = 0; + for (got, 0..) |v, i| { + const row = i / (ss * n); + const stream = i / n % ss; + const d = i % n; + const x = bf(values[row * n + d]); + const gate = bf(inject[row * ss + stream]); + const exact = bf(h[i]) + x * gate; + const old = rbf(bf(h[i]) + rbf(x * gate)); + const error_now = @abs(bf(v) - exact); + const error_old = @abs(old - exact); + try check.expect(error_now <= error_old, "writeback mode {d} row {d} stream {d} column {d}: FP64 error {e} exceeds prior {e}", .{ mode, row, stream, d, error_now, error_old }); + try check.expect(bf(v) == rbf(exact), "writeback mode {d} column {d}: one rounding of the full update", .{ mode, i }); + if (error_now < error_old) improved += 1; + worst = @max(worst, error_now); + } + try check.expect(improved > 0, "writeback mode {d}: cancellation fixture improves over double rounding", .{mode}); + std.debug.print("PASS writeback mode {d}: {d} FP64 columns, 0 worse, {d} improved, worst absolute {e}\n", .{ mode, got.len, improved, worst }); + } +} + pub fn run(gpu: check.Gpu) !void { var arena: std.heap.ArenaAllocator = .init(gpu.gpa); defer arena.deinit(); @@ -332,6 +377,7 @@ pub fn run(gpu: check.Gpu) !void { defer r.deinit(); var g = try flash.Glue.load(gpu.d); defer g.deinit(); + try writebackRounding(&r, &g); for ([_]usize{ 0, 2, 4 }) |mode| try writeback(&r, &g, mode); try readout(&r, &g); try ngram(&r, &g); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 297168af1e..6633ebed77 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -160,7 +160,8 @@ fn arg(rest: []const [:0]const u8, i: usize) ![]const u8 { fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "experts-rows")) return experts_rows_tests.run(gpu); if (std.mem.eql(u8, cmd, "route-rows")) return route_rows_tests.run(gpu); - if (std.mem.eql(u8, cmd, "flashnext-fidelity-single")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); + if (std.mem.eql(u8, cmd, "flashnext-fidelity-single")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), false); + if (std.mem.eql(u8, cmd, "flashnext-fidelity-single-serial")) return flashnext_fidelity_tests.runSingle(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), true); if (std.mem.eql(u8, cmd, "flashnext-fidelity-rank")) return flashnext_fidelity_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "flashnext-window-rank")) return flashnext_window_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "flashnext-model-rank")) return flashnext_model_rank_tests.run(gpu, try arg(rest, 0), try std.fmt.parseInt(usize, try arg(rest, 1), 10), try std.fmt.parseInt(u16, try arg(rest, 2), 10), try arg(rest, 3)); From 8f7af952c2146695dbe51201915b9e30e9f83f9c Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 06:14:13 -0700 Subject: [PATCH 091/115] fix(cuda): fuse Flash Next hidden-state writeback before BF16 rounding --- zig/kernels/cuda/flashnext_hc.cu | 2 +- zig/tests/cuda/flashnext_hc_tests.zig | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/kernels/cuda/flashnext_hc.cu b/zig/kernels/cuda/flashnext_hc.cu index 1ea5886167..2865adde56 100644 --- a/zig/kernels/cuda/flashnext_hc.cu +++ b/zig/kernels/cuda/flashnext_hc.cu @@ -92,7 +92,7 @@ __device__ __forceinline__ void hc_writeback(const __nv_bfloat16* H, __nv_bfloat const long long at = static_cast(r) * S * D + static_cast(s) * D + d; float hv = bf(H[at]); if constexpr (MODE != 0) { - hv = rbf(__fadd_rn(hv, rbf(__fmul_rn(branch, bf(INJ[r * S + s]))))); + hv = rbf(__fmaf_rn(branch, bf(INJ[r * S + s]), hv)); HOUT[at] = __float2bfloat16_rn(hv); } const float ss = warp_sum(__fmul_rn(hv, hv)); diff --git a/zig/tests/cuda/flashnext_hc_tests.zig b/zig/tests/cuda/flashnext_hc_tests.zig index 1c0767849d..974ce95d5e 100644 --- a/zig/tests/cuda/flashnext_hc_tests.zig +++ b/zig/tests/cuda/flashnext_hc_tests.zig @@ -97,7 +97,7 @@ fn writeback(r: *Rig, g: *const flash.Glue, mode: usize) !void { const branch = rbf(acc); for (0..streams) |s| { const at = row * wide + s * dims + d; - const want = rbf(bf(h[at]) + rbf(branch * bf(inj[row * streams + s]))); + const want = rbf(bf(h[at]) + branch * bf(inj[row * streams + s])); try r.within(name, bf(hout[at]), want, ulps(want, 1)); } } From 2e593165b831053f0d0f0afb16f1ae0d399cd2fd Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 06:07:35 -0700 Subject: [PATCH 092/115] cuda: refresh compiler-derived launch resource records --- docs/kernel-contract.md | 4 ++++ zig/src/core/contract/registry.zig | 12 ++++++------ 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index e625bcb2fd..0b80446a42 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -108,3 +108,7 @@ CUDA: the RMS norm rows (16-byte form and the fp32-out tree), Nemotron's router tile, lane_gemv and its split-K form (both bit-equal to it), prompt16's FP4 tile with bf16 and fp32 out, qmmf's FP4 tiles with fp32 out (K slices met in a cluster, fused, or by `reduce_kernel`), and the reduce-scatter's `tf_rank_sum`. + +## CUDA compiler resource baseline + +The registered usage arrays were re-derived with NVCC 13.3.73 (`cuda_13.3.r13.3/compiler.38244171_0`) for SM 80, 86, 89, 120 and 121. `kernel-contract-limits` reads each compiled image and reports its register count, static shared memory and thread bound. Source changes or a different compiler may change these values; re-derive the arrays and run the live audit rather than changing the audit tolerance. Launch dimensions, dynamic shared memory, device limits and planted controls remain independent checks. Only SM 120 was executed on hardware in this update; the other four records come from their compiled images, not hardware execution. diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index a14c3f58a8..d1e2bf0307 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -20,20 +20,20 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK /// CUDA kernels as launched; tf-cuda-test kernel-contract-limits re-derives the usage records. pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, - .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 }, .{ .regs = 45, .static = 132, .bound = 1024 } } }, + .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 } } }, .{ .launch = .{ .name = "router tf_nemo_router", .symbol = "tf_nemo_router", .image = "nemotron_route", .threads = 256, .dynamic = 28800 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, - .{ .launch = .{ .name = "router tf_nemo_topk", .symbol = "tf_nemo_topk", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 }, .{ .regs = 91, .static = 0, .bound = 32 } } }, + .{ .launch = .{ .name = "router tf_nemo_topk", .symbol = "tf_nemo_topk", .image = "nemotron_route", .threads = 32 }, .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "qmm_group cluster of 8", .symbol = group_symbol, .image = "qmm_group", .threads = 128, .dynamic = 35328, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmm_group one CTA a tile", .symbol = group_symbol, .image = "qmm_group", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 54, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "lane_gemv", .symbol = gemv_symbol, .image = "lane_gemv", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 57, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "lane_gemv split-K", .symbol = "tf_lane_gemv_split", .image = "lane_gemv", .threads = 128, .dynamic = 17664 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 44, .static = 128, .bound = 128 }, .{ .regs = 44, .static = 128, .bound = 128 } } }, + .{ .launch = .{ .name = "lane_gemv", .symbol = gemv_symbol, .image = "lane_gemv", .threads = 128, .dynamic = 35328 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 58, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 }, .{ .regs = 66, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "lane_gemv split-K", .symbol = "tf_lane_gemv_split", .image = "lane_gemv", .threads = 128, .dynamic = 17664 }, .usage = .{ .{ .regs = 62, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 56, .static = 128, .bound = 128 }, .{ .regs = 72, .static = 128, .bound = 128 }, .{ .regs = 72, .static = 128, .bound = 128 } } }, .{ .launch = .{ .name = "prompt16 FP4", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 0, 0 }), .image = "prompt16", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "prompt16 FP8", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 1, 0 }), .image = "prompt16", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 218, .static = 0, .bound = 128 }, .{ .regs = 218, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "prompt16 FP8", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 1, 0 }), .image = "prompt16", .threads = 128, .dynamic = 49152 }, .usage = .{ .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 243, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "prompt16 FP4 fp32 out", .symbol = std.fmt.comptimePrint(prompt_symbol, .{ 0, 1 }), .image = "prompt16", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 255, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 }, .{ .regs = 212, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 16-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 16, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 17408 }, .usage = .{ .{ .regs = 88, .static = 0, .bound = 128 }, .{ .regs = 88, .static = 0, .bound = 128 }, .{ .regs = 96, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 }, .{ .regs = 56, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 16-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 16, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 17408, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 }, .{ .regs = 48, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600 }, .usage = .{ .{ .regs = 106, .static = 0, .bound = 128 }, .{ .regs = 106, .static = 0, .bound = 128 }, .{ .regs = 128, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, - .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 63, .static = 0, .bound = 128 }, .{ .regs = 63, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "qmmf FP4 fp32 out 32-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 32, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 25600, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 80, .static = 0, .bound = 128 }, .{ .regs = 61, .static = 0, .bound = 128 }, .{ .regs = 61, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 148, .static = 0, .bound = 128 }, .{ .regs = 180, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out 64-row tile, cluster of 8", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 1, 0 }), .image = "qmmf", .threads = 128, .dynamic = 41984, .cluster = 8, .min_sm = 90 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 }, .{ .regs = 78, .static = 0, .bound = 128 } } }, .{ .launch = .{ .name = "qmmf FP4 fp32 out fused slices", .symbol = std.fmt.comptimePrint(qmmf_f32_symbol, .{ 64, 0, 1 }), .image = "qmmf", .threads = 128, .dynamic = 41984 }, .usage = .{ .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 182, .static = 0, .bound = 128 }, .{ .regs = 228, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 } } }, From b0753faaa59ff555b2674490018af24da7c5ba68 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 06:53:34 -0700 Subject: [PATCH 093/115] style: format native family integration with pinned Zig --- zig/src/core/nvfp4_rows_test.zig | 3 +-- zig/src/native/cuda.zig | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/zig/src/core/nvfp4_rows_test.zig b/zig/src/core/nvfp4_rows_test.zig index 284f13fc57..f6006c08ee 100644 --- a/zig/src/core/nvfp4_rows_test.zig +++ b/zig/src/core/nvfp4_rows_test.zig @@ -19,8 +19,7 @@ fn scale(byte: u8) f64 { const mag = byte & 127; if (mag == 127) return std.math.nan(f64); const e: i32 = @as(i32, mag >> 3) - 7; - const value = if (mag < 8) @as(f64, @floatFromInt(mag)) / 512 else - std.math.ldexp(1 + @as(f64, @floatFromInt(mag & 7)) / 8, e); + const value = if (mag < 8) @as(f64, @floatFromInt(mag)) / 512 else std.math.ldexp(1 + @as(f64, @floatFromInt(mag & 7)) / 8, e); return if (byte & 128 != 0) -value else value; } diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 82135ecab3..e645cf3b5b 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -82,7 +82,6 @@ fn pool(a: Allocator, io: std.Io, ctx: *const cuda.Context, problem: *[]const u8 return .{ .free = free, .total = total, .reserve = reserve, .limit = limit, .unified = unified }; } - /// A lone driver returns false at completion or true when yield hands it to the lane core. const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; From f141cd5dbd50c109977231abce6fe95782b3371f Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:03:30 -0700 Subject: [PATCH 094/115] test(cuda): declare native Flash Next launch resource inventory --- zig/src/core/contract/registry.zig | 34 ++++++++++++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 2 +- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index d1e2bf0307..7e643c4db6 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -19,6 +19,40 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK /// CUDA kernels as launched; tf-cuda-test kernel-contract-limits re-derives the usage records. pub const cuda = [_]CudaRecord{ + .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_gate", .symbol = "fn_attn_gate", .image = "flashnext_attn", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_merge", .symbol = "fn_attn_merge", .image = "flashnext_attn", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_front", .symbol = "fn_gdn_front", .image = "flashnext_batch", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_norm", .symbol = "fn_gdn_norm", .image = "flashnext_batch", .threads = 128 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows", .symbol = "fn_gdn_rows", .image = "flashnext_batch", .threads = 128 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_to", .symbol = "fn_gdn_rows_to", .image = "flashnext_batch", .threads = 128 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_dense_mix", .symbol = "fn_hc_dense_mix", .image = "flashnext_hc", .threads = 32 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_down", .symbol = "fn_hc_down", .image = "flashnext_hc", .threads = 128, .dynamic = 11136 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_normalize", .symbol = "fn_hc_normalize", .image = "flashnext_hc", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_reduce_act", .symbol = "fn_hc_reduce_act", .image = "flashnext_hc", .threads = 352 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_upmix", .symbol = "fn_hc_upmix", .image = "flashnext_hc", .threads = 256, .dynamic = 18688 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback0", .symbol = "fn_hc_writeback0", .image = "flashnext_hc", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_moe", .symbol = "fn_hc_writeback_moe", .image = "flashnext_hc", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_slices", .symbol = "fn_hc_writeback_slices", .image = "flashnext_hc", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_ple_conv", .symbol = "fn_ple_conv", .image = "flashnext_hc", .threads = 512 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_ple_gate", .symbol = "fn_ple_gate", .image = "flashnext_hc", .threads = 512 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_scores", .symbol = "fn_qsa_scores", .image = "flashnext_attn", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_batch fn_qsa_select", .symbol = "fn_qsa_select", .image = "flashnext_batch", .threads = 1024 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_mtp fn_row_best", .symbol = "fn_row_best", .image = "flashnext_mtp", .threads = 1024 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_shared_scatter", .symbol = "fn_shared_scatter", .image = "flashnext_hc", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_hc fn_shared_swiglu", .symbol = "fn_shared_swiglu", .image = "flashnext_hc", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_batch fn_shift_window", .symbol = "fn_shift_window", .image = "flashnext_batch", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_batch fn_sum_slices", .symbol = "fn_sum_slices", .image = "flashnext_batch", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "embed_rows tf_embed_rows_bf16_copies", .symbol = "tf_embed_rows_bf16_copies", .image = "embed_rows", .threads = 256 }, .usage = none }, + .{ .launch = .{ .name = "norms tf_rms_rows_sums", .symbol = "tf_rms_rows_sums", .image = "norms", .threads = 1024 }, .usage = none }, + .{ .launch = .{ .name = "route_rows tf_route_logits", .symbol = "tf_route_logits", .image = "route_rows", .threads = 128, .dynamic = 26112 }, .usage = none }, + .{ .launch = .{ .name = "route_rows tf_route_logits_few", .symbol = "tf_route_logits_few", .image = "route_rows", .threads = 128, .dynamic = 76032 }, .usage = none }, + .{ .launch = .{ .name = "route_rows tf_route_topk", .symbol = "tf_route_topk", .image = "route_rows", .threads = 128 }, .usage = none }, .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 } } }, .{ .launch = .{ .name = "router tf_nemo_router", .symbol = "tf_nemo_router", .image = "nemotron_route", .threads = 256, .dynamic = 28800 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index e1f58336a2..5fe7a57dc3 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -23,7 +23,7 @@ const controls = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { - inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "nemotron_route", "qmm_group", "lane_gemv", "prompt16", "qmmf", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } From af325623f045ebad636ed457bd26be18a60d5e22 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:06:18 -0700 Subject: [PATCH 095/115] test(cuda): qualify hyper-connection writeback through numerical gates --- zig/build/cuda.zig | 9 + zig/src/core/contract/registry.zig | 6 +- .../cuda/contract_flashnext_writeback.zig | 246 ++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 14 + 4 files changed, 272 insertions(+), 3 deletions(-) create mode 100644 zig/tests/cuda/contract_flashnext_writeback.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index ba6369f379..191d1d0dfc 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -260,6 +260,15 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); + const writeback_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_flashnext_writeback.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&writeback_contract.step); + b.step("test-flashnext-contract-host", "Flash Next contract stress and rounding references without a GPU").dependOn(&writeback_contract.step); const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 7e643c4db6..e4dc74cf21 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -21,8 +21,8 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = none }, .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = none }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = none }, .{ .launch = .{ .name = "flashnext_attn fn_attn_gate", .symbol = "fn_attn_gate", .image = "flashnext_attn", .threads = 256 }, .usage = none }, .{ .launch = .{ .name = "flashnext_attn fn_attn_merge", .symbol = "fn_attn_merge", .image = "flashnext_attn", .threads = 256 }, .usage = none }, .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = none }, @@ -51,7 +51,7 @@ pub const cuda = [_]CudaRecord{ .{ .launch = .{ .name = "embed_rows tf_embed_rows_bf16_copies", .symbol = "tf_embed_rows_bf16_copies", .image = "embed_rows", .threads = 256 }, .usage = none }, .{ .launch = .{ .name = "norms tf_rms_rows_sums", .symbol = "tf_rms_rows_sums", .image = "norms", .threads = 1024 }, .usage = none }, .{ .launch = .{ .name = "route_rows tf_route_logits", .symbol = "tf_route_logits", .image = "route_rows", .threads = 128, .dynamic = 26112 }, .usage = none }, - .{ .launch = .{ .name = "route_rows tf_route_logits_few", .symbol = "tf_route_logits_few", .image = "route_rows", .threads = 128, .dynamic = 76032 }, .usage = none }, + .{ .launch = .{ .name = "route_rows tf_route_logits_few", .symbol = "tf_route_logits_few", .image = "route_rows", .threads = 128, .dynamic = 76032, .opt_in = true }, .usage = none }, .{ .launch = .{ .name = "route_rows tf_route_topk", .symbol = "tf_route_topk", .image = "route_rows", .threads = 128 }, .usage = none }, .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 } } }, diff --git a/zig/tests/cuda/contract_flashnext_writeback.zig b/zig/tests/cuda/contract_flashnext_writeback.zig new file mode 100644 index 0000000000..a04842db89 --- /dev/null +++ b/zig/tests/cuda/contract_flashnext_writeback.zig @@ -0,0 +1,246 @@ +//! Hyper-connection writeback: guarded hidden rows, exact fused FP32 rounding and independent FP64 square sums. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const c = core.contract.spec; +const g = core.contract.gates; +const streams = 4; +const slots = 11; + +pub const Mode = enum { sums, moe, slices }; + +fn depth(_: c.Shape) usize { + return 13; +} + +fn spec(mode: Mode) c.Spec { + return .{ + .name = switch (mode) { + .sums => "cuda Flash Next fn_hc_writeback0", + .moe => "cuda Flash Next fn_hc_writeback_moe", + .slices => "cuda Flash Next fn_hc_writeback_slices", + }, + .accum = .f32, + .depth = depth, + .out = .f32, + .tiles = .{ .{ .size = 16 }, .{ .size = 256, .step = 256, .min = 256 }, .{} }, + .served = &.{ .{ .rows = 17, .n = 2560, .k = 2 }, .{ .rows = 3, .n = 512, .k = 8 } }, + }; +} + +/// PSS is first in the output byte stream; the hidden output follows and remains part of the repeat and guard gates. +pub const Writeback = struct { + rig: rig_.Rig, + f: cuda.Function, + mode: Mode, + arena: ?std.heap.ArenaAllocator = null, + shape: c.Shape = undefined, + h: []u16 = &.{}, + inj: []u16 = &.{}, + y: []f32 = &.{}, + wts: []f32 = &.{}, + br: []f32 = &.{}, + dh: u64 = 0, + dinj: u64 = 0, + dy: u64 = 0, + dwts: u64 = 0, + dbr: u64 = 0, + pss: u64 = 0, + hout: u64 = 0, + msg: [200]u8 = undefined, + + pub fn deinit(self: *Writeback) void { + if (self.arena) |*arena| arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *Writeback) g.Kernel { + return .{ .spec = spec(self.mode), .ctx = self, .vt = &g.hooks(Writeback) }; + } + + pub fn prepare(self: *Writeback, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.rows > 33 or shape.n == 0 or shape.n > 2560 or shape.n % 256 != 0 or shape.k == 0 or shape.k > 8) return false; + self.rig.reset(); + if (self.arena) |*arena| arena.deinit(); + self.arena = std.heap.ArenaAllocator.init(self.rig.gpu.gpa); + const a = self.arena.?.allocator(); + self.shape = shape; + const wide = shape.n * streams; + self.h = try a.alloc(u16, shape.rows * wide); + self.inj = try a.alloc(u16, shape.rows * streams); + self.y = try a.alloc(f32, shape.rows * slots * shape.n); + self.wts = try a.alloc(f32, shape.rows * slots); + self.br = try a.alloc(f32, shape.k * shape.rows * shape.n); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (self.h) |*b| b.* = c.toBf16(if (fill == .stress) @floatFromInt(1 + rng.uintLessThan(u32, 64)) else rng.float(f32) * 2 - 1); + for (self.inj) |*b| b.* = c.toBf16(if (fill == .stress) 1 else rng.float(f32) * 4 - 2); + for (self.y) |*v| v.* = if (fill == .stress) 0 else rng.float(f32) * 2 - 1; + for (self.wts) |*v| v.* = rng.float(f32) * 0.2; + for (self.br) |*v| v.* = if (fill == .stress) 0 else rng.float(f32) - 0.5; + if (fill != .stress and self.mode != .sums) { + self.h[0] = 0x3b34; + self.inj[0] = 0x3c9a; + for (0..slots) |k| self.y[k * shape.n] = 0; + self.y[0] = c.fromBf16(0xbe16); + self.wts[0] = 1; + for (0..shape.k) |k| self.br[k * shape.rows * shape.n] = 0; + self.br[0] = c.fromBf16(0xbe16); + } + if (fill == .nonfinite) for (0..shape.rows) |row| { + self.h[row * wide] = switch (row % 5) { + 0 => 0x7fc1, + 1 => 0x7f80, + 2 => 0xff80, + 3 => 0x8000, + else => 0x0001, + }; + }; + self.dh = try self.rig.input(std.mem.sliceAsBytes(self.h)); + self.dinj = try self.rig.input(std.mem.sliceAsBytes(self.inj)); + self.dy = try self.rig.input(std.mem.sliceAsBytes(self.y)); + self.dwts = try self.rig.input(std.mem.sliceAsBytes(self.wts)); + self.dbr = try self.rig.input(std.mem.sliceAsBytes(self.br)); + self.pss = try self.rig.output(self.count() * 4); + self.hout = if (self.mode == .sums) 0 else try self.rig.output(self.h.len * 2); + return true; + } + + pub fn run(self: *Writeback) !void { + const sh = self.shape; + var args: cuda.Args = .{}; + args.add(self.dh); + if (self.mode != .sums) args.add(self.hout); + args.add(self.pss); + switch (self.mode) { + .sums => {}, + .moe => { + args.add(self.dinj); + args.add(self.dy); + args.add(self.dwts); + }, + .slices => { + args.add(self.dbr); + args.add(self.dinj); + args.add(@as(c_longlong, @intCast(sh.rows * sh.n))); + }, + } + args.add(@as(c_int, @intCast(sh.n))); + args.add(@as(c_int, streams)); + if (self.mode == .moe) { + args.add(@as(c_int, slots - 1)); + args.add(@as(c_int, slots)); + } + if (self.mode == .slices) args.add(@as(c_int, @intCast(sh.k))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(sh.rows), .y = @intCast(sh.n / 256) }, .block = .{ .x = 256 } }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Writeback) usize { + return self.shape.rows * (self.shape.n / 256) * streams; + } + + fn hidden(self: *Writeback, row: usize, stream: usize, col: usize) f32 { + const sh = self.shape; + const h = c.fromBf16(self.h[(row * streams + stream) * sh.n + col]); + if (self.mode == .sums) return h; + var branch: f32 = 0; + switch (self.mode) { + .moe => for (0..slots) |k| { + const at = row * slots + k; + branch = @mulAdd(f32, self.y[at * sh.n + col], self.wts[at], branch); + }, + .slices => { + branch = self.br[row * sh.n + col]; + for (1..sh.k) |k| branch += self.br[(k * sh.rows + row) * sh.n + col]; + }, + .sums => unreachable, + } + branch = c.fromBf16(c.toBf16(branch)); + return c.fromBf16(c.toBf16(@mulAdd(f32, branch, c.fromBf16(self.inj[row * streams + stream]), h))); + } + + pub fn reference(self: *Writeback, want: []f64, mag: []f64) !void { + for (0..self.shape.rows) |row| for (0..self.shape.n / 256) |chunk| for (0..streams) |stream| { + var sum: f64 = 0; + for (0..256) |i| { + const v: f64 = self.hidden(row, stream, chunk * 256 + i); + sum += v * v; + } + const at = (row * (self.shape.n / 256) + chunk) * streams + stream; + want[at] = sum; + mag[at] = sum; + }; + } + + pub fn simulate(self: *Writeback, acc: c.Accum, order: g.Order, out: []f64) !void { + for (0..self.shape.rows) |row| for (0..self.shape.n / 256) |chunk| for (0..streams) |stream| { + var terms: [256]f64 = undefined; + for (&terms, 0..) |*term, i| { + const v: f64 = self.hidden(row, stream, chunk * 256 + i); + term.* = v * v; + } + out[(row * (self.shape.n / 256) + chunk) * streams + stream] = core.contract.rows.sumIn(acc, order, &terms); + }; + } + + pub fn check(self: *Writeback, _: g.Gate) !?[]const u8 { + const a = self.arena.?.allocator(); + const pss = try a.dupe(u8, try self.rig.bytes(self.pss, self.count() * 4)); + if (self.mode != .sums) { + const hout = try self.rig.bytes(self.hout, self.h.len * 2); + for (0..self.shape.rows) |row| for (0..streams) |stream| for (0..self.shape.n) |col| { + const at = (row * streams + stream) * self.shape.n + col; + const got = std.mem.readInt(u16, hout[at * 2 ..][0..2], .little); + const want = c.toBf16(self.hidden(row, stream, col)); + if (got != want and !(std.math.isNan(c.fromBf16(got)) and std.math.isNan(c.fromBf16(want)))) + return try std.fmt.bufPrint(&self.msg, "hidden row {d} stream {d} column {d}: {x} != fused FP32 reference {x}", .{ row, stream, col, got, want }); + }; + } + const want = try a.alloc(f64, self.count()); + const mag = try a.alloc(f64, self.count()); + try self.reference(want, mag); + for (want, mag, 0..) |w, m, i| { + const got = c.read(.f32, pss, i); + const ok = if (std.math.isFinite(w)) @abs(got - w) <= c.bound(spec(self.mode), self.shape, w, m) else c.classOf(got) == c.classOf(w); + if (!ok) return try std.fmt.bufPrint(&self.msg, "PSS {d}: {d} differs from independent FP64 sum {d}", .{ i, got, w }); + } + return null; + } +}; + +test "integer square stress requires FP32 sums in every simulated order" { + var prng = std.Random.DefaultPrng.init(0x46e3); + const rng = prng.random(); + var terms: [4][256]f64 = undefined; + var want: [4]f64 = @splat(0); + for (&terms, &want) |*row, *sum| for (row) |*term| { + const v: f64 = @floatFromInt(1 + rng.uintLessThan(u32, 64)); + term.* = v * v; + sum.* += term.*; + }; + for (std.enums.values(g.Order)) |order| { + for (terms, want) |row, w| try std.testing.expectEqual(w, core.contract.rows.sumIn(.f32, order, &row)); + for ([_]c.Accum{ .f16, .bf16 }) |acc| { + var miss: f64 = 0; + for (terms, want) |row, w| { + const got = core.contract.rows.sumIn(acc, order, &row); + const bound = c.boundExact(spec(.sums), w); + miss = @max(miss, if (std.math.isFinite(got)) @abs(got - w) / bound else std.math.inf(f64)); + } + try std.testing.expect(miss >= 100); + } + } +} + +test "the cancellation fixture distinguishes fused update from rounded product" { + const h = c.fromBf16(0x3b34); + const branch = c.fromBf16(0xbe16); + const inject = c.fromBf16(0x3c9a); + const fused = c.toBf16(@mulAdd(f32, branch, inject, h)); + const old = c.toBf16(h + c.fromBf16(c.toBf16(branch * inject))); + const exact: f64 = @as(f64, h) + @as(f64, branch) * @as(f64, inject); + try std.testing.expect(fused != old); + try std.testing.expectEqual(c.toBf16(@floatCast(exact)), fused); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 5fe7a57dc3..0d6f483cee 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -10,6 +10,7 @@ const mm_k = @import("contract_mm.zig"); const mx_k = @import("contract_mxfp8.zig"); const wide_k = @import("contract_projection_f32.zig"); const rank_k = @import("contract_rank.zig"); +const writeback_k = @import("contract_flashnext_writeback.zig"); const contract = core.contract; const g = contract.gates; @@ -104,6 +105,19 @@ pub fn run(gpu: check.Gpu) !void { var rank = rank_k.Rank{ .rig = .{ .gpu = gpu, .s = s }, .sum = &rank_sum }; defer rank.deinit(); try kernels.append(a, rank.kernel()); + var hc = try cuda.Module.load(gpu.d, kk.flashnext_hc); + defer hc.unload(); + var writebacks: [3]writeback_k.Writeback = undefined; + for (std.enums.values(writeback_k.Mode), &writebacks) |mode, *w| { + const symbol: [:0]const u8 = switch (mode) { + .sums => "fn_hc_writeback0", + .moe => "fn_hc_writeback_moe", + .slices => "fn_hc_writeback_slices", + }; + w.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try hc.function(symbol), .mode = mode }; + } + defer for (&writebacks) |*w| w.deinit(); + for (&writebacks) |*w| try kernels.append(a, w.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); From 79480c0d65f88fd5dc1d2823f59cdef64273f8f4 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:12:49 -0700 Subject: [PATCH 096/115] test(cuda): contract ordered slice sums and history copies --- zig/tests/cuda/contract_utility.zig | 167 +++++++++++++ zig/tests/cuda/contract_utility_ref.zig | 303 ++++++++++++++++++++++++ 2 files changed, 470 insertions(+) create mode 100644 zig/tests/cuda/contract_utility.zig create mode 100644 zig/tests/cuda/contract_utility_ref.zig diff --git a/zig/tests/cuda/contract_utility.zig b/zig/tests/cuda/contract_utility.zig new file mode 100644 index 0000000000..0962db2193 --- /dev/null +++ b/zig/tests/cuda/contract_utility.zig @@ -0,0 +1,167 @@ +//! Flash Next's sum and history utilities through the guarded numerical contract. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +pub const ref = @import("contract_utility_ref.zig"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Sum = struct { + rig: rig_.Rig, + f: cuda.Function, + bf16: bool, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + terms: []f32 = &.{}, + part: u64 = 0, + out: u64 = 0, + + pub fn deinit(self: *Sum) void { + self.rig.gpu.gpa.free(self.terms); + self.rig.deinit(); + } + + pub fn kernel(self: *Sum) g.Kernel { + return .{ .spec = ref.sumSpec(self.bf16), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Sum); + + pub fn prepare(self: *Sum, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!ref.sumShape(shape)) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.terms); + self.terms = &.{}; + self.shape = shape; + self.terms = try a.alloc(f32, self.count() * shape.k); + ref.sumFill(self.terms, shape, fill, seed); + self.part = try self.rig.input(std.mem.sliceAsBytes(self.terms)); + self.out = try self.rig.output(self.count() * c.size(ref.sumSpec(self.bf16).out)); + return true; + } + + pub fn run(self: *Sum) !void { + var args: cuda.Args = .{}; + args.add(self.part); + args.add(self.out); + args.add(@as(c_int, if (self.bf16) 1 else 0)); + args.add(@as(c_longlong, @intCast(self.count()))); + args.add(@as(c_int, @intCast(self.shape.k))); + try cuda.launch.launch(self.f, .{ + .grid = .{ .x = @intCast((self.count() + 255) / 256) }, + .block = .{ .x = 256 }, + }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Sum) usize { + return self.shape.rows * self.shape.n; + } + + pub fn reference(self: *Sum, want: []f64, mag: []f64) !void { + ref.sumReference(self.terms, self.shape, want, mag); + } + + pub fn simulate(self: *Sum, acc: c.Accum, order: g.Order, out: []f64) !void { + ref.sumSimulate(self.terms, self.shape, acc, order, out); + } + + /// Retain the float64 error bound while additionally checking ordered f32 and single final RNE. + pub fn check(self: *Sum, at: g.Gate) !?[]const u8 { + return ref.sumCheck(self.rig.gpu.gpa, self.terms, self.shape, self.bf16, try self.rig.outputBytes(), at == .nonfinite); + } +}; + +pub const Shift = struct { + rig: rig_.Rig, + f: cuda.Function, + padded: bool, + shape: c.Shape = .{ .rows = 0, .n = 0, .k = 0 }, + stride: usize = 0, + initial: []u16 = &.{}, + source: []u16 = &.{}, + want: []u16 = &.{}, + win: u64 = 0, + rows: u64 = 0, + + pub fn deinit(self: *Shift) void { + const a = self.rig.gpu.gpa; + a.free(self.initial); + a.free(self.source); + a.free(self.want); + self.rig.deinit(); + } + + pub fn kernel(self: *Shift) g.Kernel { + return .{ .spec = ref.shiftSpec(self.padded), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Shift); + + pub fn prepare(self: *Shift, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!ref.shiftShape(shape)) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.initial); + a.free(self.source); + a.free(self.want); + self.initial = &.{}; + self.source = &.{}; + self.want = &.{}; + self.shape = shape; + self.stride = shape.n + (if (self.padded) @as(usize, 6240) else 0); + self.initial = try a.alloc(u16, self.count()); + self.source = try a.alloc(u16, shape.rows * self.stride); + ref.copyFill(self.initial, self.source, shape, self.stride, fill, seed); + self.want = try ref.copyReference(a, self.initial, self.source, shape, self.stride); + // Mutable history is an output state, never classified as immutable input. + self.win = try self.rig.output(self.count() * 2); + self.rows = try self.rig.input(std.mem.sliceAsBytes(self.source)); + return true; + } + + pub fn run(self: *Shift) !void { + try self.rig.s.synchronize(); + const history: cuda.DeviceBuffer = .{ .d = self.rig.gpu.d, .ptr = self.win, .len = self.count() * 2 }; + // Every repeat starts from the same history; five successive shifts would test another operation. + try history.upload(0, std.mem.sliceAsBytes(self.initial)); + var args: cuda.Args = .{}; + args.add(self.win); + args.add(self.rows); + args.add(@as(c_longlong, @intCast(self.stride))); + args.add(@as(c_int, @intCast(self.shape.rows))); + args.add(@as(c_int, @intCast(self.shape.k))); + args.add(@as(c_int, @intCast(self.shape.n))); + try cuda.launch.launch(self.f, .{ + .grid = .{ .x = @intCast((self.shape.n + 255) / 256) }, + .block = .{ .x = 256 }, + }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Shift) usize { + return self.shape.k * self.shape.n; + } + + pub fn reference(self: *Shift, want: []f64, mag: []f64) !void { + for (want, mag, self.want) |*w, *m, bits| { + w.* = c.fromBf16(bits); + m.* = 0; + } + } + + pub fn check(self: *Shift, at: g.Gate) !?[]const u8 { + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.want), try self.rig.outputBytes())) + return "history bytes differ from append-and-retain oracle"; + if (at == .nonfinite) { + var nonfinite: usize = 0; + for (self.want) |bits| if (!std.math.isFinite(c.fromBf16(bits))) { + nonfinite += 1; + }; + if (nonfinite == 0 or nonfinite == self.want.len) + return "copy fixture did not retain both finite and nonfinite values"; + } + return null; + } +}; diff --git a/zig/tests/cuda/contract_utility_ref.zig b/zig/tests/cuda/contract_utility_ref.zig new file mode 100644 index 0000000000..7dfa1c46b0 --- /dev/null +++ b/zig/tests/cuda/contract_utility_ref.zig @@ -0,0 +1,303 @@ +//! Independent host data and oracles for ordered slice sums and in-place history copies. +const std = @import("std"); +const core = @import("core"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub fn depth(s: c.Shape) usize { + return s.k - 1; +} + +pub fn sumSpec(bf16: bool) c.Spec { + return .{ + .name = if (bf16) "cuda fn_sum_slices bf16" else "cuda fn_sum_slices f32", + .accum = .f32, + .depth = depth, + .out = if (bf16) .bf16 else .f32, + .nan = .{ .stated = "propagates NaN and infinities; finite outputs use ordered f32 and one final RNE" }, + .tiles = .{ .{ .size = 16 }, .{ .size = 256 }, .{ .size = 4 } }, + .served = &.{ .{ .rows = 17, .n = 10240, .k = 4 }, .{ .rows = 4, .n = 2560, .k = 8 } }, + }; +} + +pub fn shiftSpec(padded: bool) c.Spec { + return .{ + .name = if (padded) "cuda fn_shift_window strided" else "cuda fn_shift_window contiguous", + .accum = .none, + .depth = depth, + .out = .bf16, + .in_place = true, + .nan = .{ .stated = "copies all bf16 bits, including NaN payloads, infinities and signed zero" }, + // rows = appended rows, n = channels, k = retained taps; stride is independently padded. + .tiles = .{ .{ .size = 16 }, .{ .size = 256 }, .{ .size = 9 } }, + .served = &.{ + .{ .rows = 5, .n = 10240, .k = 9 }, + .{ .rows = 1, .n = 10240, .k = 3 }, + .{ .rows = 4, .n = 10240, .k = 9 }, + .{ .rows = 9, .n = 257, .k = 9 }, + .{ .rows = 16, .n = 255, .k = 16 }, + .{ .rows = 32, .n = 10240, .k = 9 }, + }, + }; +} + +/// Bounded fixture domains, not validation added to the production launchers. +pub fn sumShape(s: c.Shape) bool { + return s.rows > 0 and s.rows <= 64 and s.n > 0 and s.n <= 16384 and s.k > 0 and s.k <= 8; +} + +pub fn shiftShape(s: c.Shape) bool { + return s.rows > 0 and s.rows <= 64 and s.n > 0 and s.n <= 16384 and s.k > 0 and s.k <= 16; +} + +pub fn sumFill(t: []f32, s: c.Shape, fill: g.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + const e = s.rows * s.n; + for (t) |*v| v.* = (rng.float(f32) * 2 - 1) * + std.math.ldexp(@as(f32, 1), @as(i32, @intCast(rng.uintLessThan(u32, 17))) - 8); + switch (fill) { + .random => {}, + .nonfinite => { + t[0] = std.math.nan(f32); + t[1] = std.math.inf(f32); + t[2] = -std.math.inf(f32); + t[3] = std.math.inf(f32); + if (s.k > 1) t[e + 3] = -std.math.inf(f32); + // The initial term is loaded, not added to a synthetic positive zero. + for (0..s.k) |j| t[j * e + 4] = -0.0; + }, + .stress => for (0..s.k) |j| { + const pattern = [_]f32{ 32768, 1, -32768, 1 }; + for (0..e) |i| t[j * e + i] = pattern[j % 4]; + }, + } +} + +pub fn sumReference(t: []const f32, s: c.Shape, want: []f64, mag: []f64) void { + const e = s.rows * s.n; + for (want, mag, 0..) |*w, *m, i| { + var total: f64 = 0; + var absolute: f64 = 0; + for (0..s.k) |j| { + const term: f64 = t[j * e + i]; + total += term; + absolute += @abs(term); + } + w.* = total; + m.* = absolute; + } +} + +pub fn sumSimulate(t: []const f32, s: c.Shape, acc: c.Accum, order: g.Order, out: []f64) void { + var terms: [8]f64 = undefined; + for (out, 0..) |*v, i| { + for (terms[0..s.k], 0..) |*term, j| term.* = t[j * out.len + i]; + v.* = core.contract.rows.sumIn(acc, order, terms[0..s.k]); + } +} + +/// A separate serial f32 oracle checks the declared slice order and single final conversion. +pub fn ordered(t: []const f32, e: usize, k: usize, i: usize) f32 { + var value = t[i]; + for (1..k) |j| value += t[j * e + i]; + return value; +} + +/// Both the independent real-arithmetic bound and the documented serial f32 order must hold. +pub fn sumCheck(a: std.mem.Allocator, t: []const f32, s: c.Shape, bf16: bool, out: []const u8, nonfinite: bool) !?[]const u8 { + const e = s.rows * s.n; + const want = try a.alloc(f64, e); + defer a.free(want); + const mag = try a.alloc(f64, e); + defer a.free(mag); + sumReference(t, s, want, mag); + const spec = sumSpec(bf16); + if (out.len != e * c.size(spec.out)) return "output byte count differs from the declared shape"; + var planted: usize = 0; + for (want, mag, 0..) |w, m, i| { + const value = ordered(t, e, s.k, i); + const rounded: f32 = if (bf16) c.fromBf16(c.toBf16(value)) else value; + const got = c.read(spec.out, out, i); + if (c.classOf(w) != .finite) planted += 1; + if (c.classOf(got) != c.classOf(w)) return "an output's class differs from the float64 oracle"; + if (std.math.isNan(rounded)) continue; + if (@as(u32, @bitCast(@as(f32, @floatCast(got)))) != @as(u32, @bitCast(rounded))) + return "an output differs from ordered f32 with one final conversion"; + if (std.math.isFinite(w) and !(@abs(got - w) <= c.bound(spec, s, w, m))) + return "an output exceeds the independent float64 error bound"; + } + if (nonfinite and (planted == 0 or planted == e)) + return "nonfinite planting did not retain both finite and nonfinite outputs"; + return null; +} + +pub fn copyFill(win: []u16, rows: []u16, s: c.Shape, stride: usize, fill: g.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (win) |*v| v.* = c.toBf16((rng.float(f32) * 2 - 1) * 32); + // Padding has unique nonfinite poison; it must never become an output channel. + @memset(rows, 0x7fff); + for (0..s.rows) |r| for (rows[r * stride ..][0..s.n]) |*v| { + v.* = c.toBf16((rng.float(f32) * 2 - 1) * 32); + }; + if (fill == .nonfinite) { + const bits = [_]u16{ 0x7fc1, 0x7f80, 0xff80, 0x8000, 0x0000, 0xffc7 }; + // Plant both the retained old history and appended rows. At least one finite channel remains. + for (win, 0..) |*v, i| if (i % 13 < bits.len) { + v.* = bits[i % 13]; + }; + for (0..s.rows) |r| for (rows[r * stride ..][0..s.n], 0..) |*v, ch| { + if (ch % 13 < bits.len) v.* = bits[ch % 13]; + }; + } +} + +/// Append channel rows to a host sequence, then retain its tail; no in-place update algorithm is reused. +pub fn copyReference(a: std.mem.Allocator, win: []const u16, rows: []const u16, s: c.Shape, stride: usize) ![]u16 { + const sequence = try a.alloc(u16, (s.k + s.rows) * s.n); + defer a.free(sequence); + @memcpy(sequence[0..win.len], win); + for (0..s.rows) |r| @memcpy(sequence[win.len + r * s.n ..][0..s.n], rows[r * stride ..][0..s.n]); + return a.dupe(u16, sequence[sequence.len - win.len ..]); +} + +test "ordered sum stress proves fp32 and rejects reduced accumulators in all three orders" { + const a = std.testing.allocator; + inline for (.{ false, true }) |bf16| for ([_]usize{ 4, 8 }) |k| { + const s: c.Shape = .{ .rows = 2, .n = 3, .k = k }; + var terms: [48]f32 = undefined; + const t = terms[0 .. s.rows * s.n * k]; + sumFill(t, s, .stress, 42); + var want: [6]f64 = undefined; + var mag: [6]f64 = undefined; + sumReference(t, s, &want, &mag); + const spec = sumSpec(bf16); + for ([_]g.Order{ .sequential, .pairwise, .strided }) |order| { + for ([_]c.Accum{ .f32, .f16, .bf16 }) |acc| { + var got: [6]f64 = undefined; + sumSimulate(t, s, acc, order, &got); + for (want, got) |w, v| { + const ratio = @abs(v - w) / c.boundExact(spec, w); + if (acc == .f32) try std.testing.expect(ratio <= 1) else try std.testing.expect(ratio >= 100); + } + } + } + const shapes = try c.shapes(a, spec); + defer a.free(shapes); + for (spec.served) |served| try std.testing.expect(sumShape(served)); + }; +} + +test "sum oracle preserves initial negative zero and propagates nonfinite classes" { + const s: c.Shape = .{ .rows = 1, .n = 8, .k = 4 }; + var terms: [32]f32 = undefined; + sumFill(&terms, s, .nonfinite, 71); + try std.testing.expectEqual(c.Class.nan, c.classOf(ordered(&terms, 8, 4, 0))); + try std.testing.expectEqual(c.Class.pos_inf, c.classOf(ordered(&terms, 8, 4, 1))); + try std.testing.expectEqual(c.Class.neg_inf, c.classOf(ordered(&terms, 8, 4, 2))); + try std.testing.expectEqual(c.Class.nan, c.classOf(ordered(&terms, 8, 4, 3))); + try std.testing.expectEqual(@as(u32, 0x80000000), @as(u32, @bitCast(ordered(&terms, 8, 4, 4)))); +} + +test "history oracle covers partial replacement, exact replacement and overflow with padding" { + const a = std.testing.allocator; + const win = [_]u16{ 1, 2, 3, 4, 5, 6 }; + const rows = [_]u16{ 7, 8, 99, 9, 10, 99, 11, 12, 99, 13, 14, 99 }; + for ([_]usize{ 1, 3, 4 }) |n| { + const want = try copyReference(a, &win, &rows, .{ .rows = n, .n = 2, .k = 3 }, 3); + defer a.free(want); + const expected: []const u16 = switch (n) { + 1 => &.{ 3, 4, 5, 6, 7, 8 }, + 3 => &.{ 7, 8, 9, 10, 11, 12 }, + else => &.{ 9, 10, 11, 12, 13, 14 }, + }; + try std.testing.expectEqualSlices(u16, expected, want); + } +} + +test "fixture geometry refuses zero, oversized slice counts and unsafe local history lengths" { + try std.testing.expect(!sumShape(.{ .rows = 1, .n = 1, .k = 0 })); + try std.testing.expect(!sumShape(.{ .rows = 1, .n = 1, .k = 9 })); + try std.testing.expect(!shiftShape(.{ .rows = 0, .n = 1, .k = 1 })); + try std.testing.expect(!shiftShape(.{ .rows = 1, .n = 1, .k = 17 })); + for (shiftSpec(true).served) |s| try std.testing.expect(shiftShape(s)); +} + +test "sum checker rejects wrong order, missing terms and repeated bf16 conversion" { + const a = std.testing.allocator; + const s: c.Shape = .{ .rows = 1, .n = 1, .k = 4 }; + const t = [_]f32{ 16777216, 1, -16777216, 1 }; + const correct = [_]f32{1}; + const reordered = [_]f32{2}; + const lost_tail = [_]f32{0}; + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&correct), false) == null); + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&reordered), false) != null); + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&lost_tail), false) != null); + const small = [_]f32{ 1, 0.00390625, 0.00390625, 0 }; + const once = [_]u16{0x3f81}; + const repeated = [_]u16{0x3f80}; + try std.testing.expect(try sumCheck(a, &small, s, true, std.mem.sliceAsBytes(&once), false) == null); + try std.testing.expect(try sumCheck(a, &small, s, true, std.mem.sliceAsBytes(&repeated), false) != null); +} + +test "history copy retains raw nonfinite payloads and distinguishes padding and repeated mutation" { + const a = std.testing.allocator; + const s: c.Shape = .{ .rows = 1, .n = 3, .k = 3 }; + const win = [_]u16{ 1, 2, 3, 0x7fc1, 0xffc7, 0x8000, 0x7f80, 0xff80, 0 }; + const source = [_]u16{ 4, 5, 6, 0x7fff }; + const once = try copyReference(a, &win, &source, s, 4); + defer a.free(once); + try std.testing.expectEqualSlices(u16, &.{ 0x7fc1, 0xffc7, 0x8000, 0x7f80, 0xff80, 0, 4, 5, 6 }, once); + const twice = try copyReference(a, once, &source, s, 4); + defer a.free(twice); + try std.testing.expect(!std.mem.eql(u16, once, twice)); + try std.testing.expect(std.mem.indexOfScalar(u16, once, 0x7fff) == null); + for ([_]bool{ false, true }) |padded| { + const spec = shiftSpec(padded); + const shapes = try c.shapes(a, spec); + defer a.free(shapes); + for (shapes) |shape| try std.testing.expect(shiftShape(shape)); + } +} + +test "seeded finite and nonfinite slice outputs satisfy both oracles and reject untouched output" { + const a = std.testing.allocator; + const s: c.Shape = .{ .rows = 5, .n = 17, .k = 8 }; + var t: [680]f32 = undefined; + var f: [85]f32 = undefined; + var b: [85]u16 = undefined; + for ([_]g.Fill{ .random, .nonfinite }) |fill| { + sumFill(&t, s, fill, 0x51ce); + for (&f, &b, 0..) |*fv, *bv, i| { + fv.* = ordered(&t, f.len, s.k, i); + bv.* = c.toBf16(fv.*); + } + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&f), fill == .nonfinite) == null); + try std.testing.expect(try sumCheck(a, &t, s, true, std.mem.sliceAsBytes(&b), fill == .nonfinite) == null); + f[f.len - 1] = std.math.nan(f32); + try std.testing.expect(try sumCheck(a, &t, s, false, std.mem.sliceAsBytes(&f), fill == .nonfinite) != null); + } +} + +test "history oracle agrees with independent row-at-a-time queue including all sixteen taps" { + const a = std.testing.allocator; + for ([_]usize{ 1, 3, 9, 16 }) |taps| for ([_]usize{ 1, 5, 16, 32 }) |n| { + const s: c.Shape = .{ .rows = n, .n = 17, .k = taps }; + const stride = s.n + 7; + const win = try a.alloc(u16, taps * s.n); + defer a.free(win); + const source = try a.alloc(u16, n * stride); + defer a.free(source); + copyFill(win, source, s, stride, .nonfinite, 0x7157); + const want = try copyReference(a, win, source, s, stride); + defer a.free(want); + const queue = try a.dupe(u16, win); + defer a.free(queue); + for (0..n) |r| { + std.mem.copyForwards(u16, queue[0 .. queue.len - s.n], queue[s.n..]); + @memcpy(queue[queue.len - s.n ..], source[r * stride ..][0..s.n]); + } + try std.testing.expectEqualSlices(u16, queue, want); + }; +} From e4efbb289a554f7315e68447af761d8138651034 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:13:30 -0700 Subject: [PATCH 097/115] test(cuda): check utility guards at every shape --- zig/tests/cuda/contract_utility.zig | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/zig/tests/cuda/contract_utility.zig b/zig/tests/cuda/contract_utility.zig index 0962db2193..dfe44d065e 100644 --- a/zig/tests/cuda/contract_utility.zig +++ b/zig/tests/cuda/contract_utility.zig @@ -69,6 +69,9 @@ pub const Sum = struct { /// Retain the float64 error bound while additionally checking ordered f32 and single final RNE. pub fn check(self: *Sum, at: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a sum launch changed a guard"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.terms), try self.rig.inputBytes())) + return "a sum launch changed its immutable slice inputs"; return ref.sumCheck(self.rig.gpu.gpa, self.terms, self.shape, self.bf16, try self.rig.outputBytes(), at == .nonfinite); } }; @@ -152,6 +155,9 @@ pub const Shift = struct { } pub fn check(self: *Shift, at: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a history launch changed a guard"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.source), try self.rig.inputBytes())) + return "a history launch changed its immutable source rows"; if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.want), try self.rig.outputBytes())) return "history bytes differ from append-and-retain oracle"; if (at == .nonfinite) { From 0e78b5f484578d019457ab7f354cba517fc24046 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:14:26 -0700 Subject: [PATCH 098/115] test(cuda): run native history and slice utility contracts --- zig/build/cuda.zig | 12 +++++++++++- zig/tests/cuda/contract_flashnext_writeback.zig | 7 +++++++ zig/tests/cuda/contract_tests.zig | 13 +++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 191d1d0dfc..dc43517f40 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -268,7 +268,17 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, }) })); step.dependOn(&writeback_contract.step); - b.step("test-flashnext-contract-host", "Flash Next contract stress and rounding references without a GPU").dependOn(&writeback_contract.step); + const utility_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_utility_ref.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "core", .module = mods.core }}, + }) })); + step.dependOn(&utility_contract.step); + const flash_contract = b.step("test-flashnext-contract-host", "Flash Next contract stress and rounding references without a GPU"); + flash_contract.dependOn(&writeback_contract.step); + flash_contract.dependOn(&utility_contract.step); const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too diff --git a/zig/tests/cuda/contract_flashnext_writeback.zig b/zig/tests/cuda/contract_flashnext_writeback.zig index a04842db89..07ef708054 100644 --- a/zig/tests/cuda/contract_flashnext_writeback.zig +++ b/zig/tests/cuda/contract_flashnext_writeback.zig @@ -186,6 +186,13 @@ pub const Writeback = struct { } pub fn check(self: *Writeback, _: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "a payload guard changed"; + const inputs = try self.rig.inputBytes(); + var offset: usize = 0; + for ([_][]const u8{ std.mem.sliceAsBytes(self.h), std.mem.sliceAsBytes(self.inj), std.mem.sliceAsBytes(self.y), std.mem.sliceAsBytes(self.wts), std.mem.sliceAsBytes(self.br) }) |expected| { + if (!std.mem.eql(u8, inputs[offset..][0..expected.len], expected)) return "an immutable input changed"; + offset += expected.len; + } const a = self.arena.?.allocator(); const pss = try a.dupe(u8, try self.rig.bytes(self.pss, self.count() * 4)); if (self.mode != .sums) { diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 0d6f483cee..95686a5b68 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -11,6 +11,7 @@ const mx_k = @import("contract_mxfp8.zig"); const wide_k = @import("contract_projection_f32.zig"); const rank_k = @import("contract_rank.zig"); const writeback_k = @import("contract_flashnext_writeback.zig"); +const utility_k = @import("contract_utility.zig"); const contract = core.contract; const g = contract.gates; @@ -118,6 +119,18 @@ pub fn run(gpu: check.Gpu) !void { } defer for (&writebacks) |*w| w.deinit(); for (&writebacks) |*w| try kernels.append(a, w.kernel()); + var batch = try cuda.Module.load(gpu.d, kk.flashnext_batch); + defer batch.unload(); + var sums: [2]utility_k.Sum = undefined; + var shifts: [2]utility_k.Shift = undefined; + for ([_]bool{ false, true }, &sums, &shifts) |variant, *sum, *shift| { + sum.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try batch.function("fn_sum_slices"), .bf16 = variant }; + shift.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try batch.function("fn_shift_window"), .padded = variant }; + } + defer for (&sums) |*sum| sum.deinit(); + defer for (&shifts) |*shift| shift.deinit(); + for (&sums) |*sum| try kernels.append(a, sum.kernel()); + for (&shifts) |*shift| try kernels.append(a, shift.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); From 5049af9f3edd622dfe090daf8e0a9a800c393e7d Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:23:58 -0700 Subject: [PATCH 099/115] test(cuda): independently contract sparse attention selection --- zig/tests/cuda/contract_qsa_ref.zig | 329 +++++++++++++++++++++++++ zig/tests/cuda/contract_qsa_select.zig | 102 ++++++++ 2 files changed, 431 insertions(+) create mode 100644 zig/tests/cuda/contract_qsa_ref.zig create mode 100644 zig/tests/cuda/contract_qsa_select.zig diff --git a/zig/tests/cuda/contract_qsa_ref.zig b/zig/tests/cuda/contract_qsa_ref.zig new file mode 100644 index 0000000000..b0e83b5a5a --- /dev/null +++ b/zig/tests/cuda/contract_qsa_ref.zig @@ -0,0 +1,329 @@ +//! Independent sort-based QSA membership, block-order emission and dense-row oracle. +const std = @import("std"); +const core = @import("core"); +const c = core.contract.spec; +const g = core.contract.gates; + +fn depth(_: c.Shape) usize { + return 0; +} + +pub fn spec(padded: bool) c.Spec { + return .{ + .name = if (padded) "cuda fn_qsa_select padded scores" else "cuda fn_qsa_select packed scores", + .accum = .none, + .depth = depth, + .out = .raw, + .nan = .{ .stated = "float total order includes signed NaN payloads and signed zero; ties choose earlier blocks" }, + // rows = query rows, n = first row's complete blocks, k = its tail length plus one. + .tiles = .{ .{ .size = 32 }, .{ .size = 1024, .min = 0 }, .{ .size = 4 } }, + .served = &.{ + .{ .rows = 8, .n = 1025, .k = 1 }, + .{ .rows = 1, .n = 0, .k = 2 }, + .{ .rows = 4, .n = 0, .k = 4 }, + .{ .rows = 8, .n = 512, .k = 1 }, + .{ .rows = 1, .n = 511, .k = 4 }, + .{ .rows = 1, .n = 512, .k = 4 }, + .{ .rows = 1, .n = 513, .k = 1 }, + .{ .rows = 256, .n = 513, .k = 4 }, + .{ .rows = 3, .n = 1023, .k = 2 }, + .{ .rows = 5, .n = 1024, .k = 3 }, + .{ .rows = 7, .n = 1025, .k = 4 }, + .{ .rows = 5, .n = 4097, .k = 2 }, + .{ .rows = 1, .n = 8192, .k = 4 }, + }, + }; +} + +pub const Layout = struct { + rows: usize, + nb: usize, + pos0: i32, + top: usize = 512, + ratio: usize = 4, + idw: usize = 2052, + + pub fn valid(l: Layout) bool { + if (l.rows == 0 or l.rows > 256 or l.nb == 0 or l.nb > 8263 or l.pos0 < 0) return false; + if (l.top == 0 or l.top > 512 or l.ratio == 0 or l.ratio > 32) return false; + if (l.idw < l.top * l.ratio + l.ratio - 1 or l.idw > 16384) return false; + const last = @as(usize, @intCast(l.pos0)) + l.rows; + return last <= std.math.maxInt(i32) and last / l.ratio <= l.nb; + } +}; + +/// Fixture bounds only: the production launcher is unchanged. +pub fn layout(s: c.Shape, padded: bool) ?Layout { + if (s.rows == 0 or s.rows > 256 or s.n > 8192 or s.k == 0 or s.k > 4) return null; + const end = s.n * 4 + s.k - 1; + if (end == 0) return null; + const l: Layout = .{ + .rows = s.rows, + .nb = @max(1, (end + s.rows - 1) / 4) + (if (padded) @as(usize, 7) else 0), + .pos0 = @intCast(end - 1), + }; + return if (l.valid()) l else null; +} + +const Entry = struct { score: f32, index: usize }; + +/// Numeric comparisons plus explicit IEEE sign/NaN payload policy; no radix key transform or scan is reused. +fn before(_: void, a: Entry, b: Entry) bool { + const an = std.math.isNan(a.score); + const bn = std.math.isNan(b.score); + const ab: u32 = @bitCast(a.score); + const bb: u32 = @bitCast(b.score); + const negative_a = std.math.signbit(a.score); + const negative_b = std.math.signbit(b.score); + if (an or bn) { + if (an and bn) { + if (negative_a != negative_b) return !negative_a; + const ap = ab & 0x7fffff; + const bp = bb & 0x7fffff; + if (ap != bp) return if (negative_a) ap < bp else ap > bp; + } else return if (an) !negative_a else negative_b; + } else { + if (a.score != b.score) return a.score > b.score; + if (negative_a != negative_b) return !negative_a; + } + return a.index < b.index; +} + +pub const Result = struct { + ids: []i32, + nk: []i32, + sparse: []i32, + + pub fn deinit(r: Result, a: std.mem.Allocator) void { + a.free(r.ids); + a.free(r.nk); + a.free(r.sparse); + } + + pub fn bytes(r: Result, a: std.mem.Allocator) ![]u8 { + const out = try a.alloc(u8, (r.ids.len + r.nk.len + r.sparse.len) * 4); + var i: usize = 0; + for ([_][]const i32{ r.ids, r.nk, r.sparse }) |part| for (part) |v| { + std.mem.writeInt(i32, out[i..][0..4], v, .little); + i += 4; + }; + return out; + } +}; + +/// Choose membership by a full sort, then enumerate selected blocks in original order and append live tail tokens. +pub fn select(a: std.mem.Allocator, scores: []const f32, l: Layout) !Result { + if (!l.valid() or scores.len != l.rows * l.nb) return error.InvalidFixture; + const ids = try a.alloc(i32, l.rows * l.idw); + errdefer a.free(ids); + const nk = try a.alloc(i32, l.rows); + errdefer a.free(nk); + const sparse = try a.alloc(i32, l.rows); + errdefer a.free(sparse); + @memset(ids, -1); + const entries = try a.alloc(Entry, l.nb); + defer a.free(entries); + const chosen = try a.alloc(bool, l.nb); + defer a.free(chosen); + for (0..l.rows) |r| { + const end = @as(usize, @intCast(l.pos0)) + r + 1; + const complete = end / l.ratio; + if (complete <= l.top) { + nk[r] = @intCast(end); + sparse[r] = 0; + continue; + } + for (entries[0..complete], scores[r * l.nb ..][0..complete], 0..) |*entry, value, b| + entry.* = .{ .score = value, .index = b }; + std.mem.sort(Entry, entries[0..complete], {}, before); + @memset(chosen, false); + for (entries[0..l.top]) |entry| chosen[entry.index] = true; + var at = r * l.idw; + for (chosen[0..complete], 0..) |yes, b| if (yes) { + for (0..l.ratio) |j| { + ids[at] = @intCast(b * l.ratio + j); + at += 1; + } + }; + for (complete * l.ratio..end) |token| { + ids[at] = @intCast(token); + at += 1; + } + nk[r] = @intCast(at - r * l.idw); + sparse[r] = 1; + } + return .{ .ids = ids, .nk = nk, .sparse = sparse }; +} + +pub const special_bits = [_]u32{ + 0xffc00007, 0xff800001, 0xff800000, 0x80000000, 0x00000000, + 0x7f800000, 0x7f800001, 0x7fc00001, 0x7fc00007, +}; + +pub fn fill(scores: []f32, l: Layout, kind: g.Fill, seed: u64) void { + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // Discrete random scores deliberately create ties on both sides of the selection boundary. + for (0..l.rows) |r| { + const complete = (@as(usize, @intCast(l.pos0)) + r + 1) / l.ratio; + for (scores[r * l.nb ..][0..l.nb], 0..) |*v, b| { + v.* = if (b < complete) @floatFromInt(@as(i32, @intCast(rng.uintLessThan(u32, 33))) - 16) else @bitCast(@as(u32, 0x7fffffff)); + if (kind == .nonfinite and b < complete) v.* = switch (r % 4) { + 0 => 0, + 1 => @bitCast(@as(u32, 0x7fc00000) + @as(u32, @intCast(b % 1024 + 1))), + 2 => @bitCast(@as(u32, 0xffc00000) + @as(u32, @intCast(b % 1024 + 1))), + else => if (b % 2 == 0) 0.0 else -0.0, + }; + } + if (kind == .nonfinite) for (special_bits, 0..) |bits, i| { + if (i < complete) scores[r * l.nb + i] = @bitCast(bits); + }; + } +} + +pub fn compare(want: []const u8, got: []const u8) ?[]const u8 { + return if (std.mem.eql(u8, want, got)) null else "selected ids, row counts, sparse flags or untouched padding differ"; +} + +test "independent score order specifies infinities, signed zeros, signed NaNs and payloads" { + var entries: [special_bits.len]Entry = undefined; + for (&entries, special_bits, 0..) |*entry, bits, i| entry.* = .{ .score = @bitCast(bits), .index = i }; + std.mem.sort(Entry, &entries, {}, before); + for (entries, 0..) |entry, i| try std.testing.expectEqual(special_bits.len - 1 - i, entry.index); +} + +test "membership ties choose earliest blocks and output follows block order rather than rank order" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 1, .nb = 5, .pos0 = 6, .top = 2, .ratio = 2, .idw = 6 }; + const result = try select(a, &.{ 1, 9, 9, 999, 999 }, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 2, 3, 4, 5, 6, -1 }, result.ids); + try std.testing.expectEqualSlices(i32, &.{5}, result.nk); + try std.testing.expectEqualSlices(i32, &.{1}, result.sparse); + const ranked = try select(a, &.{ 8, 0, 9, 999, 999 }, l); + defer ranked.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 4, 5, 6, -1 }, ranked.ids); +} + +test "dense-to-sparse transition leaves dense ids untouched and excludes future poisoned scores" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 3, .nb = 6, .pos0 = 4, .top = 2, .ratio = 2, .idw = 6 }; + const scores = [_]f32{ 9, 8, 999, 999, 999, 999, 7, 8, 9, 999, 999, 999, 9, 8, 7, 999, 999, 999 }; + const result = try select(a, &scores, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, -1, -1, -1, -1 }, result.ids[0..6]); + try std.testing.expectEqualSlices(i32, &.{ 2, 3, 4, 5, -1, -1 }, result.ids[6..12]); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 6, -1 }, result.ids[12..]); + try std.testing.expectEqualSlices(i32, &.{ 5, 4, 5 }, result.nk); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 1 }, result.sparse); +} + +test "raw comparison catches omitted tails, wrong tie ids, unused padding and flag corruption" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 1, .nb = 3, .pos0 = 6, .top = 2, .ratio = 2, .idw = 6 }; + const result = try select(a, &.{ 1, 1, 1 }, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 6, -1 }, result.ids); + const want = try result.bytes(a); + defer a.free(want); + const broken = try a.dupe(u8, want); + defer a.free(broken); + for ([_]usize{ 0, 16, 20, 24, 28 }) |byte| { + @memcpy(broken, want); + broken[byte] ^= 1; + try std.testing.expect(compare(want, broken) != null); + } + try std.testing.expect(compare(want, want) == null); +} + +test "nonfinite membership places signed NaNs and signed zero on the correct side of a cutoff" { + const a = std.testing.allocator; + const l: Layout = .{ .rows = 1, .nb = 9, .pos0 = 8, .top = 2, .ratio = 1, .idw = 4 }; + var scores: [9]f32 = undefined; + for (&scores, special_bits) |*value, bits| value.* = @bitCast(bits); + const result = try select(a, &scores, l); + defer result.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 7, 8, -1, -1 }, result.ids); + const zeros = try select(a, &.{ -0.0, 0.0 }, .{ .rows = 1, .nb = 2, .pos0 = 1, .top = 1, .ratio = 1, .idw = 2 }); + defer zeros.deinit(a); + try std.testing.expectEqualSlices(i32, &.{ 1, -1 }, zeros.ids); +} + +test "finite random sort membership agrees with an independent quadratic outranker count" { + const a = std.testing.allocator; + for ([_]usize{ 1, 3, 4, 17, 33, 65 }) |blocks| { + const l: Layout = .{ .rows = 3, .nb = blocks + 8, .pos0 = @intCast(blocks * 2 - 1), .top = 3, .ratio = 2, .idw = 9 }; + const scores = try a.alloc(f32, l.rows * l.nb); + defer a.free(scores); + fill(scores, l, .random, blocks); + const result = try select(a, scores, l); + defer result.deinit(a); + for (0..l.rows) |r| { + const end = @as(usize, @intCast(l.pos0)) + r + 1; + const complete = end / l.ratio; + const ids = result.ids[r * l.idw ..][0..l.idw]; + if (complete <= l.top) { + try std.testing.expectEqual(@as(i32, 0), result.sparse[r]); + for (ids) |id| try std.testing.expectEqual(@as(i32, -1), id); + continue; + } + var at: usize = 0; + const row = scores[r * l.nb ..][0..complete]; + for (row, 0..) |value, b| { + var better: usize = 0; + for (row, 0..) |other, j| if (other > value or (other == value and j < b)) { + better += 1; + }; + if (better < l.top) for (0..l.ratio) |j| { + try std.testing.expectEqual(@as(i32, @intCast(b * l.ratio + j)), ids[at]); + at += 1; + }; + } + for (complete * l.ratio..end) |token| { + try std.testing.expectEqual(@as(i32, @intCast(token)), ids[at]); + at += 1; + } + try std.testing.expectEqual(@as(i32, @intCast(at)), result.nk[r]); + for (ids[at..]) |id| try std.testing.expectEqual(@as(i32, -1), id); + } + } +} + +test "fixture domains retain all served and valid tile edges and explicitly refuse unsafe layouts" { + const a = std.testing.allocator; + for ([_]bool{ false, true }) |padded| { + const s = spec(padded); + for (s.served) |shape| try std.testing.expect(layout(shape, padded) != null); + const edges = try c.shapes(a, s); + defer a.free(edges); + for (edges) |shape| if (shape.k <= 4 and shape.n * 4 + shape.k > 1) + try std.testing.expect(layout(shape, padded) != null); + } + try std.testing.expect(layout(.{ .rows = 1, .n = 1, .k = 5 }, false) == null); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = -1 })); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = 16 })); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = 6, .ratio = 0 })); + try std.testing.expect(!Layout.valid(.{ .rows = 1, .nb = 3, .pos0 = 6, .idw = 2048 })); +} + +test "directed positive NaN payload cutoff spans two blocks and seeded inputs repeat exactly" { + const a = std.testing.allocator; + const l = layout(.{ .rows = 8, .n = 1025, .k = 1 }, true).?; + const scores = try a.alloc(f32, l.rows * l.nb); + defer a.free(scores); + const repeated = try a.alloc(f32, scores.len); + defer a.free(repeated); + fill(scores, l, .nonfinite, 91); + for (0..5) |_| { + fill(repeated, l, .nonfinite, 91); + try std.testing.expectEqualSlices(u8, std.mem.sliceAsBytes(scores), std.mem.sliceAsBytes(repeated)); + } + const result = try select(a, scores, l); + defer result.deinit(a); + // Row 1's 512 largest positive NaN payloads are blocks 512 through 1023. + for (0..512) |b| for (0..4) |j| + try std.testing.expectEqual(@as(i32, @intCast((512 + b) * 4 + j)), result.ids[l.idw + b * 4 + j]); + try std.testing.expectEqual(@as(i32, 4100), result.ids[l.idw + 2048]); + try std.testing.expectEqual(@as(i32, 2049), result.nk[1]); + for (result.ids[l.idw + 2049 .. 2 * l.idw]) |id| try std.testing.expectEqual(@as(i32, -1), id); +} diff --git a/zig/tests/cuda/contract_qsa_select.zig b/zig/tests/cuda/contract_qsa_select.zig new file mode 100644 index 0000000000..e253a8d1a8 --- /dev/null +++ b/zig/tests/cuda/contract_qsa_select.zig @@ -0,0 +1,102 @@ +//! Guarded fn_qsa_select contracts against an independent sort-based selected-list oracle. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +pub const ref = @import("contract_qsa_ref.zig"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Select = struct { + rig: rig_.Rig, + f: cuda.Function, + padded: bool, + l: ref.Layout = .{ .rows = 0, .nb = 0, .pos0 = 0 }, + scores: []f32 = &.{}, + want: []u8 = &.{}, + score_ptr: u64 = 0, + pos_ptr: u64 = 0, + ids_ptr: u64 = 0, + nk_ptr: u64 = 0, + sparse_ptr: u64 = 0, + + pub fn deinit(self: *Select) void { + self.rig.gpu.gpa.free(self.scores); + self.rig.gpu.gpa.free(self.want); + self.rig.deinit(); + } + + pub fn kernel(self: *Select) g.Kernel { + return .{ .spec = ref.spec(self.padded), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(Select); + + pub fn prepare(self: *Select, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + const l = ref.layout(shape, self.padded) orelse return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.scores); + a.free(self.want); + self.scores = &.{}; + self.want = &.{}; + self.l = l; + self.scores = try a.alloc(f32, l.rows * l.nb); + ref.fill(self.scores, l, fill, seed); + const result = try ref.select(a, self.scores, l); + defer result.deinit(a); + self.want = try result.bytes(a); + self.score_ptr = try self.rig.input(std.mem.sliceAsBytes(self.scores)); + self.pos_ptr = try self.rig.input(std.mem.asBytes(&self.l.pos0)); + self.ids_ptr = try self.rig.output(l.rows * l.idw * 4); + self.nk_ptr = try self.rig.output(l.rows * 4); + self.sparse_ptr = try self.rig.output(l.rows * 4); + return true; + } + + pub fn run(self: *Select) !void { + var args: cuda.Args = .{}; + args.add(self.score_ptr); + args.add(self.pos_ptr); + args.add(self.ids_ptr); + args.add(self.nk_ptr); + args.add(self.sparse_ptr); + args.add(@as(c_int, @intCast(self.l.nb))); + args.add(@as(c_int, @intCast(self.l.idw))); + args.add(@as(c_int, @intCast(self.l.top))); + args.add(@as(c_int, @intCast(self.l.ratio))); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(self.l.rows) }, .block = .{ .x = 1024 } }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn count(self: *Select) usize { + return self.want.len; + } + + pub fn reference(self: *Select, want: []f64, mag: []f64) !void { + for (want, mag, self.want) |*w, *m, byte| { + w.* = @floatFromInt(byte); + m.* = 0; + } + } + + pub fn check(self: *Select, at: g.Gate) !?[]const u8 { + if (!try self.rig.intact()) return "selection changed a buffer guard"; + const inputs = try self.rig.inputBytes(); + const score_bytes = std.mem.sliceAsBytes(self.scores); + if (inputs.len != score_bytes.len + 4 or !std.mem.eql(u8, score_bytes, inputs[0..score_bytes.len]) or + !std.mem.eql(u8, std.mem.asBytes(&self.l.pos0), inputs[score_bytes.len..])) + return "selection changed scores or the position input"; + if (at == .nonfinite) { + for (ref.special_bits) |bits| { + var found = false; + for (self.scores) |value| if (@as(u32, @bitCast(value)) == bits) { + found = true; + break; + }; + if (!found) return "nonfinite fixture lacks a required signed NaN, infinity or zero"; + } + } + return ref.compare(self.want, try self.rig.outputBytes()); + } +}; From 63000ad29a533d0b525957312e8592510a8896c9 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:26:08 -0700 Subject: [PATCH 100/115] test(cuda): contract proposal maximum and log-sum-exp independently --- zig/build/cuda.zig | 9 ++ zig/tests/cuda/contract_row_best.zig | 83 ++++++++++++ zig/tests/cuda/contract_row_best_ref.zig | 166 +++++++++++++++++++++++ zig/tests/cuda/contract_tests.zig | 6 + 4 files changed, 264 insertions(+) create mode 100644 zig/tests/cuda/contract_row_best.zig create mode 100644 zig/tests/cuda/contract_row_best_ref.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index dc43517f40..a07ddce786 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -279,6 +279,15 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st const flash_contract = b.step("test-flashnext-contract-host", "Flash Next contract stress and rounding references without a GPU"); flash_contract.dependOn(&writeback_contract.step); flash_contract.dependOn(&utility_contract.step); + const row_best_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_row_best_ref.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "core", .module = mods.core }}, + }) })); + step.dependOn(&row_best_contract.step); + flash_contract.dependOn(&row_best_contract.step); const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too diff --git a/zig/tests/cuda/contract_row_best.zig b/zig/tests/cuda/contract_row_best.zig new file mode 100644 index 0000000000..3050fe90dc --- /dev/null +++ b/zig/tests/cuda/contract_row_best.zig @@ -0,0 +1,83 @@ +//! The baseline proposal row reduction through guards, first-ID checks and the independent float64 LSE oracle. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +pub const ref = @import("contract_row_best_ref.zig"); +const c = core.contract.spec; +const g = core.contract.gates; + +pub const Best = struct { + rig: rig_.Rig, + f: cuda.Function, + shape: c.Shape = undefined, + bits: []u16 = &.{}, + input: u64 = 0, + output: u64 = 0, + probe: u64 = 0, + + pub fn deinit(self: *Best) void { + self.rig.gpu.gpa.free(self.bits); + self.rig.deinit(); + } + + pub fn kernel(self: *Best) g.Kernel { + return .{ .spec = ref.spec, .ctx = self, .vt = &g.hooks(Best) }; + } + + pub fn prepare(self: *Best, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (!ref.admits(shape)) return false; + self.rig.reset(); + const a = self.rig.gpu.gpa; + a.free(self.bits); + self.bits = &.{}; + self.shape = shape; + self.bits = try a.alloc(u16, shape.rows * (shape.k + 17)); + ref.fill(self.bits, shape, fill, seed); + self.input = try self.rig.input(std.mem.sliceAsBytes(self.bits)); + self.output = try self.rig.output(self.count() * 4); + self.probe = try self.rig.scratch(16); + return true; + } + + fn launch(self: *Best, input: u64, output: u64, rows: usize) !void { + var args: cuda.Args = .{}; + args.add(input); + args.add(@as(c_int, @intCast(self.shape.k))); + args.add(@as(c_longlong, @intCast(self.shape.k + 17))); + args.add(output); + try cuda.launch.launch(self.f, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 1024 } }, self.rig.s, &args); + try self.rig.s.synchronize(); + } + + pub fn run(self: *Best) !void { + try self.launch(self.input, self.output, self.shape.rows); + } + + pub fn count(self: *Best) usize { + return self.shape.rows * 4; + } + + pub fn reference(self: *Best, want: []f64, mag: []f64) !void { + ref.reference(self.bits, self.shape, want, mag); + } + + pub fn simulate(self: *Best, acc: c.Accum, order: g.Order, out: []f64) !void { + try ref.simulate(self.rig.gpu.gpa, self.bits, self.shape, acc, order, out); + } + + pub fn check(self: *Best, gate: g.Gate) !?[]const u8 { + const a = self.rig.gpu.gpa; + const actual = try a.dupe(u8, try self.rig.outputBytes()); + defer a.free(actual); + if (ref.check(self.bits, self.shape, actual)) |why| return why; + if (gate == .shapes) for (0..self.shape.rows) |row| { + try self.launch(self.input + row * (self.shape.k + 17) * 2, self.probe, 1); + if (!std.mem.eql(u8, actual[row * 16 ..][0..16], try self.rig.bytes(self.probe, 16))) + return "a batched reduction row differs from the same input reduced alone"; + }; + if (!try self.rig.intact()) return "row reduction changed a payload guard"; + if (!std.mem.eql(u8, std.mem.sliceAsBytes(self.bits), try self.rig.inputBytes())) return "row reduction changed immutable logits"; + return null; + } +}; diff --git a/zig/tests/cuda/contract_row_best_ref.zig b/zig/tests/cuda/contract_row_best_ref.zig new file mode 100644 index 0000000000..c9739c2669 --- /dev/null +++ b/zig/tests/cuda/contract_row_best_ref.zig @@ -0,0 +1,166 @@ +//! Independent first-maximum and stable float64 log-sum-exp oracles for the proposal row reduction. +const std = @import("std"); +const core = @import("core"); +const c = core.contract.spec; +const g = core.contract.gates; + +fn depth(shape: c.Shape) usize { + return (shape.k + 1023) / 1024 + 10; +} + +const served = blk: { + const widths = [_]usize{ 31, 1, 32, 33, 1023, 1024, 1025, 79_591, 248_320 }; + var shapes: [widths.len + 1]c.Shape = undefined; + for (widths, 0..) |k, i| shapes[i] = .{ .rows = 12, .n = 4, .k = k }; + shapes[widths.len] = .{ .rows = 33, .n = 4, .k = 1025 }; + break :blk shapes; +}; + +pub const spec: c.Spec = .{ + .name = "cuda Flash Next fn_row_best", + .accum = .f32, + .depth = depth, + .ops = 2, + .out = .f32, + .nan = .{ .stated = "NaNs do not win; all NaN/-Inf retain -Inf and INT_MAX; a NaN or infinite maximum makes LSE NaN" }, + .tiles = .{ .{ .size = 32 }, .{}, .{ .size = 1024 } }, + .served = &served, +}; + +pub fn admits(shape: c.Shape) bool { + return shape.rows > 0 and shape.rows <= 64 and shape.n == 4 and shape.k > 0 and shape.k <= 248_320; +} + +pub fn fill(bits: []u16, shape: c.Shape, mode: g.Fill, seed: u64) void { + @memset(bits, 0x7fc1); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + for (0..shape.rows) |row| for (0..shape.k) |i| { + bits[row * (shape.k + 17) + i] = if (mode == .stress) 0 else switch (row % 12) { + 0 => if (i == 0 or i + 1 == shape.k) 0x4080 else 0xbf80, + 1 => if (i % 2 == 0) 0x8000 else 0, + 2 => 0x7fc1, + 3 => 0xff80, + 4 => if (i == shape.k / 3 or i + 1 == shape.k) 0x7f80 else 0xff80, + 5 => if (i % 3 == 0) 0x7fc1 else if (i % 3 == 1) 0x3f80 else 0xbf80, + 6 => if (i == shape.k / 2) 0x7f7f else 0xff7f, + 7 => if (i == shape.k / 2) 0x007f else 0x8001, + 8 => c.toBf16(rng.float(f32) * 16 - 8), + 9 => c.toBf16(rng.float(f32) * 2000 - 1000), + 10 => rng.int(u16), + else => @truncate(i *% 40503 +% 317), + }; + }; +} + +pub const Result = struct { value: f32, id: u32, lse: f64 }; + +pub fn expected(bits: []const u16) Result { + var result = Result{ .value = -std.math.inf(f32), .id = std.math.maxInt(i32), .lse = std.math.nan(f64) }; + var nan = false; + for (bits, 0..) |b, i| { + const v = c.fromBf16(b); + nan = nan or std.math.isNan(v); + if (v > result.value) { + result.value = v; + result.id = @intCast(i); + } + } + if (nan or !std.math.isFinite(result.value)) return result; + var sum: f64 = 0; + for (bits) |b| sum += @exp(@as(f64, c.fromBf16(b)) - @as(f64, result.value)); + result.lse = @as(f64, result.value) + @log(sum); + return result; +} + +pub fn reference(bits: []const u16, shape: c.Shape, want: []f64, mag: []f64) void { + for (0..shape.rows) |row| { + const e = expected(bits[row * (shape.k + 17) ..][0..shape.k]); + const at = row * 4; + want[at..][0..4].* = .{ e.value, e.lse, @as(f32, @bitCast(e.id)), 0 }; + mag[at..][0..4].* = .{ 0, 1, 0, 0 }; + } +} + +pub fn simulate(a: std.mem.Allocator, bits: []const u16, shape: c.Shape, acc: c.Accum, order: g.Order, out: []f64) !void { + const terms = try a.alloc(f64, shape.k); + defer a.free(terms); + for (0..shape.rows) |row| { + const input = bits[row * (shape.k + 17) ..][0..shape.k]; + const e = expected(input); + for (input, terms) |b, *term| term.* = @exp(@as(f64, c.fromBf16(b)) - @as(f64, e.value)); + const sum = core.contract.rows.sumIn(acc, order, terms); + const at = row * 4; + out[at..][0..4].* = .{ e.value, @as(f32, @floatCast(@as(f64, e.value) + @log(sum))), @as(f32, @bitCast(e.id)), 0 }; + } +} + +pub fn check(bits: []const u16, shape: c.Shape, output: []const u8) ?[]const u8 { + for (0..shape.rows) |row| { + const e = expected(bits[row * (shape.k + 17) ..][0..shape.k]); + const actual = output[row * 16 ..][0..16]; + if (std.mem.readInt(u32, actual[0..4], .little) != @as(u32, @bitCast(e.value))) return "maximum bits differ from first-maximum oracle"; + if (std.mem.readInt(u32, actual[8..12], .little) != e.id) return "maximum ID differs from independent first-maximum oracle"; + if (std.mem.readInt(u32, actual[12..16], .little) != 0) return "reserved output field was not zeroed"; + const got = c.read(.f32, actual, 1); + if (std.math.isFinite(e.lse)) { + if (!std.math.isFinite(got) or @abs(got - e.lse) > c.bound(spec, shape, e.lse, 1)) return "LSE differs from stable independent float64 oracle"; + } else if (c.classOf(got) != c.classOf(e.lse)) return "LSE nonfinite class differs from the declared formula"; + } + return null; +} + +test "first-maximum oracle preserves first zero bits and excludes all non-winners" { + const zero = expected(&.{ 0x8000, 0 }); + try std.testing.expectEqual(@as(u32, 0x80000000), @as(u32, @bitCast(zero.value))); + try std.testing.expectEqual(@as(u32, 0), zero.id); + const tied = expected(&.{ 0xbf80, 0x4080, 0x4080 }); + try std.testing.expectEqual(@as(u32, 1), tied.id); + for ([_][]const u16{ &.{0x7fc1}, &.{0xff80}, &.{ 0x7fc1, 0xff80 } }) |input| { + const none = expected(input); + try std.testing.expectEqual(@as(u32, std.math.maxInt(i32)), none.id); + try std.testing.expectEqual(c.Class.neg_inf, c.classOf(none.value)); + try std.testing.expect(std.math.isNan(none.lse)); + } + try std.testing.expect(std.math.isNan(expected(&.{ 0x7fc1, 0x3f80 }).lse)); + try std.testing.expect(std.math.isNan(expected(&.{0x7f80}).lse)); +} + +test "equal logits prove FP32 exponential sums and expose both lower precision accumulators" { + const a = std.testing.allocator; + const shape: c.Shape = .{ .rows = 1, .n = 4, .k = 248_320 }; + const bits = try a.alloc(u16, shape.k + 17); + defer a.free(bits); + fill(bits, shape, .stress, 0); + var want: [4]f64 = undefined; + var mag: [4]f64 = undefined; + reference(bits, shape, &want, &mag); + const allow = c.boundExact(spec, want[1]); + for (std.enums.values(g.Order)) |order| { + var got: [4]f64 = undefined; + try simulate(a, bits, shape, .f32, order, &got); + try std.testing.expect(@abs(got[1] - want[1]) <= allow); + for ([_]c.Accum{ .f16, .bf16 }) |acc| { + try simulate(a, bits, shape, acc, order, &got); + try std.testing.expect(!std.math.isFinite(got[1]) or @abs(got[1] - want[1]) / allow >= 100); + } + } +} + +test "host checker rejects a last-tie result and a corrupted log-sum-exp" { + const shape: c.Shape = .{ .rows = 1, .n = 4, .k = 2 }; + var bits: [19]u16 = @splat(0x7fc1); + bits[0..2].* = .{ 0x4080, 0x4080 }; + const e = expected(bits[0..2]); + var out: [16]u8 = @splat(0); + std.mem.writeInt(u32, out[0..4], @bitCast(e.value), .little); + std.mem.writeInt(u32, out[4..8], @bitCast(@as(f32, @floatCast(e.lse))), .little); + try std.testing.expect(check(&bits, shape, &out) == null); + std.mem.writeInt(u32, out[8..12], 1, .little); + try std.testing.expect(check(&bits, shape, &out) != null); + std.mem.writeInt(u32, out[8..12], 0, .little); + std.mem.writeInt(u32, out[4..8], @bitCast(@as(f32, 4)), .little); + try std.testing.expect(check(&bits, shape, &out) != null); + try std.testing.expect(!admits(.{ .rows = 1, .n = 4, .k = 0 })); + try std.testing.expect(!admits(.{ .rows = 1, .n = 4, .k = 248_321 })); +} diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index 95686a5b68..fc7af94f80 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -12,6 +12,7 @@ const wide_k = @import("contract_projection_f32.zig"); const rank_k = @import("contract_rank.zig"); const writeback_k = @import("contract_flashnext_writeback.zig"); const utility_k = @import("contract_utility.zig"); +const best_k = @import("contract_row_best.zig"); const contract = core.contract; const g = contract.gates; @@ -131,6 +132,11 @@ pub fn run(gpu: check.Gpu) !void { defer for (&shifts) |*shift| shift.deinit(); for (&sums) |*sum| try kernels.append(a, sum.kernel()); for (&shifts) |*shift| try kernels.append(a, shift.kernel()); + var mtp = try cuda.Module.load(gpu.d, kk.flashnext_mtp); + defer mtp.unload(); + var best = best_k.Best{ .rig = .{ .gpu = gpu, .s = s }, .f = try mtp.function("fn_row_best") }; + defer best.deinit(); + try kernels.append(a, best.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); From d4673076516756032d77f44bede0f667cb997894 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:27:29 -0700 Subject: [PATCH 101/115] test(cuda): run independent sparse selection contracts --- zig/build/cuda.zig | 9 +++++++++ zig/tests/cuda/contract_tests.zig | 6 ++++++ 2 files changed, 15 insertions(+) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index a07ddce786..39deaf8a4e 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -288,6 +288,15 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st }) })); step.dependOn(&row_best_contract.step); flash_contract.dependOn(&row_best_contract.step); + const qsa_contract = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/contract_qsa_ref.zig"), + .target = host, + .optimize = .debug, + .link_libc = true, + .imports = &.{.{ .name = "core", .module = mods.core }}, + }) })); + step.dependOn(&qsa_contract.step); + flash_contract.dependOn(&qsa_contract.step); const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index fc7af94f80..425a73a37a 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -13,6 +13,7 @@ const rank_k = @import("contract_rank.zig"); const writeback_k = @import("contract_flashnext_writeback.zig"); const utility_k = @import("contract_utility.zig"); const best_k = @import("contract_row_best.zig"); +const qsa_k = @import("contract_qsa_select.zig"); const contract = core.contract; const g = contract.gates; @@ -137,6 +138,11 @@ pub fn run(gpu: check.Gpu) !void { var best = best_k.Best{ .rig = .{ .gpu = gpu, .s = s }, .f = try mtp.function("fn_row_best") }; defer best.deinit(); try kernels.append(a, best.kernel()); + var selections: [2]qsa_k.Select = undefined; + for ([_]bool{ false, true }, &selections) |padded, *selection| + selection.* = .{ .rig = .{ .gpu = gpu, .s = s }, .f = try batch.function("fn_qsa_select"), .padded = padded }; + defer for (&selections) |*selection| selection.deinit(); + for (&selections) |*selection| try kernels.append(a, selection.kernel()); var t: g.Tally = .{}; const t0 = check.now(gpu.io); try g.runAll(a, kernels.items, 0x6b63, &t); From 8874feb54975af859c85f48af8c794e91f859ec8 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:27:47 -0700 Subject: [PATCH 102/115] test(cuda): record compiled native family launch resources --- zig/src/core/contract/registry.zig | 68 +++++++++++++++--------------- 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index e4dc74cf21..14aab5d740 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -19,40 +19,40 @@ const full72_symbol = "_ZN20tf_prefill_attention11full_kernelILi72ELi8ELi8EEEvPK /// CUDA kernels as launched; tf-cuda-test kernel-contract-limits re-derives the usage records. pub const cuda = [_]CudaRecord{ - .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_gate", .symbol = "fn_attn_gate", .image = "flashnext_attn", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_merge", .symbol = "fn_attn_merge", .image = "flashnext_attn", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_batch fn_gdn_front", .symbol = "fn_gdn_front", .image = "flashnext_batch", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_batch fn_gdn_norm", .symbol = "fn_gdn_norm", .image = "flashnext_batch", .threads = 128 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows", .symbol = "fn_gdn_rows", .image = "flashnext_batch", .threads = 128 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_to", .symbol = "fn_gdn_rows_to", .image = "flashnext_batch", .threads = 128 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_dense_mix", .symbol = "fn_hc_dense_mix", .image = "flashnext_hc", .threads = 32 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_down", .symbol = "fn_hc_down", .image = "flashnext_hc", .threads = 128, .dynamic = 11136 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_normalize", .symbol = "fn_hc_normalize", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_reduce_act", .symbol = "fn_hc_reduce_act", .image = "flashnext_hc", .threads = 352 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_upmix", .symbol = "fn_hc_upmix", .image = "flashnext_hc", .threads = 256, .dynamic = 18688 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback0", .symbol = "fn_hc_writeback0", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_moe", .symbol = "fn_hc_writeback_moe", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_slices", .symbol = "fn_hc_writeback_slices", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_ple_conv", .symbol = "fn_ple_conv", .image = "flashnext_hc", .threads = 512 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_ple_gate", .symbol = "fn_ple_gate", .image = "flashnext_hc", .threads = 512 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_attn fn_qsa_scores", .symbol = "fn_qsa_scores", .image = "flashnext_attn", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_batch fn_qsa_select", .symbol = "fn_qsa_select", .image = "flashnext_batch", .threads = 1024 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_mtp fn_row_best", .symbol = "fn_row_best", .image = "flashnext_mtp", .threads = 1024 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_shared_scatter", .symbol = "fn_shared_scatter", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_hc fn_shared_swiglu", .symbol = "fn_shared_swiglu", .image = "flashnext_hc", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_batch fn_shift_window", .symbol = "fn_shift_window", .image = "flashnext_batch", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "flashnext_batch fn_sum_slices", .symbol = "fn_sum_slices", .image = "flashnext_batch", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "embed_rows tf_embed_rows_bf16_copies", .symbol = "tf_embed_rows_bf16_copies", .image = "embed_rows", .threads = 256 }, .usage = none }, - .{ .launch = .{ .name = "norms tf_rms_rows_sums", .symbol = "tf_rms_rows_sums", .image = "norms", .threads = 1024 }, .usage = none }, - .{ .launch = .{ .name = "route_rows tf_route_logits", .symbol = "tf_route_logits", .image = "route_rows", .threads = 128, .dynamic = 26112 }, .usage = none }, - .{ .launch = .{ .name = "route_rows tf_route_logits_few", .symbol = "tf_route_logits_few", .image = "route_rows", .threads = 128, .dynamic = 76032, .opt_in = true }, .usage = none }, - .{ .launch = .{ .name = "route_rows tf_route_topk", .symbol = "tf_route_topk", .image = "route_rows", .threads = 128 }, .usage = none }, + .{ .launch = .{ .name = "flashnext_gdn flashnext_gdn_chain", .symbol = "flashnext_gdn_chain", .image = "flashnext_gdn", .threads = 1024 }, .usage = .{ .{ .regs = 51, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 56, .static = 2060, .bound = 1024 }, .{ .regs = 61, .static = 2060, .bound = 1024 }, .{ .regs = 61, .static = 2060, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_add_streams", .symbol = "fn_add_streams", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 }, .{ .regs = 35, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks", .symbol = "fn_attn_chunks", .image = "flashnext_attn", .threads = 128, .dynamic = 43520, .opt_in = true }, .usage = .{ .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 71, .static = 640, .bound = 128 }, .{ .regs = 73, .static = 640, .bound = 128 }, .{ .regs = 73, .static = 640, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_chunks_ahead", .symbol = "fn_attn_chunks_ahead", .image = "flashnext_attn", .threads = 128, .dynamic = 77312, .opt_in = true }, .usage = .{ .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 }, .{ .regs = 72, .static = 768, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_gate", .symbol = "fn_attn_gate", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 23, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 }, .{ .regs = 22, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_merge", .symbol = "fn_attn_merge", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 28, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 }, .{ .regs = 25, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_attn_prep", .symbol = "fn_attn_prep", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 34, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 }, .{ .regs = 30, .static = 128, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_front", .symbol = "fn_gdn_front", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 26, .static = 1024, .bound = 256 }, .{ .regs = 30, .static = 1024, .bound = 256 }, .{ .regs = 30, .static = 1024, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_norm", .symbol = "fn_gdn_norm", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 26, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 }, .{ .regs = 24, .static = 516, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows", .symbol = "fn_gdn_rows", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_gdn_rows_to", .symbol = "fn_gdn_rows_to", .image = "flashnext_batch", .threads = 128 }, .usage = .{ .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 }, .{ .regs = 64, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_dense_mix", .symbol = "fn_hc_dense_mix", .image = "flashnext_hc", .threads = 32 }, .usage = .{ .{ .regs = 22, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 }, .{ .regs = 36, .static = 128, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_down", .symbol = "fn_hc_down", .image = "flashnext_hc", .threads = 128, .dynamic = 11136 }, .usage = .{ .{ .regs = 56, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 54, .static = 64, .bound = 128 }, .{ .regs = 64, .static = 64, .bound = 128 }, .{ .regs = 64, .static = 64, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_normalize", .symbol = "fn_hc_normalize", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 }, .{ .regs = 40, .static = 16, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_reduce_act", .symbol = "fn_hc_reduce_act", .image = "flashnext_hc", .threads = 352 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 1024 }, .{ .regs = 63, .static = 4096, .bound = 1024 }, .{ .regs = 63, .static = 4096, .bound = 1024 }, .{ .regs = 60, .static = 4096, .bound = 1024 }, .{ .regs = 60, .static = 4096, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_upmix", .symbol = "fn_hc_upmix", .image = "flashnext_hc", .threads = 256, .dynamic = 18688 }, .usage = .{ .{ .regs = 69, .static = 2048, .bound = 256 }, .{ .regs = 64, .static = 2048, .bound = 256 }, .{ .regs = 64, .static = 2048, .bound = 256 }, .{ .regs = 62, .static = 2048, .bound = 256 }, .{ .regs = 62, .static = 2048, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback0", .symbol = "fn_hc_writeback0", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 30, .static = 256, .bound = 256 }, .{ .regs = 37, .static = 256, .bound = 256 }, .{ .regs = 37, .static = 256, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_moe", .symbol = "fn_hc_writeback_moe", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 47, .static = 256, .bound = 256 }, .{ .regs = 47, .static = 256, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_hc_writeback_slices", .symbol = "fn_hc_writeback_slices", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 38, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 }, .{ .regs = 40, .static = 256, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_ple_conv", .symbol = "fn_ple_conv", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 23, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 35, .static = 0, .bound = 512 }, .{ .regs = 36, .static = 0, .bound = 512 }, .{ .regs = 36, .static = 0, .bound = 512 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_ple_gate", .symbol = "fn_ple_gate", .image = "flashnext_hc", .threads = 512 }, .usage = .{ .{ .regs = 32, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 34, .static = 128, .bound = 512 }, .{ .regs = 33, .static = 128, .bound = 512 }, .{ .regs = 33, .static = 128, .bound = 512 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_pool", .symbol = "fn_qsa_pool", .image = "flashnext_attn", .threads = 128 }, .usage = .{ .{ .regs = 32, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 40, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 }, .{ .regs = 39, .static = 1152, .bound = 128 } } }, + .{ .launch = .{ .name = "flashnext_attn fn_qsa_scores", .symbol = "fn_qsa_scores", .image = "flashnext_attn", .threads = 256 }, .usage = .{ .{ .regs = 40, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 46, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 }, .{ .regs = 48, .static = 4096, .bound = 256 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_qsa_select", .symbol = "fn_qsa_select", .image = "flashnext_batch", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 50, .static = 1160, .bound = 1024 }, .{ .regs = 44, .static = 1160, .bound = 1024 }, .{ .regs = 44, .static = 1160, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_mtp fn_row_best", .symbol = "fn_row_best", .image = "flashnext_mtp", .threads = 1024 }, .usage = .{ .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 17, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 }, .{ .regs = 21, .static = 392, .bound = 1024 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_shared_scatter", .symbol = "fn_shared_scatter", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_hc fn_shared_swiglu", .symbol = "fn_shared_swiglu", .image = "flashnext_hc", .threads = 256 }, .usage = .{ .{ .regs = 17, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 }, .{ .regs = 18, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_shift_window", .symbol = "fn_shift_window", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "flashnext_batch fn_sum_slices", .symbol = "fn_sum_slices", .image = "flashnext_batch", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 38, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 } } }, + .{ .launch = .{ .name = "embed_rows tf_embed_rows_bf16_copies", .symbol = "tf_embed_rows_bf16_copies", .image = "embed_rows", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "norms tf_rms_rows_sums", .symbol = "tf_rms_rows_sums", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 31, .static = 132, .bound = 1024 }, .{ .regs = 37, .static = 132, .bound = 1024 }, .{ .regs = 37, .static = 132, .bound = 1024 }, .{ .regs = 34, .static = 132, .bound = 1024 }, .{ .regs = 34, .static = 132, .bound = 1024 } } }, + .{ .launch = .{ .name = "route_rows tf_route_logits", .symbol = "tf_route_logits", .image = "route_rows", .threads = 128, .dynamic = 26112 }, .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 39, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 }, .{ .regs = 35, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "route_rows tf_route_logits_few", .symbol = "tf_route_logits_few", .image = "route_rows", .threads = 128, .dynamic = 76032, .opt_in = true }, .usage = .{ .{ .regs = 31, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 40, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 }, .{ .regs = 34, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "route_rows tf_route_topk", .symbol = "tf_route_topk", .image = "route_rows", .threads = 128 }, .usage = .{ .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 39, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 }, .{ .regs = 36, .static = 64, .bound = 128 } } }, .{ .launch = .{ .name = "norms tf_rms_rows_bf16", .symbol = "tf_rms_rows_bf16", .image = "norms", .threads = 256 }, .usage = .{ .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 79, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 }, .{ .regs = 80, .static = 32, .bound = 256 } } }, .{ .launch = .{ .name = "norms tf_rms_rows", .symbol = "tf_rms_rows", .image = "norms", .threads = 1024 }, .usage = .{ .{ .regs = 32, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 46, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 }, .{ .regs = 43, .static = 132, .bound = 1024 } } }, .{ .launch = .{ .name = "router tf_nemo_router", .symbol = "tf_nemo_router", .image = "nemotron_route", .threads = 256, .dynamic = 28800 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 } } }, From 46836015d7d8fcdab4882aea8075fa69c996d05b Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:44:18 -0700 Subject: [PATCH 103/115] fix(test): retain CUDA contract hooks for the runner lifetime --- .../cuda/contract_flashnext_writeback.zig | 25 ++++++++++++++++++- zig/tests/cuda/contract_row_best.zig | 4 ++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/zig/tests/cuda/contract_flashnext_writeback.zig b/zig/tests/cuda/contract_flashnext_writeback.zig index 07ef708054..3700d915db 100644 --- a/zig/tests/cuda/contract_flashnext_writeback.zig +++ b/zig/tests/cuda/contract_flashnext_writeback.zig @@ -55,8 +55,10 @@ pub const Writeback = struct { self.rig.deinit(); } + const vt = g.hooks(Writeback); + pub fn kernel(self: *Writeback) g.Kernel { - return .{ .spec = spec(self.mode), .ctx = self, .vt = &g.hooks(Writeback) }; + return .{ .spec = spec(self.mode), .ctx = self, .vt = &vt }; } pub fn prepare(self: *Writeback, shape: c.Shape, fill: g.Fill, seed: u64) !bool { @@ -251,3 +253,24 @@ test "the cancellation fixture distinguishes fused update from rounded product" try std.testing.expect(fused != old); try std.testing.expectEqual(c.toBf16(@floatCast(exact)), fused); } + +noinline fn reuseRegistrationStack() usize { + var scratch: [4096]usize = undefined; + const words: *volatile [4096]usize = &scratch; + for (0..scratch.len) |i| words[i] = i *% 73; + return words[42]; +} + +test "registered writeback hooks survive later stack frames" { + var first = Writeback{ .rig = undefined, .f = undefined, .mode = .sums, .shape = .{ .rows = 17, .n = 2560, .k = 2 } }; + var second = Writeback{ .rig = undefined, .f = undefined, .mode = .moe, .shape = .{ .rows = 3, .n = 512, .k = 8 } }; + const Best = @import("contract_row_best.zig").Best; + var best = Best{ .rig = undefined, .f = undefined, .shape = .{ .rows = 4, .n = 4, .k = 31 } }; + const registered = [_]g.Kernel{ first.kernel(), second.kernel(), best.kernel() }; + try std.testing.expectEqual(@as(usize, 3066), reuseRegistrationStack()); + for (registered, [_]usize{ 680, 24, 16 }) |k, expected| { + try std.testing.expectEqual(expected, k.vt.count(k.ctx)); + try std.testing.expect(!try k.vt.prepare(k.ctx, .{ .rows = 0, .n = 256, .k = 1 }, .random, 51)); + try std.testing.expectEqual(expected, k.vt.count(k.ctx)); + } +} diff --git a/zig/tests/cuda/contract_row_best.zig b/zig/tests/cuda/contract_row_best.zig index 3050fe90dc..e63acda285 100644 --- a/zig/tests/cuda/contract_row_best.zig +++ b/zig/tests/cuda/contract_row_best.zig @@ -21,8 +21,10 @@ pub const Best = struct { self.rig.deinit(); } + const vt = g.hooks(Best); + pub fn kernel(self: *Best) g.Kernel { - return .{ .spec = ref.spec, .ctx = self, .vt = &g.hooks(Best) }; + return .{ .spec = ref.spec, .ctx = self, .vt = &vt }; } pub fn prepare(self: *Best, shape: c.Shape, fill: g.Fill, seed: u64) !bool { From 7b94b5e6d20a4a18b7d8d4707e6c640049044958 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:05:49 -0700 Subject: [PATCH 104/115] test: bound accepted TCP stream checks --- zig/src/server/listener.zig | 31 +++++++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/zig/src/server/listener.zig b/zig/src/server/listener.zig index a4d0896e9b..b508c71854 100644 --- a/zig/src/server/listener.zig +++ b/zig/src/server/listener.zig @@ -152,8 +152,7 @@ test "accepted IPv4 and IPv6 TCP streams disable Nagle for incremental responses var storage: Threaded.PosixAddress = undefined; const len = Threaded.addressToPosix(&peer, &storage); if (posix.errno(posix.system.connect(client, &storage.any, len)) != .SUCCESS) return error.ConnectFailed; - const stop = std.atomic.Value(bool).init(false); - const accepted = listener.accept(&stop) orelse return error.AcceptFailed; + const accepted = acceptWithin(listener, 2000) orelse return error.AcceptFailed; defer _ = posix.system.close(accepted.fd); var value: c_int = 0; var size: posix.socklen_t = @sizeOf(c_int); @@ -161,3 +160,31 @@ test "accepted IPv4 and IPv6 TCP streams disable Nagle for incremental responses try std.testing.expect(value != 0); } } + +fn acceptWithin(listener: Listener, timeout_ms: u32) ?Listener.Accepted { + const Watchdog = struct { + stop: std.atomic.Value(bool) = .init(false), + done: std.atomic.Value(bool) = .init(false), + + fn run(w: *@This(), milliseconds: u32) void { + var elapsed: u32 = 0; + while (elapsed < milliseconds and !w.done.load(.acquire)) : (elapsed += 10) { + std.Io.sleep(log.io(), .fromMilliseconds(10), .awake) catch break; + } + if (!w.done.load(.acquire)) w.stop.store(true, .release); + } + }; + var watchdog: Watchdog = .{}; + const thread = std.Thread.spawn(.{}, Watchdog.run, .{ &watchdog, timeout_ms }) catch return null; + defer { + watchdog.done.store(true, .release); + thread.join(); + } + return listener.accept(&watchdog.stop); +} + +test "the accepted-stream test stops when no usable connection arrives" { + const listener = try Listener.open(.{ .ip4 = .loopback(0) }); + defer listener.close(); + try std.testing.expect(acceptWithin(listener, 20) == null); +} From a56a005a1acbb9bf0f76b31e625c28a3d7fb1007 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= <1188270+cwschroeder@users.noreply.github.com> Date: Fri, 16 Oct 2026 16:20:32 +0200 Subject: [PATCH 105/115] flashnext cuda: refuse rank snapshot reuse without draft state --- zig/src/families/flashnext/cuda_lanes.zig | 11 +- .../flashnext/cuda_lanes_restore_test.zig | 144 ++++++++++++++++++ 2 files changed, 154 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/flashnext/cuda_lanes_restore_test.zig diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 89d6095007..6fd33fdafc 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -260,6 +260,8 @@ pub fn Namespace(comptime Tri: type) type { const kept: *snap.Snap = @ptrCast(@alignCast(saved)); if (kept.at != s.reuse.at or kept.at == 0 or kept.at >= ids.len) { s.reuse_failed = true; + } else if (self.ranks != null and l.draft != null) { + s.reuse_failed = true; } else if (if (self.ranks) |r| r.restore(&l.seq, kept) else snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, kept)) |_| { l.filled = kept.at; s.cached = @intCast(kept.at); @@ -548,6 +550,10 @@ pub fn Namespace(comptime Tri: type) type { } } + test { + _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, snapRestore, dropLane); + } + // -- the prompt cache's copies ------------------------------------------------------------------------ fn snapBytes(ptr: *anyopaque, at: u32) u64 { @@ -571,7 +577,10 @@ pub fn Namespace(comptime Tri: type) type { const self = of(ptr); const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); const l = try self.lane(s); - if (self.ranks) |r| return r.restore(&l.seq, @ptrCast(@alignCast(saved))); + if (self.ranks) |r| { + if (l.draft != null) return error.NoHeadState; + return r.restore(&l.seq, @ptrCast(@alignCast(saved))); + } try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); } diff --git a/zig/src/families/flashnext/cuda_lanes_restore_test.zig b/zig/src/families/flashnext/cuda_lanes_restore_test.zig new file mode 100644 index 0000000000..4d2734942b --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes_restore_test.zig @@ -0,0 +1,144 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const prompt = @import("cuda_prompt.zig"); +const mtp = @import("cuda_mtp.zig"); +const snap = @import("cuda_snap.zig"); +const ranks = @import("cuda_rank_engine.zig"); +const wire = @import("cuda_rank_wire.zig"); + +// Virtual device addresses; these callbacks never open CUDA or dereference device memory. +const Device = struct { + var next: u64 = 0x100000; + var copies: usize = 0; + + fn alloc(out: *u64, n: usize) callconv(.c) c_int { + out.* = next; + next += n + 256; + return 0; + } + fn free(_: u64) callconv(.c) c_int { + return 0; + } + fn fill(_: u64, _: u8, _: usize, _: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + fn copy(_: u64, _: u64, _: usize, _: cuda.abi.Stream) callconv(.c) c_int { + copies += 1; + return 0; + } + fn sync(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + fn driver() cuda.Driver { + var d: cuda.Driver = undefined; + d.api.cuMemAlloc_v2 = alloc; + d.api.cuMemFree_v2 = free; + d.api.cuMemsetD8Async = fill; + d.api.cuMemcpyDtoDAsync_v2 = copy; + d.api.cuStreamSynchronize = sync; + return d; + } +}; + +/// Exercise the actual lane open and explicit snapshot entry points with a CPU control peer. +pub fn run(comptime Cuda: type, comptime open: anytype, comptime restore: anytype, comptime drop: anytype, explicit: bool, drafting: bool) !void { + const a = std.testing.allocator; + Device.next = 0x100000; + Device.copies = 0; + var driver = Device.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.gpa = a; + engine.driver = &driver; + engine.stream = &stream; + engine.eos = 7; + engine.chunk = 16; + engine.round_graph = .{}; + var pair: [2]std.posix.socket_t = undefined; + if (std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketPair; + defer _ = std.c.close(pair[0]); + defer _ = std.c.close(pair[1]); + const nonblock: c_int = @bitCast(@as(u32, @bitCast(std.c.O{ .NONBLOCK = true }))); + for (pair) |fd| { + const flags = std.c.fcntl(fd, std.c.F.GETFL); + if (flags < 0 or std.c.fcntl(fd, std.c.F.SETFL, flags | nonblock) < 0) return error.Nonblocking; + } + const peer: wire.Link = .{ .fd = pair[1] }; + // A successful create, followed by a successful restore if the lane attempts it. + try peer.sendHeader(.{ .kind = .ack, .sequence = 1, .object = 1 }, wire.deadline(1000)); + try peer.sendHeader(.{ .kind = .ack, .sequence = 2, .object = 1 }, wire.deadline(1000)); + var rank: ranks.Session = .{ .gpa = a, .rank = 0, .link = .{ .fd = pair[0] }, .collective = undefined, .engine = &engine }; + rank.collective.failed = false; + var ids = [_]u32{1}; + var head: mtp.Head = undefined; + head.ids = &ids; + var self: Cuda = undefined; + self.gpa = a; + self.context = 1024; + self.engine = &engine; + self.stream = stream; + self.ranks = &rank; + self.head = if (drafting) &head else null; + self.lanes = .empty; + defer self.lanes.deinit(a); + var kept: snap.Snap = .{ .rank_id = 99, .buf = undefined, .at = 2, .history = .{ 8, 9 }, .head = false }; + kept.buf.ptr = 0x100; + var s = try lanes.Stream.init(a, .{ .id = "restore", .prompt = &.{ 1, 2, 3, 4 }, .max_new = 1, .drafts = drafting, .reuse = if (explicit) .{} else .{ .saved = &kept, .at = 2 } }); + defer s.deinit(a); + const lane = try open(&self, &s); + defer { + rank.stopped = true; + drop(&self, lane); + } + if (explicit) { + if (drafting) { + try std.testing.expectError(error.NoHeadState, restore(&self, &s, &kept)); + } else try restore(&self, &s, &kept); + } + const created = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectEqual(wire.Kind.create, created.kind); + try std.testing.expectEqual(@as(u64, 1), created.object); + if (drafting) { + try std.testing.expectEqual(@as(usize, 0), lane.seq.fed); + try std.testing.expectEqual(@as(usize, 0), lane.draft.?.len); + try std.testing.expectEqual(@as(usize, 0), lane.filled.?); + try std.testing.expectEqual(@as(u32, 0), s.cached); + try std.testing.expectEqual(!explicit, s.reuse_failed); + try std.testing.expectEqual(@as(u64, 2), rank.next); + try std.testing.expectEqual(@as(usize, 0), Device.copies); + try std.testing.expectEqual([2]i64{ 7, 7 }, lane.seq.history); + } else { + const resumed = try peer.receiveHeader(wire.deadline(1000)); + try std.testing.expectEqual(wire.Kind.restore, resumed.kind); + try std.testing.expectEqual(@as(u64, 1), resumed.object); + try std.testing.expectEqual(kept.rank_id, resumed.arg); + try std.testing.expectEqual(@as(usize, 2), lane.seq.fed); + try std.testing.expect(lane.draft == null); + try std.testing.expect(!s.reuse_failed); + try std.testing.expectEqual(@as(u64, 3), rank.next); + try std.testing.expect(Device.copies > 0); + try std.testing.expectEqual([2]i64{ 8, 9 }, lane.seq.history); + } +} + +/// Each case runs independently in the CUDA family's standard host test root. +pub fn Cases(comptime Cuda: type, comptime open: anytype, comptime restore: anytype, comptime drop: anytype) type { + return struct { + test "rank cache restore: drafted open refuses a plain snapshot before rank restore" { + try run(Cuda, open, restore, drop, false, true); + } + + test "rank cache restore: explicit drafted restore refuses a plain snapshot" { + try run(Cuda, open, restore, drop, true, true); + } + + test "rank cache restore: plain open retains the rank restore command" { + try run(Cuda, open, restore, drop, false, false); + } + + test "rank cache restore: explicit plain restore retains the rank command" { + try run(Cuda, open, restore, drop, true, false); + } + }; +} From df44f518c49ab13e7154f7421f92d18af01667be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= <1188270+cwschroeder@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:19:46 +0200 Subject: [PATCH 106/115] Reject non-finite errors in CUDA reference comparisons --- zig/build/cuda.zig | 12 ++++ zig/tests/cuda/quant_ref_host_tests.zig | 83 +++++++++++++++++++++++++ zig/tests/cuda/quant_ref_tests.zig | 1 + 3 files changed, 96 insertions(+) create mode 100644 zig/tests/cuda/quant_ref_host_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..8aee541747 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -240,6 +240,18 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/build/check_fatbins.zig"), .target = host }) })).step); const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); + const reference_step = b.step("test-cuda-reference", "CUDA's FP64 comparison with CPU-only readback (no GPU)"); + for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |mode| { + const reference = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/quant_ref_host_tests.zig"), + .target = host, + .optimize = mode, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&reference.step); + reference_step.dependOn(&reference.step); + } const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); const native = native_modules.engines; // the HTTP server's unit tests (zig/src/server/root.zig), which the macOS build runs, on Linux hosts too diff --git a/zig/tests/cuda/quant_ref_host_tests.zig b/zig/tests/cuda/quant_ref_host_tests.zig new file mode 100644 index 0000000000..dc13f4235e --- /dev/null +++ b/zig/tests/cuda/quant_ref_host_tests.zig @@ -0,0 +1,83 @@ +//! CPU tests of Rig.close: only stream synchronization and readback are replaced; the comparison is the GPU gate's. +const std = @import("std"); +const cuda = @import("cuda"); +const Rig = @import("quant_ref_tests.zig").RigT; + +const Readback = struct { + var syncs: usize = 0; + var copies: usize = 0; + + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + syncs += 1; + return cuda.abi.success; + } + + fn copy(dst: ?*anyopaque, src: cuda.abi.DevicePtr, n: usize) callconv(.c) cuda.abi.Result { + const out: [*]u8 = @ptrCast(dst.?); + const in: [*]const u8 = @ptrFromInt(src); + @memcpy(out[0..n], in[0..n]); + copies += 1; + return cuda.abi.success; + } +}; + +fn close(comptime T: type, got: []const T, want: []const f64, tol: []const f64) !void { + var d: cuda.Driver = undefined; + d.api.cuStreamSynchronize = Readback.sync; + d.api.cuMemcpyDtoH_v2 = Readback.copy; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var rig: Rig = .{ + .gpu = .{ .d = &d, .ctx = undefined, .gpa = std.testing.allocator, .io = std.testing.io }, + .a = arena.allocator(), + .s = .{ .d = &d, .handle = null }, + .rng = undefined, + }; + Readback.syncs = 0; + Readback.copies = 0; + const result = rig.close("CPU reference comparison", @intFromPtr(got.ptr), T == f32, want, tol); + try std.testing.expectEqual(@as(usize, 1), Readback.syncs); + try std.testing.expectEqual(@as(usize, 1), Readback.copies); + return result; +} + +const zero = [_]f64{ 0, 0, 0 }; +const tolerance = [_]f64{ 1, 1, 1 }; + +test "finite FP32 and BF16 outputs pass at the bound and fail beyond it" { + try close(f32, &.{ 0.5, 1, 0 }, &zero, &tolerance); + try close(u16, &.{ 0x3f00, 0x3f80, 0 }, &zero, &tolerance); + try std.testing.expectError(error.TestFailed, close(f32, &.{ 2, 0, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0x4000, 0, 0 }, &zero, &tolerance)); +} + +test "FP32 NaN cannot disappear before a later finite result" { + for ([_]u32{ 0x7fc00000, 0xffffffff }) |bits| { + for (0..3) |i| { + var got = [_]f32{ 0, 0, 0 }; + got[i] = @bitCast(bits); + try std.testing.expectError(error.TestFailed, close(f32, &got, &zero, &tolerance)); + } + } +} + +test "BF16 NaN cannot disappear before a later finite result" { + for ([_]u16{ 0x7fc0, 0xffff }) |bits| { + for (0..3) |i| { + var got = [_]u16{ 0, 0, 0 }; + got[i] = bits; + try std.testing.expectError(error.TestFailed, close(u16, &got, &zero, &tolerance)); + } + } +} + +test "NaN cannot erase an earlier out-of-bound or infinite result" { + try std.testing.expectError(error.TestFailed, close(f32, &.{ 0, std.math.inf(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(f32, &.{ 0, -std.math.inf(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0, 0x7f80, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0, 0xff80, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(f32, &.{ 2, std.math.nan(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(f32, &.{ std.math.inf(f32), std.math.nan(f32), 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0x4000, 0xffff, 0 }, &zero, &tolerance)); + try std.testing.expectError(error.TestFailed, close(u16, &.{ 0x7f80, 0xffff, 0 }, &zero, &tolerance)); +} diff --git a/zig/tests/cuda/quant_ref_tests.zig b/zig/tests/cuda/quant_ref_tests.zig index 657cd388b7..a0f40b664c 100644 --- a/zig/tests/cuda/quant_ref_tests.zig +++ b/zig/tests/cuda/quant_ref_tests.zig @@ -58,6 +58,7 @@ const Rig = struct { for (want, tol, 0..) |w, t, i| { const g: f64 = if (f32_out) std.mem.bytesToValue(f32, raw[4 * i ..][0..4]) else bf(std.mem.bytesToValue(u16, raw[2 * i ..][0..2])); const q = @abs(g - w) / t; + try check.expect(std.math.isFinite(q), "{s}: non-finite error at element {d} (got {e}, reference {e}, tolerance {e})", .{ name, i, g, w, t }); if (!(q <= most)) { most = q; at = i; From 37359a17e1be44fd50cd0d7b4086b07a03314601 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Schr=C3=B6der?= <1188270+cwschroeder@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:41:51 +0200 Subject: [PATCH 107/115] test: poison each FP32 projection dispatch and reduction scratch --- zig/tests/cuda/projection_f32_tests.zig | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/zig/tests/cuda/projection_f32_tests.zig b/zig/tests/cuda/projection_f32_tests.zig index 85b1160365..19a62fdf4f 100644 --- a/zig/tests/cuda/projection_f32_tests.zig +++ b/zig/tests/cuda/projection_f32_tests.zig @@ -80,9 +80,15 @@ fn projection(r: *Rig, lane: *cuda.qmmf.Lane, mode: Mode) !void { for (x[row * ldx ..][0..k]) |*v| v.* = bf16((r.rng.float(f32) - 0.5) * 2); } const dx = try r.dev(std.mem.sliceAsBytes(x)); - const out = try r.zeros(m * n * 4); - const out_bf16 = try r.zeros(m * n * 2); - const part = try r.zeros(8 * m * n * 4); + var out_buf = try cuda.DeviceBuffer.alloc(r.gpu.d, m * n * 4); + defer out_buf.free(); + var bf16_buf = try cuda.DeviceBuffer.alloc(r.gpu.d, m * n * 2); + defer bf16_buf.free(); + var part_buf = try cuda.DeviceBuffer.alloc(r.gpu.d, 8 * m * n * 4); + defer part_buf.free(); + const out = out_buf.ptr; + const out_bf16 = bf16_buf.ptr; + const part = part_buf.ptr; const want = try a.alloc(f64, m * n); const tol = try a.alloc(f64, m * n); for (0..m) |row| for (0..n) |col| { @@ -111,10 +117,14 @@ fn projection(r: *Rig, lane: *cuda.qmmf.Lane, mode: Mode) !void { for ([_]bool{ true, false }) |clusters| { lane.clusters = clusters and (try r.gpu.ctx.capability()) >= 90; for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + try out_buf.fill8(0xff, r.s.handle); + try part_buf.fill8(0xff, r.s.handle); try lane.matmulF32(r.s, dx, ldx, rows, w, out, part); try r.close("FP32 projection", out, true, want[0 .. rows * n], tol[0 .. rows * n]); const got = try r.back(out, rows * n * 4); if (baseline) |b| try check.sameBytes("FP32 row/reduction invariance", got, b[0..got.len]) else baseline = got; + try bf16_buf.fill8(0xff, r.s.handle); + try part_buf.fill8(0xff, r.s.handle); try lane.matmul(r.s, dx, ldx, rows, w, out_bf16, part); const rounded = try r.back(out_bf16, rows * n * 2); for (0..rows * n) |i| { From 4544c5c5da7c5a9b6ffa0cfa46275e32cb011397 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:08:15 -0700 Subject: [PATCH 108/115] test(cuda): reject no-op projection dispatches after stale prefixes --- zig/build/cuda.zig | 10 ++ zig/tests/cuda/projection_f32_host_tests.zig | 109 +++++++++++++++++++ zig/tests/cuda/projection_f32_tests.zig | 16 +-- 3 files changed, 128 insertions(+), 7 deletions(-) create mode 100644 zig/tests/cuda/projection_f32_host_tests.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 8aee541747..ee6685bf9d 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -241,6 +241,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st const cuda = runtime(b, host, .debug, &.{}); const mods = family(b, host, .debug, cuda, draft_ids); const reference_step = b.step("test-cuda-reference", "CUDA's FP64 comparison with CPU-only readback (no GPU)"); + const projection_step = b.step("test-cuda-projection-fixture", "Projection fixture dispatch controls with CPU-only CUDA calls"); for ([_]std.builtin.OptimizeMode{ .debug, .fast }) |mode| { const reference = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/quant_ref_host_tests.zig"), @@ -251,6 +252,15 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st }) })); step.dependOn(&reference.step); reference_step.dependOn(&reference.step); + const projection = b.addRunArtifact(b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/cuda/projection_f32_host_tests.zig"), + .target = host, + .optimize = mode, + .link_libc = true, + .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } }, + }) })); + step.dependOn(&projection.step); + projection_step.dependOn(&projection.step); } const native_modules = engines(b, host, .debug, cuda, mods.core, mods.lanes, mods.nemotron, mods.heat); const native = native_modules.engines; diff --git a/zig/tests/cuda/projection_f32_host_tests.zig b/zig/tests/cuda/projection_f32_host_tests.zig new file mode 100644 index 0000000000..5d38e2e29c --- /dev/null +++ b/zig/tests/cuda/projection_f32_host_tests.zig @@ -0,0 +1,109 @@ +//! Exercise the fixture's actual poison-and-dispatch path with CPU buffers and no-op CUDA launches. +const std = @import("std"); +const cuda = @import("cuda"); +const fixture = @import("projection_f32_tests.zig"); +const Rig = @import("quant_ref_tests.zig").RigT; + +const Mock = struct { + var destination: []u8 = &.{}; + var scratch: []u8 = &.{}; + var fills: usize = 0; + var launches: usize = 0; + var fresh = true; + + fn fill(ptr: cuda.abi.DevicePtr, byte: u8, n: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + const bytes: [*]u8 = @ptrFromInt(ptr); + @memset(bytes[0..n], byte); + fills += 1; + return cuda.abi.success; + } + + fn observe() cuda.abi.Result { + for (destination) |byte| fresh = fresh and byte == 0xff; + for (scratch) |byte| fresh = fresh and byte == 0xff; + launches += 1; + return cuda.abi.success; + } + + fn plain(_: cuda.abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: cuda.abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + return observe(); + } + + fn extended(_: *const cuda.abi.LaunchConfig, _: cuda.abi.Function, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) cuda.abi.Result { + return observe(); + } + + fn sync(_: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + return cuda.abi.success; + } + + fn copy(dst: ?*anyopaque, src: cuda.abi.DevicePtr, n: usize) callconv(.c) cuda.abi.Result { + const out: [*]u8 = @ptrCast(dst.?); + const in: [*]const u8 = @ptrFromInt(src); + @memcpy(out[0..n], in[0..n]); + return cuda.abi.success; + } +}; + +test "every F32 and BF16 fixture dispatch rejects a no-op after a valid stale prefix" { + const a = std.testing.allocator; + var driver: cuda.Driver = undefined; + driver.api.cuMemsetD8Async = Mock.fill; + driver.api.cuLaunchKernel = Mock.plain; + driver.api.cuLaunchKernelEx = Mock.extended; + driver.api.cuStreamSynchronize = Mock.sync; + driver.api.cuMemcpyDtoH_v2 = Mock.copy; + const stream: cuda.Stream = .{ .d = &driver, .handle = @ptrFromInt(1) }; + const f: cuda.Function = .{ .d = &driver, .handle = null }; + var lane: cuda.qmmf.Lane = .{ + .mod = .{ .d = &driver, .handle = null }, + .fns = @splat(@splat(f)), + .f32_fns = @splat(@splat(f)), + .reduce = @splat(f), + .clusters = false, + }; + const weight: cuda.qmmf.Weight = .{ .mode = .fp4, .codes = 0, .scales = 0, .n = 70, .k = 1024, .npad = 128 }; + const elements = 256 * weight.n; + const expected = try a.alloc(f64, elements); + defer a.free(expected); + const tolerance = try a.alloc(f64, elements); + defer a.free(tolerance); + @memset(expected, 0); + @memset(tolerance, 1); + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var rig: Rig = .{ + .gpu = .{ .d = &driver, .ctx = undefined, .gpa = a, .io = std.testing.io }, + .a = arena.allocator(), + .s = stream, + .rng = undefined, + }; + for ([_]cuda.qmmf.Out{ .f32, .bf16 }) |format| { + const width: usize = if (format == .f32) 4 else 2; + const destination = try a.alloc(u8, elements * width); + defer a.free(destination); + const scratch = try a.alloc(u8, 8 * elements * 4); + defer a.free(scratch); + const out: cuda.DeviceBuffer = .{ .d = &driver, .ptr = @intFromPtr(destination.ptr), .len = destination.len }; + const part: cuda.DeviceBuffer = .{ .d = &driver, .ptr = @intFromPtr(scratch.ptr), .len = scratch.len }; + Mock.destination = destination; + Mock.scratch = scratch; + for ([_]bool{ true, false }) |clusters| { + lane.clusters = clusters; + for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { + @memset(destination, 0); + @memset(scratch, 0); + const count = rows * weight.n; + try rig.close("valid stale prefix", out.ptr, format == .f32, expected[0..count], tolerance[0..count]); + Mock.fills = 0; + Mock.launches = 0; + Mock.fresh = true; + try fixture.dispatch(&lane, stream, 0, 1088, rows, weight, out, part, format); + try std.testing.expectEqual(@as(usize, 2), Mock.fills); + try std.testing.expect(Mock.launches > 0); + try std.testing.expect(Mock.fresh); + try std.testing.expectError(error.TestFailed, rig.close("no-op projection", out.ptr, format == .f32, expected[0..count], tolerance[0..count])); + } + } + } +} diff --git a/zig/tests/cuda/projection_f32_tests.zig b/zig/tests/cuda/projection_f32_tests.zig index 19a62fdf4f..3edd66a243 100644 --- a/zig/tests/cuda/projection_f32_tests.zig +++ b/zig/tests/cuda/projection_f32_tests.zig @@ -9,6 +9,13 @@ const bf = reference.helpers.bfOf; const bf16 = reference.helpers.bf16Of; const Mode = cuda.qmmf.Mode; +/// One fixture dispatch starts with fresh destination and reduction poison, including a reused row prefix. +pub fn dispatch(lane: *const cuda.qmmf.Lane, s: cuda.Stream, x: u64, ldx: usize, rows: usize, w: cuda.qmmf.Weight, out: cuda.DeviceBuffer, part: cuda.DeviceBuffer, format: cuda.qmmf.Out) !void { + try out.fill8(0xff, s.handle); + try part.fill8(0xff, s.handle); + try lane.matmulOut(s, x, ldx, rows, w, out.ptr, format, part.ptr); +} + pub fn run(gpu: check.Gpu) !void { var arena = std.heap.ArenaAllocator.init(gpu.gpa); defer arena.deinit(); @@ -88,7 +95,6 @@ fn projection(r: *Rig, lane: *cuda.qmmf.Lane, mode: Mode) !void { defer part_buf.free(); const out = out_buf.ptr; const out_bf16 = bf16_buf.ptr; - const part = part_buf.ptr; const want = try a.alloc(f64, m * n); const tol = try a.alloc(f64, m * n); for (0..m) |row| for (0..n) |col| { @@ -117,15 +123,11 @@ fn projection(r: *Rig, lane: *cuda.qmmf.Lane, mode: Mode) !void { for ([_]bool{ true, false }) |clusters| { lane.clusters = clusters and (try r.gpu.ctx.capability()) >= 90; for ([_]usize{ 256, 1, 7, 17, 33 }) |rows| { - try out_buf.fill8(0xff, r.s.handle); - try part_buf.fill8(0xff, r.s.handle); - try lane.matmulF32(r.s, dx, ldx, rows, w, out, part); + try dispatch(lane, r.s, dx, ldx, rows, w, out_buf, part_buf, .f32); try r.close("FP32 projection", out, true, want[0 .. rows * n], tol[0 .. rows * n]); const got = try r.back(out, rows * n * 4); if (baseline) |b| try check.sameBytes("FP32 row/reduction invariance", got, b[0..got.len]) else baseline = got; - try bf16_buf.fill8(0xff, r.s.handle); - try part_buf.fill8(0xff, r.s.handle); - try lane.matmul(r.s, dx, ldx, rows, w, out_bf16, part); + try dispatch(lane, r.s, dx, ldx, rows, w, bf16_buf, part_buf, .bf16); const rounded = try r.back(out_bf16, rows * n * 2); for (0..rows * n) |i| { const v = std.mem.readInt(u32, got[i * 4 ..][0..4], .little); From f859ed4ead62a6fc64fdeaef1018a4831c2e55ea Mon Sep 17 00:00:00 2001 From: cwschroeder <1188270+cwschroeder@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:08:20 -0700 Subject: [PATCH 109/115] test(server): cover streamed admission and drained health regressions --- zig/src/native/flashnext_host.zig | 22 +++++++++++++++ zig/src/server/health_test.zig | 34 ++++++++++++++++++++++-- zig/src/server/stream_preflight_test.zig | 18 ++++++++----- 3 files changed, 66 insertions(+), 8 deletions(-) diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index 8281b19df4..ddb3f74a68 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -534,3 +534,25 @@ test "status after a reply finishes reads no freed job" { try std.testing.expectEqual(@as(u32, 0), st.running); try std.testing.expectEqual(@as(usize, 0), st.streams); } + +test "a drained Flash Next host is unhealthy before another submit" { + const gpa = std.testing.allocator; + var host: Host = .{ .gpa = gpa, .io = std.testing.io, .eng = undefined, .info_ = .{} }; + defer host.decoded.deinit(gpa); + const engine = host.engine(); + var status: api.Status = .{}; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + engine.drain(); + engine.status(&status, &.{}); + const Sink = struct { + fn event(_: *anyopaque, _: api.Id, _: *const api.Event) void { + @panic("a drained host must refuse admission"); + } + }; + var ctx: u8 = 0; + const request: api.Request = .{ .prompt = &.{1}, .max_tokens = 1 }; + try std.testing.expectError(error.Closed, engine.submit(1, &request, .{ .ctx = &ctx, .event = Sink.event })); + try std.testing.expectEqual(@as(usize, 0), host.queued.items.len); + try std.testing.expect(!status.healthy); +} diff --git a/zig/src/server/health_test.zig b/zig/src/server/health_test.zig index 272933be59..be31991c49 100644 --- a/zig/src/server/health_test.zig +++ b/zig/src/server/health_test.zig @@ -17,6 +17,7 @@ const earlier = [_][]const u8{ "status", "model", "model_ids", "max_batch_size", /// Four lanes, some memory, nothing running. const Idle = struct { + healthy: bool = true, fn engine(e: *Idle) api.Engine { return .{ .ctx = e, .vtable = &.{ .info = info, @@ -37,8 +38,8 @@ const Idle = struct { fn cancel(_: *anyopaque, _: api.Id) void {} - fn status(_: *anyopaque, out: *api.Status, _: []u32) void { - out.* = .{}; + fn status(ctx: *anyopaque, out: *api.Status, _: []u32) void { + out.* = .{ .healthy = @as(*Idle, @ptrCast(@alignCast(ctx))).healthy }; } fn memory(_: *anyopaque, _: bool) ?api.Memory { @@ -265,3 +266,32 @@ test "with keys configured, /health sends only its status, with and without a va try std.testing.expect(std.mem.startsWith(u8, try exchange(a, port, models_open), "HTTP/1.1 401 ")); try std.testing.expect(std.mem.startsWith(u8, try exchange(a, port, models_keyed), "HTTP/1.1 200 ")); } + +test "unavailable health is minimal HTTP 503 with keys enabled or disabled and with or without a valid key" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var problem: []const u8 = ""; + var keys = try auth.Store.init(std.testing.allocator, &.{key}, "", null, false, &problem); + defer keys.deinit(); + for ([_]bool{ false, true }) |keyed| { + var engine: Idle = .{ .healthy = false }; + var text: Bytes = .{}; + const srv = try server.Server.init(std.testing.allocator, std.testing.io, engine.engine(), text.text(), .{ + .served_name = "private-model", + .model_ids = &.{"private-model"}, + .identity = .{ .version = "9.8.7", .backend = "test", .device = "private-device" }, + }, if (keyed) &keys else null); + defer srv.deinit(); + const lis = try listener.Listener.open(.{ .ip4 = .loopback(0) }); + defer lis.close(); + var stop: std.atomic.Value(bool) = .init(false); + const accept = try std.Thread.spawn(.{}, server.Server.serve, .{ srv, lis, &stop }); + defer quiesce(srv, &stop, accept); + for ([_][]const u8{ health_open, health_keyed }) |request| { + const reply = try exchange(a, lis.port(), request); + try std.testing.expect(std.mem.startsWith(u8, reply, "HTTP/1.1 503 ")); + try std.testing.expectEqualStrings("{\"status\": \"unavailable\"}", try bodyOf(reply)); + } + } +} diff --git a/zig/src/server/stream_preflight_test.zig b/zig/src/server/stream_preflight_test.zig index 25e4024d07..8bf11643a0 100644 --- a/zig/src/server/stream_preflight_test.zig +++ b/zig/src/server/stream_preflight_test.zig @@ -473,14 +473,20 @@ test "closed or busy admission is service unavailable across chat, messages and .{ .route = .messages, .body = "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":1}" }, .{ .route = .responses, .body = "{\"model\":\"test-model\",\"input\":\"x\",\"max_output_tokens\":1}" }, }; + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + var wrong: usize = 0; var engine: ClosedEngine = .{}; for ([_]api.SubmitError{ error.Closed, error.Busy }) |reason| { engine.reason = reason; - for (cases) |c| { - const seen = try sendWith(engine.engine(), c.route, c.body, .{}); - try std.testing.expectEqual(@as(?u16, 503), seen.status); - try std.testing.expect(!seen.opened); - try std.testing.expectEqual(@as(i64, 0), seen.preparing); - } + for (cases) |c| for ([_]bool{ false, true }) |streaming| { + const body = if (streaming) try std.mem.concat(arena.allocator(), u8, &.{ c.body[0 .. c.body.len - 1], ",\"stream\":true}" }) else c.body; + const seen = try sendWith(engine.engine(), c.route, body, .{}); + if (seen.status != 503 or seen.opened or seen.preparing != 0) { + wrong += 1; + std.debug.print("ADMISSION reason={s} route={s} stream={any} status={any} opened={any} events={d} preparing={d}\n", .{ @errorName(reason), @tagName(c.route), streaming, seen.status, seen.opened, seen.events, seen.preparing }); + } + }; } + try std.testing.expectEqual(@as(usize, 0), wrong); } From e8f50daaea39b21929e2f3c05e97b9944af6446b Mon Sep 17 00:00:00 2001 From: cwschroeder <1188270+cwschroeder@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:09:08 -0700 Subject: [PATCH 110/115] fix(native): report drained Flash Next hosts unavailable --- zig/src/native/flashnext_host.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/src/native/flashnext_host.zig b/zig/src/native/flashnext_host.zig index ddb3f74a68..aac4b59136 100644 --- a/zig/src/native/flashnext_host.zig +++ b/zig/src/native/flashnext_host.zig @@ -164,6 +164,7 @@ pub const Host = struct { n = 1; } out.* = .{ + .healthy = !h.closing, .running = @intFromBool(h.running != null), .waiting = @intCast(h.queued.items.len), .decode_tokens_per_second = @as(f64, @floatFromInt(tokens)) / 2.0, From 8f87428802834954826d3457f9d8edac814aefd7 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:11:42 -0700 Subject: [PATCH 111/115] fix(server): admit streams before opening and wire family readiness --- docs/recipes/adding-a-zig-family.md | 6 +++ zig/src/native/cuda.zig | 12 ++++++ zig/src/native/cuda_health_test.zig | 49 +++++++++++++++++++++++ zig/src/server/chat.zig | 3 ++ zig/src/server/compact_test.zig | 9 +++-- zig/src/server/cortheon/cortheon_test.zig | 28 ++++++++++++- zig/src/server/cortheon/loop.zig | 8 +++- zig/src/server/openai.zig | 24 ++++++++--- zig/src/server/stream_preflight_test.zig | 32 ++++++++++++++- 9 files changed, 159 insertions(+), 12 deletions(-) create mode 100644 zig/src/native/cuda_health_test.zig diff --git a/docs/recipes/adding-a-zig-family.md b/docs/recipes/adding-a-zig-family.md index 7252bd9cf7..cbc2f14d3c 100644 --- a/docs/recipes/adding-a-zig-family.md +++ b/docs/recipes/adding-a-zig-family.md @@ -194,6 +194,12 @@ On M5 tensor units. Chips before M5 need their own path. 250-500 ms, on its own queue, inside the engine process. Without it macOS unwires the set after 1-2 s idle, and the next request pays a full re-wire: 2.7 s on K3. +Native CUDA families may export `healthy(ctx: *anyopaque) bool`. The native host forwards this callback to +lane health and admission. It must read thread-safe host state without binding a GPU context, waiting on device +work or taking the inference thread's lock. Families without it retain the lane host's normal closing/stale checks. +Unavailable or busy admission returns HTTP 503 before a streaming response opens. Context-window refusals retain +their existing streamed error policy; failures after admission remain stream errors. + ## Across several Macs From Kimi K3 on four Macs: diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 254a8799a9..e444624bba 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -111,8 +111,14 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + health: ?api.HealthSource = null, recording: ?*record.Recording = null, // --record: the recorder in front of the family's backend + /// A family's readiness reads host state only: no CUDA binding or blocking device work. + fn readiness(h: *Host, comptime F: type) void { + if (@hasDecl(F, "healthy")) h.health = .{ .ctx = h.family, .read = F.healthy }; + } + fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); h.host.stop(); @@ -157,6 +163,7 @@ const Host = struct { h.host = api.LaneHost.init(h.gpa, io, &h.core, info); h.host.memory = if (h.gpu != null) .{ .read = readMemory } else null; h.host.explain = explain; + h.host.health = h.health; h.host.cache = h.store; h.host.warm_quiet = if (info.warm_turns) nemotron.native.warm_quiet else 0; if (h.lone != null) h.host.lone = .{ .ctx = h, .run = loneRun, .sampled = true, .logprobs = h.inner.vtable.first_row != null, .structures = h.inner.vtable.first_masked != null, .min_budget = nemotron.native.lone_min_budget, .resumes = true }; @@ -403,6 +410,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; + h.readiness(F); const held = loaded.stream_bytes * (streams -| free); const cache = if (@hasDecl(F, "snap_save")) budget.cacheBytes(o.prompt_cache_gib, room, held) else 0; h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ @@ -581,3 +589,7 @@ test "a cancel during a long prompt stops it, and the next request's reply is un } try std.testing.expectEqual(@as(usize, 40), restored.tokens.items.len); } + +test "family readiness reaches production lane admission without binding CUDA" { + try @import("cuda_health_test.zig").check(Host, Host.readiness, Host.serve); +} diff --git a/zig/src/native/cuda_health_test.zig b/zig/src/native/cuda_health_test.zig new file mode 100644 index 0000000000..8dbf458534 --- /dev/null +++ b/zig/src/native/cuda_health_test.zig @@ -0,0 +1,49 @@ +//! The native family's readiness source reaches the real lane host without a device. +const std = @import("std"); +const api = @import("engine_api"); +const lanes = @import("lanes"); + +pub fn check(comptime Host: type, comptime readiness: anytype, comptime serve: anytype) !void { + const gpa = std.testing.allocator; + const Family = struct { + available: std.atomic.Value(bool) = .init(true), + pub fn healthy(ctx: *anyopaque) bool { + const f: *@This() = @ptrCast(@alignCast(ctx)); + return f.available.load(.acquire); + } + fn release(_: *anyopaque) void {} + fn event(_: *anyopaque, _: api.Id, _: *const api.Event) void { + @panic("an unavailable family must refuse admission"); + } + }; + var family: Family = .{}; + var backend: lanes.fake.Fake = .{ .gpa = gpa }; + defer backend.deinit(); + const h = try gpa.create(Host); + defer gpa.destroy(h); + h.* = .{ .gpa = gpa, .gpu = null, .family = &family, .release = Family.release, .inner = backend.backend(), .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined }; + readiness(h, struct {}); // existing families need no readiness callback + try std.testing.expect(h.health == null); + readiness(h, Family); + try serve(h, std.testing.io, .{ .exact_width = 8 }, 8, .{ .lanes = 2 }, null); + defer { + h.host.stop(); + h.core.deinit(); + h.cfg.deinit(gpa); + } + const engine = h.host.engine(); + var status: api.Status = .{}; + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + family.available.store(false, .release); + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); + const request: api.Request = .{ .prompt = &.{1}, .max_tokens = 1 }; + try std.testing.expectError(error.Closed, engine.submit(1, &request, .{ .ctx = &family, .event = Family.event })); + family.available.store(true, .release); + engine.status(&status, &.{}); + try std.testing.expect(status.healthy); + engine.drain(); + engine.status(&status, &.{}); + try std.testing.expect(!status.healthy); +} diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 72ba3a8bf2..bd6e588c0b 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -41,6 +41,8 @@ pub const Input = struct { pub const Sink = struct { ctx: *anyopaque, call: *const fn (ctx: *anyopaque, delta: Value) error{Closed}!void, + /// Called only after admission, with cancellation and mailbox draining already guarded. + begin: ?*const fn (ctx: *anyopaque, prompt_len: usize) error{Closed}!void = null, }; pub const Reply = struct { @@ -204,6 +206,7 @@ pub fn generate(srv: *Server, cx: *Cx, prepared: Prepared, sink: ?Sink, gone: an defer srv.noteRequest(prepared.prompt_len, box.cached orelse 0, gen.collected.items.len, box.stats.drafted, box.stats.accepted, box.stats.rounds, received, gen.first_ns, gen.last_ns, box.stats.prefill_seconds); errdefer if (!gen.engine_done) gen.cancel(); // the engine writes to the mailbox until it says finished const result: Failure!Reply = blk: { + if (sink) |s| if (s.begin) |begin| begin(s.ctx, prepared.prompt_len) catch break :blk error.Cancelled; if (sink != null and input.tools.len > 0) gen.calls = tool_stream.Streamer.init(a, input.tools) catch |e| break :blk e; gen.loop(gone) catch |e| break :blk e; break :blk gen.finish(cx, prepared.prompt_len, received, thinking, effort, prepared.sampled, prepared.drafts); diff --git a/zig/src/server/compact_test.zig b/zig/src/server/compact_test.zig index 986fd10922..aaa3a688eb 100644 --- a/zig/src/server/compact_test.zig +++ b/zig/src/server/compact_test.zig @@ -441,6 +441,7 @@ test "an incremental request summarizes only the new middle and a user edit is t const Watch = struct { eng: *Eng, opened: bool = false, + admitted_at_open: usize = 0, saw: bool = false, fn out(w: *Watch) openai.Out { @@ -449,6 +450,7 @@ const Watch = struct { fn open(ctx: *anyopaque, _: usize) error{Closed}!void { const w: *Watch = @ptrCast(@alignCast(ctx)); + w.admitted_at_open = w.eng.prompts.items.len; w.opened = true; w.eng.opened = true; } @@ -469,7 +471,7 @@ const Watch = struct { } }; -test "a stream compacts before it opens" { +test "a stream compacts and admits its final request before it opens" { var eng: Eng = .{ .window = 2000 }; var text: Text = .{}; defer freePrompts(&eng); @@ -486,6 +488,7 @@ test "a stream compacts before it opens" { openai.run(srv, a, watch.out(), .{ .conn = &conn }, true, raw); try std.testing.expect(watch.opened); try std.testing.expect(watch.saw); - try std.testing.expect(eng.early >= 1); - try std.testing.expect(eng.late >= 1); + try std.testing.expect(watch.admitted_at_open >= 2); // compaction and the final reply both admitted + try std.testing.expectEqual(watch.admitted_at_open, eng.early); + try std.testing.expectEqual(@as(usize, 0), eng.late); } diff --git a/zig/src/server/cortheon/cortheon_test.zig b/zig/src/server/cortheon/cortheon_test.zig index 2287eaac82..15413efef1 100644 --- a/zig/src/server/cortheon/cortheon_test.zig +++ b/zig/src/server/cortheon/cortheon_test.zig @@ -63,6 +63,8 @@ const Text = struct { const Eng = struct { script: []const []const u8 = &.{}, next: usize = 0, + refuse_at: ?usize = null, + refusal: api.SubmitError = error.Closed, prompts: std.ArrayList([]u8) = .empty, fn engine(e: *Eng) api.Engine { @@ -73,6 +75,7 @@ const Eng = struct { } fn submit(ctx: *anyopaque, id: api.Id, request: *const api.Request, sink: api.Sink) api.SubmitError!void { const e: *Eng = @ptrCast(@alignCast(ctx)); + if (e.refuse_at == e.next) return e.refusal; const buf = gpa.alloc(u8, request.prompt.len) catch return error.Busy; for (request.prompt, buf) |tok, *byte| byte.* = @intCast(tok); e.prompts.append(gpa, buf) catch {}; @@ -121,17 +124,22 @@ const Web = struct { const Got = struct { events: std.ArrayList(Value) = .empty, body: ?Value = null, + opened: bool = false, + status: ?u16 = null, a: Allocator, fn out(g: *Got) openai.Out { return .{ .ctx = g, .vt = &.{ .open = open, .event = event, .reply = reply } }; } - fn open(_: *anyopaque, _: usize) error{Closed}!void {} + fn open(ctx: *anyopaque, _: usize) error{Closed}!void { + @as(*Got, @ptrCast(@alignCast(ctx))).opened = true; + } fn event(ctx: *anyopaque, payload: ?Value) error{Closed}!void { const g: *Got = @ptrCast(@alignCast(ctx)); if (payload) |p| g.events.append(g.a, json.deepCopy(g.a, p) catch return error.Closed) catch return error.Closed; } - fn reply(ctx: *anyopaque, _: u16, payload: Value) void { + fn reply(ctx: *anyopaque, status: u16, payload: Value) void { const g: *Got = @ptrCast(@alignCast(ctx)); + g.status = status; g.body = json.deepCopy(g.a, payload) catch null; } /// Every streamed delta's text under `key`, joined. @@ -245,6 +253,22 @@ test "a fresh question on current facts is researched before the model starts" { try std.testing.expectEqualStrings("The latest is 0.17.\n\nOK", message(got.body.?).strField("content").?); } +test "research progress cannot open a stream before its final request is admitted" { + for ([_]api.SubmitError{ error.Closed, error.Busy }) |reason| { + var rig: Rig = .{ .eng = .{ .script = &.{"Zig latest release notes"}, .refuse_at = 1, .refusal = reason } }; + try rig.start(true); + defer rig.stop(); + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const got = try rig.ask(arena.allocator(), "{\"stream\":true,\"messages\":[{\"role\":\"user\",\"content\":\"What is the latest release of Zig?\"}]}"); + try std.testing.expect(rig.web.searches > 0); + try std.testing.expectEqual(@as(?u16, 503), got.status); + try std.testing.expect(!got.opened); + try std.testing.expectEqual(@as(usize, 0), got.events.items.len); + try std.testing.expectEqual(@as(i64, 0), rig.srv.preparing.load(.acquire)); + } +} + test "a stream hides Cortheon's calls, shows progress as reasoning, and passes the client's calls" { var rig: Rig = .{ .eng = .{ .script = &.{ "Let me check.{\"name\": \"cortheon_fetch\", \"arguments\": {\"url\": \"https://docs.example/x\"}}", "Listing files.{\"name\": \"bash\", \"arguments\": {\"command\": \"ls\"}}" } } }; try rig.start(true); diff --git a/zig/src/server/cortheon/loop.zig b/zig/src/server/cortheon/loop.zig index 52eb5db2cc..a4fa577678 100644 --- a/zig/src/server/cortheon/loop.zig +++ b/zig/src/server/cortheon/loop.zig @@ -131,7 +131,13 @@ const Filter = struct { const Call = struct { index: i64, server: bool, client: u32 }; fn sink(f: *Filter) chat.Sink { - return .{ .ctx = f, .call = call }; + return .{ .ctx = f, .call = call, .begin = begin }; + } + + fn begin(ctx: *anyopaque, prompt_len: usize) error{Closed}!void { + const f: *Filter = @ptrCast(@alignCast(ctx)); + const inner = f.inner orelse return; + if (inner.begin) |start| try start(inner.ctx, prompt_len); } fn startRound(f: *Filter) void { diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index eca0dedd46..45cf437364 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -429,10 +429,25 @@ const Run = struct { const StreamSink = struct { run: *Run, tools: bool, + prompt_len: usize, + opened: bool = false, + pending: std.ArrayList(Value) = .empty, // research progress waits for actual engine admission + + fn begin(ctx: *anyopaque, _: usize) error{Closed}!void { + const s: *StreamSink = @ptrCast(@alignCast(ctx)); + if (s.opened) return; + const r = s.run; + try r.out.vt.open(r.out.ctx, s.prompt_len); + s.opened = true; + if (!s.tools and r.is_chat) try r.emit(r.chunk(roleDelta(r.a) catch return error.Closed, null) catch return error.Closed); + for (s.pending.items) |delta| try call(ctx, delta); + s.pending.clearRetainingCapacity(); + } fn call(ctx: *anyopaque, delta: Value) error{Closed}!void { const s: *StreamSink = @ptrCast(@alignCast(ctx)); const r = s.run; + if (!s.opened) return s.pending.append(r.a, delta) catch error.Closed; if (!s.tools) { // completions stream text only, with its logprob rows when asked if (!r.is_chat and delta != .string and (delta != .object or delta.object.get("content") == null)) return; @@ -482,7 +497,7 @@ const Run = struct { const a = r.a; const tools = r.plan.input.tools.len > 0; var cx: Cx = .{ .a = a }; - // a context-window refusal is reported inside the stream; every other refusal is still a 400 before it opens + // Context-window refusals stay streamed; other pre-admission refusals retain their HTTP status. var compaction: ?compact.Stamp = null; const prepared = if (r.srv.config.compact_at == null) chat.prepare(r.srv, &cx, r.plan.input, gone) catch |e| { if (e == error.Refused and cx.kind == .context_length) return r.streamedContext(&cx, tools, field); @@ -497,12 +512,11 @@ const Run = struct { }; var handed = false; // generate gives the preparing count back from here on defer if (!handed) chat.release(r.srv, prepared.preparing); - r.out.vt.open(r.out.ctx, prepared.prompt_len) catch return; - if (!tools and r.is_chat) r.emit(r.chunk(roleDelta(a) catch return, null) catch return) catch return; - var sink_state: StreamSink = .{ .run = r, .tools = tools }; + var sink_state: StreamSink = .{ .run = r, .tools = tools, .prompt_len = prepared.prompt_len }; handed = true; - const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call }; + const sink: chat.Sink = .{ .ctx = &sink_state, .call = StreamSink.call, .begin = StreamSink.begin }; var reply = (if (r.turn) |t| cortheon.loop.stream(t, r.srv, &cx, &r.plan.input, prepared, sink, gone) else chat.generate(r.srv, &cx, prepared, sink, gone)) catch |e| { + if (!sink_state.opened) return r.unsent(&cx, e, field); switch (e) { error.Cancelled => return, error.Refused => if (cx.kind != .other and cx.kind != .server) { diff --git a/zig/src/server/stream_preflight_test.zig b/zig/src/server/stream_preflight_test.zig index 8bf11643a0..b38f81ecc5 100644 --- a/zig/src/server/stream_preflight_test.zig +++ b/zig/src/server/stream_preflight_test.zig @@ -1,4 +1,4 @@ -//! A streamed context-window refusal is an error event, and every other refusal is a 400 before the stream opens. +//! Context-window refusals stay streamed; other refusals preserve their HTTP status before admission. const std = @import("std"); const api = @import("engine_api"); const server = @import("root.zig"); @@ -295,6 +295,36 @@ test "a stream whose client is gone when it opens gives the preparing count back try std.testing.expectEqual(@as(i64, 0), seen.preparing); } +test "disconnect at the admitted stream boundary cancels and drains its mailbox" { + const Waiting = struct { + sink: ?api.Sink = null, + submissions: usize = 0, + cancellations: usize = 0, + fn engine(e: *@This()) api.Engine { + return .{ .ctx = e, .vtable = &.{ .info = HiEngine.info, .submit = submit, .cancel = cancel, .status = HiEngine.status, .memory = HiEngine.memory } }; + } + fn submit(ctx: *anyopaque, _: api.Id, _: *const api.Request, sink: api.Sink) api.SubmitError!void { + const e: *@This() = @ptrCast(@alignCast(ctx)); + e.submissions += 1; + e.sink = sink; + } + fn cancel(ctx: *anyopaque, id: api.Id) void { + const e: *@This() = @ptrCast(@alignCast(ctx)); + e.cancellations += 1; + const sink = e.sink.?; + e.sink = null; + sink.event(sink.ctx, id, &.{ .finished = .{ .reason = .stop } }); + } + }; + var waiting: Waiting = .{}; + const seen = try sendWith(waiting.engine(), .chat, "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"stream\":true,\"max_tokens\":2}", .{ .gone_at_open = true }); + try std.testing.expectEqual(@as(usize, 1), waiting.submissions); + try std.testing.expectEqual(@as(usize, 1), waiting.cancellations); + try std.testing.expect(waiting.sink == null); + try std.testing.expect(!seen.opened); + try std.testing.expectEqual(@as(i64, 0), seen.preparing); +} + test "a reply labels a greedy request greedy and a sampled one exact" { const greedy = try send(.chat, "{\"model\":\"test-model\",\"messages\":[{\"role\":\"user\",\"content\":\"x\"}],\"max_tokens\":2,\"temperature\":0}"); try std.testing.expectEqual(@as(?u16, 200), greedy.status); From 66c9e378dbdd853beadb2822eec467a47133acab Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:14:19 -0700 Subject: [PATCH 112/115] native: keep CUDA context configuration in a focused module --- zig/src/native/cuda.zig | 31 ++----------------------------- zig/src/native/cuda_context.zig | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 29 deletions(-) create mode 100644 zig/src/native/cuda_context.zig diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index e3b7dfb806..61171f3c76 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -40,25 +40,8 @@ pub const deviceName = cuda_device.deviceName; const getenv = cuda_device.getenv; const setting = cuda_device.setting; -/// The model's window (config.json's max_position_embeddings, text_config's first), 0 when it names none. -fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { - const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return 0; - const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch return 0; - const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return 0; - if (doc != .object) return 0; - const text = if (doc.object.get("text_config")) |t| (if (t == .object) t else doc) else doc; - const limit = text.object.get("max_position_embeddings") orelse doc.object.get("max_position_embeddings") orelse return 0; - return if (limit == .integer and limit.integer > 0) limit.integer else 0; -} - -/// Prompt plus reply tokens a request may use: --context (0: the model's window), else the family default within it. -fn contextWindow(requested: ?i64, native: i64, default: i64) error{ Negative, NoNative, PastNative }!i64 { - const r = requested orelse return if (native > 0) @min(default, native) else default; - if (r < 0) return error.Negative; - if (r == 0) return if (native > 0) native else error.NoNative; - if (native > 0 and r > native) return error.PastNative; - return r; -} +const modelContext = @import("cuda_context.zig").modelContext; +const contextWindow = @import("cuda_context.zig").contextWindow; /// The pool now, read on the thread whose context is current; `problem` names a bad variable. fn pool(a: Allocator, io: std.Io, ctx: *const cuda.Context, problem: *[]const u8) !?budget.Pool { @@ -516,16 +499,6 @@ test "every registered family and d1's decision engine are listed for capabiliti try std.testing.expectEqualStrings("lfm2_vl", families[registry.len].model_type); } -test "context windows: the family default inside the model's, 0 for the model's, never past it" { - try std.testing.expectEqual(@as(i64, 16384), try contextWindow(null, 262144, 16384)); - try std.testing.expectEqual(@as(i64, 4096), try contextWindow(null, 4096, 16384)); - try std.testing.expectEqual(@as(i64, 262144), try contextWindow(0, 262144, 16384)); - try std.testing.expectEqual(@as(i64, 65536), try contextWindow(65536, 262144, 16384)); - try std.testing.expectError(error.PastNative, contextWindow(300000, 262144, 16384)); - try std.testing.expectError(error.Negative, contextWindow(-5, 262144, 16384)); - try std.testing.expectError(error.NoNative, contextWindow(0, 0, 16384)); -} - test "a cancel during a long prompt stops it, and the next request's reply is unchanged (#291)" { const gpa = std.testing.allocator; var target: lanes.fake.Fake = .{ .gpa = gpa, .prefill_chunks = 10 }; diff --git a/zig/src/native/cuda_context.zig b/zig/src/native/cuda_context.zig new file mode 100644 index 0000000000..baa74e3df9 --- /dev/null +++ b/zig/src/native/cuda_context.zig @@ -0,0 +1,33 @@ +//! Native CUDA context-window configuration and checkpoint limits. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +/// The model's window (config.json's max_position_embeddings, text_config's first), 0 when it names none. +pub fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { + const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return 0; + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch return 0; + const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return 0; + if (doc != .object) return 0; + const text = if (doc.object.get("text_config")) |t| (if (t == .object) t else doc) else doc; + const limit = text.object.get("max_position_embeddings") orelse doc.object.get("max_position_embeddings") orelse return 0; + return if (limit == .integer and limit.integer > 0) limit.integer else 0; +} + +/// Prompt plus reply tokens a request may use: --context (0: the model's window), else the family default within it. +pub fn contextWindow(requested: ?i64, native: i64, default: i64) error{ Negative, NoNative, PastNative }!i64 { + const r = requested orelse return if (native > 0) @min(default, native) else default; + if (r < 0) return error.Negative; + if (r == 0) return if (native > 0) native else error.NoNative; + if (native > 0 and r > native) return error.PastNative; + return r; +} + +test "context windows: the family default inside the model's, 0 for the model's, never past it" { + try std.testing.expectEqual(@as(i64, 16384), try contextWindow(null, 262144, 16384)); + try std.testing.expectEqual(@as(i64, 4096), try contextWindow(null, 4096, 16384)); + try std.testing.expectEqual(@as(i64, 262144), try contextWindow(0, 262144, 16384)); + try std.testing.expectEqual(@as(i64, 65536), try contextWindow(65536, 262144, 16384)); + try std.testing.expectError(error.PastNative, contextWindow(300000, 262144, 16384)); + try std.testing.expectError(error.Negative, contextWindow(-5, 262144, 16384)); + try std.testing.expectError(error.NoNative, contextWindow(0, 0, 16384)); +} From beec1aad6c967f6e32aa14a3fd21ca8358710279 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:16:38 -0700 Subject: [PATCH 113/115] test(native): retain context-window coverage in the host root --- zig/src/native/cuda.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 61171f3c76..d425d0d166 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -478,6 +478,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O const toGib = family_memory.toGib; test { + _ = @import("cuda_context.zig"); _ = budget; // cuda_memory.zig's own tests, which no other test binary reaches _ = refusals; _ = lane_backend; From ab7f7baa39661112a1cc7b190d70ee10101d8e15 Mon Sep 17 00:00:00 2001 From: balancedscorpion <175030879+balancedscorpion@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:21:50 +0000 Subject: [PATCH 114/115] test(cuda): clear experts-rows buffers on the kernels' stream Rig.zeros cleared each buffer with cuMemsetD8 on the default stream, while the kernels run on a non-blocking stream that does not wait for it. On a GB10 (sm_121) the clear of the one-row reference's output sometimes landed after the decode kernel wrote it, and experts-rows failed with an all-zero or partly zero reference: 8 of 30 runs at beec1aad6, 0 of 30 with this change, alternated on one machine. Clear on the rows' stream. --- zig/tests/cuda/experts_rows_tests.zig | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/zig/tests/cuda/experts_rows_tests.zig b/zig/tests/cuda/experts_rows_tests.zig index cb5ecd86ac..7bf59fae5d 100644 --- a/zig/tests/cuda/experts_rows_tests.zig +++ b/zig/tests/cuda/experts_rows_tests.zig @@ -31,7 +31,8 @@ const Rig = struct { fn zeros(r: *Rig, n: usize) !u64 { var b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(n, 16)); - try b.fill8(0, null); + // on the kernels' stream: it is non-blocking, so a memset on the default stream may land after them + try b.fill8(0, r.s.handle); try r.bufs.append(r.gpu.gpa, b); return b.ptr; } From b6fd0d925dbb95cef44bc25342eea469e7911080 Mon Sep 17 00:00:00 2001 From: balancedscorpion <175030879+balancedscorpion@users.noreply.github.com> Date: Sat, 17 Oct 2026 22:41:09 +0000 Subject: [PATCH 115/115] docs: measure the ModelOpt FP4-Local checkpoint on one GB10 The recipe's measurements covered only the MLX affine checkpoint. Add the ModelOpt one, measured with beec1aad6 built unpatched for sm_121: decode and concurrency from bench_openai and bench_concurrent, cold first tokens from bench_cache with exact token arrays, the served exactness checks, the seams at fixed histories, and what the 128 GB machine admits. --- docs/recipes/flash-next-cuda.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/recipes/flash-next-cuda.md b/docs/recipes/flash-next-cuda.md index 25e8a2ef89..fb0661be9b 100644 --- a/docs/recipes/flash-next-cuda.md +++ b/docs/recipes/flash-next-cuda.md @@ -124,6 +124,29 @@ candidates are close: 32 greedy tokens equal it on four of `zig/tests/cuda/flash token 15 on the 3,001-token one. Over 4,096 tokens of this repository's documentation, the next-token perplexity is 6.168 on the captured set and 6.178 on ours. +### The ModelOpt checkpoint on one GB10 + +`nvidia/Qwen3.8-Flash-Next-FP4-Local` (revision `b9ce18a6`) on a GB10 (128 GB, driver 580.178, CUDA 13.3), greedy, +`--prompt-cache-gib 0`, the tools above, three runs of the defaults at a 32,768-token window alternated with other +settings in one session: + +| | defaults | +| --- | --- | +| One stream, drafted, code / chat | 64.1-64.6 / 40.7-41.1 | +| drafted, 1 / 4 / 8 streams, greedy code | 64.1-64.4 / 44.9-45.2 / 66.2-66.8 | +| drafted, 1 / 4 / 8 streams, greedy chat | 40.6-41.1 / 35.0-35.2 / 52.4-52.8 | +| Cold first token at 2,048 / 8,192 / 32,768 / 65,536 / 131,072 tokens (s) | 1.83 / 7.11 / 29.2 / 60.1 / 131.7 | + +Every concurrent reply equalled its solo run and every drafted one its `"draft": false` run. A reply cancelled +mid-stream, or sent again after a restart, equalled a fresh one, and a second conversation turn resumed from the prompt +cache (`--no-drafts`) equalled the same turn computed without it. Prompt rows equal one-token rows bit for bit at seven +fixed histories of 20 to 6,496 tokens. Three streams draft at a time; the rest share plain rounds. + +The engine holds 73.2 GiB and keeps a 12.2 GiB reserve. A 32,768-token window then admits 7 or 8 streams of 1.21 GiB +and a 131,136-token window (the cold prompts above) 2 of 4.03 GiB; a 262,144-token stream takes 7.79 GiB, which fitted +in one start of four on this machine. After the 131,072-token prompt the process had peaked at 6.5 GB of host and +83.0 GB of device memory. Absolute speeds on this machine moved by about a tenth between sessions. + ## Not yet - The dense projections run a copy of the Python engine's group-32 lane matmul (`qmm.cu`); the core's `qmm_group` and